Replace the per-step publish with State.assign()

Per review, drop _publish_state() and the step-6 call. Actuator.step()
and the current/next double-buffer contract go back to untouched, so a
caller can again retain the pre-step state and no actuator pays a state
copy on every step.

Add Actuator.State.assign() instead, mirroring newton.State.assign():
callers that capture an odd-length region assign at its boundary in
place of the final swap. The helper walks __dict__ rather than
dataclass fields, so a controller state that allocates arrays in
__init__ is copied too, and it raises on a component or array mismatch
instead of silently skipping. An undecorated Controller.State subclass
is a dataclass with no fields, which the previous helper missed.

Document both routes: the boundary assign, and alternating one graph
per buffer orientation, which needs no Newton-side support and so works
on releases without assign(). Test both, and cover the assign contract.

Closes #4098

Author: AntoineRichard

changelog/4098.added.md

@@ -0,0 +1 @@
+Add `Actuator.State.assign()` to copy actuator state between the caller's two state objects, so a CUDA graph region holding an odd number of actuator steps can advance stateful `Delay`, `ControllerPID` and `ControllerNeuralLSTM` state by assigning at the region boundary instead of relying on the host-side swap, which a graph does not record.

changelog/4098.fixed.md

@@ -159,14 +159,16 @@ 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
 ------------------
 
 Controllers that maintain internal state (e.g. :class:`ControllerPID` with an
 integral accumulator, or :class:`ControllerNeuralLSTM` with hidden/cell state) and
 actuators with a :class:`Delay` require explicit double-buffered state
-management.  Create two state objects with :meth:`Actuator.state` and pass them
-to every step:
+management.  Create two state objects with :meth:`Actuator.state` and swap them
+after each step:
 
 .. testcode:: actuator-usage
 
@@ -178,18 +180,52 @@ to every step:
    for step in range(3):
        control.joint_f.zero_()  # zero output before stepping actuators
        model.actuators[0].step(state, control, state_0, state_1, dt=0.01)
+       state_0, state_1 = state_1, state_0
+
+That swap is a host-side rebinding of two Python names.  A CUDA graph records
+buffer addresses instead, so a captured region holding an **odd** number of
+actuator steps replays with its two state objects the wrong way round: the last
+step's update is discarded on every replay, and a region holding a single step
+never advances state at all.  An even-length region is unaffected, because its
+swaps cancel inside the graph.
+
+Either pattern below keeps an odd-length region correct.  Assign at the region
+boundary, in place of its final swap, to keep one graph at the cost of one state
+copy per replay:
+
+.. code-block:: python
+
+   steps = 3  # odd, so the region needs a boundary assign
+   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
+
+:meth:`Actuator.State.assign` mirrors :meth:`newton.State.assign`, which solves
+the same problem for an odd number of solver substeps.
+
+Or capture one graph per buffer orientation and alternate them.  This copies no
+state and needs at most two graphs, since the orientation returns to its start
+after two odd-length regions:
+
+.. code-block:: python
 
-:meth:`Actuator.step` writes the step's update into *state_1* and then publishes
-it back over *state_0* with a device copy, so the two objects hold the same
-advanced state when the step returns.  Both the exchange and the update are
-device operations, which is what makes a stateful actuator behave the same under
-CUDA graph capture as it does eagerly, whatever number of steps the captured
-region holds.
+   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)
 
-Earlier releases advanced state only through a host-side
-``state_0, state_1 = state_1, state_0`` swap after each step.  That swap is no
-longer needed; it remains correct if kept, since both objects now hold the same
-state.
+The second pattern needs no Newton-side support, so it also works on releases
+without :meth:`Actuator.State.assign`.
 
 Stateless actuators (e.g. a plain PD controller without delay) do not require
 state objects — simply omit them:
@@ -348,10 +384,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.
 
