Backport Actuator.State.assign() to release-1.6 (#4154)

Author: jcarius-nv

changelog/4098.added.md

@@ -0,0 +1 @@
+Add `Actuator.State.assign()` for preserving actuator state across odd-length CUDA graph replays with a single captured graph.

docs/concepts/actuators.rst

@@ -169,6 +169,8 @@ reusable from a custom simulator or test harness:
    actuator.step(sim_state, sim_control, state_a, state_b, dt=0.01)
 
 
+.. _stateful-actuators:
+
 Stateful Actuators
 ------------------
 
@@ -190,6 +192,51 @@ after each step:
        model.actuators[0].step(state, control, state_0, state_1, dt=0.01)
        state_0, state_1 = state_1, state_0
 
+The swap rebinds Python names, while a CUDA graph records fixed buffer
+addresses. A captured region with an odd number of actuator steps therefore
+cannot carry its final state into the next replay through the ordinary swap. A
+single-step region updates its destination buffer, but does not advance state
+across replays. Even-length regions end with the original buffer orientation
+and need no special handling.
+
+For an odd-length region, choose between assigning the state at the region
+boundary or alternating graphs for the two buffer orientations.
+
+Boundary assignment
+^^^^^^^^^^^^^^^^^^^
+
+Call :meth:`Actuator.State.assign` in place of the final swap. This uses one
+graph and copies the actuator state once at the boundary of each replay.
+
+.. code-block:: python
+
+   with wp.ScopedCapture() as capture:
+       for i in range(steps):
+           control.joint_f.zero_()
+           model.actuators[0].step(state, control, state_0, state_1, dt=0.01)
+           if steps % 2 == 1 and i == steps - 1:
+               state_0.assign(state_1)
+           else:
+               state_0, state_1 = state_1, state_0
+
+Alternating graphs
+^^^^^^^^^^^^^^^^^^
+
+Key captured graphs by the current state buffer and alternate between them.
+This avoids the boundary copy and keeps at most two graphs.
+
+.. code-block:: python
+
+   graphs = {}
+   for _ in range(replays):
+       key = id(state_0)  # one entry per buffer orientation
+       if key not in graphs:
+           with wp.ScopedCapture() as capture:
+               after = run_region(state_0, state_1)  # ordinary swapping inside
+           graphs[key] = (capture.graph, after)
+       graph, (state_0, state_1) = graphs[key]
+       wp.capture_launch(graph)
+
 Stateless actuators (e.g. a plain PD drive without delay) do not require
 state objects — simply omit them:
 
@@ -348,6 +395,12 @@ backend: ONNX checkpoints are graphable, while Torch checkpoints are not due
 to framework interop overhead.  :meth:`Actuator.is_graphable` returns ``True``
 when all components can be captured in a CUDA graph.
 
+:meth:`Actuator.is_graphable` describes the components, not the captured region.
+A stateful actuator also needs the region's state exchange to be graph-safe.  See
+:ref:`stateful-actuators` for the two patterns that keep an odd-length region
+correct. Torch-backed neural drives are not graphable and cannot be used in a
+captured region.
+
 Available Components
 --------------------
 
@@ -425,6 +478,12 @@ For example, a custom drive needs to implement
 :meth:`~newton.ModelBuilder.add_actuator` or USD schemas) to constructor
 parameters, filling in defaults where needed.
 
+A stateful custom drive also defines a dataclass subclass of
+:class:`DriveBase.State` and implements :meth:`~DriveBase.State.reset`. The
+default :meth:`Actuator.State.assign` behavior copies direct Warp array and
+Torch tensor fields without replacing their storage. States with other field
+types or nested storage implement ``assign()`` to define that copy.
+
 A custom drive works in the explicit mode with the methods above. To also
 support the implicit mode it provides three more things, because the solve
 evaluates the control law inside its own kernel rather than calling

newton/_src/actuators/actuator.py

@@ -4,7 +4,7 @@
 from __future__ import annotations
 
 import warnings
-from dataclasses import dataclass
+from dataclasses import dataclass, fields
 from typing import Any
 
 import numpy as np