-A graphable actuator is graphable for any number of steps in the captured
-region, including one.  Internal state — a :class:`Delay` buffer, a PID integral,
-LSTM hidden and cell state — is both advanced and exchanged on the device, so
-each replay carries it forward exactly as an eager loop of the same length does.
+: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.  ``ControllerNeuralLSTM`` on a Torch checkpoint is a further exception —
+it keeps hidden and cell state in Torch tensors that the controller rebinds on
+the host, so that state never advances inside a graph at all.
 
 Available Components
 --------------------

docs/concepts/actuators.rst

@@ -3,7 +3,7 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass, fields, is_dataclass
+from dataclasses import dataclass
 from typing import Any
 
 import numpy as np
@@ -32,30 +32,43 @@ def _scatter_add_kernel(
         computed_output[idx] = computed_output[idx] + computed_forces[i]
 
 
-def _publish_state(current: Any, advanced: Any) -> None:
-    """Copy an advanced state object back over the state the next step reads.
+def _assign_component_state(dst: Any, src: Any, name: str) -> None:
+    """Copy one component state (delay or controller) from *src* into *dst*.
 
-    Warp arrays are copied with :func:`warp.copy`, so the exchange is a device
-    operation that a CUDA graph records and replays.  Fields that are not Warp
-    arrays — the Torch tensors :class:`ControllerNeuralLSTM.State` holds for a
-    Torch checkpoint — are rebound instead, which is how the controller itself
-    publishes them; that path is host-side and not graphable either way.
+    Walks both objects' ``__dict__`` rather than their dataclass fields, so a
+    state that allocates its arrays in ``__init__`` is covered too.  Warp
+    arrays are copied in place, which is what a captured region records.
+    Values that are not Warp arrays — the Torch tensors
+    :class:`ControllerNeuralLSTM.State` holds for a Torch checkpoint — are
+    rebound, matching how that controller publishes them; a Torch checkpoint is
+    not graphable either way.
 
     Args:
-        current: State the next step reads from.
-        advanced: State the step just wrote.
+        dst: Component state to copy into.
+        src: Component state to copy from.
+        name: Component name, used in error messages.
+
+    Raises:
+        ValueError: The two states do not hold the same values.
     """
-    for field in fields(advanced):
-        src = getattr(advanced, field.name)
-        if src is None:
+    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.")
+    for attr in set(dst.__dict__) | set(src.__dict__):
+        val_dst = getattr(dst, attr, None)
+        val_src = getattr(src, attr, None)
+        if val_dst is None and val_src is None:
             continue
-        dst = getattr(current, field.name)
-        if isinstance(src, wp.array):
-            wp.copy(dst, src)
-        elif is_dataclass(src):
-            _publish_state(dst, src)
+        if val_dst is None or val_src is None:
+            raise ValueError(f"Cannot assign '{name}.{attr}': present in one state and missing in the other.")
+        array_dst = isinstance(val_dst, wp.array)
+        if array_dst != isinstance(val_src, wp.array):
+            raise ValueError(f"Cannot assign '{name}.{attr}': a Warp array in one state and not in the other.")
+        if array_dst:
+            val_dst.assign(val_src)
         else:
-            setattr(current, field.name, src)
+            setattr(dst, attr, val_src)
 
 
 class Actuator:
@@ -108,6 +121,33 @@ def reset(self, mask: wp.array[wp.bool] | None = None) -> None:
             if self.controller_state is not None:
                 self.controller_state.reset(mask)
 
+        def assign(self, other: Actuator.State) -> None:
+            """Copy the state held by *other* into this one.
+
+            Mirrors :meth:`newton.State.assign`.  A CUDA graph records buffer
+            addresses, not the caller's Python names, so a captured region
+            holding an odd number of actuator steps leaves its two state
+            objects the wrong way round for the next replay.  Assigning at the
+            boundary, in place of that region's final swap, costs one copy per
+            replay and keeps a single graph correct::
+
+                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.
+            """
+            _assign_component_state(self.delay_state, other.delay_state, "delay_state")
+            _assign_component_state(self.controller_state, other.controller_state, "controller_state")
+
     def __init__(
         self,
         indices: wp.array[wp.uint32],
@@ -300,14 +340,6 @@ def step(
            (e.g. ``control.joint_f.zero_()``) before looping over actuators.
         5. **State updates** — controller state update, then delay
            buffer write (push current targets into ``next_state``).
-        6. **State publish** — copy the advanced state back over
-           ``current_act_state``, so both state objects hold it.
-
-        Step 6 makes the state exchange a device operation instead of a
-        host-side rebinding of the two state objects.  A captured region
-        therefore advances state on every replay, whatever number of steps it
-        holds; the ``state_0, state_1 = state_1, state_0`` swap of earlier
-        releases is no longer required and stays correct if kept.
 
         Args:
             sim_state: Simulation state with position/velocity arrays.
@@ -400,7 +432,3 @@ def step(
                 current_act_state.delay_state,
                 next_act_state.delay_state,
             )
-
-        # --- 6. Publish the advanced state on-device ---
-        if self.is_stateful() and next_act_state is not current_act_state:
-            _publish_state(current_act_state, next_act_state)

newton/_src/actuators/actuator.py

@@ -4090,6 +4090,32 @@ def _loop(
             )
 
 
+class TestActuatorStateAssign(unittest.TestCase):
+    """Cover the :meth:`Actuator.State.assign` copy contract."""
+
+    def test_assign_copies_arrays_allocated_outside_dataclass_fields(self):
+        """Copy state a controller allocates in ``__init__`` rather than as a field."""
+
+        class _CustomState(Controller.State):
+            """Undecorated subclass: ``dataclasses.fields()`` is empty for it."""
+
+            def __init__(self, num_actuators: int, device: wp.Device):
+                self.extra = wp.zeros(num_actuators, dtype=wp.float32, device=device)
+
+        device = wp.get_device()
+        current, advanced = _CustomState(1, device), _CustomState(1, device)
+        advanced.extra.fill_(1.0)
+
+        Actuator.State(controller_state=current).assign(Actuator.State(controller_state=advanced))
+        self.assertEqual(current.extra.numpy()[0], 1.0)
+
+    def test_assign_rejects_mismatched_components(self):
+        """Raise rather than silently skip when only one state holds a component."""
+        state = ControllerPID.State(integral=wp.zeros(1, dtype=wp.float32, device=wp.get_device()))
+        with self.assertRaisesRegex(ValueError, "controller_state"):
+            Actuator.State(controller_state=state).assign(Actuator.State())
+
+
 @unittest.skipUnless(
     wp.get_device().is_cuda and wp.is_mempool_enabled(wp.get_device()),
     "CUDA graph capture requires CUDA device with memory pools",
@@ -4099,23 +4125,25 @@ class TestControllerStateGraphCapture(unittest.TestCase):
 
     A PID with only ``ki`` set, held at a constant position error by a model no
     solver moves, accumulates exactly ``ki * error * dt`` per actuator step
-    however the loop is chunked.  While only the caller's host-side swap
-    advanced state, an odd captured count above one discarded the last step's
-    update and a single captured step never advanced at all.
+    however the loop is chunked.  A graph cannot re-point its state buffers
+    between replays, so an odd-length captured region needs either a boundary
+    :meth:`Actuator.State.assign` or a second graph of the opposite parity.
     """
 
     DT = 0.01
     KI = 1.0
     TARGET = 1.0
     STEPS = 12
 
-    def _integral_after_all_steps(self, implicit: bool, steps_per_graph: int | None) -> float:
-        """Run :attr:`STEPS` actuator steps, eagerly or as replays of a captured region.
+    def _integral_after_all_steps(self, implicit: bool, steps_per_graph: int | None, parity_graphs: bool) -> float:
+        """Run :attr:`STEPS` actuator steps and return the resulting PID integral.
 
         Args:
             implicit: Solve the control law implicitly instead of explicitly.
                 Both effort modes advance the integral through the same state.
             steps_per_graph: Steps per captured region, or ``None`` to run eagerly.
+            parity_graphs: Capture one graph per state-buffer orientation and
+                alternate them, instead of assigning at the region boundary.
         """
         device = wp.get_device()
         builder = newton.ModelBuilder(gravity=(0.0, 0.0, 0.0))
@@ -4134,14 +4162,17 @@ def _integral_after_all_steps(self, implicit: bool, steps_per_graph: int | None)
         control.joint_target_q.fill_(self.TARGET)  # joint_q stays 0, so the error is constant
         s0, s1 = actuator.state(), actuator.state()
 
-        def run(s0, s1, steps):
-            """The documented stateful loop, host-side swap included."""
-            for _ in range(steps):
+        def run(s0, s1, steps, boundary_assign=False):
+            """Step the actuator, swapping state as the documented loop does."""
+            for i in range(steps):
                 control.joint_f.zero_()
                 if oracle is not None:
                     oracle.refresh(state)
                 actuator.step(state, control, s0, s1, dt=self.DT)
-                s0, s1 = s1, s0
+                if boundary_assign and steps % 2 == 1 and i == steps - 1:
+                    s0.assign(s1)  # keeps a single odd-length graph correct
+                else:
+                    s0, s1 = s1, s0
             return s0, s1
 
         if steps_per_graph is None:
@@ -4151,24 +4182,35 @@ def run(s0, s1, steps):
             s0, s1 = run(s0, s1, 1)
             s0.controller_state.integral.zero_()
             s1.controller_state.integral.zero_()
-            with wp.ScopedCapture(device) as capture:
-                s0, s1 = run(s0, s1, steps_per_graph)
+            graphs = {}
             for _ in range(self.STEPS // steps_per_graph):
-                wp.capture_launch(capture.graph)
-        wp.synchronize_device(device)
+                # Keying on the current buffer builds one graph per orientation.
+                key = id(s0) if parity_graphs else None
+                if key not in graphs:
+                    with wp.ScopedCapture(device) as capture:
+                        after = run(s0, s1, steps_per_graph, boundary_assign=not parity_graphs)
+                    graphs[key] = (capture.graph, after)
+                graph, (s0, s1) = graphs[key]
+                wp.capture_launch(graph)
+            self.assertLessEqual(len(graphs), 2, msg="alternating parity needs at most two graphs")
         return float(s0.controller_state.integral.numpy()[0])
 
-    def _assert_integral_matches(self, implicit: bool) -> None:
+    def test_pid_integral_advances_per_replay_boundary_assign(self):
+        """Assign at an odd-length region's boundary and match the eager integral."""
         expected = self.KI * self.TARGET * self.DT * self.STEPS
-        for steps_per_graph in (None, 1, 2, 3):
+        for implicit in (False, True):
+            for steps_per_graph in (None, 1, 2, 3):
+                with self.subTest(implicit=implicit, steps_per_graph=steps_per_graph):
+                    got = self._integral_after_all_steps(implicit, steps_per_graph, parity_graphs=False)
+                    self.assertAlmostEqual(got, expected, places=6)
+
+    def test_pid_integral_advances_per_replay_parity_graphs(self):
+        """Alternate one graph per buffer orientation and match the eager integral."""
+        expected = self.KI * self.TARGET * self.DT * self.STEPS
+        for steps_per_graph in (1, 2, 3):
             with self.subTest(steps_per_graph=steps_per_graph):
-                self.assertAlmostEqual(self._integral_after_all_steps(implicit, steps_per_graph), expected, places=6)
-
-    def test_pid_integral_advances_per_replay_explicit(self):
-        self._assert_integral_matches(implicit=False)
-
-    def test_pid_integral_advances_per_replay_implicit(self):
-        self._assert_integral_matches(implicit=True)
+                got = self._integral_after_all_steps(False, steps_per_graph, parity_graphs=True)
+                self.assertAlmostEqual(got, expected, places=6)
 
 
 # ---------------------------------------------------------------------------