@@ -42,6 +42,78 @@ def _scatter_add_kernel(
         computed_output[idx] = computed_output[idx] + computed_forces[i]
 
 
+def _assign_state_value(dst: Any, src: Any, name: str) -> None:
+    """Copy a supported state value without replacing its storage."""
+    if dst is None and src is None:
+        return
+    if dst is None or src is None:
+        raise ValueError(f"Cannot assign '{name}': present in one state and missing in the other.")
+
+    dst_is_warp = isinstance(dst, wp.array)
+    src_is_warp = isinstance(src, wp.array)
+    if dst_is_warp or src_is_warp:
+        if not (dst_is_warp and src_is_warp):
+            raise ValueError(f"Cannot assign '{name}': a Warp array in one state and not in the other.")
+        dst.assign(src)
+        return
+
+    dst_is_torch = type(dst).__module__.startswith("torch")
+    src_is_torch = type(src).__module__.startswith("torch")
+    if dst_is_torch or src_is_torch:
+        if not (dst_is_torch and src_is_torch):
+            raise ValueError(f"Cannot assign '{name}': a Torch tensor in one state and not in the other.")
+        if dst.shape != src.shape:
+            raise ValueError(f"Cannot assign '{name}': tensor shapes differ ({dst.shape} and {src.shape}).")
+        import torch
+
+        with torch.inference_mode():
+            dst.copy_(src)
+        return
+
+    raise ValueError(f"Cannot assign '{name}': expected Warp arrays or Torch tensors.")
+
+
+def _assign_component_state(dst: Any, src: Any, name: str) -> None:
+    """Copy one actuator component state from *src* into *dst*.
+
+    Args:
+        dst: Component state to copy into.
+        src: Component state to copy from.
+        name: Component name, used in error messages.
+
+    Raises:
+        ValueError: The two actuator states have incompatible components or
+            fields.
+        NotImplementedError: A custom state is not a dataclass and does not
+            implement ``assign()``.
+    """
+    if dst is None and src is None:
+        return
+    if dst is None or src is None:
+        raise ValueError(f"Cannot assign '{name}': one state has it allocated and the other does not.")
+    if type(dst) is not type(src):
+        raise ValueError(f"Cannot assign '{name}': state types differ ({type(dst).__name__} and {type(src).__name__}).")
+
+    custom_assign = getattr(dst, "assign", None)
+    if custom_assign is not None:
+        custom_assign(src)
+        return
+
+    if "__dataclass_fields__" not in type(dst).__dict__:
+        raise NotImplementedError(f"{type(dst).__qualname__} must be decorated with @dataclass or implement assign")
+
+    state_fields = fields(dst)
+    field_names = {field.name for field in state_fields}
+    attributes = set(getattr(dst, "__dict__", ())) | set(getattr(src, "__dict__", ()))
+    undeclared = attributes - field_names
+    if undeclared:
+        names = ", ".join(sorted(undeclared))
+        raise ValueError(f"Cannot assign '{name}': undeclared state attributes: {names}.")
+
+    for field in state_fields:
+        _assign_state_value(getattr(dst, field.name), getattr(src, field.name), f"{name}.{field.name}")
+
+
 class Actuator:
     """Composed actuator: delay → drive → clamping.
 
@@ -130,6 +202,33 @@ def reset(self, mask: wp.array[wp.bool] | None = None) -> None:
             if self.drive_state is not None:
                 self.drive_state.reset(mask)
 
+        def assign(self, other: Actuator.State) -> None:
+            """Copy the state held by *other* into this one.
+
+            A CUDA graph records buffer addresses rather than the caller's
+            Python names. Assigning at the boundary of an odd-length captured
+            region, in place of its final state swap, preserves the advanced
+            state for the next replay::
+
+                for i in range(steps):
+                    control.joint_f.zero_()
+                    actuator.step(state, control, state_0, state_1, dt=0.01)
+                    if steps % 2 == 1 and i == steps - 1:
+                        state_0.assign(state_1)
+                    else:
+                        state_0, state_1 = state_1, state_0
+
+            Args:
+                other: State to copy from.
+
+            Raises:
+                ValueError: The two states do not hold the same components.
+                NotImplementedError: A custom state does not implement
+                    assignment.
+            """
+            _assign_component_state(self.delay_state, other.delay_state, "delay_state")
+            _assign_component_state(self.drive_state, other.drive_state, "drive_state")
+
     def __init__(
         self,
         indices: wp.array[wp.uint32],