From ce4db9c4c08e711c184065d5324b784b30cdde12 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 02:03:58 -0700
Subject: [PATCH 01/48] Reset only masked environments in Newton and OVPhysX
 assets

A mask-only Articulation.reset called ActuatorCollection.reset(None), which
cleared delay buffers, network history, and Newton-native actuator state for
every environment. Newton rigid objects and collections dropped the mask
before resetting their wrench composers, clearing every environment's
persistent wrench.

ActuatorCollection.reset now takes env_mask. Isaac Lab actuator models still
reset by index, so the collection converts the mask to indices once for all
groups, and skips the conversion when no group carries reset state. Native
actuator state resets from the mask through a new
ActuatorControl.reset_native_actuators_mask hook; Newton and OVPhysX pass the
mask to the Newton actuator adapter, and the base default converts to indices
so backends that only implement the index hook stay correct.
---
 .../jichuanh-warp-scene-mask-native.added.rst |   3 +
 .../jichuanh-warp-scene-mask-native.minor     |   0
 .../isaaclab/actuators/actuator_collection.py |  25 +++-
 .../isaaclab/actuators/actuator_control.py    |  12 ++
 .../isaaclab/actuators/newton/adapter.py      |  39 +++--
 .../actuators/newton/physx_runtime.py         |  11 +-
 .../test/utils/actuator_equivalence.py        | 137 ++++++++++--------
 .../jichuanh-warp-scene-mask-native.fixed.rst |   5 +
 .../assets/articulation/actuator_control.py   |   4 +
 .../assets/articulation/articulation.py       |   2 +-
 .../assets/rigid_object/rigid_object.py       |   4 +-
 .../rigid_object_collection.py                |   4 +-
 .../assets/test_newton_actuators_newton.py    |   2 +
 .../test/assets/test_rigid_object.py          |   8 +
 .../assets/test_rigid_object_collection.py    |   8 +
 .../jichuanh-warp-scene-mask-native.fixed.rst |   2 +
 .../assets/articulation/actuator_control.py   |   4 +
 .../assets/articulation/articulation.py       |   2 +-
 .../test/assets/test_articulation.py          |  11 ++
 19 files changed, 200 insertions(+), 83 deletions(-)
 create mode 100644 source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst
 create mode 100644 source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.minor
 create mode 100644 source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
 create mode 100644 source/isaaclab_ov/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst

diff --git a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst
new file mode 100644
index 000000000000..7b04995e3abd
--- /dev/null
+++ b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst
@@ -0,0 +1,3 @@
+* Added ``env_mask`` to :meth:`~isaaclab.actuators.ActuatorCollection.reset`. Isaac Lab actuator models reset only
+  the masked environments, converting the mask to indices once for all groups, and backend-native actuator state
+  resets from the mask through :meth:`~isaaclab.actuators.ActuatorControl.reset_native_actuators_mask`.
diff --git a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.minor b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.minor
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/source/isaaclab/isaaclab/actuators/actuator_collection.py b/source/isaaclab/isaaclab/actuators/actuator_collection.py
index 59cb7797728c..c27d2aa5fce4 100644
--- a/source/isaaclab/isaaclab/actuators/actuator_collection.py
+++ b/source/isaaclab/isaaclab/actuators/actuator_collection.py
@@ -102,6 +102,13 @@ def __init__(
             for actuator_name in self._native_group_names:
                 self._groups[actuator_name] = self._resolve_newton_group_actuators(actuator_name)
         self._validate_coverage()
+        # Groups a masked reset converts the mask to indices for. Plain implicit groups hold no reset
+        # state, and Newton-executed groups reset through the backend.
+        self._mask_reset_actuators = tuple(
+            actuator
+            for actuator in self._groups.values()
+            if isinstance(actuator, ActuatorBase) and type(actuator) is not ImplicitActuator
+        )
         self._build_execution_plan()
         if self._debug_value_resolution:
             self._print_value_resolution_table()
@@ -180,12 +187,28 @@ def applied_effort(self) -> ProxyArray:
 
     # Lifecycle.
 
-    def reset(self, env_ids: Sequence[int] | slice | None = None) -> None:
+    def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: wp.array | None = None) -> None:
         """Reset all actuator group states.
 
+        Isaac Lab actuator models reset by environment index, so a mask is converted to indices
+        once for all of them, and only when a model with reset state is configured; plain
+        :class:`~isaaclab.actuators.ImplicitActuator` groups are skipped. Backend-native actuator
+        state resets from the mask directly.
+
+        .. caution::
+            If both ``env_ids`` and ``env_mask`` are provided, ``env_mask`` takes precedence.
+
         Args:
             env_ids: Environment indices to reset. Defaults to all environments.
+            env_mask: Environment mask. Shape is (num_instances,). Defaults to None.
         """
+        if env_mask is not None:
+            if self._mask_reset_actuators:
+                group_env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1)
+                for actuator in self._mask_reset_actuators:
+                    actuator.reset(group_env_ids)
+            self._control.reset_native_actuators_mask(env_mask)
+            return
         group_env_ids = self._control._normalize_index_sequence(env_ids)
         for actuator in self._groups.values():
             # Newton-executed groups are reset through the backend below.
diff --git a/source/isaaclab/isaaclab/actuators/actuator_control.py b/source/isaaclab/isaaclab/actuators/actuator_control.py
index 2ed1424e8bc2..a1793c02505b 100644
--- a/source/isaaclab/isaaclab/actuators/actuator_control.py
+++ b/source/isaaclab/isaaclab/actuators/actuator_control.py
@@ -315,6 +315,18 @@ def reset_native_actuators(self, env_ids: Sequence[int] | slice) -> None:
             env_ids: Environments to reset.
         """
 
+    def reset_native_actuators_mask(self, env_mask: wp.array) -> None:
+        """Reset backend-native actuator state for the environments selected by a mask.
+
+        The default converts the mask to indices for :meth:`reset_native_actuators`. Backends whose
+        native state resets from a mask override it.
+
+        Args:
+            env_mask: Environments to reset. Shape is (num_instances,).
+        """
+        if self.native_actuator_path_active:
+            self.reset_native_actuators(wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1))
+
 
 class ArticulationActuatorControl(ActuatorControl):
     """Shared control adapter for articulation-backed actuator collections.
diff --git a/source/isaaclab/isaaclab/actuators/newton/adapter.py b/source/isaaclab/isaaclab/actuators/newton/adapter.py
index 543a19135ce9..6b4a58e2f66f 100644
--- a/source/isaaclab/isaaclab/actuators/newton/adapter.py
+++ b/source/isaaclab/isaaclab/actuators/newton/adapter.py
@@ -171,7 +171,11 @@ def _swap_state_buffers(self) -> None:
         """Advance the actuator state ping-pong after an eager step or graph replay."""
         self._states_a, self._states_b = self._states_b, self._states_a
 
-    def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> None:
+    def reset(
+        self,
+        env_ids: Sequence[int] | torch.Tensor | slice | None = None,
+        env_mask: wp.array | None = None,
+    ) -> None:
         """Reset actuator states for the given environments.
 
         Args:
@@ -179,6 +183,8 @@ def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) ->
                 ``slice(None)``, which IsaacLab callers sometimes pass)
                 resets all environments. Otherwise expects a torch tensor
                 or sequence of int indices.
+            env_mask: Boolean environment mask with shape ``(num_envs,)``.
+                Takes precedence over :paramref:`env_ids`.
 
         Newton's :meth:`Actuator.State.reset` expects a per-DOF boolean
         mask of length ``num_actuators`` (= ``num_envs * dofs_per_actuator``),
@@ -187,7 +193,7 @@ def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) ->
         etc.). We therefore build a per-actuator per-DOF mask from the
         env mask before delegating to each state.
         """
-        if env_ids is None or env_ids == slice(None):
+        if env_mask is None and (env_ids is None or env_ids == slice(None)):
             for sa, sb in zip(self._states_a, self._states_b):
                 if sa is not None:
                     sa.reset(None)
@@ -195,20 +201,21 @@ def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) ->
                     sb.reset(None)
             return
 
-        num_envs = len(range(self._num_envs)[env_ids]) if isinstance(env_ids, slice) else len(env_ids)
-        if num_envs == 0:
-            return
-        env_mask = wp.zeros(self._num_envs, dtype=wp.bool, device=self._device)
-        if isinstance(env_ids, slice):
-            wp.to_torch(env_mask)[env_ids] = True
-        else:
-            # Torch indexed scalar assignment uploads True and synchronizes on CUDA.
-            indices = (
-                wp.from_torch(env_ids)
-                if isinstance(env_ids, torch.Tensor)
-                else wp.array(env_ids, dtype=wp.int32, device=self._device)
-            )
-            wp.launch(set_mask_kernel, dim=num_envs, inputs=[env_mask, indices], device=self._device)
+        if env_mask is None:
+            num_envs = len(range(self._num_envs)[env_ids]) if isinstance(env_ids, slice) else len(env_ids)
+            if num_envs == 0:
+                return
+            env_mask = wp.zeros(self._num_envs, dtype=wp.bool, device=self._device)
+            if isinstance(env_ids, slice):
+                wp.to_torch(env_mask)[env_ids] = True
+            else:
+                # Torch indexed scalar assignment uploads True and synchronizes on CUDA.
+                indices = (
+                    wp.from_torch(env_ids)
+                    if isinstance(env_ids, torch.Tensor)
+                    else wp.array(env_ids, dtype=wp.int32, device=self._device)
+                )
+                wp.launch(set_mask_kernel, dim=num_envs, inputs=[env_mask, indices], device=self._device)
 
         for act, sa, sb in zip(self.actuators, self._states_a, self._states_b):
             per_dof_mask = wp.zeros(act.indices.shape[0], dtype=wp.bool, device=self._device)
diff --git a/source/isaaclab/isaaclab/actuators/newton/physx_runtime.py b/source/isaaclab/isaaclab/actuators/newton/physx_runtime.py
index 76bc8a8d1885..a8455ef4c3b1 100644
--- a/source/isaaclab/isaaclab/actuators/newton/physx_runtime.py
+++ b/source/isaaclab/isaaclab/actuators/newton/physx_runtime.py
@@ -125,10 +125,15 @@ def compute(self, collection: ActuatorCollection, dt: float) -> None:
                     return
         self._run_native_actuator_kernels(collection, dt)
 
-    def reset(self, env_ids: Sequence[int] | slice) -> None:
-        """Reset selected native actuator state."""
+    def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: wp.array | None = None) -> None:
+        """Reset selected native actuator state.
+
+        Args:
+            env_ids: Environments to reset. Defaults to all environments.
+            env_mask: Environment mask. Takes precedence over ``env_ids``. Shape is (num_instances,).
+        """
         if self.adapter is not None:
-            self.adapter.reset(env_ids)
+            self.adapter.reset(env_ids, env_mask)
 
     def _run_native_actuator_kernels(self, collection: ActuatorCollection, dt: float) -> None:
         from . import kernels as actuator_kernels  # noqa: PLC0415
diff --git a/source/isaaclab/isaaclab/test/utils/actuator_equivalence.py b/source/isaaclab/isaaclab/test/utils/actuator_equivalence.py
index 7c83b950123d..a03c6b7fa24a 100644
--- a/source/isaaclab/isaaclab/test/utils/actuator_equivalence.py
+++ b/source/isaaclab/isaaclab/test/utils/actuator_equivalence.py
@@ -326,11 +326,12 @@ class ActuatorStateResetBase:
     Inspects ``adapter.actuators[i].state.delay_state.num_pushes`` directly:
 
     * After warmup, ``num_pushes > 0`` for every DOF (buffer was populated).
-    * After ``articulation.reset(env_ids=[0])``, the entries for env 0's DOFs
+    * After a partial reset of env 0, the entries for env 0's DOFs
       must be ``0`` and the entries for env 1's DOFs must remain ``> 0``.
 
-    Done independently on Lab and Newton paths. Backend-specific twins mix
-    this into a ``unittest.TestCase`` and provide :meth:`_make_sim_cfg`,
+    Done independently on Lab and Newton paths, once per selector in
+    :attr:`RESET_SELECTORS`. Backend-specific twins mix this into a
+    ``unittest.TestCase`` and provide :meth:`_make_sim_cfg`,
     :meth:`_make_articulation`, and :meth:`_get_adapter`.
     """
 
@@ -340,6 +341,8 @@ class ActuatorStateResetBase:
     DT: float = 1.0 / 120.0
     TARGET_OFFSET: float = 0.1  # [rad] added to initial joint positions
     RESET_WARMUP_STEPS: int = 3
+    RESET_SELECTORS: tuple[str, ...] = ("env_ids",)
+    """Partial-reset selectors to exercise: ``"env_ids"`` (index tensor) and ``"env_mask"`` (Warp mask)."""
 
     def _make_sim_cfg(self, use_newton_actuators: bool) -> SimulationCfg:
         """Return the backend simulation config for the run."""
@@ -372,57 +375,74 @@ def _build_and_warm(self, *, use_newton_actuators: bool):
         target_vel = torch.zeros_like(init_pos)
         articulation.set_joint_position_target_index(target=target_pos)
         articulation.set_joint_velocity_target_index(target=target_vel)
+        self._warm(sim, articulation)
+        return ctx, sim, articulation
+
+    def _warm(self, sim, articulation) -> None:
         for _ in range(self.RESET_WARMUP_STEPS):
             articulation.write_data_to_sim()
             sim.step()
             articulation.update(self.DT)
-        return ctx, sim, articulation
+
+    def _reset_env(self, articulation, selector: str) -> None:
+        """Reset :attr:`RESET_ENV` alone through the given selector."""
+        if selector == "env_ids":
+            articulation.reset(env_ids=torch.tensor([self.RESET_ENV], device=articulation.device, dtype=torch.long))
+        else:
+            env_mask = torch.zeros(self.NUM_ENVS, dtype=torch.bool, device=articulation.device)
+            env_mask[self.RESET_ENV] = True
+            articulation.reset(env_mask=wp.from_torch(env_mask))
 
     def test_newton_state_reset_isolated_to_reset_env(self):
-        """Newton: ``num_pushes`` zeroes for env 0's DOFs only after reset of [0]."""
+        """Newton: ``num_pushes`` zeroes for env 0's DOFs only after a partial reset of env 0."""
         ctx, sim, articulation = self._build_and_warm(use_newton_actuators=True)
         try:
             adapter = self._get_adapter(articulation)
             self.assertIsNotNone(adapter)
-            # Find a DelayedPD actuator (it's the only one with delay_state).
-            stateful_pairs = [
-                (act, st)
-                for act, st in zip(adapter.actuators, adapter._states_a)
-                if st is not None and getattr(st, "delay_state", None) is not None
-            ]
-            self.assertGreater(len(stateful_pairs), 0, "expected at least one DelayedPD actuator with delay_state")
-
-            for act, state in stateful_pairs:
-                pushes_before = state.delay_state.num_pushes.numpy()
-                self.assertTrue(
-                    (pushes_before > 0).all(),
-                    "expected non-zero num_pushes for all DOFs after warmup",
-                )
-
-            articulation.reset(env_ids=torch.tensor([self.RESET_ENV], device=articulation.device, dtype=torch.long))
-
-            # Map each entry of ``act.indices`` to its env via the adapter's
-            # per-env DOF count. On the Newton backend the adapter is model-wide
-            # (includes free-joint DOFs on floating-base articulations); on
-            # PhysX it is per-articulation — ``adapter.num_joints`` is the
-            # correct stride in both cases.
-            for act, state in stateful_pairs:
-                pushes_after = state.delay_state.num_pushes.numpy()
-                indices_np = act.indices.numpy()
-                for i, global_dof in enumerate(indices_np):
-                    env = int(global_dof) // adapter.num_joints
-                    if env == self.RESET_ENV:
-                        self.assertEqual(
-                            int(pushes_after[i]),
-                            0,
-                            f"DOF {i} (env {env}) should be reset to 0, got {pushes_after[i]}",
-                        )
-                    else:
-                        self.assertGreater(
-                            int(pushes_after[i]),
-                            0,
-                            f"DOF {i} (env {env}) was NOT in reset env_ids but num_pushes is 0",
+            for selector in self.RESET_SELECTORS:
+                with self.subTest(selector=selector):
+                    self._warm(sim, articulation)
+                    # Find a DelayedPD actuator (it's the only one with delay_state).
+                    stateful_pairs = [
+                        (act, st)
+                        for act, st in zip(adapter.actuators, adapter._states_a)
+                        if st is not None and getattr(st, "delay_state", None) is not None
+                    ]
+                    self.assertGreater(
+                        len(stateful_pairs), 0, "expected at least one DelayedPD actuator with delay_state"
+                    )
+
+                    for act, state in stateful_pairs:
+                        pushes_before = state.delay_state.num_pushes.numpy()
+                        self.assertTrue(
+                            (pushes_before > 0).all(),
+                            "expected non-zero num_pushes for all DOFs after warmup",
                         )
+
+                    self._reset_env(articulation, selector)
+
+                    # Map each entry of ``act.indices`` to its env via the adapter's
+                    # per-env DOF count. On the Newton backend the adapter is model-wide
+                    # (includes free-joint DOFs on floating-base articulations); on
+                    # PhysX it is per-articulation — ``adapter.num_joints`` is the
+                    # correct stride in both cases.
+                    for act, state in stateful_pairs:
+                        pushes_after = state.delay_state.num_pushes.numpy()
+                        indices_np = act.indices.numpy()
+                        for i, global_dof in enumerate(indices_np):
+                            env = int(global_dof) // adapter.num_joints
+                            if env == self.RESET_ENV:
+                                self.assertEqual(
+                                    int(pushes_after[i]),
+                                    0,
+                                    f"DOF {i} (env {env}) should be reset to 0, got {pushes_after[i]}",
+                                )
+                            else:
+                                self.assertGreater(
+                                    int(pushes_after[i]),
+                                    0,
+                                    f"DOF {i} (env {env}) was NOT in the reset selection but num_pushes is 0",
+                                )
         finally:
             ctx.__exit__(None, None, None)
 
@@ -432,23 +452,26 @@ def test_lab_state_reset_isolated_to_reset_env(self):
         try:
             commands = articulation.actuators.target_command
             old_target = commands.position.torch.clone()
-            # Episode reset samples the configured lag; construction alone leaves actuator lag at zero.
-            articulation.reset()
-            commands.set_position_index(value=old_target)
-            articulation.write_data_to_sim()
             new_target = old_target + 0.02
-            articulation.reset(env_ids=torch.tensor([self.RESET_ENV], device=articulation.device))
-            commands.set_position_index(value=new_target)
-            articulation.write_data_to_sim()
             expected_target = old_target.clone()
             expected_target[self.RESET_ENV] = new_target[self.RESET_ENV]
-            for actuator in articulation.actuators.values():
-                joints = actuator.joint_indices
-                demand = (
-                    actuator.stiffness * (expected_target[:, joints] - articulation.data.joint_pos.torch[:, joints])
-                    - actuator.damping * articulation.data.joint_vel.torch[:, joints]
-                )
-                torch.testing.assert_close(actuator.computed_effort, demand)
+            for selector in self.RESET_SELECTORS:
+                with self.subTest(selector=selector):
+                    # Episode reset samples the configured lag; construction alone leaves actuator lag at zero.
+                    articulation.reset()
+                    commands.set_position_index(value=old_target)
+                    articulation.write_data_to_sim()
+                    self._reset_env(articulation, selector)
+                    commands.set_position_index(value=new_target)
+                    articulation.write_data_to_sim()
+                    for actuator in articulation.actuators.values():
+                        joints = actuator.joint_indices
+                        demand = (
+                            actuator.stiffness
+                            * (expected_target[:, joints] - articulation.data.joint_pos.torch[:, joints])
+                            - actuator.damping * articulation.data.joint_vel.torch[:, joints]
+                        )
+                        torch.testing.assert_close(actuator.computed_effort, demand)
 
         finally:
             ctx.__exit__(None, None, None)
diff --git a/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst b/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
new file mode 100644
index 000000000000..4f283403787d
--- /dev/null
+++ b/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
@@ -0,0 +1,5 @@
+* Fixed :meth:`~isaaclab_newton.assets.Articulation.reset` with only ``env_mask`` resetting the actuator state
+  (delay buffers, network history, and Newton-native actuator state) of every environment.
+* Fixed :meth:`~isaaclab_newton.assets.RigidObject.reset` and
+  :meth:`~isaaclab_newton.assets.RigidObjectCollection.reset` ignoring ``env_mask`` and clearing the external
+  wrenches of every environment.
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/actuator_control.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/actuator_control.py
index 83609ce33d2e..e09b575e4604 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/actuator_control.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/actuator_control.py
@@ -200,6 +200,10 @@ def reset_native_actuators(self, env_ids: Sequence[int] | slice) -> None:
         if self._native_actuator_path_active and SimulationManager._adapter is not None:
             SimulationManager._adapter.reset(env_ids)
 
+    def reset_native_actuators_mask(self, env_mask: wp.array) -> None:
+        if self._native_actuator_path_active and SimulationManager._adapter is not None:
+            SimulationManager._adapter.reset(env_mask=env_mask)
+
     def _joint_dof_offset(self) -> int:
         """Return the first selected joint DOF's model offset within an environment."""
         from newton import Model as NewtonModel  # noqa: PLC0415
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
index 93cf3044f6a5..26006675ba9f 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
@@ -370,7 +370,7 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None
             env_ids = slice(None)
         # reset actuators, including backend-native actuator state. None selects all
         # environments; delayed-actuator buffers do not accept a slice.
-        self.actuators.reset(None if env_ids == slice(None) else env_ids)
+        self.actuators.reset(None if env_ids == slice(None) else env_ids, env_mask=env_mask)
         # reset external wrenches.
         self._instantaneous_wrench_composer.reset(env_ids, env_mask)
         self._permanent_wrench_composer.reset(env_ids, env_mask)
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py
index 63305d9e46d8..b1820db845a9 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py
@@ -131,8 +131,8 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None
         if (env_ids is None) or (env_ids == slice(None)):
             env_ids = slice(None)
         # reset external wrench
-        self._instantaneous_wrench_composer.reset(env_ids)
-        self._permanent_wrench_composer.reset(env_ids)
+        self._instantaneous_wrench_composer.reset(env_ids, env_mask)
+        self._permanent_wrench_composer.reset(env_ids, env_mask)
 
     def write_data_to_sim(self) -> None:
         """Write external wrench to the simulation.
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py
index 5b1cbafb9978..8c1e586616b9 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py
@@ -177,8 +177,8 @@ def reset(
         if object_ids is None:
             object_ids = self._ALL_BODY_INDICES
         # reset external wrench
-        self._instantaneous_wrench_composer.reset(env_ids)
-        self._permanent_wrench_composer.reset(env_ids)
+        self._instantaneous_wrench_composer.reset(env_ids, env_mask)
+        self._permanent_wrench_composer.reset(env_ids, env_mask)
 
     def write_data_to_sim(self) -> None:
         """Write external wrench to the simulation.
diff --git a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py
index 95e183e5163d..690c9cf0d582 100644
--- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py
+++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py
@@ -677,6 +677,8 @@ class TestActuatorStateReset(ActuatorStateResetBase, unittest.TestCase):
     this subclass provides the Newton sim config and the model-wide adapter.
     """
 
+    RESET_SELECTORS = ("env_ids", "env_mask")
+
     def _make_sim_cfg(self, use_newton_actuators: bool) -> SimulationCfg:
         return SimulationCfg(dt=DT, physics=NEWTON_CFG, use_newton_actuators=use_newton_actuators)
 
diff --git a/source/isaaclab_newton/test/assets/test_rigid_object.py b/source/isaaclab_newton/test/assets/test_rigid_object.py
index a49149b30455..7d0d6f262c1e 100644
--- a/source/isaaclab_newton/test/assets/test_rigid_object.py
+++ b/source/isaaclab_newton/test/assets/test_rigid_object.py
@@ -276,6 +276,14 @@ def simulate():
             # Second object should have fallen, so it's Z height should be less than initial height of 1.0
             assert torch.all(cube_object.data.root_pos_w.torch[1::2, 2] < 1.0)
 
+        # A masked reset clears the wrench of the selected object only.
+        env_mask = torch.zeros(num_cubes, dtype=torch.bool, device=sim.device)
+        env_mask[0] = True
+        cube_object.reset(env_mask=wp.from_torch(env_mask))
+        permanent_force_b = cube_object.permanent_wrench_composer.out_force_b.torch
+        assert torch.count_nonzero(permanent_force_b[0]) == 0
+        assert torch.count_nonzero(permanent_force_b[2]) > 0
+
 
 @pytest.mark.isaacsim_ci
 @pytest.mark.parametrize("num_cubes", [2])
diff --git a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py
index 3612013271fe..9f19d01dcf3a 100644
--- a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py
+++ b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py
@@ -380,6 +380,14 @@ def simulate():
             # Second object should have fallen, so it's Z height should be less than initial height of 1.0
             assert torch.all(object_collection.data.body_link_pos_w.torch[:, 1::2, 2] < 1.0)
 
+        # A masked reset clears the wrenches of the selected environment only.
+        env_mask = torch.zeros(num_envs, dtype=torch.bool, device=sim.device)
+        env_mask[0] = True
+        object_collection.reset(env_mask=wp.from_torch(env_mask))
+        permanent_force_b = object_collection.permanent_wrench_composer.out_force_b.torch
+        assert torch.count_nonzero(permanent_force_b[0]) == 0
+        assert torch.count_nonzero(permanent_force_b[1, 0::2]) > 0
+
 
 @pytest.mark.isaacsim_ci
 @pytest.mark.parametrize("num_envs", [3])
diff --git a/source/isaaclab_ov/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst b/source/isaaclab_ov/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
new file mode 100644
index 000000000000..776f5fc75723
--- /dev/null
+++ b/source/isaaclab_ov/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
@@ -0,0 +1,2 @@
+* Fixed :meth:`~isaaclab_ov.assets.Articulation.reset` with only ``env_mask`` resetting the actuator state of every
+  environment.
diff --git a/source/isaaclab_ov/isaaclab_ov/assets/articulation/actuator_control.py b/source/isaaclab_ov/isaaclab_ov/assets/articulation/actuator_control.py
index 9676c1a562da..8b501148aad3 100644
--- a/source/isaaclab_ov/isaaclab_ov/assets/articulation/actuator_control.py
+++ b/source/isaaclab_ov/isaaclab_ov/assets/articulation/actuator_control.py
@@ -102,6 +102,10 @@ def reset_native_actuators(self, env_ids: Sequence[int] | slice) -> None:
         if self._native_actuator_path_active and self._actuator_runtime is not None:
             self._actuator_runtime.reset(env_ids)
 
+    def reset_native_actuators_mask(self, env_mask: wp.array) -> None:
+        if self._native_actuator_path_active and self._actuator_runtime is not None:
+            self._actuator_runtime.reset(env_mask=env_mask)
+
     def stage_user_command(
         self,
         command_name: str,
diff --git a/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py b/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py
index 8bc71c184231..9e4361b9f030 100644
--- a/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py
+++ b/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py
@@ -222,7 +222,7 @@ def reset(
             env_ids = slice(None)
         # reset actuators, including backend-native actuator state. None selects all
         # environments; delayed-actuator buffers do not accept a slice.
-        self.actuators.reset(None if env_ids == slice(None) else env_ids)
+        self.actuators.reset(None if env_ids == slice(None) else env_ids, env_mask=env_mask)
         # reset external wrenches.
         self._instantaneous_wrench_composer.reset(env_ids, env_mask)
         self._permanent_wrench_composer.reset(env_ids, env_mask)
diff --git a/source/isaaclab_ov/test/assets/test_articulation.py b/source/isaaclab_ov/test/assets/test_articulation.py
index f95e7aa082a4..04880b767757 100644
--- a/source/isaaclab_ov/test/assets/test_articulation.py
+++ b/source/isaaclab_ov/test/assets/test_articulation.py
@@ -710,6 +710,17 @@ def __getitem__(self, name):
         articulation.reset(env_ids=torch.tensor([0], device=device, dtype=torch.long))
         assert stateful_pairs[0].delay_state.num_pushes.numpy().tolist() == [0, 1]
 
+        articulation.write_data_to_sim()
+        sim.step()
+        articulation.update(sim.cfg.dt)
+        articulation.reset(env_mask=wp.array([False, True], dtype=wp.bool, device=device))
+        stateful_pairs = [
+            state
+            for state in adapter._states_a
+            if state is not None and getattr(state, "delay_state", None) is not None
+        ]
+        assert stateful_pairs[0].delay_state.num_pushes.numpy().tolist() == [1, 0]
+
         env = Env(articulation)
         asset_cfg = SceneEntityCfg("robot")
         event_params = {

From 116e02167009060c44fdaa671886a036925814aa Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 02:05:55 -0700
Subject: [PATCH 02/48] Forward the reset mask to every Warp scene entity

InteractiveSceneWarp.reset passed env_ids=None to deformable objects, surface
grippers, and sensors on a masked reset, so every environment's sensor state
was cleared on each partial reset, and it never reset cable objects.

Forward the mask to every entity family in the stable scene's order. Surface
grippers use their reset_mask method. Cable objects take env_mask like the
other scene assets, so the scene does not convert the mask for a reset that
holds no state.
---
 .../jichuanh-warp-scene-mask-native.added.rst |  1 +
 .../assets/cable_object/base_cable_object.py  |  3 +-
 .../jichuanh-warp-scene-mask-native.fixed.rst |  2 +
 .../envs/interactive_scene_warp.py            | 18 ++++--
 .../test/envs/test_interactive_scene_warp.py  | 63 +++++++++++++++++++
 .../jichuanh-warp-scene-mask-native.added.rst |  1 +
 .../jichuanh-warp-scene-mask-native.minor     |  0
 .../assets/cable_object/cable_object.py       |  5 +-
 8 files changed, 85 insertions(+), 8 deletions(-)
 create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
 create mode 100644 source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py
 create mode 100644 source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.added.rst
 create mode 100644 source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.minor

diff --git a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst
index 7b04995e3abd..3329702113c2 100644
--- a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst
+++ b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst
@@ -1,3 +1,4 @@
 * Added ``env_mask`` to :meth:`~isaaclab.actuators.ActuatorCollection.reset`. Isaac Lab actuator models reset only
   the masked environments, converting the mask to indices once for all groups, and backend-native actuator state
   resets from the mask through :meth:`~isaaclab.actuators.ActuatorControl.reset_native_actuators_mask`.
+* Added ``env_mask`` to :meth:`~isaaclab.assets.BaseCableObject.reset`, matching the other scene assets.
diff --git a/source/isaaclab/isaaclab/assets/cable_object/base_cable_object.py b/source/isaaclab/isaaclab/assets/cable_object/base_cable_object.py
index 4bc2b6951f3c..49a50c93fe26 100644
--- a/source/isaaclab/isaaclab/assets/cable_object/base_cable_object.py
+++ b/source/isaaclab/isaaclab/assets/cable_object/base_cable_object.py
@@ -61,11 +61,12 @@ def root_view(self):
         raise NotImplementedError()
 
     @abstractmethod
-    def reset(self, env_ids: Sequence[int] | None = None) -> None:
+    def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None) -> None:
         """Reset the cable object.
 
         Args:
             env_ids: Environment indices. If None, all instances are used.
+            env_mask: Environment mask. If None, all instances are used. Shape is (num_instances,).
         """
         raise NotImplementedError()
 
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
new file mode 100644
index 000000000000..2d0ece0abb4e
--- /dev/null
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
@@ -0,0 +1,2 @@
+* Fixed :meth:`~isaaclab_experimental.envs.InteractiveSceneWarp.reset` with ``env_mask`` resetting deformable
+  objects, surface grippers, and sensors in every environment, and skipping cable objects.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py
index f792b994f67b..43a5047426ab 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py
@@ -24,21 +24,29 @@ class InteractiveSceneWarp(InteractiveScene):
     def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None):
         """Reset scene entities using either env_ids or a boolean env_mask.
 
+        .. caution::
+            If both ``env_ids`` and ``env_mask`` are provided, ``env_mask`` takes precedence.
+
         Args:
             env_ids: The indices of the environments to reset. Defaults to None (all instances).
             env_mask: Boolean warp mask of shape (num_envs,). Defaults to None.
         """
-        # -- assets (support env_mask)
+        # -- assets
         for articulation in self._articulations.values():
             articulation.reset(env_ids, env_mask=env_mask)
+        for cable_object in self._cable_objects.values():
+            cable_object.reset(env_ids, env_mask=env_mask)
         for deformable_object in self._deformable_objects.values():
-            deformable_object.reset(env_ids)
+            deformable_object.reset(env_ids, env_mask=env_mask)
         for rigid_object in self._rigid_objects.values():
             rigid_object.reset(env_ids, env_mask=env_mask)
         for surface_gripper in self._surface_grippers.values():
-            surface_gripper.reset(env_ids)
+            if env_mask is None:
+                surface_gripper.reset(env_ids)
+            else:
+                surface_gripper.reset_mask(wp.to_torch(env_mask))
         for rigid_object_collection in self._rigid_object_collections.values():
             rigid_object_collection.reset(env_ids, env_mask=env_mask)
-        # -- sensors (no env_mask support)
+        # -- sensors
         for sensor in self._sensors.values():
-            sensor.reset(env_ids)
+            sensor.reset(env_ids, env_mask=env_mask)
diff --git a/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py b/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py
new file mode 100644
index 000000000000..440cc82a29a9
--- /dev/null
+++ b/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py
@@ -0,0 +1,63 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tests for masked resets of the Warp interactive scene on Newton."""
+
+import torch
+import warp as wp
+from isaaclab_experimental.envs import InteractiveSceneWarp
+from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg
+
+import isaaclab.sim as sim_utils
+from isaaclab.assets import RigidObjectCfg
+from isaaclab.scene import InteractiveSceneCfg
+from isaaclab.sensors.imu import ImuCfg
+from isaaclab.sim import SimulationCfg
+from isaaclab.terrains import TerrainImporterCfg
+from isaaclab.utils import configclass
+
+
+@configclass
+class ImuSceneCfg(InteractiveSceneCfg):
+    """Scene with a cube resting on the ground and an IMU on the cube."""
+
+    env_spacing = 2.0
+    terrain = TerrainImporterCfg(prim_path="/World/ground", terrain_type="plane")
+
+    cube = RigidObjectCfg(
+        prim_path="{ENV_REGEX_NS}/Cube",
+        spawn=sim_utils.CuboidCfg(
+            size=(0.2, 0.2, 0.2),
+            rigid_props=sim_utils.UsdPhysicsRigidBodyCfg(),
+            mass_props=sim_utils.MassCfg(mass=1.0),
+            collision_props=sim_utils.UsdPhysicsCollisionCfg(),
+        ),
+        init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 0.1)),
+    )
+
+    imu = ImuCfg(prim_path="{ENV_REGEX_NS}/Cube")
+
+
+def test_masked_reset_keeps_sensor_state_of_unselected_envs():
+    """A masked scene reset zeroes the selected environment's IMU reading and keeps the other one."""
+    sim_cfg = SimulationCfg(dt=1.0 / 200.0, physics=NewtonCfg(solver_cfg=MJWarpSolverCfg(), num_substeps=1))
+    with sim_utils.build_simulation_context(sim_cfg=sim_cfg) as sim:
+        sim._app_control_on_stop_handle = None
+        scene = InteractiveSceneWarp(ImuSceneCfg(num_envs=2))
+        sim.reset()
+        scene.reset()
+        imu = scene["imu"]
+
+        for _ in range(50):
+            scene.write_data_to_sim()
+            sim.step(render=False)
+            scene.update(dt=sim.get_physics_dt())
+        pre_reset_lin_acc = imu.data.lin_acc_b.torch.clone()
+        assert (pre_reset_lin_acc[:, 2] > 5.0).all(), f"Expected a resting gravity reading, got {pre_reset_lin_acc}"
+
+        scene.reset(env_mask=wp.array([True, False], dtype=wp.bool, device=imu.device))
+        lin_acc = imu.data.lin_acc_b.torch
+        torch.testing.assert_close(lin_acc[0], torch.zeros_like(lin_acc[0]))
+        torch.testing.assert_close(lin_acc[1], pre_reset_lin_acc[1])
diff --git a/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.added.rst b/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.added.rst
new file mode 100644
index 000000000000..bce8fc255e4e
--- /dev/null
+++ b/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.added.rst
@@ -0,0 +1 @@
+* Added ``env_mask`` to :meth:`~isaaclab_newton.assets.CableObject.reset`.
diff --git a/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.minor b/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.minor
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/cable_object/cable_object.py b/source/isaaclab_newton/isaaclab_newton/assets/cable_object/cable_object.py
index 83b4522a4220..bc9bc65b72dc 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/cable_object/cable_object.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/cable_object/cable_object.py
@@ -61,13 +61,14 @@ def root_view(self) -> ArticulationView:
         """Articulation view for the cable."""
         return self._root_view
 
-    def reset(self, env_ids: Sequence[int] | None = None) -> None:
+    def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None) -> None:
         """Reset the cable object's internal buffers.
 
         Args:
             env_ids: Environment indices. Defaults to all instances.
+            env_mask: Environment mask. Defaults to all instances. Shape is (num_instances,).
         """
-        del env_ids
+        del env_ids, env_mask
 
     def write_data_to_sim(self) -> None:
         """Write buffered commands to the simulation."""

From 88b1ba2d4caf991b598be1d80267aef8a8c6b992 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 02:11:37 -0700
Subject: [PATCH 03/48] Write collection masks and reset ray-caster drift
 without host sync

RigidObjectCollection.write_body_*_to_sim_mask converted masks to indices
with torch.nonzero before delegating to the index writers, so the mask
writers synchronized the host and could not be captured in a CUDA graph,
although their docstrings name them as the graph-capable path. They now
launch masked kernels over the full environment-body grid, and body_ids
selections resolve to a mask through a preallocated scratch buffer.

BaseRayCaster.reset compacted a mask with nonzero before resampling drift.
A masked reset now resamples every environment and keeps the new values
where the mask is set; index and slice selections are unchanged.
---
 ...ichuanh-warp-scene-mask-native.changed.rst |   3 +
 .../sensors/ray_caster/base_ray_caster.py     |  29 ++-
 .../jichuanh-warp-scene-mask-native.fixed.rst |   3 +
 .../isaaclab_newton/assets/kernels.py         | 131 ++++++++++++
 .../rigid_object_collection.py                | 199 ++++++++++--------
 .../assets/test_rigid_object_collection.py    |  43 ++++
 .../sensors/test_newton_raycast_sensor.py     |  28 +++
 7 files changed, 344 insertions(+), 92 deletions(-)
 create mode 100644 source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst

diff --git a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
new file mode 100644
index 000000000000..bc83b269855f
--- /dev/null
+++ b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
@@ -0,0 +1,3 @@
+* Changed masked :meth:`~isaaclab.sensors.ray_caster.BaseRayCaster.reset` to resample drift without a host
+  synchronization. When both ``env_ids`` and ``env_mask`` are given, the mask selects the resampled environments,
+  matching :meth:`~isaaclab.sensors.SensorBase.reset`.
diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py
index 1a98a5c41fd9..283690beba4a 100644
--- a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py
+++ b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py
@@ -116,12 +116,31 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None
         self._drift_sampled |= sample_drift or sample_ray_cast_drift
         if not self._drift_sampled:
             return
+        if sample_ray_cast_drift and ray_cast_range_list != self._ray_cast_drift_range_list:
+            # Upload the ranges to the device only when the configuration changes.
+            self._ray_cast_drift_range_list = ray_cast_range_list
+            self._ray_cast_drift_ranges = torch.tensor(ray_cast_range_list, device=self.device)
+        if env_mask is not None:
+            # Resample every environment and keep the new values where the mask is set, so the
+            # masked reset needs no host synchronization.
+            mask = wp.to_torch(env_mask).view(-1, 1)
+            drift = self.drift.torch
+            new_drift = torch.zeros_like(drift)
+            if sample_drift:
+                new_drift.uniform_(*self.cfg.drift_range)
+            drift.copy_(torch.where(mask, new_drift, drift))
+            ray_cast_drift = self.ray_cast_drift.torch
+            new_ray_cast_drift = torch.zeros_like(ray_cast_drift)
+            if sample_ray_cast_drift:
+                ranges = self._ray_cast_drift_ranges
+                new_ray_cast_drift = math_utils.sample_uniform(
+                    ranges[:, 0], ranges[:, 1], tuple(ray_cast_drift.shape), device=self.device
+                )
+            ray_cast_drift.copy_(torch.where(mask, new_ray_cast_drift, ray_cast_drift))
+            return
         # determine the selected batch size
         if env_ids is not None:
             num_envs_ids = len(range(self._view_count)[env_ids]) if isinstance(env_ids, slice) else len(env_ids)
-        elif env_mask is not None:
-            env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1)
-            num_envs_ids = len(env_ids)
         else:
             env_ids = slice(None)
             num_envs_ids = self._view_count
@@ -133,10 +152,6 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None
             self.drift.torch[env_ids] = 0.0
         # resample the ray cast drift
         if sample_ray_cast_drift:
-            # Upload the ranges to the device only when the configuration changes.
-            if ray_cast_range_list != self._ray_cast_drift_range_list:
-                self._ray_cast_drift_range_list = ray_cast_range_list
-                self._ray_cast_drift_ranges = torch.tensor(ray_cast_range_list, device=self.device)
             ranges = self._ray_cast_drift_ranges
             self.ray_cast_drift.torch[env_ids] = math_utils.sample_uniform(
                 ranges[:, 0], ranges[:, 1], (num_envs_ids, 3), device=self.device
diff --git a/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst b/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
index 4f283403787d..67e8df01a9eb 100644
--- a/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
+++ b/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
@@ -3,3 +3,6 @@
 * Fixed :meth:`~isaaclab_newton.assets.RigidObject.reset` and
   :meth:`~isaaclab_newton.assets.RigidObjectCollection.reset` ignoring ``env_mask`` and clearing the external
   wrenches of every environment.
+* Fixed the ``write_body_*_to_sim_mask`` methods of :class:`~isaaclab_newton.assets.RigidObjectCollection`
+  converting masks to indices on the host, which prevented capturing them in CUDA graphs as documented. They now
+  write through masked kernels.
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/kernels.py b/source/isaaclab_newton/isaaclab_newton/assets/kernels.py
index eec3479f1fdc..6b9d0ed497fe 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/kernels.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/kernels.py
@@ -992,6 +992,137 @@ def set_body_link_velocity_to_sim_kernel(
     return _SET_BODY_LINK_VELOCITY_TO_SIM_DISPATCHER.select(env_ids, body_ids)
 
 
+@wp.kernel
+def set_body_link_pose_to_sim_mask(
+    data: wp.array2d(dtype=wp.transformf),
+    env_mask: wp.array(dtype=wp.bool),
+    body_mask: wp.array(dtype=wp.bool),
+    body_link_pose_w: wp.array2d(dtype=wp.transformf),
+):
+    """Write masked body link pose data to simulation buffers.
+
+    Launched over the full ``(num_envs, num_bodies)`` grid with an in-kernel mask test, so the
+    write needs no host-side index compaction.
+
+    Args:
+        data: Input array of body link poses. Shape is (num_envs, num_bodies).
+        env_mask: Boolean environment selection. Shape is (num_envs,).
+        body_mask: Boolean body selection. Shape is (num_bodies,).
+        body_link_pose_w: Output array where body link poses are written.
+            Shape is (num_envs, num_bodies).
+    """
+    i, j = wp.tid()
+    if env_mask[i] and body_mask[j]:
+        body_link_pose_w[i, j] = data[i, j]
+
+
+@wp.kernel
+def set_body_com_pose_to_sim_mask(
+    data: wp.array2d(dtype=wp.transformf),
+    body_com_pos_b: wp.array2d(dtype=wp.vec3f),
+    env_mask: wp.array(dtype=wp.bool),
+    body_mask: wp.array(dtype=wp.bool),
+    body_com_pose_w: wp.array2d(dtype=wp.transformf),
+    body_link_pose_w: wp.array2d(dtype=wp.transformf),
+):
+    """Write masked body COM pose data to simulation buffers.
+
+    Launched over the full ``(num_envs, num_bodies)`` grid with an in-kernel mask test, so the
+    write needs no host-side index compaction. Link poses are derived from the written COM poses
+    for the selected cells.
+
+    Args:
+        data: Input array of body COM poses. Shape is (num_envs, num_bodies).
+        body_com_pos_b: Input array of body COM positions in body frame.
+            Shape is (num_envs, num_bodies).
+        env_mask: Boolean environment selection. Shape is (num_envs,).
+        body_mask: Boolean body selection. Shape is (num_bodies,).
+        body_com_pose_w: Output array where body COM poses are written.
+            Shape is (num_envs, num_bodies).
+        body_link_pose_w: Output array where body link poses (derived from COM) are written.
+            Shape is (num_envs, num_bodies).
+    """
+    i, j = wp.tid()
+    if env_mask[i] and body_mask[j]:
+        body_com_pose_w[i, j] = data[i, j]
+        # Get the link pose from com pose
+        body_link_pose_w[i, j] = get_com_pose_in_link_frame_func(body_com_pose_w[i, j], body_com_pos_b[i, j])
+
+
+@wp.kernel
+def set_body_com_velocity_to_sim_mask(
+    data: wp.array2d(dtype=wp.spatial_vectorf),
+    env_mask: wp.array(dtype=wp.bool),
+    body_mask: wp.array(dtype=wp.bool),
+    body_com_velocity_w: wp.array2d(dtype=wp.spatial_vectorf),
+    body_acc_w: wp.array2d(dtype=wp.spatial_vectorf),
+):
+    """Write masked body COM velocity data to simulation buffers.
+
+    Launched over the full ``(num_envs, num_bodies)`` grid with an in-kernel mask test, so the
+    write needs no host-side index compaction. Accelerations of the selected cells are zeroed.
+
+    Args:
+        data: Input array of body COM spatial velocities. Shape is (num_envs, num_bodies).
+        env_mask: Boolean environment selection. Shape is (num_envs,).
+        body_mask: Boolean body selection. Shape is (num_bodies,).
+        body_com_velocity_w: Output array where body COM velocities are written.
+            Shape is (num_envs, num_bodies).
+        body_acc_w: Output array where body accelerations are zeroed.
+            Shape is (num_envs, num_bodies).
+    """
+    i, j = wp.tid()
+    if env_mask[i] and body_mask[j]:
+        body_com_velocity_w[i, j] = data[i, j]
+        # Make the acceleration zero to prevent reporting old values
+        body_acc_w[i, j] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
+
+
+@wp.kernel
+def set_body_link_velocity_to_sim_mask(
+    data: wp.array2d(dtype=wp.spatial_vectorf),
+    body_com_pos_b: wp.array2d(dtype=wp.vec3f),
+    body_link_pose_w: wp.array2d(dtype=wp.transformf),
+    env_mask: wp.array(dtype=wp.bool),
+    body_mask: wp.array(dtype=wp.bool),
+    body_link_velocity_w: wp.array2d(dtype=wp.spatial_vectorf),
+    body_com_velocity_w: wp.array2d(dtype=wp.spatial_vectorf),
+    body_acc_w: wp.array2d(dtype=wp.spatial_vectorf),
+):
+    """Write masked body link velocity data to simulation buffers.
+
+    Launched over the full ``(num_envs, num_bodies)`` grid with an in-kernel mask test, so the
+    write needs no host-side index compaction. COM velocities are derived from the written link
+    velocities and accelerations are zeroed for the selected cells.
+
+    Args:
+        data: Input array of body link spatial velocities. Shape is (num_envs, num_bodies).
+        body_com_pos_b: Input array of body COM positions in body frame.
+            Shape is (num_envs, num_bodies).
+        body_link_pose_w: Input array of body link poses in world frame.
+            Shape is (num_envs, num_bodies).
+        env_mask: Boolean environment selection. Shape is (num_envs,).
+        body_mask: Boolean body selection. Shape is (num_bodies,).
+        body_link_velocity_w: Output array where body link velocities are written.
+            Shape is (num_envs, num_bodies).
+        body_com_velocity_w: Output array where body COM velocities (derived from link) are written.
+            Shape is (num_envs, num_bodies).
+        body_acc_w: Output array where body accelerations are zeroed.
+            Shape is (num_envs, num_bodies).
+    """
+    i, j = wp.tid()
+    if env_mask[i] and body_mask[j]:
+        body_link_velocity_w[i, j] = data[i, j]
+        # Get the link velocity in the com frame
+        body_com_velocity_w[i, j] = get_link_velocity_in_com_frame_func(
+            body_link_velocity_w[i, j],
+            body_link_pose_w[i, j],
+            body_com_pos_b[i, j],
+        )
+        # Make the acceleration zero to prevent reporting old values
+        body_acc_w[i, j] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
+
+
 """
 Generic buffer-writing kernels (used by Articulation + RigidObject + RigidObjectCollection).
 """
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py
index 8c1e586616b9..41537f57e7b5 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py
@@ -23,6 +23,7 @@
 from isaaclab.physics import PhysicsEvent
 from isaaclab.utils import clone, validate
 from isaaclab.utils.warp import ProxyArray
+from isaaclab.utils.warp.utils import resolve_1d_mask
 from isaaclab.utils.wrench_composer import WrenchComposer
 
 from isaaclab_newton.assets import kernels as shared_kernels
@@ -310,16 +311,8 @@ def write_body_pose_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_ids = self._resolve_body_mask(body_mask)
-        else:
-            body_ids = self._ALL_BODY_INDICES
-        self.write_body_link_pose_to_sim_index(
-            body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward
+        self.write_body_link_pose_to_sim_mask(
+            body_poses=body_poses, body_mask=body_mask, env_mask=env_mask, skip_forward=skip_forward
         )
 
     def write_body_velocity_to_sim_index(
@@ -388,20 +381,8 @@ def write_body_velocity_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_ids = self._resolve_body_mask(body_mask)
-        else:
-            body_ids = self._ALL_BODY_INDICES
-        self.write_body_com_velocity_to_sim_index(
-            body_velocities=body_velocities,
-            env_ids=env_ids,
-            body_ids=body_ids,
-            full_data=True,
-            skip_forward=skip_forward,
+        self.write_body_com_velocity_to_sim_mask(
+            body_velocities=body_velocities, body_mask=body_mask, env_mask=env_mask, skip_forward=skip_forward
         )
 
     def write_body_link_pose_to_sim_index(
@@ -494,16 +475,27 @@ def write_body_link_pose_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
+        if env_mask is None:
             env_mask = self._ALL_ENV_MASK
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_ids = self._resolve_body_mask(body_mask)
-        self.write_body_link_pose_to_sim_index(
-            body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward
+        body_mask = self._resolve_body_selection_mask(body_ids, body_mask)
+        self.assert_shape_and_dtype(body_poses, (self.num_instances, self.num_bodies), wp.transformf, "body_poses")
+        # Write to consolidated buffer
+        wp.launch(
+            shared_kernels.set_body_link_pose_to_sim_mask,
+            dim=(self.num_instances, self.num_bodies),
+            inputs=[
+                body_poses,
+                env_mask,
+                body_mask,
+            ],
+            outputs=[
+                self.data.body_link_pose_w,
+            ],
+            device=self.device,
         )
+        # Invalidate dependent timestamps
+        if not skip_forward:
+            self.data._reset_pose(env_mask=env_mask)
 
     def write_body_com_pose_to_sim_index(
         self,
@@ -599,17 +591,29 @@ def write_body_com_pose_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
+        if env_mask is None:
             env_mask = self._ALL_ENV_MASK
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_ids = self._resolve_body_mask(body_mask)
-        # The index writer owns the invalidation (with from_link=False); forward skip_forward to it.
-        self.write_body_com_pose_to_sim_index(
-            body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward
+        body_mask = self._resolve_body_selection_mask(body_ids, body_mask)
+        self.assert_shape_and_dtype(body_poses, (self.num_instances, self.num_bodies), wp.transformf, "body_poses")
+        # Write to consolidated buffers (updates both com_pose_w and link_pose_w)
+        wp.launch(
+            shared_kernels.set_body_com_pose_to_sim_mask,
+            dim=(self.num_instances, self.num_bodies),
+            inputs=[
+                body_poses,
+                self.data.body_com_pos_b,
+                env_mask,
+                body_mask,
+            ],
+            outputs=[
+                self.data.body_com_pose_w,
+                self.data.body_link_pose_w,
+            ],
+            device=self.device,
         )
+        # Invalidate dependent timestamps. The com poses were just written, so they must not be invalidated.
+        if not skip_forward:
+            self.data._reset_pose(env_mask=env_mask, from_link=False)
 
     def write_body_com_velocity_to_sim_index(
         self,
@@ -713,21 +717,30 @@ def write_body_com_velocity_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
+        if env_mask is None:
             env_mask = self._ALL_ENV_MASK
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_ids = self._resolve_body_mask(body_mask)
-        # The index writer owns the invalidation; forward skip_forward to it.
-        self.write_body_com_velocity_to_sim_index(
-            body_velocities=body_velocities,
-            env_ids=env_ids,
-            body_ids=body_ids,
-            full_data=True,
-            skip_forward=skip_forward,
+        body_mask = self._resolve_body_selection_mask(body_ids, body_mask)
+        self.assert_shape_and_dtype(
+            body_velocities, (self.num_instances, self.num_bodies), wp.spatial_vectorf, "body_velocities"
+        )
+        # Write to consolidated buffer
+        wp.launch(
+            shared_kernels.set_body_com_velocity_to_sim_mask,
+            dim=(self.num_instances, self.num_bodies),
+            inputs=[
+                body_velocities,
+                env_mask,
+                body_mask,
+            ],
+            outputs=[
+                self.data.body_com_vel_w,
+                self.data.body_com_acc_w,
+            ],
+            device=self.device,
         )
+        # Invalidate dependent timestamps
+        if not skip_forward:
+            self.data._reset_velocity(env_mask=env_mask)
 
     def write_body_link_velocity_to_sim_index(
         self,
@@ -833,21 +846,33 @@ def write_body_link_velocity_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
+        if env_mask is None:
             env_mask = self._ALL_ENV_MASK
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_ids = self._resolve_body_mask(body_mask)
-        # The index writer owns the invalidation (with from_com=False); forward skip_forward to it.
-        self.write_body_link_velocity_to_sim_index(
-            body_velocities=body_velocities,
-            env_ids=env_ids,
-            body_ids=body_ids,
-            full_data=True,
-            skip_forward=skip_forward,
+        body_mask = self._resolve_body_selection_mask(body_ids, body_mask)
+        self.assert_shape_and_dtype(
+            body_velocities, (self.num_instances, self.num_bodies), wp.spatial_vectorf, "body_velocities"
+        )
+        # Access body_com_pos_b and body_link_pose_w to ensure they are current.
+        wp.launch(
+            shared_kernels.set_body_link_velocity_to_sim_mask,
+            dim=(self.num_instances, self.num_bodies),
+            inputs=[
+                body_velocities,
+                self.data.body_com_pos_b,
+                self.data.body_link_pose_w,
+                env_mask,
+                body_mask,
+            ],
+            outputs=[
+                self.data.body_link_vel_w,
+                self.data.body_com_vel_w,
+                self.data.body_com_acc_w,
+            ],
+            device=self.device,
         )
+        # Invalidate dependent timestamps.
+        if not skip_forward:
+            self.data._reset_velocity(env_mask=env_mask, from_com=False)
 
     """
     Operations - Setters.
@@ -1251,6 +1276,8 @@ def _create_buffers(self):
         )
         self._ALL_ENV_MASK = wp.ones((self.num_instances,), dtype=wp.bool, device=self.device)
         self._ALL_BODY_MASK = wp.ones((self.num_bodies,), dtype=wp.bool, device=self.device)
+        # reusable working buffer for converting body-index selections into masks
+        self._scratch_body_mask = wp.zeros((self.num_bodies,), dtype=wp.bool, device=self.device)
 
         # external wrench composer
         self._instantaneous_wrench_composer = WrenchComposer(self)
@@ -1321,25 +1348,27 @@ def _resolve_body_ids(self, body_ids) -> wp.array | torch.Tensor:
             )
         return body_ids
 
-    def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array | torch.Tensor:
-        """Resolve environment mask to indices via torch.nonzero."""
-        if env_mask is not None:
-            if isinstance(env_mask, wp.array):
-                env_mask = wp.to_torch(env_mask)
-            env_ids = torch.nonzero(env_mask)[:, 0]
-        else:
-            env_ids = self._ALL_ENV_INDICES
-        return env_ids
+    def _resolve_body_selection_mask(
+        self,
+        body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None,
+        body_mask: wp.array | None,
+    ) -> wp.array:
+        """Resolve a body selection to a boolean mask without host synchronization.
 
-    def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array | torch.Tensor:
-        """Resolve body mask to indices via torch.nonzero."""
-        if body_mask is not None:
-            if isinstance(body_mask, wp.array):
-                body_mask = wp.to_torch(body_mask)
-            body_ids = torch.nonzero(body_mask)[:, 0]
-        else:
-            body_ids = self._ALL_BODY_INDICES
-        return body_ids
+        Args:
+            body_ids: Body indices. If None, then all bodies are selected.
+            body_mask: Body mask. Takes precedence over :paramref:`body_ids`. Shape is (num_bodies,).
+
+        Returns:
+            A boolean Warp mask. Shape is (num_bodies,).
+        """
+        return resolve_1d_mask(
+            ids=body_ids,
+            mask=body_mask,
+            all_mask=self._ALL_BODY_MASK,
+            scratch_mask=self._scratch_body_mask,
+            device=self.device,
+        )
 
     """
     Internal simulation callbacks.
diff --git a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py
index 9f19d01dcf3a..97cf47051e41 100644
--- a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py
+++ b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py
@@ -708,3 +708,46 @@ def _fk_reset_mask_dirty() -> bool:
                 f"body_link_pose_w still aliases the pre-write pose after {writer}; the buffer was not written"
             )
             torch.testing.assert_close(body_link[..., :3], target_pose[..., :3], rtol=1e-4, atol=1e-4)
+
+
+@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA))
+def test_mask_writes_replay_from_cuda_graph_on_selected_cells(device):
+    """Mask writers record into a CUDA graph and, on replay, write only the selected environment-body cells.
+
+    The mask variants are the documented graph-capturable write path. The replay reads the input buffer's
+    contents at launch time, so the input is filled only after capture.
+    """
+    num_envs, num_cubes = 3, 2
+    with _newton_sim_context(device, gravity_enabled=False, auto_add_lighting=True) as sim:
+        sim._app_control_on_stop_handle = None
+        cube_object, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, height=0.5, device=device)
+        replicate(sim.get_clone_plan())
+        sim.reset()
+        sim.step()
+        cube_object.update(sim.cfg.dt)
+
+        env_mask = wp.array([False, True, False], dtype=wp.bool, device=device)
+        body_mask = wp.array([True, False], dtype=wp.bool, device=device)
+        selected = torch.zeros(num_envs, num_cubes, 1, dtype=torch.bool, device=device)
+        selected[1, 0] = True
+        identity_pose = torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0], device=device)
+        writers = (
+            ("write_body_link_pose_to_sim_mask", "body_poses", "body_link_pose_w"),
+            ("write_body_com_pose_to_sim_mask", "body_poses", "body_com_pose_w"),
+            ("write_body_com_velocity_to_sim_mask", "body_velocities", "body_com_vel_w"),
+            ("write_body_link_velocity_to_sim_mask", "body_velocities", "body_link_vel_w"),
+        )
+        for writer, data_arg, buffer_name in writers:
+            before = getattr(cube_object.data, buffer_name).torch.clone()
+            target = torch.zeros_like(before)
+            with wp.ScopedCapture(device) as capture:
+                getattr(cube_object, writer)(**{data_arg: target}, env_mask=env_mask, body_mask=body_mask)
+            if data_arg == "body_poses":
+                target.copy_(identity_pose.expand_as(target))
+                target[..., :3] = torch.rand(num_envs, num_cubes, 3, device=device) + 2.0
+            else:
+                target.copy_(torch.rand_like(target) + 1.0)
+            wp.capture_launch(capture.graph)
+
+            written = getattr(cube_object.data, buffer_name).torch
+            torch.testing.assert_close(written, torch.where(selected, target, before), msg=writer)
diff --git a/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py b/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py
index 6de5c0e0edc4..d73fa8695f8e 100644
--- a/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py
+++ b/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py
@@ -12,6 +12,7 @@
 
 import pytest
 import torch
+import warp as wp
 from isaaclab_newton.physics import MJWarpSolverCfg, NewtonBuilderCfg, NewtonCfg, NewtonManager
 from isaaclab_newton.renderers import NewtonWarpRendererCfg
 from isaaclab_newton.sensors import (
@@ -196,6 +197,33 @@ def test_legacy_multi_mesh_tracks_ad_hoc_regex_target(sim):
     assert sensor.data.ray_hits_w.shape[0] == 1
 
 
+@pytest.mark.parametrize("sim", [pytest.param(False, id="eager")], indirect=True)
+def test_partial_reset_resamples_drift_of_selected_envs_only(sim):
+    """Mask and slice resets resample the drift of the selected environments and keep the others."""
+    scene_cfg = RaycastTestSceneCfg(num_envs=3)
+    scene_cfg.raycast.drift_range = (0.01, 0.05)
+    scene_cfg.raycast.ray_cast_drift_range = {"z": (0.1, 0.2)}
+    scene = InteractiveScene(scene_cfg)
+    sim.reset()
+    sensor: NewtonRaycastSensor = scene["raycast"]
+    sensor.reset()
+
+    drift = sensor.drift.torch.clone()
+    ray_cast_drift = sensor.ray_cast_drift.torch.clone()
+    sensor.reset(env_mask=wp.array([False, True, False], dtype=wp.bool, device=sensor.device))
+    torch.testing.assert_close(sensor.drift.torch[[0, 2]], drift[[0, 2]])
+    torch.testing.assert_close(sensor.ray_cast_drift.torch[[0, 2]], ray_cast_drift[[0, 2]])
+    assert not torch.allclose(sensor.drift.torch[1], drift[1])
+    assert not torch.allclose(sensor.ray_cast_drift.torch[1, 2], ray_cast_drift[1, 2])
+    assert ((sensor.drift.torch >= 0.01) & (sensor.drift.torch <= 0.05)).all()
+    assert ((sensor.ray_cast_drift.torch[:, 2] >= 0.1) & (sensor.ray_cast_drift.torch[:, 2] <= 0.2)).all()
+
+    drift = sensor.drift.torch.clone()
+    sensor.reset(env_ids=slice(2, None))
+    torch.testing.assert_close(sensor.drift.torch[:2], drift[:2])
+    assert not torch.allclose(sensor.drift.torch[2], drift[2])
+
+
 def test_bvh_refit_tracks_moving_geometry(sim):
     """Sliding a box under the sensor changes the hits, proving the BVH refits live.
 

From 16ee93260a9303710e0309d3b67a1b75db8bad5b Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 02:23:05 -0700
Subject: [PATCH 04/48] Read history_order as a group setting in the Warp
 observation manager

#8079 added ObservationGroupCfg.history_order and taught the stable
observation manager to skip it as a group setting. The Warp manager's skip
list was not updated, so it parsed the field as a term and raised a
TypeError for every observation group, and no manager-based Warp
environment could be constructed.

Skip history_order like the other group settings. The Warp manager raises
NotImplementedError for any term with history, and history_order only
takes effect with history, so a time-ordered group cannot be silently
reordered wrong.
---
 .../jichuanh-warp-scene-mask-native.fixed.rst |  3 ++
 .../managers/observation_manager.py           |  1 +
 .../test/managers/test_observation_manager.py | 37 +++++++++++++++++++
 3 files changed, 41 insertions(+)
 create mode 100644 source/isaaclab_experimental/test/managers/test_observation_manager.py

diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
index 2d0ece0abb4e..84ea86a8e137 100644
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
@@ -1,2 +1,5 @@
 * Fixed :meth:`~isaaclab_experimental.envs.InteractiveSceneWarp.reset` with ``env_mask`` resetting deformable
   objects, surface grippers, and sensors in every environment, and skipping cable objects.
+* Fixed the Warp :class:`~isaaclab_experimental.managers.ObservationManager` rejecting every observation group with
+  ``Configuration for the term 'history_order' is not of type ObservationTermCfg``, which stopped all manager-based
+  Warp environments from constructing. ``history_order`` is now read as a group setting, as in the stable manager.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
index 7bed3b315aa6..051aaa8a6ff6 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
@@ -670,6 +670,7 @@ def _prepare_terms(self):  # noqa: C901
                     "concatenate_terms",
                     "history_length",
                     "flatten_history_dim",
+                    "history_order",
                     "concatenate_dim",
                 ]:
                     continue
diff --git a/source/isaaclab_experimental/test/managers/test_observation_manager.py b/source/isaaclab_experimental/test/managers/test_observation_manager.py
new file mode 100644
index 000000000000..3094a9863da0
--- /dev/null
+++ b/source/isaaclab_experimental/test/managers/test_observation_manager.py
@@ -0,0 +1,37 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tests for the Warp observation manager configuration parsing."""
+
+from types import SimpleNamespace
+
+import warp as wp
+from isaaclab_experimental.managers import ObservationManager
+
+from isaaclab.managers import ObservationGroupCfg, ObservationTermCfg
+from isaaclab.utils import configclass
+
+
+def constant_obs(env, out: wp.array, out_dim: int):
+    """Warp-first observation term; the manager reads its width from ``out_dim``."""
+    out.fill_(1.0)
+
+
+@configclass
+class PolicyCfg(ObservationGroupCfg):
+    """Group with two terms and every group-level setting at its default."""
+
+    first = ObservationTermCfg(func=constant_obs, params={"out_dim": 2})
+    second = ObservationTermCfg(func=constant_obs, params={"out_dim": 3})
+
+
+def test_group_settings_are_not_parsed_as_terms():
+    """Every :class:`ObservationGroupCfg` field, including ``history_order``, is a group setting, not a term."""
+    env = SimpleNamespace(num_envs=4, device="cpu", sim=SimpleNamespace(is_playing=lambda: True))
+
+    manager = ObservationManager({"policy": PolicyCfg()}, env)
+
+    assert manager.active_terms["policy"] == ["first", "second"]
+    assert manager.group_obs_dim["policy"] == (5,)

From 7949cb98617a36d57a0239595cd70359833423fb Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 03:34:56 -0700
Subject: [PATCH 05/48] Write fixed tendon masks on Newton articulations
 without host sync

The fixed-tendon set_*_mask methods and
write_fixed_tendon_properties_to_sim_mask converted masks to indices with
torch.nonzero and delegated to the index writers, so they synchronized the
host and could not be captured in a CUDA graph.

Write the masked cells in-kernel instead: scalar and full data go into the
staging buffers through masked kernels, and the property writer copies the
masked cells of each staging buffer into its simulation binding. The mask
compaction helpers have no remaining callers and are removed.
---
 .../jichuanh-warp-scene-mask-native.fixed.rst |   3 +
 .../assets/articulation/articulation.py       | 106 ++++++++++--------
 .../assets/articulation/kernels.py            |  20 ++++
 .../test/assets/test_articulation.py          |  57 ++++++++++
 4 files changed, 141 insertions(+), 45 deletions(-)

diff --git a/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst b/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
index 67e8df01a9eb..78b02f8e956f 100644
--- a/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
+++ b/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
@@ -6,3 +6,6 @@
 * Fixed the ``write_body_*_to_sim_mask`` methods of :class:`~isaaclab_newton.assets.RigidObjectCollection`
   converting masks to indices on the host, which prevented capturing them in CUDA graphs as documented. They now
   write through masked kernels.
+* Fixed the fixed-tendon ``set_*_mask`` methods and
+  :meth:`~isaaclab_newton.assets.Articulation.write_fixed_tendon_properties_to_sim_mask` converting masks to indices
+  on the host, which prevented capturing them in CUDA graphs. They now write through masked kernels.
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
index 26006675ba9f..f7e9af453a3d 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
@@ -2703,13 +2703,8 @@ def set_fixed_tendon_stiffness_mask(
             fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used.
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
-        # Resolve masks.
-        env_ids = self._resolve_env_mask(env_mask)
-        fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask)
-        self.set_fixed_tendon_stiffness_index(
-            stiffness=self._select_full_data(stiffness, env_ids, fixed_tendon_ids),
-            fixed_tendon_ids=fixed_tendon_ids,
-            env_ids=env_ids,
+        self._set_fixed_tendon_buffer_mask(
+            stiffness, self.data._fixed_tendon_stiffness, fixed_tendon_mask, env_mask, "stiffness"
         )
 
     def set_fixed_tendon_damping_index(
@@ -2800,13 +2795,8 @@ def set_fixed_tendon_damping_mask(
             fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used.
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
-        # Resolve masks.
-        env_ids = self._resolve_env_mask(env_mask)
-        fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask)
-        self.set_fixed_tendon_damping_index(
-            damping=self._select_full_data(damping, env_ids, fixed_tendon_ids),
-            fixed_tendon_ids=fixed_tendon_ids,
-            env_ids=env_ids,
+        self._set_fixed_tendon_buffer_mask(
+            damping, self.data._fixed_tendon_damping, fixed_tendon_mask, env_mask, "damping"
         )
 
     def set_fixed_tendon_limit_stiffness_index(
@@ -2918,12 +2908,10 @@ def set_fixed_tendon_position_limit_mask(
             fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used.
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
-        env_ids = self._resolve_env_mask(env_mask)
-        fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask)
-        self.set_fixed_tendon_position_limit_index(
-            limit=self._select_full_data(limit, env_ids, fixed_tendon_ids),
-            fixed_tendon_ids=fixed_tendon_ids,
-            env_ids=env_ids,
+        if isinstance(limit, float):
+            raise ValueError("Fixed tendon position limits must be a tensor or array, not a float.")
+        self._set_fixed_tendon_buffer_mask(
+            limit, self.data._fixed_tendon_pos_limits, fixed_tendon_mask, env_mask, "limit"
         )
 
     def set_fixed_tendon_rest_length_index(
@@ -3124,10 +3112,22 @@ def write_fixed_tendon_properties_to_sim_mask(
             fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are updated.
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
-        self.write_fixed_tendon_properties_to_sim_index(
-            fixed_tendon_ids=self._resolve_fixed_tendon_mask(fixed_tendon_mask),
-            env_ids=self._resolve_env_mask(env_mask),
-        )
+        env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK)
+        fixed_tendon_mask = self._resolve_mask(fixed_tendon_mask, self._ALL_FIXED_TENDON_MASK)
+        for staged, sim_bind in (
+            (self.data._fixed_tendon_stiffness, self.data._sim_bind_fixed_tendon_stiffness),
+            (self.data._fixed_tendon_damping, self.data._sim_bind_fixed_tendon_damping),
+            (self.data._fixed_tendon_pos_limits, self.data._sim_bind_fixed_tendon_pos_limits),
+        ):
+            wp.launch(
+                articulation_kernels.copy_2d_data_with_mask,
+                dim=staged.shape,
+                inputs=[staged, env_mask, fixed_tendon_mask],
+                outputs=[sim_bind],
+                device=self.device,
+            )
+        # the solver keeps its own copy of the tendon properties and only re-reads them when notified
+        SimulationManager.add_model_change(ModelFlags.TENDON_PROPERTIES)
 
     def set_spatial_tendon_stiffness_index(
         self,
@@ -3801,33 +3801,49 @@ def _resolve_spatial_tendon_ids(
             return self._ALL_SPATIAL_TENDON_INDICES
         return spatial_tendon_ids
 
-    def _resolve_env_mask(self, env_mask: wp.array | torch.Tensor | None) -> wp.array | torch.Tensor:
-        """Resolve an environment mask to environment indices."""
-        return self._ALL_INDICES if env_mask is None else self._mask_to_ids(env_mask)
-
-    def _resolve_fixed_tendon_mask(self, fixed_tendon_mask: wp.array | torch.Tensor | None) -> wp.array | torch.Tensor:
-        """Resolve a fixed tendon mask to fixed tendon indices."""
-        return self._ALL_FIXED_TENDON_INDICES if fixed_tendon_mask is None else self._mask_to_ids(fixed_tendon_mask)
-
-    @staticmethod
-    def _mask_to_ids(mask: wp.array | torch.Tensor) -> torch.Tensor:
-        """Convert a boolean mask to int32 indices."""
-        mask = wp.to_torch(mask) if isinstance(mask, wp.array) else mask
-        return torch.nonzero(mask)[:, 0].to(torch.int32)
-
     @staticmethod
     def _to_torch_ids(ids: wp.array | torch.Tensor) -> torch.Tensor:
         """View resolved indices as a torch tensor for indexing."""
         return (wp.to_torch(ids) if isinstance(ids, wp.array) else ids).long()
 
-    def _select_full_data(
-        self, data: float | torch.Tensor | wp.array, env_ids: wp.array | torch.Tensor, ids: wp.array | torch.Tensor
-    ) -> float | torch.Tensor:
-        """Select the rows and columns of full (num_instances, num_items, ...) data given by the indices."""
+    def _set_fixed_tendon_buffer_mask(
+        self,
+        data: float | torch.Tensor | wp.array,
+        buffer: wp.array,
+        fixed_tendon_mask: wp.array | torch.Tensor | None,
+        env_mask: wp.array | torch.Tensor | None,
+        name: str,
+    ) -> None:
+        """Write full (num_instances, num_fixed_tendons) data into the masked cells of a fixed tendon buffer.
+
+        Args:
+            data: Scalar or full data. A torch tensor for a ``wp.vec2f`` buffer has a trailing dimension of 2.
+            buffer: Fixed tendon staging buffer. Shape is (num_instances, num_fixed_tendons).
+            fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used.
+            env_mask: Environment mask. If None, then all the instances are updated.
+            name: Name of the data, used in shape errors.
+        """
+        env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK)
+        fixed_tendon_mask = self._resolve_mask(fixed_tendon_mask, self._ALL_FIXED_TENDON_MASK)
         if isinstance(data, float):
-            return data
-        data = wp.to_torch(data) if isinstance(data, wp.array) else data
-        return data[self._to_torch_ids(env_ids)[:, None], self._to_torch_ids(ids)]
+            wp.launch(
+                articulation_kernels.float_data_to_buffer_with_mask,
+                dim=buffer.shape,
+                inputs=[data, env_mask, fixed_tendon_mask],
+                outputs=[buffer],
+                device=self.device,
+            )
+            return
+        if isinstance(data, torch.Tensor):
+            data = wp.from_torch(data, dtype=buffer.dtype)
+        self.assert_shape_and_dtype(data, buffer.shape, buffer.dtype, name)
+        wp.launch(
+            articulation_kernels.copy_2d_data_with_mask,
+            dim=buffer.shape,
+            inputs=[data, env_mask, fixed_tendon_mask],
+            outputs=[buffer],
+            device=self.device,
+        )
 
     def _resolve_mask(self, mask: wp.array | torch.Tensor | None, full_mask: wp.array) -> wp.array:
         """Resolve a mask to a warp array.
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py
index a530cbe6df2a..f423c2558138 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py
@@ -542,6 +542,26 @@ def float_data_to_buffer_with_mask(
         out_data[i, j] = in_data
 
 
+@wp.kernel
+def copy_2d_data_with_mask(
+    in_data: wp.array2d(dtype=Any),
+    env_mask: wp.array(dtype=wp.bool),
+    item_mask: wp.array(dtype=wp.bool),
+    out_data: wp.array2d(dtype=Any),
+):
+    """Copy the masked cells of a 2D per-environment buffer of any element type.
+
+    Args:
+        in_data: Input array. Shape is (num_envs, num_items).
+        env_mask: Input array of environment mask. Shape is (num_envs,).
+        item_mask: Input array of item mask, e.g. fixed tendons. Shape is (num_items,).
+        out_data: Output array where the selected cells are written. Shape is (num_envs, num_items).
+    """
+    i, j = wp.tid()
+    if env_mask[i] and item_mask[j]:
+        out_data[i, j] = in_data[i, j]
+
+
 @wp.kernel
 def update_soft_joint_pos_limits(
     joint_pos_limits: wp.array2d(dtype=wp.vec2f),
diff --git a/source/isaaclab_newton/test/assets/test_articulation.py b/source/isaaclab_newton/test/assets/test_articulation.py
index 2bb728645b39..94d1bd2307ef 100644
--- a/source/isaaclab_newton/test/assets/test_articulation.py
+++ b/source/isaaclab_newton/test/assets/test_articulation.py
@@ -1651,6 +1651,63 @@ def test_fixed_tendon_properties_reach_solver(sim, num_articulations, device, ar
     np.testing.assert_allclose(solver_model.tendon_damping.numpy(), expected_damping)
 
 
+@pytest.mark.parametrize("num_articulations", [2])
+@pytest.mark.parametrize("device", ["cuda:0"])
+@pytest.mark.parametrize("articulation_type", ["shadow_hand"])
+def test_fixed_tendon_mask_writes_replay_from_cuda_graph_on_selected_cells(
+    sim, num_articulations, device, articulation_type
+):
+    """Fixed tendon mask setters and writer record into a CUDA graph and, on replay, write only the selected cells.
+
+    The mask variants are the documented graph-capturable path. The inputs are filled only after capture, so the
+    replay must read them at launch time.
+    """
+    articulation, _ = generate_articulation(
+        generate_articulation_cfg(articulation_type=articulation_type), num_articulations, device
+    )
+    replicate(sim.get_clone_plan())
+    sim.reset()
+    num_tendons = articulation.num_fixed_tendons
+    env_mask = wp.array([False, True], dtype=wp.bool, device=device)
+    tendon_mask = torch.zeros(num_tendons, dtype=torch.bool, device=device)
+    tendon_mask[1] = True
+    tendon_mask = wp.from_torch(tendon_mask)
+    selected = torch.zeros(num_articulations, num_tendons, 1, dtype=torch.bool, device=device)
+    selected[1, 1] = True
+
+    stiffness = torch.zeros(num_articulations, num_tendons, device=device)
+    damping = torch.zeros_like(stiffness)
+    limits = torch.zeros(num_articulations, num_tendons, 2, device=device)
+    data = articulation.data
+    before = {
+        "stiffness": data.fixed_tendon_stiffness.torch.clone(),
+        "damping": data.fixed_tendon_damping.torch.clone(),
+        "pos_limits": data.fixed_tendon_pos_limits.torch.clone(),
+    }
+    with wp.ScopedCapture(device) as capture:
+        articulation.set_fixed_tendon_stiffness_mask(
+            stiffness=stiffness, fixed_tendon_mask=tendon_mask, env_mask=env_mask
+        )
+        articulation.set_fixed_tendon_damping_mask(damping=damping, fixed_tendon_mask=tendon_mask, env_mask=env_mask)
+        articulation.set_fixed_tendon_position_limit_mask(
+            limit=limits, fixed_tendon_mask=tendon_mask, env_mask=env_mask
+        )
+        articulation.write_fixed_tendon_properties_to_sim_mask(fixed_tendon_mask=tendon_mask, env_mask=env_mask)
+    stiffness.fill_(12.0)
+    damping.fill_(3.0)
+    limits[..., 0] = -0.1
+    limits[..., 1] = 0.2
+    wp.capture_launch(capture.graph)
+
+    written = {"stiffness": stiffness, "damping": damping, "pos_limits": limits}
+    for name, value in written.items():
+        after = getattr(data, f"fixed_tendon_{name}").torch
+        expected = torch.where(
+            selected, value.view(*selected.shape[:2], -1), before[name].view(*selected.shape[:2], -1)
+        )
+        torch.testing.assert_close(after.view_as(expected), expected, msg=name)
+
+
 @pytest.mark.parametrize("device", ["cpu"])
 @pytest.mark.parametrize("add_ground_plane", [True])
 @pytest.mark.parametrize("articulation_type", ["anymal"])

From 12bd5a9598357b3dda6d054dea1d21ee8b70826d Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 03:43:37 -0700
Subject: [PATCH 06/48] Make masked native-actuator and ray-caster resets
 capturable

A masked Articulation.reset on the Newton-native actuator path allocated a
per-DOF mask for every actuator on each call, and a masked ray-caster
reset sampled drift with torch ops and uploaded its ranges, so neither
could run inside a CUDA graph captured on Warp's stream.

NewtonActuatorAdapter preallocates its per-DOF reset masks. BaseRayCaster
resamples masked drift in a Warp kernel from per-environment random states
seeded from torch's global generator at initialization, with the ranges
passed as launch arguments; index and slice resets keep the torch path.
---
 ...ichuanh-warp-scene-mask-native.changed.rst |  9 ++-
 .../isaaclab/actuators/newton/adapter.py      |  5 +-
 .../sensors/ray_caster/base_ray_caster.py     | 49 +++++++++------
 .../isaaclab/sensors/ray_caster/kernels.py    | 62 ++++++++++++++++++
 .../test/utils/actuator_equivalence.py        | 63 ++++++++++---------
 .../assets/test_newton_actuators_newton.py    | 22 +++++++
 .../sensors/test_newton_raycast_sensor.py     | 18 +++++-
 7 files changed, 174 insertions(+), 54 deletions(-)

diff --git a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
index bc83b269855f..ea53a4476e4e 100644
--- a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
+++ b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
@@ -1,3 +1,6 @@
-* Changed masked :meth:`~isaaclab.sensors.ray_caster.BaseRayCaster.reset` to resample drift without a host
-  synchronization. When both ``env_ids`` and ``env_mask`` are given, the mask selects the resampled environments,
-  matching :meth:`~isaaclab.sensors.SensorBase.reset`.
+* Changed masked :meth:`~isaaclab.sensors.ray_caster.BaseRayCaster.reset` to resample drift in a Warp kernel from
+  per-environment random states seeded from torch's global generator, so it neither synchronizes the host nor
+  allocates and can be captured in a CUDA graph. When both ``env_ids`` and ``env_mask`` are given, the mask selects
+  the resampled environments, matching :meth:`~isaaclab.sensors.SensorBase.reset`.
+* Changed :meth:`~isaaclab.actuators.newton.NewtonActuatorAdapter.reset` to reuse preallocated per-DOF masks, so
+  masked resets of Newton-native actuator state do not allocate and can be captured in a CUDA graph.
diff --git a/source/isaaclab/isaaclab/actuators/newton/adapter.py b/source/isaaclab/isaaclab/actuators/newton/adapter.py
index 6b4a58e2f66f..961543fa6fda 100644
--- a/source/isaaclab/isaaclab/actuators/newton/adapter.py
+++ b/source/isaaclab/isaaclab/actuators/newton/adapter.py
@@ -105,6 +105,8 @@ def __init__(
 
         self._states_a = [act.state() for act in actuators]
         self._states_b = [act.state() for act in actuators]
+        # Per-DOF reset masks, rebuilt in full on every partial reset so masked resets do not allocate.
+        self._reset_dof_masks = [wp.zeros(act.indices.shape[0], dtype=wp.bool, device=device) for act in actuators]
 
         # Pre-clamp computed effort buffer. Each Newton actuator scatter-adds
         # its raw controller output to ``sim_control.joint_computed_f`` when
@@ -217,8 +219,7 @@ def reset(
                 )
                 wp.launch(set_mask_kernel, dim=num_envs, inputs=[env_mask, indices], device=self._device)
 
-        for act, sa, sb in zip(self.actuators, self._states_a, self._states_b):
-            per_dof_mask = wp.zeros(act.indices.shape[0], dtype=wp.bool, device=self._device)
+        for act, sa, sb, per_dof_mask in zip(self.actuators, self._states_a, self._states_b, self._reset_dof_masks):
             wp.launch(
                 build_per_dof_env_mask_kernel,
                 dim=act.indices.shape[0],
diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py
index 283690beba4a..2968954b91cb 100644
--- a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py
+++ b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py
@@ -116,27 +116,23 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None
         self._drift_sampled |= sample_drift or sample_ray_cast_drift
         if not self._drift_sampled:
             return
-        if sample_ray_cast_drift and ray_cast_range_list != self._ray_cast_drift_range_list:
-            # Upload the ranges to the device only when the configuration changes.
-            self._ray_cast_drift_range_list = ray_cast_range_list
-            self._ray_cast_drift_ranges = torch.tensor(ray_cast_range_list, device=self.device)
         if env_mask is not None:
-            # Resample every environment and keep the new values where the mask is set, so the
-            # masked reset needs no host synchronization.
-            mask = wp.to_torch(env_mask).view(-1, 1)
-            drift = self.drift.torch
-            new_drift = torch.zeros_like(drift)
-            if sample_drift:
-                new_drift.uniform_(*self.cfg.drift_range)
-            drift.copy_(torch.where(mask, new_drift, drift))
-            ray_cast_drift = self.ray_cast_drift.torch
-            new_ray_cast_drift = torch.zeros_like(ray_cast_drift)
-            if sample_ray_cast_drift:
-                ranges = self._ray_cast_drift_ranges
-                new_ray_cast_drift = math_utils.sample_uniform(
-                    ranges[:, 0], ranges[:, 1], tuple(ray_cast_drift.shape), device=self.device
-                )
-            ray_cast_drift.copy_(torch.where(mask, new_ray_cast_drift, ray_cast_drift))
+            # Resample in a kernel from per-environment random states with the ranges as launch
+            # arguments, so the masked reset neither synchronizes nor allocates and can be captured.
+            # A zero range draws zero, which clears drift that was sampled before it was disabled.
+            wp.launch(
+                ray_caster_kernels.resample_drift_masked_kernel,
+                dim=self._view_count,
+                inputs=[
+                    env_mask,
+                    wp.vec2f(*self.cfg.drift_range),
+                    wp.vec3f(*(low for low, _ in ray_cast_range_list)),
+                    wp.vec3f(*(high for _, high in ray_cast_range_list)),
+                    self._drift_rng_state,
+                ],
+                outputs=[self.drift.warp, self.ray_cast_drift.warp],
+                device=self._device,
+            )
             return
         # determine the selected batch size
         if env_ids is not None:
@@ -152,6 +148,10 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None
             self.drift.torch[env_ids] = 0.0
         # resample the ray cast drift
         if sample_ray_cast_drift:
+            # Upload the ranges to the device only when the configuration changes.
+            if ray_cast_range_list != self._ray_cast_drift_range_list:
+                self._ray_cast_drift_range_list = ray_cast_range_list
+                self._ray_cast_drift_ranges = torch.tensor(ray_cast_range_list, device=self.device)
             ranges = self._ray_cast_drift_ranges
             self.ray_cast_drift.torch[env_ids] = math_utils.sample_uniform(
                 ranges[:, 0], ranges[:, 1], (num_envs_ids, 3), device=self.device
@@ -255,6 +255,15 @@ def _initialize_rays_impl(self):
         # Drift buffers are warp-first; reset uses explicit .torch views for sampling.
         self.drift = ProxyArray(wp.zeros(self._view_count, dtype=wp.vec3f, device=self._device))
         self.ray_cast_drift = ProxyArray(wp.zeros(self._view_count, dtype=wp.vec3f, device=self._device))
+        # Per-environment random states for masked drift resampling, seeded from torch's global generator.
+        self._drift_rng_state = wp.empty(self._view_count, dtype=wp.uint32, device=self._device)
+        wp.launch(
+            ray_caster_kernels.init_rng_state_kernel,
+            dim=self._view_count,
+            inputs=[int(torch.randint(0, 2**31 - 1, (1,)))],
+            outputs=[self._drift_rng_state],
+            device=self._device,
+        )
 
         # World-frame ray buffers
         self._ray_starts_w = wp.empty((self._view_count, self.num_rays), dtype=wp.vec3f, device=self._device)
diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py b/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py
index 4bcd10aba339..a337fef6f21e 100644
--- a/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py
+++ b/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py
@@ -170,6 +170,68 @@ def apply_z_drift_kernel(
     ray_hits[env, ray] = wp.vec3f(hit[0], hit[1], hit[2] + ray_cast_drift[env][2])
 
 
+@wp.kernel(enable_backward=False)
+def init_rng_state_kernel(
+    # input
+    seed: wp.int32,
+    # output
+    rng_state: wp.array(dtype=wp.uint32),
+):
+    """Initialize one random number generator state per environment.
+
+    Launch with dim=(num_envs,).
+
+    Args:
+        seed: Seed shared by all environments; each environment uses its index as the offset.
+        rng_state: Random number generator states. Shape is (num_envs,).
+    """
+    env = wp.tid()
+    rng_state[env] = wp.rand_init(seed, env)
+
+
+@wp.kernel(enable_backward=False)
+def resample_drift_masked_kernel(
+    # input
+    env_mask: wp.array(dtype=wp.bool),
+    drift_range: wp.vec2f,
+    ray_cast_drift_min: wp.vec3f,
+    ray_cast_drift_max: wp.vec3f,
+    # input/output
+    rng_state: wp.array(dtype=wp.uint32),
+    # output
+    drift: wp.array(dtype=wp.vec3f),
+    ray_cast_drift: wp.array(dtype=wp.vec3f),
+):
+    """Resample the position drift and ray-cast drift of the masked environments.
+
+    Launch with dim=(num_envs,).
+
+    Args:
+        env_mask: Boolean mask for which environments to resample. Shape is (num_envs,).
+        drift_range: Uniform range of every position drift component [m].
+        ray_cast_drift_min: Per-axis lower bound of the ray-cast drift [m].
+        ray_cast_drift_max: Per-axis upper bound of the ray-cast drift [m].
+        rng_state: Random number generator states, advanced for the resampled environments. Shape is (num_envs,).
+        drift: Position drift [m]. Shape is (num_envs,).
+        ray_cast_drift: Ray-cast drift [m]. Shape is (num_envs,).
+    """
+    env = wp.tid()
+    if not env_mask[env]:
+        return
+    state = rng_state[env]
+    drift[env] = wp.vec3f(
+        wp.randf(state, drift_range[0], drift_range[1]),
+        wp.randf(state, drift_range[0], drift_range[1]),
+        wp.randf(state, drift_range[0], drift_range[1]),
+    )
+    ray_cast_drift[env] = wp.vec3f(
+        wp.randf(state, ray_cast_drift_min[0], ray_cast_drift_max[0]),
+        wp.randf(state, ray_cast_drift_min[1], ray_cast_drift_max[1]),
+        wp.randf(state, ray_cast_drift_min[2], ray_cast_drift_max[2]),
+    )
+    rng_state[env] = state
+
+
 @wp.kernel(enable_backward=False)
 def fill_ray_hits_distance_inf_kernel(
     # input
diff --git a/source/isaaclab/isaaclab/test/utils/actuator_equivalence.py b/source/isaaclab/isaaclab/test/utils/actuator_equivalence.py
index a03c6b7fa24a..140ef1362da9 100644
--- a/source/isaaclab/isaaclab/test/utils/actuator_equivalence.py
+++ b/source/isaaclab/isaaclab/test/utils/actuator_equivalence.py
@@ -393,6 +393,39 @@ def _reset_env(self, articulation, selector: str) -> None:
             env_mask[self.RESET_ENV] = True
             articulation.reset(env_mask=wp.from_torch(env_mask))
 
+    def _stateful_pairs(self, adapter) -> list:
+        """Return ``(actuator, state)`` pairs of the adapter's delayed actuators."""
+        return [
+            (act, st)
+            for act, st in zip(adapter.actuators, adapter._states_a)
+            if st is not None and getattr(st, "delay_state", None) is not None
+        ]
+
+    def _assert_delay_state_cleared_for_reset_env_only(self, adapter, stateful_pairs) -> None:
+        """Assert ``num_pushes`` is zero for :attr:`RESET_ENV`'s DOFs and positive for every other env's DOFs."""
+        # Map each entry of ``act.indices`` to its env via the adapter's
+        # per-env DOF count. On the Newton backend the adapter is model-wide
+        # (includes free-joint DOFs on floating-base articulations); on
+        # PhysX it is per-articulation — ``adapter.num_joints`` is the
+        # correct stride in both cases.
+        for act, state in stateful_pairs:
+            pushes_after = state.delay_state.num_pushes.numpy()
+            indices_np = act.indices.numpy()
+            for i, global_dof in enumerate(indices_np):
+                env = int(global_dof) // adapter.num_joints
+                if env == self.RESET_ENV:
+                    self.assertEqual(
+                        int(pushes_after[i]),
+                        0,
+                        f"DOF {i} (env {env}) should be reset to 0, got {pushes_after[i]}",
+                    )
+                else:
+                    self.assertGreater(
+                        int(pushes_after[i]),
+                        0,
+                        f"DOF {i} (env {env}) was NOT in the reset selection but num_pushes is 0",
+                    )
+
     def test_newton_state_reset_isolated_to_reset_env(self):
         """Newton: ``num_pushes`` zeroes for env 0's DOFs only after a partial reset of env 0."""
         ctx, sim, articulation = self._build_and_warm(use_newton_actuators=True)
@@ -403,11 +436,7 @@ def test_newton_state_reset_isolated_to_reset_env(self):
                 with self.subTest(selector=selector):
                     self._warm(sim, articulation)
                     # Find a DelayedPD actuator (it's the only one with delay_state).
-                    stateful_pairs = [
-                        (act, st)
-                        for act, st in zip(adapter.actuators, adapter._states_a)
-                        if st is not None and getattr(st, "delay_state", None) is not None
-                    ]
+                    stateful_pairs = self._stateful_pairs(adapter)
                     self.assertGreater(
                         len(stateful_pairs), 0, "expected at least one DelayedPD actuator with delay_state"
                     )
@@ -420,29 +449,7 @@ def test_newton_state_reset_isolated_to_reset_env(self):
                         )
 
                     self._reset_env(articulation, selector)
-
-                    # Map each entry of ``act.indices`` to its env via the adapter's
-                    # per-env DOF count. On the Newton backend the adapter is model-wide
-                    # (includes free-joint DOFs on floating-base articulations); on
-                    # PhysX it is per-articulation — ``adapter.num_joints`` is the
-                    # correct stride in both cases.
-                    for act, state in stateful_pairs:
-                        pushes_after = state.delay_state.num_pushes.numpy()
-                        indices_np = act.indices.numpy()
-                        for i, global_dof in enumerate(indices_np):
-                            env = int(global_dof) // adapter.num_joints
-                            if env == self.RESET_ENV:
-                                self.assertEqual(
-                                    int(pushes_after[i]),
-                                    0,
-                                    f"DOF {i} (env {env}) should be reset to 0, got {pushes_after[i]}",
-                                )
-                            else:
-                                self.assertGreater(
-                                    int(pushes_after[i]),
-                                    0,
-                                    f"DOF {i} (env {env}) was NOT in the reset selection but num_pushes is 0",
-                                )
+                    self._assert_delay_state_cleared_for_reset_env_only(adapter, stateful_pairs)
         finally:
             ctx.__exit__(None, None, None)
 
diff --git a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py
index 690c9cf0d582..0482629d5bed 100644
--- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py
+++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py
@@ -698,6 +698,28 @@ def _make_articulation(self) -> Articulation:
     def _get_adapter(self, articulation):
         return SimulationManager._adapter
 
+    def test_newton_masked_reset_replays_from_cuda_graph_without_allocation(self):
+        """A masked reset records into a CUDA graph without allocating and, on replay, clears only the masked env."""
+        ctx, sim, articulation = self._build_and_warm(use_newton_actuators=True)
+        try:
+            adapter = self._get_adapter(articulation)
+            env_mask = torch.zeros(self.NUM_ENVS, dtype=torch.bool, device=articulation.device)
+            env_mask[self.RESET_ENV] = True
+            env_mask = wp.from_torch(env_mask)
+            # Allocating inside the capture raises while the memory pool is disabled.
+            with wp.ScopedMempool(articulation.device, False), wp.ScopedCapture(articulation.device) as capture:
+                articulation.reset(env_mask=env_mask)
+            self._warm(sim, articulation)
+            stateful_pairs = self._stateful_pairs(adapter)
+            self.assertGreater(len(stateful_pairs), 0, "expected at least one DelayedPD actuator with delay_state")
+            for _, state in stateful_pairs:
+                self.assertTrue((state.delay_state.num_pushes.numpy() > 0).all(), "the replay must do the reset")
+
+            wp.capture_launch(capture.graph)
+            self._assert_delay_state_cleared_for_reset_env_only(adapter, stateful_pairs)
+        finally:
+            ctx.__exit__(None, None, None)
+
 
 # ---------------------------------------------------------------------------
 # RemotizedPD actuator: PD + delay + position-based clamping lookup table
diff --git a/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py b/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py
index d73fa8695f8e..9996cf39f108 100644
--- a/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py
+++ b/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py
@@ -199,7 +199,7 @@ def test_legacy_multi_mesh_tracks_ad_hoc_regex_target(sim):
 
 @pytest.mark.parametrize("sim", [pytest.param(False, id="eager")], indirect=True)
 def test_partial_reset_resamples_drift_of_selected_envs_only(sim):
-    """Mask and slice resets resample the drift of the selected environments and keep the others."""
+    """Mask and slice resets resample the drift of the selected environments only, also on graph replay."""
     scene_cfg = RaycastTestSceneCfg(num_envs=3)
     scene_cfg.raycast.drift_range = (0.01, 0.05)
     scene_cfg.raycast.ray_cast_drift_range = {"z": (0.1, 0.2)}
@@ -223,6 +223,22 @@ def test_partial_reset_resamples_drift_of_selected_envs_only(sim):
     torch.testing.assert_close(sensor.drift.torch[:2], drift[:2])
     assert not torch.allclose(sensor.drift.torch[2], drift[2])
 
+    # The masked reset records into a CUDA graph without allocating (the disabled memory pool makes an
+    # allocation raise), and every replay draws new drift for the masked environment only.
+    env_mask = wp.array([True, False, False], dtype=wp.bool, device=sensor.device)
+    with wp.ScopedMempool(sensor.device, False), wp.ScopedCapture(sensor.device) as capture:
+        sensor.reset(env_mask=env_mask)
+    for _ in range(2):
+        drift = sensor.drift.torch.clone()
+        ray_cast_drift = sensor.ray_cast_drift.torch.clone()
+        wp.capture_launch(capture.graph)
+        torch.testing.assert_close(sensor.drift.torch[1:], drift[1:])
+        torch.testing.assert_close(sensor.ray_cast_drift.torch[1:], ray_cast_drift[1:])
+        assert not torch.allclose(sensor.drift.torch[0], drift[0])
+        assert not torch.allclose(sensor.ray_cast_drift.torch[0, 2], ray_cast_drift[0, 2])
+    assert ((sensor.drift.torch >= 0.01) & (sensor.drift.torch <= 0.05)).all()
+    assert ((sensor.ray_cast_drift.torch[:, 2] >= 0.1) & (sensor.ray_cast_drift.torch[:, 2] <= 0.2)).all()
+
 
 def test_bvh_refit_tracks_moving_geometry(sim):
     """Sliding a box under the sensor changes the hits, proving the BVH refits live.

From eb458ca6a7cc9b540448a94d9f168c61983dee7c Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 03:58:52 -0700
Subject: [PATCH 07/48] Write Newton tendon targets every step so captured
 commands apply

The Newton fixed-tendon position-target setters raised a host-side dirty
flag, and write_data_to_sim scattered the buffered targets into MuJoCo's
controls only while it was set. A setter recorded in a CUDA graph raised
the flag once, at capture time, so later replays updated the target buffer
on the device but never reached the solver.

Write the buffered targets on every step whenever the articulation has
tendon actuators, and drop the flag from the Newton setters. The write
costs about 60 us per call on Shadow Hand; gating it on a device-side flag
would need Newton's control-array mapping re-implemented in a kernel.
---
 .../jichuanh-warp-scene-mask-native.fixed.rst |  3 +++
 .../assets/articulation/articulation.py       |  7 +++--
 .../test/assets/test_articulation.py          | 27 +++++++++++++++++++
 3 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst b/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
index 78b02f8e956f..ba9d43c347f8 100644
--- a/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
+++ b/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
@@ -9,3 +9,6 @@
 * Fixed the fixed-tendon ``set_*_mask`` methods and
   :meth:`~isaaclab_newton.assets.Articulation.write_fixed_tendon_properties_to_sim_mask` converting masks to indices
   on the host, which prevented capturing them in CUDA graphs. They now write through masked kernels.
+* Fixed fixed-tendon position targets commanded inside a captured CUDA graph reaching MuJoCo's controls only at
+  capture time. :meth:`~isaaclab_newton.assets.Articulation.write_data_to_sim` now writes the buffered targets on
+  every step when the articulation has tendon actuators, instead of behind a host-side flag.
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
index f7e9af453a3d..7f87fcaa5eea 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
@@ -437,9 +437,10 @@ def write_data_to_sim(self):
 
         # Tendon submission is solver-specific: MuJoCo drives tendons through actuator controls
         # outside the articulation view, so the manager owns how a buffered target reaches the solver.
-        if self._fixed_tendon_target_dirty:
+        # The buffer is written every step rather than behind a host-side flag, so a target
+        # commanded inside a captured CUDA graph reaches the solver on every replay.
+        if self._fixed_tendon_control is not None:
             self._fixed_tendon_control.write_data_to_sim(SimulationManager.get_control())
-            self._fixed_tendon_target_dirty = False
 
     def update(self, dt: float):
         """Updates the simulation data.
@@ -2989,7 +2990,6 @@ def set_fixed_tendon_position_target_index(
         self._fixed_tendon_control.set_position_target_index(
             target=target, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids
         )
-        self._fixed_tendon_target_dirty = True
 
     def set_fixed_tendon_offset_index(
         self,
@@ -3041,7 +3041,6 @@ def set_fixed_tendon_position_target_mask(
         self._fixed_tendon_control.set_position_target_mask(
             target=target, fixed_tendon_mask=fixed_tendon_mask, env_mask=env_mask
         )
-        self._fixed_tendon_target_dirty = True
 
     def set_fixed_tendon_offset_mask(
         self,
diff --git a/source/isaaclab_newton/test/assets/test_articulation.py b/source/isaaclab_newton/test/assets/test_articulation.py
index 94d1bd2307ef..fd337e1a1cb8 100644
--- a/source/isaaclab_newton/test/assets/test_articulation.py
+++ b/source/isaaclab_newton/test/assets/test_articulation.py
@@ -1651,6 +1651,33 @@ def test_fixed_tendon_properties_reach_solver(sim, num_articulations, device, ar
     np.testing.assert_allclose(solver_model.tendon_damping.numpy(), expected_damping)
 
 
+@pytest.mark.parametrize("num_articulations", [2])
+@pytest.mark.parametrize("device", ["cuda:0"])
+@pytest.mark.parametrize("articulation_type", ["shadow_hand"])
+def test_captured_tendon_target_reaches_control_on_every_replay(sim, num_articulations, device, articulation_type):
+    """A tendon target commanded inside a CUDA graph reaches MuJoCo's controls on every replay, not only at capture."""
+    from isaaclab_newton.physics.mjwarp_tendon_control import resolve_fixed_tendon_actuator_columns
+
+    articulation, _ = generate_articulation(
+        generate_articulation_cfg(articulation_type=articulation_type), num_articulations, device
+    )
+    replicate(sim.get_clone_plan())
+    sim.reset()
+    actuator_columns, _ = resolve_fixed_tendon_actuator_columns(articulation.root_view, SimulationManager.get_model())
+
+    target = torch.zeros(num_articulations, articulation.num_fixed_tendons, device=device)
+    with wp.ScopedCapture(device) as capture:
+        articulation.set_fixed_tendon_position_target_mask(target=target)
+    articulation.write_data_to_sim()
+    for value in (0.3, 0.6):
+        target.fill_(value)
+        wp.capture_launch(capture.graph)
+        articulation.write_data_to_sim()
+        ctrl = wp.to_torch(articulation.root_view.get_attribute("mujoco.ctrl", SimulationManager.get_control()))
+        tendon_ctrl = ctrl[:, 0, actuator_columns]
+        torch.testing.assert_close(tendon_ctrl, torch.full_like(tendon_ctrl, value))
+
+
 @pytest.mark.parametrize("num_articulations", [2])
 @pytest.mark.parametrize("device", ["cuda:0"])
 @pytest.mark.parametrize("articulation_type", ["shadow_hand"])

From b358e1a73c910b23840278c6cc9d15082acbd4f2 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 02:54:05 -0700
Subject: [PATCH 08/48] Capture Warp frontend stages without an eager warm-up

Replace ManagerCallSwitch with an environment-owned WarpGraphCache that
follows the Newton core capture model: a stage is recorded through
NewtonQueries.capture_graph with no eager warm-up, so the first
execution is the launch of the new graph; recording starts at the first
step; graphs are keyed by the pointers and scalar values of the stage
arguments; and every graph is dropped when the physics backend rebinds
its buffers or stops, so the next step records against the new buffers.

Managers with a non-capturable term keep their stages eager through
register_capturability. ISAACLAB_WARP_CAPTURE=0 runs every stage
eagerly and ISAACLAB_SYNC_DEBUG=1 traps hidden host synchronizations in
eager stages. The direct environment uses the same cache in place of
capture_or_replay.
---
 .../isaaclab_experimental.utils.rst           |   8 -
 .../jichuanh-warp-frontend-cleanup.added.rst  |   4 +
 ...jichuanh-warp-frontend-cleanup.changed.rst |   6 +
 .../jichuanh-warp-frontend-cleanup.fixed.rst  |   2 +
 .../jichuanh-warp-frontend-cleanup.minor      |   0
 .../envs/direct_rl_env_warp.py                |  14 +-
 .../envs/manager_based_env_warp.py            | 153 ++------
 .../envs/manager_based_rl_env_warp.py         | 158 +++-----
 .../managers/manager_base.py                  |   8 +-
 .../isaaclab_experimental/utils/__init__.pyi  |   3 -
 .../utils/manager_call_switch.py              | 262 -------------
 .../utils/warp_graph_cache.py                 | 226 ++++++++---
 .../test/envs/mdp/test_capture_safety.py      |   6 +-
 .../test/envs/test_warp_env_capture.py        |  94 +++++
 .../test/utils/test_manager_call_switch.py    | 363 ------------------
 .../test/utils/test_warp_graph_cache.py       | 346 ++++++++---------
 16 files changed, 557 insertions(+), 1096 deletions(-)
 create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst
 create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst
 create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
 create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor
 delete mode 100644 source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py
 create mode 100644 source/isaaclab_experimental/test/envs/test_warp_env_capture.py
 delete mode 100644 source/isaaclab_experimental/test/utils/test_manager_call_switch.py

diff --git a/docs/source/api/lab_experimental/isaaclab_experimental.utils.rst b/docs/source/api/lab_experimental/isaaclab_experimental.utils.rst
index cd0989d66e31..4f00a75aeff9 100644
--- a/docs/source/api/lab_experimental/isaaclab_experimental.utils.rst
+++ b/docs/source/api/lab_experimental/isaaclab_experimental.utils.rst
@@ -13,15 +13,7 @@ The following classes are part of the public :mod:`isaaclab_experimental.utils`
 .. autosummary::
    :nosignatures:
 
-   ManagerCallMode
-   ManagerCallSwitch
    WarpGraphCache
 
-.. autoclass:: ManagerCallMode
-   :show-inheritance:
-
-.. autoclass:: ManagerCallSwitch
-   :show-inheritance:
-
 .. autoclass:: WarpGraphCache
    :show-inheritance:
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst
new file mode 100644
index 000000000000..13db10a0a56a
--- /dev/null
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst
@@ -0,0 +1,4 @@
+* Added the ``ISAACLAB_WARP_CAPTURE`` environment variable: ``ISAACLAB_WARP_CAPTURE=0`` runs every Warp
+  environment stage eagerly, for debugging and eager-versus-captured comparisons.
+* Added the ``ISAACLAB_SYNC_DEBUG`` environment variable: ``ISAACLAB_SYNC_DEBUG=1`` raises on a hidden
+  GPU-to-host synchronization inside an eagerly executed Warp environment stage.
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst
new file mode 100644
index 000000000000..3c2f2127e843
--- /dev/null
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst
@@ -0,0 +1,6 @@
+* **Breaking:** Removed ``ManagerCallSwitch``, ``ManagerCallMode`` and the ``MANAGER_CALL_CONFIG``
+  environment variable. Warp environments construct their Warp managers directly and keep a manager's
+  stages eager when one of its terms is not capturable.
+* **Breaking:** Replaced ``WarpGraphCache.capture_or_replay`` with
+  :meth:`~isaaclab_experimental.utils.WarpGraphCache.call`. Stages are recorded from the first environment
+  step without an eager warm-up, so a stage runs exactly once per call.
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
new file mode 100644
index 000000000000..3c4f09ee4e12
--- /dev/null
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
@@ -0,0 +1,2 @@
+* Fixed Warp environments replaying CUDA graphs that read freed simulation buffers after the simulation was
+  reset or stopped. The environments now record their stages again on the next step.
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
index 7fdd1ffe3562..4d095c031807 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
@@ -227,8 +227,9 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
         self.torch_reset_time_outs: torch.Tensor = None
         self.torch_episode_length_buf: torch.Tensor = None
 
-        # Warp CUDA graph cache for capture-or-replay
-        self._graph_cache = WarpGraphCache()
+        # Task stages run through this cache and record their CUDA graphs on the first step.
+        self._warp_graph_cache = WarpGraphCache(self.device)
+        self._warp_graph_cache.invalidate_on(self.sim.physics_manager)
 
         # setup the action and observation spaces for Gym
         self._configure_gym_env_spaces()
@@ -391,6 +392,8 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
             A tuple containing the observations, rewards, resets (terminated and truncated) and extras.
         """
 
+        # stages start recording at the first step after construction or an invalidation
+        self._warp_graph_cache.arm()
         action = action.to(self.device)
         # add action noise
         if self.cfg.action_noise_model:
@@ -413,7 +416,7 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
                 # set actions into buffers
                 # simulate
                 with Timer(name="apply_action", msg="Action processing step took:", enable=DEBUG_TIMERS):
-                    self._graph_cache.capture_or_replay("action", self.step_warp_action)
+                    self._warp_graph_cache.call("DirectAction_step", self.step_warp_action)
 
                 # write_data_to_sim runs outside the CUDA graph because _apply_actuator_model
                 # uses torch ops (wp.to_torch + torch arithmetic) that cross CUDA streams.
@@ -433,12 +436,12 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
 
         self.common_step_counter += 1  # total step (common for all envs)
         with Timer(name="end_pre_graph", msg="End pre-graph took:", enable=DEBUG_TIMERS):
-            self._graph_cache.capture_or_replay("end_pre", self._step_warp_end_pre)
+            self._warp_graph_cache.call("DirectEndPre_step", self._step_warp_end_pre)
         # write_data_to_sim runs uncaptured — it uses torch ops that cross CUDA streams.
         with Timer(name="write_data_to_sim_post", msg="Write data to sim (post-reset) took:", enable=DEBUG_TIMERS):
             self.scene.write_data_to_sim()
         with Timer(name="end_post_graph", msg="End post-graph took:", enable=DEBUG_TIMERS):
-            self._graph_cache.capture_or_replay("end_post", self._step_warp_end_post)
+            self._warp_graph_cache.call("DirectEndPost_step", self._step_warp_end_post)
 
         # Visualization hook — runs after CUDA graph scope. Override in subclass
         # to update markers or other non-graphable visual elements.
@@ -575,6 +578,7 @@ def close(self):
             if self.cfg.events:
                 del self.event_manager
             del self.scene
+            self._warp_graph_cache.close()
 
             # # clear callbacks and instance
             # if float(".".join(get_version()[2])) >= 5:
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
index 3e94486af610..704d101ad7ee 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
@@ -16,11 +16,9 @@
 
 # import builtins
 import contextlib
-import importlib
 import logging
 import warnings
 from collections.abc import Sequence
-from copy import deepcopy
 from typing import Any
 
 import torch
@@ -33,6 +31,7 @@
     export_articulations_data,
     export_scene_data,
 )
+from isaaclab.managers import RecorderManager
 from isaaclab.sim import SimulationContext
 from isaaclab.sim.utils import use_stage
 from isaaclab.utils import validate
@@ -40,8 +39,10 @@
 from isaaclab.utils.timer import Timer
 
 from isaaclab_experimental.envs.interactive_scene_warp import InteractiveSceneWarp as InteractiveScene
-from isaaclab_experimental.utils.manager_call_switch import ManagerCallMode, ManagerCallSwitch
+from isaaclab_experimental.managers import ActionManager, EventManager, ObservationManager
+from isaaclab_experimental.utils.torch_utils import clone_obs_buffer
 from isaaclab_experimental.utils.warp import resolve_1d_mask
+from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache
 
 # import logger
 logger = logging.getLogger(__name__)
@@ -89,11 +90,6 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
             )
         # initialize internal variables
         self._is_closed = False
-        # temporary debug runtime config for manager source/call switching.
-        cfg_source: dict | str | None = getattr(self.cfg, "manager_call_config", None)
-        max_modes: dict[str, int] | None = getattr(self.cfg, "manager_call_max_mode", None)
-        self._manager_call_switch = ManagerCallSwitch(cfg_source, max_modes=max_modes)
-        self._apply_manager_term_cfg_profile()
 
         # set the seed for the environment
         if self.cfg.seed is not None:
@@ -117,6 +113,10 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
         if "cuda" in self.device:
             torch.cuda.set_device(self.device)
 
+        # Manager stages run through this cache; managers register terms that are unsafe to capture.
+        self._warp_graph_cache = WarpGraphCache(self.device)
+        self._warp_graph_cache.invalidate_on(self.sim.physics_manager)
+
         # print useful information
         print("[INFO]: Base environment:")
         print(f"\tEnvironment device    : {self.device}")
@@ -170,7 +170,7 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
         # create event manager
         # note: this is needed here (rather than after simulation play) to allow USD-related randomization events
         #   that must happen before the simulation starts. Example: randomizing mesh scale
-        self.event_manager = self._manager_call_switch.resolve_manager_class("EventManager")(self.cfg.events, self)
+        self.event_manager = EventManager(self.cfg.events, self)
 
         # apply USD-related randomization events
         if "prestartup" in self.event_manager.available_modes:
@@ -373,22 +373,24 @@ def load_managers(self):
             :meth:`SimulationContext.reset_async` and it isn't possible to call async functions in the constructor.
 
         """
+        # recorded stages read buffers of the managers being replaced
+        self._warp_graph_cache.invalidate()
         # prepare the managers
         # -- event manager (we print it here to make the logging consistent)
         print("[INFO] Event Manager: ", self.event_manager)
         # -- recorder manager
-        self.recorder_manager = self._manager_call_switch.resolve_manager_class("RecorderManager")(
-            self.cfg.recorders, self
-        )
+        self.recorder_manager = RecorderManager(self.cfg.recorders, self)
         print("[INFO] Recorder Manager: ", self.recorder_manager)
         # -- action manager
-        self.action_manager = self._manager_call_switch.resolve_manager_class("ActionManager")(self.cfg.actions, self)
+        self.action_manager = ActionManager(self.cfg.actions, self)
         print("[INFO] Action Manager: ", self.action_manager)
         # -- observation manager
-        self.observation_manager = self._manager_call_switch.resolve_manager_class("ObservationManager")(
-            self.cfg.observations, self
-        )
+        self.observation_manager = ObservationManager(self.cfg.observations, self)
         print("[INFO] Observation Manager:", self.observation_manager)
+        # persistent action input: recorded stages read this pointer, so each step copies into it
+        self._action_in_wp = wp.zeros(
+            (self.num_envs, self.action_manager.total_action_dim), dtype=wp.float32, device=self.device
+        )
 
         # perform events at the start of the simulation
         # in-case a child implementation creates other managers, the randomization should happen
@@ -543,14 +545,14 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]:
         Returns:
             A tuple containing the observations and extras.
         """
+        # stages start recording at the first step after construction or an invalidation
+        self._warp_graph_cache.arm()
         # process actions
-        action_device = action.to(self.device)
-        if action_device.dtype != torch.float32:
-            action_device = action_device.float()
-        if not action_device.is_contiguous():
-            action_device = action_device.contiguous()
-        action_wp = wp.from_torch(action_device, dtype=wp.float32)
-        self.action_manager.process_action(action_wp)
+        action_device = action.to(device=self.device, dtype=torch.float32).contiguous()
+        wp.copy(self._action_in_wp, wp.from_torch(action_device, dtype=wp.float32))
+        self._warp_graph_cache.call(
+            "ActionManager_process_action", self.action_manager.process_action, action=self._action_in_wp
+        )
 
         self.recorder_manager.record_pre_step()
 
@@ -562,7 +564,7 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]:
         for _ in range(self.cfg.decimation):
             self._sim_step_counter += 1
             # set actions into buffers
-            self.action_manager.apply_action()
+            self._warp_graph_cache.call("ActionManager_apply_action", self.action_manager.apply_action)
             # set actions into simulator
             self.scene.write_data_to_sim()
             # simulate
@@ -577,10 +579,18 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]:
 
         # post-step: step interval event
         if "interval" in self.event_manager.available_modes:
-            self.event_manager.apply(mode="interval", dt=self.step_dt)
+            self._warp_graph_cache.call(
+                "EventManager_apply_interval", self.event_manager.apply, mode="interval", dt=self.step_dt
+            )
 
         # -- compute observations
-        self.obs_buf = self.observation_manager.compute(update_history=True)
+        self.obs_buf = self._warp_graph_cache.call(
+            "ObservationManager_compute_update_history",
+            self.observation_manager.compute,
+            update_history=True,
+            return_cloned_output=False,
+            output=clone_obs_buffer,
+        )
         self.recorder_manager.record_post_step()
 
         # return observations and extras
@@ -607,6 +617,7 @@ def close(self):
             del self.event_manager
             del self.recorder_manager
             del self.scene
+            self._warp_graph_cache.close()
 
             # self.sim.clear_all_callbacks()
             self.sim.clear_instance()
@@ -621,96 +632,6 @@ def close(self):
     Helper functions.
     """
 
-    def _resolve_stable_cfg_counterpart(self) -> ManagerBasedEnvCfg | None:
-        """Resolve a stable task config counterpart for the current experimental task config.
-
-        The lookup follows a module-name mirror convention:
-        ``isaaclab_tasks_experimental...`` -> ``isaaclab_tasks...`` with the same config class name.
-        """
-        cfg_cls = self.cfg.__class__
-        cfg_module_name = cfg_cls.__module__
-        if "isaaclab_tasks_experimental" not in cfg_module_name:
-            return None
-
-        stable_module_name = cfg_module_name.replace("isaaclab_tasks_experimental", "isaaclab_tasks", 1)
-        try:
-            stable_module = importlib.import_module(stable_module_name)
-        except Exception as exc:
-            logger.warning(
-                "Failed to import stable task cfg module '%s' for manager_call_config stable mode: %s",
-                stable_module_name,
-                exc,
-            )
-            return None
-
-        stable_cfg_cls = getattr(stable_module, cfg_cls.__name__, None)
-        if stable_cfg_cls is None:
-            logger.warning(
-                "Stable task cfg class '%s' not found in module '%s'.",
-                cfg_cls.__name__,
-                stable_module_name,
-            )
-            return None
-
-        try:
-            return stable_cfg_cls()
-        except Exception as exc:
-            logger.warning(
-                "Failed to instantiate stable task cfg '%s.%s': %s",
-                stable_module_name,
-                cfg_cls.__name__,
-                exc,
-            )
-            return None
-
-    def _apply_manager_term_cfg_profile(self) -> None:
-        """Align term configs with manager modes for stable manager selections.
-
-        When a manager is configured as STABLE (0), swap its corresponding config subtree
-        from the stable task counterpart to keep manager-term type/signature compatibility.
-        """
-        manager_to_cfg_attr = {
-            "ActionManager": "actions",
-            "ObservationManager": "observations",
-            "EventManager": "events",
-            "RecorderManager": "recorders",
-            "CommandManager": "commands",
-            "TerminationManager": "terminations",
-            "RewardManager": "rewards",
-            "CurriculumManager": "curriculum",
-        }
-
-        stable_manager_names = [
-            manager_name
-            for manager_name in manager_to_cfg_attr
-            if self._manager_call_switch.get_mode_for_manager(manager_name) == ManagerCallMode.STABLE
-        ]
-        if not stable_manager_names:
-            return
-
-        stable_cfg = self._resolve_stable_cfg_counterpart()
-        if stable_cfg is None:
-            logger.warning(
-                "Stable managers requested (%s), but no stable cfg counterpart could be resolved."
-                " Keeping experimental term configs.",
-                ", ".join(stable_manager_names),
-            )
-            return
-
-        replaced_items: list[str] = []
-        for manager_name, cfg_attr in manager_to_cfg_attr.items():
-            if self._manager_call_switch.get_mode_for_manager(manager_name) != ManagerCallMode.STABLE:
-                continue
-            if not hasattr(self.cfg, cfg_attr) or not hasattr(stable_cfg, cfg_attr):
-                continue
-            setattr(self.cfg, cfg_attr, deepcopy(getattr(stable_cfg, cfg_attr)))
-            replaced_items.append(f"{manager_name} -> cfg.{cfg_attr}")
-
-        if replaced_items:
-            print("[INFO] Applied stable term config profile for managers:")
-            for item in replaced_items:
-                print(f"  - {item}")
-
     def _reset_idx(self, env_ids: Sequence[int]):
         """Reset environments based on specified indices.
 
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
index b6b56abc09d5..6414f375cac2 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
@@ -25,11 +25,11 @@
 
 from isaaclab.envs.common import VecEnvStepReturn
 from isaaclab.envs.manager_based_rl_env_cfg import ManagerBasedRLEnvCfg
-from isaaclab.managers import CommandManager
+from isaaclab.managers import CommandManager, CurriculumManager
 from isaaclab.utils import index_fill_
 from isaaclab.utils.timer import Timer
 
-from isaaclab_experimental.utils.manager_call_switch import ManagerCallMode
+from isaaclab_experimental.managers import RewardManager, TerminationManager
 from isaaclab_experimental.utils.torch_utils import clone_obs_buffer
 
 from .manager_based_env_warp import ManagerBasedEnvWarp
@@ -119,11 +119,6 @@ def __init__(self, cfg: ManagerBasedRLEnvCfg, render_mode: str | None = None, **
         # TODO: termination manager provides the same mask, investigate whether this can be replaced.
         self.reset_mask_wp = wp.zeros(cfg.scene.num_envs, dtype=wp.bool, device=cfg.sim.device)
 
-        # Persistent action input buffer to keep pointer stable for captured graphs.
-        self._action_in_wp: wp.array = wp.zeros(
-            (self.num_envs, self.action_manager.total_action_dim), dtype=wp.float32, device=self.device
-        )
-
         # initialize data and constants
         # -- set the framerate of the gym video recorder wrapper so that the playback speed
         # of the produced video matches the simulation
@@ -163,7 +158,7 @@ def max_episode_length(self) -> int:
     def load_managers(self):
         # note: this order is important since observation manager needs to know the command and action managers
         # and the reward manager needs to know the termination manager
-        # -- command manager (stable impl — not routed through ManagerCallSwitch)
+        # -- command manager (stable implementation)
         self.command_manager = CommandManager(self.cfg.commands, self)
         print("[INFO] Command Manager: ", self.command_manager)
 
@@ -172,17 +167,13 @@ def load_managers(self):
 
         # prepare the managers
         # -- termination manager
-        self.termination_manager = self._manager_call_switch.resolve_manager_class("TerminationManager")(
-            self.cfg.terminations, self
-        )
+        self.termination_manager = TerminationManager(self.cfg.terminations, self)
         print("[INFO] Termination Manager: ", self.termination_manager)
         # -- reward manager (experimental fork; Warp-compatible rewards)
-        self.reward_manager = self._manager_call_switch.resolve_manager_class("RewardManager")(self.cfg.rewards, self)
+        self.reward_manager = RewardManager(self.cfg.rewards, self)
         print("[INFO] Reward Manager: ", self.reward_manager)
-        # -- curriculum manager
-        self.curriculum_manager = self._manager_call_switch.resolve_manager_class("CurriculumManager")(
-            self.cfg.curriculum, self
-        )
+        # -- curriculum manager (stable implementation)
+        self.curriculum_manager = CurriculumManager(self.cfg.curriculum, self)
         print("[INFO] Curriculum Manager: ", self.curriculum_manager)
 
         # setup the action and observation spaces for Gym
@@ -217,7 +208,7 @@ def invalidate_wp_graphs(self) -> None:
 
         Call this if the captured launch topology changes (e.g. different term list, shapes, etc.).
         """
-        self._manager_call_switch.invalidate_graphs()
+        self._warp_graph_cache.invalidate()
 
     def step_warp_termination_compute(self) -> None:
         """Captured stage: compute terminations (env-step frequency)."""
@@ -245,19 +236,20 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         Returns:
             A tuple containing the observations, rewards, resets (terminated and truncated) and extras.
         """
+        # stages start recording at the first step after construction or an invalidation
+        self._warp_graph_cache.arm()
         # process actions
         # NOTE: keep a persistent action input buffer for graph pointer stability.
         # IMPORTANT: Do NOT re-wrap/replace the `wp.array` used by captured graphs each step.
         # Instead, copy the latest actions into the persistent buffer.
         with Timer(name="action_preprocess", msg="Action preprocessing took:", enable=DEBUG_TIMER_STEP, time_unit="us"):
-            if self._action_in_wp is None:
-                raise RuntimeError("Action buffer not initialized. Call reset() before step().")
-            action_device = action.to(self.device)
+            action_device = action.to(device=self.device, dtype=torch.float32).contiguous()
             wp.copy(self._action_in_wp, wp.from_torch(action_device, dtype=wp.float32))
 
-        self._manager_call_switch.call_stage(
-            stage="ActionManager_process_action",
-            warp_call={"fn": self.action_manager.process_action, "kwargs": {"action": self._action_in_wp}},
+        self._warp_graph_cache.call(
+            "ActionManager_process_action",
+            self.action_manager.process_action,
+            action=self._action_in_wp,
             timer=DEBUG_TIMER_STEP,
         )
 
@@ -271,16 +263,14 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         for _ in range(self.cfg.decimation):
             self._sim_step_counter += 1
             # set actions into buffers
-            self._manager_call_switch.call_stage(
-                stage="ActionManager_apply_action",
-                warp_call={"fn": self.action_manager.apply_action},
-                timer=DEBUG_TIMER_STEP,
-            )
-            self._manager_call_switch.call_stage(
-                stage="Scene_write_data_to_sim",
-                warp_call={"fn": self.scene.write_data_to_sim},
-                timer=DEBUG_TIMER_STEP,
+            self._warp_graph_cache.call(
+                "ActionManager_apply_action", self.action_manager.apply_action, timer=DEBUG_TIMER_STEP
             )
+            # scene writes cross the actuator and asset host boundaries, so they stay eager
+            with Timer(
+                name="Scene_write_data_to_sim", msg="Scene write took:", enable=DEBUG_TIMER_STEP, time_unit="us"
+            ):
+                self.scene.write_data_to_sim()
 
             # simulate
             with Timer(name="simulate", msg="Newton simulation step took:", enable=DEBUG_TIMER_STEP, time_unit="us"):
@@ -306,22 +296,19 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         self.common_step_counter += 1  # total step (common for all envs)
 
         # -- post-processing (termination + reward) as independently configurable stages
-        self._manager_call_switch.call_stage(
-            stage="TerminationManager_compute",
-            warp_call={"fn": self.step_warp_termination_compute},
-            timer=DEBUG_TIMER_STEP,
+        self._warp_graph_cache.call(
+            "TerminationManager_compute", self.step_warp_termination_compute, timer=DEBUG_TIMER_STEP
         )
-        self.reward_buf = self._manager_call_switch.call_stage(
-            stage="RewardManager_compute",
-            warp_call={"fn": self.reward_manager.compute, "kwargs": {"dt": float(self.step_dt)}},
-            timer=DEBUG_TIMER_STEP,
+        self.reward_buf = self._warp_graph_cache.call(
+            "RewardManager_compute", self.reward_manager.compute, dt=float(self.step_dt), timer=DEBUG_TIMER_STEP
         )
 
         if len(self.recorder_manager.active_terms) > 0:
             # update observations for recording if needed
-            self._manager_call_switch.call_stage(
-                stage="ObservationManager_compute_no_history",
-                warp_call={"fn": self.observation_manager.compute, "kwargs": {"return_cloned_output": False}},
+            self._warp_graph_cache.call(
+                "ObservationManager_compute_no_history",
+                self.observation_manager.compute,
+                return_cloned_output=False,
                 timer=DEBUG_TIMER_STEP,
             )
             self.recorder_manager.record_post_step()
@@ -336,14 +323,7 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
             enable=DEBUG_TIMER_STEP,
             time_unit="us",
         ):
-            # Keep the reset-mask handoff fully in Warp when experimental termination buffers exist.
-            # Stable termination manager path exposes torch-only dones/reset buffers.
-            termination_manager_mode = self._manager_call_switch.get_mode_for_manager("TerminationManager")
-            if termination_manager_mode == ManagerCallMode.STABLE:
-                # copy still needed as mask will be used if manager is set to mode > 0
-                wp.copy(self.reset_mask_wp, wp.from_torch(self.reset_buf, dtype=wp.bool))
-            else:
-                wp.copy(self.reset_mask_wp, self.termination_manager.dones_wp)
+            wp.copy(self.reset_mask_wp, self.termination_manager.dones_wp)
             reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1)
         if len(reset_env_ids) > 0:
             # trigger recorder terms for pre-reset calls
@@ -370,21 +350,22 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
 
         # -- step interval events
         if "interval" in self.event_manager.available_modes:
-            self._manager_call_switch.call_stage(
-                stage="EventManager_apply_interval",
-                warp_call={"fn": self.event_manager.apply, "kwargs": {"mode": "interval", "dt": float(self.step_dt)}},
+            self._warp_graph_cache.call(
+                "EventManager_apply_interval",
+                self.event_manager.apply,
+                mode="interval",
+                dt=float(self.step_dt),
                 timer=DEBUG_TIMER_STEP,
             )
 
         # -- compute observations
         # note: done after reset to get the correct observations for reset envs
-        self.obs_buf = self._manager_call_switch.call_stage(
-            stage="ObservationManager_compute_update_history",
-            warp_call={
-                "fn": self.observation_manager.compute,
-                "kwargs": {"update_history": True, "return_cloned_output": False},
-                "output": lambda r: clone_obs_buffer(r),
-            },
+        self.obs_buf = self._warp_graph_cache.call(
+            "ObservationManager_compute_update_history",
+            self.observation_manager.compute,
+            update_history=True,
+            return_cloned_output=False,
+            output=clone_obs_buffer,
             timer=DEBUG_TIMER_STEP,
         )
         # return observations, rewards, resets and extras
@@ -523,24 +504,17 @@ def _reset_idx(
             self.curriculum_manager.compute(env_ids=env_ids)
 
         # reset the internal buffers of the scene elements
-        self._manager_call_switch.call_stage(
-            stage="Scene_reset",
-            warp_call={"fn": self.scene.reset, "kwargs": {"env_ids": env_ids, "env_mask": env_mask}},
-            timer=DEBUG_TIMER_RESET,
-        )
+        with Timer(name="Scene_reset", msg="Scene reset took:", enable=DEBUG_TIMER_RESET, time_unit="us"):
+            self.scene.reset(env_ids=env_ids, env_mask=env_mask)
 
         if "reset" in self.event_manager.available_modes:
             self._global_env_step_count_wp.fill_(self._sim_step_counter // self.cfg.decimation)
-            self._manager_call_switch.call_stage(
-                stage="EventManager_apply_reset",
-                warp_call={
-                    "fn": self.event_manager.apply,
-                    "kwargs": {
-                        "mode": "reset",
-                        "env_mask_wp": env_mask,
-                        "global_env_step_count": self._global_env_step_count_wp,
-                    },
-                },
+            self._warp_graph_cache.call(
+                "EventManager_apply_reset",
+                self.event_manager.apply,
+                mode="reset",
+                env_mask_wp=env_mask,
+                global_env_step_count=self._global_env_step_count_wp,
                 timer=DEBUG_TIMER_RESET,
             )
 
@@ -548,20 +522,14 @@ def _reset_idx(
         # this returns a dictionary of information which is stored in the extras
         # note: This is order-sensitive! Certain things need be reset before others.
         # -- observation manager + action + reward managers
-        obs_info = self._manager_call_switch.call_stage(
-            stage="ObservationManager_reset",
-            warp_call={"fn": self.observation_manager.reset, "kwargs": {"env_mask": env_mask}},
-            timer=DEBUG_TIMER_RESET,
+        obs_info = self._warp_graph_cache.call(
+            "ObservationManager_reset", self.observation_manager.reset, env_mask=env_mask, timer=DEBUG_TIMER_RESET
         )
-        action_info = self._manager_call_switch.call_stage(
-            stage="ActionManager_reset",
-            warp_call={"fn": self.action_manager.reset, "kwargs": {"env_mask": env_mask}},
-            timer=DEBUG_TIMER_RESET,
+        action_info = self._warp_graph_cache.call(
+            "ActionManager_reset", self.action_manager.reset, env_mask=env_mask, timer=DEBUG_TIMER_RESET
         )
-        reward_info = self._manager_call_switch.call_stage(
-            stage="RewardManager_reset",
-            warp_call={"fn": self.reward_manager.reset, "kwargs": {"env_mask": env_mask}},
-            timer=DEBUG_TIMER_RESET,
+        reward_info = self._warp_graph_cache.call(
+            "RewardManager_reset", self.reward_manager.reset, env_mask=env_mask, timer=DEBUG_TIMER_RESET
         )
 
         # -- curriculum manager
@@ -575,17 +543,11 @@ def _reset_idx(
 
         # -- command + event + termination managers
         command_info = self.command_manager.reset(env_ids=env_ids)
-        event_info = self._manager_call_switch.call_stage(
-            stage="EventManager_reset",
-            warp_call={"fn": self.event_manager.reset, "kwargs": {"env_mask": env_mask}},
-            stable_call={"fn": self.event_manager.reset, "kwargs": {"env_ids": env_ids}},
-            timer=DEBUG_TIMER_RESET,
+        event_info = self._warp_graph_cache.call(
+            "EventManager_reset", self.event_manager.reset, env_mask=env_mask, timer=DEBUG_TIMER_RESET
         )
-        termination_info = self._manager_call_switch.call_stage(
-            stage="TerminationManager_reset",
-            warp_call={"fn": self.termination_manager.reset, "kwargs": {"env_mask": env_mask}},
-            stable_call={"fn": self.termination_manager.reset, "kwargs": {"env_ids": env_ids}},
-            timer=DEBUG_TIMER_RESET,
+        termination_info = self._warp_graph_cache.call(
+            "TerminationManager_reset", self.termination_manager.reset, env_mask=env_mask, timer=DEBUG_TIMER_RESET
         )
 
         # -- recorder manager
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py
index 4fca43ee07c8..3e3325138f3f 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py
@@ -411,11 +411,11 @@ def _resolve_common_term_cfg(self, term_name: str, term_cfg: ManagerTermBaseCfg,
                     f" and optional parameters: {args_with_defaults}, but received: {term_params}."
                 )
 
-        # register non-capturable terms with the call switch for mode=2 fallback
+        # a non-capturable term keeps every stage of this manager eager
         if not is_warp_capturable(term_cfg.func):
-            switch = getattr(self._env, "_manager_call_switch", None)
-            if switch is not None:
-                switch.register_manager_capturability(type(self).__name__, False)
+            graph_cache = getattr(self._env, "_warp_graph_cache", None)
+            if graph_cache is not None:
+                graph_cache.register_capturability(type(self).__name__, False)
 
         # process attributes at runtime
         # these properties are only resolvable once the simulation starts playing
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/__init__.pyi b/source/isaaclab_experimental/isaaclab_experimental/utils/__init__.pyi
index de2860732514..00660ba9775e 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/utils/__init__.pyi
+++ b/source/isaaclab_experimental/isaaclab_experimental/utils/__init__.pyi
@@ -4,8 +4,6 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 __all__ = [
-    "ManagerCallMode",
-    "ManagerCallSwitch",
     "WarpGraphCache",
     "clone_obs_buffer",
     "buffers",
@@ -14,7 +12,6 @@ __all__ = [
     "warp",
 ]
 
-from .manager_call_switch import ManagerCallMode, ManagerCallSwitch
 from .torch_utils import clone_obs_buffer
 from .warp_graph_cache import WarpGraphCache
 from . import buffers, modifiers, noise, warp
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py b/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py
deleted file mode 100644
index 89924c789bc3..000000000000
--- a/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py
+++ /dev/null
@@ -1,262 +0,0 @@
-# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
-# All rights reserved.
-#
-# SPDX-License-Identifier: BSD-3-Clause
-
-"""Manager call switch for routing manager stage calls through stable/warp/captured paths."""
-
-from __future__ import annotations
-
-import importlib
-import json
-import os
-from enum import IntEnum
-from typing import Any
-
-from isaaclab.utils.timer import Timer
-
-from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache
-
-
-class ManagerCallMode(IntEnum):
-    """Execution mode for manager stage calls.
-
-    * ``STABLE``  (0): Call stable Python manager implementations from :mod:`isaaclab.managers`.
-    * ``WARP_NOT_CAPTURED`` (1): Call Warp-compatible implementations without CUDA graph capture.
-    * ``WARP_CAPTURED`` (2): Call Warp implementations with CUDA graph capture/replay.
-    """
-
-    STABLE = 0
-    WARP_NOT_CAPTURED = 1
-    WARP_CAPTURED = 2
-
-
-class ManagerCallSwitch:
-    """Per-manager call switch for stable/warp/captured execution.
-
-    Routes each manager stage call through the configured execution path:
-    stable Python, Warp (eager), or Warp (captured CUDA graph). Optionally
-    wraps each call in a :class:`Timer` context for profiling.
-    """
-
-    DEFAULT_CONFIG: dict[str, int] = {"default": 2}
-    DEFAULT_KEY = "default"
-    MANAGER_NAMES: tuple[str, ...] = (
-        "ActionManager",
-        "ObservationManager",
-        "EventManager",
-        "RecorderManager",
-        "CommandManager",
-        "TerminationManager",
-        "RewardManager",
-        "CurriculumManager",
-        "Scene",
-    )
-    # FIXME: Scene_write_data_to_sim calls articulation._apply_actuator_model which
-    #  uses wp.to_torch + torch indexing -- not capture-safe on this branch.
-    #  Cap Scene stages to WARP_NOT_CAPTURED until the articulation layer is capture-ready.
-    MAX_MODE_OVERRIDES: dict[str, int] = {"Scene": ManagerCallMode.WARP_NOT_CAPTURED}
-
-    ENV_VAR = "MANAGER_CALL_CONFIG"
-    """Environment variable name for the JSON config string.
-
-    Example usage::
-
-        MANAGER_CALL_CONFIG='{"RewardManager": 0, "default": 2}' python train.py ...
-    """
-
-    def __init__(
-        self,
-        cfg_source: dict | str | None = None,
-        *,
-        max_modes: dict[str, int] | None = None,
-    ):
-        self._graph_cache = WarpGraphCache()
-        # Merge caller-supplied max_modes with the class-level MAX_MODE_OVERRIDES.
-        self._max_modes = dict(self.MAX_MODE_OVERRIDES)
-        if max_modes is not None:
-            self._max_modes.update(max_modes)
-        # Resolve config: prefer explicit cfg_source, fall back to env var.
-        if cfg_source is None:
-            cfg_source = os.environ.get(self.ENV_VAR)
-        self._cfg = self._load_cfg(cfg_source)
-        print("[INFO] ManagerCallSwitch configuration:")
-        print(f"  - {self.DEFAULT_KEY}: {self._cfg[self.DEFAULT_KEY]}")
-        for manager_name in self.MANAGER_NAMES:
-            mode = int(self.get_mode_for_manager(manager_name))
-            cap = self._max_modes.get(manager_name)
-            cap_str = f" (cap={cap})" if cap is not None else ""
-            print(f"  - {manager_name}: {mode}{cap_str}")
-
-    # ------------------------------------------------------------------
-    # Graph management
-    # ------------------------------------------------------------------
-
-    def invalidate_graphs(self) -> None:
-        """Invalidate cached capture graphs and their cached return values."""
-        self._graph_cache.invalidate()
-
-    # ------------------------------------------------------------------
-    # Stage dispatch
-    # ------------------------------------------------------------------
-
-    def call_stage(
-        self,
-        *,
-        stage: str,
-        warp_call: dict[str, Any],
-        stable_call: dict[str, Any] | None = None,
-        timer: bool = False,
-    ) -> Any:
-        """Run the stage according to configured mode, optionally wrapped in a :class:`Timer`.
-
-        A call spec dict supports the following keys:
-
-        * ``fn`` (required): The callable to invoke.
-        * ``args`` (optional): Positional arguments tuple.
-        * ``kwargs`` (optional): Keyword arguments dict.
-        * ``output`` (optional): A ``Callable[[Any], Any]`` that transforms the raw
-          return value into the final output. For captured stages the raw value is
-          ``None``. When omitted, the raw return value is used as-is.
-
-        Args:
-            stage: Stage identifier in the form ``"ManagerName_function_name"``.
-            warp_call: Call spec for the warp path (eager or captured).
-            stable_call: Call spec for the stable (torch) path. Defaults to ``None``.
-            timer: Whether to wrap execution in a :class:`Timer`. Defaults to ``True``
-                (controlled by the global :attr:`Timer.enable` class-level toggle).
-                Pass a module-level flag like ``TIMER_ENABLED_STEP`` to make timing
-                conditional on that flag.
-
-        Returns:
-            The (possibly transformed) return value of the stage.
-        """
-        with Timer(name=stage, msg=f"{stage} took:", enable=timer, time_unit="us"):
-            return self._dispatch(stage, stable_call, warp_call)
-
-    def _dispatch(
-        self,
-        stage: str,
-        stable_call: dict[str, Any] | None,
-        warp_call: dict[str, Any],
-    ) -> Any:
-        """Select call path based on mode, execute, and apply output."""
-        mode = self.get_mode_for_manager(self._manager_name_from_stage(stage))
-        if mode == ManagerCallMode.STABLE:
-            if stable_call is None:
-                raise ValueError(f"Stage '{stage}' is configured as STABLE (mode=0) but no stable_call was provided.")
-            call, result = stable_call, self._run_call(stable_call)
-        elif mode == ManagerCallMode.WARP_CAPTURED:
-            call, result = warp_call, self._wp_capture_or_launch(stage, warp_call)
-        else:
-            call, result = warp_call, self._run_call(warp_call)
-
-        output_fn = call.get("output")
-        return output_fn(result) if output_fn is not None else result
-
-    # ------------------------------------------------------------------
-    # Manager resolution
-    # ------------------------------------------------------------------
-
-    def _manager_name_from_stage(self, stage: str) -> str:
-        if "_" not in stage:
-            raise ValueError(f"Invalid stage '{stage}'. Expected '{{manager_name}}_{{function_name}}'.")
-        return stage.split("_", 1)[0]
-
-    def get_mode_for_manager(self, manager_name: str) -> ManagerCallMode:
-        """Return the resolved execution mode for the given manager.
-
-        Looks up the manager in the config dict, falls back to the default,
-        then caps by :attr:`_max_modes` (static overrides + dynamic registrations).
-        """
-        default_key = next(iter(self.DEFAULT_CONFIG))
-        mode_value = self._cfg.get(manager_name, self._cfg[default_key])
-        cap = self._max_modes.get(manager_name)
-        if cap is not None:
-            mode_value = min(mode_value, cap)
-        return ManagerCallMode(mode_value)
-
-    def resolve_manager_class(self, manager_name: str, mode_override: ManagerCallMode | int | None = None) -> type:
-        """Import and return the manager class for the configured mode."""
-        mode = self.get_mode_for_manager(manager_name) if mode_override is None else ManagerCallMode(mode_override)
-        module_name = "isaaclab.managers" if mode == ManagerCallMode.STABLE else "isaaclab_experimental.managers"
-        module = importlib.import_module(module_name)
-        if not hasattr(module, manager_name):
-            raise AttributeError(f"Manager '{manager_name}' not found in module '{module_name}'.")
-        return getattr(module, manager_name)
-
-    def register_manager_capturability(self, manager_name: str, capturable: bool) -> None:
-        """Register that a manager has non-capturable terms, capping its mode.
-
-        Called by :class:`ManagerBase` during term preparation when a term
-        is decorated with ``@warp_capturable(False)``.
-        """
-        if not capturable:
-            self._max_modes[manager_name] = min(
-                self._max_modes.get(manager_name, ManagerCallMode.WARP_CAPTURED),
-                ManagerCallMode.WARP_NOT_CAPTURED,
-            )
-
-    # ------------------------------------------------------------------
-    # Internal helpers
-    # ------------------------------------------------------------------
-
-    def _run_call(self, call: dict[str, Any]) -> Any:
-        """Execute a single call spec eagerly."""
-        return call["fn"](*call.get("args", ()), **call.get("kwargs", {}))
-
-    def _wp_capture_or_launch(self, stage: str, call: dict[str, Any]) -> Any:
-        """Capture Warp CUDA graph on first call, then replay.
-
-        Delegates to :class:`WarpGraphCache` which handles warm-up, capture,
-        caching the return value, and replay.
-        """
-        return self._graph_cache.capture_or_replay(
-            stage,
-            call["fn"],
-            args=call.get("args", ()),
-            kwargs=call.get("kwargs", {}),
-        )
-
-    def _load_cfg(self, cfg_source: dict | str | None) -> dict[str, int]:
-        if cfg_source is None:
-            cfg = dict(self.DEFAULT_CONFIG)
-        elif isinstance(cfg_source, dict):
-            cfg = dict(cfg_source)
-            if self.DEFAULT_KEY not in cfg:
-                cfg[self.DEFAULT_KEY] = self.DEFAULT_CONFIG[self.DEFAULT_KEY]
-        elif isinstance(cfg_source, str):
-            if cfg_source.strip() == "":
-                cfg = dict(self.DEFAULT_CONFIG)
-            else:
-                parsed = json.loads(cfg_source)
-                if not isinstance(parsed, dict):
-                    raise TypeError("manager_call_config must decode to a dict.")
-                cfg = dict(parsed)
-                if self.DEFAULT_KEY not in cfg:
-                    cfg[self.DEFAULT_KEY] = self.DEFAULT_CONFIG[self.DEFAULT_KEY]
-        else:
-            raise TypeError(f"cfg_source must be a dict, string, or None, got: {type(cfg_source)}")
-
-        # validation
-        for manager_name, mode_value in cfg.items():
-            if not isinstance(mode_value, int):
-                raise TypeError(
-                    f"manager_call_config value for '{manager_name}' must be int (0/1/2), got: {type(mode_value)}"
-                )
-            try:
-                ManagerCallMode(mode_value)
-            except ValueError as exc:
-                raise ValueError(
-                    f"Invalid manager_call_config value for '{manager_name}': {mode_value}. Expected 0/1/2."
-                ) from exc
-
-        # Apply max mode caps: bake caps into the resolved config so
-        # get_mode_for_manager never needs per-call branching.
-        default_mode = cfg[self.DEFAULT_KEY]
-        for name, max_mode in self._max_modes.items():
-            resolved = cfg.get(name, default_mode)
-            if resolved > max_mode:
-                cfg[name] = max_mode
-
-        return cfg
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py
index d446be28d099..647d8cb3f2b7 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py
@@ -3,79 +3,201 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
-"""Warp CUDA graph capture-or-replay utility."""
+"""CUDA graph capture and replay for Warp frontend stages."""
 
+from __future__ import annotations
+
+import os
 from collections.abc import Callable
 from typing import Any
 
+import torch
 import warp as wp
+from isaaclab_newton.physics import NewtonQueries
 
+from isaaclab.physics import PhysicsEvent, PhysicsManager
+from isaaclab.utils.timer import Timer
+from isaaclab.utils.version import has_kit
 
-class WarpGraphCache:
-    """Caches Warp CUDA graphs by stage name: captures on first call, replays after.
+CAPTURE_ENV_VAR = "ISAACLAB_WARP_CAPTURE"
+"""Set ``ISAACLAB_WARP_CAPTURE=0`` to run every stage eagerly (debugging and eager-vs-captured comparisons)."""
+
+SYNC_DEBUG_ENV_VAR = "ISAACLAB_SYNC_DEBUG"
+"""Set ``ISAACLAB_SYNC_DEBUG=1`` to run eager stage calls under ``torch.cuda.set_sync_debug_mode("error")``,
+so a hidden GPU-to-host synchronization inside a stage raises instead of stalling silently."""
 
-    On the very first call for a given stage, an **eager warm-up** run
-    executes *before* graph capture.  This lets one-time initialisation
-    code (memory allocations, torch dtype casts, ``hasattr`` guards, etc.)
-    run outside the capture context.  Only the steady-state kernel
-    launches are then recorded into the graph.
 
-    The return value from the capture run is cached and returned on every
-    subsequent replay, ensuring captured stages return the same references
-    (e.g. tensor views) as eager stages.
+class WarpGraphCache:
+    """Run Warp frontend stages eagerly or through cached CUDA graphs.
+
+    Graphs follow the physics backend's capture lifecycle:
 
-    Usage::
+    * A stage is recorded without an eager warm-up: its first execution is the launch of the
+      newly recorded graph, so a stateful stage advances exactly once per call. Stage owners
+      allocate their buffers before the first call.
+    * Recording waits until :meth:`arm`. Environments arm the cache on each step, so stages
+      called only while constructing or resetting the environment run eagerly until then.
+    * A graph is keyed by the pointers and scalar values of the stage arguments. Passing a
+      reallocated array, or a different scalar, records the stage again.
+    * :meth:`invalidate` drops every graph and holds recording until the next :meth:`arm`.
+      :meth:`invalidate_on` does so whenever the physics backend rebinds its buffers or stops.
 
-        cache = WarpGraphCache()
-        result = cache.capture_or_replay("my_stage", my_warp_function)
-        # uncaptured work here ...
-        result2 = cache.capture_or_replay("my_stage_post", my_other_function)
+    A stage group is the stage name up to its first underscore. A group registered as not
+    capturable, a CPU device, or ``ISAACLAB_WARP_CAPTURE=0`` runs its stages eagerly.
     """
 
-    def __init__(self):
-        self._graphs: dict[str, Any] = {}
-        self._results: dict[str, Any] = {}
+    def __init__(self, device: wp.DeviceLike = None):
+        """Initialize the cache.
+
+        Args:
+            device: Warp device that runs the stages. Defaults to the current Warp device.
+        """
+        self._device = wp.get_device(device)
+        self._enabled = self._device.is_cuda and os.environ.get(CAPTURE_ENV_VAR, "1") != "0"
+        self._sync_debug = os.environ.get(SYNC_DEBUG_ENV_VAR, "0") == "1"
+        self._armed = False
+        self._graphs: dict[str, tuple[tuple, wp.Graph, Any]] = {}
+        self._capturable: dict[str, bool] = {}
+        self._callback_handles = []
+
+    @property
+    def captured_stages(self) -> tuple[str, ...]:
+        """Stages currently backed by a recorded graph, sorted by name."""
+        return tuple(sorted(self._graphs))
 
-    def capture_or_replay(
+    def call(
         self,
         stage: str,
         fn: Callable[..., Any],
-        args: tuple = (),
-        kwargs: dict[str, Any] | None = None,
+        /,
+        *args: Any,
+        output: Callable[[Any], Any] | None = None,
+        timer: bool = False,
+        **kwargs: Any,
     ) -> Any:
-        """Capture *fn* into a CUDA graph on the first call, then replay.
+        """Run a stage eagerly or through its recorded graph.
 
         Args:
-            stage: Unique name identifying this captured scope.
-            fn: The callable to capture. Must contain only CUDA-graph-safe
-                operations (pure warp kernels, no Python-level branching on
-                GPU data).
-            args: Positional arguments forwarded to *fn*. Defaults to ``()``.
-            kwargs: Keyword arguments forwarded to *fn*. Defaults to ``None``.
+            stage: Stage identifier in the form ``"GroupName_function_name"``.
+            fn: Callable implementing the stage. Its GPU work must be capturable when the group is.
+            *args: Positional arguments forwarded to :paramref:`fn`.
+            output: Transform applied to the stage result on every call, outside the graph.
+            timer: Whether to time the stage.
+            **kwargs: Keyword arguments forwarded to :paramref:`fn`.
 
         Returns:
-            The cached return value from the first (capture) invocation.
+            The stage result, optionally transformed by :paramref:`output`. A replayed stage returns
+            the result of the call that recorded it.
+        """
+        with Timer(name=stage, msg=f"{stage} took:", enable=timer, time_unit="us"):
+            if not self._enabled or not self.is_capturable(stage.partition("_")[0]):
+                result = self._run_eager(fn, args, kwargs)
+            else:
+                key = _argument_key(args, kwargs)
+                recorded = self._graphs.get(stage)
+                if recorded is not None and recorded[0] == key:
+                    wp.capture_launch(recorded[1])
+                    result = recorded[2]
+                elif self._armed:
+                    graph, result = self._record(fn, args, kwargs)
+                    self._graphs[stage] = (key, graph, result)
+                    wp.capture_launch(graph)
+                else:
+                    result = self._run_eager(fn, args, kwargs)
+        return output(result) if output is not None else result
+
+    def arm(self) -> None:
+        """Allow capturable stages to record on their next call."""
+        self._armed = True
+
+    def invalidate(self, group: str | None = None) -> None:
+        """Drop recorded graphs.
+
+        Args:
+            group: Stage group whose graphs are dropped; they record again on their next armed call.
+                Defaults to None, which drops every graph and holds recording until the next :meth:`arm`.
+        """
+        stages = [stage for stage in self._graphs if group is None or stage.partition("_")[0] == group]
+        if not stages:
+            if group is None:
+                self._armed = False
+            return
+        # Graphs may still be executing; drain the device before releasing them.
+        wp.synchronize_device(self._device)
+        for stage in stages:
+            del self._graphs[stage]
+        if group is None:
+            self._armed = False
+
+    def invalidate_on(self, physics_manager: type[PhysicsManager]) -> None:
+        """Invalidate every graph when the physics backend rebinds its buffers or stops.
+
+        Args:
+            physics_manager: Physics manager class that dispatches the lifecycle events.
         """
-        if kwargs is None:
-            kwargs = {}
-        graph = self._graphs.get(stage)
-        if graph is not None:
-            wp.capture_launch(graph)
-            return self._results[stage]
-        # Warm-up: run eagerly to flush first-call allocations / hasattr guards.
-        fn(*args, **kwargs)
-        # Capture: allocations already done, only wp.launch calls are recorded.
-        with wp.ScopedCapture() as capture:
+        for event in (PhysicsEvent.PHYSICS_READY, PhysicsEvent.STOP):
+            self._callback_handles.append(physics_manager.register_callback(self._on_physics_event, event))
+
+    def close(self) -> None:
+        """Deregister the physics callbacks and drop every graph."""
+        for handle in self._callback_handles:
+            handle.deregister()
+        self._callback_handles.clear()
+        self.invalidate()
+
+    def register_capturability(self, group: str, capturable: bool) -> None:
+        """Register whether a stage group is safe to record.
+
+        One unsafe registration keeps the group eager.
+
+        Args:
+            group: Stage group identifier.
+            capturable: Whether the registered work is safe during CUDA graph capture.
+        """
+        self._capturable[group] = self._capturable.get(group, True) and capturable
+
+    def is_capturable(self, group: str) -> bool:
+        """Return whether a stage group is eligible for recording."""
+        return self._capturable.get(group, True)
+
+    def _on_physics_event(self, payload: Any) -> None:
+        self.invalidate()
+
+    def _record(self, fn: Callable[..., Any], args: tuple, kwargs: dict[str, Any]) -> tuple[wp.Graph, Any]:
+        """Record one stage call without executing it."""
+        result = None
+
+        def capture_target() -> None:
+            nonlocal result
             result = fn(*args, **kwargs)
-        self._graphs[stage] = capture.graph
-        self._results[stage] = result
-        return result
-
-    def invalidate(self, stage: str | None = None) -> None:
-        """Drop cached graph(s). If *stage* is ``None``, drop all."""
-        if stage is None:
-            self._graphs.clear()
-            self._results.clear()
-        else:
-            self._graphs.pop(stage, None)
-            self._results.pop(stage, None)
+
+        graph = NewtonQueries.capture_graph(str(self._device), capture_target, relaxed=has_kit())
+        return graph, result
+
+    def _run_eager(self, fn: Callable[..., Any], args: tuple, kwargs: dict[str, Any]) -> Any:
+        """Run a stage eagerly, under the synchronization trap when requested."""
+        if not self._sync_debug:
+            return fn(*args, **kwargs)
+        previous = torch.cuda.get_sync_debug_mode()
+        torch.cuda.set_sync_debug_mode("error")
+        try:
+            return fn(*args, **kwargs)
+        finally:
+            torch.cuda.set_sync_debug_mode(previous)
+
+
+def _argument_key(args: tuple, kwargs: dict[str, Any]) -> tuple:
+    """Key a stage call by the pointers of its arrays and the values of its scalars."""
+    return tuple(_value_key(value) for value in args) + tuple(
+        (name, _value_key(value)) for name, value in kwargs.items()
+    )
+
+
+def _value_key(value: Any) -> Any:
+    if isinstance(value, wp.array):
+        return value.ptr
+    if isinstance(value, torch.Tensor):
+        return value.data_ptr()
+    if value is None or isinstance(value, (bool, int, float, str)):
+        return value
+    return id(value)
diff --git a/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py b/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py
index ea69b3308fc4..f0afa42058e3 100644
--- a/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py
+++ b/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py
@@ -5,9 +5,9 @@
 
 """CUDA-graph capture safety for warp MDP terms.
 
-The manager call switch runs every manager graph-captured by default
-(``ManagerCallSwitch.DEFAULT_CONFIG``), so a warp MDP term that is not capture-safe is not an
-edge case — it is the normal execution path. A captured graph records launch topology and
+The Warp environments record every capturable manager stage into a CUDA graph by default
+(:class:`~isaaclab_experimental.utils.WarpGraphCache`), so a warp MDP term that is not capture-safe
+is not an edge case — it is the normal execution path. A captured graph records launch topology and
 pointer values and replays them, so a term must launch a fixed dim, touch only buffers that
 were allocated before capture and never move, and do all of its work in kernels: host-side
 Python between launches is simply not replayed.
diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_capture.py b/source/isaaclab_experimental/test/envs/test_warp_env_capture.py
new file mode 100644
index 000000000000..64bf18f84ef1
--- /dev/null
+++ b/source/isaaclab_experimental/test/envs/test_warp_env_capture.py
@@ -0,0 +1,94 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Eager and CUDA-graph execution of the Warp environments must produce the same rollout.
+
+Each task runs the same seeded action sequence twice, once with ``ISAACLAB_WARP_CAPTURE=0`` and
+once with stage capture on. For manager-based tasks, a hard :meth:`SimulationContext.reset` halfway
+through rebuilds the Newton model and reallocates every simulation buffer; recorded graphs still
+point at the freed buffers, so the captured rollout only stays equal when the environment records
+its stages again.
+"""
+
+from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg
+
+from isaaclab.sim import SimulationCfg
+from isaaclab.test.utils import launch_test_simulation
+
+launch_test_simulation(SimulationCfg(physics=NewtonCfg(solver_cfg=MJWarpSolverCfg())))
+
+import isaaclab_tasks_experimental  # noqa: F401
+import pytest
+import torch
+from isaaclab_experimental.envs.frontend import WarpFrontend
+from isaaclab_experimental.utils.warp_graph_cache import CAPTURE_ENV_VAR
+
+import isaaclab.sim as sim_utils
+
+import isaaclab_tasks  # noqa: F401
+from isaaclab_tasks.utils import resolve_task_config
+
+_NUM_ENVS = 32
+_STEPS = 80
+_REBIND_STEP = 40
+
+
+def _rollout(task_id: str, capture: bool, rebind: bool, monkeypatch: pytest.MonkeyPatch) -> dict:
+    """Run the seeded action sequence and return the rollout with the recorded stages."""
+    monkeypatch.setenv(CAPTURE_ENV_VAR, "1" if capture else "0")
+    env_cfg, _ = resolve_task_config(task_id, "", overrides=("physics=newton_mjwarp",))
+    env_cfg.seed = 7
+    env_cfg.scene.num_envs = _NUM_ENVS
+    # contacts are only reproducible between runs with deterministic contact ordering
+    env_cfg.sim.physics.deterministic_mode = "run_to_run"
+    env_cfg.sim.physics.solver_cfg.disable_sensors = True
+    sim_utils.create_new_stage()
+    env = WarpFrontend.build_env(env_cfg, task_id).unwrapped
+    rollout = {"obs": [], "reward": [], "terminated": [], "truncated": []}
+    try:
+        env.reset()
+        generator = torch.Generator().manual_seed(123)
+        # a per-env bias drives environments to terminate at different steps
+        bias = torch.rand((_NUM_ENVS, env.action_space.shape[-1]), generator=generator) * 2.0 - 1.0
+        for step in range(_STEPS):
+            if rebind and step == _REBIND_STEP:
+                rollout["stages_before_rebind"] = env._warp_graph_cache.captured_stages
+                env.sim.reset()
+                rollout["stages_at_rebind"] = env._warp_graph_cache.captured_stages
+                env.reset()
+            noise = torch.rand((_NUM_ENVS, bias.shape[1]), generator=generator) * 2.0 - 1.0
+            obs, reward, terminated, truncated, _ = env.step((bias + 0.5 * noise).to(env.device))
+            rollout["obs"].append(obs["policy"].clone())
+            rollout["reward"].append(reward.clone())
+            rollout["terminated"].append(terminated.clone())
+            rollout["truncated"].append(truncated.clone())
+        rollout["stages_at_end"] = env._warp_graph_cache.captured_stages
+    finally:
+        env.close()
+    for key in ("obs", "reward", "terminated", "truncated"):
+        rollout[key] = torch.stack(rollout[key])
+    return rollout
+
+
+# Direct tasks keep the simulation arrays they bind in ``__init__``, so they do not survive a rebind.
+@pytest.mark.parametrize(
+    ("task_id", "rebind"),
+    [("Isaac-Cartpole", True), ("Isaac-Cartpole-Direct", False)],
+    ids=["manager-cartpole", "direct-cartpole"],
+)
+def test_captured_rollout_matches_eager(task_id: str, rebind: bool, monkeypatch: pytest.MonkeyPatch):
+    eager = _rollout(task_id, capture=False, rebind=rebind, monkeypatch=monkeypatch)
+    captured = _rollout(task_id, capture=True, rebind=rebind, monkeypatch=monkeypatch)
+
+    assert eager["stages_at_end"] == ()
+    assert captured["stages_at_end"], "no stage was recorded; the comparison proves nothing"
+    if rebind:
+        assert captured["stages_at_rebind"] == (), "the rebind must drop graphs that read freed buffers"
+        assert captured["stages_at_end"] == captured["stages_before_rebind"]
+    for window in (slice(0, _REBIND_STEP), slice(_REBIND_STEP, _STEPS)):
+        dones = captured["terminated"][window] | captured["truncated"][window]
+        assert dones.any(), "each half must reset environments so the reset stages replay"
+    for key in ("obs", "reward", "terminated", "truncated"):
+        assert torch.equal(captured[key], eager[key]), f"{key} differs between eager and captured execution"
diff --git a/source/isaaclab_experimental/test/utils/test_manager_call_switch.py b/source/isaaclab_experimental/test/utils/test_manager_call_switch.py
deleted file mode 100644
index 8f4bb904c7c1..000000000000
--- a/source/isaaclab_experimental/test/utils/test_manager_call_switch.py
+++ /dev/null
@@ -1,363 +0,0 @@
-# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
-# All rights reserved.
-#
-# SPDX-License-Identifier: BSD-3-Clause
-
-"""Tests for ManagerCallSwitch config loading, mode resolution, and dispatch."""
-
-from __future__ import annotations
-
-import json
-import os
-import unittest
-
-import warp as wp
-from isaaclab_experimental.utils.manager_call_switch import ManagerCallMode, ManagerCallSwitch
-
-
-@wp.kernel
-def _add_one(a: wp.array(dtype=wp.float32), b: wp.array(dtype=wp.float32)):
-    i = wp.tid()
-    b[i] = a[i] + 1.0
-
-
-# ======================================================================
-# Config loading
-# ======================================================================
-
-
-class TestConfigLoading(unittest.TestCase):
-    """Tests for ManagerCallSwitch config parsing from dict, JSON, env var, and None."""
-
-    def test_none_uses_default(self):
-        switch = ManagerCallSwitch(cfg_source=None)
-        self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_CAPTURED)
-
-    def test_dict_config(self):
-        switch = ManagerCallSwitch(cfg_source={"default": 1})
-        self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_NOT_CAPTURED)
-
-    def test_dict_per_manager_override(self):
-        switch = ManagerCallSwitch(cfg_source={"default": 2, "RewardManager": 0})
-        self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.STABLE)
-        self.assertEqual(switch.get_mode_for_manager("ActionManager"), ManagerCallMode.WARP_CAPTURED)
-
-    def test_dict_without_default_key(self):
-        """A dict missing 'default' should inherit from DEFAULT_CONFIG."""
-        switch = ManagerCallSwitch(cfg_source={"RewardManager": 0})
-        # default should be 2 (from DEFAULT_CONFIG)
-        self.assertEqual(switch.get_mode_for_manager("ActionManager"), ManagerCallMode.WARP_CAPTURED)
-        self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.STABLE)
-
-    def test_json_string_config(self):
-        cfg_str = json.dumps({"default": 1, "ObservationManager": 0})
-        switch = ManagerCallSwitch(cfg_source=cfg_str)
-        self.assertEqual(switch.get_mode_for_manager("ObservationManager"), ManagerCallMode.STABLE)
-        self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_NOT_CAPTURED)
-
-    def test_empty_string_uses_default(self):
-        switch = ManagerCallSwitch(cfg_source="")
-        self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_CAPTURED)
-
-    def test_env_var_fallback(self):
-        """When cfg_source is None, should read from MANAGER_CALL_CONFIG env var."""
-        old = os.environ.get(ManagerCallSwitch.ENV_VAR)
-        try:
-            os.environ[ManagerCallSwitch.ENV_VAR] = json.dumps({"default": 0})
-            switch = ManagerCallSwitch(cfg_source=None)
-            self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.STABLE)
-        finally:
-            if old is None:
-                os.environ.pop(ManagerCallSwitch.ENV_VAR, None)
-            else:
-                os.environ[ManagerCallSwitch.ENV_VAR] = old
-
-    def test_invalid_json_raises(self):
-        with self.assertRaises(json.JSONDecodeError):
-            ManagerCallSwitch(cfg_source="not valid json")
-
-    def test_invalid_mode_value_raises(self):
-        with self.assertRaises(ValueError):
-            ManagerCallSwitch(cfg_source={"default": 99})
-
-    def test_non_int_mode_raises(self):
-        with self.assertRaises(TypeError):
-            ManagerCallSwitch(cfg_source={"default": "fast"})
-
-    def test_invalid_cfg_type_raises(self):
-        with self.assertRaises(TypeError):
-            ManagerCallSwitch(cfg_source=42)
-
-
-# ======================================================================
-# Mode resolution and capping
-# ======================================================================
-
-
-class TestModeResolution(unittest.TestCase):
-    """Tests for mode resolution, MAX_MODE_OVERRIDES, and dynamic capturability."""
-
-    def test_max_mode_override_caps_default(self):
-        """Scene is capped to WARP_NOT_CAPTURED by MAX_MODE_OVERRIDES."""
-        switch = ManagerCallSwitch(cfg_source={"default": 2})
-        # Scene has a static cap of WARP_NOT_CAPTURED
-        self.assertEqual(
-            switch.get_mode_for_manager("Scene"),
-            ManagerCallMode.WARP_NOT_CAPTURED,
-        )
-        # Other managers should not be capped
-        self.assertEqual(
-            switch.get_mode_for_manager("RewardManager"),
-            ManagerCallMode.WARP_CAPTURED,
-        )
-
-    def test_max_modes_kwarg(self):
-        """max_modes kwarg should cap managers beyond MAX_MODE_OVERRIDES."""
-        switch = ManagerCallSwitch(
-            cfg_source={"default": 2},
-            max_modes={"RewardManager": ManagerCallMode.WARP_NOT_CAPTURED},
-        )
-        self.assertEqual(
-            switch.get_mode_for_manager("RewardManager"),
-            ManagerCallMode.WARP_NOT_CAPTURED,
-        )
-
-    def test_register_manager_capturability_downgrades(self):
-        """register_manager_capturability(False) should cap a manager to WARP_NOT_CAPTURED."""
-        switch = ManagerCallSwitch(cfg_source={"default": 2})
-        self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_CAPTURED)
-
-        switch.register_manager_capturability("RewardManager", capturable=False)
-        self.assertEqual(
-            switch.get_mode_for_manager("RewardManager"),
-            ManagerCallMode.WARP_NOT_CAPTURED,
-        )
-
-    def test_register_capturability_true_is_noop(self):
-        """register_manager_capturability(True) should not change anything."""
-        switch = ManagerCallSwitch(cfg_source={"default": 2})
-        switch.register_manager_capturability("RewardManager", capturable=True)
-        self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_CAPTURED)
-
-    def test_register_capturability_does_not_upgrade(self):
-        """Once capped to NOT_CAPTURED, a later capturable=True registration must not lift the cap."""
-        switch = ManagerCallSwitch(cfg_source={"default": 2})
-        switch.register_manager_capturability("RewardManager", capturable=False)
-        switch.register_manager_capturability("RewardManager", capturable=True)
-        self.assertEqual(
-            switch.get_mode_for_manager("RewardManager"),
-            ManagerCallMode.WARP_NOT_CAPTURED,
-        )
-
-    def test_mode_0_not_affected_by_cap(self):
-        """A manager explicitly set to STABLE (0) should stay at 0 even with cap at 1."""
-        switch = ManagerCallSwitch(cfg_source={"default": 2, "RewardManager": 0})
-        switch.register_manager_capturability("RewardManager", capturable=False)
-        self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.STABLE)
-
-
-# ======================================================================
-# Stage dispatch
-# ======================================================================
-
-
-class TestStageDispatch(unittest.TestCase):
-    """Tests for call_stage routing through stable / warp-eager / warp-captured paths."""
-
-    def test_stable_mode_calls_stable_fn(self):
-        switch = ManagerCallSwitch(cfg_source={"default": 0})
-        called = {"stable": False, "warp": False}
-
-        def stable_fn():
-            called["stable"] = True
-            return "stable_result"
-
-        def warp_fn():
-            called["warp"] = True
-            return "warp_result"
-
-        result = switch.call_stage(
-            stage="RewardManager_compute",
-            warp_call={"fn": warp_fn},
-            stable_call={"fn": stable_fn},
-        )
-        self.assertTrue(called["stable"])
-        self.assertFalse(called["warp"])
-        self.assertEqual(result, "stable_result")
-
-    def test_stable_mode_without_stable_call_raises(self):
-        switch = ManagerCallSwitch(cfg_source={"default": 0})
-        with self.assertRaises(ValueError):
-            switch.call_stage(
-                stage="RewardManager_compute",
-                warp_call={"fn": lambda: None},
-                stable_call=None,
-            )
-
-    def test_warp_eager_mode_calls_warp_fn(self):
-        """WARP_NOT_CAPTURED calls the warp fn exactly once per call (no warm-up, no stable call)."""
-        switch = ManagerCallSwitch(cfg_source={"default": 1})
-        called = {"stable": 0, "warp": 0}
-
-        def stable_fn():
-            called["stable"] += 1
-            return "stable_result"
-
-        def warp_fn():
-            called["warp"] += 1
-            return "warp_result"
-
-        call = {"stage": "RewardManager_compute", "warp_call": {"fn": warp_fn}, "stable_call": {"fn": stable_fn}}
-        result = switch.call_stage(**call)
-        self.assertEqual(called, {"stable": 0, "warp": 1}, "Eager mode should call the warp fn exactly once")
-        self.assertEqual(result, "warp_result")
-
-        switch.call_stage(**call)
-        self.assertEqual(called["warp"], 2, "Eager mode should call fn again each time")
-
-    def test_output_transform(self):
-        """The 'output' key in call spec should transform the return value."""
-        switch = ManagerCallSwitch(cfg_source={"default": 1})
-
-        def warp_fn():
-            return [1, 2, 3]
-
-        result = switch.call_stage(
-            stage="RewardManager_compute",
-            warp_call={"fn": warp_fn, "output": len},
-        )
-        self.assertEqual(result, 3)
-
-    def test_args_and_kwargs_forwarded(self):
-        """call_stage should forward args and kwargs from the call spec."""
-        switch = ManagerCallSwitch(cfg_source={"default": 1})
-        received = {}
-
-        def warp_fn(a, b, key=None):
-            received["a"] = a
-            received["b"] = b
-            received["key"] = key
-
-        switch.call_stage(
-            stage="RewardManager_compute",
-            warp_call={"fn": warp_fn, "args": (1, 2), "kwargs": {"key": "val"}},
-        )
-        self.assertEqual(received, {"a": 1, "b": 2, "key": "val"})
-
-
-class TestStageDispatchCaptured(unittest.TestCase):
-    """Tests for WARP_CAPTURED mode (requires GPU)."""
-
-    def setUp(self):
-        self.device = "cuda:0"
-
-    def test_captured_warmup_call_count(self):
-        """WARP_CAPTURED first call should invoke fn exactly 2 times (warm-up + capture)."""
-        switch = ManagerCallSwitch(cfg_source={"default": 2})
-        src = wp.zeros(4, dtype=wp.float32, device=self.device)
-        dst = wp.zeros(4, dtype=wp.float32, device=self.device)
-        call_count = [0]
-
-        def warp_fn():
-            call_count[0] += 1
-            wp.launch(_add_one, dim=4, inputs=[src, dst], device=self.device)
-            return dst
-
-        # First call_stage: warm-up (1) + capture (2)
-        switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn})
-        self.assertEqual(call_count[0], 2, "First captured call should invoke fn twice (warm-up + capture)")
-
-        # Second call_stage: replay only — no new fn invocation
-        switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn})
-        self.assertEqual(call_count[0], 2, "Replay should not invoke fn again")
-
-        # Third call_stage: still replay
-        switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn})
-        self.assertEqual(call_count[0], 2)
-
-
-# ======================================================================
-# Graph invalidation
-# ======================================================================
-
-
-class TestGraphInvalidation(unittest.TestCase):
-    """Tests for invalidate_graphs on ManagerCallSwitch."""
-
-    def setUp(self):
-        self.device = "cuda:0"
-
-    def test_invalidate_forces_recapture(self):
-        switch = ManagerCallSwitch(cfg_source={"default": 2})
-        src = wp.full(4, value=1.0, dtype=wp.float32, device=self.device)
-        dst = wp.zeros(4, dtype=wp.float32, device=self.device)
-
-        call_count = [0]
-
-        def warp_fn():
-            call_count[0] += 1
-            wp.launch(_add_one, dim=4, inputs=[src, dst], device=self.device)
-            return dst
-
-        # First capture
-        switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn})
-        self.assertEqual(call_count[0], 2)  # warm-up + capture
-
-        # Replay — no new calls
-        switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn})
-        self.assertEqual(call_count[0], 2)
-
-        # Invalidate
-        switch.invalidate_graphs()
-
-        # Should re-warm-up + re-capture
-        switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn})
-        self.assertEqual(call_count[0], 4)
-
-
-# ======================================================================
-# resolve_manager_class
-# ======================================================================
-
-
-class TestResolveManagerClass(unittest.TestCase):
-    """Tests for resolve_manager_class."""
-
-    def test_stable_resolves_to_isaaclab_managers(self):
-        switch = ManagerCallSwitch(cfg_source={"default": 0})
-        cls = switch.resolve_manager_class("RewardManager")
-        self.assertTrue(cls.__module__.startswith("isaaclab.managers"))
-
-    def test_warp_resolves_to_experimental_managers(self):
-        switch = ManagerCallSwitch(cfg_source={"default": 1})
-        cls = switch.resolve_manager_class("RewardManager")
-        self.assertTrue(cls.__module__.startswith("isaaclab_experimental"))
-
-    def test_mode_override(self):
-        """mode_override should bypass the config for class resolution."""
-        switch = ManagerCallSwitch(cfg_source={"default": 2})
-        cls = switch.resolve_manager_class("RewardManager", mode_override=ManagerCallMode.STABLE)
-        self.assertTrue(cls.__module__.startswith("isaaclab.managers"))
-
-    def test_invalid_manager_raises(self):
-        switch = ManagerCallSwitch(cfg_source={"default": 0})
-        with self.assertRaises(AttributeError):
-            switch.resolve_manager_class("NonexistentManager")
-
-
-# ======================================================================
-# Manager name parsing
-# ======================================================================
-
-
-class TestManagerNameParsing(unittest.TestCase):
-    """Tests for stage name → manager name extraction."""
-
-    def test_invalid_stage_name_raises(self):
-        switch = ManagerCallSwitch(cfg_source={"default": 1})
-        with self.assertRaises(ValueError):
-            switch.call_stage(stage="nounderscore", warp_call={"fn": lambda: None})
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py b/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py
index 7a715c639d44..6ed2c1c27a54 100644
--- a/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py
+++ b/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py
@@ -3,195 +3,177 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
-"""Tests for WarpGraphCache capture/replay and warm-up behavior."""
+"""Tests for the WarpGraphCache capture lifecycle."""
 
 from __future__ import annotations
 
-import unittest
-
+import pytest
+import torch
 import warp as wp
-from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache
+from isaaclab_experimental.utils.warp_graph_cache import CAPTURE_ENV_VAR, SYNC_DEBUG_ENV_VAR, WarpGraphCache
+
+wp.init()
+pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required")
+
+DEVICE = "cuda:0"
 
 
 @wp.kernel
-def _add_one(a: wp.array(dtype=wp.float32), b: wp.array(dtype=wp.float32)):
-    """Simple warp kernel: b[i] = a[i] + 1."""
+def _increment(values: wp.array(dtype=wp.int32)):
     i = wp.tid()
-    b[i] = a[i] + 1.0
-
-
-class TestWarpGraphCache(unittest.TestCase):
-    """Tests for :class:`WarpGraphCache`."""
-
-    def setUp(self):
-        self.device = "cuda:0"
-        self.cache = WarpGraphCache()
-
-    # ------------------------------------------------------------------
-    # Warm-up
-    # ------------------------------------------------------------------
-
-    def test_warmup_runs_before_capture(self):
-        """The function should be called eagerly (warm-up) before graph capture.
-
-        We verify this by counting total invocations on the first call.
-        Warm-up = 1, capture = 1, so fn should be called exactly 2 times.
-        """
-        call_count = [0]
-        src = wp.zeros(4, dtype=wp.float32, device=self.device)
-        dst = wp.zeros(4, dtype=wp.float32, device=self.device)
-
-        def counted_launch():
-            call_count[0] += 1
-            wp.launch(_add_one, dim=4, inputs=[src, dst], device=self.device)
-            return dst
-
-        # First call: warm-up + capture
-        self.cache.capture_or_replay("stage_a", counted_launch)
-        self.assertEqual(call_count[0], 2, "First call should invoke fn twice (warm-up + capture)")
-
-        # Second call: replay only
-        self.cache.capture_or_replay("stage_a", counted_launch)
-        self.assertEqual(call_count[0], 2, "Replay should NOT invoke fn again")
-
-    # ------------------------------------------------------------------
-    # Capture / replay correctness
-    # ------------------------------------------------------------------
-
-    def test_replay_uses_updated_input(self):
-        """Replay should re-read from the same input buffer (pointer-stable).
-
-        CUDA graph replay re-executes the same kernel on the same memory
-        addresses. If we update the input buffer in-place, the output
-        should reflect the new values.
-        """
-        src = wp.full(4, value=1.0, dtype=wp.float32, device=self.device)
-        dst = wp.zeros(4, dtype=wp.float32, device=self.device)
-
-        def my_fn():
-            wp.launch(_add_one, dim=4, inputs=[src, dst], device=self.device)
-            return dst
-
-        # Capture
-        first = self.cache.capture_or_replay("replay_test", my_fn)
-
-        # Update input in-place
-        wp.copy(src, wp.full(4, value=10.0, dtype=wp.float32, device=self.device))
-
-        # Replay — should see updated input
-        result = self.cache.capture_or_replay("replay_test", my_fn)
-        # Replay returns the cached object from the capture call, not a new buffer.
-        self.assertIs(result, first)
-        result_np = result.numpy()
-        for val in result_np:
-            self.assertAlmostEqual(val, 11.0, places=5)
-
-    def test_multiple_stages_independent(self):
-        """Different stages should be captured and replayed independently."""
-        src_a = wp.full(4, value=1.0, dtype=wp.float32, device=self.device)
-        dst_a = wp.zeros(4, dtype=wp.float32, device=self.device)
-        src_b = wp.full(4, value=5.0, dtype=wp.float32, device=self.device)
-        dst_b = wp.zeros(4, dtype=wp.float32, device=self.device)
-
-        def fn_a():
-            wp.launch(_add_one, dim=4, inputs=[src_a, dst_a], device=self.device)
-            return dst_a
-
-        def fn_b():
-            wp.launch(_add_one, dim=4, inputs=[src_b, dst_b], device=self.device)
-            return dst_b
-
-        result_a = self.cache.capture_or_replay("stage_a", fn_a)
-        result_b = self.cache.capture_or_replay("stage_b", fn_b)
-
-        self.assertAlmostEqual(result_a.numpy()[0], 2.0, places=5)
-        self.assertAlmostEqual(result_b.numpy()[0], 6.0, places=5)
-
-        # Replay both
-        result_a2 = self.cache.capture_or_replay("stage_a", fn_a)
-        result_b2 = self.cache.capture_or_replay("stage_b", fn_b)
-        self.assertIs(result_a, result_a2)
-        self.assertIs(result_b, result_b2)
-
-    # ------------------------------------------------------------------
-    # Invalidation
-    # ------------------------------------------------------------------
-
-    def test_invalidate_all(self):
-        """invalidate() with no args should drop all cached graphs."""
-        src = wp.zeros(4, dtype=wp.float32, device=self.device)
-        dst = wp.zeros(4, dtype=wp.float32, device=self.device)
-
-        call_count = [0]
-
-        def my_fn():
-            call_count[0] += 1
-            wp.launch(_add_one, dim=4, inputs=[src, dst], device=self.device)
-            return dst
-
-        self.cache.capture_or_replay("s1", my_fn)
-        self.assertEqual(call_count[0], 2)  # warm-up + capture
-
-        self.cache.invalidate()
-
-        # After invalidation, next call should re-warm-up and re-capture
-        self.cache.capture_or_replay("s1", my_fn)
-        self.assertEqual(call_count[0], 4)  # 2 more (warm-up + capture)
-
-    def test_invalidate_single_stage(self):
-        """invalidate(stage) should only drop the named stage."""
-        src = wp.zeros(4, dtype=wp.float32, device=self.device)
-        dst_a = wp.zeros(4, dtype=wp.float32, device=self.device)
-        dst_b = wp.zeros(4, dtype=wp.float32, device=self.device)
-
-        count_a = [0]
-        count_b = [0]
-
-        def fn_a():
-            count_a[0] += 1
-            wp.launch(_add_one, dim=4, inputs=[src, dst_a], device=self.device)
-            return dst_a
-
-        def fn_b():
-            count_b[0] += 1
-            wp.launch(_add_one, dim=4, inputs=[src, dst_b], device=self.device)
-            return dst_b
-
-        self.cache.capture_or_replay("a", fn_a)
-        self.cache.capture_or_replay("b", fn_b)
-        self.assertEqual(count_a[0], 2)
-        self.assertEqual(count_b[0], 2)
-
-        # Invalidate only "a"; a stage that was never captured is a no-op
-        self.cache.invalidate("a")
-        self.cache.invalidate("nonexistent")
-
-        self.cache.capture_or_replay("a", fn_a)
-        self.cache.capture_or_replay("b", fn_b)
-        self.assertEqual(count_a[0], 4, "Stage 'a' should re-capture after invalidation")
-        self.assertEqual(count_b[0], 2, "Stage 'b' should replay (not re-capture)")
-
-    # ------------------------------------------------------------------
-    # Args / kwargs forwarding
-    # ------------------------------------------------------------------
-
-    def test_args_and_kwargs_forwarded(self):
-        """capture_or_replay should forward args and kwargs to fn."""
-        src = wp.full(4, value=2.0, dtype=wp.float32, device=self.device)
-        dst = wp.zeros(4, dtype=wp.float32, device=self.device)
-
-        def my_fn(a, b, device="cuda:0"):
-            wp.launch(_add_one, dim=4, inputs=[a, b], device=device)
-            return b
-
-        result = self.cache.capture_or_replay(
-            "args_test",
-            my_fn,
-            args=(src, dst),
-            kwargs={"device": self.device},
-        )
-        self.assertAlmostEqual(result.numpy()[0], 3.0, places=5)
+    values[i] = values[i] + 1
+
+
+class _CountingStage:
+    """Stage that increments a device counter and counts its Python invocations."""
+
+    def __init__(self):
+        self.calls = 0
+
+    def __call__(self, values: wp.array) -> wp.array:
+        self.calls += 1
+        wp.launch(_increment, dim=values.shape[0], inputs=[values], device=DEVICE)
+        return values
+
+
+def _read(values: wp.array) -> list[int]:
+    wp.synchronize_device(DEVICE)
+    return values.numpy().tolist()
+
+
+@pytest.fixture
+def cache(monkeypatch):
+    monkeypatch.delenv(CAPTURE_ENV_VAR, raising=False)
+    monkeypatch.delenv(SYNC_DEBUG_ENV_VAR, raising=False)
+    cache = WarpGraphCache(DEVICE)
+    yield cache
+    cache.close()
+
+
+def test_first_armed_call_records_and_runs_the_stage_once(cache):
+    """Recording has no eager warm-up: a stateful stage advances once per call, including the first."""
+    stage = _CountingStage()
+    values = wp.zeros(2, dtype=wp.int32, device=DEVICE)
+    cache.arm()
+
+    result = cache.call("Group_stage", stage, values)
+    assert _read(values) == [1, 1]
+    assert result is values
+
+    cache.call("Group_stage", stage, values)
+    assert _read(values) == [2, 2]
+    assert stage.calls == 1, "the second call must replay the recorded graph"
+    assert cache.captured_stages == ("Group_stage",)
+
+
+def test_stages_run_eagerly_until_armed(cache):
+    stage = _CountingStage()
+    values = wp.zeros(1, dtype=wp.int32, device=DEVICE)
+
+    cache.call("Group_stage", stage, values)
+    cache.call("Group_stage", stage, values)
+    assert (stage.calls, _read(values), cache.captured_stages) == (2, [2], ())
+
+    cache.arm()
+    cache.call("Group_stage", stage, values)
+    cache.call("Group_stage", stage, values)
+    assert (stage.calls, _read(values), cache.captured_stages) == (3, [4], ("Group_stage",))
+
+
+def test_reallocated_argument_records_the_stage_again(cache):
+    """A graph replays recorded pointers, so a new array must not reuse the old graph."""
+    stage = _CountingStage()
+    old = wp.zeros(1, dtype=wp.int32, device=DEVICE)
+    cache.arm()
+    cache.call("Group_stage", stage, old)
+
+    new = wp.zeros(1, dtype=wp.int32, device=DEVICE)
+    cache.call("Group_stage", stage, new)
+    cache.call("Group_stage", stage, new)
+
+    assert (_read(old), _read(new), stage.calls) == ([1], [2], 2)
+
+
+def test_changed_scalar_argument_records_the_stage_again(cache):
+    values = wp.zeros(1, dtype=wp.float32, device=DEVICE)
+
+    def add(values: wp.array, amount: float) -> None:
+        wp.launch(_add_scalar, dim=1, inputs=[values, amount], device=DEVICE)
+
+    cache.arm()
+    cache.call("Group_stage", add, values, 1.0)
+    cache.call("Group_stage", add, values, 10.0)
+    assert _read(values) == [11.0]
+
+
+@wp.kernel
+def _add_scalar(values: wp.array(dtype=wp.float32), amount: wp.float32):
+    values[wp.tid()] = values[wp.tid()] + amount
+
+
+def test_invalidate_drops_graphs_and_waits_for_the_next_arm(cache):
+    stage = _CountingStage()
+    values = wp.zeros(1, dtype=wp.int32, device=DEVICE)
+    cache.arm()
+    cache.call("Group_stage", stage, values)
+
+    cache.invalidate()
+    cache.call("Group_stage", stage, values)
+    assert (stage.calls, cache.captured_stages) == (2, ())
+
+    cache.arm()
+    cache.call("Group_stage", stage, values)
+    cache.call("Group_stage", stage, values)
+    assert (stage.calls, _read(values), cache.captured_stages) == (3, [4], ("Group_stage",))
+
+
+def test_group_invalidation_keeps_other_groups_recorded(cache):
+    values = wp.zeros(1, dtype=wp.int32, device=DEVICE)
+    cache.arm()
+    cache.call("First_stage", _CountingStage(), values)
+    cache.call("Second_stage", _CountingStage(), values)
+
+    cache.invalidate("First")
+
+    assert cache.captured_stages == ("Second_stage",)
+    stage = _CountingStage()
+    cache.call("First_stage", stage, values)
+    assert (stage.calls, cache.captured_stages) == (1, ("First_stage", "Second_stage"))
+
+
+def test_non_capturable_group_stays_eager(cache):
+    stage = _CountingStage()
+    values = wp.zeros(1, dtype=wp.int32, device=DEVICE)
+    cache.register_capturability("Group", True)
+    cache.register_capturability("Group", False)
+    cache.register_capturability("Group", True)
+    cache.arm()
+
+    cache.call("Group_stage", stage, values)
+    cache.call("Group_stage", stage, values)
+
+    assert (stage.calls, _read(values), cache.captured_stages) == (2, [2], ())
+
+
+def test_capture_env_var_forces_eager_execution(monkeypatch):
+    monkeypatch.setenv(CAPTURE_ENV_VAR, "0")
+    cache = WarpGraphCache(DEVICE)
+    stage = _CountingStage()
+    values = wp.zeros(1, dtype=wp.int32, device=DEVICE)
+    cache.arm()
+
+    cache.call("Group_stage", stage, values)
+    cache.call("Group_stage", stage, values)
+
+    assert (stage.calls, _read(values), cache.captured_stages) == (2, [2], ())
+
 
+def test_sync_debug_trap_raises_on_a_hidden_host_sync(monkeypatch):
+    monkeypatch.setenv(CAPTURE_ENV_VAR, "0")
+    monkeypatch.setenv(SYNC_DEBUG_ENV_VAR, "1")
+    cache = WarpGraphCache(DEVICE)
+    flags = torch.ones(4, dtype=torch.bool, device=DEVICE)
 
-if __name__ == "__main__":
-    unittest.main()
+    with pytest.raises(RuntimeError, match="synchroniz"):
+        cache.call("Group_stage", lambda: flags.any().item())
+    assert torch.cuda.get_sync_debug_mode() == 0, "the trap must be released after the stage"

From d0b411a36522ed46dc4f2bbc8ff016d7eb7b84f5 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 03:01:33 -0700
Subject: [PATCH 09/48] Reset Warp manager-based environments by mask

The manager-based Warp environments now reset through one mask-first
pipeline: every reset selection is copied into an environment-owned mask
that the recorded reset stages read, the scene and the Warp managers
reset from that mask, and the episode-length buffer is cleared by a
masked kernel. The mask is compacted to environment indices once per
reset, and only when a stable command, curriculum or recorder manager has
active terms. ManagerBasedEnvWarp.reset accepts env_mask.

Each step makes one host check to decide whether any environment
terminated: the index compaction when a manager needs indices, otherwise
reset_buf.any(). DirectRLEnvWarp moves its reset out of the end-of-step
graph into its own stage behind the same predicate.

A static gate pins the sanctioned host synchronizations of the mask-first
path, which also removes the device read from CircularBuffer.max_length.
The Episode_Termination metric keeps averaging over every environment,
as the stable manager does.
---
 .../jichuanh-warp-frontend-cleanup.added.rst  |   1 +
 ...jichuanh-warp-frontend-cleanup.changed.rst |   4 +
 .../jichuanh-warp-frontend-cleanup.fixed.rst  |   1 +
 .../envs/direct_rl_env_warp.py                |   8 +-
 .../envs/manager_based_env_warp.py            |  70 +++++--
 .../envs/manager_based_rl_env_warp.py         | 116 +++++------
 .../utils/buffers/circular_buffer.py          |   5 +-
 .../utils/warp/__init__.py                    |   2 +-
 .../utils/warp/kernels.py                     |   8 +
 .../test/envs/test_warp_env_reset.py          |  49 +++++
 .../test/managers/test_termination_manager.py |  39 ++++
 .../test/utils/test_mask_boundary_gate.py     | 181 ++++++++++++++++++
 12 files changed, 398 insertions(+), 86 deletions(-)
 create mode 100644 source/isaaclab_experimental/test/envs/test_warp_env_reset.py
 create mode 100644 source/isaaclab_experimental/test/managers/test_termination_manager.py
 create mode 100644 source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py

diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst
index 13db10a0a56a..dee68bc25e57 100644
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst
@@ -2,3 +2,4 @@
   environment stage eagerly, for debugging and eager-versus-captured comparisons.
 * Added the ``ISAACLAB_SYNC_DEBUG`` environment variable: ``ISAACLAB_SYNC_DEBUG=1`` raises on a hidden
   GPU-to-host synchronization inside an eagerly executed Warp environment stage.
+* Added a boolean ``env_mask`` argument to :meth:`~isaaclab_experimental.envs.ManagerBasedEnvWarp.reset`.
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst
index 3c2f2127e843..1d4e33069b81 100644
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst
@@ -4,3 +4,7 @@
 * **Breaking:** Replaced ``WarpGraphCache.capture_or_replay`` with
   :meth:`~isaaclab_experimental.utils.WarpGraphCache.call`. Stages are recorded from the first environment
   step without an eager warm-up, so a stage runs exactly once per call.
+* Changed the Warp manager-based environments to reset environments by boolean mask. The mask is converted
+  to environment indices once per reset, and only when a command, curriculum or recorder term is active.
+* Changed :class:`~isaaclab_experimental.envs.DirectRLEnvWarp` to skip its reset stage on steps where no
+  environment terminated.
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
index 3c4f09ee4e12..0699dcb75532 100644
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
@@ -1,2 +1,3 @@
 * Fixed Warp environments replaying CUDA graphs that read freed simulation buffers after the simulation was
   reset or stopped. The environments now record their stages again on the next step.
+* Fixed :attr:`~isaaclab_experimental.utils.buffers.CircularBuffer.max_length` synchronizing with the GPU.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
index 4d095c031807..981c64dfd344 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
@@ -218,6 +218,7 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
         self.reset_terminated = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device)
         self.reset_time_outs = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device)
         self.reset_buf = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device)
+        self._torch_reset_buf = wp.to_torch(self.reset_buf)
         self._ALL_ENV_MASK = wp.ones(self.num_envs, dtype=wp.bool, device=self.device)
 
         # Expected bindings:
@@ -437,6 +438,10 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         self.common_step_counter += 1  # total step (common for all envs)
         with Timer(name="end_pre_graph", msg="End pre-graph took:", enable=DEBUG_TIMERS):
             self._warp_graph_cache.call("DirectEndPre_step", self._step_warp_end_pre)
+        # one host predicate per step skips the reset stage when no environment terminated
+        if self._torch_reset_buf.any().item():
+            with Timer(name="reset_graph", msg="Reset graph took:", enable=DEBUG_TIMERS):
+                self._warp_graph_cache.call("DirectReset_step", self._step_warp_reset)
         # write_data_to_sim runs uncaptured — it uses torch ops that cross CUDA streams.
         with Timer(name="write_data_to_sim_post", msg="Write data to sim (post-reset) took:", enable=DEBUG_TIMERS):
             self.scene.write_data_to_sim()
@@ -488,7 +493,8 @@ def _step_warp_end_pre(self) -> None:
         self._get_dones()
         self._get_rewards()
 
-        # -- reset envs that terminated/timed-out and log the episode information
+    def _step_warp_reset(self) -> None:
+        """Capturable reset of the environments that terminated or timed out."""
         self._reset_idx(mask=self.reset_buf)
 
     def _step_warp_end_post(self) -> None:
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
index 704d101ad7ee..e2b9a3ad4f48 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
@@ -163,6 +163,8 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
         # Pre-allocated env masks (shared across managers/terms via `env`).
         self.ALL_ENV_MASK = wp.ones((self.num_envs,), dtype=wp.bool, device=self.device)
         self.ENV_MASK = wp.zeros((self.num_envs,), dtype=wp.bool, device=self.device)
+        # Reset stages read this mask; every reset selection is copied into it.
+        self.reset_mask_wp = wp.zeros((self.num_envs,), dtype=wp.bool, device=self.device)
 
         # Persistent scalar buffer for global env step count (stable pointer for capture).
         self._global_env_step_count_wp = wp.zeros((1,), dtype=wp.int32, device=self.device)
@@ -391,6 +393,8 @@ def load_managers(self):
         self._action_in_wp = wp.zeros(
             (self.num_envs, self.action_manager.total_action_dim), dtype=wp.float32, device=self.device
         )
+        # whether a reset has to compact its mask for a manager that resets by environment index
+        self._resets_by_index = bool(self.recorder_manager.active_terms)
 
         # perform events at the start of the simulation
         # in-case a child implementation creates other managers, the randomization should happen
@@ -415,11 +419,15 @@ def setup_manager_visualizers(self):
     """
 
     def reset(
-        self, seed: int | None = None, env_ids: Sequence[int] | None = None, options: dict[str, Any] | None = None
+        self,
+        seed: int | None = None,
+        env_ids: Sequence[int] | None = None,
+        options: dict[str, Any] | None = None,
+        env_mask: wp.array | torch.Tensor | None = None,
     ) -> tuple[VecEnvObs, dict]:
         """Resets the specified environments and returns observations.
 
-        This function calls the :meth:`_reset_idx` function to reset the specified environments.
+        This function calls the :meth:`_reset_mask` function to reset the specified environments.
         However, certain operations, such as procedural terrain generation, that happened during initialization
         are not repeated.
 
@@ -431,11 +439,14 @@ def reset(
                 Note:
                     This argument is used for compatibility with Gymnasium environment definition.
 
+            env_mask: Boolean mask of the environments to reset, of shape ``(num_envs,)``. Takes precedence
+                over :paramref:`env_ids`. Defaults to None.
+
         Returns:
             A tuple containing the observations and extras.
         """
-        if env_ids is None:
-            env_ids = torch.arange(self.num_envs, dtype=torch.int64, device=self.device)
+        env_mask = self.resolve_env_mask(env_ids=env_ids, env_mask=env_mask)
+        env_ids = self._reset_env_ids(env_mask)
 
         # trigger recorder terms for pre-reset calls
         self.recorder_manager.record_pre_reset(env_ids)
@@ -454,7 +465,7 @@ def reset(
             )
 
         # reset state of scene
-        self._reset_idx(env_ids)
+        self._reset_mask(env_mask, env_ids)
 
         # update articulation kinematics
         self.scene.write_data_to_sim()
@@ -508,7 +519,7 @@ def reset_to(
         if seed is not None:
             self.seed(seed)
 
-        self._reset_idx(env_ids)
+        self._reset_mask(self.resolve_env_mask(env_ids=env_ids), env_ids)
 
         # set the state
         self.scene.reset_to(state, env_ids, is_relative=is_relative)
@@ -632,36 +643,63 @@ def close(self):
     Helper functions.
     """
 
-    def _reset_idx(self, env_ids: Sequence[int]):
-        """Reset environments based on specified indices.
+    def _reset_env_ids(self, env_mask: wp.array) -> torch.Tensor | None:
+        """Compact a reset mask for the managers that reset by environment index.
+
+        This is the host boundary of the mask-first reset: it synchronizes with the device.
+
+        Args:
+            env_mask: Boolean mask of the environments to reset.
+
+        Returns:
+            Indices of the selected environments, or None when every active manager resets by mask.
+        """
+        if not self._resets_by_index:
+            return None
+        return wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1)
+
+    def _reset_mask(self, env_mask: wp.array, env_ids: torch.Tensor | None = None) -> None:
+        """Reset the selected environments.
 
         Args:
-            env_ids: List of environment ids which must be reset
+            env_mask: Boolean mask of the environments to reset.
+            env_ids: Indices of the same environments for the managers that reset by index, or None
+                when every active manager resets by mask.
         """
+        # recorded reset stages read the environment-owned mask
+        if env_mask is not self.reset_mask_wp:
+            wp.copy(self.reset_mask_wp, env_mask)
+        env_mask = self.reset_mask_wp
+
         # reset the internal buffers of the scene elements
-        self.scene.reset(env_ids)
+        self.scene.reset(env_mask=env_mask)
 
         # apply events such as randomization for environments that need a reset
         if "reset" in self.event_manager.available_modes:
             env_step_count = self._sim_step_counter // self.cfg.decimation
             self._global_env_step_count_wp.fill_(env_step_count)
-            self.event_manager.apply(
-                mode="reset", env_ids=env_ids, global_env_step_count=self._global_env_step_count_wp
+            self._warp_graph_cache.call(
+                "EventManager_apply_reset",
+                self.event_manager.apply,
+                mode="reset",
+                env_mask_wp=env_mask,
+                global_env_step_count=self._global_env_step_count_wp,
             )
 
         # iterate over all managers and reset them
         # this returns a dictionary of information which is stored in the extras
         # note: This is order-sensitive! Certain things need be reset before others.
         self.extras["log"] = dict()
-        env_mask = self.resolve_env_mask(env_ids=env_ids)
         # -- observation manager
-        info = self.observation_manager.reset(env_mask=env_mask)
+        info = self._warp_graph_cache.call(
+            "ObservationManager_reset", self.observation_manager.reset, env_mask=env_mask
+        )
         self.extras["log"].update(info)
         # -- action manager
-        info = self.action_manager.reset(env_mask=env_mask)
+        info = self._warp_graph_cache.call("ActionManager_reset", self.action_manager.reset, env_mask=env_mask)
         self.extras["log"].update(info)
         # -- event manager
-        info = self.event_manager.reset(env_mask=env_mask)
+        info = self._warp_graph_cache.call("EventManager_reset", self.event_manager.reset, env_mask=env_mask)
         self.extras["log"].update(info)
         # -- recorder manager
         info = self.recorder_manager.reset(env_ids)
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
index 6414f375cac2..c720de01ca5f 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
@@ -15,7 +15,6 @@
 
 import math
 import os
-from collections.abc import Sequence
 from typing import Any, ClassVar
 
 import gymnasium as gym
@@ -26,11 +25,11 @@
 from isaaclab.envs.common import VecEnvStepReturn
 from isaaclab.envs.manager_based_rl_env_cfg import ManagerBasedRLEnvCfg
 from isaaclab.managers import CommandManager, CurriculumManager
-from isaaclab.utils import index_fill_
 from isaaclab.utils.timer import Timer
 
 from isaaclab_experimental.managers import RewardManager, TerminationManager
 from isaaclab_experimental.utils.torch_utils import clone_obs_buffer
+from isaaclab_experimental.utils.warp import zero_masked_int64
 
 from .manager_based_env_warp import ManagerBasedEnvWarp
 
@@ -114,11 +113,6 @@ def __init__(self, cfg: ManagerBasedRLEnvCfg, render_mode: str | None = None, **
         # store the render mode
         self.render_mode = render_mode
 
-        # The persistent reset mask needed for warp capture
-        # The intended use is to copy into this mask whenever capture is needed
-        # TODO: termination manager provides the same mask, investigate whether this can be replaced.
-        self.reset_mask_wp = wp.zeros(cfg.scene.num_envs, dtype=wp.bool, device=cfg.sim.device)
-
         # initialize data and constants
         # -- set the framerate of the gym video recorder wrapper so that the playback speed
         # of the produced video matches the simulation
@@ -176,6 +170,13 @@ def load_managers(self):
         self.curriculum_manager = CurriculumManager(self.cfg.curriculum, self)
         print("[INFO] Curriculum Manager: ", self.curriculum_manager)
 
+        # stable command and curriculum managers reset by environment index
+        self._resets_by_index = (
+            self._resets_by_index
+            or bool(self.command_manager.active_terms)
+            or bool(self.curriculum_manager.active_terms)
+        )
+
         # setup the action and observation spaces for Gym
         self._configure_gym_env_spaces()
 
@@ -314,36 +315,7 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
             self.recorder_manager.record_post_step()
 
         # -- reset envs that terminated/timed-out and log the episode information
-        # NOTE: Interim path (intentional).
-        # We still compact `reset_buf` into `env_ids` here because several reset-time managers/recorders
-        # are still `env_ids`-based. Do NOT remove/replace this until mask-based reset is end-to-end.
-        with Timer(
-            name="reset_selection",
-            msg="Reset selection took:",
-            enable=DEBUG_TIMER_STEP,
-            time_unit="us",
-        ):
-            wp.copy(self.reset_mask_wp, self.termination_manager.dones_wp)
-            reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1)
-        if len(reset_env_ids) > 0:
-            # trigger recorder terms for pre-reset calls
-            self.recorder_manager.record_pre_reset(reset_env_ids)
-
-            with Timer(
-                name="reset_idx",
-                msg="Reset idx took:",
-                enable=DEBUG_TIMER_STEP,
-                time_unit="us",
-            ):
-                self._reset_idx(env_ids=reset_env_ids, env_mask=self.reset_mask_wp)
-
-            # if sensors are added to the scene, make sure we render to reflect changes in reset
-            if self.has_rtx_sensors and self.cfg.num_rerenders_on_reset > 0:
-                for _ in range(self.cfg.num_rerenders_on_reset):
-                    self.sim.render()
-
-            # trigger recorder terms for post-reset calls
-            self.recorder_manager.record_post_reset(reset_env_ids)
+        self._reset_terminated_envs()
 
         # -- update command
         self.command_manager.compute(dt=float(self.step_dt))
@@ -460,39 +432,47 @@ def _configure_gym_env_spaces(self):
         self.observation_space = gym.vector.utils.batch_space(self.single_observation_space, self.num_envs)
         self.action_space = gym.vector.utils.batch_space(self.single_action_space, self.num_envs)
 
-    def _reset_idx(
-        self,
-        env_ids: Sequence[int] | slice | torch.Tensor,
-        *,
-        env_mask: wp.array | None = None,
-    ):
-        """Reset environments based on specified indices.
+    def _reset_terminated_envs(self) -> None:
+        """Reset the environments that terminated or timed out in this step."""
+        with Timer(name="reset_selection", msg="Reset selection took:", enable=DEBUG_TIMER_STEP, time_unit="us"):
+            # The reset pipeline advances curriculum and logging state even for an empty mask, so one
+            # host synchronization per step decides whether it runs. When a manager resets by index,
+            # the same synchronization also compacts the mask.
+            if self._resets_by_index:
+                reset_env_ids = self._reset_env_ids(self.termination_manager.dones_wp)
+                has_resets = len(reset_env_ids) > 0
+            else:
+                reset_env_ids = None
+                has_resets = bool(self.reset_buf.any().item())
+        if not has_resets:
+            return
 
-        IMPORTANT:
-            This function always uses the **TerminationManager-produced Warp env mask** (`self.reset_buf`) to select
-            which envs to reset. The ids/mask conversion is performed in `step()` before calling this function.
+        # trigger recorder terms for pre-reset calls
+        self.recorder_manager.record_pre_reset(reset_env_ids)
 
-            In other words:
-            - If `env_mask` is provided, it **must** be `self.reset_buf` (Warp bool mask)
-            - If `env_mask` is not provided, this function will populate `self.reset_buf` from `env_ids`
-            - When `env_mask` is provided, `env_ids` **must** correspond to the same mask
+        with Timer(name="reset_mask", msg="Reset mask took:", enable=DEBUG_TIMER_STEP, time_unit="us"):
+            self._reset_mask(self.termination_manager.dones_wp, reset_env_ids)
+
+        # if sensors are added to the scene, make sure we render to reflect changes in reset
+        if self.has_rtx_sensors and self.cfg.num_rerenders_on_reset > 0:
+            for _ in range(self.cfg.num_rerenders_on_reset):
+                self.sim.render()
+
+        # trigger recorder terms for post-reset calls
+        self.recorder_manager.record_post_reset(reset_env_ids)
+
+    def _reset_mask(self, env_mask: wp.array, env_ids: torch.Tensor | None = None) -> None:
+        """Reset the selected environments.
 
         Args:
-            env_ids: Environment indices to reset.
-            env_mask: Warp boolean env mask selecting envs to reset. Must be `self.reset_buf`.
-                If None, uses and populates `self.reset_buf` from `env_ids`.
+            env_mask: Boolean mask of the environments to reset.
+            env_ids: Indices of the same environments for the managers that reset by index, or None
+                when every active manager resets by mask.
         """
-        if env_mask is None:
-            # Base `reset()` / `reset_to()` call-path provides only `env_ids`.
-            # Populate the stable TerminationManager-owned mask (`self.reset_buf`) from ids.
-            env_mask = self.reset_mask_wp
-            # Use the centralized env-id/mask resolution from the base Warp env, then copy into the
-            # stable TerminationManager-owned buffer (`self.reset_buf`) used by captured graphs.
-            resolved_mask = self.resolve_env_mask(env_ids=env_ids)
-            wp.copy(env_mask, resolved_mask)
-
-        if not isinstance(env_mask, wp.array):
-            raise TypeError(f"env_mask must be a wp.array (got {type(env_mask)}).")
+        # recorded reset stages read the environment-owned mask
+        if env_mask is not self.reset_mask_wp:
+            wp.copy(self.reset_mask_wp, env_mask)
+        env_mask = self.reset_mask_wp
 
         # update the curriculum for environments that need a reset
         with Timer(
@@ -505,7 +485,7 @@ def _reset_idx(
 
         # reset the internal buffers of the scene elements
         with Timer(name="Scene_reset", msg="Scene reset took:", enable=DEBUG_TIMER_RESET, time_unit="us"):
-            self.scene.reset(env_ids=env_ids, env_mask=env_mask)
+            self.scene.reset(env_mask=env_mask)
 
         if "reset" in self.event_manager.available_modes:
             self._global_env_step_count_wp.fill_(self._sim_step_counter // self.cfg.decimation)
@@ -554,7 +534,9 @@ def _reset_idx(
         recorder_info = self.recorder_manager.reset(env_ids=env_ids)
 
         # reset the episode length buffer
-        index_fill_(self.episode_length_buf, env_ids, 0)
+        wp.launch(
+            zero_masked_int64, dim=self.num_envs, inputs=[env_mask, self._episode_length_buf_wp], device=self.device
+        )
 
         # aggregate logging info
         log: dict[str, Any] = {}
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/buffers/circular_buffer.py b/source/isaaclab_experimental/isaaclab_experimental/utils/buffers/circular_buffer.py
index 3eba7590624c..ad9ba04950a5 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/utils/buffers/circular_buffer.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/utils/buffers/circular_buffer.py
@@ -44,6 +44,9 @@ def __init__(self, max_len: int, batch_size: int, device: str):
         self._device = device
         self._ALL_INDICES = torch.arange(batch_size, device=device)
 
+        # max length as a host int: the property is read on the append path, so it must not
+        # synchronize with the device tensor below.
+        self._max_len_int = int(max_len)
         # max length tensor for comparisons
         self._max_len = torch.full((batch_size,), max_len, dtype=torch.int, device=device)
         # number of data pushes passed since the last call to :meth:`reset`
@@ -71,7 +74,7 @@ def device(self) -> str:
     @property
     def max_length(self) -> int:
         """The maximum length of the ring buffer."""
-        return int(self._max_len[0].item())
+        return self._max_len_int
 
     @property
     def current_length(self) -> torch.Tensor:
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py
index 2d071b823a46..a91c051b0d08 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py
@@ -5,7 +5,7 @@
 
 """Warp utility functions and shared kernels for isaaclab_experimental."""
 
-from .kernels import compute_reset_scale, count_masked
+from .kernels import compute_reset_scale, count_masked, zero_masked_int64
 from .utils import (
     WarpCapturable,
     is_warp_capturable,
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/kernels.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/kernels.py
index 8d3f9e65d49a..6320240740ce 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/kernels.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/kernels.py
@@ -10,6 +10,14 @@
 import warp as wp
 
 
+@wp.kernel
+def zero_masked_int64(mask: wp.array(dtype=wp.bool), values: wp.array(dtype=wp.int64)):
+    """Zero the entries of a one-dimensional int64 array selected by ``mask``."""
+    index = wp.tid()
+    if mask[index]:
+        values[index] = wp.int64(0)
+
+
 @wp.kernel
 def count_masked(
     mask: wp.array(dtype=wp.bool),
diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_reset.py b/source/isaaclab_experimental/test/envs/test_warp_env_reset.py
new file mode 100644
index 000000000000..b5148ca7275c
--- /dev/null
+++ b/source/isaaclab_experimental/test/envs/test_warp_env_reset.py
@@ -0,0 +1,49 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Mask-first resets of the manager-based Warp environment."""
+
+from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg
+
+from isaaclab.sim import SimulationCfg
+from isaaclab.test.utils import launch_test_simulation
+
+launch_test_simulation(SimulationCfg(physics=NewtonCfg(solver_cfg=MJWarpSolverCfg())))
+
+import torch
+from isaaclab_experimental.envs.frontend import WarpFrontend
+
+import isaaclab.sim as sim_utils
+
+import isaaclab_tasks  # noqa: F401
+from isaaclab_tasks.utils import resolve_task_config
+
+_NUM_ENVS = 8
+
+
+def test_reset_with_mask_resets_only_selected_environments():
+    env_cfg, _ = resolve_task_config("Isaac-Cartpole", "", overrides=("physics=newton_mjwarp",))
+    env_cfg.seed = 7
+    env_cfg.scene.num_envs = _NUM_ENVS
+    sim_utils.create_new_stage()
+    env = WarpFrontend.build_env(env_cfg, "Isaac-Cartpole").unwrapped
+    try:
+        env.reset()
+        actions = torch.ones((_NUM_ENVS, env.action_space.shape[-1]), device=env.device)
+        for _ in range(5):
+            env.step(actions)
+        robot = env.scene["robot"]
+        joint_pos = robot.data.joint_pos.torch.clone()
+        selected = torch.arange(_NUM_ENVS, device=env.device) % 2 == 0
+
+        env.reset(env_mask=selected)
+
+        assert torch.equal(env.episode_length_buf[selected], torch.zeros_like(env.episode_length_buf[selected]))
+        assert torch.equal(env.episode_length_buf[~selected], torch.full_like(env.episode_length_buf[~selected], 5))
+        assert torch.equal(robot.data.joint_pos.torch[~selected], joint_pos[~selected])
+        # the reset events sample new joint positions for every selected environment
+        assert not torch.isclose(robot.data.joint_pos.torch[selected], joint_pos[selected]).all(dim=1).any()
+    finally:
+        env.close()
diff --git a/source/isaaclab_experimental/test/managers/test_termination_manager.py b/source/isaaclab_experimental/test/managers/test_termination_manager.py
new file mode 100644
index 000000000000..8d60dbad3329
--- /dev/null
+++ b/source/isaaclab_experimental/test/managers/test_termination_manager.py
@@ -0,0 +1,39 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tests for the Warp termination manager."""
+
+from types import SimpleNamespace
+
+import pytest
+import warp as wp
+from isaaclab_experimental.managers import TerminationManager
+
+from isaaclab.managers import TerminationTermCfg
+
+
+def copy_flags(env, out: wp.array):
+    """Warp-first termination term that reports the environment's preset flags."""
+    wp.copy(out, env.flags)
+
+
+def test_episode_termination_metric_averages_over_all_environments():
+    """Like the stable manager, the metric is the share of all environments whose last episode the term ended.
+
+    Only environment 0 is reset. It terminated, so averaging over the reset selection would log 1.0.
+    """
+    env = SimpleNamespace(
+        num_envs=4,
+        device="cpu",
+        sim=SimpleNamespace(is_playing=lambda: True),
+        flags=wp.array([True, True, False, False], dtype=wp.bool, device="cpu"),
+    )
+    manager = TerminationManager({"fell": TerminationTermCfg(func=copy_flags)}, env)
+    manager.compute()
+
+    env_mask = wp.array([True, False, False, False], dtype=wp.bool, device="cpu")
+    extras = manager.reset(env_mask=env_mask)
+
+    assert float(extras["Episode_Termination/fell"]) == pytest.approx(0.5)
diff --git a/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py
new file mode 100644
index 000000000000..a6c7fd74511d
--- /dev/null
+++ b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py
@@ -0,0 +1,181 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Static gate for the mask-first execution path.
+
+Warp frontend code performs no data-dependent GPU-to-host synchronization outside named boundaries:
+
+* **Scope**: every function in the experimental production tree, plus every ``*mask*``-named function in
+  the core and Newton trees (the functions that are the Warp path there). Torch-first core code is out of
+  scope.
+* **Sanctioned boundaries**: the functions in :data:`SANCTIONED_BOUNDARIES`. The list is exact; a stale
+  entry fails the gate.
+* **Runtime backstop**: ``ISAACLAB_SYNC_DEBUG=1`` runs eager Warp stages under
+  ``torch.cuda.set_sync_debug_mode("error")``, trapping synchronizations this scan cannot see.
+
+A second test pins every ``@WarpCapturable(False)`` opt-out, so non-capturable terms are listed in one place.
+"""
+
+from __future__ import annotations
+
+import ast
+from pathlib import Path
+
+_REPO_ROOT = Path(__file__).resolve().parents[4]
+
+# Attribute calls that synchronize with, or copy from, device memory.
+SYNC_ATTRS = {"nonzero", "argwhere", "item", "cpu", "tolist", "numpy"}
+
+# Trees scanned in full (production code only).
+FULL_SCAN_ROOTS = [
+    "source/isaaclab_experimental/isaaclab_experimental/envs",
+    "source/isaaclab_experimental/isaaclab_experimental/managers",
+    "source/isaaclab_experimental/isaaclab_experimental/utils",
+]
+
+# Trees where only ``*mask*``-named functions are in scope.
+MASK_SCAN_ROOTS = [
+    "source/isaaclab/isaaclab/actuators",
+    "source/isaaclab/isaaclab/scene",
+    "source/isaaclab/isaaclab/sensors",
+    "source/isaaclab/isaaclab/terrains",
+    "source/isaaclab_newton/isaaclab_newton",
+]
+
+# Functions that never run on the per-step path: construction, reporting, and debug surfaces.
+NON_STEP_FUNCTIONS = {
+    "__init__",
+    "__post_init__",
+    "__str__",
+    "__repr__",
+    "__del__",
+    "_prepare_terms",
+    "serialize",
+    "get_active_iterable_terms",
+    "_debug_vis_callback",
+    "_set_debug_vis_impl",
+    # deprecated IO-descriptor export surfaces
+    "IO_descriptor",
+    "_collect_io_descriptors",
+    "get_term_cfg",
+}
+
+# Sanctioned host-synchronization boundaries, by (path suffix, function name).
+SANCTIONED_BOUNDARIES = {
+    # one host predicate per step decides whether the reset pipeline runs (W7)
+    ("envs/manager_based_rl_env_warp.py", "_reset_terminated_envs"),
+    ("envs/direct_rl_env_warp.py", "step"),
+    # mask to index compaction for the managers that reset by index (W11)
+    ("envs/manager_based_env_warp.py", "_reset_env_ids"),
+    # startup-only stable randomization terms take environment indices
+    ("envs/mdp/events.py", "_mask_to_env_ids"),
+    # the camera's empty-reset predicate
+    ("isaaclab/sensors/camera/camera.py", "_env_mask_has_any"),
+    # backends without a mask-native actuator reset fall back to indices; Newton and OVPhysX override it
+    ("isaaclab/actuators/actuator_control.py", "reset_native_actuators_mask"),
+    # joint-limit and tendon-property writes mutate the solver model (event-driven, not per-step)
+    ("isaaclab_newton/assets/articulation/articulation.py", "write_joint_position_limit_to_sim_mask"),
+    ("isaaclab_newton/assets/articulation/articulation.py", "_mask_to_ids"),
+}
+
+
+def _sync_calls(func_node: ast.AST) -> list[str]:
+    """Names of synchronizing attribute calls inside a function body."""
+    return [
+        node.func.attr
+        for node in ast.walk(func_node)
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr in SYNC_ATTRS
+    ]
+
+
+def _scan(root: Path, mask_only: bool) -> dict[tuple[str, str], list[str]]:
+    """Map (relative path, function) to synchronizing calls for the in-scope functions under ``root``."""
+    findings: dict[tuple[str, str], list[str]] = {}
+    for path in sorted(root.rglob("*.py")):
+        rel = path.relative_to(_REPO_ROOT).as_posix()
+        if "/test/" in rel:
+            continue
+        for node in ast.walk(ast.parse(path.read_text(), filename=rel)):
+            if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) or node.name in NON_STEP_FUNCTIONS:
+                continue
+            if mask_only and "mask" not in node.name:
+                continue
+            hits = _sync_calls(node)
+            if hits:
+                findings[(rel, node.name)] = hits
+    return findings
+
+
+def _is_sanctioned(rel: str, name: str) -> bool:
+    return any(rel.endswith(suffix) and name == func for suffix, func in SANCTIONED_BOUNDARIES)
+
+
+def test_mask_first_code_has_no_unsanctioned_host_syncs():
+    """Every synchronization in scope is one of the named, sanctioned boundaries."""
+    findings: dict[tuple[str, str], list[str]] = {}
+    for root in FULL_SCAN_ROOTS:
+        findings.update(_scan(_REPO_ROOT / root, mask_only=False))
+    for root in MASK_SCAN_ROOTS:
+        findings.update(_scan(_REPO_ROOT / root, mask_only=True))
+
+    violations = {key: hits for key, hits in findings.items() if not _is_sanctioned(*key)}
+    assert not violations, "Unsanctioned host syncs on the mask-first path:\n" + "\n".join(
+        f"  {rel}::{name} -> {hits}" for (rel, name), hits in sorted(violations.items())
+    )
+
+    matched = {key for key in findings if _is_sanctioned(*key)}
+    stale = {
+        (suffix, func)
+        for suffix, func in SANCTIONED_BOUNDARIES
+        if not any(rel.endswith(suffix) and name == func for rel, name in matched)
+    }
+    assert not stale, f"Stale sanctioned boundaries (function gone or no longer syncs): {sorted(stale)}"
+
+
+# ``@WarpCapturable(False)`` opt-outs, by (path suffix, name).
+EXPECTED_NON_CAPTURABLE = {
+    ("isaaclab_experimental/envs/mdp/events.py", "randomize_rigid_body_com"),
+}
+
+NON_CAPTURABLE_SCAN_ROOTS = [
+    "source/isaaclab_experimental/isaaclab_experimental",
+    "source/isaaclab_tasks_experimental/isaaclab_tasks_experimental",
+]
+
+
+def _non_capturable_targets(tree: ast.Module) -> list[str]:
+    """Names of functions and classes decorated with ``@WarpCapturable(False, ...)``."""
+    targets = []
+    for node in ast.walk(tree):
+        if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+            continue
+        for decorator in node.decorator_list:
+            if (
+                isinstance(decorator, ast.Call)
+                and isinstance(decorator.func, ast.Name)
+                and decorator.func.id == "WarpCapturable"
+                and decorator.args
+                and isinstance(decorator.args[0], ast.Constant)
+                and decorator.args[0].value is False
+            ):
+                targets.append(node.name)
+    return targets
+
+
+def test_non_capturable_terms_are_inventoried():
+    """Every ``@WarpCapturable(False)`` opt-out is listed here, and only these."""
+    found = set()
+    for scan_root in NON_CAPTURABLE_SCAN_ROOTS:
+        for path in sorted((_REPO_ROOT / scan_root).rglob("*.py")):
+            rel = path.relative_to(_REPO_ROOT).as_posix()
+            if "/test/" in rel:
+                continue
+            for name in _non_capturable_targets(ast.parse(path.read_text(), filename=rel)):
+                found.add((rel, name))
+    expected = {(f"source/{suffix.split('/')[0]}/{suffix}", name) for suffix, name in EXPECTED_NON_CAPTURABLE}
+    assert found == expected, (
+        f"Non-capturable inventory drift.\n  unexpected: {sorted(found - expected)}\n"
+        f"  missing:    {sorted(expected - found)}"
+    )

From d528675c99dcf6add06d096760e0a76176c1ced0 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 03:12:38 -0700
Subject: [PATCH 10/48] Make stateful Warp event terms classes

push_by_setting_velocity, apply_external_force_torque,
reset_root_state_uniform and randomize_rigid_body_com kept scratch
arrays and parsed ranges as attributes of the module-level function. The
first call allocated them, which inside a recorded stage becomes
graph-owned memory; the parsed ranges froze the first configuration; and
the state was shared by every environment in the process.

They are now ManagerTermBase classes that allocate in __init__ and read
their ranges on every call. randomize_rigid_body_com offsets the center of
mass it finds on its first call instead of accumulating, and
apply_external_force_torque with zero ranges leaves the permanent wrench
composer untouched.

The material and mass adapters reject modes other than startup and are
marked non-capturable. Startup terms no longer keep the event stages
eager, since they never run inside one, and set_term_cfg re-resolves the
term and drops the recorded event stages.
---
 ...jichuanh-warp-frontend-cleanup.changed.rst |   7 +
 .../jichuanh-warp-frontend-cleanup.fixed.rst  |   6 +
 .../isaaclab_experimental/envs/mdp/events.py  | 397 ++++++++++--------
 .../managers/event_manager.py                 |  50 ++-
 .../managers/manager_base.py                  |  17 +-
 .../isaaclab_experimental/utils/warp/utils.py |  25 +-
 .../test/envs/mdp/parity_helpers.py           |   8 +-
 .../test/envs/mdp/test_events_warp_parity.py  | 154 +++++--
 .../test/envs/test_warp_env_reset.py          |  36 +-
 .../test/managers/test_event_manager.py       |  32 ++
 .../test/utils/test_mask_boundary_gate.py     |   2 +
 11 files changed, 507 insertions(+), 227 deletions(-)
 create mode 100644 source/isaaclab_experimental/test/managers/test_event_manager.py

diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst
index 1d4e33069b81..562e349a8a02 100644
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst
@@ -8,3 +8,10 @@
   to environment indices once per reset, and only when a command, curriculum or recorder term is active.
 * Changed :class:`~isaaclab_experimental.envs.DirectRLEnvWarp` to skip its reset stage on steps where no
   environment terminated.
+* Changed the Warp ``push_by_setting_velocity``, ``apply_external_force_torque``, ``reset_root_state_uniform``
+  and ``randomize_rigid_body_com`` event terms to classes that allocate their buffers when the term is
+  created. They read their ranges on every call, so a changed range applies to the next event.
+* Changed the Warp event manager to record its reset and interval stages when a non-capturable term runs
+  only at startup, such as ``randomize_rigid_body_com`` in the velocity tasks.
+* Changed the Warp ``randomize_rigid_body_material`` and ``randomize_rigid_body_mass`` adapters to reject
+  event modes other than ``startup``.
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
index 0699dcb75532..bdb84a7707a3 100644
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
@@ -1,3 +1,9 @@
 * Fixed Warp environments replaying CUDA graphs that read freed simulation buffers after the simulation was
   reset or stopped. The environments now record their stages again on the next step.
 * Fixed :attr:`~isaaclab_experimental.utils.buffers.CircularBuffer.max_length` synchronizing with the GPU.
+* Fixed the Warp ``randomize_rigid_body_com`` event accumulating offsets across calls. It now offsets the
+  center of mass the term finds on its first call, like the stable term.
+* Fixed the Warp ``apply_external_force_torque`` event clearing the permanent wrenches of reset environments
+  when its force and torque ranges are zero.
+* Fixed :meth:`EventManager.set_term_cfg <isaaclab_experimental.managers.EventManager.set_term_cfg>`
+  replaying recorded event stages with the replaced configuration.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py
index 90e5696a6b10..2f40fd98e61e 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py
@@ -17,6 +17,10 @@
 These Warp-first implementations avoid that by writing directly into the sim-bound Warp state buffers
 (`asset.data.joint_pos` / `asset.data.joint_vel`) for the selected envs/joints.
 
+Terms that need scratch buffers are classes that allocate them in ``__init__``: a recorded stage runs its
+terms for the first time inside CUDA graph capture, where an allocation would become graph-owned memory.
+Ranges are read from the call arguments, so a changed term configuration takes effect on the next call.
+
 Notes:
 - These terms assume the Newton/Warp backend (Warp arrays are available for joint state and defaults).
 - For best performance, pass :class:`isaaclab_experimental.managers.SceneEntityCfg` so `joint_ids_wp` is cached.
@@ -32,14 +36,29 @@
 from isaaclab.envs.mdp.events import randomize_rigid_body_mass as _StableRandomizeRigidBodyMass
 from isaaclab.envs.mdp.events import randomize_rigid_body_material as _StableRandomizeRigidBodyMaterial
 
-from isaaclab_experimental.managers import ManagerTermBase as _WarpManagerTermBase
-from isaaclab_experimental.managers import SceneEntityCfg
+from isaaclab_experimental.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg
 from isaaclab_experimental.utils.warp import WarpCapturable
 
 if TYPE_CHECKING:
-    from isaaclab.assets import Articulation
+    from isaaclab.assets import Articulation, RigidObject
     from isaaclab.envs import ManagerBasedEnv
 
+
+def _resolve_body_ids(asset_cfg: SceneEntityCfg, num_bodies: int) -> list[int]:
+    """Resolve a configured body selection to a list of body indices."""
+    if asset_cfg.body_ids is None or asset_cfg.body_ids == slice(None):
+        return list(range(num_bodies))
+    if isinstance(asset_cfg.body_ids, int):
+        return [asset_cfg.body_ids]
+    return list(asset_cfg.body_ids)
+
+
+def _vec3_range(ranges: dict[str, tuple[float, float]], keys: tuple[str, str, str]) -> tuple[wp.vec3f, wp.vec3f]:
+    """Lower and upper bounds of three named ranges; a missing key is the zero range."""
+    bounds = [ranges.get(key, (0.0, 0.0)) for key in keys]
+    return wp.vec3f(*(bound[0] for bound in bounds)), wp.vec3f(*(bound[1] for bound in bounds))
+
+
 # ---------------------------------------------------------------------------
 # Randomize rigid body center of mass
 # ---------------------------------------------------------------------------
@@ -49,66 +68,75 @@
 def _randomize_com_kernel(
     env_mask: wp.array(dtype=wp.bool),
     rng_state: wp.array(dtype=wp.uint32),
+    default_body_com_pos_b: wp.array(dtype=wp.vec3f, ndim=2),
     body_com_pos_b: wp.array(dtype=wp.vec3f, ndim=2),
     body_ids: wp.array(dtype=wp.int32),
     com_lo: wp.vec3f,
     com_hi: wp.vec3f,
 ):
-    """Add random offset to center of mass positions for selected bodies."""
+    """Offset the default center of mass of the selected bodies by one random offset per environment."""
     env_id = wp.tid()
     if not env_mask[env_id]:
         return
 
     state = rng_state[env_id]
+    dx = wp.randf(state, com_lo[0], com_hi[0])
+    dy = wp.randf(state, com_lo[1], com_hi[1])
+    dz = wp.randf(state, com_lo[2], com_hi[2])
     for k in range(body_ids.shape[0]):
         b = body_ids[k]
-        v = body_com_pos_b[env_id, b]
-        dx = wp.randf(state, com_lo[0], com_hi[0])
-        dy = wp.randf(state, com_lo[1], com_hi[1])
-        dz = wp.randf(state, com_lo[2], com_hi[2])
+        v = default_body_com_pos_b[env_id, b]
         body_com_pos_b[env_id, b] = wp.vec3f(v[0] + dx, v[1] + dy, v[2] + dz)
     rng_state[env_id] = state
 
 
 @WarpCapturable(False, reason="set_coms_mask calls SimulationManager.add_model_change")
-def randomize_rigid_body_com(
-    env,
-    env_mask: wp.array,
-    com_range: dict[str, tuple[float, float]],
-    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
-):
-    """Randomize the center of mass (CoM) of rigid bodies by adding random offsets.
+class randomize_rigid_body_com(ManagerTermBase):
+    """Randomize the center of mass (CoM) of rigid bodies around their default positions.
 
-    Warp-first override of :func:`isaaclab.envs.mdp.events.randomize_rigid_body_com`.
-    Writes directly into the sim-bound ``body_com_pos_b`` buffer, then notifies the solver
-    via :meth:`set_coms_mask` so it recomputes inertial properties.
+    Warp-first override of :class:`isaaclab.envs.mdp.events.randomize_rigid_body_com`. Writes directly into
+    the sim-bound ``body_com_pos_b`` buffer, then notifies the solver via :meth:`set_coms_mask` so it
+    recomputes inertial properties.
     """
-    asset: Articulation = env.scene[asset_cfg.name]
-
-    fn = randomize_rigid_body_com
-    if not getattr(fn, "_is_warmed_up", False) or fn._asset_name != asset_cfg.name:
-        fn._asset_name = asset_cfg.name
-        r = [com_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]]
-        fn._com_lo = wp.vec3f(r[0][0], r[1][0], r[2][0])
-        fn._com_hi = wp.vec3f(r[0][1], r[1][1], r[2][1])
-        fn._is_warmed_up = True
 
-    wp.launch(
-        kernel=_randomize_com_kernel,
-        dim=env.num_envs,
-        inputs=[
-            env_mask,
-            env.rng_state_wp,
-            asset.data.body_com_pos_b.warp,
-            asset_cfg.body_ids_wp,
-            fn._com_lo,
-            fn._com_hi,
-        ],
-        device=env.device,
-    )
+    def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv):
+        """Initialize the term.
+
+        Args:
+            cfg: The configuration of the event term.
+            env: The environment instance.
+        """
+        super().__init__(cfg, env)
+        asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"]
+        self._asset: RigidObject | Articulation = env.scene[asset_cfg.name]
+        self._body_ids = wp.array(
+            _resolve_body_ids(asset_cfg, self._asset.num_bodies), dtype=wp.int32, device=env.device
+        )
+        # Other startup terms may move the CoM after this term is built, so the defaults are copied on the
+        # first call. Randomizing from them keeps repeated calls from accumulating offsets.
+        self._default_com = wp.zeros_like(self._asset.data.body_com_pos_b.warp)
+        self._has_default_com = False
 
-    # Notify the solver that inertial properties changed (COM position affects inertia).
-    asset.set_coms_mask(coms=asset.data.body_com_pos_b.warp, env_mask=env_mask)
+    def __call__(
+        self,
+        env: ManagerBasedEnv,
+        env_mask: wp.array,
+        com_range: dict[str, tuple[float, float]],
+        asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
+    ):
+        """Offset the selected environments' CoM positions [m] by a sample from ``com_range``."""
+        body_com_pos_b = self._asset.data.body_com_pos_b.warp
+        if not self._has_default_com:
+            wp.copy(self._default_com, body_com_pos_b)
+            self._has_default_com = True
+        com_lo, com_hi = _vec3_range(com_range, ("x", "y", "z"))
+        wp.launch(
+            kernel=_randomize_com_kernel,
+            dim=env.num_envs,
+            inputs=[env_mask, env.rng_state_wp, self._default_com, body_com_pos_b, self._body_ids, com_lo, com_hi],
+            device=env.device,
+        )
+        self._asset.set_coms_mask(coms=body_com_pos_b, env_mask=env_mask)
 
 
 # ---------------------------------------------------------------------------
@@ -120,6 +148,7 @@ def randomize_rigid_body_com(
 def _apply_external_force_torque_kernel(
     env_mask: wp.array(dtype=wp.bool),
     rng_state: wp.array(dtype=wp.uint32),
+    body_ids: wp.array(dtype=wp.int32),
     force_out: wp.array(dtype=wp.vec3f, ndim=2),
     torque_out: wp.array(dtype=wp.vec3f, ndim=2),
     force_lo: float,
@@ -127,16 +156,14 @@ def _apply_external_force_torque_kernel(
     torque_lo: float,
     torque_hi: float,
 ):
+    """Sample uniform forces [N] and torques [N·m] for the selected bodies of the selected environments."""
     env_id = wp.tid()
     if not env_mask[env_id]:
-        # zero out unmasked envs so they don't accumulate stale forces
-        for b in range(force_out.shape[1]):
-            force_out[env_id, b] = wp.vec3f(0.0, 0.0, 0.0)
-            torque_out[env_id, b] = wp.vec3f(0.0, 0.0, 0.0)
         return
 
     state = rng_state[env_id]
-    for b in range(force_out.shape[1]):
+    for k in range(body_ids.shape[0]):
+        b = body_ids[k]
         force_out[env_id, b] = wp.vec3f(
             wp.randf(state, force_lo, force_hi),
             wp.randf(state, force_lo, force_hi),
@@ -150,50 +177,61 @@ def _apply_external_force_torque_kernel(
     rng_state[env_id] = state
 
 
-def apply_external_force_torque(
-    env,
-    env_mask: wp.array,
-    force_range: tuple[float, float],
-    torque_range: tuple[float, float],
-    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
-):
-    """Randomize external forces and torques applied to the asset's bodies.
+class apply_external_force_torque(ManagerTermBase):
+    """Randomize the external forces and torques applied to the asset's bodies.
 
     Warp-first override of :func:`isaaclab.envs.mdp.events.apply_external_force_torque`.
     """
-    asset: Articulation = env.scene[asset_cfg.name]
 
-    # First-call: allocate scratch and pre-convert constant arguments.
-    if not getattr(apply_external_force_torque, "_is_warmed_up", False):
-        apply_external_force_torque._scratch_forces = wp.zeros(
-            (env.num_envs, asset.num_bodies), dtype=wp.vec3f, device=env.device
+    def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv):
+        """Initialize the term.
+
+        Args:
+            cfg: The configuration of the event term.
+            env: The environment instance.
+        """
+        super().__init__(cfg, env)
+        asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"]
+        self._asset: RigidObject | Articulation = env.scene[asset_cfg.name]
+        num_bodies = self._asset.num_bodies
+        body_ids = _resolve_body_ids(asset_cfg, num_bodies)
+        self._body_ids = wp.array(body_ids, dtype=wp.int32, device=env.device)
+        self._body_mask = wp.array([i in body_ids for i in range(num_bodies)], dtype=wp.bool, device=env.device)
+        self._forces = wp.zeros((env.num_envs, num_bodies), dtype=wp.vec3f, device=env.device)
+        self._torques = wp.zeros((env.num_envs, num_bodies), dtype=wp.vec3f, device=env.device)
+
+    def __call__(
+        self,
+        env: ManagerBasedEnv,
+        env_mask: wp.array,
+        force_range: tuple[float, float],
+        torque_range: tuple[float, float],
+        asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
+    ):
+        """Set the selected bodies' wrenches to samples from ``force_range`` [N] and ``torque_range`` [N·m]."""
+        # Zero ranges configure a placeholder term. Skipping it keeps wrenches that other sources wrote to the
+        # permanent composer instead of clearing them on every reset.
+        if tuple(force_range) == (0.0, 0.0) and tuple(torque_range) == (0.0, 0.0):
+            return
+        wp.launch(
+            kernel=_apply_external_force_torque_kernel,
+            dim=env.num_envs,
+            inputs=[
+                env_mask,
+                env.rng_state_wp,
+                self._body_ids,
+                self._forces,
+                self._torques,
+                float(force_range[0]),
+                float(force_range[1]),
+                float(torque_range[0]),
+                float(torque_range[1]),
+            ],
+            device=env.device,
         )
-        apply_external_force_torque._scratch_torques = wp.zeros(
-            (env.num_envs, asset.num_bodies), dtype=wp.vec3f, device=env.device
+        self._asset.permanent_wrench_composer.set_forces_and_torques_mask(
+            forces=self._forces, torques=self._torques, body_mask=self._body_mask, env_mask=env_mask
         )
-        apply_external_force_torque._is_warmed_up = True
-
-    wp.launch(
-        kernel=_apply_external_force_torque_kernel,
-        dim=env.num_envs,
-        inputs=[
-            env_mask,
-            env.rng_state_wp,
-            apply_external_force_torque._scratch_forces,
-            apply_external_force_torque._scratch_torques,
-            force_range[0],
-            force_range[1],
-            torque_range[0],
-            torque_range[1],
-        ],
-        device=env.device,
-    )
-
-    asset.permanent_wrench_composer.set_forces_and_torques_mask(
-        forces=apply_external_force_torque._scratch_forces,
-        torques=apply_external_force_torque._scratch_torques,
-        env_mask=env_mask,
-    )
 
 
 # ---------------------------------------------------------------------------
@@ -231,45 +269,49 @@ def _push_by_setting_velocity_kernel(
     rng_state[env_id] = state
 
 
-def push_by_setting_velocity(
-    env,
-    env_mask: wp.array,
-    velocity_range: dict[str, tuple[float, float]],
-    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
-):
-    """Push the asset by setting the root velocity to a random value within the given ranges.
+class push_by_setting_velocity(ManagerTermBase):
+    """Push the asset by adding a random velocity to its root velocity.
 
     Warp-first override of :func:`isaaclab.envs.mdp.events.push_by_setting_velocity`.
     """
-    asset: Articulation = env.scene[asset_cfg.name]
 
-    # First-call: allocate scratch and pre-parse constant range arguments.
-    if not getattr(push_by_setting_velocity, "_is_warmed_up", False):
-        push_by_setting_velocity._scratch_vel = wp.zeros((env.num_envs,), dtype=wp.spatial_vectorf, device=env.device)
-        r = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]]
-        push_by_setting_velocity._lin_lo = wp.vec3f(r[0][0], r[1][0], r[2][0])
-        push_by_setting_velocity._lin_hi = wp.vec3f(r[0][1], r[1][1], r[2][1])
-        push_by_setting_velocity._ang_lo = wp.vec3f(r[3][0], r[4][0], r[5][0])
-        push_by_setting_velocity._ang_hi = wp.vec3f(r[3][1], r[4][1], r[5][1])
-        push_by_setting_velocity._is_warmed_up = True
+    def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv):
+        """Initialize the term.
 
-    wp.launch(
-        kernel=_push_by_setting_velocity_kernel,
-        dim=env.num_envs,
-        inputs=[
-            env_mask,
-            env.rng_state_wp,
-            asset.data.root_vel_w.warp,
-            push_by_setting_velocity._scratch_vel,
-            push_by_setting_velocity._lin_lo,
-            push_by_setting_velocity._lin_hi,
-            push_by_setting_velocity._ang_lo,
-            push_by_setting_velocity._ang_hi,
-        ],
-        device=env.device,
-    )
+        Args:
+            cfg: The configuration of the event term.
+            env: The environment instance.
+        """
+        super().__init__(cfg, env)
+        self._asset: RigidObject | Articulation = env.scene[cfg.params["asset_cfg"].name]
+        self._velocity = wp.zeros(env.num_envs, dtype=wp.spatial_vectorf, device=env.device)
 
-    asset.write_root_velocity_to_sim_mask(root_velocity=push_by_setting_velocity._scratch_vel, env_mask=env_mask)
+    def __call__(
+        self,
+        env: ManagerBasedEnv,
+        env_mask: wp.array,
+        velocity_range: dict[str, tuple[float, float]],
+        asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
+    ):
+        """Add samples from ``velocity_range`` [m/s, rad/s] to the selected environments' root velocity."""
+        lin_lo, lin_hi = _vec3_range(velocity_range, ("x", "y", "z"))
+        ang_lo, ang_hi = _vec3_range(velocity_range, ("roll", "pitch", "yaw"))
+        wp.launch(
+            kernel=_push_by_setting_velocity_kernel,
+            dim=env.num_envs,
+            inputs=[
+                env_mask,
+                env.rng_state_wp,
+                self._asset.data.root_vel_w.warp,
+                self._velocity,
+                lin_lo,
+                lin_hi,
+                ang_lo,
+                ang_hi,
+            ],
+            device=env.device,
+        )
+        self._asset.write_root_velocity_to_sim_mask(root_velocity=self._velocity, env_mask=env_mask)
 
 
 # ---------------------------------------------------------------------------
@@ -341,62 +383,61 @@ def _reset_root_state_uniform_kernel(
     rng_state[env_id] = state
 
 
-def reset_root_state_uniform(
-    env,
-    env_mask: wp.array,
-    pose_range: dict[str, tuple[float, float]],
-    velocity_range: dict[str, tuple[float, float]],
-    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
-):
+class reset_root_state_uniform(ManagerTermBase):
     """Reset the asset root state to a random position and velocity uniformly within the given ranges.
 
     Warp-first override of :func:`isaaclab.envs.mdp.events.reset_root_state_uniform`.
     """
-    asset: Articulation = env.scene[asset_cfg.name]
 
-    # First-call: allocate scratch and pre-parse range dicts.
-    if not getattr(reset_root_state_uniform, "_is_warmed_up", False):
-        reset_root_state_uniform._scratch_pose = wp.zeros((env.num_envs,), dtype=wp.transformf, device=env.device)
-        reset_root_state_uniform._scratch_vel = wp.zeros((env.num_envs,), dtype=wp.spatial_vectorf, device=env.device)
-        # Pre-parse pose_range dict
-        p = [pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]]
-        reset_root_state_uniform._pos_lo = wp.vec3f(p[0][0], p[1][0], p[2][0])
-        reset_root_state_uniform._pos_hi = wp.vec3f(p[0][1], p[1][1], p[2][1])
-        reset_root_state_uniform._rot_lo = wp.vec3f(p[3][0], p[4][0], p[5][0])
-        reset_root_state_uniform._rot_hi = wp.vec3f(p[3][1], p[4][1], p[5][1])
-        # Pre-parse velocity_range dict
-        v = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]]
-        reset_root_state_uniform._vel_lin_lo = wp.vec3f(v[0][0], v[1][0], v[2][0])
-        reset_root_state_uniform._vel_lin_hi = wp.vec3f(v[0][1], v[1][1], v[2][1])
-        reset_root_state_uniform._vel_ang_lo = wp.vec3f(v[3][0], v[4][0], v[5][0])
-        reset_root_state_uniform._vel_ang_hi = wp.vec3f(v[3][1], v[4][1], v[5][1])
-        reset_root_state_uniform._is_warmed_up = True
+    def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv):
+        """Initialize the term.
 
-    wp.launch(
-        kernel=_reset_root_state_uniform_kernel,
-        dim=env.num_envs,
-        inputs=[
-            env_mask,
-            env.rng_state_wp,
-            asset.data.default_root_pose.warp,
-            asset.data.default_root_vel.warp,
-            env.env_origins_wp,
-            reset_root_state_uniform._scratch_pose,
-            reset_root_state_uniform._scratch_vel,
-            reset_root_state_uniform._pos_lo,
-            reset_root_state_uniform._pos_hi,
-            reset_root_state_uniform._rot_lo,
-            reset_root_state_uniform._rot_hi,
-            reset_root_state_uniform._vel_lin_lo,
-            reset_root_state_uniform._vel_lin_hi,
-            reset_root_state_uniform._vel_ang_lo,
-            reset_root_state_uniform._vel_ang_hi,
-        ],
-        device=env.device,
-    )
+        Args:
+            cfg: The configuration of the event term.
+            env: The environment instance.
+        """
+        super().__init__(cfg, env)
+        self._asset: RigidObject | Articulation = env.scene[cfg.params["asset_cfg"].name]
+        self._pose = wp.zeros(env.num_envs, dtype=wp.transformf, device=env.device)
+        self._velocity = wp.zeros(env.num_envs, dtype=wp.spatial_vectorf, device=env.device)
 
-    asset.write_root_pose_to_sim_mask(root_pose=reset_root_state_uniform._scratch_pose, env_mask=env_mask)
-    asset.write_root_velocity_to_sim_mask(root_velocity=reset_root_state_uniform._scratch_vel, env_mask=env_mask)
+    def __call__(
+        self,
+        env: ManagerBasedEnv,
+        env_mask: wp.array,
+        pose_range: dict[str, tuple[float, float]],
+        velocity_range: dict[str, tuple[float, float]],
+        asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
+    ):
+        """Reset the selected environments' root pose [m, rad] and velocity [m/s, rad/s]."""
+        pos_lo, pos_hi = _vec3_range(pose_range, ("x", "y", "z"))
+        rot_lo, rot_hi = _vec3_range(pose_range, ("roll", "pitch", "yaw"))
+        vel_lin_lo, vel_lin_hi = _vec3_range(velocity_range, ("x", "y", "z"))
+        vel_ang_lo, vel_ang_hi = _vec3_range(velocity_range, ("roll", "pitch", "yaw"))
+        wp.launch(
+            kernel=_reset_root_state_uniform_kernel,
+            dim=env.num_envs,
+            inputs=[
+                env_mask,
+                env.rng_state_wp,
+                self._asset.data.default_root_pose.warp,
+                self._asset.data.default_root_vel.warp,
+                env.env_origins_wp,
+                self._pose,
+                self._velocity,
+                pos_lo,
+                pos_hi,
+                rot_lo,
+                rot_hi,
+                vel_lin_lo,
+                vel_lin_hi,
+                vel_ang_lo,
+                vel_ang_hi,
+            ],
+            device=env.device,
+        )
+        self._asset.write_root_pose_to_sim_mask(root_pose=self._pose, env_mask=env_mask)
+        self._asset.write_root_velocity_to_sim_mask(root_velocity=self._velocity, env_mask=env_mask)
 
 
 # ---------------------------------------------------------------------------
@@ -601,7 +642,7 @@ def reset_joints_by_scale(
 # physics backend (Newton included); only their calling convention differs —
 # the warp EventManager invokes terms with a Warp env-mask while the stable
 # class terms expect torch env indices. Both terms run in ``startup`` mode,
-# once at construction and outside any captured path, so a host-side
+# once at construction and outside any recorded stage, so a host-side
 # mask-to-ids conversion is acceptable.
 
 
@@ -610,7 +651,14 @@ def _mask_to_env_ids(env_mask: wp.array) -> torch.Tensor:
     return torch.nonzero(wp.to_torch(env_mask), as_tuple=False).squeeze(-1)
 
 
-class randomize_rigid_body_material(_StableRandomizeRigidBodyMaterial, _WarpManagerTermBase):
+def _require_startup_mode(cfg: EventTermCfg, term: type) -> None:
+    """Reject modes that would run a host-synchronizing adapter on every reset or interval."""
+    if cfg.mode != "startup":
+        raise ValueError(f"'{term.__name__}' supports only the 'startup' event mode, got '{cfg.mode}'.")
+
+
+@WarpCapturable(False, reason="converts the environment mask to indices on the host")
+class randomize_rigid_body_material(_StableRandomizeRigidBodyMaterial, ManagerTermBase):
     """Warp adapter for the stable, backend-dispatched material randomization term.
 
     Converts the warp event manager's env-mask calling convention to the stable
@@ -618,6 +666,10 @@ class randomize_rigid_body_material(_StableRandomizeRigidBodyMaterial, _WarpMana
     warp :class:`ManagerTermBase` so the warp managers accept it as a class term.
     """
 
+    def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv):
+        _require_startup_mode(cfg, type(self))
+        super().__init__(cfg, env)
+
     def __call__(
         self,
         env: ManagerBasedEnv,
@@ -645,7 +697,8 @@ def reset(self, env_mask: wp.array | None = None) -> None:
         return
 
 
-class randomize_rigid_body_mass(_StableRandomizeRigidBodyMass, _WarpManagerTermBase):
+@WarpCapturable(False, reason="converts the environment mask to indices on the host")
+class randomize_rigid_body_mass(_StableRandomizeRigidBodyMass, ManagerTermBase):
     """Warp adapter for the stable, backend-dispatched mass randomization term.
 
     Converts the warp event manager's env-mask calling convention to the stable
@@ -653,6 +706,10 @@ class randomize_rigid_body_mass(_StableRandomizeRigidBodyMass, _WarpManagerTermB
     warp :class:`ManagerTermBase` so the warp managers accept it as a class term.
     """
 
+    def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv):
+        _require_startup_mode(cfg, type(self))
+        super().__init__(cfg, env)
+
     def __call__(
         self,
         env: ManagerBasedEnv,
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
index 757d9a24d547..2b5329d92eca 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
@@ -31,7 +31,7 @@
 
 from isaaclab.managers.manager_term_cfg import EventTermCfg
 
-from .manager_base import ManagerBase
+from .manager_base import ManagerBase, ManagerTermBase
 
 logger = logging.getLogger(__name__)
 
@@ -223,14 +223,43 @@ def available_modes(self) -> list[str]:
         return list(self._mode_term_names.keys())
 
     def set_term_cfg(self, term_name: str, cfg: EventTermCfg):
-        term_found = False
+        """Set the configuration of an existing event term.
+
+        Recorded event stages read the replaced configuration, so they are recorded again on their next call.
+
+        Args:
+            term_name: Name of the event term.
+            cfg: Replacement event term configuration.
+
+        Raises:
+            ValueError: If the term does not exist, or the replacement changes its mode or interval timing.
+        """
         for mode, terms in self._mode_term_names.items():
-            if term_name in terms:
-                self._mode_term_cfgs[mode][terms.index(term_name)] = cfg
-                term_found = True
-                break
-        if not term_found:
-            raise ValueError(f"Event term '{term_name}' not found.")
+            if term_name not in terms:
+                continue
+            if cfg.mode != mode:
+                raise ValueError(f"Event term '{term_name}' has mode '{mode}', but the replacement has '{cfg.mode}'.")
+            term_index = terms.index(term_name)
+            if mode == "interval":
+                if cfg.interval_range_s is None:
+                    raise ValueError(f"Event term '{term_name}' has mode 'interval' but 'interval_range_s' is not set.")
+                if cfg.is_global_time != self._mode_term_cfgs[mode][term_index].is_global_time:
+                    raise ValueError(f"Event term '{term_name}' cannot change 'is_global_time' at runtime.")
+            self._resolve_common_term_cfg(term_name, cfg, min_argc=2)
+            self._mode_term_cfgs[mode][term_index] = cfg
+            if mode == "interval":
+                lower, upper = cfg.interval_range_s
+                self._interval_term_ranges[term_index] = (float(lower), float(upper))
+            self._mode_class_term_cfgs[mode] = [
+                term_cfg
+                for term_cfg in self._mode_term_cfgs[mode]
+                if inspect.isclass(term_cfg.func) or isinstance(term_cfg.func, ManagerTermBase)
+            ]
+            graph_cache = getattr(self._env, "_warp_graph_cache", None)
+            if graph_cache is not None:
+                graph_cache.invalidate(type(self).__name__)
+            return
+        raise ValueError(f"Event term '{term_name}' not found.")
 
     def get_term_cfg(self, term_name: str) -> EventTermCfg:
         for mode, terms in self._mode_term_names.items():
@@ -395,6 +424,11 @@ def _apply_reset(self, env_mask_wp: wp.array, global_env_step_count_wp: wp.array
             )
             term_cfg.func(self._env, self._scratch_term_mask_wp, **term_cfg.params)
 
+    def _register_term_capturability(self, term_cfg: EventTermCfg) -> None:
+        # startup terms run once while the environment is built, outside the recorded stages
+        if term_cfg.mode not in ("prestartup", "startup"):
+            super()._register_term_capturability(term_cfg)
+
     def _prepare_terms(self):
         # check if config is dict already
         if isinstance(self.cfg, dict):
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py
index 3e3325138f3f..f30d62040506 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py
@@ -411,17 +411,24 @@ def _resolve_common_term_cfg(self, term_name: str, term_cfg: ManagerTermBaseCfg,
                     f" and optional parameters: {args_with_defaults}, but received: {term_params}."
                 )
 
-        # a non-capturable term keeps every stage of this manager eager
-        if not is_warp_capturable(term_cfg.func):
-            graph_cache = getattr(self._env, "_warp_graph_cache", None)
-            if graph_cache is not None:
-                graph_cache.register_capturability(type(self).__name__, False)
+        self._register_term_capturability(term_cfg)
 
         # process attributes at runtime
         # these properties are only resolvable once the simulation starts playing
         if self._env.sim.is_playing():
             self._process_term_cfg_at_play(term_name, term_cfg)
 
+    def _register_term_capturability(self, term_cfg: ManagerTermBaseCfg) -> None:
+        """Keep every stage of this manager eager when the term is not capturable.
+
+        Args:
+            term_cfg: The term configuration.
+        """
+        if not is_warp_capturable(term_cfg.func):
+            graph_cache = getattr(self._env, "_warp_graph_cache", None)
+            if graph_cache is not None:
+                graph_cache.register_capturability(type(self).__name__, False)
+
     def _process_term_cfg_at_play(self, term_name: str, term_cfg: ManagerTermBaseCfg):
         """Process the term configuration at runtime.
 
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py
index a90f6b2d6fc6..cb249af9fa02 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py
@@ -146,26 +146,33 @@ class WarpCapturable:
         def reset_root_state_uniform(env, env_mask, ...):
             ...
 
-        @WarpCapturable(False, reason="calls write_root_pose_to_sim")
-        def push_by_setting_velocity(env, env_mask, ...):
+        @WarpCapturable(False, reason="notifies the solver of a model change")
+        class randomize_rigid_body_com(ManagerTermBase):
             ...
 
     - ``@WarpCapturable(True)`` or no decorator: capturable, returned unwrapped.
-    - ``@WarpCapturable(False)``: sets ``func._warp_capturable = False``, wraps with
+    - ``@WarpCapturable(False)`` on a function: sets ``func._warp_capturable = False`` and wraps it with a
       runtime guard that raises if ``wp.get_device().is_capturing`` is ``True``.
+    - ``@WarpCapturable(False)`` on a class term: annotates the class and guards its ``__call__``.
     """
 
     def __init__(self, capturable: bool, *, reason: str | None = None):
         self._capturable = capturable
         self._reason = reason
 
-    def __call__(self, func):
-        """Decorate *func* with capture safety annotation and optional runtime guard."""
-        import functools
-
-        func._warp_capturable = self._capturable
+    def __call__(self, target):
+        """Decorate a term function or class with the capture safety annotation and runtime guard."""
+        target._warp_capturable = self._capturable
         if self._capturable:
-            return func
+            return target
+        if isinstance(target, type):
+            target.__call__ = self._guarded(target.__call__)
+            return target
+        return self._guarded(target)
+
+    def _guarded(self, func):
+        """Wrap *func* so that calling it during CUDA graph capture raises."""
+        import functools
 
         reason = self._reason
 
diff --git a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py
index 758b2f9ecbee..e9e07bca5e9f 100644
--- a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py
+++ b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py
@@ -279,15 +279,19 @@ def resolve_joint_mask(self, joint_ids=None):
 
 
 class MockWrenchComposer:
-    """Mock wrench composer that records the last forces and torques it was handed."""
+    """Mock wrench composer that records the last wrenches and selection it was handed."""
 
     def __init__(self):
         self.last_forces = None
         self.last_torques = None
+        self.last_body_mask = None
+        self.last_env_mask = None
 
-    def set_forces_and_torques_mask(self, forces=None, torques=None, **kw):
+    def set_forces_and_torques_mask(self, forces=None, torques=None, body_mask=None, env_mask=None, **kw):
         self.last_forces = forces
         self.last_torques = torques
+        self.last_body_mask = body_mask
+        self.last_env_mask = env_mask
 
 
 class MockArticulation:
diff --git a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py
index e93e73267938..d7451d300933 100644
--- a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py
+++ b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py
@@ -7,6 +7,8 @@
 
 from __future__ import annotations
 
+from types import SimpleNamespace
+
 import numpy as np
 import pytest
 import torch
@@ -17,6 +19,7 @@
 pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required")
 
 import isaaclab_experimental.envs.mdp.events as warp_evt
+from isaaclab_experimental.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg
 from parity_helpers import (
     DEVICE,
     NUM_ACTIONS,
@@ -36,25 +39,6 @@
 # ============================================================================
 
 
-@pytest.fixture(autouse=True)
-def _clear_function_caches():
-    """Clear first-call caches on warp MDP functions so each test starts fresh.
-
-    Functions that cache warp views via the ``hasattr`` pattern need clearing
-    between tests to avoid stale references from prior fixtures.
-    """
-    yield
-    for fn in (
-        warp_evt.push_by_setting_velocity,
-        warp_evt.apply_external_force_torque,
-        warp_evt.reset_root_state_uniform,
-        warp_evt.randomize_rigid_body_com,
-    ):
-        for attr in list(vars(fn)):
-            if attr.startswith("_"):
-                delattr(fn, attr)
-
-
 @pytest.fixture()
 def art_data():
     return MockArticulationData(NUM_ENVS, NUM_JOINTS, DEVICE)
@@ -101,6 +85,7 @@ class _Env:
     env.episode_length_buf = episode_length_buf
     env.step_dt = 0.02
     env.max_episode_length_s = 10.0
+    env.env_origins_wp = scene.env_origins
     # RNG state for events (seeded deterministically)
     env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
@@ -111,6 +96,16 @@ def all_joints_cfg():
     return MockSceneEntityCfg("robot", list(range(NUM_JOINTS)), NUM_JOINTS, DEVICE)
 
 
+def _make_event_term(term_type: type[ManagerTermBase], env, mode: str = "reset", **params) -> ManagerTermBase:
+    """Build a class event term the way the event manager does."""
+    params.setdefault("asset_cfg", SceneEntityCfg("robot"))
+    return term_type(EventTermCfg(func=term_type, mode=mode, params=params), env)
+
+
+def _all_envs() -> wp.array:
+    return wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE)
+
+
 # ============================================================================
 # Event capture-mutate-replay tests
 # ============================================================================
@@ -192,7 +187,7 @@ def test_reset_joints_by_scale(self, warp_env, art_data, all_joints_cfg):
 
     def test_push_by_setting_velocity(self, warp_env, art_data, all_joints_cfg):
         """With zero-width velocity range, the written root velocity == root_vel_w.  Mutate root_vel_w -> tracks."""
-        mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE)
+        mask = _all_envs()
         zero_range = {
             "x": (0.0, 0.0),
             "y": (0.0, 0.0),
@@ -201,10 +196,10 @@ def test_push_by_setting_velocity(self, warp_env, art_data, all_joints_cfg):
             "pitch": (0.0, 0.0),
             "yaw": (0.0, 0.0),
         }
+        term = _make_event_term(warp_evt.push_by_setting_velocity, warp_env, mode="interval", velocity_range=zero_range)
 
-        warp_evt.push_by_setting_velocity(warp_env, mask, velocity_range=zero_range)
         with wp.ScopedCapture() as cap:
-            warp_evt.push_by_setting_velocity(warp_env, mask, velocity_range=zero_range)
+            term(warp_env, mask, **term.cfg.params)
 
         # Mutate root_vel_w
         new_vel = np.tile([1.0, 2.0, 3.0, 0.1, 0.2, 0.3], (NUM_ENVS, 1)).astype(np.float32)
@@ -220,26 +215,34 @@ def test_push_by_setting_velocity(self, warp_env, art_data, all_joints_cfg):
     # -- apply_external_force_torque --------------------------------------------
 
     def test_apply_external_force_torque(self, warp_env, art_data, all_joints_cfg):
-        """A degenerate non-zero range reaches masked envs through the replayed kernel; unmasked envs get zero."""
+        """A degenerate non-zero range reaches masked envs through the replayed kernel; the composer gets the mask."""
         mask = wp.array([i < NUM_ENVS // 2 for i in range(NUM_ENVS)], dtype=wp.bool, device=DEVICE)
+        term = _make_event_term(
+            warp_evt.apply_external_force_torque, warp_env, force_range=(2.0, 2.0), torque_range=(3.0, 3.0)
+        )
 
-        warp_evt.apply_external_force_torque(warp_env, mask, force_range=(2.0, 2.0), torque_range=(3.0, 3.0))
         with wp.ScopedCapture() as cap:
-            warp_evt.apply_external_force_torque(warp_env, mask, force_range=(2.0, 2.0), torque_range=(3.0, 3.0))
+            term(warp_env, mask, **term.cfg.params)
         composer = warp_env.scene["robot"].permanent_wrench_composer
         forces = wp.to_torch(composer.last_forces)
         torques = wp.to_torch(composer.last_torques)
-        # clear the buffers so only the replayed kernel can produce the expected values
-        forces.fill_(-1.0)
-        torques.fill_(-1.0)
         wp.capture_launch(cap.graph)
         wp.synchronize()
 
         half = NUM_ENVS // 2
         assert_close(forces[:half], torch.full_like(forces[:half], 2.0))
         assert_close(torques[:half], torch.full_like(torques[:half], 3.0))
-        assert_close(forces[half:], torch.zeros_like(forces[half:]))
-        assert_close(torques[half:], torch.zeros_like(torques[half:]))
+        assert composer.last_env_mask is mask
+
+    def test_apply_external_force_torque_zero_ranges_are_no_op(self, warp_env, art_data, all_joints_cfg):
+        """Zero ranges keep the composer's wrenches: a placeholder term must not clear them on every reset."""
+        term = _make_event_term(
+            warp_evt.apply_external_force_torque, warp_env, force_range=(0.0, 0.0), torque_range=(0.0, 0.0)
+        )
+
+        term(warp_env, _all_envs(), **term.cfg.params)
+
+        assert warp_env.scene["robot"].permanent_wrench_composer.last_forces is None
 
     # -- env_mask selectivity ---------------------------------------------------
 
@@ -267,3 +270,94 @@ def test_reset_joints_mask_selectivity(self, warp_env, art_data, all_joints_cfg)
         assert_close(result[: NUM_ENVS // 2], torch.zeros(NUM_ENVS // 2, NUM_JOINTS, device=DEVICE))
         # Unmasked envs: still 999.0
         assert_close(result[NUM_ENVS // 2 :], torch.full((NUM_ENVS // 2, NUM_JOINTS), 999.0, device=DEVICE))
+
+
+# ============================================================================
+# Class event terms: persistent state and live configuration
+# ============================================================================
+
+
+def _make_event_env(seed: int, num_bodies: int = 1):
+    """An independent Warp event environment with its own asset."""
+    data = MockArticulationData(NUM_ENVS, NUM_JOINTS, DEVICE, seed=seed, num_bodies=num_bodies)
+    asset = MockArticulation(data, num_bodies=num_bodies)
+    origins = wp.zeros(NUM_ENVS, dtype=wp.vec3f, device=DEVICE)
+    scene = MockScene({"robot": asset}, origins)
+    env = SimpleNamespace(
+        scene=scene,
+        num_envs=NUM_ENVS,
+        device=DEVICE,
+        env_origins_wp=origins,
+        rng_state_wp=wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + seed, device=DEVICE),
+    )
+    return env, data, asset
+
+
+@pytest.mark.parametrize(
+    ("term_type", "params"),
+    [
+        (warp_evt.push_by_setting_velocity, {"velocity_range": {"x": (1.0, 1.0)}}),
+        (warp_evt.apply_external_force_torque, {"force_range": (1.0, 1.0), "torque_range": (2.0, 2.0)}),
+        (warp_evt.reset_root_state_uniform, {"pose_range": {}, "velocity_range": {}}),
+    ],
+    ids=["push", "external_wrench", "root_reset"],
+)
+def test_class_terms_allocate_only_during_initialization(term_type, params, monkeypatch):
+    """A recorded stage runs its terms for the first time inside CUDA graph capture, so a call must not allocate."""
+    env, _, _ = _make_event_env(81)
+    term = _make_event_term(term_type, env, **params)
+
+    def fail(*args, **kwargs):
+        pytest.fail("the term allocated during a call")
+
+    monkeypatch.setattr(wp, "zeros", fail)
+    monkeypatch.setattr(wp, "empty", fail)
+    monkeypatch.setattr(wp, "clone", fail)
+    term(env, _all_envs(), **term.cfg.params)
+    term(env, _all_envs(), **term.cfg.params)
+
+
+def test_push_range_change_applies_on_the_next_call():
+    """A curriculum that edits the term's range in place changes the next push, as in the stable term."""
+    env, data, asset = _make_event_env(277)
+    copy_np_to_wp(data.root_vel_w, np.zeros((NUM_ENVS, 6), dtype=np.float32))
+    term = _make_event_term(warp_evt.push_by_setting_velocity, env, mode="interval", velocity_range={"x": (1.0, 1.0)})
+    term(env, _all_envs(), **term.cfg.params)
+
+    term.cfg.params["velocity_range"]["x"] = (4.0, 4.0)
+    term(env, _all_envs(), **term.cfg.params)
+    wp.synchronize()
+
+    assert_close(wp.to_torch(asset.last_root_velocity)[:, 0], torch.full((NUM_ENVS,), 4.0, device=DEVICE))
+
+
+def test_com_randomization_uses_the_first_call_baseline_without_accumulating():
+    """Startup terms may move the CoM after construction; each call offsets that baseline, not the last result."""
+    env, data, asset = _make_event_env(84, num_bodies=2)
+    baseline = np.zeros((NUM_ENVS, 2, 3), dtype=np.float32)
+    copy_np_to_wp(data.body_com_pos_b, baseline)
+    asset.set_coms_mask = lambda **kwargs: None
+    term = _make_event_term(
+        warp_evt.randomize_rigid_body_com,
+        env,
+        mode="startup",
+        com_range={"x": (1.0, 1.0)},
+        asset_cfg=SceneEntityCfg("robot", body_ids=[0, 1]),
+    )
+    # another startup term moves the CoM between construction and the first call
+    baseline[..., 0] = 0.75
+    copy_np_to_wp(data.body_com_pos_b, baseline)
+
+    term(env, _all_envs(), **term.cfg.params)
+    term(env, _all_envs(), **term.cfg.params)
+    wp.synchronize()
+
+    assert_close(data.body_com_pos_b.torch[..., 0], torch.full((NUM_ENVS, 2), 1.75, device=DEVICE))
+
+
+@pytest.mark.parametrize("term_type", [warp_evt.randomize_rigid_body_material, warp_evt.randomize_rigid_body_mass])
+def test_mask_adapters_reject_modes_other_than_startup(term_type):
+    """The adapters convert the mask to indices on the host, which a reset or interval stage must not do."""
+    cfg = EventTermCfg(func=term_type, mode="reset", params={"asset_cfg": SceneEntityCfg("robot")})
+    with pytest.raises(ValueError, match="startup"):
+        term_type(cfg, SimpleNamespace())
diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_reset.py b/source/isaaclab_experimental/test/envs/test_warp_env_reset.py
index b5148ca7275c..c75dea09e091 100644
--- a/source/isaaclab_experimental/test/envs/test_warp_env_reset.py
+++ b/source/isaaclab_experimental/test/envs/test_warp_env_reset.py
@@ -3,7 +3,7 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
-"""Mask-first resets of the manager-based Warp environment."""
+"""Resets of the manager-based Warp environment."""
 
 from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg
 
@@ -23,12 +23,16 @@
 _NUM_ENVS = 8
 
 
-def test_reset_with_mask_resets_only_selected_environments():
+def _make_cartpole():
     env_cfg, _ = resolve_task_config("Isaac-Cartpole", "", overrides=("physics=newton_mjwarp",))
     env_cfg.seed = 7
     env_cfg.scene.num_envs = _NUM_ENVS
     sim_utils.create_new_stage()
-    env = WarpFrontend.build_env(env_cfg, "Isaac-Cartpole").unwrapped
+    return WarpFrontend.build_env(env_cfg, "Isaac-Cartpole").unwrapped
+
+
+def test_reset_with_mask_resets_only_selected_environments():
+    env = _make_cartpole()
     try:
         env.reset()
         actions = torch.ones((_NUM_ENVS, env.action_space.shape[-1]), device=env.device)
@@ -47,3 +51,29 @@ def test_reset_with_mask_resets_only_selected_environments():
         assert not torch.isclose(robot.data.joint_pos.torch[selected], joint_pos[selected]).all(dim=1).any()
     finally:
         env.close()
+
+
+def test_set_term_cfg_applies_to_recorded_reset_events():
+    """A replaced event term configuration must not replay the scalars recorded from the old one."""
+    env = _make_cartpole()
+    try:
+        env.reset()
+        env.step(torch.zeros((_NUM_ENVS, env.action_space.shape[-1]), device=env.device))
+        selected = torch.arange(_NUM_ENVS, device=env.device) % 2 == 0
+        # records the reset events with the configured random cart offsets
+        env.reset(env_mask=selected)
+        assert "EventManager_apply_reset" in env._warp_graph_cache.captured_stages
+
+        term_cfg = env.event_manager.get_term_cfg("reset_cart_position")
+        term_cfg.params["position_range"] = (0.5, 0.5)
+        term_cfg.params["velocity_range"] = (0.0, 0.0)
+        env.event_manager.set_term_cfg("reset_cart_position", term_cfg)
+        env.reset(env_mask=selected)
+
+        robot = env.scene["robot"]
+        cart = robot.find_joints("slider_to_cart")[0][0]
+        cart_pos = robot.data.joint_pos.torch[selected, cart]
+        default_pos = robot.data.default_joint_pos.torch[selected, cart]
+        assert torch.allclose(cart_pos, default_pos + 0.5)
+    finally:
+        env.close()
diff --git a/source/isaaclab_experimental/test/managers/test_event_manager.py b/source/isaaclab_experimental/test/managers/test_event_manager.py
new file mode 100644
index 000000000000..3c6fb4199660
--- /dev/null
+++ b/source/isaaclab_experimental/test/managers/test_event_manager.py
@@ -0,0 +1,32 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tests for the Warp event manager."""
+
+from types import SimpleNamespace
+
+import pytest
+import warp as wp
+from isaaclab_experimental.managers import EventManager, EventTermCfg
+from isaaclab_experimental.utils import WarpGraphCache
+from isaaclab_experimental.utils.warp import WarpCapturable
+
+
+@WarpCapturable(False, reason="synchronizes with the host")
+def host_synchronizing_term(env, env_mask: wp.array):
+    """Event term that cannot run inside a recorded stage."""
+
+
+@pytest.mark.parametrize(("mode", "stages_capturable"), [("startup", True), ("reset", False)])
+def test_non_capturable_term_keeps_event_stages_eager_only_when_it_runs_in_them(mode: str, stages_capturable: bool):
+    """Startup terms run once while the environment is built, so they must not keep the reset stages eager."""
+    graph_cache = WarpGraphCache("cpu")
+    env = SimpleNamespace(
+        num_envs=4, device="cpu", sim=SimpleNamespace(is_playing=lambda: True), _warp_graph_cache=graph_cache
+    )
+
+    EventManager({"term": EventTermCfg(func=host_synchronizing_term, mode=mode)}, env)
+
+    assert graph_cache.is_capturable("EventManager") is stages_capturable
diff --git a/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py
index a6c7fd74511d..cb193e398be3 100644
--- a/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py
+++ b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py
@@ -137,6 +137,8 @@ def test_mask_first_code_has_no_unsanctioned_host_syncs():
 # ``@WarpCapturable(False)`` opt-outs, by (path suffix, name).
 EXPECTED_NON_CAPTURABLE = {
     ("isaaclab_experimental/envs/mdp/events.py", "randomize_rigid_body_com"),
+    ("isaaclab_experimental/envs/mdp/events.py", "randomize_rigid_body_mass"),
+    ("isaaclab_experimental/envs/mdp/events.py", "randomize_rigid_body_material"),
 }
 
 NON_CAPTURABLE_SCAN_ROOTS = [

From 80cacc6070323a5c20ca32ffde2c360817e41c87 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 03:28:37 -0700
Subject: [PATCH 11/48] Test event capture through the recorded reset stage

---
 .../test/managers/test_event_manager.py       | 40 +++++++++++++++----
 1 file changed, 32 insertions(+), 8 deletions(-)

diff --git a/source/isaaclab_experimental/test/managers/test_event_manager.py b/source/isaaclab_experimental/test/managers/test_event_manager.py
index 3c6fb4199660..ae052362cdb7 100644
--- a/source/isaaclab_experimental/test/managers/test_event_manager.py
+++ b/source/isaaclab_experimental/test/managers/test_event_manager.py
@@ -13,20 +13,44 @@
 from isaaclab_experimental.utils import WarpGraphCache
 from isaaclab_experimental.utils.warp import WarpCapturable
 
+wp.init()
+pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required")
+
+DEVICE = "cuda:0"
+NUM_ENVS = 4
+
 
 @WarpCapturable(False, reason="synchronizes with the host")
 def host_synchronizing_term(env, env_mask: wp.array):
     """Event term that cannot run inside a recorded stage."""
 
 
-@pytest.mark.parametrize(("mode", "stages_capturable"), [("startup", True), ("reset", False)])
-def test_non_capturable_term_keeps_event_stages_eager_only_when_it_runs_in_them(mode: str, stages_capturable: bool):
-    """Startup terms run once while the environment is built, so they must not keep the reset stages eager."""
-    graph_cache = WarpGraphCache("cpu")
+def device_term(env, env_mask: wp.array):
+    """Event term whose work is entirely on the device."""
+
+
+def test_startup_only_term_does_not_keep_the_reset_stage_eager():
+    """Startup terms run once while the environment is built, so they must not stop reset events from recording."""
+    graph_cache = WarpGraphCache(DEVICE)
     env = SimpleNamespace(
-        num_envs=4, device="cpu", sim=SimpleNamespace(is_playing=lambda: True), _warp_graph_cache=graph_cache
+        num_envs=NUM_ENVS, device=DEVICE, sim=SimpleNamespace(is_playing=lambda: True), _warp_graph_cache=graph_cache
+    )
+    manager = EventManager(
+        {
+            "startup_term": EventTermCfg(func=host_synchronizing_term, mode="startup"),
+            "reset_term": EventTermCfg(func=device_term, mode="reset"),
+        },
+        env,
+    )
+    graph_cache.arm()
+
+    graph_cache.call(
+        "EventManager_apply_reset",
+        manager.apply,
+        mode="reset",
+        env_mask_wp=wp.ones(NUM_ENVS, dtype=wp.bool, device=DEVICE),
+        global_env_step_count=wp.zeros(1, dtype=wp.int32, device=DEVICE),
     )
 
-    EventManager({"term": EventTermCfg(func=host_synchronizing_term, mode=mode)}, env)
-
-    assert graph_cache.is_capturable("EventManager") is stages_capturable
+    assert graph_cache.captured_stages == ("EventManager_apply_reset",)
+    graph_cache.close()

From ca4eefbd2d55235098b892d8b8c66030519ebe4e Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 03:46:19 -0700
Subject: [PATCH 12/48] Decide Warp stage capture per term

A manager with one non-capturable term used to run every stage eagerly.
Managers now describe each stage as ordered steps, one per term plus
their own setup and finalize kernels, and WarpGraphCache.call_steps
records each run of consecutive capturable steps into one graph while
the terms marked @WarpCapturable(False) run eagerly in place.

Reward and termination terms write their own output columns and the
finalize kernel combines the columns in a fixed order, so these stages
evaluate the capturable terms first and need at most one eager group.
Observation noise and events draw from the shared per-environment random
state, and action terms are applied in order, so those stages keep the
configured term order and split into contiguous runs.

Observation terms with history run as eager steps instead of raising
inside capture. Replacing a reward term rebuilds its steps and drops the
recorded reward stages.
---
 .../jichuanh-warp-frontend-cleanup.added.rst  |   3 +
 ...jichuanh-warp-frontend-cleanup.changed.rst |  11 +-
 .../jichuanh-warp-frontend-cleanup.fixed.rst  |   5 +-
 .../envs/manager_based_env_warp.py            |  33 +--
 .../envs/manager_based_rl_env_warp.py         |  71 ++++---
 .../managers/action_manager.py                |  65 ++++--
 .../managers/event_manager.py                 | 195 +++++++++--------
 .../managers/manager_base.py                  |  68 ++++--
 .../managers/observation_manager.py           | 199 +++++++++++-------
 .../managers/reward_manager.py                |  95 ++++++---
 .../managers/termination_manager.py           |  71 +++++--
 .../utils/warp_graph_cache.py                 | 107 ++++++++--
 .../test/envs/test_warp_env_capture.py        |  45 +++-
 .../test/utils/test_warp_graph_cache.py       |  33 +++
 14 files changed, 672 insertions(+), 329 deletions(-)

diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst
index dee68bc25e57..37982f7d807d 100644
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst
@@ -3,3 +3,6 @@
 * Added the ``ISAACLAB_SYNC_DEBUG`` environment variable: ``ISAACLAB_SYNC_DEBUG=1`` raises on a hidden
   GPU-to-host synchronization inside an eagerly executed Warp environment stage.
 * Added a boolean ``env_mask`` argument to :meth:`~isaaclab_experimental.envs.ManagerBasedEnvWarp.reset`.
+* Added :meth:`~isaaclab_experimental.utils.WarpGraphCache.call_steps` and
+  :meth:`ManagerBase.stage_steps <isaaclab_experimental.managers.ManagerBase.stage_steps>`, which split a
+  manager stage into per-term steps that are recorded into CUDA graphs or run eagerly.
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst
index 562e349a8a02..895b48a3a557 100644
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst
@@ -1,6 +1,11 @@
 * **Breaking:** Removed ``ManagerCallSwitch``, ``ManagerCallMode`` and the ``MANAGER_CALL_CONFIG``
-  environment variable. Warp environments construct their Warp managers directly and keep a manager's
-  stages eager when one of its terms is not capturable.
+  environment variable. Warp environments construct their Warp managers directly.
+* Changed the Warp environments to decide CUDA graph capture per term: a stage records its capturable
+  terms and runs each term decorated with ``@WarpCapturable(False)`` eagerly, instead of running the whole
+  stage eagerly. Reward and termination stages evaluate their capturable terms first; observation, action
+  and event stages keep the configured term order.
+* **Breaking:** Removed ``ManagerBasedRLEnvWarp.step_warp_termination_compute``. The termination stage runs
+  :meth:`TerminationManager.compute <isaaclab_experimental.managers.TerminationManager.compute>` directly.
 * **Breaking:** Replaced ``WarpGraphCache.capture_or_replay`` with
   :meth:`~isaaclab_experimental.utils.WarpGraphCache.call`. Stages are recorded from the first environment
   step without an eager warm-up, so a stage runs exactly once per call.
@@ -11,7 +16,5 @@
 * Changed the Warp ``push_by_setting_velocity``, ``apply_external_force_torque``, ``reset_root_state_uniform``
   and ``randomize_rigid_body_com`` event terms to classes that allocate their buffers when the term is
   created. They read their ranges on every call, so a changed range applies to the next event.
-* Changed the Warp event manager to record its reset and interval stages when a non-capturable term runs
-  only at startup, such as ``randomize_rigid_body_com`` in the velocity tasks.
 * Changed the Warp ``randomize_rigid_body_material`` and ``randomize_rigid_body_mass`` adapters to reject
   event modes other than ``startup``.
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
index bdb84a7707a3..9fbaf0c2ae3c 100644
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
@@ -5,5 +5,6 @@
   center of mass the term finds on its first call, like the stable term.
 * Fixed the Warp ``apply_external_force_torque`` event clearing the permanent wrenches of reset environments
   when its force and torque ranges are zero.
-* Fixed :meth:`EventManager.set_term_cfg <isaaclab_experimental.managers.EventManager.set_term_cfg>`
-  replaying recorded event stages with the replaced configuration.
+* Fixed :meth:`EventManager.set_term_cfg <isaaclab_experimental.managers.EventManager.set_term_cfg>` and
+  :meth:`RewardManager.set_term_cfg <isaaclab_experimental.managers.RewardManager.set_term_cfg>` replaying
+  recorded stages with the replaced term.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
index e2b9a3ad4f48..1c30edde04f2 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
@@ -561,8 +561,8 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]:
         # process actions
         action_device = action.to(device=self.device, dtype=torch.float32).contiguous()
         wp.copy(self._action_in_wp, wp.from_torch(action_device, dtype=wp.float32))
-        self._warp_graph_cache.call(
-            "ActionManager_process_action", self.action_manager.process_action, action=self._action_in_wp
+        self._warp_graph_cache.call_steps(
+            "ActionManager_process_action", self.action_manager.stage_steps("process_action"), action=self._action_in_wp
         )
 
         self.recorder_manager.record_pre_step()
@@ -575,7 +575,9 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]:
         for _ in range(self.cfg.decimation):
             self._sim_step_counter += 1
             # set actions into buffers
-            self._warp_graph_cache.call("ActionManager_apply_action", self.action_manager.apply_action)
+            self._warp_graph_cache.call_steps(
+                "ActionManager_apply_action", self.action_manager.stage_steps("apply_action")
+            )
             # set actions into simulator
             self.scene.write_data_to_sim()
             # simulate
@@ -590,14 +592,14 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]:
 
         # post-step: step interval event
         if "interval" in self.event_manager.available_modes:
-            self._warp_graph_cache.call(
-                "EventManager_apply_interval", self.event_manager.apply, mode="interval", dt=self.step_dt
+            self._warp_graph_cache.call_steps(
+                "EventManager_apply_interval", self.event_manager.stage_steps("apply_interval"), dt=self.step_dt
             )
 
         # -- compute observations
-        self.obs_buf = self._warp_graph_cache.call(
+        self.obs_buf = self._warp_graph_cache.call_steps(
             "ObservationManager_compute_update_history",
-            self.observation_manager.compute,
+            self.observation_manager.stage_steps("compute"),
             update_history=True,
             return_cloned_output=False,
             output=clone_obs_buffer,
@@ -678,10 +680,9 @@ def _reset_mask(self, env_mask: wp.array, env_ids: torch.Tensor | None = None) -
         if "reset" in self.event_manager.available_modes:
             env_step_count = self._sim_step_counter // self.cfg.decimation
             self._global_env_step_count_wp.fill_(env_step_count)
-            self._warp_graph_cache.call(
+            self._warp_graph_cache.call_steps(
                 "EventManager_apply_reset",
-                self.event_manager.apply,
-                mode="reset",
+                self.event_manager.stage_steps("apply_reset"),
                 env_mask_wp=env_mask,
                 global_env_step_count=self._global_env_step_count_wp,
             )
@@ -691,15 +692,19 @@ def _reset_mask(self, env_mask: wp.array, env_ids: torch.Tensor | None = None) -
         # note: This is order-sensitive! Certain things need be reset before others.
         self.extras["log"] = dict()
         # -- observation manager
-        info = self._warp_graph_cache.call(
-            "ObservationManager_reset", self.observation_manager.reset, env_mask=env_mask
+        info = self._warp_graph_cache.call_steps(
+            "ObservationManager_reset", self.observation_manager.stage_steps("reset"), env_mask=env_mask
         )
         self.extras["log"].update(info)
         # -- action manager
-        info = self._warp_graph_cache.call("ActionManager_reset", self.action_manager.reset, env_mask=env_mask)
+        info = self._warp_graph_cache.call_steps(
+            "ActionManager_reset", self.action_manager.stage_steps("reset"), env_mask=env_mask
+        )
         self.extras["log"].update(info)
         # -- event manager
-        info = self._warp_graph_cache.call("EventManager_reset", self.event_manager.reset, env_mask=env_mask)
+        info = self._warp_graph_cache.call_steps(
+            "EventManager_reset", self.event_manager.stage_steps("reset"), env_mask=env_mask
+        )
         self.extras["log"].update(info)
         # -- recorder manager
         info = self.recorder_manager.reset(env_ids)
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
index c720de01ca5f..9470c10dbe5f 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
@@ -211,12 +211,6 @@ def invalidate_wp_graphs(self) -> None:
         """
         self._warp_graph_cache.invalidate()
 
-    def step_warp_termination_compute(self) -> None:
-        """Captured stage: compute terminations (env-step frequency)."""
-        self.reset_buf = self.termination_manager.compute()
-        self.reset_terminated = self.termination_manager.terminated
-        self.reset_time_outs = self.termination_manager.time_outs
-
     @Timer(name="env_step", msg="Step took:", enable=DEBUG_TIMER_STEP, time_unit="us")
     def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         """Execute one time-step of the environment's dynamics and reset terminated environments.
@@ -247,9 +241,9 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
             action_device = action.to(device=self.device, dtype=torch.float32).contiguous()
             wp.copy(self._action_in_wp, wp.from_torch(action_device, dtype=wp.float32))
 
-        self._warp_graph_cache.call(
+        self._warp_graph_cache.call_steps(
             "ActionManager_process_action",
-            self.action_manager.process_action,
+            self.action_manager.stage_steps("process_action"),
             action=self._action_in_wp,
             timer=DEBUG_TIMER_STEP,
         )
@@ -264,8 +258,8 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         for _ in range(self.cfg.decimation):
             self._sim_step_counter += 1
             # set actions into buffers
-            self._warp_graph_cache.call(
-                "ActionManager_apply_action", self.action_manager.apply_action, timer=DEBUG_TIMER_STEP
+            self._warp_graph_cache.call_steps(
+                "ActionManager_apply_action", self.action_manager.stage_steps("apply_action"), timer=DEBUG_TIMER_STEP
             )
             # scene writes cross the actuator and asset host boundaries, so they stay eager
             with Timer(
@@ -297,18 +291,23 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         self.common_step_counter += 1  # total step (common for all envs)
 
         # -- post-processing (termination + reward) as independently configurable stages
-        self._warp_graph_cache.call(
-            "TerminationManager_compute", self.step_warp_termination_compute, timer=DEBUG_TIMER_STEP
+        self.reset_buf = self._warp_graph_cache.call_steps(
+            "TerminationManager_compute", self.termination_manager.stage_steps("compute"), timer=DEBUG_TIMER_STEP
         )
-        self.reward_buf = self._warp_graph_cache.call(
-            "RewardManager_compute", self.reward_manager.compute, dt=float(self.step_dt), timer=DEBUG_TIMER_STEP
+        self.reset_terminated = self.termination_manager.terminated
+        self.reset_time_outs = self.termination_manager.time_outs
+        self.reward_buf = self._warp_graph_cache.call_steps(
+            "RewardManager_compute",
+            self.reward_manager.stage_steps("compute"),
+            dt=float(self.step_dt),
+            timer=DEBUG_TIMER_STEP,
         )
 
         if len(self.recorder_manager.active_terms) > 0:
             # update observations for recording if needed
-            self._warp_graph_cache.call(
+            self._warp_graph_cache.call_steps(
                 "ObservationManager_compute_no_history",
-                self.observation_manager.compute,
+                self.observation_manager.stage_steps("compute"),
                 return_cloned_output=False,
                 timer=DEBUG_TIMER_STEP,
             )
@@ -322,19 +321,18 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
 
         # -- step interval events
         if "interval" in self.event_manager.available_modes:
-            self._warp_graph_cache.call(
+            self._warp_graph_cache.call_steps(
                 "EventManager_apply_interval",
-                self.event_manager.apply,
-                mode="interval",
+                self.event_manager.stage_steps("apply_interval"),
                 dt=float(self.step_dt),
                 timer=DEBUG_TIMER_STEP,
             )
 
         # -- compute observations
         # note: done after reset to get the correct observations for reset envs
-        self.obs_buf = self._warp_graph_cache.call(
+        self.obs_buf = self._warp_graph_cache.call_steps(
             "ObservationManager_compute_update_history",
-            self.observation_manager.compute,
+            self.observation_manager.stage_steps("compute"),
             update_history=True,
             return_cloned_output=False,
             output=clone_obs_buffer,
@@ -489,10 +487,9 @@ def _reset_mask(self, env_mask: wp.array, env_ids: torch.Tensor | None = None) -
 
         if "reset" in self.event_manager.available_modes:
             self._global_env_step_count_wp.fill_(self._sim_step_counter // self.cfg.decimation)
-            self._warp_graph_cache.call(
+            self._warp_graph_cache.call_steps(
                 "EventManager_apply_reset",
-                self.event_manager.apply,
-                mode="reset",
+                self.event_manager.stage_steps("apply_reset"),
                 env_mask_wp=env_mask,
                 global_env_step_count=self._global_env_step_count_wp,
                 timer=DEBUG_TIMER_RESET,
@@ -502,14 +499,17 @@ def _reset_mask(self, env_mask: wp.array, env_ids: torch.Tensor | None = None) -
         # this returns a dictionary of information which is stored in the extras
         # note: This is order-sensitive! Certain things need be reset before others.
         # -- observation manager + action + reward managers
-        obs_info = self._warp_graph_cache.call(
-            "ObservationManager_reset", self.observation_manager.reset, env_mask=env_mask, timer=DEBUG_TIMER_RESET
+        obs_info = self._warp_graph_cache.call_steps(
+            "ObservationManager_reset",
+            self.observation_manager.stage_steps("reset"),
+            env_mask=env_mask,
+            timer=DEBUG_TIMER_RESET,
         )
-        action_info = self._warp_graph_cache.call(
-            "ActionManager_reset", self.action_manager.reset, env_mask=env_mask, timer=DEBUG_TIMER_RESET
+        action_info = self._warp_graph_cache.call_steps(
+            "ActionManager_reset", self.action_manager.stage_steps("reset"), env_mask=env_mask, timer=DEBUG_TIMER_RESET
         )
-        reward_info = self._warp_graph_cache.call(
-            "RewardManager_reset", self.reward_manager.reset, env_mask=env_mask, timer=DEBUG_TIMER_RESET
+        reward_info = self._warp_graph_cache.call_steps(
+            "RewardManager_reset", self.reward_manager.stage_steps("reset"), env_mask=env_mask, timer=DEBUG_TIMER_RESET
         )
 
         # -- curriculum manager
@@ -523,11 +523,14 @@ def _reset_mask(self, env_mask: wp.array, env_ids: torch.Tensor | None = None) -
 
         # -- command + event + termination managers
         command_info = self.command_manager.reset(env_ids=env_ids)
-        event_info = self._warp_graph_cache.call(
-            "EventManager_reset", self.event_manager.reset, env_mask=env_mask, timer=DEBUG_TIMER_RESET
+        event_info = self._warp_graph_cache.call_steps(
+            "EventManager_reset", self.event_manager.stage_steps("reset"), env_mask=env_mask, timer=DEBUG_TIMER_RESET
         )
-        termination_info = self._warp_graph_cache.call(
-            "TerminationManager_reset", self.termination_manager.reset, env_mask=env_mask, timer=DEBUG_TIMER_RESET
+        termination_info = self._warp_graph_cache.call_steps(
+            "TerminationManager_reset",
+            self.termination_manager.stage_steps("reset"),
+            env_mask=env_mask,
+            timer=DEBUG_TIMER_RESET,
         )
 
         # -- recorder manager
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py
index fa0ce31f4219..32ef81c03bc0 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py
@@ -10,7 +10,8 @@
 import inspect
 import re
 from abc import abstractmethod
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
+from functools import partial
 from typing import TYPE_CHECKING, Any
 
 import torch
@@ -22,6 +23,8 @@
 from isaaclab.managers.manager_term_cfg import ActionTermCfg
 from isaaclab.utils import instantiate
 
+from isaaclab_experimental.utils.warp import is_warp_capturable
+
 from .manager_base import ManagerBase, ManagerTermBase
 
 if TYPE_CHECKING:
@@ -391,8 +394,10 @@ def reset(
                     "Do not pass env_ids on captured paths."
                 )
             env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask)
+        return self._run_steps("reset", env_mask=env_mask)
 
-        # reset the action history
+    def _reset_action_history(self, env_mask: wp.array) -> dict[str, Any]:
+        """Zero the current and previous actions of the selected environments."""
         if env_mask is None:
             self._prev_action.fill_(0.0)
             self._action.fill_(0.0)
@@ -409,10 +414,12 @@ def reset(
                 inputs=[env_mask, self._action],
                 device=self.device,
             )
+        # nothing to log here
+        return {}
 
-        # reset all action terms
-        for term in self._terms.values():
-            term.reset(env_mask=env_mask)
+    @staticmethod
+    def _reset_term(term: ActionTerm, env_mask: wp.array) -> dict[str, Any]:
+        term.reset(env_mask=env_mask)
         # nothing to log here
         return {}
 
@@ -425,19 +432,7 @@ def process_action(self, action: wp.array):
         Args:
             action: The actions to process. Shape is (num_envs, total_action_dim).
         """
-        # check if action dimension is valid
-        if self.total_action_dim != action.shape[1]:
-            raise ValueError(f"Invalid action shape, expected: {self.total_action_dim}, received: {action.shape[1]}.")
-
-        # store the input actions
-        wp.copy(self._prev_action, self._action)
-        wp.copy(self._action, action)
-
-        # split the actions and apply to each term
-        idx = 0
-        for term in self._terms.values():
-            term.process_actions(self._action, idx)
-            idx += term.action_dim
+        self._run_steps("process_action", action)
 
     def apply_action(self) -> None:
         """Applies the actions to the environment/simulation.
@@ -445,8 +440,38 @@ def apply_action(self) -> None:
         Note:
             This should be called at every simulation step.
         """
-        for term in self._terms.values():
-            term.apply_actions()
+        self._run_steps("apply_action")
+
+    """
+    Operations - Stage steps.
+    """
+
+    def _build_stage_steps(self, operation: str) -> list[tuple[bool, Callable[..., Any]]]:
+        terms = list(self._terms.values())
+        if operation == "process_action":
+            offsets = [sum(term.action_dim for term in terms[:index]) for index in range(len(terms))]
+            steps = [
+                (is_warp_capturable(term), partial(self._process_term_actions, term, offset))
+                for term, offset in zip(terms, offsets)
+            ]
+            return [(True, self._store_actions), *steps]
+        if operation == "apply_action":
+            return [(is_warp_capturable(term), term.apply_actions) for term in terms]
+        if operation == "reset":
+            steps = [(is_warp_capturable(term), partial(self._reset_term, term)) for term in terms]
+            return [(True, self._reset_action_history), *steps]
+        return super()._build_stage_steps(operation)
+
+    def _store_actions(self, action: wp.array) -> None:
+        """Check the action shape and store the actions, keeping the previous ones."""
+        if self.total_action_dim != action.shape[1]:
+            raise ValueError(f"Invalid action shape, expected: {self.total_action_dim}, received: {action.shape[1]}.")
+        wp.copy(self._prev_action, self._action)
+        wp.copy(self._action, action)
+
+    def _process_term_actions(self, term: ActionTerm, offset: int, action: wp.array) -> None:
+        """Hand a term its slice of the stored actions."""
+        term.process_actions(self._action, offset)
 
     def get_term(self, name: str) -> ActionTerm:
         """Returns the action term with the specified name.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
index 2b5329d92eca..fbeaaeb63524 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
@@ -23,7 +23,9 @@
 
 import inspect
 import logging
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
+from functools import partial
+from typing import Any
 
 import torch
 import warp as wp
@@ -31,6 +33,8 @@
 
 from isaaclab.managers.manager_term_cfg import EventTermCfg
 
+from isaaclab_experimental.utils.warp import is_warp_capturable
+
 from .manager_base import ManagerBase, ManagerTermBase
 
 logger = logging.getLogger(__name__)
@@ -255,9 +259,7 @@ def set_term_cfg(self, term_name: str, cfg: EventTermCfg):
                 for term_cfg in self._mode_term_cfgs[mode]
                 if inspect.isclass(term_cfg.func) or isinstance(term_cfg.func, ManagerTermBase)
             ]
-            graph_cache = getattr(self._env, "_warp_graph_cache", None)
-            if graph_cache is not None:
-                graph_cache.invalidate(type(self).__name__)
+            self._clear_stage_steps()
             return
         raise ValueError(f"Event term '{term_name}' not found.")
 
@@ -283,32 +285,7 @@ def reset(
                 )
             env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask)
 
-        # reset class terms (mask-based)
-        for mode_cfg in self._mode_class_term_cfgs.values():
-            for term_cfg in mode_cfg:
-                term_cfg.func.reset(env_mask=env_mask)
-
-        # reset interval timers for non-global interval events
-        # note: terms with resample_interval_on_reset=False keep their per-environment counter
-        #   across resets, so we do not resample them either
-        if "interval" in self._mode_term_cfgs:
-            for i, term_cfg in enumerate(self._mode_term_cfgs["interval"]):
-                if term_cfg.is_global_time or not term_cfg.resample_interval_on_reset:
-                    continue
-                lower, upper = self._interval_term_ranges[i]
-                wp.launch(
-                    kernel=_interval_reset_selected,
-                    dim=self.num_envs,
-                    inputs=[
-                        env_mask,
-                        self._interval_term_time_left_wp[i],
-                        self._env.rng_state_wp,
-                        float(lower),
-                        float(upper),
-                    ],
-                    device=self.device,
-                )
-        return {}
+        return self._run_steps("reset", env_mask=env_mask)
 
     def apply(
         self,
@@ -362,72 +339,124 @@ def apply(
     def _apply_interval(self, dt: float) -> None:
         if self._env.rng_state_wp is None:
             raise RuntimeError("EventManager._apply_interval: env.rng_state_wp is not initialized.")
-
-        # iterate over all the interval terms (fixed list; captured graph-friendly)
-        for i, term_cfg in enumerate(self._mode_term_cfgs["interval"]):
-            lower, upper = self._interval_term_ranges[i]
-            if self._interval_term_is_global[i]:
-                if self._interval_global_rng_state_wp is None:
-                    raise RuntimeError(
-                        "EventManager._apply_interval: _interval_global_rng_state_wp is not initialized."
-                    )
-                # update scalar time_left and scalar flag (mask is a broadcast view of the flag)
-                wp.launch(
-                    kernel=_interval_step_global,
-                    dim=1,
-                    inputs=[
-                        self._interval_term_time_left_wp[i],
-                        self._interval_global_rng_state_wp,
-                        self._scratch_interval_trigger_flag_wp,
-                        float(dt),
-                        float(lower),
-                        float(upper),
-                    ],
-                    device=self.device,
-                )
-                term_cfg.func(self._env, self._scratch_interval_trigger_mask_view_wp, **term_cfg.params)
-            else:
-                wp.launch(
-                    kernel=_interval_step_per_env,
-                    dim=self.num_envs,
-                    inputs=[
-                        self._interval_term_time_left_wp[i],
-                        self._env.rng_state_wp,
-                        self._scratch_term_mask_wp,
-                        float(dt),
-                        float(lower),
-                        float(upper),
-                    ],
-                    device=self.device,
-                )
-                term_cfg.func(self._env, self._scratch_term_mask_wp, **term_cfg.params)
+        self._run_steps("apply_interval", dt=dt)
 
     def _apply_reset(self, env_mask_wp: wp.array, global_env_step_count_wp: wp.array) -> None:
         if self._scratch_term_mask_wp is None:
             raise RuntimeError("EventManager._apply_reset: _scratch_term_mask_wp is not initialized.")
-
-        # iterate over all the reset terms
-        for index, term_cfg in enumerate(self._mode_term_cfgs["reset"]):
-            min_step_count = int(term_cfg.min_step_count_between_reset)
+        self._run_steps("apply_reset", env_mask_wp=env_mask_wp, global_env_step_count=global_env_step_count_wp)
+
+    """
+    Operations - Stage steps.
+    """
+
+    def _build_stage_steps(self, operation: str) -> list[tuple[bool, Callable[..., Any]]]:
+        # Terms draw from the environment's shared random state, so every stage keeps the configured order.
+        if operation == "apply_interval":
+            return [
+                (is_warp_capturable(term_cfg.func), partial(self._apply_interval_term, index))
+                for index, term_cfg in enumerate(self._mode_term_cfgs.get("interval", []))
+            ]
+        if operation == "apply_reset":
+            return [
+                (is_warp_capturable(term_cfg.func), partial(self._apply_reset_term, index))
+                for index, term_cfg in enumerate(self._mode_term_cfgs.get("reset", []))
+            ]
+        if operation == "reset":
+            steps = [
+                (is_warp_capturable(term_cfg.func), partial(self._reset_class_term, term_cfg))
+                for mode_cfgs in self._mode_class_term_cfgs.values()
+                for term_cfg in mode_cfgs
+            ]
+            # note: terms with resample_interval_on_reset=False keep their per-environment counter
+            #   across resets, so we do not resample them either
+            steps += [
+                (True, partial(self._resample_interval, index))
+                for index, term_cfg in enumerate(self._mode_term_cfgs.get("interval", []))
+                if not term_cfg.is_global_time and term_cfg.resample_interval_on_reset
+            ]
+            return [*steps, (True, self._reset_output)]
+        return super()._build_stage_steps(operation)
+
+    def _apply_interval_term(self, index: int, dt: float) -> None:
+        """Advance one interval term's timers and apply it to the environments whose timer expired."""
+        term_cfg = self._mode_term_cfgs["interval"][index]
+        lower, upper = self._interval_term_ranges[index]
+        if self._interval_term_is_global[index]:
+            if self._interval_global_rng_state_wp is None:
+                raise RuntimeError("EventManager._apply_interval: _interval_global_rng_state_wp is not initialized.")
+            # update scalar time_left and scalar flag (mask is a broadcast view of the flag)
             wp.launch(
-                kernel=_reset_compute_valid_mask,
+                kernel=_interval_step_global,
+                dim=1,
+                inputs=[
+                    self._interval_term_time_left_wp[index],
+                    self._interval_global_rng_state_wp,
+                    self._scratch_interval_trigger_flag_wp,
+                    float(dt),
+                    float(lower),
+                    float(upper),
+                ],
+                device=self.device,
+            )
+            term_cfg.func(self._env, self._scratch_interval_trigger_mask_view_wp, **term_cfg.params)
+        else:
+            wp.launch(
+                kernel=_interval_step_per_env,
                 dim=self.num_envs,
                 inputs=[
-                    env_mask_wp,
-                    self._reset_term_last_triggered_step_wp[index],
-                    self._reset_term_triggered_once_wp[index],
+                    self._interval_term_time_left_wp[index],
+                    self._env.rng_state_wp,
                     self._scratch_term_mask_wp,
-                    global_env_step_count_wp,
-                    int(min_step_count),
+                    float(dt),
+                    float(lower),
+                    float(upper),
                 ],
                 device=self.device,
             )
             term_cfg.func(self._env, self._scratch_term_mask_wp, **term_cfg.params)
 
-    def _register_term_capturability(self, term_cfg: EventTermCfg) -> None:
-        # startup terms run once while the environment is built, outside the recorded stages
-        if term_cfg.mode not in ("prestartup", "startup"):
-            super()._register_term_capturability(term_cfg)
+    def _apply_reset_term(self, index: int, env_mask_wp: wp.array, global_env_step_count: wp.array) -> None:
+        """Apply one reset term to the selected environments its trigger spacing allows."""
+        term_cfg = self._mode_term_cfgs["reset"][index]
+        wp.launch(
+            kernel=_reset_compute_valid_mask,
+            dim=self.num_envs,
+            inputs=[
+                env_mask_wp,
+                self._reset_term_last_triggered_step_wp[index],
+                self._reset_term_triggered_once_wp[index],
+                self._scratch_term_mask_wp,
+                global_env_step_count,
+                int(term_cfg.min_step_count_between_reset),
+            ],
+            device=self.device,
+        )
+        term_cfg.func(self._env, self._scratch_term_mask_wp, **term_cfg.params)
+
+    @staticmethod
+    def _reset_class_term(term_cfg: EventTermCfg, env_mask: wp.array) -> None:
+        term_cfg.func.reset(env_mask=env_mask)
+
+    def _resample_interval(self, index: int, env_mask: wp.array) -> None:
+        """Draw new interval timers for the selected environments."""
+        lower, upper = self._interval_term_ranges[index]
+        wp.launch(
+            kernel=_interval_reset_selected,
+            dim=self.num_envs,
+            inputs=[
+                env_mask,
+                self._interval_term_time_left_wp[index],
+                self._env.rng_state_wp,
+                float(lower),
+                float(upper),
+            ],
+            device=self.device,
+        )
+
+    @staticmethod
+    def _reset_output(env_mask: wp.array) -> dict[str, float]:
+        return {}
 
     def _prepare_terms(self):
         # check if config is dict already
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py
index f30d62040506..9ad3064aec94 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py
@@ -20,7 +20,7 @@
 import inspect
 import logging
 from abc import ABC, abstractmethod
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from typing import TYPE_CHECKING, Any
 
 import warp as wp
@@ -195,6 +195,9 @@ def __init__(self, cfg: object, env: ManagerBasedEnv):
         # else:
         #     self._resolve_terms_handle = None
         self._resolve_terms_handle = None
+        # steps of each operation for the environment's graph cache, built on first use
+        self._stage_steps: dict[str, tuple[tuple[bool, Callable[..., Any]], ...]] = {}
+        self._all_terms_capturable = True
 
         # parse config to create terms information
         if self.cfg:
@@ -283,10 +286,57 @@ def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequenc
         """
         raise NotImplementedError
 
+    def stage_steps(self, operation: str) -> tuple[tuple[bool, Callable[..., Any]], ...]:
+        """Ordered steps of an operation, each marked as capturable or not.
+
+        Calling the steps in order with the operation's arguments performs the operation. The Warp
+        environments record each run of consecutive capturable steps into a CUDA graph and run the other
+        steps eagerly in between (see :meth:`~isaaclab_experimental.utils.WarpGraphCache.call_steps`). A
+        term decorated with ``@WarpCapturable(False)`` makes its own step eager.
+
+        Args:
+            operation: Name of the operation, such as ``"compute"`` or ``"reset"``.
+
+        Returns:
+            The steps, reused until a term changes.
+        """
+        steps = self._stage_steps.get(operation)
+        if steps is None:
+            steps = self._stage_steps[operation] = tuple(self._build_stage_steps(operation))
+        return steps
+
     """
     Implementation specific.
     """
 
+    def _build_stage_steps(self, operation: str) -> list[tuple[bool, Callable[..., Any]]]:
+        """Build the steps of an operation.
+
+        The default runs the whole method as one step, capturable when every term is. Managers that
+        split an operation per term override this.
+
+        Args:
+            operation: Name of the operation.
+
+        Returns:
+            The ordered steps.
+        """
+        return [(self._all_terms_capturable, getattr(self, operation))]
+
+    def _run_steps(self, operation: str, *args: Any, **kwargs: Any) -> Any:
+        """Run the steps of an operation eagerly and return the last step's result."""
+        result = None
+        for _, step in self.stage_steps(operation):
+            result = step(*args, **kwargs)
+        return result
+
+    def _clear_stage_steps(self) -> None:
+        """Rebuild the steps on next use and drop the graphs recorded from them."""
+        self._stage_steps.clear()
+        graph_cache = getattr(self._env, "_warp_graph_cache", None)
+        if graph_cache is not None:
+            graph_cache.invalidate(type(self).__name__)
+
     @abstractmethod
     def _prepare_terms(self):
         """Prepare terms information from the configuration object."""
@@ -321,6 +371,8 @@ def _resolve_terms_callback(self, event):
 
         # set the flag
         self._is_scene_entities_resolved = True
+        # class terms were just created
+        self._clear_stage_steps()
 
     """
     Internal functions.
@@ -411,24 +463,14 @@ def _resolve_common_term_cfg(self, term_name: str, term_cfg: ManagerTermBaseCfg,
                     f" and optional parameters: {args_with_defaults}, but received: {term_params}."
                 )
 
-        self._register_term_capturability(term_cfg)
+        if not is_warp_capturable(term_cfg.func):
+            self._all_terms_capturable = False
 
         # process attributes at runtime
         # these properties are only resolvable once the simulation starts playing
         if self._env.sim.is_playing():
             self._process_term_cfg_at_play(term_name, term_cfg)
 
-    def _register_term_capturability(self, term_cfg: ManagerTermBaseCfg) -> None:
-        """Keep every stage of this manager eager when the term is not capturable.
-
-        Args:
-            term_cfg: The term configuration.
-        """
-        if not is_warp_capturable(term_cfg.func):
-            graph_cache = getattr(self._env, "_warp_graph_cache", None)
-            if graph_cache is not None:
-                graph_cache.register_capturability(type(self).__name__, False)
-
     def _process_term_cfg_at_play(self, term_name: str, term_cfg: ManagerTermBaseCfg):
         """Process the term configuration at runtime.
 
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
index 051aaa8a6ff6..0fafea0f3829 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
@@ -47,7 +47,8 @@
 from __future__ import annotations
 
 import inspect
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
+from functools import partial
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
@@ -62,6 +63,7 @@
 from isaaclab_experimental.utils import modifiers, noise
 from isaaclab_experimental.utils.buffers import CircularBuffer
 from isaaclab_experimental.utils.torch_utils import clone_obs_buffer
+from isaaclab_experimental.utils.warp import is_warp_capturable
 
 from .manager_base import ManagerBase, ManagerTermBase
 
@@ -400,21 +402,7 @@ def reset(
                     "Do not pass env_ids on captured paths."
                 )
             env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask)
-
-        # call all terms that are classes
-        for group_name, group_cfg in self._group_obs_class_term_cfgs.items():
-            for term_cfg in group_cfg:
-                term_cfg.func.reset(env_mask=env_mask)
-            # reset terms with history
-            for term_name in self._group_obs_term_names[group_name]:
-                if term_name in self._group_obs_term_history_buffer[group_name]:
-                    self._group_obs_term_history_buffer[group_name][term_name].reset(env_mask=env_mask)
-        # call all modifiers/noise models that are classes
-        for mod in self._group_obs_class_instances:
-            mod.reset(env_mask=env_mask)
-
-        # nothing to log here
-        return {}
+        return self._run_steps("reset", env_mask=env_mask)
 
     def compute(
         self, update_history: bool = False, return_cloned_output: bool = True
@@ -436,9 +424,12 @@ def compute(
             The observations are either concatenated into a single tensor or returned as a dictionary
             with keys corresponding to the term's name.
         """
-        # Launch kernels for every group (writes into persistent buffers in-place).
-        for group_name in self._group_obs_term_names:
-            self.compute_group(group_name, update_history=update_history)
+        return self._run_steps("compute", update_history=update_history, return_cloned_output=return_cloned_output)
+
+    def _compute_output(
+        self, update_history: bool = False, return_cloned_output: bool = True
+    ) -> dict[str, torch.Tensor | dict[str, torch.Tensor]]:
+        """Return the persistent observation buffer, or a clone of it."""
         # Build the obs buffer once (persistent refs to in-place-updated tensors/dicts).
         if self._obs_buffer is None:
             self._obs_buffer = {
@@ -497,74 +488,122 @@ def compute_group(self, group_name: str, update_history: bool = False) -> torch.
         # iterate over all the terms in each group
         group_term_names = self._group_obs_term_names[group_name]
 
-        # Persistent per-term obs dict (pre-allocated in _prepare_terms).
+        # evaluate terms: compute, add noise, clip, scale, custom modifiers
+        for index in range(len(group_term_names)):
+            self._compute_term(group_name, index, update_history=update_history)
+
+        # return persistent output (updated in-place by kernels above)
+        if self._group_use_warp_concat[group_name]:
+            return self._group_out_torch[group_name]
+        return self._group_obs_dict[group_name]
+
+    def _compute_term(
+        self, group_name: str, index: int, update_history: bool = False, return_cloned_output: bool = True
+    ) -> None:
+        """Compute one observation term into its slice of the group buffer."""
+        term_name = self._group_obs_term_names[group_name][index]
+        term_cfg = self._group_obs_term_cfgs[group_name][index]
         group_obs = self._group_obs_dict[group_name]
+        # compute term's value into pre-allocated Warp output
+        term_cfg.func(self._env, term_cfg.out_wp, **term_cfg.params)
+
+        # apply custom modifiers (in-place on out_wp)
+        if term_cfg.modifiers is not None:
+            for modifier in term_cfg.modifiers:
+                modifier.func(term_cfg.out_wp, **modifier.params)
+
+        # apply noise (Warp in-place on out_wp)
+        if isinstance(term_cfg.noise, noise.NoiseCfg):
+            term_cfg.noise.func(term_cfg.out_wp, term_cfg.noise)
+        elif isinstance(term_cfg.noise, noise.NoiseModelCfg) and term_cfg.noise.func is not None:
+            term_cfg.noise.func(term_cfg.out_wp)
+
+        # clip then scale (stable semantics); implementation may use Warp kernels
+        if term_cfg.clip is not None:
+            wp.launch(
+                kernel=_apply_clip,
+                dim=self.num_envs,
+                inputs=[term_cfg.out_wp, float(term_cfg.clip[0]), float(term_cfg.clip[1])],
+                device=self.device,
+            )
+        if term_cfg.scale is not None:
+            wp.launch(
+                kernel=_apply_scale,
+                dim=self.num_envs,
+                inputs=[term_cfg.out_wp, term_cfg.scale_wp],
+                device=self.device,
+            )
 
-        # evaluate terms: compute, add noise, clip, scale, custom modifiers
-        for term_name, term_cfg in zip(group_term_names, self._group_obs_term_cfgs[group_name]):
-            # compute term's value into pre-allocated Warp output
-            term_cfg.func(self._env, term_cfg.out_wp, **term_cfg.params)
-
-            # apply custom modifiers (in-place on out_wp)
-            if term_cfg.modifiers is not None:
-                for modifier in term_cfg.modifiers:
-                    modifier.func(term_cfg.out_wp, **modifier.params)
-
-            # apply noise (Warp in-place on out_wp)
-            if isinstance(term_cfg.noise, noise.NoiseCfg):
-                term_cfg.noise.func(term_cfg.out_wp, term_cfg.noise)
-            elif isinstance(term_cfg.noise, noise.NoiseModelCfg) and term_cfg.noise.func is not None:
-                term_cfg.noise.func(term_cfg.out_wp)
-
-            # clip then scale (stable semantics); implementation may use Warp kernels
-            if term_cfg.clip is not None:
-                wp.launch(
-                    kernel=_apply_clip,
-                    dim=self.num_envs,
-                    inputs=[term_cfg.out_wp, float(term_cfg.clip[0]), float(term_cfg.clip[1])],
-                    device=self.device,
+        # TODO(jichuanh): This is not migrated yet. Need revisit.
+        # Update the history buffer if observation term has history enabled
+        if term_cfg.history_length > 0:
+            # circular buffer is not capture safe
+            if wp.get_device().is_capturing:
+                raise RuntimeError(
+                    "Observation terms with history (circular buffer) are not CUDA-graph-capture-safe yet. "
+                    "Disable history for observation terms used inside a captured graph, or restructure "
+                    "the graph to exclude history-buffered terms."
                 )
-            if term_cfg.scale is not None:
-                wp.launch(
-                    kernel=_apply_scale,
-                    dim=self.num_envs,
-                    inputs=[term_cfg.out_wp, term_cfg.scale_wp],
-                    device=self.device,
+            circular_buffer = self._group_obs_term_history_buffer[group_name][term_name]
+            if update_history:
+                circular_buffer.append(wp.to_torch(term_cfg.out_wp))
+            elif circular_buffer._buffer is None:
+                # because circular buffer only exits after the simulation steps,
+                # this guards history buffer from corruption by external calls before simulation start
+                circular_buffer = CircularBuffer(
+                    max_len=circular_buffer.max_length,
+                    batch_size=circular_buffer.batch_size,
+                    device=circular_buffer.device,
                 )
+                self._group_obs_term_history_buffer[group_name][term_name] = circular_buffer
+                circular_buffer.append(wp.to_torch(term_cfg.out_wp))
 
-            # TODO(jichuanh): This is not migrated yet. Need revisit.
-            # Update the history buffer if observation term has history enabled
-            if term_cfg.history_length > 0:
-                # circular buffer is not capture safe
-                if wp.get_device().is_capturing:
-                    raise RuntimeError(
-                        "Observation terms with history (circular buffer) are not CUDA-graph-capture-safe yet. "
-                        "Disable history for observation terms used inside a captured graph, or restructure "
-                        "the graph to exclude history-buffered terms."
-                    )
-                circular_buffer = self._group_obs_term_history_buffer[group_name][term_name]
-                if update_history:
-                    circular_buffer.append(wp.to_torch(term_cfg.out_wp))
-                elif circular_buffer._buffer is None:
-                    # because circular buffer only exits after the simulation steps,
-                    # this guards history buffer from corruption by external calls before simulation start
-                    circular_buffer = CircularBuffer(
-                        max_len=circular_buffer.max_length,
-                        batch_size=circular_buffer.batch_size,
-                        device=circular_buffer.device,
-                    )
-                    self._group_obs_term_history_buffer[group_name][term_name] = circular_buffer
-                    circular_buffer.append(wp.to_torch(term_cfg.out_wp))
+            if term_cfg.flatten_history_dim:
+                group_obs[term_name] = circular_buffer.buffer.reshape(self._env.num_envs, -1)
+            else:
+                group_obs[term_name] = circular_buffer.buffer
 
-                if term_cfg.flatten_history_dim:
-                    group_obs[term_name] = circular_buffer.buffer.reshape(self._env.num_envs, -1)
-                else:
-                    group_obs[term_name] = circular_buffer.buffer
+    """
+    Operations - Stage steps.
+    """
 
-        # return persistent output (updated in-place by kernels above)
-        if self._group_use_warp_concat[group_name]:
-            return self._group_out_torch[group_name]
-        return group_obs
+    def _build_stage_steps(self, operation: str) -> list[tuple[bool, Callable[..., Any]]]:
+        if operation == "compute":
+            # Noise draws from the environment's shared random state, so terms keep their order.
+            steps = [
+                (
+                    is_warp_capturable(term_cfg.func) and term_cfg.history_length == 0,
+                    partial(self._compute_term, group_name, index),
+                )
+                for group_name, term_cfgs in self._group_obs_term_cfgs.items()
+                for index, term_cfg in enumerate(term_cfgs)
+            ]
+            return [*steps, (True, self._compute_output)]
+        if operation == "reset":
+            steps = []
+            for group_name, term_cfgs in self._group_obs_class_term_cfgs.items():
+                steps += [(is_warp_capturable(cfg.func), partial(self._reset_class_term, cfg)) for cfg in term_cfgs]
+                # history buffers are not capture-safe
+                steps += [
+                    (False, partial(self._reset_history, group_name, term_name))
+                    for term_name in self._group_obs_term_history_buffer[group_name]
+                ]
+            steps += [(True, instance.reset) for instance in self._group_obs_class_instances]
+            return [*steps, (True, self._reset_output)]
+        return super()._build_stage_steps(operation)
+
+    @staticmethod
+    def _reset_class_term(term_cfg: ObservationTermCfg, env_mask: wp.array) -> None:
+        term_cfg.func.reset(env_mask=env_mask)
+
+    def _reset_history(self, group_name: str, term_name: str, env_mask: wp.array) -> None:
+        # computing a term before the first step may replace its buffer, so look it up on every call
+        self._group_obs_term_history_buffer[group_name][term_name].reset(env_mask=env_mask)
+
+    @staticmethod
+    def _reset_output(env_mask: wp.array) -> dict:
+        # nothing to log here
+        return {}
 
     def serialize(self) -> dict:
         """Serialize the observation term configurations for all active groups.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py
index 7dd6b85bd6ed..6b2cb0f27978 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py
@@ -11,8 +11,9 @@
 
 from __future__ import annotations
 
-from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from collections.abc import Callable, Sequence
+from functools import partial
+from typing import TYPE_CHECKING, Any
 
 import torch
 import warp as wp
@@ -20,6 +21,7 @@
 
 from isaaclab.managers.manager_term_cfg import RewardTermCfg
 
+from isaaclab_experimental.utils.warp import is_warp_capturable
 from isaaclab_experimental.utils.warp.kernels import compute_reset_scale, count_masked
 
 from .manager_base import ManagerBase, ManagerTermBase
@@ -251,7 +253,46 @@ def reset(
                     "Do not pass env_ids on captured paths."
                 )
             env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask)
+        return self._run_steps("reset", env_mask=env_mask)
 
+    def compute(self, dt: float) -> torch.Tensor:
+        """Computes the reward signal as a weighted sum of individual terms.
+
+        This function calls each reward term managed by the class and adds them to compute the net
+        reward signal. It also updates the episodic sums corresponding to individual reward terms.
+
+        Args:
+            dt: The time-step interval of the environment.
+
+        Returns:
+            The net reward signal of shape (num_envs,).
+        """
+        return self._run_steps("compute", dt=dt)
+
+    """
+    Operations - Stage steps.
+    """
+
+    def _build_stage_steps(self, operation: str) -> list[tuple[bool, Callable[..., Any]]]:
+        if operation == "compute":
+            term_steps = [
+                (is_warp_capturable(term_cfg.func), partial(self._compute_term, index))
+                for index, term_cfg in enumerate(self._term_cfgs)
+            ]
+            # Terms write their own output columns and the finalize step sums them in a fixed order, so
+            # running the capturable terms first leaves the reward unchanged bitwise.
+            term_steps.sort(key=lambda step: not step[0])
+            return [(True, self._reset_step_buffers), *term_steps, (True, self._finalize_step)]
+        if operation == "reset":
+            term_steps = [
+                (is_warp_capturable(term_cfg.func), partial(self._reset_term, term_cfg))
+                for term_cfg in self._class_term_cfgs
+            ]
+            return [(True, self._reset_episode_sums), *term_steps]
+        return super()._build_stage_steps(operation)
+
+    def _reset_episode_sums(self, env_mask: wp.array) -> dict[str, torch.Tensor]:
+        """Average the episodic sums of the selected environments and zero them."""
         self._episode_sum_avg_wp.zero_()
         self._reset_count_wp.zero_()
         self._reset_scale_wp.zero_()
@@ -271,48 +312,38 @@ def reset(
                 inputs=[env_mask, self._reset_scale_wp, self._episode_sums_wp, self._episode_sum_avg_wp],
                 device=self.device,
             )
-
-        # reset all the reward terms; class terms may expose logging values as
-        # persistent tensor views (see :meth:`ManagerTermBase.reset`), which are
-        # merged into the manager's reset extras. Merging the same views again on
-        # subsequent resets is a no-op, so this stays capture-safe.
-        for term_cfg in self._class_term_cfgs:
-            term_extras = term_cfg.func.reset(env_mask=env_mask)
-            if term_extras:
-                self._reset_extras.update(term_extras)
-
         return self._reset_extras
 
-    def compute(self, dt: float) -> torch.Tensor:
-        """Computes the reward signal as a weighted sum of individual terms.
-
-        This function calls each reward term managed by the class and adds them to compute the net
-        reward signal. It also updates the episodic sums corresponding to individual reward terms.
-
-        Args:
-            dt: The time-step interval of the environment.
+    def _reset_term(self, term_cfg: RewardTermCfg, env_mask: wp.array) -> dict[str, torch.Tensor]:
+        """Reset a class term and merge its logging values into the reset extras.
 
-        Returns:
-            The net reward signal of shape (num_envs,).
+        Class terms may expose logging values as persistent tensor views (see :meth:`ManagerTermBase.reset`).
+        Merging the same views again on later resets is a no-op, so this stays capture-safe.
         """
+        term_extras = term_cfg.func.reset(env_mask=env_mask)
+        if term_extras:
+            self._reset_extras.update(term_extras)
+        return self._reset_extras
+
+    def _reset_step_buffers(self, dt: float) -> None:
+        """Zero the reward, step reward and term output buffers in a single kernel launch."""
         # TODO: Investigate performance diff between two .fill_ and kernel launch
-        # reset computation (Warp buffers) in a single kernel launch
         wp.launch(
             kernel=_reward_pre_compute_reset,
             dim=self.num_envs,
             inputs=[self._reward_wp, self._step_reward_wp, self._term_outs_wp],
             device=self.device,
         )
-        # iterate over all the reward terms (Python loop; per-term math is warp)
-        for term_cfg in self._term_cfgs:
-            # skip if weight is zero (kind of a micro-optimization)
-            if term_cfg.weight == 0.0:
-                continue
-            # compute term into the persistent warp buffer (raw, unweighted)
-            # NOTE: `out` is pre-zeroed every step by `_reward_pre_compute_reset`.
+
+    def _compute_term(self, index: int, dt: float) -> None:
+        """Compute one term into its pre-zeroed output column (raw, unweighted)."""
+        term_cfg = self._term_cfgs[index]
+        # skip if weight is zero (kind of a micro-optimization)
+        if term_cfg.weight != 0.0:
             term_cfg.func(self._env, term_cfg.out, **term_cfg.params)
 
-        # update total reward, episodic sums and step rewards in a single kernel launch
+    def _finalize_step(self, dt: float) -> torch.Tensor:
+        """Update the total reward, episodic sums and step rewards in a single kernel launch."""
         wp.launch(
             kernel=_reward_finalize,
             dim=self.num_envs,
@@ -353,6 +384,8 @@ def set_term_cfg(self, term_name: str, cfg: RewardTermCfg):
         self._term_cfgs[term_idx] = cfg
         # keep on-device weights in sync (call this to update weights used in compute)
         self._term_weights_tensor_view[term_idx] = float(cfg.weight)
+        # the replaced term may differ in capturability
+        self._clear_stage_steps()
 
     def get_term_cfg(self, term_name: str) -> RewardTermCfg:
         """Gets the configuration for the specified term.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py
index 368d9ca5d130..5055a59b0f5e 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py
@@ -15,8 +15,9 @@
 
 from __future__ import annotations
 
-from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from collections.abc import Callable, Sequence
+from functools import partial
+from typing import TYPE_CHECKING, Any
 
 import torch
 import warp as wp
@@ -24,6 +25,8 @@
 
 from isaaclab.managers.manager_term_cfg import TerminationTermCfg
 
+from isaaclab_experimental.utils.warp import is_warp_capturable
+
 from .manager_base import ManagerBase, ManagerTermBase
 
 if TYPE_CHECKING:
@@ -268,17 +271,7 @@ def reset(
                     "Do not pass env_ids on captured paths."
                 )
             env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask)
-        if len(self._term_names) > 0:
-            self._term_done_avg_wp.zero_()
-            wp.launch(
-                kernel=_termination_reset_mean_all_2d,
-                dim=(self.num_envs, len(self._term_names)),
-                inputs=[self._last_episode_dones_wp, self._term_done_avg_wp],
-                device=self.device,
-            )
-        for term_cfg in self._class_term_cfgs:
-            term_cfg.func.reset(env_mask=env_mask)
-        return self._reset_extras
+        return self._run_steps("reset", env_mask=env_mask)
 
     @property
     def episode_termination_extras(self) -> dict[str, torch.Tensor]:
@@ -291,7 +284,48 @@ def compute(self) -> torch.Tensor:
         Returns:
             The combined termination signal of shape (num_envs,).
         """
-        # reset computation (Warp buffers) in a single kernel launch
+        return self._run_steps("compute")
+
+    """
+    Operations - Stage steps.
+    """
+
+    def _build_stage_steps(self, operation: str) -> list[tuple[bool, Callable[..., Any]]]:
+        if operation == "compute":
+            term_steps = [
+                (is_warp_capturable(term_cfg.func), partial(self._compute_term, term_cfg))
+                for term_cfg in self._term_cfgs
+            ]
+            # Terms write their own done columns and the finalize step combines them in a fixed order, so
+            # running the capturable terms first leaves the dones unchanged.
+            term_steps.sort(key=lambda step: not step[0])
+            return [(True, self._reset_step_buffers), *term_steps, (True, self._finalize_step)]
+        if operation == "reset":
+            term_steps = [
+                (is_warp_capturable(term_cfg.func), partial(self._reset_term, term_cfg))
+                for term_cfg in self._class_term_cfgs
+            ]
+            return [(True, self._average_last_episode_dones), *term_steps]
+        return super()._build_stage_steps(operation)
+
+    def _average_last_episode_dones(self, env_mask: wp.array) -> dict[str, torch.Tensor]:
+        """Average each term's last-episode done flags over all environments, as the stable manager does."""
+        if len(self._term_names) > 0:
+            self._term_done_avg_wp.zero_()
+            wp.launch(
+                kernel=_termination_reset_mean_all_2d,
+                dim=(self.num_envs, len(self._term_names)),
+                inputs=[self._last_episode_dones_wp, self._term_done_avg_wp],
+                device=self.device,
+            )
+        return self._reset_extras
+
+    def _reset_term(self, term_cfg: TerminationTermCfg, env_mask: wp.array) -> dict[str, torch.Tensor]:
+        term_cfg.func.reset(env_mask=env_mask)
+        return self._reset_extras
+
+    def _reset_step_buffers(self) -> None:
+        """Clear the done buffers in a single kernel launch."""
         wp.launch(
             kernel=_termination_pre_compute_reset,
             dim=self.num_envs,
@@ -299,11 +333,11 @@ def compute(self) -> torch.Tensor:
             device=self.device,
         )
 
-        # iterate over all the termination terms (fixed list; per-term math is Warp)
-        for term_cfg in self._term_cfgs:
-            term_cfg.func(self._env, term_cfg.out, **term_cfg.params)
+    def _compute_term(self, term_cfg: TerminationTermCfg) -> None:
+        term_cfg.func(self._env, term_cfg.out, **term_cfg.params)
 
-        # finalize dones and update last-episode term flags (single kernel launch)
+    def _finalize_step(self) -> torch.Tensor:
+        """Finalize the dones and update the last-episode term flags in a single kernel launch."""
         wp.launch(
             kernel=_termination_finalize,
             dim=self.num_envs,
@@ -317,7 +351,6 @@ def compute(self) -> torch.Tensor:
             ],
             device=self.device,
         )
-
         return self._dones_tensor_view
 
     def get_term(self, name: str) -> torch.Tensor:
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py
index 647d8cb3f2b7..8307158fa018 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py
@@ -26,6 +26,9 @@
 """Set ``ISAACLAB_SYNC_DEBUG=1`` to run eager stage calls under ``torch.cuda.set_sync_debug_mode("error")``,
 so a hidden GPU-to-host synchronization inside a stage raises instead of stalling silently."""
 
+StageStep = tuple[bool, Callable[..., Any]]
+"""One step of a stage: whether its work may be recorded into a CUDA graph, and the callable that runs it."""
+
 
 class WarpGraphCache:
     """Run Warp frontend stages eagerly or through cached CUDA graphs.
@@ -41,6 +44,8 @@ class WarpGraphCache:
       reallocated array, or a different scalar, records the stage again.
     * :meth:`invalidate` drops every graph and holds recording until the next :meth:`arm`.
       :meth:`invalidate_on` does so whenever the physics backend rebinds its buffers or stops.
+    * A stage made of steps (:meth:`call_steps`) records each run of consecutive capturable steps
+      as its own graph and runs the other steps eagerly in between, keeping the step order.
 
     A stage group is the stage name up to its first underscore. A group registered as not
     capturable, a CPU device, or ``ISAACLAB_WARP_CAPTURE=0`` runs its stages eagerly.
@@ -57,6 +62,9 @@ def __init__(self, device: wp.DeviceLike = None):
         self._sync_debug = os.environ.get(SYNC_DEBUG_ENV_VAR, "0") == "1"
         self._armed = False
         self._graphs: dict[str, tuple[tuple, wp.Graph, Any]] = {}
+        self._runs: dict[
+            int, tuple[tuple[StageStep, ...], tuple[tuple[bool, tuple[Callable[..., Any], ...]], ...]]
+        ] = {}
         self._capturable: dict[str, bool] = {}
         self._callback_handles = []
 
@@ -75,7 +83,7 @@ def call(
         timer: bool = False,
         **kwargs: Any,
     ) -> Any:
-        """Run a stage eagerly or through its recorded graph.
+        """Run a single-step stage eagerly or through its recorded graph.
 
         Args:
             stage: Stage identifier in the form ``"GroupName_function_name"``.
@@ -89,21 +97,65 @@ def call(
             The stage result, optionally transformed by :paramref:`output`. A replayed stage returns
             the result of the call that recorded it.
         """
+        return self._run(stage, ((True, (fn,)),), args, kwargs, output, timer)
+
+    def call_steps(
+        self,
+        stage: str,
+        steps: tuple[StageStep, ...],
+        /,
+        *args: Any,
+        output: Callable[[Any], Any] | None = None,
+        timer: bool = False,
+        **kwargs: Any,
+    ) -> Any:
+        """Run the steps of a stage in order, recording each run of consecutive capturable steps.
+
+        Every step receives the stage arguments. A run of consecutive capturable steps is recorded into
+        one graph, named ``"<stage>[<run index>]"`` when the stage has more than one run. Steps that are
+        not capturable run eagerly at their place in the order.
+
+        Args:
+            stage: Stage identifier in the form ``"GroupName_function_name"``.
+            steps: Ordered steps of the stage. Pass the same tuple on every call.
+            *args: Positional arguments forwarded to every step.
+            output: Transform applied to the stage result on every call, outside the graphs.
+            timer: Whether to time the stage.
+            **kwargs: Keyword arguments forwarded to every step.
+
+        Returns:
+            The result of the last step, optionally transformed by :paramref:`output`. A replayed run
+            returns the result of the call that recorded it.
+        """
+        return self._run(stage, self._split_runs(steps), args, kwargs, output, timer)
+
+    def _run(
+        self,
+        stage: str,
+        runs: tuple[tuple[bool, tuple[Callable[..., Any], ...]], ...],
+        args: tuple,
+        kwargs: dict[str, Any],
+        output: Callable[[Any], Any] | None,
+        timer: bool,
+    ) -> Any:
+        """Replay, record, or eagerly run each run of a stage in order."""
         with Timer(name=stage, msg=f"{stage} took:", enable=timer, time_unit="us"):
-            if not self._enabled or not self.is_capturable(stage.partition("_")[0]):
-                result = self._run_eager(fn, args, kwargs)
-            else:
-                key = _argument_key(args, kwargs)
-                recorded = self._graphs.get(stage)
+            capture = self._enabled and self.is_capturable(stage.partition("_")[0])
+            key = _argument_key(args, kwargs) if capture else None
+            result = None
+            for index, (capturable, fns) in enumerate(runs):
+                name = stage if len(runs) == 1 else f"{stage}[{index}]"
+                recorded = self._graphs.get(name) if capture and capturable else None
                 if recorded is not None and recorded[0] == key:
                     wp.capture_launch(recorded[1])
                     result = recorded[2]
-                elif self._armed:
-                    graph, result = self._record(fn, args, kwargs)
-                    self._graphs[stage] = (key, graph, result)
+                elif capture and capturable and self._armed:
+                    graph, result = self._record(fns, args, kwargs)
+                    self._graphs[name] = (key, graph, result)
                     wp.capture_launch(graph)
                 else:
-                    result = self._run_eager(fn, args, kwargs)
+                    for fn in fns:
+                        result = self._run_eager(fn, args, kwargs)
         return output(result) if output is not None else result
 
     def arm(self) -> None:
@@ -118,15 +170,13 @@ def invalidate(self, group: str | None = None) -> None:
                 Defaults to None, which drops every graph and holds recording until the next :meth:`arm`.
         """
         stages = [stage for stage in self._graphs if group is None or stage.partition("_")[0] == group]
-        if not stages:
-            if group is None:
-                self._armed = False
-            return
-        # Graphs may still be executing; drain the device before releasing them.
-        wp.synchronize_device(self._device)
-        for stage in stages:
-            del self._graphs[stage]
+        if stages:
+            # Graphs may still be executing; drain the device before releasing them.
+            wp.synchronize_device(self._device)
+            for stage in stages:
+                del self._graphs[stage]
         if group is None:
+            self._runs.clear()
             self._armed = False
 
     def invalidate_on(self, physics_manager: type[PhysicsManager]) -> None:
@@ -163,13 +213,28 @@ def is_capturable(self, group: str) -> bool:
     def _on_physics_event(self, payload: Any) -> None:
         self.invalidate()
 
-    def _record(self, fn: Callable[..., Any], args: tuple, kwargs: dict[str, Any]) -> tuple[wp.Graph, Any]:
-        """Record one stage call without executing it."""
+    def _split_runs(self, steps: tuple[StageStep, ...]) -> tuple[tuple[bool, tuple[Callable[..., Any], ...]], ...]:
+        """Group consecutive steps with the same capturability, memoized per steps tuple."""
+        memo = self._runs.get(id(steps))
+        if memo is None or memo[0] is not steps:
+            runs: list[tuple[bool, list[Callable[..., Any]]]] = []
+            for capturable, fn in steps:
+                if runs and runs[-1][0] == capturable:
+                    runs[-1][1].append(fn)
+                else:
+                    runs.append((capturable, [fn]))
+            memo = (steps, tuple((capturable, tuple(fns)) for capturable, fns in runs))
+            self._runs[id(steps)] = memo
+        return memo[1]
+
+    def _record(self, fns: tuple[Callable[..., Any], ...], args: tuple, kwargs: dict[str, Any]) -> tuple[wp.Graph, Any]:
+        """Record one run of steps without executing it."""
         result = None
 
         def capture_target() -> None:
             nonlocal result
-            result = fn(*args, **kwargs)
+            for fn in fns:
+                result = fn(*args, **kwargs)
 
         graph = NewtonQueries.capture_graph(str(self._device), capture_target, relaxed=has_kit())
         return graph, result
diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_capture.py b/source/isaaclab_experimental/test/envs/test_warp_env_capture.py
index 64bf18f84ef1..3887a4292fb8 100644
--- a/source/isaaclab_experimental/test/envs/test_warp_env_capture.py
+++ b/source/isaaclab_experimental/test/envs/test_warp_env_capture.py
@@ -9,7 +9,8 @@
 once with stage capture on. For manager-based tasks, a hard :meth:`SimulationContext.reset` halfway
 through rebuilds the Newton model and reallocates every simulation buffer; recorded graphs still
 point at the freed buffers, so the captured rollout only stays equal when the environment records
-its stages again.
+its stages again. A variant marks one term of every manager as not capturable: its stage must record
+the other terms and run that term eagerly in its place.
 """
 
 from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg
@@ -23,6 +24,7 @@
 import pytest
 import torch
 from isaaclab_experimental.envs.frontend import WarpFrontend
+from isaaclab_experimental.utils.warp import WarpCapturable
 from isaaclab_experimental.utils.warp_graph_cache import CAPTURE_ENV_VAR
 
 import isaaclab.sim as sim_utils
@@ -34,8 +36,18 @@
 _STEPS = 80
 _REBIND_STEP = 40
 
+# One function term per staged manager of the manager-based Cartpole.
+_CARTPOLE_EAGER_TERMS = (
+    "observations.policy.joint_pos_rel",
+    "rewards.pole_pos",
+    "terminations.cart_out_of_bounds",
+    "events.reset_cart_position",
+)
+
 
-def _rollout(task_id: str, capture: bool, rebind: bool, monkeypatch: pytest.MonkeyPatch) -> dict:
+def _rollout(
+    task_id: str, capture: bool, rebind: bool, monkeypatch: pytest.MonkeyPatch, eager_terms: tuple[str, ...] = ()
+) -> dict:
     """Run the seeded action sequence and return the rollout with the recorded stages."""
     monkeypatch.setenv(CAPTURE_ENV_VAR, "1" if capture else "0")
     env_cfg, _ = resolve_task_config(task_id, "", overrides=("physics=newton_mjwarp",))
@@ -44,6 +56,13 @@ def _rollout(task_id: str, capture: bool, rebind: bool, monkeypatch: pytest.Monk
     # contacts are only reproducible between runs with deterministic contact ordering
     env_cfg.sim.physics.deterministic_mode = "run_to_run"
     env_cfg.sim.physics.solver_cfg.disable_sensors = True
+    WarpFrontend.adapt_cfg(env_cfg)
+    for path in eager_terms:
+        term_cfg = env_cfg
+        for name in path.split("."):
+            term_cfg = getattr(term_cfg, name)
+        # the guard raises if the term is ever recorded into a graph
+        term_cfg.func = WarpCapturable(False, reason="marked eager by the test")(term_cfg.func)
     sim_utils.create_new_stage()
     env = WarpFrontend.build_env(env_cfg, task_id).unwrapped
     rollout = {"obs": [], "reward": [], "terminated": [], "truncated": []}
@@ -74,16 +93,26 @@ def _rollout(task_id: str, capture: bool, rebind: bool, monkeypatch: pytest.Monk
 
 # Direct tasks keep the simulation arrays they bind in ``__init__``, so they do not survive a rebind.
 @pytest.mark.parametrize(
-    ("task_id", "rebind"),
-    [("Isaac-Cartpole", True), ("Isaac-Cartpole-Direct", False)],
-    ids=["manager-cartpole", "direct-cartpole"],
+    ("task_id", "rebind", "eager_terms"),
+    [
+        ("Isaac-Cartpole", True, ()),
+        ("Isaac-Cartpole", False, _CARTPOLE_EAGER_TERMS),
+        ("Isaac-Cartpole-Direct", False, ()),
+    ],
+    ids=["manager-cartpole", "manager-cartpole-eager-terms", "direct-cartpole"],
 )
-def test_captured_rollout_matches_eager(task_id: str, rebind: bool, monkeypatch: pytest.MonkeyPatch):
-    eager = _rollout(task_id, capture=False, rebind=rebind, monkeypatch=monkeypatch)
-    captured = _rollout(task_id, capture=True, rebind=rebind, monkeypatch=monkeypatch)
+def test_captured_rollout_matches_eager(
+    task_id: str, rebind: bool, eager_terms: tuple[str, ...], monkeypatch: pytest.MonkeyPatch
+):
+    eager = _rollout(task_id, capture=False, rebind=rebind, monkeypatch=monkeypatch, eager_terms=eager_terms)
+    captured = _rollout(task_id, capture=True, rebind=rebind, monkeypatch=monkeypatch, eager_terms=eager_terms)
 
     assert eager["stages_at_end"] == ()
     assert captured["stages_at_end"], "no stage was recorded; the comparison proves nothing"
+    if eager_terms:
+        # every stage holding an eager term still records the terms around it
+        recorded_groups = {stage.partition("_")[0] for stage in captured["stages_at_end"]}
+        assert {"ObservationManager", "RewardManager", "TerminationManager", "EventManager"} <= recorded_groups
     if rebind:
         assert captured["stages_at_rebind"] == (), "the rebind must drop graphs that read freed buffers"
         assert captured["stages_at_end"] == captured["stages_before_rebind"]
diff --git a/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py b/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py
index 6ed2c1c27a54..794106aa0933 100644
--- a/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py
+++ b/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py
@@ -177,3 +177,36 @@ def test_sync_debug_trap_raises_on_a_hidden_host_sync(monkeypatch):
     with pytest.raises(RuntimeError, match="synchroniz"):
         cache.call("Group_stage", lambda: flags.any().item())
     assert torch.cuda.get_sync_debug_mode() == 0, "the trap must be released after the stage"
+
+
+@wp.kernel
+def _affine(values: wp.array(dtype=wp.int32), scale: wp.int32, shift: wp.int32):
+    values[wp.tid()] = values[wp.tid()] * scale + shift
+
+
+class _AffineStep:
+    """Stage step that applies ``x * scale + shift`` on the device and counts its Python invocations."""
+
+    def __init__(self, scale: int, shift: int):
+        self.scale, self.shift, self.calls = scale, shift, 0
+
+    def __call__(self, values: wp.array) -> wp.array:
+        self.calls += 1
+        wp.launch(_affine, dim=values.shape[0], inputs=[values, self.scale, self.shift], device=DEVICE)
+        return values
+
+
+def test_steps_record_capturable_runs_and_run_the_rest_eagerly_in_order(cache):
+    """Each run of capturable steps becomes one graph; an eager step runs between them on every call."""
+    double, increment, triple = _AffineStep(2, 0), _AffineStep(1, 1), _AffineStep(3, 0)
+    steps = ((True, double), (False, increment), (True, triple))
+    values = wp.ones(1, dtype=wp.int32, device=DEVICE)
+    cache.arm()
+
+    cache.call_steps("Group_stage", steps, values)
+    cache.call_steps("Group_stage", steps, values)
+
+    # the non-commuting steps only give 57 in order: ((1 * 2 + 1) * 3 = 9) -> ((9 * 2 + 1) * 3 = 57)
+    assert _read(values) == [57]
+    assert (double.calls, increment.calls, triple.calls) == (1, 2, 1)
+    assert cache.captured_stages == ("Group_stage[0]", "Group_stage[2]")

From eb61420a9efa2f0a4cabdf97d655d0bee6d149d5 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 03:47:31 -0700
Subject: [PATCH 13/48] Drop the removed fixed-tendon mask conversion from the
 sync gate

---
 .../test/utils/test_mask_boundary_gate.py                      | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py
index cb193e398be3..34c74ac87e61 100644
--- a/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py
+++ b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py
@@ -75,9 +75,8 @@
     ("isaaclab/sensors/camera/camera.py", "_env_mask_has_any"),
     # backends without a mask-native actuator reset fall back to indices; Newton and OVPhysX override it
     ("isaaclab/actuators/actuator_control.py", "reset_native_actuators_mask"),
-    # joint-limit and tendon-property writes mutate the solver model (event-driven, not per-step)
+    # joint-limit writes mutate the solver model (event-driven, not per-step)
     ("isaaclab_newton/assets/articulation/articulation.py", "write_joint_position_limit_to_sim_mask"),
-    ("isaaclab_newton/assets/articulation/articulation.py", "_mask_to_ids"),
 }
 
 

From bf583953a941ded9d2fed4f5f87cf1184599b4c8 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 04:05:46 -0700
Subject: [PATCH 14/48] Decide Warp term capturability from each term's
 parameters

A term can be capture-safe for one configuration and not another: a
reward that reads a ray-caster refreshes the sensor on the host and
launches its query graph, which cannot run inside a recorded stage,
while the same reward without a sensor can.

WarpCapturable accepts a predicate over the term's parameters, and its
guard raises only for the parameters the predicate rejects. The Warp
managers resolve capturability with each term's parameters, and a class
term may set _warp_capturable in __init__ to decide per instance. The
sync gate's opt-out inventory also lists predicate annotations.
---
 .../jichuanh-warp-frontend-cleanup.added.rst  |  4 +
 .../managers/event_manager.py                 |  6 +-
 .../managers/manager_base.py                  |  2 +-
 .../managers/observation_manager.py           |  7 +-
 .../managers/reward_manager.py                |  4 +-
 .../managers/termination_manager.py           |  4 +-
 .../isaaclab_experimental/utils/warp/utils.py | 65 +++++++++------
 .../test/managers/test_reward_manager.py      | 79 +++++++++++++++++++
 .../test/utils/test_mask_boundary_gate.py     |  9 +--
 9 files changed, 140 insertions(+), 40 deletions(-)
 create mode 100644 source/isaaclab_experimental/test/managers/test_reward_manager.py

diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst
index 37982f7d807d..37c0aacc0377 100644
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst
@@ -6,3 +6,7 @@
 * Added :meth:`~isaaclab_experimental.utils.WarpGraphCache.call_steps` and
   :meth:`ManagerBase.stage_steps <isaaclab_experimental.managers.ManagerBase.stage_steps>`, which split a
   manager stage into per-term steps that are recorded into CUDA graphs or run eagerly.
+* Added parameter predicates to :class:`~isaaclab_experimental.utils.warp.WarpCapturable`:
+  ``@WarpCapturable(lambda params: ...)`` marks a term capturable only for the parameters the predicate
+  accepts. A class term can also set ``self._warp_capturable`` in ``__init__``. The Warp managers decide
+  capture for each configured term instance.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
index fbeaaeb63524..a961e5b2fb16 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
@@ -354,17 +354,17 @@ def _build_stage_steps(self, operation: str) -> list[tuple[bool, Callable[..., A
         # Terms draw from the environment's shared random state, so every stage keeps the configured order.
         if operation == "apply_interval":
             return [
-                (is_warp_capturable(term_cfg.func), partial(self._apply_interval_term, index))
+                (is_warp_capturable(term_cfg.func, term_cfg.params), partial(self._apply_interval_term, index))
                 for index, term_cfg in enumerate(self._mode_term_cfgs.get("interval", []))
             ]
         if operation == "apply_reset":
             return [
-                (is_warp_capturable(term_cfg.func), partial(self._apply_reset_term, index))
+                (is_warp_capturable(term_cfg.func, term_cfg.params), partial(self._apply_reset_term, index))
                 for index, term_cfg in enumerate(self._mode_term_cfgs.get("reset", []))
             ]
         if operation == "reset":
             steps = [
-                (is_warp_capturable(term_cfg.func), partial(self._reset_class_term, term_cfg))
+                (is_warp_capturable(term_cfg.func, term_cfg.params), partial(self._reset_class_term, term_cfg))
                 for mode_cfgs in self._mode_class_term_cfgs.values()
                 for term_cfg in mode_cfgs
             ]
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py
index 9ad3064aec94..95c7d45f1c0c 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py
@@ -463,7 +463,7 @@ def _resolve_common_term_cfg(self, term_name: str, term_cfg: ManagerTermBaseCfg,
                     f" and optional parameters: {args_with_defaults}, but received: {term_params}."
                 )
 
-        if not is_warp_capturable(term_cfg.func):
+        if not is_warp_capturable(term_cfg.func, term_cfg.params):
             self._all_terms_capturable = False
 
         # process attributes at runtime
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
index 0fafea0f3829..80d89688c413 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
@@ -572,7 +572,7 @@ def _build_stage_steps(self, operation: str) -> list[tuple[bool, Callable[..., A
             # Noise draws from the environment's shared random state, so terms keep their order.
             steps = [
                 (
-                    is_warp_capturable(term_cfg.func) and term_cfg.history_length == 0,
+                    is_warp_capturable(term_cfg.func, term_cfg.params) and term_cfg.history_length == 0,
                     partial(self._compute_term, group_name, index),
                 )
                 for group_name, term_cfgs in self._group_obs_term_cfgs.items()
@@ -582,7 +582,10 @@ def _build_stage_steps(self, operation: str) -> list[tuple[bool, Callable[..., A
         if operation == "reset":
             steps = []
             for group_name, term_cfgs in self._group_obs_class_term_cfgs.items():
-                steps += [(is_warp_capturable(cfg.func), partial(self._reset_class_term, cfg)) for cfg in term_cfgs]
+                steps += [
+                    (is_warp_capturable(cfg.func, cfg.params), partial(self._reset_class_term, cfg))
+                    for cfg in term_cfgs
+                ]
                 # history buffers are not capture-safe
                 steps += [
                     (False, partial(self._reset_history, group_name, term_name))
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py
index 6b2cb0f27978..f415ce03c907 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py
@@ -276,7 +276,7 @@ def compute(self, dt: float) -> torch.Tensor:
     def _build_stage_steps(self, operation: str) -> list[tuple[bool, Callable[..., Any]]]:
         if operation == "compute":
             term_steps = [
-                (is_warp_capturable(term_cfg.func), partial(self._compute_term, index))
+                (is_warp_capturable(term_cfg.func, term_cfg.params), partial(self._compute_term, index))
                 for index, term_cfg in enumerate(self._term_cfgs)
             ]
             # Terms write their own output columns and the finalize step sums them in a fixed order, so
@@ -285,7 +285,7 @@ def _build_stage_steps(self, operation: str) -> list[tuple[bool, Callable[..., A
             return [(True, self._reset_step_buffers), *term_steps, (True, self._finalize_step)]
         if operation == "reset":
             term_steps = [
-                (is_warp_capturable(term_cfg.func), partial(self._reset_term, term_cfg))
+                (is_warp_capturable(term_cfg.func, term_cfg.params), partial(self._reset_term, term_cfg))
                 for term_cfg in self._class_term_cfgs
             ]
             return [(True, self._reset_episode_sums), *term_steps]
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py
index 5055a59b0f5e..32174f3eb33c 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py
@@ -293,7 +293,7 @@ def compute(self) -> torch.Tensor:
     def _build_stage_steps(self, operation: str) -> list[tuple[bool, Callable[..., Any]]]:
         if operation == "compute":
             term_steps = [
-                (is_warp_capturable(term_cfg.func), partial(self._compute_term, term_cfg))
+                (is_warp_capturable(term_cfg.func, term_cfg.params), partial(self._compute_term, term_cfg))
                 for term_cfg in self._term_cfgs
             ]
             # Terms write their own done columns and the finalize step combines them in a fixed order, so
@@ -302,7 +302,7 @@ def _build_stage_steps(self, operation: str) -> list[tuple[bool, Callable[..., A
             return [(True, self._reset_step_buffers), *term_steps, (True, self._finalize_step)]
         if operation == "reset":
             term_steps = [
-                (is_warp_capturable(term_cfg.func), partial(self._reset_term, term_cfg))
+                (is_warp_capturable(term_cfg.func, term_cfg.params), partial(self._reset_term, term_cfg))
                 for term_cfg in self._class_term_cfgs
             ]
             return [(True, self._average_last_episode_dones), *term_steps]
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py
index cb249af9fa02..f43dd0f425f9 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py
@@ -5,7 +5,8 @@
 
 from __future__ import annotations
 
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
+from typing import Any
 
 import torch
 import warp as wp
@@ -113,17 +114,23 @@ def decorator(func):
     return decorator
 
 
-def is_warp_capturable(func) -> bool:
-    """Check if a term function is CUDA-graph-capturable.
+def is_warp_capturable(func, params: dict[str, Any] | None = None) -> bool:
+    """Check if a term is CUDA-graph-capturable.
 
-    Checks ``_warp_capturable`` on the function and its ``__wrapped__`` target.
-    Returns True (capturable) by default if no annotation is found.
+    Checks ``_warp_capturable`` on the term (a function, a class, or a class instance) and on its
+    ``__wrapped__`` target. The annotation is a bool, or a predicate that decides from the term's
+    parameters (see :class:`WarpCapturable`). A class term may also set ``self._warp_capturable`` in
+    ``__init__`` to decide per instance. Returns True (capturable) if no annotation is found.
+
+    Args:
+        func: The term function, class, or class instance.
+        params: The term's parameters, passed to a predicate annotation. Defaults to None (no parameters).
     """
     for f in (func, getattr(func, "__wrapped__", None)):
         if f is not None:
             val = getattr(f, "_warp_capturable", None)
             if val is not None:
-                return val
+                return val if isinstance(val, bool) else bool(val(params or {}))
     return True
 
 
@@ -150,56 +157,64 @@ def reset_root_state_uniform(env, env_mask, ...):
         class randomize_rigid_body_com(ManagerTermBase):
             ...
 
+        @WarpCapturable(lambda params: params.get("sensor_cfg") is None, reason="refreshes the sensor on the host")
+        def base_height_l2(env, out, target_height, asset_cfg=..., sensor_cfg=None):
+            ...
+
     - ``@WarpCapturable(True)`` or no decorator: capturable, returned unwrapped.
     - ``@WarpCapturable(False)`` on a function: sets ``func._warp_capturable = False`` and wraps it with a
       runtime guard that raises if ``wp.get_device().is_capturing`` is ``True``.
     - ``@WarpCapturable(False)`` on a class term: annotates the class and guards its ``__call__``.
+    - ``@WarpCapturable(predicate)``: the term is capturable for the parameters the predicate accepts. The
+      predicate receives the term's parameters as a dict, and the guard raises only for rejected ones.
+
+    A class term can also decide per instance by setting ``self._warp_capturable`` in ``__init__``.
     """
 
-    def __init__(self, capturable: bool, *, reason: str | None = None):
+    def __init__(self, capturable: bool | Callable[[dict[str, Any]], bool], *, reason: str | None = None):
         self._capturable = capturable
         self._reason = reason
 
     def __call__(self, target):
         """Decorate a term function or class with the capture safety annotation and runtime guard."""
-        target._warp_capturable = self._capturable
-        if self._capturable:
-            return target
+        if isinstance(self._capturable, bool):
+            target._warp_capturable = self._capturable
+            if self._capturable:
+                return target
+        else:
+            # a plain function stored on a class would bind to instances
+            target._warp_capturable = staticmethod(self._capturable) if isinstance(target, type) else self._capturable
         if isinstance(target, type):
             target.__call__ = self._guarded(target.__call__)
             return target
         return self._guarded(target)
 
     def _guarded(self, func):
-        """Wrap *func* so that calling it during CUDA graph capture raises."""
+        """Wrap *func* so that calling it during CUDA graph capture with rejected parameters raises."""
         import functools
 
+        capturable = self._capturable
         reason = self._reason
 
         @functools.wraps(func)
         def wrapper(*args, **kwargs):
-            if wp.get_device().is_capturing:
-                msg = f"'{func.__qualname__}' is marked @WarpCapturable(False) but called during CUDA graph capture."
+            if wp.get_device().is_capturing and (isinstance(capturable, bool) or not capturable(kwargs)):
+                msg = (
+                    f"'{func.__qualname__}' is not capturable with these parameters but was called during CUDA"
+                    " graph capture."
+                )
                 if reason:
                     msg = f"{msg} {reason}"
                 raise RuntimeError(msg)
             return func(*args, **kwargs)
 
-        wrapper._warp_capturable = False
+        wrapper._warp_capturable = capturable
         return wrapper
 
     @staticmethod
-    def is_capturable(func) -> bool:
-        """Check capturability annotation. Default: ``True``.
-
-        Checks ``__wrapped__`` for decorated functions to handle stacked decorators.
-        """
-        for f in (func, getattr(func, "__wrapped__", None)):
-            if f is not None:
-                val = getattr(f, "_warp_capturable", None)
-                if val is not None:
-                    return val
-        return True
+    def is_capturable(func, params: dict[str, Any] | None = None) -> bool:
+        """Check the capturability annotation. See :func:`is_warp_capturable`."""
+        return is_warp_capturable(func, params)
 
 
 @wp.kernel
diff --git a/source/isaaclab_experimental/test/managers/test_reward_manager.py b/source/isaaclab_experimental/test/managers/test_reward_manager.py
new file mode 100644
index 000000000000..5af3755aeee5
--- /dev/null
+++ b/source/isaaclab_experimental/test/managers/test_reward_manager.py
@@ -0,0 +1,79 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tests for the Warp reward manager."""
+
+from collections import Counter
+from types import SimpleNamespace
+
+import pytest
+import torch
+import warp as wp
+from isaaclab_experimental.managers import ManagerTermBase, RewardManager
+from isaaclab_experimental.utils import WarpGraphCache
+from isaaclab_experimental.utils.warp import WarpCapturable
+
+from isaaclab.managers import RewardTermCfg
+
+wp.init()
+pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required")
+
+DEVICE = "cuda:0"
+NUM_ENVS = 4
+CALLS: Counter[float] = Counter()
+
+
+@WarpCapturable(lambda params: params.get("sensor_cfg") is None, reason="refreshes the sensor on the host")
+def constant_reward(env, out: wp.array, value: float, sensor_cfg: str | None = None):
+    """Function term that is capturable only without a sensor."""
+    CALLS[value] += 1
+    out.fill_(value)
+
+
+class ConstantReward(ManagerTermBase):
+    """Class term that decides its capturability per instance."""
+
+    def __init__(self, cfg: RewardTermCfg, env):
+        super().__init__(cfg, env)
+        self._warp_capturable = cfg.params["capturable"]
+
+    def __call__(self, env, out: wp.array, value: float, capturable: bool):
+        CALLS[value] += 1
+        out.fill_(value)
+
+
+def test_capturability_is_decided_per_term_instance():
+    """The same term is recorded for one parameter set and runs eagerly for another.
+
+    An eager instance runs its Python on every step; a recorded one only while its graph is recorded.
+    The guard of the predicate term raises if the instance with a sensor is ever recorded.
+    """
+    graph_cache = WarpGraphCache(DEVICE)
+    env = SimpleNamespace(
+        num_envs=NUM_ENVS,
+        device=DEVICE,
+        sim=SimpleNamespace(is_playing=lambda: True),
+        max_episode_length_s=1.0,
+        _warp_graph_cache=graph_cache,
+    )
+    cfg = {
+        "function_recorded": RewardTermCfg(func=constant_reward, weight=1.0, params={"value": 1.0}),
+        "function_eager": RewardTermCfg(
+            func=constant_reward, weight=1.0, params={"value": 2.0, "sensor_cfg": "height_scanner"}
+        ),
+        "class_recorded": RewardTermCfg(func=ConstantReward, weight=1.0, params={"value": 4.0, "capturable": True}),
+        "class_eager": RewardTermCfg(func=ConstantReward, weight=1.0, params={"value": 8.0, "capturable": False}),
+    }
+    manager = RewardManager(cfg, env)
+    graph_cache.arm()
+    CALLS.clear()
+
+    for _ in range(3):
+        reward = graph_cache.call_steps("RewardManager_compute", manager.stage_steps("compute"), dt=0.5)
+    wp.synchronize()
+
+    assert CALLS == {1.0: 1, 2.0: 3, 4.0: 1, 8.0: 3}
+    assert torch.equal(reward, torch.full((NUM_ENVS,), (1.0 + 2.0 + 4.0 + 8.0) * 0.5, device=DEVICE))
+    graph_cache.close()
diff --git a/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py
index 34c74ac87e61..66dd929894d5 100644
--- a/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py
+++ b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py
@@ -133,7 +133,7 @@ def test_mask_first_code_has_no_unsanctioned_host_syncs():
     assert not stale, f"Stale sanctioned boundaries (function gone or no longer syncs): {sorted(stale)}"
 
 
-# ``@WarpCapturable(False)`` opt-outs, by (path suffix, name).
+# ``@WarpCapturable(False)`` and parameter-predicate opt-outs, by (path suffix, name).
 EXPECTED_NON_CAPTURABLE = {
     ("isaaclab_experimental/envs/mdp/events.py", "randomize_rigid_body_com"),
     ("isaaclab_experimental/envs/mdp/events.py", "randomize_rigid_body_mass"),
@@ -147,7 +147,7 @@ def test_mask_first_code_has_no_unsanctioned_host_syncs():
 
 
 def _non_capturable_targets(tree: ast.Module) -> list[str]:
-    """Names of functions and classes decorated with ``@WarpCapturable(False, ...)``."""
+    """Names of functions and classes decorated with ``@WarpCapturable(False)`` or a parameter predicate."""
     targets = []
     for node in ast.walk(tree):
         if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
@@ -158,15 +158,14 @@ def _non_capturable_targets(tree: ast.Module) -> list[str]:
                 and isinstance(decorator.func, ast.Name)
                 and decorator.func.id == "WarpCapturable"
                 and decorator.args
-                and isinstance(decorator.args[0], ast.Constant)
-                and decorator.args[0].value is False
+                and not (isinstance(decorator.args[0], ast.Constant) and decorator.args[0].value is True)
             ):
                 targets.append(node.name)
     return targets
 
 
 def test_non_capturable_terms_are_inventoried():
-    """Every ``@WarpCapturable(False)`` opt-out is listed here, and only these."""
+    """Every ``@WarpCapturable(False)`` or predicate opt-out is listed here, and only these."""
     found = set()
     for scan_root in NON_CAPTURABLE_SCAN_ROOTS:
         for path in sorted((_REPO_ROOT / scan_root).rglob("*.py")):

From ab709ebdbfca444d21c952aa186d631228878972 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 11:01:48 -0700
Subject: [PATCH 15/48] Let physics own decimation in the Warp environments

Stable environments tell the physics backend their decimation and, when the
backend runs the whole decimation loop, apply actions, step the simulation and
update the scene once per environment step. The Warp environments still looped
once per physics step, so contact sensors sampled four times as often on Go2
and the contact history, air times and the rewards and terminations reading
them differed from the stable environment with the same configuration.

Mirror the stable loop in the manager-based and direct Warp environments. The
per-physics-step loop remains for backends that do not run the decimation loop.
---
 .../jichuanh-warp-frontend-cleanup.fixed.rst  |   4 +
 .../envs/direct_rl_env_warp.py                |  15 +-
 .../envs/manager_based_env_warp.py            |  15 +-
 .../envs/manager_based_rl_env_warp.py         |  10 +-
 .../test/envs/test_warp_env_decimation.py     | 160 ++++++++++++++++++
 5 files changed, 189 insertions(+), 15 deletions(-)
 create mode 100644 source/isaaclab_experimental/test/envs/test_warp_env_decimation.py

diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
index 9fbaf0c2ae3c..a72e668bb213 100644
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
@@ -1,5 +1,9 @@
 * Fixed Warp environments replaying CUDA graphs that read freed simulation buffers after the simulation was
   reset or stopped. The environments now record their stages again on the next step.
+* Fixed the Warp environments applying actions, stepping the simulation and updating the scene once per physics
+  step when the physics backend runs the whole decimation loop. Like the stable environments, they now do it
+  once per environment step, so contact histories, air times and the rewards and terminations reading them
+  match the stable environments.
 * Fixed :attr:`~isaaclab_experimental.utils.buffers.CircularBuffer.max_length` synchronizing with the GPU.
 * Fixed the Warp ``randomize_rigid_body_com`` event accumulating offsets across calls. It now offsets the
   center of mass the term finds on its first call, like the stable term.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
index 981c64dfd344..cac82aca4e65 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
@@ -122,6 +122,7 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
         self.render_mode = render_mode
         # initialize internal variables
         self._is_closed = False
+        self._physics_handles_decimation = False
 
         # set the seed for the environment
         if self.cfg.seed is not None:
@@ -189,6 +190,10 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
             # this shouldn't cause an issue since later on, users do a reset over all the
             # environments so the lazy buffers would be reset.
             self.scene.update(dt=self.physics_dt)
+        # let the physics backend know about the env decimation so it can
+        # fold the full loop into a single step() when possible
+        self.sim.physics_manager.set_decimation(self.cfg.decimation)
+        self._physics_handles_decimation = self.sim.physics_manager.handles_decimation()
 
         # check if debug visualization is has been implemented by the environment
         source_code = inspect.getsource(self._set_debug_vis_impl)
@@ -410,10 +415,11 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         # note: hoisted out of the decimation loop; is_rendering does live settings lookups
         is_rendering = self.sim.is_rendering
 
-        # perform physics stepping
+        # physics-owned decimation covers all substeps in one call
+        steps_per_call = self.cfg.decimation if self._physics_handles_decimation else 1
         with Timer(name="physics_loop", msg="Physics loop took:", enable=DEBUG_TIMERS):
-            for _ in range(self.cfg.decimation):
-                self._sim_step_counter += 1
+            for _ in range(self.cfg.decimation // steps_per_call):
+                self._sim_step_counter += steps_per_call
                 # set actions into buffers
                 # simulate
                 with Timer(name="apply_action", msg="Action processing step took:", enable=DEBUG_TIMERS):
@@ -431,9 +437,8 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
                 #    If a camera needs rendering at a faster frequency, this will lead to unexpected behavior.
                 if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering:
                     self.sim.render()
-                # update buffers at sim dt
                 with Timer(name="scene_update", msg="Scene update took:", enable=DEBUG_TIMERS):
-                    self.scene.update(dt=self.physics_dt)
+                    self.scene.update(dt=self.physics_dt * steps_per_call)
 
         self.common_step_counter += 1  # total step (common for all envs)
         with Timer(name="end_pre_graph", msg="End pre-graph took:", enable=DEBUG_TIMERS):
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
index 1c30edde04f2..6fc70b825df8 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
@@ -90,6 +90,7 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
             )
         # initialize internal variables
         self._is_closed = False
+        self._physics_handles_decimation = False
 
         # set the seed for the environment
         if self.cfg.seed is not None:
@@ -193,6 +194,10 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
             # this shouldn't cause an issue since later on, users do a reset over all the environments so the lazy
             # buffers would be reset.
             self.scene.update(dt=self.physics_dt)
+        # let the physics backend know about the env decimation so it can
+        # fold the full loop into a single step() when possible
+        self.sim.physics_manager.set_decimation(self.cfg.decimation)
+        self._physics_handles_decimation = self.sim.physics_manager.handles_decimation()
 
         # TODO(jichuanh): This is a temporary solution for event_manager only, but it should be general for all managers
         # Resolve SceneEntityCfg-dependent term params once before any captured event paths.
@@ -571,9 +576,10 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]:
         # note: hoisted out of the decimation loop; is_rendering does live settings lookups
         is_rendering = self.sim.is_rendering
 
-        # perform physics stepping
-        for _ in range(self.cfg.decimation):
-            self._sim_step_counter += 1
+        # physics-owned decimation covers all substeps in one call
+        steps_per_call = self.cfg.decimation if self._physics_handles_decimation else 1
+        for _ in range(self.cfg.decimation // steps_per_call):
+            self._sim_step_counter += steps_per_call
             # set actions into buffers
             self._warp_graph_cache.call_steps(
                 "ActionManager_apply_action", self.action_manager.stage_steps("apply_action")
@@ -587,8 +593,7 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]:
             #    If a camera needs rendering at a faster frequency, this will lead to unexpected behavior.
             if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering:
                 self.sim.render()
-            # update buffers at sim dt
-            self.scene.update(dt=self.physics_dt)
+            self.scene.update(dt=self.physics_dt * steps_per_call)
 
         # post-step: step interval event
         if "interval" in self.event_manager.available_modes:
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
index 9470c10dbe5f..bbf58910687e 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
@@ -254,9 +254,10 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         # note: checked here once to avoid multiple checks within the loop
         is_rendering = self.sim.is_rendering
 
-        # perform physics stepping
-        for _ in range(self.cfg.decimation):
-            self._sim_step_counter += 1
+        # physics-owned decimation covers all substeps in one call
+        steps_per_call = self.cfg.decimation if self._physics_handles_decimation else 1
+        for _ in range(self.cfg.decimation // steps_per_call):
+            self._sim_step_counter += steps_per_call
             # set actions into buffers
             self._warp_graph_cache.call_steps(
                 "ActionManager_apply_action", self.action_manager.stage_steps("apply_action"), timer=DEBUG_TIMER_STEP
@@ -276,14 +277,13 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
             #    If a camera needs rendering at a faster frequency, this will lead to unexpected behavior.
             if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering:
                 self.sim.render()
-            # update buffers at sim dt
             with Timer(
                 name="scene.update",
                 msg="Scene.update took:",
                 enable=DEBUG_TIMER_STEP,
                 time_unit="us",
             ):
-                self.scene.update(dt=self.physics_dt)
+                self.scene.update(dt=self.physics_dt * steps_per_call)
 
         # post-step:
         # -- update env counters (used for curriculum generation)
diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_decimation.py b/source/isaaclab_experimental/test/envs/test_warp_env_decimation.py
new file mode 100644
index 000000000000..9ef6e9a3e8ed
--- /dev/null
+++ b/source/isaaclab_experimental/test/envs/test_warp_env_decimation.py
@@ -0,0 +1,160 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""The Warp manager-based environment steps physics at the stable environment's cadence.
+
+When every actuator runs inside Newton, the physics backend folds the whole decimation loop into one
+step: actions are applied, the simulation stepped and the scene updated once per environment step.
+Contact sensors then sample once per environment step, so their history, the air times and every
+reward and termination reading them depend on that cadence. The same seeded Go2 rollout on the stable
+and the Warp frontend must make the same calls and produce the same contact data.
+"""
+
+from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg
+
+from isaaclab.sim import SimulationCfg
+from isaaclab.test.utils import launch_test_simulation
+
+launch_test_simulation(SimulationCfg(physics=NewtonCfg(solver_cfg=MJWarpSolverCfg())))
+
+import math
+
+import isaaclab_tasks_experimental  # noqa: F401
+import pytest
+import torch
+from isaaclab_experimental.envs.frontend import WarpFrontend
+
+import isaaclab.sim as sim_utils
+from isaaclab.envs import ManagerBasedRLEnv
+
+import isaaclab_tasks  # noqa: F401
+from isaaclab_tasks.core.velocity.velocity_env_cfg import RewardsCfg
+from isaaclab_tasks.utils import resolve_task_config
+
+_TASK_ID = "Isaac-Velocity-Flat-UnitreeGo2"
+_NUM_ENVS = 64
+_STEPS = 200
+
+
+def _go2_cfg():
+    """Go2 flat with every random draw removed, so both frontends see identical physics inputs."""
+    env_cfg, _ = resolve_task_config(_TASK_ID, "", overrides=("physics=newton_mjwarp",))
+    env_cfg.seed = 7
+    env_cfg.scene.num_envs = _NUM_ENVS
+    # contacts are only reproducible between runs with deterministic contact ordering
+    env_cfg.sim.physics.deterministic_mode = "run_to_run"
+    env_cfg.sim.physics.solver_cfg.disable_sensors = True
+    env_cfg.events.add_base_mass = None
+    env_cfg.events.base_com = None
+    env_cfg.events.push_robot = None
+    env_cfg.events.reset_base.params["pose_range"] = {}
+    env_cfg.events.reset_base.params["velocity_range"] = {}
+    env_cfg.events.reset_robot_joints.params["position_range"] = (1.0, 1.0)
+    command = env_cfg.commands.base_velocity
+    command.rel_standing_envs = 0.0
+    command.rel_heading_envs = 0.0
+    command.heading_command = False
+    command.debug_vis = False
+    command.ranges.lin_vel_x = (0.5, 0.5)
+    command.ranges.lin_vel_y = (0.0, 0.0)
+    command.ranges.ang_vel_z = (0.0, 0.0)
+    # Go2 drops the thigh-contact penalty; restore it to cover a second contact-history reader
+    env_cfg.rewards.undesired_contacts = RewardsCfg().undesired_contacts
+    env_cfg.rewards.undesired_contacts.params["sensor_cfg"].body_names = ".*_thigh"
+    return env_cfg
+
+
+def _rollout(frontend: str) -> dict:
+    """Step a fixed action sequence and record the physics calls and the contact data of every step."""
+    env_cfg = _go2_cfg()
+    sim_utils.create_new_stage()
+    if frontend == "warp":
+        env = WarpFrontend.build_env(env_cfg, _TASK_ID).unwrapped
+    else:
+        env = ManagerBasedRLEnv(cfg=env_cfg)
+    calls = {"sim_step": 0, "apply_action": 0}
+    sim_step = env.sim.step
+
+    def count_sim_step(*args, **kwargs):
+        calls["sim_step"] += 1
+        return sim_step(*args, **kwargs)
+
+    env.sim.step = count_sim_step
+    if frontend == "warp":
+        call_steps = env._warp_graph_cache.call_steps
+
+        def count_apply_action(stage, *args, **kwargs):
+            calls["apply_action"] += stage == "ActionManager_apply_action"
+            return call_steps(stage, *args, **kwargs)
+
+        env._warp_graph_cache.call_steps = count_apply_action
+        step_reward = env.reward_manager._step_reward_tensor_view
+        term_dones = env.termination_manager._term_dones_tensor_view
+    else:
+        apply_action = env.action_manager.apply_action
+
+        def count_apply_action():
+            calls["apply_action"] += 1
+            apply_action()
+
+        env.action_manager.apply_action = count_apply_action
+        step_reward = env.reward_manager._step_reward
+        term_dones = env.termination_manager._term_dones
+    reward_names = env.reward_manager._term_names
+    base_contact = env.termination_manager._term_names.index("base_contact")
+    sensor = env.scene["contact_forces"].data
+    rollout = {"calls": [], "handles_decimation": env._physics_handles_decimation}
+    records = {
+        "net_forces_w_history": lambda: sensor.net_forces_w_history.torch,
+        "current_air_time": lambda: sensor.current_air_time.torch,
+        "last_air_time": lambda: sensor.last_air_time.torch,
+        "undesired_contacts": lambda: step_reward[:, reward_names.index("undesired_contacts")],
+        "feet_air_time": lambda: step_reward[:, reward_names.index("feet_air_time")],
+        "base_contact": lambda: term_dones[:, base_contact],
+    }
+    try:
+        env.reset()
+        action_dim = env.action_space.shape[-1]
+        envs = torch.arange(_NUM_ENVS, dtype=torch.float32)[:, None]
+        joints = torch.arange(action_dim, dtype=torch.float32)[None, :]
+        # growing amplitudes make environments fall, and reset, at different steps
+        amplitude = 1.0 + 3.0 * envs / (_NUM_ENVS - 1)
+        for step in range(_STEPS):
+            if step < 20:
+                action = torch.zeros((_NUM_ENVS, action_dim))
+            else:
+                action = amplitude * torch.sin(2 * math.pi * step / 25.0 + 0.7 * joints + 0.05 * envs)
+            calls["sim_step"] = calls["apply_action"] = 0
+            env.step(action.to(env.device))
+            rollout["calls"].append((calls["sim_step"], calls["apply_action"]))
+            for key, read in records.items():
+                rollout.setdefault(key, []).append(read().clone())
+    finally:
+        env.close()
+    for key in records:
+        rollout[key] = torch.stack(rollout[key])
+    return rollout
+
+
+@pytest.fixture(scope="module")
+def rollouts() -> dict[str, dict]:
+    return {frontend: _rollout(frontend) for frontend in ("stable", "warp")}
+
+
+def test_physics_loop_cadence_matches_stable(rollouts: dict[str, dict]):
+    stable, warp = rollouts["stable"], rollouts["warp"]
+    assert stable["handles_decimation"], "Go2 runs its actuators inside Newton, so physics must own decimation"
+    assert warp["handles_decimation"]
+    assert warp["calls"] == stable["calls"] == [(1, 1)] * _STEPS
+
+
+def test_contact_data_matches_stable(rollouts: dict[str, dict]):
+    stable, warp = rollouts["stable"], rollouts["warp"]
+    assert stable["base_contact"].any(), "no environment terminated; the resets are not exercised"
+    for key in ("net_forces_w_history", "current_air_time", "last_air_time", "undesired_contacts", "base_contact"):
+        assert torch.equal(warp[key], stable[key]), key
+    # stable stores ``raw * weight * dt / dt`` and reduces the feet with torch, the Warp manager stores
+    # ``raw * weight`` and its term sums the feet in a loop: the two differ in the last bit only
+    torch.testing.assert_close(warp["feet_air_time"], stable["feet_air_time"], rtol=0.0, atol=1e-6)

From 3ac3db42dfea287f844cc5ebd56c2797ce33da24 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 15:32:12 -0700
Subject: [PATCH 16/48] Store terminal observations in the Warp environments

The stable environments honor compute_final_obs: a step that resets
environments stores the observations from before the reset in
extras["final_obs"] and declares Same-Step autoreset in the metadata. The
Warp environments ignored the flag, so torchrl, which sets it, bootstrapped
time-outs from the post-reset observation.

Compute the observations before the reset stage and keep a copy, because the
reset overwrites the persistent observation buffers. Declare Same-Step
autoreset in both Warp environment classes.
---
 .../jichuanh-warp-frontend-cleanup.fixed.rst  |   4 +
 .../envs/direct_rl_env_warp.py                |   6 +
 .../envs/manager_based_rl_env_warp.py         |  11 ++
 .../test/envs/test_warp_env_final_obs.py      | 129 ++++++++++++++++++
 4 files changed, 150 insertions(+)
 create mode 100644 source/isaaclab_experimental/test/envs/test_warp_env_final_obs.py

diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
index a72e668bb213..474b02a3fc9f 100644
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
@@ -4,6 +4,10 @@
   step when the physics backend runs the whole decimation loop. Like the stable environments, they now do it
   once per environment step, so contact histories, air times and the rewards and terminations reading them
   match the stable environments.
+* Fixed :class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp` and
+  :class:`~isaaclab_experimental.envs.DirectRLEnvWarp` ignoring ``compute_final_obs``. A step that resets
+  environments now stores the observations from before the reset in ``extras["final_obs"]``, and the
+  environments declare Same-Step autoreset in their metadata, like the stable environments.
 * Fixed :attr:`~isaaclab_experimental.utils.buffers.CircularBuffer.max_length` synchronizing with the GPU.
 * Fixed the Warp ``randomize_rigid_body_com`` event accumulating offsets across calls. It now offsets the
   center of mass the term finds on its first call, like the stable term.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
index cac82aca4e65..8204c3449491 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
@@ -98,6 +98,7 @@ class DirectRLEnvWarp(DirectRLEnv):
     """Whether the environment is a vectorized environment."""
     metadata: ClassVar[dict[str, Any]] = {
         "render_modes": [None, "human", "rgb_array"],
+        "autoreset_mode": gym.vector.AutoresetMode.SAME_STEP,
         # "isaac_sim_version": get_version(),
     }
     """Metadata for the environment."""
@@ -445,6 +446,11 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
             self._warp_graph_cache.call("DirectEndPre_step", self._step_warp_end_pre)
         # one host predicate per step skips the reset stage when no environment terminated
         if self._torch_reset_buf.any().item():
+            # capture the terminal observation before reset and expose it for Same-Step autoreset.
+            # the reset overwrites the persistent observation buffer, so keep a copy
+            if self.cfg.compute_final_obs:
+                self._warp_graph_cache.call("DirectFinalObs_step", self._get_observations)
+                self.extras["final_obs"] = {"policy": self.torch_obs_buf.clone()}
             with Timer(name="reset_graph", msg="Reset graph took:", enable=DEBUG_TIMERS):
                 self._warp_graph_cache.call("DirectReset_step", self._step_warp_reset)
         # write_data_to_sim runs uncaptured — it uses torch ops that cross CUDA streams.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
index bbf58910687e..4776099dbca2 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
@@ -77,6 +77,7 @@ class ManagerBasedRLEnvWarp(ManagerBasedEnvWarp, gym.Env):
     """Whether the environment is a vectorized environment."""
     metadata: ClassVar[dict[str, Any]] = {
         "render_modes": [None, "human", "rgb_array"],
+        "autoreset_mode": gym.vector.AutoresetMode.SAME_STEP,
         # "isaac_sim_version": get_version(),
     }
     """Metadata for the environment."""
@@ -445,6 +446,16 @@ def _reset_terminated_envs(self) -> None:
         if not has_resets:
             return
 
+        # capture the terminal observation before reset and expose it for Same-Step autoreset.
+        # the reset overwrites the persistent observation buffers, so keep a copy
+        if self.cfg.compute_final_obs:
+            final_obs = self._warp_graph_cache.call_steps(
+                "ObservationManager_compute_no_history",
+                self.observation_manager.stage_steps("compute"),
+                return_cloned_output=False,
+                timer=DEBUG_TIMER_STEP,
+            )
+            self.extras["final_obs"] = clone_obs_buffer(final_obs)
         # trigger recorder terms for pre-reset calls
         self.recorder_manager.record_pre_reset(reset_env_ids)
 
diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_final_obs.py b/source/isaaclab_experimental/test/envs/test_warp_env_final_obs.py
new file mode 100644
index 000000000000..07c2e4413244
--- /dev/null
+++ b/source/isaaclab_experimental/test/envs/test_warp_env_final_obs.py
@@ -0,0 +1,129 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Terminal observations of the Warp environments match the stable environments.
+
+With ``compute_final_obs`` set, a step that resets environments exposes the observation computed before the
+reset in ``extras["final_obs"]`` and returns the observation after the reset, as the stable environments do
+for Same-Step autoreset.
+"""
+
+from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg
+
+from isaaclab.sim import SimulationCfg
+from isaaclab.test.utils import launch_test_simulation
+
+launch_test_simulation(SimulationCfg(physics=NewtonCfg(solver_cfg=MJWarpSolverCfg())))
+
+import gymnasium as gym
+import isaaclab_tasks_experimental  # noqa: F401
+import pytest
+import torch
+from isaaclab_experimental.envs.frontend import WarpFrontend
+
+import isaaclab.sim as sim_utils
+
+import isaaclab_tasks  # noqa: F401
+from isaaclab_tasks.utils import resolve_task_config
+
+_NUM_ENVS = 16
+_STEPS = 120
+_CART_BOUND = 3.0
+_RESET_CART_POS = 0.5
+_TASK_IDS = ("Isaac-Cartpole", "Isaac-Cartpole-Direct")
+
+
+def _cartpole_cfg(task_id: str, compute_final_obs: bool):
+    """Cartpole with fixed reset states, so the stable and the Warp random draws agree."""
+    env_cfg, _ = resolve_task_config(task_id, "", overrides=("physics=newton_mjwarp",))
+    env_cfg.seed = 7
+    env_cfg.scene.num_envs = _NUM_ENVS
+    env_cfg.compute_final_obs = compute_final_obs
+    if task_id == "Isaac-Cartpole":
+        env_cfg.events.reset_cart_position.params.update(
+            position_range=(_RESET_CART_POS, _RESET_CART_POS), velocity_range=(0.0, 0.0)
+        )
+        env_cfg.events.reset_pole_position.params.update(position_range=(0.1, 0.1), velocity_range=(0.0, 0.0))
+    else:
+        env_cfg.initial_cart_position_range = (_RESET_CART_POS, _RESET_CART_POS)
+        env_cfg.initial_cart_velocity_range = (0.0, 0.0)
+        env_cfg.initial_pole_angle_range = (0.1, 0.1)
+        env_cfg.initial_pole_velocity_range = (0.0, 0.0)
+    return env_cfg
+
+
+def _rollout(task_id: str, frontend: str, compute_final_obs: bool) -> dict:
+    """Push every cart with a constant force and record the observations of every step."""
+    env_cfg = _cartpole_cfg(task_id, compute_final_obs)
+    sim_utils.create_new_stage()
+    if frontend == "warp":
+        env = WarpFrontend.build_env(env_cfg, task_id).unwrapped
+    else:
+        env = gym.make(task_id, cfg=env_cfg).unwrapped
+    rollout = {"obs": [], "dones": [], "final_obs": {}, "final_obs_key": []}
+    try:
+        env.reset()
+        # different pushes drive the carts out of bounds at different steps
+        push = torch.linspace(-0.3, 0.3, _NUM_ENVS, device=env.device)[:, None]
+        for step in range(_STEPS):
+            obs, _, terminated, truncated, extras = env.step(push)
+            dones = terminated | truncated
+            rollout["obs"].append(obs["policy"].clone())
+            rollout["dones"].append(dones.clone())
+            rollout["final_obs_key"].append("final_obs" in extras)
+            if dones.any() and "final_obs" in extras:
+                rollout["final_obs"][step] = extras["final_obs"]["policy"].clone()
+        rollout["cart_index"] = (
+            0 if task_id != "Isaac-Cartpole" else env.scene["robot"].find_joints("slider_to_cart")[0][0]
+        )
+    finally:
+        env.close()
+    rollout["obs"] = torch.stack(rollout["obs"])
+    rollout["dones"] = torch.stack(rollout["dones"])
+    return rollout
+
+
+@pytest.fixture(scope="module")
+def rollouts() -> dict[tuple[str, str, bool], dict]:
+    arms = [(task_id, frontend, True) for task_id in _TASK_IDS for frontend in ("stable", "warp")]
+    arms += [(task_id, "warp", False) for task_id in _TASK_IDS]
+    return {arm: _rollout(*arm) for arm in arms}
+
+
+@pytest.mark.parametrize("task_id", _TASK_IDS)
+def test_final_obs_is_the_observation_before_reset(task_id: str, rollouts: dict):
+    rollout = rollouts[(task_id, "warp", True)]
+    cart = rollout["cart_index"]
+    reset_steps = torch.nonzero(rollout["dones"].any(dim=1)).squeeze(-1).tolist()
+    assert reset_steps, "no environment terminated; the terminal observation is not exercised"
+    assert sorted(rollout["final_obs"]) == reset_steps
+    for step in reset_steps:
+        dones = rollout["dones"][step]
+        final_obs, obs = rollout["final_obs"][step], rollout["obs"][step]
+        # the terminal observation still has the cart out of bounds, the returned one is reset
+        assert (final_obs[dones, cart].abs() > _CART_BOUND).all()
+        assert torch.equal(obs[dones, cart], torch.full_like(obs[dones, cart], _RESET_CART_POS))
+        assert torch.equal(final_obs[~dones], obs[~dones])
+
+
+@pytest.mark.parametrize("task_id", _TASK_IDS)
+def test_final_obs_is_absent_when_disabled(task_id: str, rollouts: dict):
+    rollout = rollouts[(task_id, "warp", False)]
+    assert rollout["dones"].any()
+    assert not any(rollout["final_obs_key"])
+
+
+@pytest.mark.parametrize(
+    ("task_id", "atol"),
+    # the direct twin differs from the stable task in the last bit only (at most 4.8e-7 at |obs| < 8)
+    [("Isaac-Cartpole", 0.0), ("Isaac-Cartpole-Direct", 1e-6)],
+)
+def test_final_obs_matches_stable(task_id: str, atol: float, rollouts: dict):
+    stable, warp = rollouts[(task_id, "stable", True)], rollouts[(task_id, "warp", True)]
+    assert torch.equal(warp["dones"], stable["dones"])
+    torch.testing.assert_close(warp["obs"], stable["obs"], rtol=0.0, atol=atol)
+    assert sorted(warp["final_obs"]) == sorted(stable["final_obs"])
+    for step, final_obs in stable["final_obs"].items():
+        torch.testing.assert_close(warp["final_obs"][step], final_obs, rtol=0.0, atol=atol, msg=f"step {step}")

From e6fcbbab71e7fedd9aed0e8c9d895cd165c60a38 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 15:41:24 -0700
Subject: [PATCH 17/48] Return copies of the logged values from the Warp
 environments

Warp managers and tasks log persistent buffers that later steps overwrite.
Loggers such as the rsl_rl logger keep extras["log"] of every step and average
it at the end of an iteration, so every entry showed the latest value. The
stable environments log fresh tensors.

The manager-based environment copies the tensor values when it builds the
reset log. The direct environment keeps the log the task writes and returns a
copy of it from every step, so task code that is recorded into a CUDA graph
keeps refreshing its own buffers.
---
 .../jichuanh-warp-frontend-cleanup.fixed.rst  |  4 +
 .../envs/direct_rl_env_warp.py                | 13 ++++
 .../envs/manager_based_rl_env_warp.py         |  4 +-
 .../test/envs/test_warp_env_log.py            | 78 +++++++++++++++++++
 4 files changed, 98 insertions(+), 1 deletion(-)
 create mode 100644 source/isaaclab_experimental/test/envs/test_warp_env_log.py

diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
index 474b02a3fc9f..a3454630d1fe 100644
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
@@ -8,6 +8,10 @@
   :class:`~isaaclab_experimental.envs.DirectRLEnvWarp` ignoring ``compute_final_obs``. A step that resets
   environments now stores the observations from before the reset in ``extras["final_obs"]``, and the
   environments declare Same-Step autoreset in their metadata, like the stable environments.
+* Fixed :class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp` and
+  :class:`~isaaclab_experimental.envs.DirectRLEnvWarp` returning ``extras["log"]`` values that later steps
+  overwrite. Loggers that keep the logs of an iteration and average them, such as the rsl_rl logger, reported
+  the latest value for every step. Each step now returns copies, like the stable environments.
 * Fixed :attr:`~isaaclab_experimental.utils.buffers.CircularBuffer.max_length` synchronizing with the GPU.
 * Fixed the Warp ``randomize_rigid_body_com`` event accumulating offsets across calls. It now offsets the
   center of mass the term finds on its first call, like the stable term.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
index 8204c3449491..eb9cc6dc257b 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
@@ -212,6 +212,8 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
 
         # allocate dictionary to store metrics
         self.extras = {}
+        # the log the task writes; the environment returns copies of it (see :meth:`step`)
+        self._task_log: dict | None = None
 
         # initialize data and constants
         # -- counter for simulation steps
@@ -401,6 +403,9 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
 
         # stages start recording at the first step after construction or an invalidation
         self._warp_graph_cache.arm()
+        # task code that runs during this step writes into its own log, not into the copy returned last step
+        if self._task_log is not None:
+            self.extras["log"] = self._task_log
         action = action.to(self.device)
         # add action noise
         if self.cfg.action_noise_model:
@@ -464,6 +469,14 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         with Timer(name="visualize", msg="Visualize took:", enable=DEBUG_TIMERS):
             self._post_step_visualize()
 
+        # tasks log persistent buffers that later steps overwrite, and RL libraries keep the log of every step
+        # until they average it, so hand out copies like the fresh tensors of stable tasks
+        self._task_log = self.extras.get("log")
+        if self._task_log:
+            self.extras["log"] = {
+                key: value.clone() if torch.is_tensor(value) else value for key, value in self._task_log.items()
+            }
+
         # return observations, rewards, resets and extras
         # store the returned buffer so RslRlVecEnvWrapper.get_observations() can read env.obs_buf
         self.obs_buf = {"policy": self.torch_obs_buf.clone()}
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
index 4776099dbca2..5ee14a667dbc 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
@@ -565,4 +565,6 @@ def _reset_mask(self, env_mask: wp.array, env_ids: torch.Tensor | None = None) -
             recorder_info,
         ):
             log.update(info)
-        self.extras["log"] = log
+        # the managers log persistent buffers that the next reset overwrites, and RL libraries keep the log of
+        # every step until they average it, so hand out copies like the stable managers' fresh tensors
+        self.extras["log"] = {key: value.clone() if torch.is_tensor(value) else value for key, value in log.items()}
diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_log.py b/source/isaaclab_experimental/test/envs/test_warp_env_log.py
new file mode 100644
index 000000000000..51324adcaf30
--- /dev/null
+++ b/source/isaaclab_experimental/test/envs/test_warp_env_log.py
@@ -0,0 +1,78 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Logged values of the Warp environments keep the value of the step that returned them.
+
+Warp managers and tasks log persistent buffers that later steps overwrite. RL libraries keep the
+``extras["log"]`` of every step and average them at the end of an iteration, so the environments must return
+copies, like the fresh tensors of the stable environments.
+"""
+
+from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg
+
+from isaaclab.sim import SimulationCfg
+from isaaclab.test.utils import launch_test_simulation
+
+launch_test_simulation(SimulationCfg(physics=NewtonCfg(solver_cfg=MJWarpSolverCfg())))
+
+import isaaclab_tasks_experimental  # noqa: F401
+import torch
+from isaaclab_experimental.envs.frontend import WarpFrontend
+
+import isaaclab.sim as sim_utils
+
+import isaaclab_tasks  # noqa: F401
+from isaaclab_tasks.utils import resolve_task_config
+
+_NUM_ENVS = 16
+_STEPS = 120
+
+
+def _make_env(task_id: str):
+    env_cfg, _ = resolve_task_config(task_id, "", overrides=("physics=newton_mjwarp",))
+    env_cfg.seed = 7
+    env_cfg.scene.num_envs = _NUM_ENVS
+    sim_utils.create_new_stage()
+    return WarpFrontend.build_env(env_cfg, task_id).unwrapped
+
+
+def test_reset_logs_keep_their_values():
+    env = _make_env("Isaac-Cartpole")
+    logged = []
+    try:
+        env.reset()
+        # different pushes drive the carts out of bounds at different steps, so episode sums differ
+        push = torch.linspace(-0.3, 0.3, _NUM_ENVS, device=env.device)[:, None]
+        for _ in range(_STEPS):
+            _, _, terminated, truncated, extras = env.step(push)
+            if (terminated | truncated).any():
+                log = extras["log"]
+                logged.append((log, {key: value.clone() for key, value in log.items() if torch.is_tensor(value)}))
+    finally:
+        env.close()
+
+    first_values = logged[0][1]
+    assert any(not torch.equal(values[key], first_values[key]) for _, values in logged[1:] for key in first_values)
+    for log, values in logged:
+        for key, value in values.items():
+            assert torch.equal(log[key], value), key
+
+
+def test_direct_logs_keep_their_values():
+    env = _make_env("Isaac-Cartpole-Direct")
+    logged = []
+    try:
+        env.reset()
+        # a task logs a persistent buffer that its kernels refresh
+        metric = torch.zeros((), device=env.device)
+        env.extras["log"] = {"Metrics/value": metric}
+        for step in range(5):
+            metric.fill_(step)
+            _, _, _, _, extras = env.step(torch.zeros((_NUM_ENVS, 1), device=env.device))
+            logged.append(extras["log"])
+    finally:
+        env.close()
+
+    assert [log["Metrics/value"].item() for log in logged] == [0.0, 1.0, 2.0, 3.0, 4.0]

From a1743b7c0d6476472cb91d4b96bff2f7cd12fb19 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 15:44:08 -0700
Subject: [PATCH 18/48] Return copies of the reset log from the base Warp
 environment

ManagerBasedEnvWarp built its reset log from the same persistent manager
buffers as the RL environment, so a logged value changed with the next reset.
Copy the tensor values like the RL environment does.
---
 .../jichuanh-warp-frontend-cleanup.fixed.rst  |  3 ++-
 .../envs/manager_based_env_warp.py            |  5 ++++
 .../test/envs/test_warp_env_log.py            | 24 +++++++++++++++++++
 3 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
index a3454630d1fe..f2186f9b5823 100644
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
@@ -8,7 +8,8 @@
   :class:`~isaaclab_experimental.envs.DirectRLEnvWarp` ignoring ``compute_final_obs``. A step that resets
   environments now stores the observations from before the reset in ``extras["final_obs"]``, and the
   environments declare Same-Step autoreset in their metadata, like the stable environments.
-* Fixed :class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp` and
+* Fixed :class:`~isaaclab_experimental.envs.ManagerBasedEnvWarp`,
+  :class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp` and
   :class:`~isaaclab_experimental.envs.DirectRLEnvWarp` returning ``extras["log"]`` values that later steps
   overwrite. Loggers that keep the logs of an iteration and average them, such as the rsl_rl logger, reported
   the latest value for every step. Each step now returns copies, like the stable environments.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
index 6fc70b825df8..b0e2d07cb8eb 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
@@ -714,3 +714,8 @@ def _reset_mask(self, env_mask: wp.array, env_ids: torch.Tensor | None = None) -
         # -- recorder manager
         info = self.recorder_manager.reset(env_ids)
         self.extras["log"].update(info)
+        # the managers log persistent buffers that the next reset overwrites, and RL libraries keep the log of
+        # every step until they average it, so hand out copies like the stable managers' fresh tensors
+        self.extras["log"] = {
+            key: value.clone() if torch.is_tensor(value) else value for key, value in self.extras["log"].items()
+        }
diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_log.py b/source/isaaclab_experimental/test/envs/test_warp_env_log.py
index 51324adcaf30..fd5a88e0549a 100644
--- a/source/isaaclab_experimental/test/envs/test_warp_env_log.py
+++ b/source/isaaclab_experimental/test/envs/test_warp_env_log.py
@@ -18,7 +18,9 @@
 launch_test_simulation(SimulationCfg(physics=NewtonCfg(solver_cfg=MJWarpSolverCfg())))
 
 import isaaclab_tasks_experimental  # noqa: F401
+import pytest
 import torch
+from isaaclab_experimental.envs import ManagerBasedEnvWarp
 from isaaclab_experimental.envs.frontend import WarpFrontend
 
 import isaaclab.sim as sim_utils
@@ -60,6 +62,28 @@ def test_reset_logs_keep_their_values():
             assert torch.equal(log[key], value), key
 
 
+def test_base_env_reset_logs_keep_their_values(monkeypatch: pytest.MonkeyPatch):
+    env_cfg, _ = resolve_task_config("Isaac-Cartpole", "", overrides=("physics=newton_mjwarp",))
+    env_cfg.seed = 7
+    env_cfg.scene.num_envs = _NUM_ENVS
+    WarpFrontend.adapt_cfg(env_cfg)
+    sim_utils.create_new_stage()
+    env = ManagerBasedEnvWarp(cfg=env_cfg)
+    logged = []
+    try:
+        # a manager logs a persistent buffer that it refreshes on every reset
+        metric = torch.zeros((), device=env.device)
+        monkeypatch.setattr(env.recorder_manager, "reset", lambda env_ids=None: {"Metrics/value": metric})
+        for value in range(3):
+            metric.fill_(value)
+            _, extras = env.reset()
+            logged.append(extras["log"])
+    finally:
+        env.close()
+
+    assert [log["Metrics/value"].item() for log in logged] == [0.0, 1.0, 2.0]
+
+
 def test_direct_logs_keep_their_values():
     env = _make_env("Isaac-Cartpole-Direct")
     logged = []

From e70010ee90bf15e0521e6dd20bebe4e208968cba Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 03:43:00 -0700
Subject: [PATCH 19/48] Fix the Warp reorient twin's orientation error and
 object velocity

The shared Warp reorientation environment measured the orientation error
from the (y, z, w) components of the (x, y, z, w) quaternion difference, so
aligned orientations read as ~pi: its orientation reward, goal success, and
goal resampling diverged from the stable task. Its full observation also
read the object's [lin_vel, ang_vel] root velocity as [ang, lin].

Add a same-state parity test that recomputes the stable task's observations
and rewards from the twin's own state after every captured step.
---
 .../jichuanh-task-cleanup-dex-warp-exp.skip   |   0
 .../test/envs/test_reorient_warp_env.py       | 109 ++++++++++++++++++
 ...chuanh-task-cleanup-dex-warp-exp.fixed.rst |   5 +
 .../jichuanh-task-cleanup-dex-warp-exp.minor  |   0
 .../core/reorient/reorient_warp_env.py        |  23 ++--
 5 files changed, 124 insertions(+), 13 deletions(-)
 create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.skip
 create mode 100644 source/isaaclab_experimental/test/envs/test_reorient_warp_env.py
 create mode 100644 source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.fixed.rst
 create mode 100644 source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.minor

diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.skip b/source/isaaclab_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.skip
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/source/isaaclab_experimental/test/envs/test_reorient_warp_env.py b/source/isaaclab_experimental/test/envs/test_reorient_warp_env.py
new file mode 100644
index 000000000000..381c462c1d20
--- /dev/null
+++ b/source/isaaclab_experimental/test/envs/test_reorient_warp_env.py
@@ -0,0 +1,109 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Same-state parity of the Warp in-hand reorientation twins against the stable Direct task.
+
+Each twin is built from its stable registration through :class:`WarpFrontend` and stepped on its
+default captured path. After every step, the stable torch implementation recomputes observations
+and rewards from the twin's own state, so the comparison isolates the twin's kernels from
+simulation noise.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import warp as wp
+from isaaclab_experimental.envs.frontend import WarpFrontend
+
+from isaaclab.sim import SimulationContext
+from isaaclab.utils.math import quat_error_magnitude
+
+import isaaclab_tasks  # noqa: F401
+from isaaclab_tasks.core.reorient.reorient_direct_env import ReorientDirectEnv, reorient_reward
+from isaaclab_tasks.utils import resolve_task_config
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="The Warp twins capture CUDA graphs.")
+
+_TASKS = ["Isaac-Reorient-Cube-Allegro-Direct"]
+
+
+@pytest.fixture(scope="module", params=_TASKS)
+def env(request):
+    """A 16-environment twin of the parametrized stable task, shared by this module's tests."""
+    cfg, _ = resolve_task_config(request.param, "", overrides=("physics=newton_mjwarp",))
+    cfg.scene.num_envs = 16
+    cfg.seed = 42
+    env = WarpFrontend.build_env(cfg, request.param)
+    yield env
+    env.close()
+    SimulationContext.clear_instance()
+
+
+def _stable_state(env, actions: torch.Tensor) -> SimpleNamespace:
+    """The state the stable task derives its observations and rewards from, read off the twin."""
+    state = SimpleNamespace(
+        cfg=env.cfg,
+        num_envs=env.num_envs,
+        hand=env.hand,
+        object=env.object,
+        scene=env.scene,
+        finger_bodies=env.hand.find_bodies(env.cfg.fingertip_body_names)[0],
+        actions=actions,
+        goal_rot=wp.to_torch(env.goal_rot),
+    )
+    state.num_fingertips = len(state.finger_bodies)
+    ReorientDirectEnv._compute_intermediate_values(state)
+    joint_limits = env.hand.data.joint_limits.torch
+    state.hand_dof_lower_limits, state.hand_dof_upper_limits = joint_limits[..., 0], joint_limits[..., 1]
+    state.in_hand_pos = env.object.data.default_root_pose.torch[:, :3] + torch.tensor(
+        env.cfg.in_hand_pos_offset, device=env.device
+    )
+    return state
+
+
+def test_observations_and_rewards_match_stable_task(env):
+    """Observations and rewards equal the stable task's on the same state.
+
+    Rewards are computed before the step's resets, so they are compared on the environments whose
+    state no episode reset or goal resample changed afterwards.
+    """
+    env.reset()
+    generator = torch.Generator(device=env.device).manual_seed(0)
+    num_compared = 0
+    for _ in range(40):
+        goal_rot = wp.to_torch(env.goal_rot).clone()
+        # unit Gaussian: a policy's samples leave the [-1, 1] action range too
+        actions = torch.randn((env.num_envs, env.cfg.action_space), device=env.device, generator=generator)
+        obs, reward, terminated, truncated, _ = env.step(actions)
+        stable = _stable_state(env, actions)
+
+        torch.testing.assert_close(obs["policy"], ReorientDirectEnv.compute_full_observations(stable))
+
+        untouched = ~(terminated | truncated) & (stable.goal_rot == goal_rot).all(dim=-1)
+        error = quat_error_magnitude(stable.object_rot, stable.goal_rot)
+        no_reset = torch.zeros_like(untouched)
+        expected, *_ = reorient_reward(
+            no_reset,
+            no_reset,
+            torch.zeros_like(error),
+            torch.zeros(1, device=env.device),
+            stable.object_pos,
+            stable.in_hand_pos,
+            error <= env.cfg.success_tolerance,
+            error,
+            actions,
+            env.cfg.dist_reward_scale,
+            env.cfg.rot_reward_scale,
+            env.cfg.rot_eps,
+            env.cfg.action_penalty_scale,
+            env.cfg.reach_goal_bonus,
+            env.cfg.fall_dist,
+            env.cfg.fall_penalty,
+            env.cfg.av_factor,
+        )
+        torch.testing.assert_close(reward[untouched], expected[untouched], rtol=1e-4, atol=1e-4)
+        num_compared += int(untouched.sum())
+    assert num_compared > 0
diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.fixed.rst b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.fixed.rst
new file mode 100644
index 000000000000..393a1e9d9837
--- /dev/null
+++ b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.fixed.rst
@@ -0,0 +1,5 @@
+* Fixed the ``Isaac-Reorient-Cube-Allegro-Direct`` Warp environment measuring the object's
+  orientation error from the ``(y, z, w)`` quaternion components instead of ``(x, y, z)``, which
+  made its orientation reward, goal success, and goal resampling differ from the stable task.
+* Fixed the ``Isaac-Reorient-Cube-Allegro-Direct`` Warp environment swapping the object's linear
+  and angular velocity in its full observation.
diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.minor b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.minor
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
index c2a34edc9da6..778660d86591 100644
--- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
+++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
@@ -359,16 +359,15 @@ def compute_full_observations(
     observations[env_id, offset + 3] = obj_rot[3]
     offset += 4
 
-    # spatial_vectorf layout: [0:3]=angular, [3:6]=linear
-    # torch reference order: linear (unscaled) first, then angular (scaled)
-    observations[env_id, offset + 0] = object_vels[env_id][3]
-    observations[env_id, offset + 1] = object_vels[env_id][4]
-    observations[env_id, offset + 2] = object_vels[env_id][5]
+    # root velocities are laid out [lin_vel, ang_vel]; only the angular part is scaled
+    observations[env_id, offset + 0] = object_vels[env_id][0]
+    observations[env_id, offset + 1] = object_vels[env_id][1]
+    observations[env_id, offset + 2] = object_vels[env_id][2]
     offset += 3
 
-    observations[env_id, offset + 0] = vel_obs_scale * object_vels[env_id][0]
-    observations[env_id, offset + 1] = vel_obs_scale * object_vels[env_id][1]
-    observations[env_id, offset + 2] = vel_obs_scale * object_vels[env_id][2]
+    observations[env_id, offset + 0] = vel_obs_scale * object_vels[env_id][3]
+    observations[env_id, offset + 1] = vel_obs_scale * object_vels[env_id][4]
+    observations[env_id, offset + 2] = vel_obs_scale * object_vels[env_id][5]
     offset += 3
 
     # goal
@@ -535,12 +534,10 @@ def randomize_rotation(rand0: wp.float32, rand1: wp.float32, x_axis: wp.vec3f, y
 
 @wp.func
 def rotation_distance(object_rot: wp.quatf, target_rot: wp.quatf) -> wp.float32:
-    # Orientation alignment for the cube in hand and goal cube
+    """Angle [rad] between two ``(x, y, z, w)`` orientations, as :func:`~isaaclab.utils.math.quat_error_magnitude`."""
     quat_diff = object_rot * wp.quat_inverse(target_rot)
-    # Match Torch env convention: uses indices [1:4] for the vector part (see `rotation_distance` in Torch env).
-    v_norm = wp.sqrt(quat_diff[1] * quat_diff[1] + quat_diff[2] * quat_diff[2] + quat_diff[3] * quat_diff[3])
-    v_norm = wp.min(v_norm, wp.float32(1.0))
-    return wp.float32(2.0) * wp.asin(v_norm)
+    v_norm = wp.length(wp.vec3f(quat_diff[0], quat_diff[1], quat_diff[2]))
+    return wp.float32(2.0) * wp.asin(wp.min(v_norm, wp.float32(1.0)))
 
 
 class ReorientDirectWarpEnv(DirectRLEnvWarp):

From b7707006a0326bfbf9a32e94424d564b4cb043dd Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 03:47:43 -0700
Subject: [PATCH 20/48] Share the Warp reorient twin's success and reset paths
 with other hands

Prepare the Warp reorientation environment for hands beyond Allegro:

- Evaluate the orientation error and goal success once per step in
  _get_dones and reuse them in the reward, as the stable task does,
  instead of recomputing the quaternion error in the reward kernel.
- Apply reset states through the assets' mask writers, which invalidate
  derived data and map public joint order onto the solver's, instead of
  writing the data buffers directly.
- Read the in-hand anchor and goal-marker position from the task cfg and
  command joints through the actuator target API.
- Log Metrics/success_rate on reset, matching the stable task.
- Reject asymmetric_obs configurations: this environment has no critic
  observations.
---
 ...chuanh-task-cleanup-dex-warp-exp.added.rst |   2 +
 ...uanh-task-cleanup-dex-warp-exp.changed.rst |   3 +
 .../core/reorient/reorient_warp_env.py        | 175 +++++++++---------
 3 files changed, 97 insertions(+), 83 deletions(-)
 create mode 100644 source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.added.rst
 create mode 100644 source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.changed.rst

diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.added.rst b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.added.rst
new file mode 100644
index 000000000000..ec9ef0e39417
--- /dev/null
+++ b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.added.rst
@@ -0,0 +1,2 @@
+* Added the ``Metrics/success_rate`` episode metric to the Warp in-hand reorientation environment,
+  matching the stable Direct task.
diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.changed.rst b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.changed.rst
new file mode 100644
index 000000000000..20484dbf7694
--- /dev/null
+++ b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.changed.rst
@@ -0,0 +1,3 @@
+* Changed the Warp in-hand reorientation environment to raise :class:`NotImplementedError` for
+  configurations with ``asymmetric_obs`` enabled instead of silently omitting the critic
+  observations.
diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
index 778660d86591..6b40bfd36d12 100644
--- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
+++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
@@ -6,7 +6,6 @@
 
 from __future__ import annotations
 
-from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
 import torch
@@ -106,7 +105,7 @@ def reset_object(
     root_pose_w: wp.array(dtype=wp.transformf),
     root_vel_w: wp.array(dtype=wp.spatial_vectorf),
 ):
-    """Reset masked envs' object root pose with position noise and a randomized rotation, zeroing its velocity."""
+    """Sample masked envs' object root pose with position noise and a randomized rotation, and a zero velocity."""
     env_id = wp.tid()
     if env_mask[env_id]:
         nx = wp.randf(rng_state[env_id], wp.float32(-1.0), wp.float32(1.0))
@@ -126,7 +125,6 @@ def reset_object(
         rng_state[env_id] += wp.uint32(1)
         rot_w = randomize_rotation(rand0, rand1, x_unit_vec, y_unit_vec)
 
-        # The following should be equivalent, but consider using write_root_pose_to_sim and write_root_velocity_to_sim
         root_pose_w[env_id] = wp.transform(pos_w, rot_w)
         root_vel_w[env_id] = wp.spatial_vectorf(
             wp.float32(0.0), wp.float32(0.0), wp.float32(0.0), wp.float32(0.0), wp.float32(0.0), wp.float32(0.0)
@@ -151,9 +149,8 @@ def reset_hand(
     joint_vel: wp.array2d(dtype=wp.float32),
     prev_targets: wp.array2d(dtype=wp.float32),
     cur_targets: wp.array2d(dtype=wp.float32),
-    hand_dof_targets: wp.array2d(dtype=wp.float32),
 ):
-    """Reset masked envs' hand joint positions and velocities with noise and initialize their position targets."""
+    """Sample masked envs' reset joint positions and velocities with noise and initialize their position targets."""
     env_id = wp.tid()
     if env_mask[env_id]:
         # Each env runs sequentially inside this kernel (avoids RNG races across DOFs).
@@ -170,29 +167,47 @@ def reset_hand(
             rng_state[env_id] += wp.uint32(1)
             vel = default_joint_vel[env_id, dof_id] + reset_dof_vel_noise * dof_vel_noise
 
-            # The following lines should be equivalent to the following:
-            # self.hand.write_joint_state_to_sim(dof_pos, dof_vel, env_ids=env_ids)
             joint_pos[env_id, dof_id] = pos
             joint_vel[env_id, dof_id] = vel
-
             prev_targets[env_id, dof_id] = pos
             cur_targets[env_id, dof_id] = pos
-            hand_dof_targets[env_id, dof_id] = pos
 
 
 @wp.kernel
 def reset_successes(
     # input
     env_mask: wp.array(dtype=wp.bool),
-    # output
+    # input/output
     successes: wp.array(dtype=wp.float32),
+    success_rate_stats: wp.array(dtype=wp.float32),
 ):
-    """Zero the consecutive-success counter for masked envs."""
+    """Accumulate masked envs' finished-episode success rate into ``success_rate_stats``, then zero their counts.
+
+    Reaching a goal draws a replacement, so an episode that reached ``n`` goals attempted ``n + 1``.
+    ``success_rate_stats`` must be zeroed before the launch; ``[0]`` receives the summed per-episode
+    rates and ``[1]`` the number of finished episodes.
+    """
     env_id = wp.tid()
     if env_mask[env_id]:
+        goals = successes[env_id]
+        wp.atomic_add(success_rate_stats, 0, goals / (goals + wp.float32(1.0)))
+        wp.atomic_add(success_rate_stats, 1, wp.float32(1.0))
         successes[env_id] = wp.float32(0.0)
 
 
+@wp.kernel
+def update_success_rate_from_stats(
+    # input
+    success_rate_stats: wp.array(dtype=wp.float32),
+    # output
+    success_rate: wp.array(dtype=wp.float32),
+):
+    """Average the accumulated episode success rates; keep the last value when no episode finished."""
+    # single-thread kernel (dim=1)
+    if success_rate_stats[1] > wp.float32(0.0):
+        success_rate[0] = success_rate_stats[0] / success_rate_stats[1]
+
+
 @wp.kernel
 def compute_intermediate_values(
     # input
@@ -242,11 +257,17 @@ def get_dones(
     # input/output
     episode_length_buf: wp.array(dtype=wp.int32),
     # output
+    orientation_error: wp.array(dtype=wp.float32),
+    goal_reached: wp.array(dtype=wp.bool),
     out_of_reach: wp.array(dtype=wp.bool),
     time_out: wp.array(dtype=wp.bool),
     reset: wp.array(dtype=wp.bool),
 ):
-    """Flag object-fall and time-out/max-success termination, reset progress on goal, write the reset union."""
+    """Evaluate goal success once per step, flag object-fall and time-out/max-success termination.
+
+    The orientation error [rad] and success flags are written for the reward to reuse. Progress is
+    reset on reaching a goal when a consecutive-success cap is configured.
+    """
     env_id = wp.tid()
 
     object_pos = wp.transform_get_translation(object_pose[env_id])
@@ -255,11 +276,14 @@ def get_dones(
     goal_dist = wp.length(object_pos - in_hand_pos[env_id])
     out_of_reach[env_id] = goal_dist >= fall_dist
 
+    error = rotation_distance(object_rot, goal_rot[env_id])
+    reached = error <= success_tolerance
+    orientation_error[env_id] = error
+    goal_reached[env_id] = reached
+
     max_success_reached = False
     if max_consecutive_success > 0:
-        # Reset progress (episode length buf) on goal envs if max_consecutive_success > 0
-        rot_dist = rotation_distance(object_rot, goal_rot[env_id])
-        if wp.abs(rot_dist) <= success_tolerance:
+        if reached:
             episode_length_buf[env_id] = 0
         max_success_reached = successes[env_id] >= wp.float32(max_consecutive_success)
 
@@ -437,13 +461,13 @@ def compute_rewards(
     reset_buf: wp.array(dtype=wp.bool),
     object_pose: wp.array(dtype=wp.transformf),
     target_pos: wp.array(dtype=wp.vec3f),
-    target_rot: wp.array(dtype=wp.quatf),
+    goal_reached: wp.array(dtype=wp.bool),
+    orientation_error: wp.array(dtype=wp.float32),
     dist_reward_scale: wp.float32,
     rot_reward_scale: wp.float32,
     rot_eps: wp.float32,
     actions: wp.array2d(dtype=wp.float32),
     action_penalty_scale: wp.float32,
-    success_tolerance: wp.float32,
     reach_goal_bonus: wp.float32,
     fall_dist: wp.float32,
     fall_penalty: wp.float32,
@@ -460,13 +484,10 @@ def compute_rewards(
     env_id = wp.tid()
 
     obj_pos = wp.transform_get_translation(object_pose[env_id])
-    obj_rot = wp.transform_get_rotation(object_pose[env_id])
-
     goal_dist = wp.length(obj_pos - target_pos[env_id])
-    rot_dist = rotation_distance(obj_rot, target_rot[env_id])
 
     dist_rew = goal_dist * dist_reward_scale
-    rot_rew = wp.float32(1.0) / (wp.abs(rot_dist) + rot_eps) * rot_reward_scale
+    rot_rew = rot_reward_scale / (orientation_error[env_id] + rot_eps)
 
     action_penalty = wp.float32(0.0)
     for i in range(action_dim):
@@ -475,15 +496,11 @@ def compute_rewards(
     # Total reward is: position distance + orientation alignment + action regularization + success bonus + fall penalty
     reward = dist_rew + rot_rew + action_penalty * action_penalty_scale
 
-    # Find out which envs hit the goal and update successes count
-    reached = wp.abs(rot_dist) <= success_tolerance
-    goal_resets = reached or reset_goal_buf[env_id]
+    # a goal stays flagged until it is resampled, which happens later this step
+    goal_resets = goal_reached[env_id] or reset_goal_buf[env_id]
     reset_goal_buf[env_id] = goal_resets
     if goal_resets:
         successes[env_id] = successes[env_id] + wp.float32(1.0)
-
-    # Success bonus: orientation is within `success_tolerance` of goal orientation
-    if goal_resets:
         reward = reward + reach_goal_bonus
 
     # Fall penalty: distance to the goal is larger than a threshold
@@ -541,10 +558,13 @@ def rotation_distance(object_rot: wp.quatf, target_rot: wp.quatf) -> wp.float32:
 
 
 class ReorientDirectWarpEnv(DirectRLEnvWarp):
-    cfg: AllegroHandEnvCfg  # | ShadowHandEnvCfg
+    """Warp twin of :class:`~isaaclab_tasks.core.reorient.reorient_direct_env.ReorientDirectEnv`."""
+
+    cfg: AllegroHandEnvCfg
 
-    # def __init__(self, cfg: AllegroHandEnvCfg | ShadowHandEnvCfg, render_mode: str | None = None, **kwargs):
     def __init__(self, cfg: AllegroHandEnvCfg, render_mode: str | None = None, **kwargs):
+        if cfg.asymmetric_obs:
+            raise NotImplementedError("The Warp reorientation environment has no asymmetric critic observations.")
         super().__init__(cfg, render_mode, **kwargs)
         self.hand, self.object, self.goal_markers = [self.scene[name] for name in ("robot", "object", "goal_object")]
 
@@ -561,13 +581,7 @@ def __init__(self, cfg: AllegroHandEnvCfg, render_mode: str | None = None, **kwa
             actuated_dof_indices.append(self.hand.joint_names.index(joint_name))
         actuated_dof_indices.sort()
         self.num_actuated_dofs = len(actuated_dof_indices)
-
-        # Warp index/mask helpers for kernels and articulation APIs.
         self.actuated_dof_indices = wp.array(actuated_dof_indices, dtype=wp.int32, device=self.device)
-        actuated_mask = [False] * self.num_hand_dofs
-        for idx in actuated_dof_indices:
-            actuated_mask[idx] = True
-        self.actuated_dof_mask = wp.array(actuated_mask, dtype=wp.bool, device=self.device)
 
         # finger bodies
         finger_bodies: list[int] = list()
@@ -584,7 +598,6 @@ def __init__(self, cfg: AllegroHandEnvCfg, render_mode: str | None = None, **kwa
         # unit vectors
         self.x_unit_vec = wp.vec3f(1.0, 0.0, 0.0)
         self.y_unit_vec = wp.vec3f(0.0, 1.0, 0.0)
-        self.z_unit_vec = wp.vec3f(0.0, 0.0, 1.0)
 
         # Per-env origins (Warp view for kernels; Torch env uses `self.scene.env_origins` directly).
         self.env_origins = wp.from_torch(self.scene.env_origins, dtype=wp.vec3f)
@@ -594,10 +607,19 @@ def __init__(self, cfg: AllegroHandEnvCfg, render_mode: str | None = None, **kwa
         # ---------------------------------------------------------------------
 
         # buffers for position targets
-        self.hand_dof_targets = wp.zeros((self.num_envs, self.num_hand_dofs), dtype=wp.float32, device=self.device)
         self.prev_targets = wp.zeros((self.num_envs, self.num_hand_dofs), dtype=wp.float32, device=self.device)
         self.cur_targets = wp.zeros((self.num_envs, self.num_hand_dofs), dtype=wp.float32, device=self.device)
 
+        # reset states sampled per env; the asset writers apply the masked rows
+        self.reset_joint_pos = wp.zeros((self.num_envs, self.num_hand_dofs), dtype=wp.float32, device=self.device)
+        self.reset_joint_vel = wp.zeros((self.num_envs, self.num_hand_dofs), dtype=wp.float32, device=self.device)
+        self.reset_object_pose = wp.zeros(self.num_envs, dtype=wp.transformf, device=self.device)
+        self.reset_object_vel = wp.zeros(self.num_envs, dtype=wp.spatial_vectorf, device=self.device)
+
+        # per-step goal evaluation, written in `_get_dones` and reused by the reward
+        self.orientation_error = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device)
+        self.goal_reached = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device)
+
         # track goal resets
         self.reset_goal_buf = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device)
         # used to compare object position
@@ -608,12 +630,12 @@ def __init__(self, cfg: AllegroHandEnvCfg, render_mode: str | None = None, **kwa
         self.goal_pos_w = wp.zeros(self.num_envs, dtype=wp.vec3f, device=self.device)
 
         # Initialize goal constants from Torch (avoid a one-off kernel launch).
-        default_root_pose = self.object.data.default_root_pose.torch.to(self.device)
-        in_hand_pos = default_root_pose[:, 0:3].clone()
-        in_hand_pos[:, 2] -= 0.04
+        in_hand_pos = self.object.data.default_root_pose.torch[:, 0:3] + torch.tensor(
+            self.cfg.in_hand_pos_offset, device=self.device
+        )
         self.in_hand_pos.assign(wp.from_torch(in_hand_pos, dtype=wp.vec3f))
 
-        goal_pos = torch.tensor([-0.2, -0.45, 0.68], device=self.device, dtype=torch.float32).repeat((self.num_envs, 1))
+        goal_pos = torch.tensor(self.cfg.goal_marker_position, device=self.device).repeat((self.num_envs, 1))
         self.goal_pos.assign(wp.from_torch(goal_pos, dtype=wp.vec3f))
 
         goal_rot = torch.zeros((self.num_envs, 4), device=self.device, dtype=torch.float32)
@@ -626,6 +648,8 @@ def __init__(self, cfg: AllegroHandEnvCfg, render_mode: str | None = None, **kwa
         # track successes
         self.successes = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device)
         self.consecutive_successes = wp.zeros(1, dtype=wp.float32, device=self.device)
+        self._success_rate_stats = wp.zeros(2, dtype=wp.float32, device=self.device)
+        self.success_rate = wp.zeros(1, dtype=wp.float32, device=self.device)
 
         # Persistent RL buffers (Warp).
         self.actions = wp.zeros((self.num_envs, self.cfg.action_space), dtype=wp.float32, device=self.device)
@@ -689,17 +713,10 @@ def _apply_action(self) -> None:
             device=self.device,
         )
 
-        # Apply position targets using mask method (CUDA graph safe).
-        # All joints are actuated for Allegro, so default masks (None = all) are correct.
-        self.hand.set_joint_position_target_mask(target=self.cur_targets)
+        # unactuated joints keep the targets their last reset wrote
+        self.hand.actuators.target_command.set_position_mask(value=self.cur_targets)
 
     def _get_observations(self) -> dict:
-        # if self.cfg.asymmetric_obs:
-        #    self.fingertip_force_sensors = self.hand.root_physx_view.get_link_incoming_joint_force()[
-        #        :, self.finger_bodies
-        #    ]
-        # NOTE: if re-enabled, `self.finger_bodies` holds public-order indices (from `body_names.index`)
-        # and must be translated to backend view order before indexing this `root_physx_view` array.
         if self.cfg.obs_type == "openai":
             self.compute_reduced_observations()
         elif self.cfg.obs_type == "full":
@@ -710,8 +727,6 @@ def _get_observations(self) -> dict:
 
     def _get_rewards(self) -> None:
         # Clear reduction buffers before launching the reward kernel.
-        # wp.assign(self._num_resets, 0.0)
-        # wp.assign(self._finished_cons_successes, 0.0)
         self._num_resets.zero_()
         self._finished_cons_successes.zero_()
         wp.launch(
@@ -721,13 +736,13 @@ def _get_rewards(self) -> None:
                 self.reset_buf,
                 self.object_pose,
                 self.in_hand_pos,
-                self.goal_rot,
+                self.goal_reached,
+                self.orientation_error,
                 self.cfg.dist_reward_scale,
                 self.cfg.rot_reward_scale,
                 self.cfg.rot_eps,
                 self.actions,
                 self.cfg.action_penalty_scale,
-                self.cfg.success_tolerance,
                 self.cfg.reach_goal_bonus,
                 self.cfg.fall_dist,
                 self.cfg.fall_penalty,
@@ -759,9 +774,7 @@ def _get_rewards(self) -> None:
         # .mean() cannot be called here as it causes problems on stream
         self.extras["log"]["consecutive_successes"] = wp.to_torch(self.consecutive_successes)
 
-        # Reset goals for envs that reached the target (mask is `reset_goal_buf`).
-        # This avoids Torch-side index extraction and keeps the step graphable.
-        self._reset_target_pose(mask=self.reset_goal_buf)
+        self._reset_target_pose(self.reset_goal_buf)
 
     def _get_dones(self) -> None:
         self._compute_intermediate_values()
@@ -779,6 +792,8 @@ def _get_dones(self) -> None:
                 self.cfg.max_consecutive_success,
                 self.successes,
                 self._episode_length_buf_wp,
+                self.orientation_error,
+                self.goal_reached,
                 self.reset_terminated,
                 self.reset_time_outs,
                 self.reset_buf,
@@ -793,10 +808,8 @@ def _reset_idx(self, mask: wp.array | None = None):
         # resets articulation and rigid body attributes
         super()._reset_idx(mask)
 
-        # reset goals
-        self._reset_target_pose(mask=mask)
+        self._reset_target_pose(mask)
 
-        # reset object
         wp.launch(
             reset_object,
             dim=self.num_envs,
@@ -808,13 +821,14 @@ def _reset_idx(self, mask: wp.array | None = None):
                 self.y_unit_vec,
                 mask,
                 self.rng_state,
-                self.object.data.root_link_pose_w.warp,
-                self.object.data.root_com_vel_w.warp,
+                self.reset_object_pose,
+                self.reset_object_vel,
             ],
             device=self.device,
         )
+        self.object.write_root_pose_to_sim_mask(root_pose=self.reset_object_pose, env_mask=mask)
+        self.object.write_root_velocity_to_sim_mask(root_velocity=self.reset_object_vel, env_mask=mask)
 
-        # reset hand
         wp.launch(
             reset_hand,
             dim=self.num_envs,
@@ -828,40 +842,35 @@ def _reset_idx(self, mask: wp.array | None = None):
                 mask,
                 self.num_hand_dofs,
                 self.rng_state,
-                self.hand.data.joint_pos.warp,
-                self.hand.data.joint_vel.warp,
+                self.reset_joint_pos,
+                self.reset_joint_vel,
                 self.prev_targets,
                 self.cur_targets,
-                self.hand_dof_targets,
             ],
             device=self.device,
         )
+        self.hand.actuators.target_command.set_position_mask(value=self.cur_targets, env_mask=mask)
+        self.hand.write_joint_position_to_sim_mask(position=self.reset_joint_pos, env_mask=mask)
+        self.hand.write_joint_velocity_to_sim_mask(velocity=self.reset_joint_vel, env_mask=mask)
 
-        self.hand.set_joint_position_target_mask(target=self.cur_targets, env_mask=mask)
-
+        self._success_rate_stats.zero_()
         wp.launch(
             reset_successes,
             dim=self.num_envs,
-            inputs=[
-                mask,
-                self.successes,
-            ],
+            inputs=[mask, self.successes, self._success_rate_stats],
             device=self.device,
         )
+        wp.launch(
+            update_success_rate_from_stats,
+            dim=1,
+            inputs=[self._success_rate_stats, self.success_rate],
+            device=self.device,
+        )
+        self.extras.setdefault("log", {})["Metrics/success_rate"] = wp.to_torch(self.success_rate)
 
         self._compute_intermediate_values()
 
-    def _reset_target_pose(self, env_ids: Sequence[int] | None = None, mask: wp.array | None = None):
-        # reset goal rotation
-        if mask is None:
-            if env_ids is None:
-                return
-            env_mask_list = [False] * self.num_envs
-            for env_id in env_ids:
-                env_mask_list[int(env_id)] = True
-            mask = wp.array(env_mask_list, dtype=wp.bool, device=self.device)
-
-        # update goal pose and markers
+    def _reset_target_pose(self, mask: wp.array):
         wp.launch(
             reset_target_pose,
             dim=self.num_envs,

From 910a0ba9b90956b900df2de7d8e659d35e0f2e49 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Wed, 30 Sep 2026 03:54:21 -0700
Subject: [PATCH 21/48] Add a Warp twin of the direct Shadow Hand reorientation
 task

Isaac-Reorient-Cube-Shadow-Direct now runs under --frontend warp: the
stable ShadowHandDirectEnv entry point mirrors to ShadowHandDirectWarpEnv,
which extends the shared Warp reorientation environment the way the
stable task extends ReorientDirectEnv. The last four action columns
command the finger tendons through the articulation's tendon target
writer, unscaled to the configured range and clamped without smoothing.

The tendon command is issued in _pre_physics_step rather than in the
captured action stage: the articulation marks a buffered tendon target
for submission on the host, and a graph replay does not repeat that.

Tests extend the same-state parity check to the Shadow task and add a
tendon-response check and an assigned-episode-length partial-reset check.
---
 .../source/_static/css/environment-browser.js |  2 +-
 docs/source/concepts/warp_environments.rst    |  1 +
 .../test/envs/test_frontend_cfg_conversion.py |  1 +
 .../test/envs/test_reorient_warp_env.py       | 52 ++++++++++-
 ...chuanh-task-cleanup-dex-warp-exp.added.rst |  3 +
 .../core/reorient/config/__init__.py          |  6 ++
 .../reorient/config/shadow_hand/__init__.py   | 11 +++
 .../shadow_hand/shadow_hand_warp_env.py       | 86 +++++++++++++++++++
 8 files changed, 160 insertions(+), 2 deletions(-)
 create mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/__init__.py
 create mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/shadow_hand/__init__.py
 create mode 100644 source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/shadow_hand/shadow_hand_warp_env.py

diff --git a/docs/source/_static/css/environment-browser.js b/docs/source/_static/css/environment-browser.js
index 1bd8f03505b5..06a3d6f837f3 100644
--- a/docs/source/_static/css/environment-browser.js
+++ b/docs/source/_static/css/environment-browser.js
@@ -39,7 +39,7 @@
             ["Isaac-RenderBenchmark-Franka-Cabinet", "", "isaacsim_physx,newton_mjwarp,ovphysx", "isaacsim_rtx,newton_renderer,ovrtx", "albedo,depth,rgb,simple_shading_constant_diffuse,simple_shading_diffuse_mdl,simple_shading_full_mdl", "", false, {"*": ["rgb"]}],
             ["Isaac-Reorient-Cube-Allegro-Direct", "rl_games,rsl_rl,skrl", "isaacsim_physx,newton_mjwarp,ovphysx", "", "", "tasks/manipulation/allegro_cube.jpg", true],
             ["Isaac-Reorient-Cube-Allegro", "rl_games,rsl_rl,skrl", "isaacsim_physx,newton_mjwarp,ovphysx", "", "randomized,reset_only", "tasks/manipulation/allegro_cube.jpg", false, {"*": ["reset_only"]}],
-            ["Isaac-Reorient-Cube-Shadow-Direct", "rl_games,rsl_rl,skrl", "isaacsim_physx,newton_mjwarp,ovphysx", "", "", "tasks/manipulation/shadow_cube.jpg"],
+            ["Isaac-Reorient-Cube-Shadow-Direct", "rl_games,rsl_rl,skrl", "isaacsim_physx,newton_mjwarp,ovphysx", "", "", "tasks/manipulation/shadow_cube.jpg", true],
             ["Isaac-Reorient-Cube-Shadow", "rl_games,rsl_rl,skrl", "isaacsim_physx,newton_mjwarp,ovphysx", "", "asymmetric,randomized"],
             ["Isaac-Reorient-Cube-Shadow-Camera-Direct", "rl_games,rsl_rl", "isaacsim_physx,newton_mjwarp,ovphysx", "isaacsim_rtx,newton_renderer,ovrtx", "albedo,depth,full,rgb,semantic_segmentation,simple_shading_constant_diffuse,simple_shading_diffuse_mdl,simple_shading_full_mdl", "tasks/manipulation/shadow_cube.jpg", false, {"*": ["full"]}],
             ["Isaac-Reorient-Cube-Shadow-Camera", "rl_games,rsl_rl", "isaacsim_physx,newton_mjwarp,ovphysx", "isaacsim_rtx,newton_renderer,ovrtx", "albedo,depth,full,randomized,rgb,semantic_segmentation,simple_shading_constant_diffuse,simple_shading_diffuse_mdl,simple_shading_full_mdl", "", false, {"*": ["full"]}],
diff --git a/docs/source/concepts/warp_environments.rst b/docs/source/concepts/warp_environments.rst
index 59e56aeddb1d..78ae444e8f25 100644
--- a/docs/source/concepts/warp_environments.rst
+++ b/docs/source/concepts/warp_environments.rst
@@ -47,6 +47,7 @@ implementation:
 - ``Isaac-Ant-Direct`` — Ant locomotion
 - ``Isaac-Humanoid-Direct`` — Humanoid locomotion
 - ``Isaac-Reorient-Cube-Allegro-Direct`` — Allegro hand cube reorient
+- ``Isaac-Reorient-Cube-Shadow-Direct`` — Shadow hand cube reorient
 
 
 Manager-Based Warp Execution (``--frontend warp``)
diff --git a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py
index 7496d4c7824f..3a0026522462 100644
--- a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py
+++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py
@@ -201,6 +201,7 @@ def test_no_warp_task_registrations_remain():
     "Isaac-Ant-Direct",
     "Isaac-Humanoid-Direct",
     "Isaac-Reorient-Cube-Allegro-Direct",
+    "Isaac-Reorient-Cube-Shadow-Direct",
 ]
 
 
diff --git a/source/isaaclab_experimental/test/envs/test_reorient_warp_env.py b/source/isaaclab_experimental/test/envs/test_reorient_warp_env.py
index 381c462c1d20..85c57918bb03 100644
--- a/source/isaaclab_experimental/test/envs/test_reorient_warp_env.py
+++ b/source/isaaclab_experimental/test/envs/test_reorient_warp_env.py
@@ -27,7 +27,8 @@
 
 pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="The Warp twins capture CUDA graphs.")
 
-_TASKS = ["Isaac-Reorient-Cube-Allegro-Direct"]
+_SHADOW_TASK = "Isaac-Reorient-Cube-Shadow-Direct"
+_TASKS = ["Isaac-Reorient-Cube-Allegro-Direct", _SHADOW_TASK]
 
 
 @pytest.fixture(scope="module", params=_TASKS)
@@ -107,3 +108,52 @@ def test_observations_and_rewards_match_stable_task(env):
         torch.testing.assert_close(reward[untouched], expected[untouched], rtol=1e-4, atol=1e-4)
         num_compared += int(untouched.sum())
     assert num_compared > 0
+
+
+@pytest.mark.parametrize("env", [_SHADOW_TASK], indirect=True)
+def test_tendon_actions_curl_the_coupled_finger_joints(env):
+    """The tendon action columns drive the middle and distal finger joints, step after step.
+
+    Every finger is first opened, then half of the environments curl theirs; a tendon command
+    that never reached the solver, or reached it only once, leaves both halves open.
+    """
+    env.reset()
+    coupled_joints, _ = env.hand.find_joints("rh_(FF|MF|RF|LF)J[12]")
+    num_joint_actions = len(env.cfg.actuated_joint_names)
+    curling = torch.arange(env.num_envs, device=env.device) < env.num_envs // 2
+    actions = torch.zeros((env.num_envs, env.cfg.action_space), device=env.device)
+    actions[:, num_joint_actions:] = -1.0
+    for _ in range(10):
+        env.step(actions)
+    actions[curling, num_joint_actions:] = 1.0
+    was_reset = torch.zeros_like(curling)
+    for _ in range(30):
+        _, _, terminated, truncated, _ = env.step(actions)
+        was_reset |= terminated | truncated
+
+    curl = env.hand.data.joint_pos.torch[:, coupled_joints].sum(dim=-1)
+    assert (curl[curling & ~was_reset].min() - curl[~curling & ~was_reset].max()) > 1.0
+
+
+@pytest.mark.parametrize("env", [_SHADOW_TASK], indirect=True)
+def test_assigned_episode_lengths_time_out_only_their_envs(env):
+    """An episode-length buffer assigned from outside, as RSL-RL does, times out only the envs it ends.
+
+    The masked reset restarts those episodes and resamples their goals, leaving the others' alone.
+    """
+    env.reset()
+    episode_lengths = torch.zeros_like(env.episode_length_buf)
+    episode_lengths[0] = env.max_episode_length - 2
+    env.episode_length_buf = episode_lengths
+    goal_rot = wp.to_torch(env.goal_rot).clone()
+
+    _, _, terminated, truncated, _ = env.step(torch.zeros((env.num_envs, env.cfg.action_space), device=env.device))
+
+    assert truncated[0] and not truncated[1:].any()
+    continuing = ~(terminated | truncated)
+    assert env.episode_length_buf[0] == 0
+    assert (env.episode_length_buf[continuing] == 1).all()
+    new_goal_rot = wp.to_torch(env.goal_rot)
+    assert not torch.equal(new_goal_rot[0], goal_rot[0])
+    unreached = continuing & ~wp.to_torch(env.goal_reached)
+    torch.testing.assert_close(new_goal_rot[unreached], goal_rot[unreached])
diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.added.rst b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.added.rst
index ec9ef0e39417..0b608386bcdd 100644
--- a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.added.rst
+++ b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.added.rst
@@ -1,2 +1,5 @@
+* Added :class:`~isaaclab_tasks_experimental.core.reorient.config.shadow_hand.shadow_hand_warp_env.ShadowHandDirectWarpEnv`,
+  so ``Isaac-Reorient-Cube-Shadow-Direct`` runs under ``--frontend warp``. Like the stable task, it
+  commands the four finger tendons from the last four action columns.
 * Added the ``Metrics/success_rate`` episode metric to the Warp in-hand reorientation environment,
   matching the stable Direct task.
diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/__init__.py
new file mode 100644
index 000000000000..f758624246e0
--- /dev/null
+++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/__init__.py
@@ -0,0 +1,6 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Warp implementations of the hand-specific in-hand reorientation tasks."""
diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/shadow_hand/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/shadow_hand/__init__.py
new file mode 100644
index 000000000000..6abd1f868033
--- /dev/null
+++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/shadow_hand/__init__.py
@@ -0,0 +1,11 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Warp implementation of the stable direct Shadow Hand reorientation task.
+
+``Isaac-Reorient-Cube-Shadow-Direct`` resolves
+:class:`~isaaclab_tasks_experimental.core.reorient.config.shadow_hand.shadow_hand_warp_env.ShadowHandDirectWarpEnv`
+by name from the stable env class; run the stable id with ``--frontend warp`` and ``presets=newton_mjwarp``.
+"""
diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/shadow_hand/shadow_hand_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/shadow_hand/shadow_hand_warp_env.py
new file mode 100644
index 000000000000..349bfda33074
--- /dev/null
+++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/shadow_hand/shadow_hand_warp_env.py
@@ -0,0 +1,86 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import warp as wp
+
+from isaaclab_tasks.core.reorient.utils import resolve_actuated_tendons
+
+from ...reorient_warp_env import ReorientDirectWarpEnv, scale
+
+if TYPE_CHECKING:
+    from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_direct_env_cfg import ShadowHandEnvCfg
+
+
+@wp.kernel
+def apply_actions_to_tendon_targets(
+    # input
+    actions: wp.array2d(dtype=wp.float32),
+    first_tendon_action: wp.int32,
+    lower_limits: wp.array2d(dtype=wp.float32),
+    upper_limits: wp.array2d(dtype=wp.float32),
+    tendon_ids: wp.array(dtype=wp.int32),
+    # output
+    tendon_targets: wp.array2d(dtype=wp.float32),
+):
+    """Scale tendon actions to each tendon's command range and clamp them to it, without smoothing."""
+    env_id, i = wp.tid()
+    lower = lower_limits[env_id, i]
+    upper = upper_limits[env_id, i]
+    target = scale(actions[env_id, first_tendon_action + i], lower, upper)
+    tendon_targets[env_id, tendon_ids[i]] = wp.clamp(target, lower, upper)
+
+
+class ShadowHandDirectWarpEnv(ReorientDirectWarpEnv):
+    """Warp twin of the stable Shadow Hand reorientation environment.
+
+    Implements :class:`~isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_direct_env.ShadowHandDirectEnv`:
+    four of this hand's twenty motors pull a tendon spanning a finger's middle and distal joints.
+    Actions are ordered joints first, then tendons.
+    """
+
+    cfg: ShadowHandEnvCfg
+
+    def __init__(self, cfg: ShadowHandEnvCfg, render_mode: str | None = None, **kwargs):
+        super().__init__(cfg, render_mode, **kwargs)
+        tendon_ids, lower_limits, upper_limits = resolve_actuated_tendons(
+            self.hand, cfg.actuated_tendon_names, self.num_envs, self.device, cfg.actuated_tendon_position_limits
+        )
+        self.num_actuated_tendons = len(tendon_ids)
+        self.actuated_tendon_ids = wp.array(tendon_ids, dtype=wp.int32, device=self.device)
+        self.actuated_tendon_mask = wp.array(
+            [tendon_id in tendon_ids for tendon_id in range(self.hand.num_fixed_tendons)],
+            dtype=wp.bool,
+            device=self.device,
+        )
+        self.tendon_lower_limits = wp.from_torch(lower_limits)
+        self.tendon_upper_limits = wp.from_torch(upper_limits)
+        self.tendon_targets = wp.zeros(
+            (self.num_envs, self.hand.num_fixed_tendons), dtype=wp.float32, device=self.device
+        )
+
+    def _pre_physics_step(self, actions: wp.array) -> None:
+        super()._pre_physics_step(actions)
+        # Commanded here, outside the captured action stage: the articulation flags a buffered
+        # tendon target for submission on the host, which a graph replay would not repeat.
+        wp.launch(
+            apply_actions_to_tendon_targets,
+            dim=(self.num_envs, self.num_actuated_tendons),
+            inputs=[
+                self.actions,
+                self.num_actuated_dofs,
+                self.tendon_lower_limits,
+                self.tendon_upper_limits,
+                self.actuated_tendon_ids,
+                self.tendon_targets,
+            ],
+            device=self.device,
+        )
+        self.hand.set_fixed_tendon_position_target_mask(
+            target=self.tendon_targets, fixed_tendon_mask=self.actuated_tendon_mask
+        )

From 2bb8e9e3068049308e52c71ec6943bd7b9896f69 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Thu, 1 Oct 2026 00:37:32 -0700
Subject: [PATCH 22/48] Resolve adapter reset ids through preallocated masks

---
 .../isaaclab/actuators/newton/adapter.py      | 27 +++++++------------
 .../isaaclab/actuators/newton/kernels.py      | 10 +------
 2 files changed, 11 insertions(+), 26 deletions(-)

diff --git a/source/isaaclab/isaaclab/actuators/newton/adapter.py b/source/isaaclab/isaaclab/actuators/newton/adapter.py
index 961543fa6fda..4790fd16b897 100644
--- a/source/isaaclab/isaaclab/actuators/newton/adapter.py
+++ b/source/isaaclab/isaaclab/actuators/newton/adapter.py
@@ -31,12 +31,8 @@
 from newton.selection import ArticulationView
 
 from ...utils import index_fill_
-from .kernels import (
-    build_implicit_dof_mask,
-    build_per_dof_env_mask_kernel,
-    set_mask_kernel,
-    zero_at_indices_kernel,
-)
+from ...utils.warp.utils import resolve_1d_mask
+from .kernels import build_implicit_dof_mask, build_per_dof_env_mask_kernel, zero_at_indices_kernel
 
 if TYPE_CHECKING:
     from .. import ActuatorCollection
@@ -107,6 +103,8 @@ def __init__(
         self._states_b = [act.state() for act in actuators]
         # Per-DOF reset masks, rebuilt in full on every partial reset so masked resets do not allocate.
         self._reset_dof_masks = [wp.zeros(act.indices.shape[0], dtype=wp.bool, device=device) for act in actuators]
+        self._all_env_mask = wp.ones(num_envs, dtype=wp.bool, device=device)
+        self._scratch_env_mask = wp.zeros(num_envs, dtype=wp.bool, device=device)
 
         # Pre-clamp computed effort buffer. Each Newton actuator scatter-adds
         # its raw controller output to ``sim_control.joint_computed_f`` when
@@ -207,17 +205,12 @@ def reset(
             num_envs = len(range(self._num_envs)[env_ids]) if isinstance(env_ids, slice) else len(env_ids)
             if num_envs == 0:
                 return
-            env_mask = wp.zeros(self._num_envs, dtype=wp.bool, device=self._device)
-            if isinstance(env_ids, slice):
-                wp.to_torch(env_mask)[env_ids] = True
-            else:
-                # Torch indexed scalar assignment uploads True and synchronizes on CUDA.
-                indices = (
-                    wp.from_torch(env_ids)
-                    if isinstance(env_ids, torch.Tensor)
-                    else wp.array(env_ids, dtype=wp.int32, device=self._device)
-                )
-                wp.launch(set_mask_kernel, dim=num_envs, inputs=[env_mask, indices], device=self._device)
+            env_mask = resolve_1d_mask(
+                ids=env_ids,
+                all_mask=self._all_env_mask,
+                scratch_mask=self._scratch_env_mask,
+                device=self._device,
+            )
 
         for act, sa, sb, per_dof_mask in zip(self.actuators, self._states_a, self._states_b, self._reset_dof_masks):
             wp.launch(
diff --git a/source/isaaclab/isaaclab/actuators/newton/kernels.py b/source/isaaclab/isaaclab/actuators/newton/kernels.py
index 5272de15f539..c88da490a00e 100644
--- a/source/isaaclab/isaaclab/actuators/newton/kernels.py
+++ b/source/isaaclab/isaaclab/actuators/newton/kernels.py
@@ -6,7 +6,6 @@
 """Shared Warp kernels for the Newton actuator fast path."""
 
 from collections.abc import Sequence
-from typing import Any
 
 import torch
 import warp as wp
@@ -14,7 +13,7 @@
 from ...utils import index_fill_
 
 # ---------------------------------------------------------------------------
-# Adapter / per-actuator helper kernels: per-DOF zeroing, env-mask building,
+# Adapter / per-actuator helper kernels: per-DOF zeroing
 # and per-DOF env-mask projection (used by :meth:`NewtonActuatorAdapter.reset`).
 # Parameter reads and writes go through Newton's selection API instead.
 # ---------------------------------------------------------------------------
@@ -27,13 +26,6 @@ def zero_at_indices_kernel(data: wp.array(dtype=wp.float32), indices: wp.array(d
     data[indices[i]] = 0.0
 
 
-@wp.kernel(enable_backward=False)
-def set_mask_kernel(mask: wp.array(dtype=wp.bool), indices: wp.array(dtype=Any)):
-    """Set selected mask entries without converting or uploading the indices."""
-    i = wp.tid()
-    mask[indices[i]] = True
-
-
 @wp.kernel(enable_backward=False)
 def build_per_dof_env_mask_kernel(
     indices: wp.array(dtype=wp.uint32),

From edc61c8a5e5453ee623cee0834a13256349bf583 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Thu, 1 Oct 2026 00:46:22 -0700
Subject: [PATCH 23/48] Reset Isaac Lab actuator models from the environment
 mask

---
 .../jichuanh-warp-scene-mask-native.added.rst |  5 +-
 ...ichuanh-warp-scene-mask-native.changed.rst |  6 ++
 .../isaaclab/actuators/actuator_base.py       |  8 ++-
 .../isaaclab/actuators/actuator_collection.py | 29 ++++-----
 .../isaaclab/actuators/actuator_pd.py         | 11 +++-
 .../isaaclab/utils/buffers/delay_buffer.py    | 33 ++++++++---
 .../actuators/test_delayed_pd_actuator.py     | 59 +++++++++++++++++++
 .../isaaclab/test/utils/test_delay_buffer.py  | 35 +++++++++++
 8 files changed, 155 insertions(+), 31 deletions(-)
 create mode 100644 source/isaaclab/test/actuators/test_delayed_pd_actuator.py

diff --git a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst
index 3329702113c2..7cb97ba76bb4 100644
--- a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst
+++ b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst
@@ -1,4 +1,7 @@
 * Added ``env_mask`` to :meth:`~isaaclab.actuators.ActuatorCollection.reset`. Isaac Lab actuator models reset only
-  the masked environments, converting the mask to indices once for all groups, and backend-native actuator state
+  the masked environments from the mask itself, without converting it to indices, and backend-native actuator state
   resets from the mask through :meth:`~isaaclab.actuators.ActuatorControl.reset_native_actuators_mask`.
+* Added boolean-mask selection to :meth:`~isaaclab.utils.buffers.DelayBuffer.set_time_lag` and
+  :meth:`~isaaclab.utils.buffers.DelayBuffer.reset`, which select batches on the device without synchronizing the
+  host. Tensor lags set through a mask are full-sized and are not range-checked.
 * Added ``env_mask`` to :meth:`~isaaclab.assets.BaseCableObject.reset`, matching the other scene assets.
diff --git a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
index ea53a4476e4e..6100151584ba 100644
--- a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
+++ b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
@@ -4,3 +4,9 @@
   the resampled environments, matching :meth:`~isaaclab.sensors.SensorBase.reset`.
 * Changed :meth:`~isaaclab.actuators.newton.NewtonActuatorAdapter.reset` to reuse preallocated per-DOF masks, so
   masked resets of Newton-native actuator state do not allocate and can be captured in a CUDA graph.
+* Changed :meth:`~isaaclab.actuators.ActuatorBase.reset` to also receive a 1-D boolean mask of the environments to
+  reset, which masked resets pass in place of indices. Custom actuator models must accept a mask, for example by
+  resetting their buffers through :func:`~isaaclab.utils.array.index_fill_`.
+* Changed :class:`~isaaclab.actuators.DelayedPDActuator` to raise :class:`ValueError` at construction unless
+  ``0 <= min_delay <= max_delay``, instead of at a reset, because masked resets apply the sampled delays without
+  checking them on the host.
diff --git a/source/isaaclab/isaaclab/actuators/actuator_base.py b/source/isaaclab/isaaclab/actuators/actuator_base.py
index 8b031f6fb534..d7adda51f19b 100644
--- a/source/isaaclab/isaaclab/actuators/actuator_base.py
+++ b/source/isaaclab/isaaclab/actuators/actuator_base.py
@@ -279,11 +279,15 @@ def joint_indices(self) -> slice | torch.Tensor:
     """
 
     @abstractmethod
-    def reset(self, env_ids: Sequence[int]):
+    def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None):
         """Reset the internals within the group.
 
+        Implementations must accept a boolean mask without synchronizing the host, as
+        :func:`~isaaclab.utils.array.index_fill_` does.
+
         Args:
-            env_ids: List of environment IDs to reset.
+            env_ids: Environment indices, a 1-D boolean mask of shape (num_envs,), or a slice selecting the
+                environments to reset. None resets all environments.
         """
         raise NotImplementedError
 
diff --git a/source/isaaclab/isaaclab/actuators/actuator_collection.py b/source/isaaclab/isaaclab/actuators/actuator_collection.py
index c27d2aa5fce4..b6840e6b3be0 100644
--- a/source/isaaclab/isaaclab/actuators/actuator_collection.py
+++ b/source/isaaclab/isaaclab/actuators/actuator_collection.py
@@ -102,13 +102,6 @@ def __init__(
             for actuator_name in self._native_group_names:
                 self._groups[actuator_name] = self._resolve_newton_group_actuators(actuator_name)
         self._validate_coverage()
-        # Groups a masked reset converts the mask to indices for. Plain implicit groups hold no reset
-        # state, and Newton-executed groups reset through the backend.
-        self._mask_reset_actuators = tuple(
-            actuator
-            for actuator in self._groups.values()
-            if isinstance(actuator, ActuatorBase) and type(actuator) is not ImplicitActuator
-        )
         self._build_execution_plan()
         if self._debug_value_resolution:
             self._print_value_resolution_table()
@@ -190,10 +183,9 @@ def applied_effort(self) -> ProxyArray:
     def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: wp.array | None = None) -> None:
         """Reset all actuator group states.
 
-        Isaac Lab actuator models reset by environment index, so a mask is converted to indices
-        once for all of them, and only when a model with reset state is configured; plain
-        :class:`~isaaclab.actuators.ImplicitActuator` groups are skipped. Backend-native actuator
-        state resets from the mask directly.
+        Isaac Lab actuator models receive a mask as a zero-copy boolean tensor, which
+        :meth:`~isaaclab.actuators.ActuatorBase.reset` accepts in place of indices. Backend-native
+        actuator state resets from the mask directly.
 
         .. caution::
             If both ``env_ids`` and ``env_mask`` are provided, ``env_mask`` takes precedence.
@@ -203,18 +195,17 @@ def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: wp.array
             env_mask: Environment mask. Shape is (num_instances,). Defaults to None.
         """
         if env_mask is not None:
-            if self._mask_reset_actuators:
-                group_env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1)
-                for actuator in self._mask_reset_actuators:
-                    actuator.reset(group_env_ids)
-            self._control.reset_native_actuators_mask(env_mask)
-            return
-        group_env_ids = self._control._normalize_index_sequence(env_ids)
+            group_env_ids = wp.to_torch(env_mask)
+        else:
+            group_env_ids = self._control._normalize_index_sequence(env_ids)
         for actuator in self._groups.values():
             # Newton-executed groups are reset through the backend below.
             if isinstance(actuator, ActuatorBase):
                 actuator.reset(group_env_ids)
-        self._control.reset_native_actuators(slice(None) if group_env_ids is None else group_env_ids)
+        if env_mask is not None:
+            self._control.reset_native_actuators_mask(env_mask)
+        else:
+            self._control.reset_native_actuators(slice(None) if group_env_ids is None else group_env_ids)
 
     def compute(self, dt: float = 0.0) -> None:
         """Compute processed actuator commands and telemetry.
diff --git a/source/isaaclab/isaaclab/actuators/actuator_pd.py b/source/isaaclab/isaaclab/actuators/actuator_pd.py
index bc20ea721fb7..1885ccbb7fea 100644
--- a/source/isaaclab/isaaclab/actuators/actuator_pd.py
+++ b/source/isaaclab/isaaclab/actuators/actuator_pd.py
@@ -451,18 +451,25 @@ class DelayedPDActuator(IdealPDActuator):
 
     def __init__(self, cfg: DelayedPDActuatorCfg, *args, **kwargs):
         super().__init__(cfg, *args, **kwargs)
+        # masked resets apply the sampled delays without range-checking them on the host
+        if not 0 <= cfg.min_delay <= cfg.max_delay:
+            raise ValueError(f"Expected 0 <= min_delay <= max_delay, received {cfg.min_delay} and {cfg.max_delay}.")
         # instantiate the delay buffers
         self.positions_delay_buffer = DelayBuffer(cfg.max_delay, self._num_envs, device=self._device)
         self.velocities_delay_buffer = DelayBuffer(cfg.max_delay, self._num_envs, device=self._device)
         self.efforts_delay_buffer = DelayBuffer(cfg.max_delay, self._num_envs, device=self._device)
         self._delay_buffers = (self.positions_delay_buffer, self.velocities_delay_buffer, self.efforts_delay_buffer)
 
-    def reset(self, env_ids: Sequence[int]):
+    def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None):
         super().reset(env_ids)
         # number of environments (since env_ids can be a slice)
         if env_ids is None:
             env_ids = slice(None)
-        num_envs = len(range(self._num_envs)[env_ids]) if isinstance(env_ids, slice) else len(env_ids)
+        if isinstance(env_ids, torch.Tensor) and env_ids.dtype == torch.bool:
+            # a mask samples every environment and the delay buffers keep the masked ones on the device
+            num_envs = self._num_envs
+        else:
+            num_envs = len(range(self._num_envs)[env_ids]) if isinstance(env_ids, slice) else len(env_ids)
         # set a new random delay for environments in env_ids
         time_lags = torch.randint(
             low=self.cfg.min_delay,
diff --git a/source/isaaclab/isaaclab/utils/buffers/delay_buffer.py b/source/isaaclab/isaaclab/utils/buffers/delay_buffer.py
index 17398b58869e..edbc864b3d8d 100644
--- a/source/isaaclab/isaaclab/utils/buffers/delay_buffer.py
+++ b/source/isaaclab/isaaclab/utils/buffers/delay_buffer.py
@@ -122,7 +122,7 @@ def time_lags(self) -> torch.Tensor:
     Operations.
     """
 
-    def set_time_lag(self, time_lag: int | torch.Tensor, batch_ids: Sequence[int] | None = None):
+    def set_time_lag(self, time_lag: int | torch.Tensor, batch_ids: Sequence[int] | torch.Tensor | None = None):
         """Sets the time lag for the delay buffer across the provided batch indices.
 
         Args:
@@ -130,10 +130,12 @@ def set_time_lag(self, time_lag: int | torch.Tensor, batch_ids: Sequence[int] |
 
               * If an integer is provided, the same delay is set for the provided batch indices.
               * If a tensor is provided, the delay is set for each batch index separately. The shape of the tensor
-                should be (len(batch_ids),).
+                should be (len(batch_ids),), or (batch_size,) when ``batch_ids`` is a boolean mask.
 
-            batch_ids: The batch indices for which the time lag is set. Default is None, which sets the time lag
-                for all batch indices.
+            batch_ids: The batch indices, or a boolean mask of shape (batch_size,), for which the time lag is set.
+                Default is None, which sets the time lag for all batch indices. A mask selects the batches on the
+                device without synchronizing the host, so tensor lags are not range-checked and must lie in
+                ``[0, history_length]``.
 
         Raises:
             TypeError: If the type of the :attr:`time_lag` is not int or integer tensor.
@@ -150,6 +152,10 @@ def set_time_lag(self, time_lag: int | torch.Tensor, batch_ids: Sequence[int] |
             # check valid dtype for time_lag: must be int or long
             if time_lag.dtype not in [torch.int, torch.long]:
                 raise TypeError(f"Invalid dtype for time_lag: {time_lag.dtype}. Expected torch.int or torch.long.")
+            if isinstance(batch_ids, torch.Tensor) and batch_ids.dtype == torch.bool:
+                time_lag = time_lag.to(device=self.device, dtype=self._time_lags.dtype)
+                torch.where(batch_ids, time_lag, self._time_lags, out=self._time_lags)
+                return
             min_time_lag = int(time_lag.min().item()) if time_lag.numel() else 0
             max_time_lag = int(time_lag.max().item()) if time_lag.numel() else 0
         else:
@@ -165,18 +171,31 @@ def set_time_lag(self, time_lag: int | torch.Tensor, batch_ids: Sequence[int] |
         else:
             index_fill_(self._time_lags, batch_ids, time_lag)
 
-    def reset(self, batch_ids: Sequence[int] | None = None):
+    def reset(self, batch_ids: Sequence[int] | torch.Tensor | None = None):
         """Reset the data in the delay buffer at the specified batch indices.
 
         Automatically sampled lags are redrawn for those batches regardless of ``hold_prob``.
         Externally selected lags are preserved.
 
         Args:
-            batch_ids: Elements to reset in the batch dimension. Default is None, which resets all the batch indices.
+            batch_ids: Elements to reset in the batch dimension, as indices or a boolean mask of shape (batch_size,).
+                Default is None, which resets all the batch indices.
         """
         indices = slice(None) if batch_ids is None else batch_ids
         index_fill_(self._num_pushes, indices, 0)
-        if self._hold_prob is not None:
+        if self._hold_prob is None:
+            return
+        if isinstance(indices, torch.Tensor) and indices.dtype == torch.bool:
+            # Draw for every batch and keep the masked ones, so the selection stays on the device.
+            lags = torch.randint(
+                self._min_lag,
+                self.history_length + 1,
+                (self.batch_size,),
+                dtype=self._time_lags.dtype,
+                device=self.device,
+            )
+            torch.where(indices, lags, self._time_lags, out=self._time_lags)
+        else:
             self._time_lags[indices] = torch.randint(
                 self._min_lag,
                 self.history_length + 1,
diff --git a/source/isaaclab/test/actuators/test_delayed_pd_actuator.py b/source/isaaclab/test/actuators/test_delayed_pd_actuator.py
new file mode 100644
index 000000000000..d3ffb048cac2
--- /dev/null
+++ b/source/isaaclab/test/actuators/test_delayed_pd_actuator.py
@@ -0,0 +1,59 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+import pytest
+import torch
+
+from isaaclab.actuators import DelayedPDActuatorCfg
+from isaaclab.test.utils import DeviceScope, test_devices
+from isaaclab.utils import instantiate
+from isaaclab.utils.types import ArticulationActions
+
+pytestmark = pytest.mark.unit
+
+NUM_ENVS, NUM_JOINTS = 4, 2
+
+
+@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA))
+def test_masked_reset_matches_index_reset(device):
+    """A boolean-mask reset delays the same environments as the equivalent index reset, without synchronizing CUDA."""
+    joint_names = [f"joint_{j}" for j in range(NUM_JOINTS)]
+    cfg = DelayedPDActuatorCfg(joint_names_expr=joint_names, stiffness=10.0, damping=1.0, min_delay=2, max_delay=2)
+    by_ids, by_mask = (
+        instantiate(cfg, joint_names=joint_names, joint_ids=slice(None), num_envs=NUM_ENVS, device=device)
+        for _ in range(2)
+    )
+    env_ids = torch.tensor([1, 2], device=device)
+    env_mask = torch.zeros(NUM_ENVS, dtype=torch.bool, device=device)
+    env_mask[env_ids] = True
+    zeros = torch.zeros(NUM_ENVS, NUM_JOINTS, device=device)
+
+    def step(actuator, target: float) -> torch.Tensor:
+        action = ArticulationActions(joint_positions=zeros + target, joint_velocities=zeros, joint_efforts=zeros)
+        actuator.compute(action, joint_pos=zeros, joint_vel=zeros)
+        return actuator.computed_effort.clone()
+
+    # A lag of one in every environment differs from the configured lag the reset samples.
+    for actuator in (by_ids, by_mask):
+        for delay_buffer in (
+            actuator.positions_delay_buffer,
+            actuator.velocities_delay_buffer,
+            actuator.efforts_delay_buffer,
+        ):
+            delay_buffer.set_time_lag(1)
+        for target in range(3):
+            step(actuator, float(target))
+
+    by_ids.reset(env_ids)
+    torch.cuda.synchronize(device)
+    previous = torch.cuda.get_sync_debug_mode()
+    torch.cuda.set_sync_debug_mode("error")
+    try:
+        by_mask.reset(env_mask)
+    finally:
+        torch.cuda.set_sync_debug_mode(previous)
+
+    for target in range(3, 7):
+        torch.testing.assert_close(step(by_mask, float(target)), step(by_ids, float(target)))
diff --git a/source/isaaclab/test/utils/test_delay_buffer.py b/source/isaaclab/test/utils/test_delay_buffer.py
index d2e46d19d13d..84ffb9a365cd 100644
--- a/source/isaaclab/test/utils/test_delay_buffer.py
+++ b/source/isaaclab/test/utils/test_delay_buffer.py
@@ -101,6 +101,41 @@ def test_random_time_lags(delay_buffer, hold_prob, monkeypatch):
         torch.testing.assert_close(delay_buffer.compute(data + 100, update_history=False), expected)
 
 
+@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA))
+def test_mask_selects_like_indices(device, monkeypatch):
+    """A boolean mask sets lags and resets histories of the same batches as indices, without synchronizing CUDA."""
+    batch_ids = [1, 3]
+    by_ids, by_mask = (DelayBuffer(4, 10, device, hold_prob=1.0) for _ in range(2))
+    mask = torch.zeros(by_mask.batch_size, dtype=torch.bool, device=device)
+    mask[batch_ids] = True
+    # Masked lags are full-sized; unmasked entries differ from the current lags so applying them shows.
+    masked_lags = torch.ones(by_mask.batch_size, dtype=torch.long, device=device)
+    masked_lags[batch_ids] = torch.tensor([4, 0], device=device)
+    monkeypatch.setattr(torch, "randint", lambda low, high, size, **kwargs: torch.full(size, 2, **kwargs))
+    for delay_buffer in (by_ids, by_mask):
+        delay_buffer.set_time_lag(torch.arange(delay_buffer.batch_size, device=device) % 5)
+        for data in _generate_data(delay_buffer.batch_size, 3, device):
+            delay_buffer.compute(data)
+
+    def without_sync(call):
+        torch.cuda.synchronize(device)
+        previous = torch.cuda.get_sync_debug_mode()
+        torch.cuda.set_sync_debug_mode("error")
+        try:
+            call()
+        finally:
+            torch.cuda.set_sync_debug_mode(previous)
+
+    by_ids.set_time_lag(masked_lags[batch_ids], batch_ids)
+    without_sync(lambda: by_mask.set_time_lag(masked_lags, mask))
+    torch.testing.assert_close(by_mask.time_lags, by_ids.time_lags)
+    by_ids.reset(batch_ids)
+    without_sync(lambda: by_mask.reset(mask))
+    torch.testing.assert_close(by_mask.time_lags, by_ids.time_lags)
+    for data in _generate_data(by_mask.batch_size, 4, device):
+        torch.testing.assert_close(by_mask.compute(data), by_ids.compute(data))
+
+
 @pytest.mark.parametrize(
     ("time_lag", "batch_ids"),
     [

From e89284aa1acb4ca32c50e8767bcc6c1675b5fa42 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Thu, 1 Oct 2026 00:46:59 -0700
Subject: [PATCH 24/48] Reset surface grippers by index in the Warp scene

---
 .../changelog.d/jichuanh-warp-scene-mask-native.fixed.rst    | 2 +-
 .../isaaclab_experimental/envs/interactive_scene_warp.py     | 5 +----
 2 files changed, 2 insertions(+), 5 deletions(-)

diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
index 84ea86a8e137..7154f9e9454a 100644
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
@@ -1,5 +1,5 @@
 * Fixed :meth:`~isaaclab_experimental.envs.InteractiveSceneWarp.reset` with ``env_mask`` resetting deformable
-  objects, surface grippers, and sensors in every environment, and skipping cable objects.
+  objects and sensors in every environment, and skipping cable objects.
 * Fixed the Warp :class:`~isaaclab_experimental.managers.ObservationManager` rejecting every observation group with
   ``Configuration for the term 'history_order' is not of type ObservationTermCfg``, which stopped all manager-based
   Warp environments from constructing. ``history_order`` is now read as a group setting, as in the stable manager.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py
index 43a5047426ab..5fa1a2eda227 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py
@@ -41,10 +41,7 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None
         for rigid_object in self._rigid_objects.values():
             rigid_object.reset(env_ids, env_mask=env_mask)
         for surface_gripper in self._surface_grippers.values():
-            if env_mask is None:
-                surface_gripper.reset(env_ids)
-            else:
-                surface_gripper.reset_mask(wp.to_torch(env_mask))
+            surface_gripper.reset(env_ids)
         for rigid_object_collection in self._rigid_object_collections.values():
             rigid_object_collection.reset(env_ids, env_mask=env_mask)
         # -- sensors

From 8b2a16165ab301d85656e3f3e9c8675c1f58b0bc Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Thu, 1 Oct 2026 00:48:22 -0700
Subject: [PATCH 25/48] Drop float from the fixed-tendon position limit type

---
 .../jichuanh-warp-scene-mask-native.changed.rst        |  4 ++++
 .../isaaclab/assets/articulation/base_articulation.py  | 10 ++++++----
 .../assets/articulation/articulation.py                |  4 ++--
 .../changelog.d/jichuanh-warp-scene-mask-native.skip   |  0
 .../isaaclab_physx/assets/articulation/articulation.py |  4 ++--
 5 files changed, 14 insertions(+), 8 deletions(-)
 create mode 100644 source/isaaclab_physx/changelog.d/jichuanh-warp-scene-mask-native.skip

diff --git a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
index 6100151584ba..fe555c95ab08 100644
--- a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
+++ b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
@@ -10,3 +10,7 @@
 * Changed :class:`~isaaclab.actuators.DelayedPDActuator` to raise :class:`ValueError` at construction unless
   ``0 <= min_delay <= max_delay``, instead of at a reset, because masked resets apply the sampled delays without
   checking them on the host.
+* Changed the ``limit`` annotation of
+  :meth:`~isaaclab.assets.BaseArticulation.set_fixed_tendon_position_limit_index` and
+  :meth:`~isaaclab.assets.BaseArticulation.set_fixed_tendon_position_limit_mask` to exclude ``float``, matching the
+  backends, which raise :class:`ValueError` for a float because it cannot hold both bounds.
diff --git a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py
index d0f34188e445..f205e736b96f 100644
--- a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py
+++ b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py
@@ -2089,7 +2089,7 @@ def set_fixed_tendon_limit_stiffness_mask(
     def set_fixed_tendon_position_limit_index(
         self,
         *,
-        limit: float | torch.Tensor | wp.array,
+        limit: torch.Tensor | wp.array,
         fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None,
         env_ids: Sequence[int] | torch.Tensor | wp.array | None = None,
     ) -> None:
@@ -2107,7 +2107,8 @@ def set_fixed_tendon_position_limit_index(
             Some backends may provide optimized implementations for masks / indices.
 
         Args:
-            limit: Fixed tendon limit. Shape is (len(env_ids), len(fixed_tendon_ids)).
+            limit: Fixed tendon position limits ``[lower, upper]`` [m]. Shape is (len(env_ids), len(fixed_tendon_ids))
+                with dtype wp.vec2f. A torch tensor has a trailing dimension of 2.
             fixed_tendon_ids: The tendon indices to set the limit for. Defaults to None (all fixed tendons).
             env_ids: The environment indices to set the limit for. Defaults to None (all instances).
         """
@@ -2117,7 +2118,7 @@ def set_fixed_tendon_position_limit_index(
     def set_fixed_tendon_position_limit_mask(
         self,
         *,
-        limit: float | torch.Tensor | wp.array,
+        limit: torch.Tensor | wp.array,
         fixed_tendon_mask: wp.array | None = None,
         env_mask: wp.array | None = None,
     ) -> None:
@@ -2135,7 +2136,8 @@ def set_fixed_tendon_position_limit_mask(
             Some backends may provide optimized implementations for masks / indices.
 
         Args:
-            limit: Fixed tendon limit. Shape is (num_instances, num_fixed_tendons).
+            limit: Fixed tendon position limits ``[lower, upper]`` [m]. Shape is (num_instances, num_fixed_tendons)
+                with dtype wp.vec2f. A torch tensor has a trailing dimension of 2.
             fixed_tendon_mask: Fixed tendon mask. If None, then all the fixed tendons are updated.
                 Shape is (num_fixed_tendons,).
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
index a0b3c003c6fe..913a41bd54c2 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
@@ -2851,7 +2851,7 @@ def set_fixed_tendon_limit_stiffness_mask(
     def set_fixed_tendon_position_limit_index(
         self,
         *,
-        limit: float | torch.Tensor | wp.array,
+        limit: torch.Tensor | wp.array,
         fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None,
         env_ids: Sequence[int] | torch.Tensor | wp.array | None = None,
     ) -> None:
@@ -2889,7 +2889,7 @@ def set_fixed_tendon_position_limit_index(
     def set_fixed_tendon_position_limit_mask(
         self,
         *,
-        limit: float | torch.Tensor | wp.array,
+        limit: torch.Tensor | wp.array,
         fixed_tendon_mask: wp.array | None = None,
         env_mask: wp.array | None = None,
     ) -> None:
diff --git a/source/isaaclab_physx/changelog.d/jichuanh-warp-scene-mask-native.skip b/source/isaaclab_physx/changelog.d/jichuanh-warp-scene-mask-native.skip
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py
index c023134d3a95..86b05a3e4181 100644
--- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py
+++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py
@@ -3018,7 +3018,7 @@ def set_fixed_tendon_limit_stiffness_mask(
     def set_fixed_tendon_position_limit_index(
         self,
         *,
-        limit: float | torch.Tensor | wp.array,
+        limit: torch.Tensor | wp.array,
         fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None,
         env_ids: Sequence[int] | torch.Tensor | wp.array | None = None,
         full_data: bool = False,
@@ -3062,7 +3062,7 @@ def set_fixed_tendon_position_limit_index(
     def set_fixed_tendon_position_limit_mask(
         self,
         *,
-        limit: float | torch.Tensor | wp.array,
+        limit: torch.Tensor | wp.array,
         fixed_tendon_mask: wp.array | None = None,
         env_mask: wp.array | None = None,
     ) -> None:

From 4fb4e0145c594085f4e4125df57058003f506913 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Thu, 1 Oct 2026 00:48:56 -0700
Subject: [PATCH 26/48] Rename the fixed-tendon mask helper and move it beside
 its setters

---
 .../assets/articulation/articulation.py       | 84 +++++++++----------
 1 file changed, 42 insertions(+), 42 deletions(-)

diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
index 913a41bd54c2..4809a2ac6a38 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
@@ -2622,6 +2622,45 @@ def set_inertias_mask(
     Operations - Tendons.
     """
 
+    def _set_fixed_tendon_property_mask(
+        self,
+        data: float | torch.Tensor | wp.array,
+        buffer: wp.array,
+        fixed_tendon_mask: wp.array | torch.Tensor | None,
+        env_mask: wp.array | torch.Tensor | None,
+        name: str,
+    ) -> None:
+        """Write full (num_instances, num_fixed_tendons) data into the masked cells of a fixed tendon buffer.
+
+        Args:
+            data: Scalar or full data. A torch tensor for a ``wp.vec2f`` buffer has a trailing dimension of 2.
+            buffer: Fixed tendon staging buffer. Shape is (num_instances, num_fixed_tendons).
+            fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used.
+            env_mask: Environment mask. If None, then all the instances are updated.
+            name: Name of the data, used in shape errors.
+        """
+        env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK)
+        fixed_tendon_mask = self._resolve_mask(fixed_tendon_mask, self._ALL_FIXED_TENDON_MASK)
+        if isinstance(data, float):
+            wp.launch(
+                articulation_kernels.float_data_to_buffer_with_mask,
+                dim=buffer.shape,
+                inputs=[data, env_mask, fixed_tendon_mask],
+                outputs=[buffer],
+                device=self.device,
+            )
+            return
+        if isinstance(data, torch.Tensor):
+            data = wp.from_torch(data, dtype=buffer.dtype)
+        self.assert_shape_and_dtype(data, buffer.shape, buffer.dtype, name)
+        wp.launch(
+            articulation_kernels.copy_2d_data_with_mask,
+            dim=buffer.shape,
+            inputs=[data, env_mask, fixed_tendon_mask],
+            outputs=[buffer],
+            device=self.device,
+        )
+
     def set_fixed_tendon_stiffness_index(
         self,
         *,
@@ -2709,7 +2748,7 @@ def set_fixed_tendon_stiffness_mask(
             fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used.
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
-        self._set_fixed_tendon_buffer_mask(
+        self._set_fixed_tendon_property_mask(
             stiffness, self.data._fixed_tendon_stiffness, fixed_tendon_mask, env_mask, "stiffness"
         )
 
@@ -2801,7 +2840,7 @@ def set_fixed_tendon_damping_mask(
             fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used.
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
-        self._set_fixed_tendon_buffer_mask(
+        self._set_fixed_tendon_property_mask(
             damping, self.data._fixed_tendon_damping, fixed_tendon_mask, env_mask, "damping"
         )
 
@@ -2916,7 +2955,7 @@ def set_fixed_tendon_position_limit_mask(
         """
         if isinstance(limit, float):
             raise ValueError("Fixed tendon position limits must be a tensor or array, not a float.")
-        self._set_fixed_tendon_buffer_mask(
+        self._set_fixed_tendon_property_mask(
             limit, self.data._fixed_tendon_pos_limits, fixed_tendon_mask, env_mask, "limit"
         )
 
@@ -3810,45 +3849,6 @@ def _to_torch_ids(ids: wp.array | torch.Tensor) -> torch.Tensor:
         """View resolved indices as a torch tensor for indexing."""
         return (wp.to_torch(ids) if isinstance(ids, wp.array) else ids).long()
 
-    def _set_fixed_tendon_buffer_mask(
-        self,
-        data: float | torch.Tensor | wp.array,
-        buffer: wp.array,
-        fixed_tendon_mask: wp.array | torch.Tensor | None,
-        env_mask: wp.array | torch.Tensor | None,
-        name: str,
-    ) -> None:
-        """Write full (num_instances, num_fixed_tendons) data into the masked cells of a fixed tendon buffer.
-
-        Args:
-            data: Scalar or full data. A torch tensor for a ``wp.vec2f`` buffer has a trailing dimension of 2.
-            buffer: Fixed tendon staging buffer. Shape is (num_instances, num_fixed_tendons).
-            fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used.
-            env_mask: Environment mask. If None, then all the instances are updated.
-            name: Name of the data, used in shape errors.
-        """
-        env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK)
-        fixed_tendon_mask = self._resolve_mask(fixed_tendon_mask, self._ALL_FIXED_TENDON_MASK)
-        if isinstance(data, float):
-            wp.launch(
-                articulation_kernels.float_data_to_buffer_with_mask,
-                dim=buffer.shape,
-                inputs=[data, env_mask, fixed_tendon_mask],
-                outputs=[buffer],
-                device=self.device,
-            )
-            return
-        if isinstance(data, torch.Tensor):
-            data = wp.from_torch(data, dtype=buffer.dtype)
-        self.assert_shape_and_dtype(data, buffer.shape, buffer.dtype, name)
-        wp.launch(
-            articulation_kernels.copy_2d_data_with_mask,
-            dim=buffer.shape,
-            inputs=[data, env_mask, fixed_tendon_mask],
-            outputs=[buffer],
-            device=self.device,
-        )
-
     def _resolve_mask(self, mask: wp.array | torch.Tensor | None, full_mask: wp.array) -> wp.array:
         """Resolve a mask to a warp array.
 

From 3c5e2de914cd8548a35f3a018ea9d179a144e61c Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Thu, 1 Oct 2026 00:54:23 -0700
Subject: [PATCH 27/48] Reuse the shared masked-copy kernel for tendon and
 link-pose writes

---
 .../assets/articulation/articulation.py       |  4 +-
 .../assets/articulation/kernels.py            | 20 ----------
 .../isaaclab_newton/assets/kernels.py         | 39 ++++---------------
 .../rigid_object_collection.py                |  6 ++-
 .../physics/mjwarp_tendon_control.py          |  3 ++
 5 files changed, 16 insertions(+), 56 deletions(-)

diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
index 4809a2ac6a38..481ff0dbd85c 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
@@ -2654,7 +2654,7 @@ def _set_fixed_tendon_property_mask(
             data = wp.from_torch(data, dtype=buffer.dtype)
         self.assert_shape_and_dtype(data, buffer.shape, buffer.dtype, name)
         wp.launch(
-            articulation_kernels.copy_2d_data_with_mask,
+            shared_kernels.write_2d_data_to_buffer_with_mask,
             dim=buffer.shape,
             inputs=[data, env_mask, fixed_tendon_mask],
             outputs=[buffer],
@@ -3163,7 +3163,7 @@ def write_fixed_tendon_properties_to_sim_mask(
             (self.data._fixed_tendon_pos_limits, self.data._sim_bind_fixed_tendon_pos_limits),
         ):
             wp.launch(
-                articulation_kernels.copy_2d_data_with_mask,
+                shared_kernels.write_2d_data_to_buffer_with_mask,
                 dim=staged.shape,
                 inputs=[staged, env_mask, fixed_tendon_mask],
                 outputs=[sim_bind],
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py
index f423c2558138..a530cbe6df2a 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py
@@ -542,26 +542,6 @@ def float_data_to_buffer_with_mask(
         out_data[i, j] = in_data
 
 
-@wp.kernel
-def copy_2d_data_with_mask(
-    in_data: wp.array2d(dtype=Any),
-    env_mask: wp.array(dtype=wp.bool),
-    item_mask: wp.array(dtype=wp.bool),
-    out_data: wp.array2d(dtype=Any),
-):
-    """Copy the masked cells of a 2D per-environment buffer of any element type.
-
-    Args:
-        in_data: Input array. Shape is (num_envs, num_items).
-        env_mask: Input array of environment mask. Shape is (num_envs,).
-        item_mask: Input array of item mask, e.g. fixed tendons. Shape is (num_items,).
-        out_data: Output array where the selected cells are written. Shape is (num_envs, num_items).
-    """
-    i, j = wp.tid()
-    if env_mask[i] and item_mask[j]:
-        out_data[i, j] = in_data[i, j]
-
-
 @wp.kernel
 def update_soft_joint_pos_limits(
     joint_pos_limits: wp.array2d(dtype=wp.vec2f),
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/kernels.py b/source/isaaclab_newton/isaaclab_newton/assets/kernels.py
index 6b9d0ed497fe..685de41b7799 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/kernels.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/kernels.py
@@ -992,30 +992,6 @@ def set_body_link_velocity_to_sim_kernel(
     return _SET_BODY_LINK_VELOCITY_TO_SIM_DISPATCHER.select(env_ids, body_ids)
 
 
-@wp.kernel
-def set_body_link_pose_to_sim_mask(
-    data: wp.array2d(dtype=wp.transformf),
-    env_mask: wp.array(dtype=wp.bool),
-    body_mask: wp.array(dtype=wp.bool),
-    body_link_pose_w: wp.array2d(dtype=wp.transformf),
-):
-    """Write masked body link pose data to simulation buffers.
-
-    Launched over the full ``(num_envs, num_bodies)`` grid with an in-kernel mask test, so the
-    write needs no host-side index compaction.
-
-    Args:
-        data: Input array of body link poses. Shape is (num_envs, num_bodies).
-        env_mask: Boolean environment selection. Shape is (num_envs,).
-        body_mask: Boolean body selection. Shape is (num_bodies,).
-        body_link_pose_w: Output array where body link poses are written.
-            Shape is (num_envs, num_bodies).
-    """
-    i, j = wp.tid()
-    if env_mask[i] and body_mask[j]:
-        body_link_pose_w[i, j] = data[i, j]
-
-
 @wp.kernel
 def set_body_com_pose_to_sim_mask(
     data: wp.array2d(dtype=wp.transformf),
@@ -1166,21 +1142,20 @@ def write_2d_data_to_buffer_with_indices_kernel(
 
 @wp.kernel
 def write_2d_data_to_buffer_with_mask(
-    in_data: wp.array2d(dtype=wp.float32),
+    in_data: wp.array2d(dtype=Any),
     env_mask: wp.array(dtype=wp.bool),
     joint_mask: wp.array(dtype=wp.bool),
-    out_data: wp.array2d(dtype=wp.float32),
+    out_data: wp.array2d(dtype=Any),
 ):
-    """Write 2D float data to a buffer at specified indices.
+    """Copy the masked cells of a 2D per-environment buffer of any element type.
 
-    This kernel copies float data from an input array to an output buffer at the specified
-    environment and joint/body indices.
+    The arguments must be Warp arrays: a generic kernel cannot infer the element type of a torch tensor.
 
     Args:
-        in_data: Input array containing float data. Shape is (num_instances, num_joints).
+        in_data: Input array. Shape is (num_instances, num_items).
         env_mask: Input array of environment mask. Shape is (num_instances,).
-        joint_mask: Input array of joint/body mask. Shape is (num_instances, num_joints).
-        out_data: Output array where data is written. Shape is (num_instances, num_joints).
+        joint_mask: Input array of item mask, e.g. joints, bodies, or fixed tendons. Shape is (num_items,).
+        out_data: Output array where the selected cells are written. Shape is (num_instances, num_items).
     """
     i, j = wp.tid()
     if env_mask[i] and joint_mask[j]:
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py
index 41537f57e7b5..7ea3fb417aee 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py
@@ -479,9 +479,11 @@ def write_body_link_pose_to_sim_mask(
             env_mask = self._ALL_ENV_MASK
         body_mask = self._resolve_body_selection_mask(body_ids, body_mask)
         self.assert_shape_and_dtype(body_poses, (self.num_instances, self.num_bodies), wp.transformf, "body_poses")
+        if isinstance(body_poses, torch.Tensor):
+            body_poses = wp.from_torch(body_poses, dtype=wp.transformf)
         # Write to consolidated buffer
         wp.launch(
-            shared_kernels.set_body_link_pose_to_sim_mask,
+            shared_kernels.write_2d_data_to_buffer_with_mask,
             dim=(self.num_instances, self.num_bodies),
             inputs=[
                 body_poses,
@@ -489,7 +491,7 @@ def write_body_link_pose_to_sim_mask(
                 body_mask,
             ],
             outputs=[
-                self.data.body_link_pose_w,
+                self.data.body_link_pose_w.warp,
             ],
             device=self.device,
         )
diff --git a/source/isaaclab_newton/isaaclab_newton/physics/mjwarp_tendon_control.py b/source/isaaclab_newton/isaaclab_newton/physics/mjwarp_tendon_control.py
index 0d8245b5e158..504d20ce5d0d 100644
--- a/source/isaaclab_newton/isaaclab_newton/physics/mjwarp_tendon_control.py
+++ b/source/isaaclab_newton/isaaclab_newton/physics/mjwarp_tendon_control.py
@@ -163,6 +163,9 @@ def set_position_target_mask(
         articulation = self._articulation
         buffer = articulation.data._fixed_tendon_position_target
         articulation.assert_shape_and_dtype(target, buffer.shape, wp.float32, "target")
+        # the masked-copy kernel is generic, so it cannot take a torch tensor directly
+        if isinstance(target, torch.Tensor):
+            target = wp.from_torch(target, dtype=wp.float32)
         env_mask = articulation._resolve_mask(env_mask, articulation._ALL_ENV_MASK)
         fixed_tendon_mask = articulation._resolve_mask(fixed_tendon_mask, articulation._ALL_FIXED_TENDON_MASK)
         wp.launch(

From ea94e2008b3752acb8af56d7eeab168b2fc959bf Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Thu, 1 Oct 2026 00:55:06 -0700
Subject: [PATCH 28/48] Rename the collection body-mask resolver to
 _resolve_body_mask

---
 .../rigid_object_collection/rigid_object_collection.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py
index 7ea3fb417aee..2298424be7f5 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py
@@ -477,7 +477,7 @@ def write_body_link_pose_to_sim_mask(
         """
         if env_mask is None:
             env_mask = self._ALL_ENV_MASK
-        body_mask = self._resolve_body_selection_mask(body_ids, body_mask)
+        body_mask = self._resolve_body_mask(body_ids, body_mask)
         self.assert_shape_and_dtype(body_poses, (self.num_instances, self.num_bodies), wp.transformf, "body_poses")
         if isinstance(body_poses, torch.Tensor):
             body_poses = wp.from_torch(body_poses, dtype=wp.transformf)
@@ -595,7 +595,7 @@ def write_body_com_pose_to_sim_mask(
         """
         if env_mask is None:
             env_mask = self._ALL_ENV_MASK
-        body_mask = self._resolve_body_selection_mask(body_ids, body_mask)
+        body_mask = self._resolve_body_mask(body_ids, body_mask)
         self.assert_shape_and_dtype(body_poses, (self.num_instances, self.num_bodies), wp.transformf, "body_poses")
         # Write to consolidated buffers (updates both com_pose_w and link_pose_w)
         wp.launch(
@@ -721,7 +721,7 @@ def write_body_com_velocity_to_sim_mask(
         """
         if env_mask is None:
             env_mask = self._ALL_ENV_MASK
-        body_mask = self._resolve_body_selection_mask(body_ids, body_mask)
+        body_mask = self._resolve_body_mask(body_ids, body_mask)
         self.assert_shape_and_dtype(
             body_velocities, (self.num_instances, self.num_bodies), wp.spatial_vectorf, "body_velocities"
         )
@@ -850,7 +850,7 @@ def write_body_link_velocity_to_sim_mask(
         """
         if env_mask is None:
             env_mask = self._ALL_ENV_MASK
-        body_mask = self._resolve_body_selection_mask(body_ids, body_mask)
+        body_mask = self._resolve_body_mask(body_ids, body_mask)
         self.assert_shape_and_dtype(
             body_velocities, (self.num_instances, self.num_bodies), wp.spatial_vectorf, "body_velocities"
         )
@@ -1350,7 +1350,7 @@ def _resolve_body_ids(self, body_ids) -> wp.array | torch.Tensor:
             )
         return body_ids
 
-    def _resolve_body_selection_mask(
+    def _resolve_body_mask(
         self,
         body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None,
         body_mask: wp.array | None,

From 3f26cfa533a16284e8a86cf8beb119eec1087f5b Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Thu, 1 Oct 2026 11:47:31 -0700
Subject: [PATCH 29/48] Pass masks explicitly to actuator models, native
 actuators and delay buffers

ActuatorBase.reset, ActuatorControl.reset_native_actuators and DelayBuffer take an explicit env_mask or batch_mask, replacing reset_native_actuators_mask. DelayedPD validates its delay range at construction because the masked path skips the host-side check.
---
 .../jichuanh-warp-scene-mask-native.added.rst |  8 +--
 ...ichuanh-warp-scene-mask-native.changed.rst |  7 +--
 ... => jichuanh-warp-scene-mask-native.major} |  0
 .../isaaclab/actuators/actuator_base.py       | 12 +++--
 .../isaaclab/actuators/actuator_collection.py | 19 +++----
 .../isaaclab/actuators/actuator_control.py    | 18 ++-----
 .../isaaclab/actuators/actuator_net.py        | 14 ++---
 .../isaaclab/actuators/actuator_pd.py         | 14 ++---
 .../isaaclab/actuators/newton/adapter.py      |  4 +-
 .../isaaclab/utils/buffers/delay_buffer.py    | 51 ++++++++++++-------
 .../actuators/test_delayed_pd_actuator.py     |  2 +-
 .../isaaclab/test/utils/test_delay_buffer.py  |  4 +-
 .../assets/articulation/actuator_control.py   |  8 +--
 .../assets/articulation/actuator_control.py   |  8 +--
 .../jichuanh-warp-scene-mask-native.added.rst |  3 ++
 .../jichuanh-warp-scene-mask-native.skip      |  0
 .../assets/articulation/actuator_control.py   |  6 ++-
 17 files changed, 91 insertions(+), 87 deletions(-)
 rename source/isaaclab/changelog.d/{jichuanh-warp-scene-mask-native.minor => jichuanh-warp-scene-mask-native.major} (100%)
 create mode 100644 source/isaaclab_physx/changelog.d/jichuanh-warp-scene-mask-native.added.rst
 delete mode 100644 source/isaaclab_physx/changelog.d/jichuanh-warp-scene-mask-native.skip

diff --git a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst
index 7cb97ba76bb4..37303ca182c2 100644
--- a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst
+++ b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst
@@ -1,7 +1,7 @@
-* Added ``env_mask`` to :meth:`~isaaclab.actuators.ActuatorCollection.reset`. Isaac Lab actuator models reset only
-  the masked environments from the mask itself, without converting it to indices, and backend-native actuator state
-  resets from the mask through :meth:`~isaaclab.actuators.ActuatorControl.reset_native_actuators_mask`.
-* Added boolean-mask selection to :meth:`~isaaclab.utils.buffers.DelayBuffer.set_time_lag` and
+* Added ``env_mask`` to :meth:`~isaaclab.actuators.ActuatorCollection.reset` and
+  :meth:`~isaaclab.actuators.ActuatorControl.reset_native_actuators`. Isaac Lab actuator models and backend-native
+  actuator state reset only the masked environments from the mask itself, without converting it to indices.
+* Added ``batch_mask`` to :meth:`~isaaclab.utils.buffers.DelayBuffer.set_time_lag` and
   :meth:`~isaaclab.utils.buffers.DelayBuffer.reset`, which select batches on the device without synchronizing the
   host. Tensor lags set through a mask are full-sized and are not range-checked.
 * Added ``env_mask`` to :meth:`~isaaclab.assets.BaseCableObject.reset`, matching the other scene assets.
diff --git a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
index fe555c95ab08..ae017575fe71 100644
--- a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
+++ b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
@@ -4,9 +4,10 @@
   the resampled environments, matching :meth:`~isaaclab.sensors.SensorBase.reset`.
 * Changed :meth:`~isaaclab.actuators.newton.NewtonActuatorAdapter.reset` to reuse preallocated per-DOF masks, so
   masked resets of Newton-native actuator state do not allocate and can be captured in a CUDA graph.
-* Changed :meth:`~isaaclab.actuators.ActuatorBase.reset` to also receive a 1-D boolean mask of the environments to
-  reset, which masked resets pass in place of indices. Custom actuator models must accept a mask, for example by
-  resetting their buffers through :func:`~isaaclab.utils.array.index_fill_`.
+* **Breaking:** Changed :meth:`~isaaclab.actuators.ActuatorBase.reset` to take an ``env_mask`` argument, the
+  boolean mask of the environments to reset that masked resets pass instead of indices. Custom actuator models must
+  accept ``env_mask`` and reset from it without synchronizing the host, for example through
+  :func:`~isaaclab.utils.array.index_fill_` or :func:`torch.where`.
 * Changed :class:`~isaaclab.actuators.DelayedPDActuator` to raise :class:`ValueError` at construction unless
   ``0 <= min_delay <= max_delay``, instead of at a reset, because masked resets apply the sampled delays without
   checking them on the host.
diff --git a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.minor b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.major
similarity index 100%
rename from source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.minor
rename to source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.major
diff --git a/source/isaaclab/isaaclab/actuators/actuator_base.py b/source/isaaclab/isaaclab/actuators/actuator_base.py
index d7adda51f19b..a75221b05574 100644
--- a/source/isaaclab/isaaclab/actuators/actuator_base.py
+++ b/source/isaaclab/isaaclab/actuators/actuator_base.py
@@ -279,15 +279,17 @@ def joint_indices(self) -> slice | torch.Tensor:
     """
 
     @abstractmethod
-    def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None):
+    def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: torch.Tensor | None = None):
         """Reset the internals within the group.
 
-        Implementations must accept a boolean mask without synchronizing the host, as
-        :func:`~isaaclab.utils.array.index_fill_` does.
+        .. caution::
+            If both ``env_ids`` and ``env_mask`` are provided, ``env_mask`` takes precedence.
 
         Args:
-            env_ids: Environment indices, a 1-D boolean mask of shape (num_envs,), or a slice selecting the
-                environments to reset. None resets all environments.
+            env_ids: Environment indices or a slice selecting the environments to reset. None resets all
+                environments.
+            env_mask: Boolean mask of the environments to reset. Shape is (num_envs,). Implementations reset from it
+                without synchronizing the host, for example through :func:`~isaaclab.utils.array.index_fill_`.
         """
         raise NotImplementedError
 
diff --git a/source/isaaclab/isaaclab/actuators/actuator_collection.py b/source/isaaclab/isaaclab/actuators/actuator_collection.py
index b6840e6b3be0..cddb26627499 100644
--- a/source/isaaclab/isaaclab/actuators/actuator_collection.py
+++ b/source/isaaclab/isaaclab/actuators/actuator_collection.py
@@ -183,9 +183,9 @@ def applied_effort(self) -> ProxyArray:
     def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: wp.array | None = None) -> None:
         """Reset all actuator group states.
 
-        Isaac Lab actuator models receive a mask as a zero-copy boolean tensor, which
-        :meth:`~isaaclab.actuators.ActuatorBase.reset` accepts in place of indices. Backend-native
-        actuator state resets from the mask directly.
+        Isaac Lab actuator models receive the mask as a zero-copy boolean tensor through the ``env_mask``
+        argument of :meth:`~isaaclab.actuators.ActuatorBase.reset`. Backend-native actuator state resets from
+        the mask directly.
 
         .. caution::
             If both ``env_ids`` and ``env_mask`` are provided, ``env_mask`` takes precedence.
@@ -194,18 +194,13 @@ def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: wp.array
             env_ids: Environment indices to reset. Defaults to all environments.
             env_mask: Environment mask. Shape is (num_instances,). Defaults to None.
         """
-        if env_mask is not None:
-            group_env_ids = wp.to_torch(env_mask)
-        else:
-            group_env_ids = self._control._normalize_index_sequence(env_ids)
+        group_env_ids = self._control._normalize_index_sequence(env_ids)
+        group_env_mask = None if env_mask is None else wp.to_torch(env_mask)
         for actuator in self._groups.values():
             # Newton-executed groups are reset through the backend below.
             if isinstance(actuator, ActuatorBase):
-                actuator.reset(group_env_ids)
-        if env_mask is not None:
-            self._control.reset_native_actuators_mask(env_mask)
-        else:
-            self._control.reset_native_actuators(slice(None) if group_env_ids is None else group_env_ids)
+                actuator.reset(group_env_ids, env_mask=group_env_mask)
+        self._control.reset_native_actuators(slice(None) if group_env_ids is None else group_env_ids, env_mask)
 
     def compute(self, dt: float = 0.0) -> None:
         """Compute processed actuator commands and telemetry.
diff --git a/source/isaaclab/isaaclab/actuators/actuator_control.py b/source/isaaclab/isaaclab/actuators/actuator_control.py
index a1793c02505b..b1aab97f58ca 100644
--- a/source/isaaclab/isaaclab/actuators/actuator_control.py
+++ b/source/isaaclab/isaaclab/actuators/actuator_control.py
@@ -308,24 +308,16 @@ def submit_commands(self, collection: ActuatorCollection) -> None:
         """
         raise NotImplementedError
 
-    def reset_native_actuators(self, env_ids: Sequence[int] | slice) -> None:
+    def reset_native_actuators(self, env_ids: Sequence[int] | slice, env_mask: wp.array | None = None) -> None:
         """Reset backend-native actuator state.
 
-        Args:
-            env_ids: Environments to reset.
-        """
-
-    def reset_native_actuators_mask(self, env_mask: wp.array) -> None:
-        """Reset backend-native actuator state for the environments selected by a mask.
-
-        The default converts the mask to indices for :meth:`reset_native_actuators`. Backends whose
-        native state resets from a mask override it.
+        .. caution::
+            If both ``env_ids`` and ``env_mask`` are provided, ``env_mask`` takes precedence.
 
         Args:
-            env_mask: Environments to reset. Shape is (num_instances,).
+            env_ids: Environments to reset.
+            env_mask: Environment mask. Shape is (num_instances,). Defaults to None.
         """
-        if self.native_actuator_path_active:
-            self.reset_native_actuators(wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1))
 
 
 class ArticulationActuatorControl(ActuatorControl):
diff --git a/source/isaaclab/isaaclab/actuators/actuator_net.py b/source/isaaclab/isaaclab/actuators/actuator_net.py
index c5dca36a609d..9a026fe07d5a 100644
--- a/source/isaaclab/isaaclab/actuators/actuator_net.py
+++ b/source/isaaclab/isaaclab/actuators/actuator_net.py
@@ -68,11 +68,12 @@ def __init__(self, cfg: ActuatorNetLSTMCfg, *args, **kwargs):
     Operations.
     """
 
-    def reset(self, env_ids: Sequence[int]):
+    def reset(self, env_ids: Sequence[int] | None = None, env_mask: torch.Tensor | None = None):
         # reset the hidden and cell states for the specified environments
+        selection = env_ids if env_mask is None else env_mask
         with torch.no_grad():
-            index_fill_(self.sea_hidden_state_per_env, env_ids, 0.0, dim=1)
-            index_fill_(self.sea_cell_state_per_env, env_ids, 0.0, dim=1)
+            index_fill_(self.sea_hidden_state_per_env, selection, 0.0, dim=1)
+            index_fill_(self.sea_cell_state_per_env, selection, 0.0, dim=1)
 
     def compute(
         self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor
@@ -138,10 +139,11 @@ def __init__(self, cfg: ActuatorNetMLPCfg, *args, **kwargs):
     Operations.
     """
 
-    def reset(self, env_ids: Sequence[int]):
+    def reset(self, env_ids: Sequence[int] | None = None, env_mask: torch.Tensor | None = None):
         # reset the history for the specified environments
-        index_fill_(self._joint_pos_error_history, env_ids, 0.0)
-        index_fill_(self._joint_vel_history, env_ids, 0.0)
+        selection = env_ids if env_mask is None else env_mask
+        index_fill_(self._joint_pos_error_history, selection, 0.0)
+        index_fill_(self._joint_vel_history, selection, 0.0)
 
     def compute(
         self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor
diff --git a/source/isaaclab/isaaclab/actuators/actuator_pd.py b/source/isaaclab/isaaclab/actuators/actuator_pd.py
index 1885ccbb7fea..c2d2e5a2c192 100644
--- a/source/isaaclab/isaaclab/actuators/actuator_pd.py
+++ b/source/isaaclab/isaaclab/actuators/actuator_pd.py
@@ -190,7 +190,7 @@ def _warn_articulation_owned_write(self, name: str, writer_name: str) -> None:
     Operations.
     """
 
-    def reset(self, env_ids: Sequence[int] | None = None):
+    def reset(self, env_ids: Sequence[int] | None = None, env_mask: torch.Tensor | None = None):
         # This is a no-op. There is no state to reset for implicit actuators.
         pass
 
@@ -313,7 +313,7 @@ def __init__(
     Operations.
     """
 
-    def reset(self, env_ids: Sequence[int]):
+    def reset(self, env_ids: Sequence[int] | None = None, env_mask: torch.Tensor | None = None):
         pass
 
     def compute(
@@ -460,12 +460,12 @@ def __init__(self, cfg: DelayedPDActuatorCfg, *args, **kwargs):
         self.efforts_delay_buffer = DelayBuffer(cfg.max_delay, self._num_envs, device=self._device)
         self._delay_buffers = (self.positions_delay_buffer, self.velocities_delay_buffer, self.efforts_delay_buffer)
 
-    def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None):
-        super().reset(env_ids)
+    def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: torch.Tensor | None = None):
+        super().reset(env_ids, env_mask)
         # number of environments (since env_ids can be a slice)
         if env_ids is None:
             env_ids = slice(None)
-        if isinstance(env_ids, torch.Tensor) and env_ids.dtype == torch.bool:
+        if env_mask is not None:
             # a mask samples every environment and the delay buffers keep the masked ones on the device
             num_envs = self._num_envs
         else:
@@ -480,8 +480,8 @@ def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None):
         )
         # set delays and reset buffers
         for delay_buffer in self._delay_buffers:
-            delay_buffer.set_time_lag(time_lags, env_ids)
-            delay_buffer.reset(env_ids)
+            delay_buffer.set_time_lag(time_lags, env_ids, batch_mask=env_mask)
+            delay_buffer.reset(env_ids, batch_mask=env_mask)
 
     def compute(
         self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor
diff --git a/source/isaaclab/isaaclab/actuators/newton/adapter.py b/source/isaaclab/isaaclab/actuators/newton/adapter.py
index 4790fd16b897..bb1ace3438f5 100644
--- a/source/isaaclab/isaaclab/actuators/newton/adapter.py
+++ b/source/isaaclab/isaaclab/actuators/newton/adapter.py
@@ -101,8 +101,10 @@ def __init__(
 
         self._states_a = [act.state() for act in actuators]
         self._states_b = [act.state() for act in actuators]
-        # Per-DOF reset masks, rebuilt in full on every partial reset so masked resets do not allocate.
+        # Newton actuator states reset by a per-DOF mask, which every partial reset expands from the env mask.
+        # Allocated once and fully rewritten on each reset, so a reset does not allocate.
         self._reset_dof_masks = [wp.zeros(act.indices.shape[0], dtype=wp.bool, device=device) for act in actuators]
+        # Env masks for resolving env ids in :meth:`reset`.
         self._all_env_mask = wp.ones(num_envs, dtype=wp.bool, device=device)
         self._scratch_env_mask = wp.zeros(num_envs, dtype=wp.bool, device=device)
 
diff --git a/source/isaaclab/isaaclab/utils/buffers/delay_buffer.py b/source/isaaclab/isaaclab/utils/buffers/delay_buffer.py
index edbc864b3d8d..9ab8c652994f 100644
--- a/source/isaaclab/isaaclab/utils/buffers/delay_buffer.py
+++ b/source/isaaclab/isaaclab/utils/buffers/delay_buffer.py
@@ -122,20 +122,29 @@ def time_lags(self) -> torch.Tensor:
     Operations.
     """
 
-    def set_time_lag(self, time_lag: int | torch.Tensor, batch_ids: Sequence[int] | torch.Tensor | None = None):
-        """Sets the time lag for the delay buffer across the provided batch indices.
+    def set_time_lag(
+        self,
+        time_lag: int | torch.Tensor,
+        batch_ids: Sequence[int] | None = None,
+        batch_mask: torch.Tensor | None = None,
+    ):
+        """Sets the time lag for the delay buffer across the provided batch indices or mask.
+
+        .. caution::
+            If both ``batch_ids`` and ``batch_mask`` are provided, ``batch_mask`` takes precedence.
 
         Args:
             time_lag: The desired delay for the buffer.
 
-              * If an integer is provided, the same delay is set for the provided batch indices.
-              * If a tensor is provided, the delay is set for each batch index separately. The shape of the tensor
-                should be (len(batch_ids),), or (batch_size,) when ``batch_ids`` is a boolean mask.
+              * If an integer is provided, the same delay is set for the selected batches.
+              * If a tensor is provided, the delay is set for each batch separately. The shape of the tensor
+                should be (len(batch_ids),), or (batch_size,) with ``batch_mask``.
 
-            batch_ids: The batch indices, or a boolean mask of shape (batch_size,), for which the time lag is set.
-                Default is None, which sets the time lag for all batch indices. A mask selects the batches on the
-                device without synchronizing the host, so tensor lags are not range-checked and must lie in
-                ``[0, history_length]``.
+            batch_ids: The batch indices for which the time lag is set. Default is None, which sets the time lag
+                for all batch indices.
+            batch_mask: Boolean mask of the batches for which the time lag is set. Shape is (batch_size,).
+                It selects on the device without synchronizing the host, so tensor lags are not range-checked and
+                must lie in ``[0, history_length]``.
 
         Raises:
             TypeError: If the type of the :attr:`time_lag` is not int or integer tensor.
@@ -152,9 +161,9 @@ def set_time_lag(self, time_lag: int | torch.Tensor, batch_ids: Sequence[int] |
             # check valid dtype for time_lag: must be int or long
             if time_lag.dtype not in [torch.int, torch.long]:
                 raise TypeError(f"Invalid dtype for time_lag: {time_lag.dtype}. Expected torch.int or torch.long.")
-            if isinstance(batch_ids, torch.Tensor) and batch_ids.dtype == torch.bool:
+            if batch_mask is not None:
                 time_lag = time_lag.to(device=self.device, dtype=self._time_lags.dtype)
-                torch.where(batch_ids, time_lag, self._time_lags, out=self._time_lags)
+                torch.where(batch_mask, time_lag, self._time_lags, out=self._time_lags)
                 return
             min_time_lag = int(time_lag.min().item()) if time_lag.numel() else 0
             max_time_lag = int(time_lag.max().item()) if time_lag.numel() else 0
@@ -169,23 +178,27 @@ def set_time_lag(self, time_lag: int | torch.Tensor, batch_ids: Sequence[int] |
         if isinstance(time_lag, torch.Tensor):
             self._time_lags[batch_ids] = time_lag.to(device=self.device, dtype=self._time_lags.dtype)
         else:
-            index_fill_(self._time_lags, batch_ids, time_lag)
+            index_fill_(self._time_lags, batch_ids if batch_mask is None else batch_mask, time_lag)
 
-    def reset(self, batch_ids: Sequence[int] | torch.Tensor | None = None):
-        """Reset the data in the delay buffer at the specified batch indices.
+    def reset(self, batch_ids: Sequence[int] | None = None, batch_mask: torch.Tensor | None = None):
+        """Reset the data in the delay buffer at the specified batch indices or mask.
 
         Automatically sampled lags are redrawn for those batches regardless of ``hold_prob``.
         Externally selected lags are preserved.
 
+        .. caution::
+            If both ``batch_ids`` and ``batch_mask`` are provided, ``batch_mask`` takes precedence.
+
         Args:
-            batch_ids: Elements to reset in the batch dimension, as indices or a boolean mask of shape (batch_size,).
-                Default is None, which resets all the batch indices.
+            batch_ids: Elements to reset in the batch dimension. Default is None, which resets all the batch indices.
+            batch_mask: Boolean mask of the batches to reset. Shape is (batch_size,). It selects on the device
+                without synchronizing the host.
         """
         indices = slice(None) if batch_ids is None else batch_ids
-        index_fill_(self._num_pushes, indices, 0)
+        index_fill_(self._num_pushes, indices if batch_mask is None else batch_mask, 0)
         if self._hold_prob is None:
             return
-        if isinstance(indices, torch.Tensor) and indices.dtype == torch.bool:
+        if batch_mask is not None:
             # Draw for every batch and keep the masked ones, so the selection stays on the device.
             lags = torch.randint(
                 self._min_lag,
@@ -194,7 +207,7 @@ def reset(self, batch_ids: Sequence[int] | torch.Tensor | None = None):
                 dtype=self._time_lags.dtype,
                 device=self.device,
             )
-            torch.where(indices, lags, self._time_lags, out=self._time_lags)
+            torch.where(batch_mask, lags, self._time_lags, out=self._time_lags)
         else:
             self._time_lags[indices] = torch.randint(
                 self._min_lag,
diff --git a/source/isaaclab/test/actuators/test_delayed_pd_actuator.py b/source/isaaclab/test/actuators/test_delayed_pd_actuator.py
index d3ffb048cac2..5c2d5a6cd38c 100644
--- a/source/isaaclab/test/actuators/test_delayed_pd_actuator.py
+++ b/source/isaaclab/test/actuators/test_delayed_pd_actuator.py
@@ -51,7 +51,7 @@ def step(actuator, target: float) -> torch.Tensor:
     previous = torch.cuda.get_sync_debug_mode()
     torch.cuda.set_sync_debug_mode("error")
     try:
-        by_mask.reset(env_mask)
+        by_mask.reset(env_mask=env_mask)
     finally:
         torch.cuda.set_sync_debug_mode(previous)
 
diff --git a/source/isaaclab/test/utils/test_delay_buffer.py b/source/isaaclab/test/utils/test_delay_buffer.py
index 84ffb9a365cd..6055fc0ffde0 100644
--- a/source/isaaclab/test/utils/test_delay_buffer.py
+++ b/source/isaaclab/test/utils/test_delay_buffer.py
@@ -127,10 +127,10 @@ def without_sync(call):
             torch.cuda.set_sync_debug_mode(previous)
 
     by_ids.set_time_lag(masked_lags[batch_ids], batch_ids)
-    without_sync(lambda: by_mask.set_time_lag(masked_lags, mask))
+    without_sync(lambda: by_mask.set_time_lag(masked_lags, batch_mask=mask))
     torch.testing.assert_close(by_mask.time_lags, by_ids.time_lags)
     by_ids.reset(batch_ids)
-    without_sync(lambda: by_mask.reset(mask))
+    without_sync(lambda: by_mask.reset(batch_mask=mask))
     torch.testing.assert_close(by_mask.time_lags, by_ids.time_lags)
     for data in _generate_data(by_mask.batch_size, 4, device):
         torch.testing.assert_close(by_mask.compute(data), by_ids.compute(data))
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/actuator_control.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/actuator_control.py
index e09b575e4604..f2e99903dc24 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/actuator_control.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/actuator_control.py
@@ -196,13 +196,9 @@ def submit_commands(self, collection: ActuatorCollection) -> None:
                 articulation._ALL_ENV_MASK,
             )
 
-    def reset_native_actuators(self, env_ids: Sequence[int] | slice) -> None:
+    def reset_native_actuators(self, env_ids: Sequence[int] | slice, env_mask: wp.array | None = None) -> None:
         if self._native_actuator_path_active and SimulationManager._adapter is not None:
-            SimulationManager._adapter.reset(env_ids)
-
-    def reset_native_actuators_mask(self, env_mask: wp.array) -> None:
-        if self._native_actuator_path_active and SimulationManager._adapter is not None:
-            SimulationManager._adapter.reset(env_mask=env_mask)
+            SimulationManager._adapter.reset(env_ids, env_mask)
 
     def _joint_dof_offset(self) -> int:
         """Return the first selected joint DOF's model offset within an environment."""
diff --git a/source/isaaclab_ov/isaaclab_ov/assets/articulation/actuator_control.py b/source/isaaclab_ov/isaaclab_ov/assets/articulation/actuator_control.py
index 8b501148aad3..15ddafab9200 100644
--- a/source/isaaclab_ov/isaaclab_ov/assets/articulation/actuator_control.py
+++ b/source/isaaclab_ov/isaaclab_ov/assets/articulation/actuator_control.py
@@ -98,13 +98,9 @@ def compute_native_actuators(self, collection: ActuatorCollection, dt: float) ->
         self._actuator_runtime.compute(collection, dt)
         return True
 
-    def reset_native_actuators(self, env_ids: Sequence[int] | slice) -> None:
+    def reset_native_actuators(self, env_ids: Sequence[int] | slice, env_mask: wp.array | None = None) -> None:
         if self._native_actuator_path_active and self._actuator_runtime is not None:
-            self._actuator_runtime.reset(env_ids)
-
-    def reset_native_actuators_mask(self, env_mask: wp.array) -> None:
-        if self._native_actuator_path_active and self._actuator_runtime is not None:
-            self._actuator_runtime.reset(env_mask=env_mask)
+            self._actuator_runtime.reset(env_ids, env_mask)
 
     def stage_user_command(
         self,
diff --git a/source/isaaclab_physx/changelog.d/jichuanh-warp-scene-mask-native.added.rst b/source/isaaclab_physx/changelog.d/jichuanh-warp-scene-mask-native.added.rst
new file mode 100644
index 000000000000..e51504073ab7
--- /dev/null
+++ b/source/isaaclab_physx/changelog.d/jichuanh-warp-scene-mask-native.added.rst
@@ -0,0 +1,3 @@
+* Added ``env_mask`` support to the Newton-native actuator reset of PhysX articulations, so
+  :meth:`~isaaclab.actuators.ActuatorCollection.reset` with a mask resets the native actuator state of only the masked
+  environments, without converting the mask to indices.
diff --git a/source/isaaclab_physx/changelog.d/jichuanh-warp-scene-mask-native.skip b/source/isaaclab_physx/changelog.d/jichuanh-warp-scene-mask-native.skip
deleted file mode 100644
index e69de29bb2d1..000000000000
diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/actuator_control.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/actuator_control.py
index da1714701482..13f0d39ef883 100644
--- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/actuator_control.py
+++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/actuator_control.py
@@ -22,6 +22,8 @@
 from isaaclab.sim.utils.stage import get_current_stage
 
 if TYPE_CHECKING:
+    import warp as wp
+
     from .articulation import Articulation
 
 logger = logging.getLogger(__name__)
@@ -155,6 +157,6 @@ def submit_commands(self, collection: ActuatorCollection) -> None:
             articulation.root_view.set_dof_position_targets(pos_target, articulation._ALL_INDICES)
             articulation.root_view.set_dof_velocity_targets(vel_target, articulation._ALL_INDICES)
 
-    def reset_native_actuators(self, env_ids: Sequence[int] | slice) -> None:
+    def reset_native_actuators(self, env_ids: Sequence[int] | slice, env_mask: wp.array | None = None) -> None:
         if self._native_actuator_path_active and self._actuator_runtime is not None:
-            self._actuator_runtime.reset(env_ids)
+            self._actuator_runtime.reset(env_ids, env_mask)

From 75de1d92de9e2718dcc86f6a92193d32f7641ef4 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Thu, 1 Oct 2026 11:47:31 -0700
Subject: [PATCH 30/48] Forward env_mask through PhysX asset resets

PhysX articulation, rigid object and rigid object collection resets pass env_mask to their actuators and wrench composers instead of resetting every environment.
---
 .../jichuanh-warp-scene-mask-native.fixed.rst |  5 ++++
 .../assets/articulation/articulation.py       |  2 +-
 .../assets/rigid_object/rigid_object.py       |  4 ++--
 .../rigid_object_collection.py                |  4 ++--
 .../assets/test_newton_actuators_physx.py     |  2 ++
 .../test/assets/test_rigid_object.py          | 23 +++++++++++--------
 .../assets/test_rigid_object_collection.py    | 23 +++++++++++--------
 7 files changed, 40 insertions(+), 23 deletions(-)
 create mode 100644 source/isaaclab_physx/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst

diff --git a/source/isaaclab_physx/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst b/source/isaaclab_physx/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
new file mode 100644
index 000000000000..fcfe2bc98ec3
--- /dev/null
+++ b/source/isaaclab_physx/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
@@ -0,0 +1,5 @@
+* Fixed :meth:`~isaaclab_physx.assets.Articulation.reset` with only ``env_mask`` resetting the actuator state of every
+  environment.
+* Fixed :meth:`~isaaclab_physx.assets.RigidObject.reset` and
+  :meth:`~isaaclab_physx.assets.RigidObjectCollection.reset` ignoring ``env_mask`` and clearing the external
+  wrenches of every environment.
diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py
index 86b05a3e4181..025dd97492af 100644
--- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py
+++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py
@@ -229,7 +229,7 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None
             env_ids = slice(None)
         # reset actuators, including backend-native actuator state. None selects all
         # environments; delayed-actuator buffers do not accept a slice.
-        self.actuators.reset(None if env_ids == slice(None) else env_ids)
+        self.actuators.reset(None if env_ids == slice(None) else env_ids, env_mask=env_mask)
         # reset external wrenches.
         self._instantaneous_wrench_composer.reset(env_ids, env_mask)
         self._permanent_wrench_composer.reset(env_ids, env_mask)
diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py
index cb95f52bdc62..546b699e8543 100644
--- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py
+++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py
@@ -138,8 +138,8 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None
         if (env_ids is None) or (env_ids == slice(None)):
             env_ids = slice(None)
         # reset external wrench
-        self._instantaneous_wrench_composer.reset(env_ids)
-        self._permanent_wrench_composer.reset(env_ids)
+        self._instantaneous_wrench_composer.reset(env_ids, env_mask)
+        self._permanent_wrench_composer.reset(env_ids, env_mask)
 
     def write_data_to_sim(self) -> None:
         """Write external wrench to the simulation.
diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py
index 341d60b62bbb..4b320d20baa6 100644
--- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py
+++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py
@@ -175,8 +175,8 @@ def reset(
         if object_ids is None:
             object_ids = self._ALL_BODY_INDICES
         # reset external wrench
-        self._instantaneous_wrench_composer.reset(env_ids)
-        self._permanent_wrench_composer.reset(env_ids)
+        self._instantaneous_wrench_composer.reset(env_ids, env_mask)
+        self._permanent_wrench_composer.reset(env_ids, env_mask)
 
     def write_data_to_sim(self) -> None:
         """Write external wrench to the simulation.
diff --git a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py
index 266cd318d3a6..2324ee8325bd 100644
--- a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py
+++ b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py
@@ -713,6 +713,8 @@ class TestActuatorStateReset(ActuatorStateResetBase, unittest.TestCase):
     adapter (``articulation.newton_actuator_adapter``).
     """
 
+    RESET_SELECTORS = ("env_ids", "env_mask")
+
     def _make_sim_cfg(self, use_newton_actuators: bool) -> SimulationCfg:
         return SimulationCfg(dt=DT, physics=PhysxCfg(), use_newton_actuators=use_newton_actuators)
 
diff --git a/source/isaaclab_physx/test/assets/test_rigid_object.py b/source/isaaclab_physx/test/assets/test_rigid_object.py
index ac84e269f1b6..8de895c6a9b8 100644
--- a/source/isaaclab_physx/test/assets/test_rigid_object.py
+++ b/source/isaaclab_physx/test/assets/test_rigid_object.py
@@ -281,16 +281,21 @@ def test_reset_rigid_object(num_cubes, device):
 
         composers = (cube_object.instantaneous_wrench_composer, cube_object.permanent_wrench_composer)
         ones = torch.ones((num_cubes, 1, 3), device=device)
-        cube_object.permanent_wrench_composer.set_forces_and_torques_index(forces=ones, torques=ones)
-        cube_object.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones)
 
-        # A partial reset clears only the selected environment
-        cube_object.reset(env_ids=torch.tensor([0], device=device))
-        for composer in composers:
-            assert composer.active
-            for buffer in (composer.out_force_b.torch, composer.out_torque_b.torch):
-                assert torch.count_nonzero(buffer[0]) == 0
-                assert torch.count_nonzero(buffer[1:]) == buffer[1:].numel()
+        # A partial reset by indices or by mask clears only the selected environment
+        selections = {
+            "env_ids": torch.tensor([0], device=device),
+            "env_mask": wp.from_torch(torch.arange(num_cubes, device=device) == 0),
+        }
+        for selector, selection in selections.items():
+            cube_object.permanent_wrench_composer.set_forces_and_torques_index(forces=ones, torques=ones)
+            cube_object.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones)
+            cube_object.reset(**{selector: selection})
+            for composer in composers:
+                assert composer.active
+                for buffer in (composer.out_force_b.torch, composer.out_torque_b.torch):
+                    assert torch.count_nonzero(buffer[0]) == 0, selector
+                    assert torch.count_nonzero(buffer[1:]) == buffer[1:].numel(), selector
 
         # A full reset clears every environment and deactivates the composers
         cube_object.reset()
diff --git a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py
index 2e0af9a4be8b..d31ceeebf9c4 100644
--- a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py
+++ b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py
@@ -599,16 +599,21 @@ def test_reset_object_collection(sim, num_envs, num_cubes, device):
 
     composers = (object_collection.instantaneous_wrench_composer, object_collection.permanent_wrench_composer)
     ones = torch.ones((num_envs, num_cubes, 3), device=device)
-    object_collection.permanent_wrench_composer.set_forces_and_torques_index(forces=ones, torques=ones)
-    object_collection.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones)
 
-    # A partial reset clears only the selected environment
-    object_collection.reset(env_ids=torch.tensor([0], device=device))
-    for composer in composers:
-        assert composer.active
-        for buffer in (composer.out_force_b.torch, composer.out_torque_b.torch):
-            assert torch.count_nonzero(buffer[0]) == 0
-            assert torch.count_nonzero(buffer[1:]) == buffer[1:].numel()
+    # A partial reset by indices or by mask clears only the selected environment
+    selections = {
+        "env_ids": torch.tensor([0], device=device),
+        "env_mask": wp.from_torch(torch.arange(num_envs, device=device) == 0),
+    }
+    for selector, selection in selections.items():
+        object_collection.permanent_wrench_composer.set_forces_and_torques_index(forces=ones, torques=ones)
+        object_collection.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones)
+        object_collection.reset(**{selector: selection})
+        for composer in composers:
+            assert composer.active
+            for buffer in (composer.out_force_b.torch, composer.out_torque_b.torch):
+                assert torch.count_nonzero(buffer[0]) == 0, selector
+                assert torch.count_nonzero(buffer[1:]) == buffer[1:].numel(), selector
 
     # A full reset clears every environment
     object_collection.reset()

From 6ac6061a1375c016192792d2d7d8f905a100c5a3 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Thu, 1 Oct 2026 11:47:31 -0700
Subject: [PATCH 31/48] Share one per-environment Warp RNG state through
 WarpRng

Environments initialize WarpRng once at construction and configure_seed reseeds it in place. The ray caster, experimental managers, events, noise models and Warp tasks read WarpRng.state instead of keeping their own copies; rng_state_wp is removed.
---
 .../jichuanh-warp-scene-mask-native.added.rst |  3 +
 ...ichuanh-warp-scene-mask-native.changed.rst |  6 +-
 .../sensors/ray_caster/base_ray_caster.py     | 14 +----
 .../isaaclab/sensors/ray_caster/kernels.py    | 19 ------
 source/isaaclab/isaaclab/utils/seed.py        | 58 ++++++++++++++++-
 .../isaaclab/isaaclab/utils/warp/kernels.py   | 25 ++++++++
 source/isaaclab/test/utils/test_warp_rng.py   | 63 +++++++++++++++++++
 ...ichuanh-warp-scene-mask-native.changed.rst |  4 ++
 ...ichuanh-warp-scene-mask-native.removed.rst |  2 +
 .../envs/direct_rl_env_warp.py                |  9 ++-
 .../envs/manager_based_env_warp.py            | 39 +++---------
 .../isaaclab_experimental/envs/mdp/events.py  | 23 +++----
 .../managers/command_manager.py               |  5 +-
 .../managers/event_manager.py                 | 14 ++---
 .../managers/observation_manager.py           |  6 --
 .../utils/noise/noise_model.py                | 12 ++--
 .../test/envs/mdp/test_events_warp_parity.py  |  5 +-
 .../envs/mdp/test_observations_warp_parity.py |  7 ++-
 .../test/envs/mdp/test_rewards_warp_parity.py |  7 ++-
 .../envs/mdp/test_terminations_warp_parity.py |  7 ++-
 .../sensors/test_newton_raycast_sensor.py     |  4 ++
 ...ichuanh-warp-scene-mask-native.changed.rst |  2 +
 .../core/cartpole/cartpole_warp_env.py        | 27 +-------
 .../core/locomotion/locomotion_env_warp.py    | 20 ++----
 .../core/reorient/reorient_warp_env.py        | 34 ++--------
 25 files changed, 232 insertions(+), 183 deletions(-)
 create mode 100644 source/isaaclab/test/utils/test_warp_rng.py
 create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
 create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.removed.rst
 create mode 100644 source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-scene-mask-native.changed.rst

diff --git a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst
index 37303ca182c2..df8793cd86bb 100644
--- a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst
+++ b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst
@@ -5,3 +5,6 @@
   :meth:`~isaaclab.utils.buffers.DelayBuffer.reset`, which select batches on the device without synchronizing the
   host. Tensor lags set through a mask are full-sized and are not range-checked.
 * Added ``env_mask`` to :meth:`~isaaclab.assets.BaseCableObject.reset`, matching the other scene assets.
+* Added :class:`~isaaclab.utils.seed.WarpRng`, the per-environment Warp random number generator state that
+  environments and sensors share, one per process. :func:`~isaaclab.utils.seed.configure_seed` reseeds it in place,
+  and environments call :meth:`~isaaclab.utils.seed.WarpRng.initialize` at construction.
diff --git a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
index ae017575fe71..1467738a803b 100644
--- a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
+++ b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
@@ -1,7 +1,7 @@
 * Changed masked :meth:`~isaaclab.sensors.ray_caster.BaseRayCaster.reset` to resample drift in a Warp kernel from
-  per-environment random states seeded from torch's global generator, so it neither synchronizes the host nor
-  allocates and can be captured in a CUDA graph. When both ``env_ids`` and ``env_mask`` are given, the mask selects
-  the resampled environments, matching :meth:`~isaaclab.sensors.SensorBase.reset`.
+  the per-environment random states of :class:`~isaaclab.utils.seed.WarpRng`, so it neither synchronizes the host
+  nor allocates and can be captured in a CUDA graph. When both ``env_ids`` and ``env_mask`` are given, the mask
+  selects the resampled environments, matching :meth:`~isaaclab.sensors.SensorBase.reset`.
 * Changed :meth:`~isaaclab.actuators.newton.NewtonActuatorAdapter.reset` to reuse preallocated per-DOF masks, so
   masked resets of Newton-native actuator state do not allocate and can be captured in a CUDA graph.
 * **Breaking:** Changed :meth:`~isaaclab.actuators.ActuatorBase.reset` to take an ``env_mask`` argument, the
diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py
index 2968954b91cb..fe0514c39acc 100644
--- a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py
+++ b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py
@@ -19,6 +19,7 @@
 from ...markers import VisualizationMarkers
 from ...terrains.trimesh.utils import make_plane
 from ...utils import math as math_utils
+from ...utils.seed import WarpRng
 from ...utils.warp import ProxyArray, convert_to_warp_mesh
 from ...utils.warp.kernels import raycast_mesh_masked_kernel
 from ..sensor_base import SensorBase
@@ -117,7 +118,7 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None
         if not self._drift_sampled:
             return
         if env_mask is not None:
-            # Resample in a kernel from per-environment random states with the ranges as launch
+            # Resample in a kernel from the shared per-environment random states with the ranges as launch
             # arguments, so the masked reset neither synchronizes nor allocates and can be captured.
             # A zero range draws zero, which clears drift that was sampled before it was disabled.
             wp.launch(
@@ -128,7 +129,7 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None
                     wp.vec2f(*self.cfg.drift_range),
                     wp.vec3f(*(low for low, _ in ray_cast_range_list)),
                     wp.vec3f(*(high for _, high in ray_cast_range_list)),
-                    self._drift_rng_state,
+                    WarpRng.state,
                 ],
                 outputs=[self.drift.warp, self.ray_cast_drift.warp],
                 device=self._device,
@@ -255,15 +256,6 @@ def _initialize_rays_impl(self):
         # Drift buffers are warp-first; reset uses explicit .torch views for sampling.
         self.drift = ProxyArray(wp.zeros(self._view_count, dtype=wp.vec3f, device=self._device))
         self.ray_cast_drift = ProxyArray(wp.zeros(self._view_count, dtype=wp.vec3f, device=self._device))
-        # Per-environment random states for masked drift resampling, seeded from torch's global generator.
-        self._drift_rng_state = wp.empty(self._view_count, dtype=wp.uint32, device=self._device)
-        wp.launch(
-            ray_caster_kernels.init_rng_state_kernel,
-            dim=self._view_count,
-            inputs=[int(torch.randint(0, 2**31 - 1, (1,)))],
-            outputs=[self._drift_rng_state],
-            device=self._device,
-        )
 
         # World-frame ray buffers
         self._ray_starts_w = wp.empty((self._view_count, self.num_rays), dtype=wp.vec3f, device=self._device)
diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py b/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py
index a337fef6f21e..149ed48a41af 100644
--- a/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py
+++ b/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py
@@ -170,25 +170,6 @@ def apply_z_drift_kernel(
     ray_hits[env, ray] = wp.vec3f(hit[0], hit[1], hit[2] + ray_cast_drift[env][2])
 
 
-@wp.kernel(enable_backward=False)
-def init_rng_state_kernel(
-    # input
-    seed: wp.int32,
-    # output
-    rng_state: wp.array(dtype=wp.uint32),
-):
-    """Initialize one random number generator state per environment.
-
-    Launch with dim=(num_envs,).
-
-    Args:
-        seed: Seed shared by all environments; each environment uses its index as the offset.
-        rng_state: Random number generator states. Shape is (num_envs,).
-    """
-    env = wp.tid()
-    rng_state[env] = wp.rand_init(seed, env)
-
-
 @wp.kernel(enable_backward=False)
 def resample_drift_masked_kernel(
     # input
diff --git a/source/isaaclab/isaaclab/utils/seed.py b/source/isaaclab/isaaclab/utils/seed.py
index 6b2a8ff97adf..5023e8a6237c 100644
--- a/source/isaaclab/isaaclab/utils/seed.py
+++ b/source/isaaclab/isaaclab/utils/seed.py
@@ -11,6 +11,62 @@
 import warp as wp
 
 
+class WarpRng:
+    """Process-wide per-environment random number generator state for Warp kernels.
+
+    Warp has no global generator: :func:`warp.rand_init` hashes a seed and an offset into a state that lives only
+    for one kernel launch. Kernels that draw random numbers per environment therefore read ``state[env]``, draw, and
+    store the advanced state back, so the next launch or CUDA-graph replay draws new numbers. This class holds that
+    array once per process, so environments, managers and sensors share it instead of each allocating a copy.
+
+    :meth:`initialize` must be called once the number of environments and the device are known; environments do this
+    at construction, so users read :attr:`state` directly. :meth:`seed` is optional; without it, the state is seeded
+    from :func:`torch.initial_seed`. Only kernels launched over environments, one thread per environment, may advance
+    the state.
+    """
+
+    state: wp.array | None = None
+    """The per-environment random number generator state, or None before :meth:`initialize`. Shape is (num_envs,)."""
+
+    _seed: int | None = None
+
+    @classmethod
+    def seed(cls, seed: int) -> None:
+        """Set the seed, and reseed the state in place if it is initialized.
+
+        Args:
+            seed: The random seed value.
+        """
+        cls._seed = seed
+        if cls.state is not None:
+            cls._reset_state()
+
+    @classmethod
+    def initialize(cls, num_envs: int, device: str) -> None:
+        """Allocate and seed a new state. Environments call it once at construction.
+
+        Args:
+            num_envs: The number of environments.
+            device: The device of the state.
+        """
+        cls.state = wp.empty(num_envs, dtype=wp.uint32, device=device)
+        cls._reset_state()
+
+    @classmethod
+    def _reset_state(cls) -> None:
+        """Seed every environment's state from the set seed, or from torch's initial seed without one."""
+        from .warp.kernels import initialize_rng_state
+
+        seed = cls._seed if cls._seed is not None else torch.initial_seed()
+        wp.launch(
+            initialize_rng_state,
+            dim=cls.state.shape[0],
+            inputs=[seed % 2**31],
+            outputs=[cls.state],
+            device=cls.state.device,
+        )
+
+
 def configure_seed(seed: int | None, torch_deterministic: bool = False) -> int:
     """Set seed across all random number generators (torch, numpy, random, warp).
 
@@ -30,7 +86,7 @@ def configure_seed(seed: int | None, torch_deterministic: bool = False) -> int:
     os.environ["PYTHONHASHSEED"] = str(seed)
     torch.cuda.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
-    wp.rand_init(seed)
+    WarpRng.seed(seed)
 
     if torch_deterministic:
         # refer to https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility
diff --git a/source/isaaclab/isaaclab/utils/warp/kernels.py b/source/isaaclab/isaaclab/utils/warp/kernels.py
index 3d046abe2bba..a7818691165e 100644
--- a/source/isaaclab/isaaclab/utils/warp/kernels.py
+++ b/source/isaaclab/isaaclab/utils/warp/kernels.py
@@ -780,3 +780,28 @@ def gather_subtract_2d(
 def gather_subtract_2d_kernel(env_ids: "wp.array | torch.Tensor", item_ids: "wp.array | torch.Tensor") -> wp.Kernel:
     """Select the gather-subtract worker for the selector dtypes."""
     return _GATHER_SUBTRACT_2D_DISPATCHER.select(env_ids, item_ids)
+
+
+##
+# Random number generation.
+##
+
+
+@wp.kernel(enable_backward=False)
+def initialize_rng_state(
+    # input
+    seed: wp.int32,
+    # output
+    state: wp.array(dtype=wp.uint32),
+):
+    """Initialize one random number generator state per environment.
+
+    Launch with ``dim=(num_envs,)``. Kernels that draw from ``state[i]`` should store the advanced state back,
+    so the next launch or graph replay draws new numbers.
+
+    Args:
+        seed: Seed shared by all environments; each environment uses its index as the offset.
+        state: Random number generator states. Shape is (num_envs,).
+    """
+    env_id = wp.tid()
+    state[env_id] = wp.rand_init(seed, env_id)
diff --git a/source/isaaclab/test/utils/test_warp_rng.py b/source/isaaclab/test/utils/test_warp_rng.py
new file mode 100644
index 000000000000..2bc3815a2b8c
--- /dev/null
+++ b/source/isaaclab/test/utils/test_warp_rng.py
@@ -0,0 +1,63 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tests for the process-wide per-environment Warp random number generator state."""
+
+import numpy as np
+import pytest
+import torch
+
+from isaaclab.test.utils import test_devices
+from isaaclab.utils.seed import WarpRng, configure_seed
+
+
+@pytest.fixture(autouse=True)
+def fresh_warp_rng():
+    """Start and end every test without a state or a set seed."""
+    WarpRng.state = None
+    WarpRng._seed = None
+    yield
+    WarpRng.state = None
+    WarpRng._seed = None
+
+
+@pytest.mark.parametrize("device", test_devices())
+def test_initialize_allocates_a_new_state(device):
+    """Every :meth:`WarpRng.initialize` call allocates and seeds a new state of the given size."""
+    WarpRng.initialize(4, device)
+    state = WarpRng.state
+    WarpRng.initialize(5, device)
+    assert WarpRng.state is not state
+    assert WarpRng.state.shape == (5,)
+
+
+@pytest.mark.parametrize("device", test_devices())
+def test_configure_seed_reseeds_in_place(device):
+    """configure_seed reseeds the existing state in place and reproducibly."""
+    configure_seed(7)
+    WarpRng.initialize(4, device)
+    state = WarpRng.state
+    seeded = state.numpy().copy()
+
+    state.fill_(0)
+    configure_seed(7)
+    assert WarpRng.state.ptr == state.ptr
+    np.testing.assert_array_equal(state.numpy(), seeded)
+
+    configure_seed(8)
+    assert not np.array_equal(state.numpy(), seeded)
+
+
+@pytest.mark.parametrize("device", test_devices())
+def test_default_seed_follows_torch(device):
+    """Without a set seed, the state is seeded from torch's initial seed."""
+    torch.manual_seed(5)
+    WarpRng.initialize(4, device)
+    default = WarpRng.state.numpy().copy()
+
+    WarpRng.state = None
+    WarpRng.seed(5)
+    WarpRng.initialize(4, device)
+    np.testing.assert_array_equal(WarpRng.state.numpy(), default)
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.changed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
new file mode 100644
index 000000000000..aa642663b908
--- /dev/null
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
@@ -0,0 +1,4 @@
+* Changed :class:`~isaaclab_experimental.envs.ManagerBasedEnvWarp` and
+  :class:`~isaaclab_experimental.envs.DirectRLEnvWarp` to initialize the process-wide
+  :class:`~isaaclab.utils.seed.WarpRng` state, which the Warp MDP terms, managers and noise models now draw from.
+  An environment created without a seed now seeds it from :func:`torch.initial_seed` instead of ``-1``.
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.removed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.removed.rst
new file mode 100644
index 000000000000..8e486bedf7bd
--- /dev/null
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.removed.rst
@@ -0,0 +1,2 @@
+* **Breaking:** Removed ``ManagerBasedEnvWarp.rng_state_wp``. Read :attr:`~isaaclab.utils.seed.WarpRng.state`
+  instead.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
index 7fdd1ffe3562..ef5eaa9c739f 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
@@ -32,7 +32,7 @@
 from isaaclab.sim.utils import use_stage
 from isaaclab.utils import instantiate, validate
 from isaaclab.utils.noise import NoiseModel
-from isaaclab.utils.seed import configure_seed
+from isaaclab.utils.seed import WarpRng, configure_seed
 from isaaclab.utils.timer import Timer
 
 from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache
@@ -125,7 +125,7 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
 
         # set the seed for the environment
         if self.cfg.seed is not None:
-            self.cfg.seed = self.seed(self.cfg.seed)
+            self.cfg.seed = configure_seed(self.cfg.seed)
         else:
             logger.warning("Seed not set for the environment. The environment creation may not be deterministic.")
 
@@ -164,6 +164,9 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
                 # attach_stage_to_usd_context()
         print("[INFO]: Scene manager: ", self.scene)
 
+        # Process-wide per-env Warp RNG state (isaaclab.utils.seed.WarpRng), shared by terms, tasks and sensors.
+        WarpRng.initialize(self.num_envs, self.device)
+
         # create event manager
         # note: this is needed here (rather than after simulation play) to allow USD-related randomization events
         #   that must happen before the simulation starts. Example: randomizing mesh scale
@@ -343,7 +346,7 @@ def reset(self, seed: int | None = None, options: dict[str, Any] | None = None)
         """
         # set the seed
         if seed is not None:
-            self.seed(seed)
+            configure_seed(seed)
 
         # reset state of scene
         self._reset_idx(self._ALL_ENV_MASK)
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
index 3e94486af610..5e27e8319c99 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
@@ -36,7 +36,7 @@
 from isaaclab.sim import SimulationContext
 from isaaclab.sim.utils import use_stage
 from isaaclab.utils import validate
-from isaaclab.utils.seed import configure_seed
+from isaaclab.utils.seed import WarpRng, configure_seed
 from isaaclab.utils.timer import Timer
 
 from isaaclab_experimental.envs.interactive_scene_warp import InteractiveSceneWarp as InteractiveScene
@@ -47,17 +47,6 @@
 logger = logging.getLogger(__name__)
 
 
-@wp.kernel
-def initialize_rng_state(
-    # input
-    seed: wp.int32,
-    # output
-    state: wp.array(dtype=wp.uint32),
-):
-    env_id = wp.tid()
-    state[env_id] = wp.rand_init(seed, wp.int32(env_id))
-
-
 class ManagerBasedEnvWarp:
     """The base environment for the manager-based workflow (experimental fork).
 
@@ -97,7 +86,7 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
 
         # set the seed for the environment
         if self.cfg.seed is not None:
-            self.cfg.seed = self.seed(self.cfg.seed)
+            self.cfg.seed = configure_seed(self.cfg.seed)
         else:
             logger.warning("Seed not set for the environment. The environment creation may not be deterministic.")
 
@@ -147,16 +136,8 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
                 # attach_stage_to_usd_context()
         print("[INFO]: Scene manager: ", self.scene)
 
-        # Shared per-env Warp RNG state (accessible to all managers/terms via `env`).
-        # This is a single stream per env (no lookup) and is initialized once when `num_envs` is known.
-        self.rng_state_wp = wp.zeros((self.num_envs,), dtype=wp.uint32, device=self.device)
-        seed_val = int(self.cfg.seed) if self.cfg.seed is not None else -1
-        wp.launch(
-            kernel=initialize_rng_state,
-            dim=self.num_envs,
-            inputs=[seed_val, self.rng_state_wp],
-            device=self.device,
-        )
+        # Process-wide per-env Warp RNG state (isaaclab.utils.seed.WarpRng), shared by terms, tasks and sensors.
+        WarpRng.initialize(self.num_envs, self.device)
 
         # TODO(jichuanh): this is problematic as warp capture requires stable pointers,
         #                 using different masks for different managers/terms will cause problems.
@@ -440,16 +421,10 @@ def reset(
 
         # set the seed
         if seed is not None:
-            used_seed = self.seed(seed)
+            # also reseeds the per-env Warp RNG state in place (stable pointer for capture)
+            used_seed = configure_seed(seed)
             # keep cfg seed in sync for downstream users
             self.cfg.seed = used_seed
-            # re-initialize per-env Warp RNG state without reallocating (stable pointer for capture)
-            wp.launch(
-                kernel=initialize_rng_state,
-                dim=self.num_envs,
-                inputs=[int(used_seed), self.rng_state_wp],
-                device=self.device,
-            )
 
         # reset state of scene
         self._reset_idx(env_ids)
@@ -504,7 +479,7 @@ def reset_to(
 
         # set the seed
         if seed is not None:
-            self.seed(seed)
+            configure_seed(seed)
 
         self._reset_idx(env_ids)
 
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py
index 90e5696a6b10..a64af1d309ce 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py
@@ -31,6 +31,7 @@
 
 from isaaclab.envs.mdp.events import randomize_rigid_body_mass as _StableRandomizeRigidBodyMass
 from isaaclab.envs.mdp.events import randomize_rigid_body_material as _StableRandomizeRigidBodyMaterial
+from isaaclab.utils.seed import WarpRng
 
 from isaaclab_experimental.managers import ManagerTermBase as _WarpManagerTermBase
 from isaaclab_experimental.managers import SceneEntityCfg
@@ -98,7 +99,7 @@ def randomize_rigid_body_com(
         dim=env.num_envs,
         inputs=[
             env_mask,
-            env.rng_state_wp,
+            WarpRng.state,
             asset.data.body_com_pos_b.warp,
             asset_cfg.body_ids_wp,
             fn._com_lo,
@@ -178,7 +179,7 @@ def apply_external_force_torque(
         dim=env.num_envs,
         inputs=[
             env_mask,
-            env.rng_state_wp,
+            WarpRng.state,
             apply_external_force_torque._scratch_forces,
             apply_external_force_torque._scratch_torques,
             force_range[0],
@@ -258,7 +259,7 @@ def push_by_setting_velocity(
         dim=env.num_envs,
         inputs=[
             env_mask,
-            env.rng_state_wp,
+            WarpRng.state,
             asset.data.root_vel_w.warp,
             push_by_setting_velocity._scratch_vel,
             push_by_setting_velocity._lin_lo,
@@ -377,7 +378,7 @@ def reset_root_state_uniform(
         dim=env.num_envs,
         inputs=[
             env_mask,
-            env.rng_state_wp,
+            WarpRng.state,
             asset.data.default_root_pose.warp,
             asset.data.default_root_vel.warp,
             env.env_origins_wp,
@@ -470,11 +471,6 @@ def reset_joints_by_offset(
             f"but got None for asset '{asset_cfg.name}'. "
             "Use isaaclab_experimental.managers.SceneEntityCfg and ensure joint_names are set."
         )
-    if not hasattr(env, "rng_state_wp") or env.rng_state_wp is None:
-        raise AttributeError(
-            "reset_joints_by_offset requires env.rng_state_wp to be initialized. "
-            "Use ManagerBasedEnvWarp or ManagerBasedRLEnvWarp as the base environment."
-        )
 
     wp.launch(
         kernel=_reset_joints_by_offset_kernel,
@@ -482,7 +478,7 @@ def reset_joints_by_offset(
         inputs=[
             env_mask,
             asset_cfg.joint_ids_wp,
-            env.rng_state_wp,
+            WarpRng.state,
             asset.data.default_joint_pos.warp,
             asset.data.default_joint_vel.warp,
             asset.data.joint_pos.warp,
@@ -561,11 +557,6 @@ def reset_joints_by_scale(
             f"but got None for asset '{asset_cfg.name}'. "
             "Use isaaclab_experimental.managers.SceneEntityCfg and ensure joint_names are set."
         )
-    if not hasattr(env, "rng_state_wp") or env.rng_state_wp is None:
-        raise AttributeError(
-            "reset_joints_by_scale requires env.rng_state_wp to be initialized. "
-            "Use ManagerBasedEnvWarp or ManagerBasedRLEnvWarp as the base environment."
-        )
 
     wp.launch(
         kernel=_reset_joints_by_scale_kernel,
@@ -573,7 +564,7 @@ def reset_joints_by_scale(
         inputs=[
             env_mask,
             asset_cfg.joint_ids_wp,
-            env.rng_state_wp,
+            WarpRng.state,
             asset.data.default_joint_pos.warp,
             asset.data.default_joint_vel.warp,
             asset.data.joint_pos.warp,
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py
index e7d8b4b746b0..2b4cd8de1828 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py
@@ -18,6 +18,7 @@
 
 from isaaclab.managers.manager_term_cfg import CommandTermCfg
 from isaaclab.utils import instantiate
+from isaaclab.utils.seed import WarpRng
 
 from isaaclab_experimental.utils.warp.kernels import compute_reset_scale, count_masked
 
@@ -306,8 +307,6 @@ def _resample(self, env_mask: wp.array):
             raise TypeError(f"env_mask must be a wp.array (got {type(env_mask)}).")
         if env_mask.dtype != wp.bool or env_mask.ndim != 1:
             raise TypeError(f"env_mask must be wp.bool 1D (got dtype={env_mask.dtype}, ndim={env_mask.ndim}).")
-        if self._env.rng_state_wp is None:
-            raise RuntimeError("Environment rng_state_wp is not initialized.")
 
         # resample time-left and increment command-counter for masked envs
         wp.launch(
@@ -317,7 +316,7 @@ def _resample(self, env_mask: wp.array):
                 env_mask,
                 self.time_left_wp,
                 self.command_counter_wp,
-                self._env.rng_state_wp,
+                WarpRng.state,
                 float(self.cfg.resampling_time_range[0]),
                 float(self.cfg.resampling_time_range[1]),
             ],
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
index 757d9a24d547..4e67ffedcbc2 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
@@ -30,6 +30,7 @@
 from prettytable import PrettyTable
 
 from isaaclab.managers.manager_term_cfg import EventTermCfg
+from isaaclab.utils.seed import WarpRng
 
 from .manager_base import ManagerBase
 
@@ -273,7 +274,7 @@ def reset(
                     inputs=[
                         env_mask,
                         self._interval_term_time_left_wp[i],
-                        self._env.rng_state_wp,
+                        WarpRng.state,
                         float(lower),
                         float(upper),
                     ],
@@ -331,9 +332,6 @@ def apply(
             term_cfg.func(self._env, env_mask_wp, **term_cfg.params)
 
     def _apply_interval(self, dt: float) -> None:
-        if self._env.rng_state_wp is None:
-            raise RuntimeError("EventManager._apply_interval: env.rng_state_wp is not initialized.")
-
         # iterate over all the interval terms (fixed list; captured graph-friendly)
         for i, term_cfg in enumerate(self._mode_term_cfgs["interval"]):
             lower, upper = self._interval_term_ranges[i]
@@ -363,7 +361,7 @@ def _apply_interval(self, dt: float) -> None:
                     dim=self.num_envs,
                     inputs=[
                         self._interval_term_time_left_wp[i],
-                        self._env.rng_state_wp,
+                        WarpRng.state,
                         self._scratch_term_mask_wp,
                         float(dt),
                         float(lower),
@@ -460,13 +458,11 @@ def _prepare_terms(self):
                 if term_cfg.is_global_time:
                     # allocate and seed scalar global RNG state if needed (avoid consuming env0 RNG stream)
                     if self._interval_global_rng_state_wp is None:
-                        if self._env.rng_state_wp is None:
-                            raise RuntimeError("EventManager._prepare_terms: env.rng_state_wp is not initialized.")
                         self._interval_global_rng_state_wp = wp.zeros((1,), dtype=wp.uint32, device=self.device)
                         wp.launch(
                             kernel=_seed_global_rng_from_env_rng,
                             dim=1,
-                            inputs=[self._env.rng_state_wp, self._interval_global_rng_state_wp],
+                            inputs=[WarpRng.state, self._interval_global_rng_state_wp],
                             device=self.device,
                         )
                     time_left = wp.zeros((1,), dtype=wp.float32, device=self.device)
@@ -483,7 +479,7 @@ def _prepare_terms(self):
                     wp.launch(
                         kernel=_interval_init_per_env,
                         dim=self.num_envs,
-                        inputs=[time_left, self._env.rng_state_wp, float(lower), float(upper)],
+                        inputs=[time_left, WarpRng.state, float(lower), float(upper)],
                         device=self.device,
                     )
                     self._interval_term_time_left_wp.append(time_left)
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
index 051aaa8a6ff6..6ec29aa0e921 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
@@ -780,14 +780,8 @@ def _prepare_terms(self):  # noqa: C901
                                     f" and optional parameters: {args_with_defaults}, but received: {term_params}."
                                 )
 
-                # plumb the shared per-env RNG state so Warp noise kernels can consume it
-                # (function-style NoiseCfg kernels read it off the cfg at call time)
-                if term_cfg.noise is not None and isinstance(term_cfg.noise, noise.NoiseCfg):
-                    term_cfg.noise.rng_state_wp = self._env.rng_state_wp
                 # prepare noise model classes
                 if term_cfg.noise is not None and isinstance(term_cfg.noise, noise.NoiseModelCfg):
-                    # plumb the shared per-env RNG state so Warp noise kernels can consume it
-                    term_cfg.noise.rng_state_wp = self._env.rng_state_wp
                     noise_model_cls = term_cfg.noise.class_type
                     if not issubclass(noise_model_cls, noise.NoiseModel):
                         raise TypeError(
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/noise/noise_model.py b/source/isaaclab_experimental/isaaclab_experimental/utils/noise/noise_model.py
index acba926ea488..cbc933324830 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/utils/noise/noise_model.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/utils/noise/noise_model.py
@@ -10,10 +10,8 @@
 
     noise_cfg.func(data_wp, noise_cfg) -> None
 
-Random noise kernels (gaussian, uniform) consume the shared per-env Warp RNG state
-(``rng_state_wp``) that is set on the config at manager prep time from
-``env.rng_state_wp``.  See :func:`initialize_rng_state` in
-``isaaclab_experimental.envs.manager_based_env_warp`` for the initialization pattern.
+Random noise kernels (gaussian, uniform) draw from the process-wide per-env Warp RNG state
+:attr:`isaaclab.utils.seed.WarpRng.state`.
 """
 
 from __future__ import annotations
@@ -22,6 +20,8 @@
 
 import warp as wp
 
+from isaaclab.utils.seed import WarpRng
+
 if TYPE_CHECKING:
     from . import noise_cfg
 
@@ -108,7 +108,7 @@ def uniform_noise(data: wp.array, cfg: noise_cfg.UniformNoiseCfg) -> None:
     wp.launch(
         _apply_uniform_noise,
         dim=data.shape[0],
-        inputs=[data, cfg.rng_state_wp, float(cfg.n_min), float(cfg.n_max), _OPERATION_MAP[cfg.operation]],
+        inputs=[data, WarpRng.state, float(cfg.n_min), float(cfg.n_max), _OPERATION_MAP[cfg.operation]],
         device=data.device,
     )
 
@@ -149,7 +149,7 @@ def gaussian_noise(data: wp.array, cfg: noise_cfg.GaussianNoiseCfg) -> None:
     wp.launch(
         _apply_gaussian_noise,
         dim=data.shape[0],
-        inputs=[data, cfg.rng_state_wp, float(cfg.mean), float(cfg.std), _OPERATION_MAP[cfg.operation]],
+        inputs=[data, WarpRng.state, float(cfg.mean), float(cfg.std), _OPERATION_MAP[cfg.operation]],
         device=data.device,
     )
 
diff --git a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py
index e93e73267938..e910f2d9fc1a 100644
--- a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py
+++ b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py
@@ -12,6 +12,8 @@
 import torch
 import warp as wp
 
+from isaaclab.utils.seed import WarpRng
+
 # Skip entire module if no CUDA device available
 wp.init()
 pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required")
@@ -44,6 +46,7 @@ def _clear_function_caches():
     between tests to avoid stale references from prior fixtures.
     """
     yield
+    WarpRng.state = None
     for fn in (
         warp_evt.push_by_setting_velocity,
         warp_evt.apply_external_force_torque,
@@ -102,7 +105,7 @@ class _Env:
     env.step_dt = 0.02
     env.max_episode_length_s = 10.0
     # RNG state for events (seeded deterministically)
-    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
diff --git a/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py
index 0f014cc2b3fd..91a5f8fe138f 100644
--- a/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py
+++ b/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py
@@ -12,6 +12,8 @@
 import torch
 import warp as wp
 
+from isaaclab.utils.seed import WarpRng
+
 # Skip entire module if no CUDA device available
 wp.init()
 pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required")
@@ -50,6 +52,7 @@
 @pytest.fixture(autouse=True)
 def _clear_caches():
     yield
+    WarpRng.state = None
     for fn in [warp_obs.generated_commands]:
         for attr in list(vars(fn)):
             if attr.startswith("_"):
@@ -130,7 +133,7 @@ class _Env:
     env.episode_length_buf = episode_length_buf
     env.step_dt = 0.02
     env.max_episode_length_s = 10.0
-    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
@@ -169,7 +172,7 @@ class _Env:
     env.step_dt = 0.02
     env.max_episode_length = 500
     env.max_episode_length_s = 10.0
-    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
diff --git a/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py
index 684b6f052dbd..5fe1982f883d 100644
--- a/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py
+++ b/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py
@@ -12,6 +12,8 @@
 import torch
 import warp as wp
 
+from isaaclab.utils.seed import WarpRng
+
 # Skip entire module if no CUDA device available
 wp.init()
 pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required")
@@ -56,6 +58,7 @@
 @pytest.fixture(autouse=True)
 def _clear_caches():
     yield
+    WarpRng.state = None
     for fn in [warp_rew.track_lin_vel_xy_exp, warp_rew.track_ang_vel_z_exp, warp_rew.undesired_contacts]:
         for attr in list(vars(fn)):
             if attr.startswith("_"):
@@ -142,7 +145,7 @@ class _Env:
     env.episode_length_buf = episode_length_buf
     env.step_dt = 0.02
     env.max_episode_length_s = 10.0
-    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
@@ -182,7 +185,7 @@ class _Env:
     env.step_dt = 0.02
     env.max_episode_length = 500
     env.max_episode_length_s = 10.0
-    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
diff --git a/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py
index 9676f2b06044..e0fbee333530 100644
--- a/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py
+++ b/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py
@@ -14,6 +14,8 @@
 import torch
 import warp as wp
 
+from isaaclab.utils.seed import WarpRng
+
 # Skip entire module if no CUDA device available
 wp.init()
 pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required")
@@ -59,6 +61,7 @@
 @pytest.fixture(autouse=True)
 def _clear_caches():
     yield
+    WarpRng.state = None
     for fn in [warp_term.illegal_contact]:
         for attr in list(vars(fn)):
             if attr.startswith("_"):
@@ -139,7 +142,7 @@ class _Env:
     env.episode_length_buf = episode_length_buf
     env.step_dt = 0.02
     env.max_episode_length_s = 10.0
-    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
@@ -179,7 +182,7 @@ class _Env:
     env.step_dt = 0.02
     env.max_episode_length = 500
     env.max_episode_length_s = 10.0
-    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
diff --git a/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py b/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py
index 9996cf39f108..3ed68b3cee8c 100644
--- a/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py
+++ b/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py
@@ -40,6 +40,7 @@
 from isaaclab.sim import SimulationCfg
 from isaaclab.terrains import TerrainImporterCfg
 from isaaclab.utils import configclass
+from isaaclab.utils.seed import WarpRng
 
 SENSOR_HEIGHT = 2.0
 RAY_OFFSET = 0.2
@@ -208,6 +209,9 @@ def test_partial_reset_resamples_drift_of_selected_envs_only(sim):
     sensor: NewtonRaycastSensor = scene["raycast"]
     sensor.reset()
 
+    # A masked reset draws from the process-wide Warp RNG state, which an environment initializes.
+    WarpRng.initialize(scene.num_envs, sensor.device)
+
     drift = sensor.drift.torch.clone()
     ray_cast_drift = sensor.ray_cast_drift.torch.clone()
     sensor.reset(env_mask=wp.array([False, True, False], dtype=wp.bool, device=sensor.device))
diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-scene-mask-native.changed.rst b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
new file mode 100644
index 000000000000..59b2cee59f4f
--- /dev/null
+++ b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
@@ -0,0 +1,2 @@
+* Changed the Cartpole, locomotion and reorient Warp environments to draw from the process-wide
+  :class:`~isaaclab.utils.seed.WarpRng` state instead of allocating and seeding their own.
diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py
index d4379ad13c9a..5eac6025c18d 100644
--- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py
+++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py
@@ -11,6 +11,8 @@
 from isaaclab_experimental.envs import DirectRLEnvWarp
 from isaaclab_experimental.utils.warp.utils import wrap_to_pi
 
+from isaaclab.utils.seed import WarpRng
+
 if TYPE_CHECKING:
     from isaaclab_tasks.core.cartpole.cartpole_direct_env_cfg import CartpoleEnvCfg
 
@@ -174,16 +176,6 @@ def reset(
         state[env_index] = rng_state
 
 
-@wp.kernel
-def initialize_state(
-    state: wp.array(dtype=wp.uint32),
-    seed: wp.int32,
-):
-    """Initialize each env's random number generator state from the seed."""
-    env_index = wp.tid()
-    state[env_index] = wp.rand_init(seed, env_index)
-
-
 class CartpoleWarpEnv(DirectRLEnvWarp):
     cfg: CartpoleEnvCfg
 
@@ -206,19 +198,6 @@ def __init__(self, cfg: CartpoleEnvCfg, render_mode: str | None = None, **kwargs
         self.observations = wp.zeros((self.num_envs), dtype=wp.vec4f, device=self.device)
         self.actions = wp.zeros((self.num_envs, self.cartpole.num_joints), dtype=wp.float32, device=self.device)
         self.rewards = wp.zeros((self.num_envs), dtype=wp.float32, device=self.device)
-        self.states = wp.zeros((self.num_envs), dtype=wp.uint32, device=self.device)
-
-        if self.cfg.seed is None:
-            self.cfg.seed = -1
-
-        wp.launch(
-            initialize_state,
-            dim=self.num_envs,
-            inputs=[
-                self.states,
-                self.cfg.seed,
-            ],
-        )
 
         # Bind torch buffers to warp buffers
         self.torch_obs_buf = wp.to_torch(self.observations)
@@ -317,6 +296,6 @@ def _reset_idx(self, mask: wp.array | None = None) -> None:
                 self.cfg.initial_pole_angle_range,
                 self.cfg.initial_pole_velocity_range,
                 mask,
-                self.states,
+                WarpRng.state,
             ],
         )
diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py
index dfcfad6e0f55..b7852ac644e0 100644
--- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py
+++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py
@@ -9,6 +9,7 @@
 from isaaclab_experimental.envs import DirectRLEnvWarp
 
 from isaaclab.envs import DirectRLEnvCfg
+from isaaclab.utils.seed import WarpRng
 from isaaclab.utils.string import resolve_matching_names_values
 
 
@@ -416,15 +417,12 @@ def survival_rate(
 
 
 @wp.kernel
-def initialize_state(
+def initialize_targets(
     env_origins: wp.array(dtype=wp.vec3f),
     targets: wp.array(dtype=wp.vec3f),
-    state: wp.array(dtype=wp.uint32),
-    seed: wp.int32,
 ):
-    """Initialize each env's RNG state and place its target far along +x from the env origin."""
+    """Place each env's target far along +x from the env origin."""
     env_index = wp.tid()
-    state[env_index] = wp.rand_init(seed, env_index)
     targets[env_index] = env_origins[env_index]
     targets[env_index] += wp.static(wp.vec3f(1000.0, 0.0, 0.0))
 
@@ -480,7 +478,6 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
         )
         self.rewards = wp.zeros((self.num_envs), dtype=wp.float32, device=self.sim.device)
         self.actions = wp.zeros((self.num_envs, self.robot.num_joints), dtype=wp.float32, device=self.sim.device)
-        self.states = wp.zeros((self.num_envs), dtype=wp.uint32, device=self.sim.device)
         self.potentials = wp.zeros(self.num_envs, dtype=wp.float32, device=self.sim.device)
         self.prev_potentials = wp.zeros_like(self.potentials)
         self.targets = wp.zeros((self.num_envs), dtype=wp.vec3f, device=self.sim.device)
@@ -496,18 +493,13 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
         self.env_origins = wp.from_torch(self.scene.env_origins, dtype=wp.vec3f)
         self.actions_mapped = wp.zeros((self.num_envs, self.robot.num_joints), dtype=wp.float32, device=self.sim.device)
 
-        # Initial states and targets
-        if self.cfg.seed is None:
-            self.cfg.seed = -1
-
+        # Initial targets
         wp.launch(
-            initialize_state,
+            initialize_targets,
             dim=self.num_envs,
             inputs=[
                 self.env_origins,
                 self.targets,
-                self.states,
-                self.cfg.seed,
             ],
         )
 
@@ -684,7 +676,7 @@ def _reset_idx(self, mask: wp.array | None = None):
                 self.soft_joint_vel_limits,
                 self.joint_pos,
                 self.joint_vel,
-                self.states,
+                WarpRng.state,
                 self.cfg.initial_joint_pos_range[0],
                 self.cfg.initial_joint_pos_range[1],
                 self.cfg.initial_joint_vel_range[0],
diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
index c2a34edc9da6..f98a9e0e3375 100644
--- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
+++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
@@ -13,22 +13,12 @@
 import warp as wp
 from isaaclab_experimental.envs import DirectRLEnvWarp
 
+from isaaclab.utils.seed import WarpRng
+
 if TYPE_CHECKING:
     from isaaclab_tasks.core.reorient.config.allegro_hand.allegro_hand_direct_env_cfg import AllegroHandEnvCfg
 
 
-@wp.kernel
-def initialize_rng_state(
-    # input
-    seed: wp.int32,
-    # output
-    state: wp.array(dtype=wp.uint32),
-):
-    """Initialize each env's random number generator state from the seed."""
-    env_id = wp.tid()
-    state[env_id] = wp.rand_init(seed, wp.int32(env_id))
-
-
 @wp.kernel
 def apply_actions_to_targets(
     # input
@@ -647,20 +637,6 @@ def __init__(self, cfg: AllegroHandEnvCfg, render_mode: str | None = None, **kwa
         self.object_pose = wp.zeros(self.num_envs, dtype=wp.transformf, device=self.device)
         self.object_vels = wp.zeros(self.num_envs, dtype=wp.spatial_vectorf, device=self.device)
 
-        # RNG state (per-env) for randomizations in reset/goal resets.
-        self.rng_state = wp.zeros(self.num_envs, dtype=wp.uint32, device=self.device)
-        if self.cfg.seed is None:
-            self.cfg.seed = -1
-        wp.launch(
-            initialize_rng_state,
-            dim=self.num_envs,
-            inputs=[
-                self.cfg.seed,
-                self.rng_state,
-            ],
-            device=self.device,
-        )
-
         # ---------------------------------------------------------------------
         # Torch views / aliases
         # ---------------------------------------------------------------------
@@ -810,7 +786,7 @@ def _reset_idx(self, mask: wp.array | None = None):
                 self.x_unit_vec,
                 self.y_unit_vec,
                 mask,
-                self.rng_state,
+                WarpRng.state,
                 self.object.data.root_link_pose_w.warp,
                 self.object.data.root_com_vel_w.warp,
             ],
@@ -830,7 +806,7 @@ def _reset_idx(self, mask: wp.array | None = None):
                 self.cfg.reset_dof_vel_noise,
                 mask,
                 self.num_hand_dofs,
-                self.rng_state,
+                WarpRng.state,
                 self.hand.data.joint_pos.warp,
                 self.hand.data.joint_vel.warp,
                 self.prev_targets,
@@ -874,7 +850,7 @@ def _reset_target_pose(self, env_ids: Sequence[int] | None = None, mask: wp.arra
                 self.y_unit_vec,
                 self.env_origins,
                 self.goal_pos,
-                self.rng_state,
+                WarpRng.state,
                 self.goal_rot,
                 self.reset_goal_buf,
                 self.goal_pos_w,

From 69c13a6fafe9c038ba9b518f42ff3f2a3ced7522 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Thu, 1 Oct 2026 11:50:15 -0700
Subject: [PATCH 32/48] Point the Warp migration guide at WarpRng.state

---
 docs/source/concepts/warp_environment_migration.rst | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/docs/source/concepts/warp_environment_migration.rst b/docs/source/concepts/warp_environment_migration.rst
index f1c3c5fc9067..41d205a16148 100644
--- a/docs/source/concepts/warp_environment_migration.rst
+++ b/docs/source/concepts/warp_environment_migration.rst
@@ -158,7 +158,8 @@ checks the mask to skip non-selected environments:
            return
        # ... modify state for selected envs only
 
-- RNG uses per-env ``env.rng_state_wp`` (``wp.uint32``) instead of ``torch.rand``
+- RNG uses the per-env :attr:`WarpRng.state <isaaclab.utils.seed.WarpRng.state>` (``wp.uint32``) instead of
+  ``torch.rand``
 - **Startup/prestartup** events use the torch convention ``(env, env_ids, **params)``
 - **Reset/interval** events use the warp convention ``(env, env_mask, **params)``
 

From f8463720c06a033d2fed8f9e07c4ec442ff1e3b3 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Thu, 1 Oct 2026 11:54:09 -0700
Subject: [PATCH 33/48] Allocate the scratch body mask in the Newton collection
 test helper

---
 .../test/assets/_rigid_object_collection_iface_test_utils.py     | 1 +
 1 file changed, 1 insertion(+)

diff --git a/source/isaaclab/test/assets/_rigid_object_collection_iface_test_utils.py b/source/isaaclab/test/assets/_rigid_object_collection_iface_test_utils.py
index e7593e8c6b8e..7b89bf69762a 100644
--- a/source/isaaclab/test/assets/_rigid_object_collection_iface_test_utils.py
+++ b/source/isaaclab/test/assets/_rigid_object_collection_iface_test_utils.py
@@ -219,6 +219,7 @@ def create_newton_rigid_object_collection(
     object.__setattr__(collection, "_ALL_BODY_INDICES", wp.array(np.arange(num_bodies, dtype=np.int32), device=device))
     object.__setattr__(collection, "_ALL_ENV_MASK", wp.ones((num_instances,), dtype=wp.bool, device=device))
     object.__setattr__(collection, "_ALL_BODY_MASK", wp.ones((num_bodies,), dtype=wp.bool, device=device))
+    object.__setattr__(collection, "_scratch_body_mask", wp.zeros((num_bodies,), dtype=wp.bool, device=device))
 
     return collection, mock_view
 

From 11e36c47629ad41f8678276337cb9105bf3697c1 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Thu, 1 Oct 2026 12:23:28 -0700
Subject: [PATCH 34/48] Reset Multirotor thrusters and test actuator stubs from
 the environment mask

Articulation resets now forward env_mask to their actuators. The Multirotor thruster collection and the test_reset actuator stubs still took only env_ids and raised TypeError; the collection now resets the masked environments' thrusters.
---
 .../jichuanh-warp-scene-mask-native.fixed.rst        |  2 ++
 .../isaaclab_contrib/assets/multirotor/multirotor.py |  9 ++++++++-
 .../isaaclab_contrib/test/assets/test_multirotor.py  | 12 ++++++++++++
 .../isaaclab_newton/test/assets/test_articulation.py |  4 ++--
 .../isaaclab_physx/test/assets/test_articulation.py  |  4 ++--
 5 files changed, 26 insertions(+), 5 deletions(-)
 create mode 100644 source/isaaclab_contrib/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst

diff --git a/source/isaaclab_contrib/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst b/source/isaaclab_contrib/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
new file mode 100644
index 000000000000..177eb50253d4
--- /dev/null
+++ b/source/isaaclab_contrib/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
@@ -0,0 +1,2 @@
+* Fixed :class:`~isaaclab_contrib.assets.Multirotor` resets with only ``env_mask`` resetting the thrusters of every
+  environment.
diff --git a/source/isaaclab_contrib/isaaclab_contrib/assets/multirotor/multirotor.py b/source/isaaclab_contrib/isaaclab_contrib/assets/multirotor/multirotor.py
index 0402d8b8b925..5db699bc3393 100644
--- a/source/isaaclab_contrib/isaaclab_contrib/assets/multirotor/multirotor.py
+++ b/source/isaaclab_contrib/isaaclab_contrib/assets/multirotor/multirotor.py
@@ -42,12 +42,19 @@ class _ThrusterCollection(dict):
     exposing the :meth:`reset` entry point that :meth:`isaaclab.assets.Articulation.reset` invokes.
     """
 
-    def reset(self, env_ids: Sequence[int] | slice | None = None) -> None:
+    def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: wp.array | None = None) -> None:
         """Reset every thruster actuator for the given environments.
 
+        .. caution::
+            If both ``env_ids`` and ``env_mask`` are provided, ``env_mask`` takes precedence.
+
         Args:
             env_ids: Environment indices to reset. Defaults to None (all environments).
+            env_mask: Environment mask. Shape is (num_instances,). Defaults to None.
         """
+        if env_mask is not None:
+            # thrusters index their buffers by environment
+            env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1)
         for actuator in self.values():
             actuator.reset(env_ids)
 
diff --git a/source/isaaclab_contrib/test/assets/test_multirotor.py b/source/isaaclab_contrib/test/assets/test_multirotor.py
index 5d195eeffe7a..d0bf04b2baf4 100644
--- a/source/isaaclab_contrib/test/assets/test_multirotor.py
+++ b/source/isaaclab_contrib/test/assets/test_multirotor.py
@@ -18,12 +18,14 @@
 
 import pytest
 import torch
+import warp as wp
 from isaaclab_physx.physics import PhysxCfg
 
 import isaaclab.sim.utils.prims as prim_utils
 from isaaclab.sim import SimulationCfg, build_simulation_context
 
 from isaaclab_contrib.assets import Multirotor, MultirotorCfg
+from isaaclab_contrib.assets.multirotor.multirotor import _ThrusterCollection
 
 # Best-effort: suppress unraisable destructor warnings emitted during
 # teardown of partially-constructed assets in CI/dev environments. We still
@@ -156,6 +158,16 @@ def test_set_thrust_target_env_slice_unit():
     assert torch.allclose(m._data.thrust_target[3], original[3])
 
 
+def test_thruster_collection_resets_masked_envs_unit():
+    """A mask-only reset, as the articulation forwards it, resets the thrusters of the masked envs only."""
+    reset_env_ids = []
+    thrusters = _ThrusterCollection(thrusters=types.SimpleNamespace(reset=reset_env_ids.append))
+
+    thrusters.reset(env_mask=wp.array([True, False, True, False], dtype=wp.bool, device="cpu"))
+
+    assert [env_ids.tolist() for env_ids in reset_env_ids] == [[0, 2]]
+
+
 def generate_multirotor(
     multirotor_cfg: MultirotorCfg, num_multirotors: int, device: str
 ) -> tuple[Multirotor, torch.Tensor]:
diff --git a/source/isaaclab_newton/test/assets/test_articulation.py b/source/isaaclab_newton/test/assets/test_articulation.py
index 918660767818..87259a04500a 100644
--- a/source/isaaclab_newton/test/assets/test_articulation.py
+++ b/source/isaaclab_newton/test/assets/test_articulation.py
@@ -2163,9 +2163,9 @@ def test_reset(sim, num_articulations, device, articulation_type, monkeypatch):
     actuator_reset = actuator.reset
     reset_env_ids = []
 
-    def record_actuator_reset(env_ids=None):
+    def record_actuator_reset(env_ids=None, env_mask=None):
         reset_env_ids.append(env_ids)
-        actuator_reset(env_ids)
+        actuator_reset(env_ids, env_mask=env_mask)
 
     monkeypatch.setattr(actuator, "reset", record_actuator_reset)
     articulation.reset()
diff --git a/source/isaaclab_physx/test/assets/test_articulation.py b/source/isaaclab_physx/test/assets/test_articulation.py
index e0a81ef3aa97..969ce2692127 100644
--- a/source/isaaclab_physx/test/assets/test_articulation.py
+++ b/source/isaaclab_physx/test/assets/test_articulation.py
@@ -1516,9 +1516,9 @@ def test_reset(sim, num_articulations, device, monkeypatch):
     actuator_reset = actuator.reset
     reset_env_ids = []
 
-    def record_actuator_reset(env_ids=None):
+    def record_actuator_reset(env_ids=None, env_mask=None):
         reset_env_ids.append(env_ids)
-        actuator_reset(env_ids)
+        actuator_reset(env_ids, env_mask=env_mask)
 
     monkeypatch.setattr(actuator, "reset", record_actuator_reset)
 

From 9328d9813809ebbb7cf437ab1ecd621f624f7a56 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Thu, 1 Oct 2026 12:23:29 -0700
Subject: [PATCH 35/48] Defer annotations in the seed module so the docs import
 it with Warp mocked

---
 source/isaaclab/isaaclab/utils/seed.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/source/isaaclab/isaaclab/utils/seed.py b/source/isaaclab/isaaclab/utils/seed.py
index 5023e8a6237c..fda91ad9734e 100644
--- a/source/isaaclab/isaaclab/utils/seed.py
+++ b/source/isaaclab/isaaclab/utils/seed.py
@@ -3,6 +3,8 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
+from __future__ import annotations
+
 import os
 import random
 

From ee9e8adbb9dc6524f34b53dfd99db858213abffa Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Fri, 2 Oct 2026 15:57:44 -0700
Subject: [PATCH 36/48] Record Warp stages per call and step the Warp
 environments like stable

- Replace WarpGraphCache with CapturedStage, @captured and eager(): a stage records on its first call and
  replays afterwards; a call that cannot be recorded runs eagerly at its place in the recording.
- Call the managers in the stable environments' order; remove ManagerCallSwitch and the step_warp_* entry
  points. Recordings pause the garbage collector and collect the youngest generation afterwards.
- Decide DirectRLEnvWarp resets on the GPU with a conditional graph node, so a step no longer synchronizes
  with the host.
- Take a torch action in ActionManager.process_action, as stable does.
- Judge a class term's reset by the annotation on its reset method; a class annotation covers its call.
- Turn the stateful Warp event terms into classes; the CoM term copies its defaults behind a device flag.
- Read interval ranges from the term configurations, so a replaced range applies.
- Keep a dedicated reset mask on the base environment; drop redundant float()/int() conversions.
---
 .../isaaclab_experimental.utils.rst           |  10 +-
 .../jichuanh-warp-frontend-cleanup.added.rst  |  17 +-
 ...jichuanh-warp-frontend-cleanup.changed.rst |  31 +-
 .../jichuanh-warp-frontend-cleanup.fixed.rst  |   2 +-
 ...r => jichuanh-warp-frontend-cleanup.major} |   0
 .../envs/direct_rl_env_warp.py                | 176 +++++------
 .../isaaclab_experimental/envs/frontend.py    |   4 +-
 .../envs/manager_based_env_warp.py            | 101 ++----
 .../envs/manager_based_rl_env_warp.py         | 241 ++++----------
 .../isaaclab_experimental/envs/mdp/events.py  |  79 ++---
 .../isaaclab_experimental/envs/mdp/rewards.py |  12 +-
 .../envs/mdp/terminations.py                  |  27 +-
 .../managers/action_manager.py                | 137 ++++----
 .../managers/event_manager.py                 | 281 +++++++----------
 .../managers/manager_base.py                  | 116 +++----
 .../managers/observation_manager.py           | 129 ++++----
 .../managers/reward_manager.py                | 152 +++------
 .../managers/scene_entity_cfg.py              |  22 +-
 .../managers/termination_manager.py           | 124 ++------
 .../isaaclab_experimental/utils/__init__.pyi  |   6 +-
 .../utils/warp/__init__.py                    |   2 +-
 .../utils/warp/kernels.py                     |   8 -
 .../isaaclab_experimental/utils/warp/utils.py |  13 +-
 .../utils/warp_capture.py                     | 298 ++++++++++++++++++
 .../utils/warp_graph_cache.py                 | 268 ----------------
 .../test/envs/mdp/parity_helpers.py           |   4 +-
 .../test/envs/mdp/test_capture_safety.py      |   2 +-
 .../test/envs/mdp/test_events_warp_parity.py  |  66 +++-
 .../test/envs/test_frontend.py                |   2 +-
 .../test/envs/test_warp_env_capture.py        |  55 +++-
 .../test/envs/test_warp_env_decimation.py     |  19 +-
 .../test/envs/test_warp_env_reset.py          |   3 +-
 .../test/managers/test_event_manager.py       |  77 ++++-
 .../test/managers/test_reward_manager.py      |  38 +--
 .../test/utils/test_mask_boundary_gate.py     |  13 +-
 .../test/utils/test_warp_capture.py           | 241 ++++++++++++++
 .../test/utils/test_warp_graph_cache.py       | 212 -------------
 ...jichuanh-warp-frontend-cleanup.changed.rst |   4 +
 .../core/cartpole/cartpole_warp_env.py        |   6 +
 .../core/cartpole/mdp/rewards.py              |   4 +-
 .../core/locomotion/locomotion_env_warp.py    |   6 +
 .../core/reorient/reorient_warp_env.py        |   6 +
 42 files changed, 1393 insertions(+), 1621 deletions(-)
 rename source/isaaclab_experimental/changelog.d/{jichuanh-warp-frontend-cleanup.minor => jichuanh-warp-frontend-cleanup.major} (100%)
 create mode 100644 source/isaaclab_experimental/isaaclab_experimental/utils/warp_capture.py
 delete mode 100644 source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py
 create mode 100644 source/isaaclab_experimental/test/utils/test_warp_capture.py
 delete mode 100644 source/isaaclab_experimental/test/utils/test_warp_graph_cache.py
 create mode 100644 source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst

diff --git a/docs/source/api/lab_experimental/isaaclab_experimental.utils.rst b/docs/source/api/lab_experimental/isaaclab_experimental.utils.rst
index 4f00a75aeff9..bb9a863b2b4f 100644
--- a/docs/source/api/lab_experimental/isaaclab_experimental.utils.rst
+++ b/docs/source/api/lab_experimental/isaaclab_experimental.utils.rst
@@ -13,7 +13,13 @@ The following classes are part of the public :mod:`isaaclab_experimental.utils`
 .. autosummary::
    :nosignatures:
 
-   WarpGraphCache
+   CapturedStage
+   captured
+   eager
 
-.. autoclass:: WarpGraphCache
+.. autoclass:: CapturedStage
    :show-inheritance:
+
+.. autofunction:: captured
+
+.. autofunction:: eager
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst
index 37c0aacc0377..03346fb5b32e 100644
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst
@@ -1,12 +1,15 @@
-* Added the ``ISAACLAB_WARP_CAPTURE`` environment variable: ``ISAACLAB_WARP_CAPTURE=0`` runs every Warp
-  environment stage eagerly, for debugging and eager-versus-captured comparisons.
+* Added :class:`~isaaclab_experimental.utils.CapturedStage` and the :func:`~isaaclab_experimental.utils.captured`
+  decorator, which run a Warp manager or environment stage eagerly or through CUDA graphs recorded from its own
+  calls. A stage records on its first call without an eager warm-up, records again when an argument array or
+  scalar changes or the physics backend rebinds its buffers, and runs eagerly on CPU or while
+  :attr:`CapturedStage.enabled <isaaclab_experimental.utils.CapturedStage.enabled>` is False. The Warp
+  environments enable capture.
+* Added :func:`~isaaclab_experimental.utils.eager` for calls that cannot be recorded inside a captured stage.
+  While the stage records, the call ends the current CUDA graph and runs eagerly before the next graph starts.
+  Every replay repeats the call at the same place with the same arguments.
 * Added the ``ISAACLAB_SYNC_DEBUG`` environment variable: ``ISAACLAB_SYNC_DEBUG=1`` raises on a hidden
   GPU-to-host synchronization inside an eagerly executed Warp environment stage.
 * Added a boolean ``env_mask`` argument to :meth:`~isaaclab_experimental.envs.ManagerBasedEnvWarp.reset`.
-* Added :meth:`~isaaclab_experimental.utils.WarpGraphCache.call_steps` and
-  :meth:`ManagerBase.stage_steps <isaaclab_experimental.managers.ManagerBase.stage_steps>`, which split a
-  manager stage into per-term steps that are recorded into CUDA graphs or run eagerly.
 * Added parameter predicates to :class:`~isaaclab_experimental.utils.warp.WarpCapturable`:
   ``@WarpCapturable(lambda params: ...)`` marks a term capturable only for the parameters the predicate
-  accepts. A class term can also set ``self._warp_capturable`` in ``__init__``. The Warp managers decide
-  capture for each configured term instance.
+  accepts.
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst
index 895b48a3a557..7eb3c5d84f4a 100644
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst
@@ -1,20 +1,29 @@
 * **Breaking:** Removed ``ManagerCallSwitch``, ``ManagerCallMode`` and the ``MANAGER_CALL_CONFIG``
   environment variable. Warp environments construct their Warp managers directly.
-* Changed the Warp environments to decide CUDA graph capture per term: a stage records its capturable
-  terms and runs each term decorated with ``@WarpCapturable(False)`` eagerly, instead of running the whole
-  stage eagerly. Reward and termination stages evaluate their capturable terms first; observation, action
-  and event stages keep the configured term order.
-* **Breaking:** Removed ``ManagerBasedRLEnvWarp.step_warp_termination_compute``. The termination stage runs
+* **Breaking:** Removed ``WarpGraphCache``. The Warp managers and environments record their stages through
+  :class:`~isaaclab_experimental.utils.CapturedStage`, and the environments call the managers as the stable
+  environments do.
+* **Breaking:** Removed ``ManagerBasedRLEnvWarp.step_warp_termination_compute``. The environment calls
   :meth:`TerminationManager.compute <isaaclab_experimental.managers.TerminationManager.compute>` directly.
-* **Breaking:** Replaced ``WarpGraphCache.capture_or_replay`` with
-  :meth:`~isaaclab_experimental.utils.WarpGraphCache.call`. Stages are recorded from the first environment
-  step without an eager warm-up, so a stage runs exactly once per call.
+* **Breaking:** Removed ``DirectRLEnvWarp.step_warp_action``. :class:`~isaaclab_experimental.envs.DirectRLEnvWarp`
+  calls the task methods in the order of the stable :class:`~isaaclab.envs.DirectRLEnv` and no longer records them.
+  To keep recording them as CUDA graphs, decorate a task's ``_apply_action``, ``_get_dones``, ``_get_rewards``,
+  ``_get_observations`` and ``_reset_idx`` with :func:`~isaaclab_experimental.utils.captured`.
+* **Breaking:** Changed the Warp
+  :meth:`ActionManager.process_action <isaaclab_experimental.managers.ActionManager.process_action>` to take a
+  torch tensor, as the stable :meth:`~isaaclab.managers.ActionManager.process_action` does.
+* Changed the Warp managers to decide CUDA graph capture per term: an operation records its capturable terms
+  and runs each term decorated with ``@WarpCapturable(False)``, and each observation term with history, eagerly
+  between the recorded parts, instead of running the whole operation eagerly. A class term's ``reset`` is
+  recorded unless the ``reset`` method itself is decorated with ``@WarpCapturable(False)``.
 * Changed the Warp manager-based environments to reset environments by boolean mask. The mask is converted
   to environment indices once per reset, and only when a command, curriculum or recorder term is active.
 * Changed :class:`~isaaclab_experimental.envs.DirectRLEnvWarp` to skip its reset stage on steps where no
-  environment terminated.
+  environment terminated, deciding on the GPU (a conditional CUDA graph node, CUDA 12.4+), so a step no longer
+  synchronizes with the host. With ``compute_final_obs``, ``extras["final_obs"]`` is now set on every step.
 * Changed the Warp ``push_by_setting_velocity``, ``apply_external_force_torque``, ``reset_root_state_uniform``
   and ``randomize_rigid_body_com`` event terms to classes that allocate their buffers when the term is
   created. They read their ranges on every call, so a changed range applies to the next event.
-* Changed the Warp ``randomize_rigid_body_material`` and ``randomize_rigid_body_mass`` adapters to reject
-  event modes other than ``startup``.
+* **Breaking:** Renamed :attr:`SceneEntityCfg.joint_mask <isaaclab_experimental.managers.SceneEntityCfg.joint_mask_wp>`
+  to ``joint_mask_wp``, matching the ``_wp`` suffix of its other Warp fields, and added ``body_mask_wp``, the
+  boolean mask of the selected bodies.
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
index f2186f9b5823..f295460afdc2 100644
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst
@@ -1,5 +1,5 @@
 * Fixed Warp environments replaying CUDA graphs that read freed simulation buffers after the simulation was
-  reset or stopped. The environments now record their stages again on the next step.
+  reset or stopped. The stages now record again on their next call.
 * Fixed the Warp environments applying actions, stepping the simulation and updating the scene once per physics
   step when the physics backend runs the whole decimation loop. Like the stable environments, they now do it
   once per environment step, so contact histories, air times and the rewards and terminations reading them
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.major
similarity index 100%
rename from source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor
rename to source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.major
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
index 04a86b5c9202..363893654612 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
@@ -9,7 +9,6 @@
 import inspect
 import logging
 import math
-import os
 from abc import abstractmethod
 from dataclasses import MISSING
 from typing import Any, ClassVar
@@ -35,7 +34,8 @@
 from isaaclab.utils.seed import WarpRng, configure_seed
 from isaaclab.utils.timer import Timer
 
-from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache
+from isaaclab_experimental.utils.warp import count_masked
+from isaaclab_experimental.utils.warp_capture import CapturedStage, captured, reset_captured_stages
 
 # from isaacsim.core.simulation_manager import SimulationManager
 # from isaacsim.core.version import get_version
@@ -44,12 +44,6 @@
 # import logger
 logger = logging.getLogger(__name__)
 
-DEBUG_TIMER_STEP = os.environ.get("DEBUG_TIMER_STEP", "0") == "1"
-"""Enable outer step() timer only. Set DEBUG_TIMER_STEP=1 env var to enable."""
-
-DEBUG_TIMERS = os.environ.get("DEBUG_TIMERS", "0") == "1"
-"""Enable all fine-grained inner timers (adds wp.synchronize per sub-phase). Set DEBUG_TIMERS=1 env var to enable."""
-
 
 @wp.kernel
 def zero_mask_int32(
@@ -86,6 +80,10 @@ class DirectRLEnvWarp(DirectRLEnv):
     here and leave it up to library-defined wrappers to take care of wrapping this
     environment for their agents.
 
+    Tasks record their GPU work as CUDA graphs by decorating :meth:`_apply_action`, :meth:`_get_dones`,
+    :meth:`_get_rewards`, :meth:`_get_observations` and :meth:`_reset_idx` with
+    :func:`~isaaclab_experimental.utils.captured`.
+
     Note:
         For vectorized environments, it is recommended to **only** call the :meth:`reset`
         method once before the first call to :meth:`step`, i.e. after the environment is created.
@@ -123,7 +121,6 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
         self.render_mode = render_mode
         # initialize internal variables
         self._is_closed = False
-        self._physics_handles_decimation = False
 
         # set the seed for the environment
         if self.cfg.seed is not None:
@@ -229,7 +226,8 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
         self.reset_terminated = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device)
         self.reset_time_outs = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device)
         self.reset_buf = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device)
-        self._torch_reset_buf = wp.to_torch(self.reset_buf)
+        # the number of environments to reset, counted on the device so the step does not synchronize with the host
+        self._reset_count_wp = wp.zeros(1, dtype=wp.int32, device=self.device)
         self._ALL_ENV_MASK = wp.ones(self.num_envs, dtype=wp.bool, device=self.device)
 
         # Expected bindings:
@@ -239,9 +237,9 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
         self.torch_reset_time_outs: torch.Tensor = None
         self.torch_episode_length_buf: torch.Tensor = None
 
-        # Task stages run through this cache and record their CUDA graphs on the first step.
-        self._warp_graph_cache = WarpGraphCache(self.device)
-        self._warp_graph_cache.invalidate_on(self.sim.physics_manager)
+        # the task stages record CUDA graphs, and record again after the physics buffers rebind
+        CapturedStage.enabled = True
+        CapturedStage.invalidate_on(self.sim.physics_manager)
 
         # setup the action and observation spaces for Gym
         self._configure_gym_env_spaces()
@@ -378,7 +376,6 @@ def reset(self, seed: int | None = None, options: dict[str, Any] | None = None)
         self.obs_buf = {"policy": self.torch_obs_buf.clone()}
         return self.obs_buf, self.extras
 
-    @Timer(name="env_step", msg="Step took:", enable=DEBUG_TIMER_STEP or DEBUG_TIMERS)
     def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         """Execute one time-step of the environment's dynamics.
 
@@ -404,8 +401,6 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
             A tuple containing the observations, rewards, resets (terminated and truncated) and extras.
         """
 
-        # stages start recording at the first step after construction or an invalidation
-        self._warp_graph_cache.arm()
         # task code that runs during this step writes into its own log, not into the copy returned last step
         if self._task_log is not None:
             self.extras["log"] = self._task_log
@@ -414,11 +409,8 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         if self.cfg.action_noise_model:
             action = self._action_noise_model(action)
 
-        # process actions, #TODO pass the torch tensor directly.
-        with Timer(name="pre_physics", msg="Pre-physics step took:", enable=DEBUG_TIMERS):
-            self._pre_physics_step(
-                wp.from_torch(action)
-            )  # Creates a tensor and discards it. Not graphable unless training loop reuses the same pointer.
+        # process actions
+        self._pre_physics_step(wp.from_torch(action))
 
         # check if we need to do rendering within the physics loop
         # note: hoisted out of the decimation loop; is_rendering does live settings lookups
@@ -426,51 +418,57 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
 
         # physics-owned decimation covers all substeps in one call
         steps_per_call = self.cfg.decimation if self._physics_handles_decimation else 1
-        with Timer(name="physics_loop", msg="Physics loop took:", enable=DEBUG_TIMERS):
-            for _ in range(self.cfg.decimation // steps_per_call):
-                self._sim_step_counter += steps_per_call
-                # set actions into buffers
-                # simulate
-                with Timer(name="apply_action", msg="Action processing step took:", enable=DEBUG_TIMERS):
-                    self._warp_graph_cache.call("DirectAction_step", self.step_warp_action)
-
-                # write_data_to_sim runs outside the CUDA graph because _apply_actuator_model
-                # uses torch ops (wp.to_torch + torch arithmetic) that cross CUDA streams.
-                with Timer(name="write_data_to_sim_loop", msg="Write data to sim (loop) took:", enable=DEBUG_TIMERS):
-                    self.scene.write_data_to_sim()
-
-                with Timer(name="simulate", msg="Newton simulation step took:", enable=DEBUG_TIMERS):
-                    self.sim.step(render=False)
-                # render between steps only if the GUI or an RTX sensor needs it
-                # note: we assume the render interval to be the shortest accepted rendering interval.
-                #    If a camera needs rendering at a faster frequency, this will lead to unexpected behavior.
-                if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering:
-                    self.sim.render()
-                with Timer(name="scene_update", msg="Scene update took:", enable=DEBUG_TIMERS):
-                    self.scene.update(dt=self.physics_dt * steps_per_call)
+        for _ in range(self.cfg.decimation // steps_per_call):
+            self._sim_step_counter += steps_per_call
+            self._apply_action()
+            # write_data_to_sim runs outside the CUDA graph because _apply_actuator_model
+            # uses torch ops (wp.to_torch + torch arithmetic) that cross CUDA streams.
+            self.scene.write_data_to_sim()
 
+            self.sim.step(render=False)
+            # render between steps only if the GUI or an RTX sensor needs it
+            # note: we assume the render interval to be the shortest accepted rendering interval.
+            #    If a camera needs rendering at a faster frequency, this will lead to unexpected behavior.
+            if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering:
+                self.sim.render()
+            self.scene.update(dt=self.physics_dt * steps_per_call)
+
+        # post-step:
+        # -- update env counters (used for curriculum generation)
+        wp.launch(add_to_env, dim=self.num_envs, inputs=[self._episode_length_buf_wp, 1])  # step in current episode
         self.common_step_counter += 1  # total step (common for all envs)
-        with Timer(name="end_pre_graph", msg="End pre-graph took:", enable=DEBUG_TIMERS):
-            self._warp_graph_cache.call("DirectEndPre_step", self._step_warp_end_pre)
-        # one host predicate per step skips the reset stage when no environment terminated
-        if self._torch_reset_buf.any().item():
-            # capture the terminal observation before reset and expose it for Same-Step autoreset.
-            # the reset overwrites the persistent observation buffer, so keep a copy
-            if self.cfg.compute_final_obs:
-                self._warp_graph_cache.call("DirectFinalObs_step", self._get_observations)
-                self.extras["final_obs"] = {"policy": self.torch_obs_buf.clone()}
-            with Timer(name="reset_graph", msg="Reset graph took:", enable=DEBUG_TIMERS):
-                self._warp_graph_cache.call("DirectReset_step", self._step_warp_reset)
+
+        self._get_dones()
+        self._get_rewards()
+
+        # capture the terminal observation before reset and expose it for Same-Step autoreset.
+        # whether any environment resets is only known on the device, so it is kept on every step
+        if self.cfg.compute_final_obs:
+            self._get_observations()
+            self.extras["final_obs"] = {"policy": self.torch_obs_buf.clone()}
+        # -- reset envs that terminated/timed-out
+        self._reset_terminated()
+        # if sensors are added to the scene, make sure we render to reflect changes in reset
+        # if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset:
+        #    self.sim.render()
         # write_data_to_sim runs uncaptured — it uses torch ops that cross CUDA streams.
-        with Timer(name="write_data_to_sim_post", msg="Write data to sim (post-reset) took:", enable=DEBUG_TIMERS):
-            self.scene.write_data_to_sim()
-        with Timer(name="end_post_graph", msg="End post-graph took:", enable=DEBUG_TIMERS):
-            self._warp_graph_cache.call("DirectEndPost_step", self._step_warp_end_post)
+        self.scene.write_data_to_sim()
+
+        # post-step: step interval event
+        # if self.cfg.events:
+        #    if "interval" in self.event_manager.available_modes:
+        #        self.event_manager.apply(mode="interval", dt=self.step_dt)
+
+        self._get_observations()
+
+        # add observation noise
+        # note: we apply no noise to the state space (since it is used for critic networks)
+        # if self.cfg.observation_noise_model:
+        #    self.obs_buf["policy"] = self._observation_noise_model(self.obs_buf["policy"])
 
         # Visualization hook — runs after CUDA graph scope. Override in subclass
         # to update markers or other non-graphable visual elements.
-        with Timer(name="visualize", msg="Visualize took:", enable=DEBUG_TIMERS):
-            self._post_step_visualize()
+        self._post_step_visualize()
 
         # tasks log persistent buffers that later steps overwrite, and RL libraries keep the log of every step
         # until they average it, so hand out copies like the fresh tensors of stable tasks
@@ -500,50 +498,6 @@ def _post_step_visualize(self) -> None:
         """
         pass
 
-    def step_warp_action(self) -> None:
-        self._apply_action()
-        # Note: scene.write_data_to_sim() is called separately outside the CUDA graph
-        # capture scope because it invokes _apply_actuator_model() which uses torch
-        # arithmetic (wp.to_torch + torch ops). This would cause a CUDA stream crossing
-        # error during graph capture. Moving it outside is safe since it runs every step.
-
-    def _step_warp_end_pre(self) -> None:
-        """Capturable portion before write_data_to_sim (pure warp kernels)."""
-        wp.launch(
-            add_to_env,
-            dim=self.num_envs,
-            inputs=[
-                self._episode_length_buf_wp,
-                1,
-            ],
-        )
-        self._get_dones()
-        self._get_rewards()
-
-    def _step_warp_reset(self) -> None:
-        """Capturable reset of the environments that terminated or timed out."""
-        self._reset_idx(mask=self.reset_buf)
-
-    def _step_warp_end_post(self) -> None:
-        """Capturable portion after write_data_to_sim (pure warp kernels)."""
-        # if sensors are added to the scene, make sure we render to reflect changes in reset
-        # if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset:
-        #    self.sim.render()
-
-        # TODO We could split it out.
-        # post-step: step interval event
-        # if self.cfg.events:
-        #    if "interval" in self.event_manager.available_modes:
-        #        self.event_manager.apply(mode="interval", dt=self.step_dt)
-
-        # update observations
-        self._get_observations()
-
-        # add observation noise
-        # note: we apply no noise to the state space (since it is used for critic networks)
-        # if self.cfg.observation_noise_model:
-        #    self.obs_buf["policy"] = self._observation_noise_model(self.obs_buf["policy"])
-
     @staticmethod
     def seed(seed: int = -1) -> int:
         """Set the seed for the environment.
@@ -611,7 +565,8 @@ def close(self):
             if self.cfg.events:
                 del self.event_manager
             del self.scene
-            self._warp_graph_cache.close()
+            CapturedStage.disable()
+            reset_captured_stages(self)
 
             # # clear callbacks and instance
             # if float(".".join(get_version()[2])) >= 5:
@@ -700,6 +655,17 @@ def _configure_gym_env_spaces(self):
         # instantiate actions (needed for tasks for which the observations computation is dependent on the actions)
         self.actions = sample_space(self.single_action_space, self.sim.device, batch_size=self.num_envs, fill_value=0)
 
+    @captured
+    def _reset_terminated(self) -> None:
+        """Reset the environments in :attr:`reset_buf`, deciding on the device whether any must reset.
+
+        Recorded, the decision is a conditional graph node (CUDA 12.4+), so the step does not synchronize with the host;
+        run eagerly, :func:`warp.capture_if` reads the count on the host.
+        """
+        self._reset_count_wp.zero_()
+        wp.launch(count_masked, dim=self.num_envs, inputs=[self.reset_buf, self._reset_count_wp], device=self.device)
+        wp.capture_if(self._reset_count_wp, on_true=lambda: self._reset_idx(mask=self.reset_buf))
+
     def _reset_idx(self, mask: wp.array | None = None):
         """Reset environments based on a boolean mask.
 
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py
index cf2183c9310c..e97a796eadda 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py
@@ -182,7 +182,7 @@ def adapt_cfg(cls, cfg: Any) -> None:
         2. :meth:`_promote_scene_entity_cfgs` — replace stable
            :class:`~isaaclab.managers.SceneEntityCfg` instances under each
            term's ``params`` with the warp variant (which adds warp-cached
-           ``joint_mask``, ``joint_ids_wp``, ``body_ids_wp`` fields).
+           ``joint_mask_wp``, ``joint_ids_wp``, ``body_ids_wp`` fields).
         3. :meth:`_swap_mdp` — for every MDP term found anywhere in the cfg tree
            (discovered by :meth:`_walk_terms` via :class:`ManagerTermBaseCfg`
            subclassing, not by hard-coded attribute names), replace any stable
@@ -247,7 +247,7 @@ def _promote_scene_entity_cfgs(cls, cfg: Any) -> None:
         rebuilds any stable :class:`SceneEntityCfg` value under ``term.params``
         through :meth:`isaaclab_experimental.managers.SceneEntityCfg.from_stable`.
         The warp variant subclasses the stable one, so type checks elsewhere
-        stay valid; the new fields (``joint_mask`` / ``joint_ids_wp`` /
+        stay valid; the new fields (``joint_mask_wp`` / ``joint_ids_wp`` /
         ``body_ids_wp``) are filled at :meth:`resolve` time by the warp scene.
         """
         from isaaclab_experimental.managers.scene_entity_cfg import SceneEntityCfg as _WarpSceneEntityCfg
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
index 6056f678ac43..d40ccf1ba131 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
@@ -40,9 +40,8 @@
 
 from isaaclab_experimental.envs.interactive_scene_warp import InteractiveSceneWarp as InteractiveScene
 from isaaclab_experimental.managers import ActionManager, EventManager, ObservationManager
-from isaaclab_experimental.utils.torch_utils import clone_obs_buffer
 from isaaclab_experimental.utils.warp import resolve_1d_mask
-from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache
+from isaaclab_experimental.utils.warp_capture import CapturedStage
 
 # import logger
 logger = logging.getLogger(__name__)
@@ -79,7 +78,6 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
             )
         # initialize internal variables
         self._is_closed = False
-        self._physics_handles_decimation = False
 
         # set the seed for the environment
         if self.cfg.seed is not None:
@@ -103,9 +101,9 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
         if "cuda" in self.device:
             torch.cuda.set_device(self.device)
 
-        # Manager stages run through this cache; managers register terms that are unsafe to capture.
-        self._warp_graph_cache = WarpGraphCache(self.device)
-        self._warp_graph_cache.invalidate_on(self.sim.physics_manager)
+        # the managers' stages record CUDA graphs, and record again after the physics buffers rebind
+        CapturedStage.enabled = True
+        CapturedStage.invalidate_on(self.sim.physics_manager)
 
         # print useful information
         print("[INFO]: Base environment:")
@@ -140,8 +138,6 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
         # Process-wide per-env Warp RNG state (isaaclab.utils.seed.WarpRng), shared by terms, tasks and sensors.
         WarpRng.initialize(self.num_envs, self.device)
 
-        # TODO(jichuanh): this is problematic as warp capture requires stable pointers,
-        #                 using different masks for different managers/terms will cause problems.
         # Pre-allocated env masks (shared across managers/terms via `env`).
         self.ALL_ENV_MASK = wp.ones((self.num_envs,), dtype=wp.bool, device=self.device)
         self.ENV_MASK = wp.zeros((self.num_envs,), dtype=wp.bool, device=self.device)
@@ -361,8 +357,6 @@ def load_managers(self):
             :meth:`SimulationContext.reset_async` and it isn't possible to call async functions in the constructor.
 
         """
-        # recorded stages read buffers of the managers being replaced
-        self._warp_graph_cache.invalidate()
         # prepare the managers
         # -- event manager (we print it here to make the logging consistent)
         print("[INFO] Event Manager: ", self.event_manager)
@@ -375,12 +369,6 @@ def load_managers(self):
         # -- observation manager
         self.observation_manager = ObservationManager(self.cfg.observations, self)
         print("[INFO] Observation Manager:", self.observation_manager)
-        # persistent action input: recorded stages read this pointer, so each step copies into it
-        self._action_in_wp = wp.zeros(
-            (self.num_envs, self.action_manager.total_action_dim), dtype=wp.float32, device=self.device
-        )
-        # whether a reset has to compact its mask for a manager that resets by environment index
-        self._resets_by_index = bool(self.recorder_manager.active_terms)
 
         # perform events at the start of the simulation
         # in-case a child implementation creates other managers, the randomization should happen
@@ -431,8 +419,11 @@ def reset(
         Returns:
             A tuple containing the observations and extras.
         """
+        if env_mask is not None and self.recorder_manager.active_terms:
+            # the stable recorder manager hooks take indices; a mask takes precedence over the caller's ids
+            env_mask = self.resolve_env_mask(env_mask=env_mask)
+            env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1)
         env_mask = self.resolve_env_mask(env_ids=env_ids, env_mask=env_mask)
-        env_ids = self._reset_env_ids(env_mask)
 
         # trigger recorder terms for pre-reset calls
         self.recorder_manager.record_pre_reset(env_ids)
@@ -445,7 +436,7 @@ def reset(
             self.cfg.seed = used_seed
 
         # reset state of scene
-        self._reset_mask(env_mask, env_ids)
+        self._reset_mask(env_mask)
 
         # update articulation kinematics
         self.scene.write_data_to_sim()
@@ -499,7 +490,7 @@ def reset_to(
         if seed is not None:
             configure_seed(seed)
 
-        self._reset_mask(self.resolve_env_mask(env_ids=env_ids), env_ids)
+        self._reset_mask(self.resolve_env_mask(env_ids=env_ids))
 
         # set the state
         self.scene.reset_to(state, env_ids, is_relative=is_relative)
@@ -536,14 +527,8 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]:
         Returns:
             A tuple containing the observations and extras.
         """
-        # stages start recording at the first step after construction or an invalidation
-        self._warp_graph_cache.arm()
         # process actions
-        action_device = action.to(device=self.device, dtype=torch.float32).contiguous()
-        wp.copy(self._action_in_wp, wp.from_torch(action_device, dtype=wp.float32))
-        self._warp_graph_cache.call_steps(
-            "ActionManager_process_action", self.action_manager.stage_steps("process_action"), action=self._action_in_wp
-        )
+        self.action_manager.process_action(action.to(self.device))
 
         self.recorder_manager.record_pre_step()
 
@@ -556,9 +541,7 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]:
         for _ in range(self.cfg.decimation // steps_per_call):
             self._sim_step_counter += steps_per_call
             # set actions into buffers
-            self._warp_graph_cache.call_steps(
-                "ActionManager_apply_action", self.action_manager.stage_steps("apply_action")
-            )
+            self.action_manager.apply_action()
             # set actions into simulator
             self.scene.write_data_to_sim()
             # simulate
@@ -572,18 +555,10 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]:
 
         # post-step: step interval event
         if "interval" in self.event_manager.available_modes:
-            self._warp_graph_cache.call_steps(
-                "EventManager_apply_interval", self.event_manager.stage_steps("apply_interval"), dt=self.step_dt
-            )
+            self.event_manager.apply(mode="interval", dt=self.step_dt)
 
         # -- compute observations
-        self.obs_buf = self._warp_graph_cache.call_steps(
-            "ObservationManager_compute_update_history",
-            self.observation_manager.stage_steps("compute"),
-            update_history=True,
-            return_cloned_output=False,
-            output=clone_obs_buffer,
-        )
+        self.obs_buf = self.observation_manager.compute(update_history=True)
         self.recorder_manager.record_post_step()
 
         # return observations and extras
@@ -610,7 +585,7 @@ def close(self):
             del self.event_manager
             del self.recorder_manager
             del self.scene
-            self._warp_graph_cache.close()
+            CapturedStage.disable()
 
             # self.sim.clear_all_callbacks()
             self.sim.clear_instance()
@@ -625,33 +600,20 @@ def close(self):
     Helper functions.
     """
 
-    def _reset_env_ids(self, env_mask: wp.array) -> torch.Tensor | None:
-        """Compact a reset mask for the managers that reset by environment index.
-
-        This is the host boundary of the mask-first reset: it synchronizes with the device.
+    def _reset_mask(self, env_mask: wp.array):
+        """Reset the environments selected by a boolean mask.
 
         Args:
-            env_mask: Boolean mask of the environments to reset.
-
-        Returns:
-            Indices of the selected environments, or None when every active manager resets by mask.
-        """
-        if not self._resets_by_index:
-            return None
-        return wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1)
-
-    def _reset_mask(self, env_mask: wp.array, env_ids: torch.Tensor | None = None) -> None:
-        """Reset the selected environments.
-
-        Args:
-            env_mask: Boolean mask of the environments to reset.
-            env_ids: Indices of the same environments for the managers that reset by index, or None
-                when every active manager resets by mask.
+            env_mask: Boolean mask of the environments to reset, of shape ``(num_envs,)``.
         """
         # recorded reset stages read the environment-owned mask
         if env_mask is not self.reset_mask_wp:
             wp.copy(self.reset_mask_wp, env_mask)
         env_mask = self.reset_mask_wp
+        # the stable recorder manager resets by index
+        env_ids = (
+            wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) if self.recorder_manager.active_terms else None
+        )
 
         # reset the internal buffers of the scene elements
         self.scene.reset(env_mask=env_mask)
@@ -660,11 +622,8 @@ def _reset_mask(self, env_mask: wp.array, env_ids: torch.Tensor | None = None) -
         if "reset" in self.event_manager.available_modes:
             env_step_count = self._sim_step_counter // self.cfg.decimation
             self._global_env_step_count_wp.fill_(env_step_count)
-            self._warp_graph_cache.call_steps(
-                "EventManager_apply_reset",
-                self.event_manager.stage_steps("apply_reset"),
-                env_mask_wp=env_mask,
-                global_env_step_count=self._global_env_step_count_wp,
+            self.event_manager.apply(
+                mode="reset", env_mask_wp=env_mask, global_env_step_count=self._global_env_step_count_wp
             )
 
         # iterate over all managers and reset them
@@ -672,19 +631,13 @@ def _reset_mask(self, env_mask: wp.array, env_ids: torch.Tensor | None = None) -
         # note: This is order-sensitive! Certain things need be reset before others.
         self.extras["log"] = dict()
         # -- observation manager
-        info = self._warp_graph_cache.call_steps(
-            "ObservationManager_reset", self.observation_manager.stage_steps("reset"), env_mask=env_mask
-        )
+        info = self.observation_manager.reset(env_mask=env_mask)
         self.extras["log"].update(info)
         # -- action manager
-        info = self._warp_graph_cache.call_steps(
-            "ActionManager_reset", self.action_manager.stage_steps("reset"), env_mask=env_mask
-        )
+        info = self.action_manager.reset(env_mask=env_mask)
         self.extras["log"].update(info)
         # -- event manager
-        info = self._warp_graph_cache.call_steps(
-            "EventManager_reset", self.event_manager.stage_steps("reset"), env_mask=env_mask
-        )
+        info = self.event_manager.reset(env_mask=env_mask)
         self.extras["log"].update(info)
         # -- recorder manager
         info = self.recorder_manager.reset(env_ids)
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
index 5ee14a667dbc..4eb60beffed4 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py
@@ -14,7 +14,6 @@
 from __future__ import annotations
 
 import math
-import os
 from typing import Any, ClassVar
 
 import gymnasium as gym
@@ -25,23 +24,11 @@
 from isaaclab.envs.common import VecEnvStepReturn
 from isaaclab.envs.manager_based_rl_env_cfg import ManagerBasedRLEnvCfg
 from isaaclab.managers import CommandManager, CurriculumManager
-from isaaclab.utils.timer import Timer
 
 from isaaclab_experimental.managers import RewardManager, TerminationManager
-from isaaclab_experimental.utils.torch_utils import clone_obs_buffer
-from isaaclab_experimental.utils.warp import zero_masked_int64
 
 from .manager_based_env_warp import ManagerBasedEnvWarp
 
-DEBUG_TIMERS = os.environ.get("DEBUG_TIMERS", "0") == "1"
-"""Enable outer step() timer. Set DEBUG_TIMERS=1 env var to enable."""
-
-DEBUG_TIMER_STEP = os.environ.get("DEBUG_TIMER_STEP", "0") == "1"
-"""Enable step sub-phase timers. Set DEBUG_TIMER_STEP=1 env var to enable."""
-
-DEBUG_TIMER_RESET = os.environ.get("DEBUG_TIMER_RESET", "0") == "1"
-"""Enable reset sub-phase timers. Set DEBUG_TIMER_RESET=1 env var to enable."""
-
 
 class ManagerBasedRLEnvWarp(ManagerBasedEnvWarp, gym.Env):
     """The superclass for the manager-based workflow reinforcement learning-based environments.
@@ -171,13 +158,6 @@ def load_managers(self):
         self.curriculum_manager = CurriculumManager(self.cfg.curriculum, self)
         print("[INFO] Curriculum Manager: ", self.curriculum_manager)
 
-        # stable command and curriculum managers reset by environment index
-        self._resets_by_index = (
-            self._resets_by_index
-            or bool(self.command_manager.active_terms)
-            or bool(self.curriculum_manager.active_terms)
-        )
-
         # setup the action and observation spaces for Gym
         self._configure_gym_env_spaces()
 
@@ -205,14 +185,6 @@ def setup_manager_visualizers(self):
     Operations - MDP
     """
 
-    def invalidate_wp_graphs(self) -> None:
-        """Invalidate all cached Warp graphs.
-
-        Call this if the captured launch topology changes (e.g. different term list, shapes, etc.).
-        """
-        self._warp_graph_cache.invalidate()
-
-    @Timer(name="env_step", msg="Step took:", enable=DEBUG_TIMER_STEP, time_unit="us")
     def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         """Execute one time-step of the environment's dynamics and reset terminated environments.
 
@@ -232,22 +204,8 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         Returns:
             A tuple containing the observations, rewards, resets (terminated and truncated) and extras.
         """
-        # stages start recording at the first step after construction or an invalidation
-        self._warp_graph_cache.arm()
         # process actions
-        # NOTE: keep a persistent action input buffer for graph pointer stability.
-        # IMPORTANT: Do NOT re-wrap/replace the `wp.array` used by captured graphs each step.
-        # Instead, copy the latest actions into the persistent buffer.
-        with Timer(name="action_preprocess", msg="Action preprocessing took:", enable=DEBUG_TIMER_STEP, time_unit="us"):
-            action_device = action.to(device=self.device, dtype=torch.float32).contiguous()
-            wp.copy(self._action_in_wp, wp.from_torch(action_device, dtype=wp.float32))
-
-        self._warp_graph_cache.call_steps(
-            "ActionManager_process_action",
-            self.action_manager.stage_steps("process_action"),
-            action=self._action_in_wp,
-            timer=DEBUG_TIMER_STEP,
-        )
+        self.action_manager.process_action(action.to(self.device))
 
         self.recorder_manager.record_pre_step()
 
@@ -260,85 +218,74 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         for _ in range(self.cfg.decimation // steps_per_call):
             self._sim_step_counter += steps_per_call
             # set actions into buffers
-            self._warp_graph_cache.call_steps(
-                "ActionManager_apply_action", self.action_manager.stage_steps("apply_action"), timer=DEBUG_TIMER_STEP
-            )
+            self.action_manager.apply_action()
             # scene writes cross the actuator and asset host boundaries, so they stay eager
-            with Timer(
-                name="Scene_write_data_to_sim", msg="Scene write took:", enable=DEBUG_TIMER_STEP, time_unit="us"
-            ):
-                self.scene.write_data_to_sim()
+            self.scene.write_data_to_sim()
 
             # simulate
-            with Timer(name="simulate", msg="Newton simulation step took:", enable=DEBUG_TIMER_STEP, time_unit="us"):
-                self.sim.step(render=False)
+            self.sim.step(render=False)
             self.recorder_manager.record_post_physics_decimation_step()
             # render between steps only if the GUI or an RTX sensor needs it
             # note: we assume the render interval to be the shortest accepted rendering interval.
             #    If a camera needs rendering at a faster frequency, this will lead to unexpected behavior.
             if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering:
                 self.sim.render()
-            with Timer(
-                name="scene.update",
-                msg="Scene.update took:",
-                enable=DEBUG_TIMER_STEP,
-                time_unit="us",
-            ):
-                self.scene.update(dt=self.physics_dt * steps_per_call)
+            self.scene.update(dt=self.physics_dt * steps_per_call)
 
         # post-step:
         # -- update env counters (used for curriculum generation)
         self.episode_length_buf += 1  # step in current episode (per env)
         self.common_step_counter += 1  # total step (common for all envs)
 
-        # -- post-processing (termination + reward) as independently configurable stages
-        self.reset_buf = self._warp_graph_cache.call_steps(
-            "TerminationManager_compute", self.termination_manager.stage_steps("compute"), timer=DEBUG_TIMER_STEP
-        )
+        # -- post-processing: check terminations
+        self.reset_buf = self.termination_manager.compute()
         self.reset_terminated = self.termination_manager.terminated
         self.reset_time_outs = self.termination_manager.time_outs
-        self.reward_buf = self._warp_graph_cache.call_steps(
-            "RewardManager_compute",
-            self.reward_manager.stage_steps("compute"),
-            dt=float(self.step_dt),
-            timer=DEBUG_TIMER_STEP,
-        )
+        # -- reward computation
+        self.reward_buf = self.reward_manager.compute(dt=self.step_dt)
 
         if len(self.recorder_manager.active_terms) > 0:
             # update observations for recording if needed
-            self._warp_graph_cache.call_steps(
-                "ObservationManager_compute_no_history",
-                self.observation_manager.stage_steps("compute"),
-                return_cloned_output=False,
-                timer=DEBUG_TIMER_STEP,
-            )
+            self.observation_manager.compute(return_cloned_output=False)
             self.recorder_manager.record_post_step()
 
         # -- reset envs that terminated/timed-out and log the episode information
-        self._reset_terminated_envs()
+        # the reset pipeline advances curriculum and logging state even for an empty mask, so one
+        # host synchronization per step decides whether it runs
+        if self.reset_buf.any().item():
+            reset_mask = self.termination_manager.dones_wp
+            # the stable recorder manager hooks take indices
+            reset_env_ids = (
+                wp.to_torch(reset_mask).nonzero(as_tuple=False).squeeze(-1)
+                if self.recorder_manager.active_terms
+                else None
+            )
+            # capture the terminal observation before reset and expose it for Same-Step autoreset.
+            # the reset overwrites the persistent observation buffers, so keep a copy
+            if self.cfg.compute_final_obs:
+                self.extras["final_obs"] = self.observation_manager.compute()
+            # trigger recorder terms for pre-reset calls
+            self.recorder_manager.record_pre_reset(reset_env_ids)
+            self._reset_mask(reset_mask)
+
+            # if sensors are added to the scene, make sure we render to reflect changes in reset
+            if self.has_rtx_sensors and self.cfg.num_rerenders_on_reset > 0:
+                for _ in range(self.cfg.num_rerenders_on_reset):
+                    self.sim.render()
+
+            # trigger recorder terms for post-reset calls
+            self.recorder_manager.record_post_reset(reset_env_ids)
 
         # -- update command
-        self.command_manager.compute(dt=float(self.step_dt))
+        self.command_manager.compute(dt=self.step_dt)
 
         # -- step interval events
         if "interval" in self.event_manager.available_modes:
-            self._warp_graph_cache.call_steps(
-                "EventManager_apply_interval",
-                self.event_manager.stage_steps("apply_interval"),
-                dt=float(self.step_dt),
-                timer=DEBUG_TIMER_STEP,
-            )
+            self.event_manager.apply(mode="interval", dt=self.step_dt)
 
         # -- compute observations
         # note: done after reset to get the correct observations for reset envs
-        self.obs_buf = self._warp_graph_cache.call_steps(
-            "ObservationManager_compute_update_history",
-            self.observation_manager.stage_steps("compute"),
-            update_history=True,
-            return_cloned_output=False,
-            output=clone_obs_buffer,
-            timer=DEBUG_TIMER_STEP,
-        )
+        self.obs_buf = self.observation_manager.compute(update_history=True)
         # return observations, rewards, resets and extras
         return self.obs_buf, self.reward_buf, self.reset_terminated, self.reset_time_outs, self.extras
 
@@ -431,126 +378,58 @@ def _configure_gym_env_spaces(self):
         self.observation_space = gym.vector.utils.batch_space(self.single_observation_space, self.num_envs)
         self.action_space = gym.vector.utils.batch_space(self.single_action_space, self.num_envs)
 
-    def _reset_terminated_envs(self) -> None:
-        """Reset the environments that terminated or timed out in this step."""
-        with Timer(name="reset_selection", msg="Reset selection took:", enable=DEBUG_TIMER_STEP, time_unit="us"):
-            # The reset pipeline advances curriculum and logging state even for an empty mask, so one
-            # host synchronization per step decides whether it runs. When a manager resets by index,
-            # the same synchronization also compacts the mask.
-            if self._resets_by_index:
-                reset_env_ids = self._reset_env_ids(self.termination_manager.dones_wp)
-                has_resets = len(reset_env_ids) > 0
-            else:
-                reset_env_ids = None
-                has_resets = bool(self.reset_buf.any().item())
-        if not has_resets:
-            return
-
-        # capture the terminal observation before reset and expose it for Same-Step autoreset.
-        # the reset overwrites the persistent observation buffers, so keep a copy
-        if self.cfg.compute_final_obs:
-            final_obs = self._warp_graph_cache.call_steps(
-                "ObservationManager_compute_no_history",
-                self.observation_manager.stage_steps("compute"),
-                return_cloned_output=False,
-                timer=DEBUG_TIMER_STEP,
-            )
-            self.extras["final_obs"] = clone_obs_buffer(final_obs)
-        # trigger recorder terms for pre-reset calls
-        self.recorder_manager.record_pre_reset(reset_env_ids)
-
-        with Timer(name="reset_mask", msg="Reset mask took:", enable=DEBUG_TIMER_STEP, time_unit="us"):
-            self._reset_mask(self.termination_manager.dones_wp, reset_env_ids)
-
-        # if sensors are added to the scene, make sure we render to reflect changes in reset
-        if self.has_rtx_sensors and self.cfg.num_rerenders_on_reset > 0:
-            for _ in range(self.cfg.num_rerenders_on_reset):
-                self.sim.render()
-
-        # trigger recorder terms for post-reset calls
-        self.recorder_manager.record_post_reset(reset_env_ids)
-
-    def _reset_mask(self, env_mask: wp.array, env_ids: torch.Tensor | None = None) -> None:
-        """Reset the selected environments.
+    def _reset_mask(self, env_mask: wp.array):
+        """Reset the environments selected by a boolean mask.
 
         Args:
-            env_mask: Boolean mask of the environments to reset.
-            env_ids: Indices of the same environments for the managers that reset by index, or None
-                when every active manager resets by mask.
+            env_mask: Boolean mask of the environments to reset, of shape ``(num_envs,)``.
         """
         # recorded reset stages read the environment-owned mask
         if env_mask is not self.reset_mask_wp:
             wp.copy(self.reset_mask_wp, env_mask)
         env_mask = self.reset_mask_wp
+        # the stable command, curriculum and recorder managers reset by index
+        env_ids = (
+            wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1)
+            if self.command_manager.active_terms
+            or self.curriculum_manager.active_terms
+            or self.recorder_manager.active_terms
+            else None
+        )
 
         # update the curriculum for environments that need a reset
-        with Timer(
-            name="curriculum_manager.compute_reset",
-            msg="CurriculumManager.compute (reset) took:",
-            enable=DEBUG_TIMER_RESET,
-            time_unit="us",
-        ):
-            self.curriculum_manager.compute(env_ids=env_ids)
+        self.curriculum_manager.compute(env_ids=env_ids)
 
         # reset the internal buffers of the scene elements
-        with Timer(name="Scene_reset", msg="Scene reset took:", enable=DEBUG_TIMER_RESET, time_unit="us"):
-            self.scene.reset(env_mask=env_mask)
+        self.scene.reset(env_mask=env_mask)
 
         if "reset" in self.event_manager.available_modes:
             self._global_env_step_count_wp.fill_(self._sim_step_counter // self.cfg.decimation)
-            self._warp_graph_cache.call_steps(
-                "EventManager_apply_reset",
-                self.event_manager.stage_steps("apply_reset"),
-                env_mask_wp=env_mask,
-                global_env_step_count=self._global_env_step_count_wp,
-                timer=DEBUG_TIMER_RESET,
+            self.event_manager.apply(
+                mode="reset", env_mask_wp=env_mask, global_env_step_count=self._global_env_step_count_wp
             )
 
         # iterate over all managers and reset them
         # this returns a dictionary of information which is stored in the extras
         # note: This is order-sensitive! Certain things need be reset before others.
         # -- observation manager + action + reward managers
-        obs_info = self._warp_graph_cache.call_steps(
-            "ObservationManager_reset",
-            self.observation_manager.stage_steps("reset"),
-            env_mask=env_mask,
-            timer=DEBUG_TIMER_RESET,
-        )
-        action_info = self._warp_graph_cache.call_steps(
-            "ActionManager_reset", self.action_manager.stage_steps("reset"), env_mask=env_mask, timer=DEBUG_TIMER_RESET
-        )
-        reward_info = self._warp_graph_cache.call_steps(
-            "RewardManager_reset", self.reward_manager.stage_steps("reset"), env_mask=env_mask, timer=DEBUG_TIMER_RESET
-        )
+        obs_info = self.observation_manager.reset(env_mask=env_mask)
+        action_info = self.action_manager.reset(env_mask=env_mask)
+        reward_info = self.reward_manager.reset(env_mask=env_mask)
 
         # -- curriculum manager
-        with Timer(
-            name="curriculum_manager.reset",
-            msg="CurriculumManager.reset took:",
-            enable=DEBUG_TIMER_RESET,
-            time_unit="us",
-        ):
-            curriculum_info = self.curriculum_manager.reset(env_ids=env_ids)
+        curriculum_info = self.curriculum_manager.reset(env_ids=env_ids)
 
         # -- command + event + termination managers
         command_info = self.command_manager.reset(env_ids=env_ids)
-        event_info = self._warp_graph_cache.call_steps(
-            "EventManager_reset", self.event_manager.stage_steps("reset"), env_mask=env_mask, timer=DEBUG_TIMER_RESET
-        )
-        termination_info = self._warp_graph_cache.call_steps(
-            "TerminationManager_reset",
-            self.termination_manager.stage_steps("reset"),
-            env_mask=env_mask,
-            timer=DEBUG_TIMER_RESET,
-        )
+        event_info = self.event_manager.reset(env_mask=env_mask)
+        termination_info = self.termination_manager.reset(env_mask=env_mask)
 
         # -- recorder manager
         recorder_info = self.recorder_manager.reset(env_ids=env_ids)
 
         # reset the episode length buffer
-        wp.launch(
-            zero_masked_int64, dim=self.num_envs, inputs=[env_mask, self._episode_length_buf_wp], device=self.device
-        )
+        self.episode_length_buf.masked_fill_(wp.to_torch(env_mask), 0)
 
         # aggregate logging info
         log: dict[str, Any] = {}
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py
index 870bac6a34b5..8e7cf9c9ffe0 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py
@@ -45,15 +45,6 @@
     from isaaclab.envs import ManagerBasedEnv
 
 
-def _resolve_body_ids(asset_cfg: SceneEntityCfg, num_bodies: int) -> list[int]:
-    """Resolve a configured body selection to a list of body indices."""
-    if asset_cfg.body_ids is None or asset_cfg.body_ids == slice(None):
-        return list(range(num_bodies))
-    if isinstance(asset_cfg.body_ids, int):
-        return [asset_cfg.body_ids]
-    return list(asset_cfg.body_ids)
-
-
 def _vec3_range(ranges: dict[str, tuple[float, float]], keys: tuple[str, str, str]) -> tuple[wp.vec3f, wp.vec3f]:
     """Lower and upper bounds of three named ranges; a missing key is the zero range."""
     bounds = [ranges.get(key, (0.0, 0.0)) for key in keys]
@@ -68,6 +59,7 @@ def _vec3_range(ranges: dict[str, tuple[float, float]], keys: tuple[str, str, st
 @wp.kernel
 def _randomize_com_kernel(
     env_mask: wp.array(dtype=wp.bool),
+    needs_default: wp.array(dtype=wp.bool),
     rng_state: wp.array(dtype=wp.uint32),
     default_body_com_pos_b: wp.array(dtype=wp.vec3f, ndim=2),
     body_com_pos_b: wp.array(dtype=wp.vec3f, ndim=2),
@@ -77,6 +69,9 @@ def _randomize_com_kernel(
 ):
     """Offset the default center of mass of the selected bodies by one random offset per environment."""
     env_id = wp.tid()
+    if needs_default[0]:
+        for k in range(body_ids.shape[0]):
+            default_body_com_pos_b[env_id, body_ids[k]] = body_com_pos_b[env_id, body_ids[k]]
     if not env_mask[env_id]:
         return
 
@@ -110,13 +105,10 @@ def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv):
         super().__init__(cfg, env)
         asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"]
         self._asset: RigidObject | Articulation = env.scene[asset_cfg.name]
-        self._body_ids = wp.array(
-            _resolve_body_ids(asset_cfg, self._asset.num_bodies), dtype=wp.int32, device=env.device
-        )
         # Other startup terms may move the CoM after this term is built, so the defaults are copied on the
         # first call. Randomizing from them keeps repeated calls from accumulating offsets.
         self._default_com = wp.zeros_like(self._asset.data.body_com_pos_b.warp)
-        self._has_default_com = False
+        self._needs_default_com = wp.ones(1, dtype=wp.bool, device=env.device)
 
     def __call__(
         self,
@@ -127,16 +119,25 @@ def __call__(
     ):
         """Offset the selected environments' CoM positions [m] by a sample from ``com_range``."""
         body_com_pos_b = self._asset.data.body_com_pos_b.warp
-        if not self._has_default_com:
-            wp.copy(self._default_com, body_com_pos_b)
-            self._has_default_com = True
         com_lo, com_hi = _vec3_range(com_range, ("x", "y", "z"))
         wp.launch(
             kernel=_randomize_com_kernel,
             dim=env.num_envs,
-            inputs=[env_mask, WarpRng.state, self._default_com, body_com_pos_b, self._body_ids, com_lo, com_hi],
+            inputs=[
+                env_mask,
+                self._needs_default_com,
+                WarpRng.state,
+                self._default_com,
+                body_com_pos_b,
+                asset_cfg.body_ids_wp,
+                com_lo,
+                com_hi,
+            ],
             device=env.device,
         )
+        # the kernel only reads the flag, so it is cleared after the launch
+        self._needs_default_com.zero_()
+        # Notify the solver that inertial properties changed (COM position affects inertia).
         self._asset.set_coms_mask(coms=body_com_pos_b, env_mask=env_mask)
 
 
@@ -194,12 +195,8 @@ def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv):
         super().__init__(cfg, env)
         asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"]
         self._asset: RigidObject | Articulation = env.scene[asset_cfg.name]
-        num_bodies = self._asset.num_bodies
-        body_ids = _resolve_body_ids(asset_cfg, num_bodies)
-        self._body_ids = wp.array(body_ids, dtype=wp.int32, device=env.device)
-        self._body_mask = wp.array([i in body_ids for i in range(num_bodies)], dtype=wp.bool, device=env.device)
-        self._forces = wp.zeros((env.num_envs, num_bodies), dtype=wp.vec3f, device=env.device)
-        self._torques = wp.zeros((env.num_envs, num_bodies), dtype=wp.vec3f, device=env.device)
+        self._forces = wp.zeros((env.num_envs, self._asset.num_bodies), dtype=wp.vec3f, device=env.device)
+        self._torques = wp.zeros((env.num_envs, self._asset.num_bodies), dtype=wp.vec3f, device=env.device)
 
     def __call__(
         self,
@@ -220,7 +217,7 @@ def __call__(
             inputs=[
                 env_mask,
                 WarpRng.state,
-                self._body_ids,
+                asset_cfg.body_ids_wp,
                 self._forces,
                 self._torques,
                 float(force_range[0]),
@@ -231,7 +228,7 @@ def __call__(
             device=env.device,
         )
         self._asset.permanent_wrench_composer.set_forces_and_torques_mask(
-            forces=self._forces, torques=self._torques, body_mask=self._body_mask, env_mask=env_mask
+            forces=self._forces, torques=self._torques, body_mask=asset_cfg.body_mask_wp, env_mask=env_mask
         )
 
 
@@ -632,9 +629,9 @@ def reset_joints_by_scale(
 # The stable material/mass randomization terms already dispatch to the active
 # physics backend (Newton included); only their calling convention differs —
 # the warp EventManager invokes terms with a Warp env-mask while the stable
-# class terms expect torch env indices. Both terms run in ``startup`` mode,
-# once at construction and outside any recorded stage, so a host-side
-# mask-to-ids conversion is acceptable.
+# class terms expect torch env indices. Both terms are marked non-capturable, so
+# in reset or interval mode the event stages run them eagerly and the host-side
+# mask-to-ids conversion stays outside any recorded graph.
 
 
 def _mask_to_env_ids(env_mask: wp.array) -> torch.Tensor:
@@ -642,25 +639,15 @@ def _mask_to_env_ids(env_mask: wp.array) -> torch.Tensor:
     return torch.nonzero(wp.to_torch(env_mask), as_tuple=False).squeeze(-1)
 
 
-def _require_startup_mode(cfg: EventTermCfg, term: type) -> None:
-    """Reject modes that would run a host-synchronizing adapter on every reset or interval."""
-    if cfg.mode != "startup":
-        raise ValueError(f"'{term.__name__}' supports only the 'startup' event mode, got '{cfg.mode}'.")
-
-
 @WarpCapturable(False, reason="converts the environment mask to indices on the host")
 class randomize_rigid_body_material(_StableRandomizeRigidBodyMaterial, ManagerTermBase):
     """Warp adapter for the stable, backend-dispatched material randomization term.
 
     Converts the warp event manager's env-mask calling convention to the stable
-    term's env-ids convention and delegates. Startup mode only. Inherits the
-    warp :class:`ManagerTermBase` so the warp managers accept it as a class term.
+    term's env-ids convention and delegates. Inherits the warp :class:`ManagerTermBase`
+    so the warp managers accept it as a class term.
     """
 
-    def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv):
-        _require_startup_mode(cfg, type(self))
-        super().__init__(cfg, env)
-
     def __call__(
         self,
         env: ManagerBasedEnv,
@@ -684,7 +671,7 @@ def __call__(
         )
 
     def reset(self, env_mask: wp.array | None = None) -> None:
-        """Nothing to reset; materials are randomized once at startup."""
+        """Nothing to reset; the term keeps no per-environment state."""
         return
 
 
@@ -693,14 +680,10 @@ class randomize_rigid_body_mass(_StableRandomizeRigidBodyMass, ManagerTermBase):
     """Warp adapter for the stable, backend-dispatched mass randomization term.
 
     Converts the warp event manager's env-mask calling convention to the stable
-    term's env-ids convention and delegates. Startup mode only. Inherits the
-    warp :class:`ManagerTermBase` so the warp managers accept it as a class term.
+    term's env-ids convention and delegates. Inherits the warp :class:`ManagerTermBase`
+    so the warp managers accept it as a class term.
     """
 
-    def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv):
-        _require_startup_mode(cfg, type(self))
-        super().__init__(cfg, env)
-
     def __call__(
         self,
         env: ManagerBasedEnv,
@@ -724,5 +707,5 @@ def __call__(
         )
 
     def reset(self, env_mask: wp.array | None = None) -> None:
-        """Nothing to reset; masses are randomized once at startup."""
+        """Nothing to reset; the term keeps no per-environment state."""
         return
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py
index 4cc597989021..db9488e09d92 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py
@@ -268,7 +268,7 @@ def joint_torques_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = Sc
     wp.launch(
         kernel=_sum_sq_masked_kernel,
         dim=env.num_envs,
-        inputs=[asset.actuators.applied_effort.warp, asset_cfg.joint_mask, out],
+        inputs=[asset.actuators.applied_effort.warp, asset_cfg.joint_mask_wp, out],
         device=env.device,
     )
 
@@ -292,7 +292,7 @@ def joint_vel_l1(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg) -> None
     wp.launch(
         kernel=_sum_abs_masked_kernel,
         dim=env.num_envs,
-        inputs=[asset.data.joint_vel.warp, asset_cfg.joint_mask, out],
+        inputs=[asset.data.joint_vel.warp, asset_cfg.joint_mask_wp, out],
         device=env.device,
     )
 
@@ -303,7 +303,7 @@ def joint_vel_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = SceneE
     wp.launch(
         kernel=_sum_sq_masked_kernel,
         dim=env.num_envs,
-        inputs=[asset.data.joint_vel.warp, asset_cfg.joint_mask, out],
+        inputs=[asset.data.joint_vel.warp, asset_cfg.joint_mask_wp, out],
         device=env.device,
     )
 
@@ -314,7 +314,7 @@ def joint_acc_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = SceneE
     wp.launch(
         kernel=_sum_sq_masked_kernel,
         dim=env.num_envs,
-        inputs=[asset.data.joint_acc.warp, asset_cfg.joint_mask, out],
+        inputs=[asset.data.joint_acc.warp, asset_cfg.joint_mask_wp, out],
         device=env.device,
     )
 
@@ -341,7 +341,7 @@ def joint_deviation_l1(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg =
     wp.launch(
         kernel=_sum_abs_diff_masked_kernel,
         dim=env.num_envs,
-        inputs=[asset.data.joint_pos.warp, asset.data.default_joint_pos.warp, asset_cfg.joint_mask, out],
+        inputs=[asset.data.joint_pos.warp, asset.data.default_joint_pos.warp, asset_cfg.joint_mask_wp, out],
         device=env.device,
     )
 
@@ -379,7 +379,7 @@ def joint_pos_limits(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = Sc
     wp.launch(
         kernel=_joint_pos_limits_kernel,
         dim=env.num_envs,
-        inputs=[asset.data.joint_pos.warp, asset.data.soft_joint_pos_limits.warp, asset_cfg.joint_mask, out],
+        inputs=[asset.data.joint_pos.warp, asset.data.soft_joint_pos_limits.warp, asset_cfg.joint_mask_wp, out],
         device=env.device,
     )
 
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py
index 6d538e86bb1a..5b25f6382f75 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py
@@ -199,12 +199,15 @@ def _joint_pos_out_of_manual_limit_kernel(
     upper: float,
     out: wp.array(dtype=wp.bool),
 ):
-    """2D kernel (num_envs, num_joints). ``out`` is pre-zeroed; only writes True."""
-    i, j = wp.tid()
-    if joint_mask[j]:
-        v = joint_pos[i, j]
-        if v < lower or v > upper:
-            out[i] = True
+    """1D kernel (num_envs,): whether any selected joint is outside the bounds."""
+    i = wp.tid()
+    violated = bool(False)
+    for j in range(joint_pos.shape[1]):
+        if joint_mask[j]:
+            v = joint_pos[i, j]
+            if v < lower or v > upper:
+                violated = True
+    out[i] = violated
 
 
 def joint_pos_out_of_manual_limit(
@@ -212,20 +215,20 @@ def joint_pos_out_of_manual_limit(
 ) -> None:
     """Terminate when joint positions are outside configured bounds. Writes into ``out``."""
     asset: Articulation = env.scene[asset_cfg.name]
-    if asset_cfg.joint_mask is None:
+    if asset_cfg.joint_mask_wp is None:
         raise ValueError(
-            f"joint_pos_out_of_manual_limit requires SceneEntityCfg with resolved joint_mask, "
+            f"joint_pos_out_of_manual_limit requires SceneEntityCfg with resolved joint_mask_wp, "
             f"but got None for asset '{asset_cfg.name}'."
         )
-    if asset.data.joint_pos.warp.shape[1] != asset_cfg.joint_mask.shape[0]:
+    if asset.data.joint_pos.warp.shape[1] != asset_cfg.joint_mask_wp.shape[0]:
         raise ValueError(
-            f"joint_mask length ({asset_cfg.joint_mask.shape[0]}) does not match "
+            f"joint_mask_wp length ({asset_cfg.joint_mask_wp.shape[0]}) does not match "
             f"joint_pos dim ({asset.data.joint_pos.warp.shape[1]}) for asset '{asset_cfg.name}'."
         )
     wp.launch(
         kernel=_joint_pos_out_of_manual_limit_kernel,
-        dim=(env.num_envs, asset.data.joint_pos.warp.shape[1]),
-        inputs=[asset.data.joint_pos.warp, asset_cfg.joint_mask, bounds[0], bounds[1], out],
+        dim=env.num_envs,
+        inputs=[asset.data.joint_pos.warp, asset_cfg.joint_mask_wp, bounds[0], bounds[1], out],
         device=env.device,
     )
 
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py
index 32ef81c03bc0..e201b2f68cbf 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py
@@ -10,8 +10,7 @@
 import inspect
 import re
 from abc import abstractmethod
-from collections.abc import Callable, Sequence
-from functools import partial
+from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
 import torch
@@ -24,8 +23,9 @@
 from isaaclab.utils import instantiate
 
 from isaaclab_experimental.utils.warp import is_warp_capturable
+from isaaclab_experimental.utils.warp_capture import captured, eager
 
-from .manager_base import ManagerBase, ManagerTermBase
+from .manager_base import ManagerBase, ManagerTermBase, split_terms
 
 if TYPE_CHECKING:
     from isaaclab.envs import ManagerBasedEnv
@@ -230,6 +230,15 @@ def __init__(self, cfg: object, env: ManagerBasedEnv):
         self._action_torch = wp.to_torch(self._action)
         self._prev_action_torch = wp.to_torch(self._prev_action)
 
+        terms = list(self._terms.values())
+        self._term_split = split_terms(terms, is_warp_capturable)
+        # an annotation on a term class covers its processing; its reset is judged by its own annotation
+        self._reset_split = split_terms(terms, lambda term: is_warp_capturable(term.reset))
+        # each term's slice of the action vector
+        self._term_offsets = {
+            term: sum(other.action_dim for other in terms[:index]) for index, term in enumerate(terms)
+        }
+
         # check if any term has debug visualization implemented
         self.cfg.debug_vis = False
         for term in self._terms.values():
@@ -369,61 +378,39 @@ def set_debug_vis(self, debug_vis: bool):
         for term in self._terms.values():
             term.set_debug_vis(debug_vis)
 
-    def reset(
-        self,
-        env_ids: Sequence[int] | torch.Tensor | None = None,
-        *,
-        env_mask: wp.array | None = None,
-    ) -> dict[str, Any]:
+    @captured
+    def _reset(self, env_mask: wp.array) -> dict[str, Any]:
         """Resets the action history.
 
         Args:
-            env_ids: The specific environment indices to reset.
-                If None, all environments are considered.
             env_mask: Boolean Warp mask of shape (num_envs,) indicating which envs to reset.
-                If provided, takes precedence over ``env_ids``.
 
         Returns:
             An empty dictionary.
         """
-        # Mask-first path: captured callers must provide env_mask.
-        if env_mask is None or not isinstance(env_mask, wp.array):
-            if wp.get_device().is_capturing:
-                raise RuntimeError(
-                    "ActionManager.reset requires env_mask(wp.array[bool]) during capture. "
-                    "Do not pass env_ids on captured paths."
-                )
-            env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask)
-        return self._run_steps("reset", env_mask=env_mask)
-
-    def _reset_action_history(self, env_mask: wp.array) -> dict[str, Any]:
-        """Zero the current and previous actions of the selected environments."""
-        if env_mask is None:
-            self._prev_action.fill_(0.0)
-            self._action.fill_(0.0)
-        else:
-            wp.launch(
-                kernel=_zero_masked_2d,
-                dim=(self.num_envs, self.total_action_dim),
-                inputs=[env_mask, self._prev_action],
-                device=self.device,
-            )
-            wp.launch(
-                kernel=_zero_masked_2d,
-                dim=(self.num_envs, self.total_action_dim),
-                inputs=[env_mask, self._action],
-                device=self.device,
-            )
+        # reset the action history
+        wp.launch(
+            kernel=_zero_masked_2d,
+            dim=(self.num_envs, self.total_action_dim),
+            inputs=[env_mask, self._prev_action],
+            device=self.device,
+        )
+        wp.launch(
+            kernel=_zero_masked_2d,
+            dim=(self.num_envs, self.total_action_dim),
+            inputs=[env_mask, self._action],
+            device=self.device,
+        )
+
+        # reset all action terms
+        for term in self._reset_split.eager:
+            eager(term.reset, env_mask=env_mask)
+        for term in self._reset_split.captured:
+            term.reset(env_mask=env_mask)
         # nothing to log here
         return {}
 
-    @staticmethod
-    def _reset_term(term: ActionTerm, env_mask: wp.array) -> dict[str, Any]:
-        term.reset(env_mask=env_mask)
-        # nothing to log here
-        return {}
-
-    def process_action(self, action: wp.array):
+    def process_action(self, action: torch.Tensor):
         """Processes the actions sent to the environment.
 
         Note:
@@ -432,46 +419,36 @@ def process_action(self, action: wp.array):
         Args:
             action: The actions to process. Shape is (num_envs, total_action_dim).
         """
-        self._run_steps("process_action", action)
+        # check if action dimension is valid
+        if self.total_action_dim != action.shape[1]:
+            raise ValueError(f"Invalid action shape, expected: {self.total_action_dim}, received: {action.shape[1]}.")
+
+        # store the input actions
+        wp.copy(self._prev_action, self._action)
+        wp.copy(self._action, wp.from_torch(action.float()))
 
+        # split the actions and apply to each term
+        self._process_action()
+
+    @captured
+    def _process_action(self) -> None:
+        """Processes each term's slice of the stored actions."""
+        for term in self._term_split.eager:
+            eager(term.process_actions, self._action, self._term_offsets[term])
+        for term in self._term_split.captured:
+            term.process_actions(self._action, self._term_offsets[term])
+
+    @captured
     def apply_action(self) -> None:
         """Applies the actions to the environment/simulation.
 
         Note:
             This should be called at every simulation step.
         """
-        self._run_steps("apply_action")
-
-    """
-    Operations - Stage steps.
-    """
-
-    def _build_stage_steps(self, operation: str) -> list[tuple[bool, Callable[..., Any]]]:
-        terms = list(self._terms.values())
-        if operation == "process_action":
-            offsets = [sum(term.action_dim for term in terms[:index]) for index in range(len(terms))]
-            steps = [
-                (is_warp_capturable(term), partial(self._process_term_actions, term, offset))
-                for term, offset in zip(terms, offsets)
-            ]
-            return [(True, self._store_actions), *steps]
-        if operation == "apply_action":
-            return [(is_warp_capturable(term), term.apply_actions) for term in terms]
-        if operation == "reset":
-            steps = [(is_warp_capturable(term), partial(self._reset_term, term)) for term in terms]
-            return [(True, self._reset_action_history), *steps]
-        return super()._build_stage_steps(operation)
-
-    def _store_actions(self, action: wp.array) -> None:
-        """Check the action shape and store the actions, keeping the previous ones."""
-        if self.total_action_dim != action.shape[1]:
-            raise ValueError(f"Invalid action shape, expected: {self.total_action_dim}, received: {action.shape[1]}.")
-        wp.copy(self._prev_action, self._action)
-        wp.copy(self._action, action)
-
-    def _process_term_actions(self, term: ActionTerm, offset: int, action: wp.array) -> None:
-        """Hand a term its slice of the stored actions."""
-        term.process_actions(self._action, offset)
+        for term in self._term_split.eager:
+            eager(term.apply_actions)
+        for term in self._term_split.captured:
+            term.apply_actions()
 
     def get_term(self, name: str) -> ActionTerm:
         """Returns the action term with the specified name.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
index 627199f18600..7b959c9b4521 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
@@ -23,9 +23,7 @@
 
 import inspect
 import logging
-from collections.abc import Callable, Sequence
-from functools import partial
-from typing import Any
+from collections.abc import Sequence
 
 import torch
 import warp as wp
@@ -35,8 +33,9 @@
 from isaaclab.utils.seed import WarpRng
 
 from isaaclab_experimental.utils.warp import is_warp_capturable
+from isaaclab_experimental.utils.warp_capture import captured, eager, reset_captured_stages
 
-from .manager_base import ManagerBase, ManagerTermBase
+from .manager_base import ManagerBase, ManagerTermBase, split_resets
 
 logger = logging.getLogger(__name__)
 
@@ -175,7 +174,6 @@ def __init__(self, cfg: object, env):
 
         # Warp buffers for interval/reset modes (populated in _prepare_terms)
         self._interval_term_time_left_wp: list[wp.array] = []
-        self._interval_term_ranges: list[tuple[float, float]] = []
         self._interval_term_is_global: list[bool] = []
         # Scalar RNG state for global interval timers (allocated lazily if needed).
         self._interval_global_rng_state_wp: wp.array | None = None
@@ -184,6 +182,7 @@ def __init__(self, cfg: object, env):
         self._reset_term_triggered_once_wp: list[wp.array] = []
 
         super().__init__(cfg, env)
+        self._split_terms()
 
         # persistent scratch mask for per-term interval/reset triggering (must be stable pointer for capture)
         self._scratch_term_mask_wp = wp.zeros((self.num_envs,), dtype=wp.bool, device=self.device)
@@ -228,41 +227,30 @@ def available_modes(self) -> list[str]:
         return list(self._mode_term_names.keys())
 
     def set_term_cfg(self, term_name: str, cfg: EventTermCfg):
-        """Set the configuration of an existing event term.
+        """Sets the configuration of the specified term into the manager.
 
-        Recorded event stages read the replaced configuration, so they are recorded again on their next call.
+        The method finds the term by name by searching through all the modes.
+        It then updates the configuration of the term with the first matching name.
 
         Args:
-            term_name: Name of the event term.
-            cfg: Replacement event term configuration.
+            term_name: The name of the event term.
+            cfg: The configuration for the event term.
 
         Raises:
-            ValueError: If the term does not exist, or the replacement changes its mode or interval timing.
+            ValueError: If the term name is not found.
         """
+        term_found = False
         for mode, terms in self._mode_term_names.items():
-            if term_name not in terms:
-                continue
-            if cfg.mode != mode:
-                raise ValueError(f"Event term '{term_name}' has mode '{mode}', but the replacement has '{cfg.mode}'.")
-            term_index = terms.index(term_name)
-            if mode == "interval":
-                if cfg.interval_range_s is None:
-                    raise ValueError(f"Event term '{term_name}' has mode 'interval' but 'interval_range_s' is not set.")
-                if cfg.is_global_time != self._mode_term_cfgs[mode][term_index].is_global_time:
-                    raise ValueError(f"Event term '{term_name}' cannot change 'is_global_time' at runtime.")
-            self._resolve_common_term_cfg(term_name, cfg, min_argc=2)
-            self._mode_term_cfgs[mode][term_index] = cfg
-            if mode == "interval":
-                lower, upper = cfg.interval_range_s
-                self._interval_term_ranges[term_index] = (float(lower), float(upper))
-            self._mode_class_term_cfgs[mode] = [
-                term_cfg
-                for term_cfg in self._mode_term_cfgs[mode]
-                if inspect.isclass(term_cfg.func) or isinstance(term_cfg.func, ManagerTermBase)
-            ]
-            self._clear_stage_steps()
-            return
-        raise ValueError(f"Event term '{term_name}' not found.")
+            if term_name in terms:
+                cfg.capturable = is_warp_capturable(cfg.func, cfg.params)
+                self._mode_term_cfgs[mode][terms.index(term_name)] = cfg
+                term_found = True
+                break
+        if not term_found:
+            raise ValueError(f"Event term '{term_name}' not found.")
+        # the class terms derive from the term configurations, and the recorded stages hold the replaced term
+        self._split_terms()
+        reset_captured_stages(self)
 
     def get_term_cfg(self, term_name: str) -> EventTermCfg:
         for mode, terms in self._mode_term_names.items():
@@ -270,23 +258,36 @@ def get_term_cfg(self, term_name: str) -> EventTermCfg:
                 return self._mode_term_cfgs[mode][terms.index(term_name)]
         raise ValueError(f"Event term '{term_name}' not found.")
 
-    def reset(
-        self,
-        env_ids: Sequence[int] | slice | torch.Tensor | wp.array | None = None,
-        *,
-        env_mask: wp.array | torch.Tensor | None = None,
-    ) -> dict[str, float]:
-        # Mask-first path: captured callers must provide env_mask.
-        if env_mask is None or not isinstance(env_mask, wp.array):
-            # Keep all id->mask resolution strictly outside capture.
-            if wp.get_device().is_capturing:
-                raise RuntimeError(
-                    "EventManager.reset requires env_mask(wp.array[bool]) during capture. "
-                    "Do not pass env_ids on captured paths."
+    @captured
+    def _reset(self, env_mask: wp.array) -> dict[str, float]:
+        """Resets the class terms and the interval timers of the environments selected by a boolean mask."""
+        # reset class terms (mask-based)
+        for term_cfg in self._class_term_split.eager:
+            eager(term_cfg.func.reset, env_mask=env_mask)
+        for term_cfg in self._class_term_split.captured:
+            term_cfg.func.reset(env_mask=env_mask)
+
+        # reset interval timers for non-global interval events
+        # note: terms with resample_interval_on_reset=False keep their per-environment counter
+        #   across resets, so we do not resample them either
+        if "interval" in self._mode_term_cfgs:
+            for i, term_cfg in enumerate(self._mode_term_cfgs["interval"]):
+                if term_cfg.is_global_time or not term_cfg.resample_interval_on_reset:
+                    continue
+                lower, upper = term_cfg.interval_range_s
+                wp.launch(
+                    kernel=_interval_reset_selected,
+                    dim=self.num_envs,
+                    inputs=[
+                        env_mask,
+                        self._interval_term_time_left_wp[i],
+                        WarpRng.state,
+                        lower,
+                        upper,
+                    ],
+                    device=self.device,
                 )
-            env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask)
-
-        return self._run_steps("reset", env_mask=env_mask)
+        return {}
 
     def apply(
         self,
@@ -318,7 +319,7 @@ def apply(
                 raise ValueError(
                     f"Event mode '{mode}' does not require environment indices. This is an undefined behavior."
                 )
-            self._apply_interval(float(dt))
+            self._apply_interval(dt)
             return
 
         # resolve the environment mask
@@ -337,125 +338,80 @@ def apply(
         for term_cfg in self._mode_term_cfgs[mode]:
             term_cfg.func(self._env, env_mask_wp, **term_cfg.params)
 
+    @captured
     def _apply_interval(self, dt: float) -> None:
-        self._run_steps("apply_interval", dt=dt)
+        # iterate over all the interval terms (fixed list; captured graph-friendly)
+        for i, term_cfg in enumerate(self._mode_term_cfgs["interval"]):
+            lower, upper = term_cfg.interval_range_s
+            if self._interval_term_is_global[i]:
+                if self._interval_global_rng_state_wp is None:
+                    raise RuntimeError(
+                        "EventManager._apply_interval: _interval_global_rng_state_wp is not initialized."
+                    )
+                # update scalar time_left and scalar flag (mask is a broadcast view of the flag)
+                wp.launch(
+                    kernel=_interval_step_global,
+                    dim=1,
+                    inputs=[
+                        self._interval_term_time_left_wp[i],
+                        self._interval_global_rng_state_wp,
+                        self._scratch_interval_trigger_flag_wp,
+                        dt,
+                        lower,
+                        upper,
+                    ],
+                    device=self.device,
+                )
+                term_mask = self._scratch_interval_trigger_mask_view_wp
+            else:
+                wp.launch(
+                    kernel=_interval_step_per_env,
+                    dim=self.num_envs,
+                    inputs=[
+                        self._interval_term_time_left_wp[i],
+                        WarpRng.state,
+                        self._scratch_term_mask_wp,
+                        dt,
+                        lower,
+                        upper,
+                    ],
+                    device=self.device,
+                )
+                term_mask = self._scratch_term_mask_wp
+            if term_cfg.capturable:
+                term_cfg.func(self._env, term_mask, **term_cfg.params)
+            else:
+                eager(term_cfg.func, self._env, term_mask, **term_cfg.params)
 
+    @captured
     def _apply_reset(self, env_mask_wp: wp.array, global_env_step_count_wp: wp.array) -> None:
-        if self._scratch_term_mask_wp is None:
-            raise RuntimeError("EventManager._apply_reset: _scratch_term_mask_wp is not initialized.")
-        self._run_steps("apply_reset", env_mask_wp=env_mask_wp, global_env_step_count=global_env_step_count_wp)
-
-    """
-    Operations - Stage steps.
-    """
-
-    def _build_stage_steps(self, operation: str) -> list[tuple[bool, Callable[..., Any]]]:
-        # Terms draw from the environment's shared random state, so every stage keeps the configured order.
-        if operation == "apply_interval":
-            return [
-                (is_warp_capturable(term_cfg.func, term_cfg.params), partial(self._apply_interval_term, index))
-                for index, term_cfg in enumerate(self._mode_term_cfgs.get("interval", []))
-            ]
-        if operation == "apply_reset":
-            return [
-                (is_warp_capturable(term_cfg.func, term_cfg.params), partial(self._apply_reset_term, index))
-                for index, term_cfg in enumerate(self._mode_term_cfgs.get("reset", []))
-            ]
-        if operation == "reset":
-            steps = [
-                (is_warp_capturable(term_cfg.func, term_cfg.params), partial(self._reset_class_term, term_cfg))
-                for mode_cfgs in self._mode_class_term_cfgs.values()
-                for term_cfg in mode_cfgs
-            ]
-            # note: terms with resample_interval_on_reset=False keep their per-environment counter
-            #   across resets, so we do not resample them either
-            steps += [
-                (True, partial(self._resample_interval, index))
-                for index, term_cfg in enumerate(self._mode_term_cfgs.get("interval", []))
-                if not term_cfg.is_global_time and term_cfg.resample_interval_on_reset
-            ]
-            return [*steps, (True, self._reset_output)]
-        return super()._build_stage_steps(operation)
-
-    def _apply_interval_term(self, index: int, dt: float) -> None:
-        """Advance one interval term's timers and apply it to the environments whose timer expired."""
-        term_cfg = self._mode_term_cfgs["interval"][index]
-        lower, upper = self._interval_term_ranges[index]
-        if self._interval_term_is_global[index]:
-            if self._interval_global_rng_state_wp is None:
-                raise RuntimeError("EventManager._apply_interval: _interval_global_rng_state_wp is not initialized.")
-            # update scalar time_left and scalar flag (mask is a broadcast view of the flag)
+        # iterate over all the reset terms
+        for index, term_cfg in enumerate(self._mode_term_cfgs["reset"]):
             wp.launch(
-                kernel=_interval_step_global,
-                dim=1,
-                inputs=[
-                    self._interval_term_time_left_wp[index],
-                    self._interval_global_rng_state_wp,
-                    self._scratch_interval_trigger_flag_wp,
-                    float(dt),
-                    float(lower),
-                    float(upper),
-                ],
-                device=self.device,
-            )
-            term_cfg.func(self._env, self._scratch_interval_trigger_mask_view_wp, **term_cfg.params)
-        else:
-            wp.launch(
-                kernel=_interval_step_per_env,
+                kernel=_reset_compute_valid_mask,
                 dim=self.num_envs,
                 inputs=[
-                    self._interval_term_time_left_wp[index],
-                    WarpRng.state,
+                    env_mask_wp,
+                    self._reset_term_last_triggered_step_wp[index],
+                    self._reset_term_triggered_once_wp[index],
                     self._scratch_term_mask_wp,
-                    float(dt),
-                    float(lower),
-                    float(upper),
+                    global_env_step_count_wp,
+                    term_cfg.min_step_count_between_reset,
                 ],
                 device=self.device,
             )
-            term_cfg.func(self._env, self._scratch_term_mask_wp, **term_cfg.params)
-
-    def _apply_reset_term(self, index: int, env_mask_wp: wp.array, global_env_step_count: wp.array) -> None:
-        """Apply one reset term to the selected environments its trigger spacing allows."""
-        term_cfg = self._mode_term_cfgs["reset"][index]
-        wp.launch(
-            kernel=_reset_compute_valid_mask,
-            dim=self.num_envs,
-            inputs=[
-                env_mask_wp,
-                self._reset_term_last_triggered_step_wp[index],
-                self._reset_term_triggered_once_wp[index],
-                self._scratch_term_mask_wp,
-                global_env_step_count,
-                int(term_cfg.min_step_count_between_reset),
-            ],
-            device=self.device,
-        )
-        term_cfg.func(self._env, self._scratch_term_mask_wp, **term_cfg.params)
-
-    @staticmethod
-    def _reset_class_term(term_cfg: EventTermCfg, env_mask: wp.array) -> None:
-        term_cfg.func.reset(env_mask=env_mask)
-
-    def _resample_interval(self, index: int, env_mask: wp.array) -> None:
-        """Draw new interval timers for the selected environments."""
-        lower, upper = self._interval_term_ranges[index]
-        wp.launch(
-            kernel=_interval_reset_selected,
-            dim=self.num_envs,
-            inputs=[
-                env_mask,
-                self._interval_term_time_left_wp[index],
-                WarpRng.state,
-                float(lower),
-                float(upper),
-            ],
-            device=self.device,
-        )
+            if term_cfg.capturable:
+                term_cfg.func(self._env, self._scratch_term_mask_wp, **term_cfg.params)
+            else:
+                eager(term_cfg.func, self._env, self._scratch_term_mask_wp, **term_cfg.params)
 
-    @staticmethod
-    def _reset_output(env_mask: wp.array) -> dict[str, float]:
-        return {}
+    def _split_terms(self) -> None:
+        """Collect the class terms of every mode and split them by whether their resets can be recorded."""
+        self._mode_class_term_cfgs = {
+            mode: [cfg for cfg in cfgs if inspect.isclass(cfg.func) or isinstance(cfg.func, ManagerTermBase)]
+            for mode, cfgs in self._mode_term_cfgs.items()
+        }
+        self._class_term_split = split_resets([cfg for cfgs in self._mode_class_term_cfgs.values() for cfg in cfgs])
 
     def _prepare_terms(self):
         # check if config is dict already
@@ -502,14 +458,10 @@ def _prepare_terms(self):
             if term_cfg.mode not in self._mode_term_names:
                 self._mode_term_names[term_cfg.mode] = []
                 self._mode_term_cfgs[term_cfg.mode] = []
-                self._mode_class_term_cfgs[term_cfg.mode] = []
             # add term name and parameters
             self._mode_term_names[term_cfg.mode].append(term_name)
             self._mode_term_cfgs[term_cfg.mode].append(term_cfg)
 
-            if inspect.isclass(term_cfg.func):
-                self._mode_class_term_cfgs[term_cfg.mode].append(term_cfg)
-
             # per-mode Warp buffers
             if term_cfg.mode == "interval":
                 if term_cfg.interval_range_s is None:
@@ -517,7 +469,6 @@ def _prepare_terms(self):
                         f"Event term '{term_name}' has mode 'interval' but 'interval_range_s' is not specified."
                     )
                 lower, upper = term_cfg.interval_range_s
-                self._interval_term_ranges.append((float(lower), float(upper)))
 
                 if term_cfg.is_global_time:
                     # allocate and seed scalar global RNG state if needed (avoid consuming env0 RNG stream)
@@ -533,7 +484,7 @@ def _prepare_terms(self):
                     wp.launch(
                         kernel=_interval_init_global,
                         dim=1,
-                        inputs=[time_left, self._interval_global_rng_state_wp, float(lower), float(upper)],
+                        inputs=[time_left, self._interval_global_rng_state_wp, lower, upper],
                         device=self.device,
                     )
                     self._interval_term_time_left_wp.append(time_left)
@@ -543,7 +494,7 @@ def _prepare_terms(self):
                     wp.launch(
                         kernel=_interval_init_per_env,
                         dim=self.num_envs,
-                        inputs=[time_left, WarpRng.state, float(lower), float(upper)],
+                        inputs=[time_left, WarpRng.state, lower, upper],
                         device=self.device,
                     )
                     self._interval_term_time_left_wp.append(time_left)
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py
index 87795f7814a8..36f5cd4a033c 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py
@@ -21,7 +21,7 @@
 import logging
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Sequence
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, NamedTuple
 
 import warp as wp
 
@@ -31,6 +31,7 @@
 from isaaclab.utils import string_to_callable, to_dict
 
 from isaaclab_experimental.utils.warp import is_warp_capturable
+from isaaclab_experimental.utils.warp_capture import reset_captured_stages
 
 from .scene_entity_cfg import SceneEntityCfg
 
@@ -46,6 +47,38 @@
 logger = logging.getLogger(__name__)
 
 
+class TermSplit(NamedTuple):
+    """Terms of one operation, split by whether they can be recorded into a CUDA graph."""
+
+    eager: list
+    """Terms that run eagerly, before the recorded part."""
+
+    captured: list
+    """Terms recorded with the rest of the operation."""
+
+
+def split_terms(terms: Sequence[Any], capturable: Callable[[Any], bool] = lambda term: term.capturable) -> TermSplit:
+    """Split terms by capturability, keeping their order.
+
+    Args:
+        terms: The terms (term configurations by default).
+        capturable: Whether a term can be recorded. Defaults to the term configuration's ``capturable`` flag.
+    """
+    return TermSplit([term for term in terms if not capturable(term)], [term for term in terms if capturable(term)])
+
+
+def split_resets(term_cfgs: Sequence[Any]) -> TermSplit:
+    """Split class terms by whether their ``reset`` can be recorded, keeping their order.
+
+    An annotation on a term class covers its ``__call__``; the ``reset`` is judged by the annotation on the
+    ``reset`` method itself.
+
+    Args:
+        term_cfgs: The configurations of the class terms.
+    """
+    return split_terms(term_cfgs, lambda term_cfg: is_warp_capturable(term_cfg.func.reset, term_cfg.params))
+
+
 class ManagerTermBase(ABC):
     """Base class for manager terms.
 
@@ -198,9 +231,6 @@ def __init__(self, cfg: object, env: ManagerBasedEnv):
         # else:
         #     self._resolve_terms_handle = None
         self._resolve_terms_handle = None
-        # steps of each operation for the environment's graph cache, built on first use
-        self._stage_steps: dict[str, tuple[tuple[bool, Callable[..., Any]], ...]] = {}
-        self._all_terms_capturable = True
 
         # parse config to create terms information
         if self.cfg:
@@ -240,17 +270,34 @@ def active_terms(self) -> list[str] | dict[str, list[str]]:
     Operations.
     """
 
-    def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None) -> dict[str, float]:
+    def reset(
+        self,
+        env_ids: Sequence[int] | torch.Tensor | None = None,
+        *,
+        env_mask: wp.array | None = None,
+    ) -> dict[str, float]:
         """Resets the manager and returns logging information for the current time-step.
 
+        The ids resolve to a mask on the host, which a CUDA graph cannot record, so the recorded part is :meth:`_reset`.
+
         Args:
             env_ids: The environment ids for which to log data.
                 Defaults None, which logs data for all environments.
+            env_mask: Boolean Warp mask of shape (num_envs,) selecting the environments.
+                If provided, takes precedence over ``env_ids``.
 
         Returns:
             Dictionary containing the logging information.
         """
-        return {}
+        # Mask-first path: captured callers must provide env_mask.
+        if env_mask is None or not isinstance(env_mask, wp.array):
+            if wp.get_device().is_capturing:
+                raise RuntimeError(
+                    f"{type(self).__name__}.reset requires env_mask(wp.array[bool]) during capture. "
+                    "Do not pass env_ids on captured paths."
+                )
+            env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask)
+        return self._reset(env_mask)
 
     def find_terms(self, name_keys: str | Sequence[str]) -> list[str]:
         """Find terms in the manager based on the names.
@@ -289,62 +336,19 @@ def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequenc
         """
         raise NotImplementedError
 
-    def stage_steps(self, operation: str) -> tuple[tuple[bool, Callable[..., Any]], ...]:
-        """Ordered steps of an operation, each marked as capturable or not.
-
-        Calling the steps in order with the operation's arguments performs the operation. The Warp
-        environments record each run of consecutive capturable steps into a CUDA graph and run the other
-        steps eagerly in between (see :meth:`~isaaclab_experimental.utils.WarpGraphCache.call_steps`). A
-        term decorated with ``@WarpCapturable(False)`` makes its own step eager.
-
-        Args:
-            operation: Name of the operation, such as ``"compute"`` or ``"reset"``.
-
-        Returns:
-            The steps, reused until a term changes.
-        """
-        steps = self._stage_steps.get(operation)
-        if steps is None:
-            steps = self._stage_steps[operation] = tuple(self._build_stage_steps(operation))
-        return steps
-
     """
     Implementation specific.
     """
 
-    def _build_stage_steps(self, operation: str) -> list[tuple[bool, Callable[..., Any]]]:
-        """Build the steps of an operation.
-
-        The default runs the whole method as one step, capturable when every term is. Managers that
-        split an operation per term override this.
-
-        Args:
-            operation: Name of the operation.
-
-        Returns:
-            The ordered steps.
-        """
-        return [(self._all_terms_capturable, getattr(self, operation))]
-
-    def _run_steps(self, operation: str, *args: Any, **kwargs: Any) -> Any:
-        """Run the steps of an operation eagerly and return the last step's result."""
-        result = None
-        for _, step in self.stage_steps(operation):
-            result = step(*args, **kwargs)
-        return result
-
-    def _clear_stage_steps(self) -> None:
-        """Rebuild the steps on next use and drop the graphs recorded from them."""
-        self._stage_steps.clear()
-        graph_cache = getattr(self._env, "_warp_graph_cache", None)
-        if graph_cache is not None:
-            graph_cache.invalidate(type(self).__name__)
-
     @abstractmethod
     def _prepare_terms(self):
         """Prepare terms information from the configuration object."""
         raise NotImplementedError
 
+    def _reset(self, env_mask: wp.array) -> dict[str, float]:
+        """Resets the environments selected by a boolean mask of shape (num_envs,) and returns logging information."""
+        return {}
+
     """
     Internal callbacks.
     """
@@ -375,7 +379,7 @@ def _resolve_terms_callback(self, event):
         # set the flag
         self._is_scene_entities_resolved = True
         # class terms were just created
-        self._clear_stage_steps()
+        reset_captured_stages(self)
 
     """
     Internal functions.
@@ -459,8 +463,8 @@ def _resolve_common_term_cfg(self, term_name: str, term_cfg: ManagerTermBaseCfg,
                     f" and optional parameters: {args_with_defaults}, but received: {term_params}."
                 )
 
-        if not is_warp_capturable(term_cfg.func, term_cfg.params):
-            self._all_terms_capturable = False
+        # whether the term can be recorded into a CUDA graph, resolved once per term
+        term_cfg.capturable = is_warp_capturable(term_cfg.func, term_cfg.params)
 
         # process attributes at runtime
         # these properties are only resolvable once the simulation starts playing
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
index b756d210a173..a8f4f3bdea6a 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
@@ -47,8 +47,7 @@
 from __future__ import annotations
 
 import inspect
-from collections.abc import Callable, Sequence
-from functools import partial
+from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
@@ -63,9 +62,9 @@
 from isaaclab_experimental.utils import modifiers, noise
 from isaaclab_experimental.utils.buffers import CircularBuffer
 from isaaclab_experimental.utils.torch_utils import clone_obs_buffer
-from isaaclab_experimental.utils.warp import is_warp_capturable
+from isaaclab_experimental.utils.warp_capture import captured, eager
 
-from .manager_base import ManagerBase, ManagerTermBase
+from .manager_base import ManagerBase, ManagerTermBase, split_resets, split_terms
 
 if TYPE_CHECKING:
     from isaaclab.envs import ManagerBasedEnv
@@ -196,6 +195,14 @@ def __init__(self, cfg: object, env: ManagerBasedEnv):
 
         # Stores the latest observations.
         self._obs_buffer: dict[str, torch.Tensor | dict[str, torch.Tensor]] | None = None
+        # (group, index) of every term, and the class terms, split by capturability
+        terms = [
+            (group_name, index) for group_name, cfgs in self._group_obs_term_cfgs.items() for index in range(len(cfgs))
+        ]
+        self._term_split = split_terms(terms, lambda term: self._group_obs_term_cfgs[term[0]][term[1]].capturable)
+        self._class_term_split = split_resets(
+            [cfg for cfgs in self._group_obs_class_term_cfgs.values() for cfg in cfgs]
+        )
         # Note: Persistent Warp output buffers (`_group_out_wp` / `_group_out_torch`) and per-term post-processing
         # buffers are allocated during `_prepare_terms()` since they are per-term/per-group setup.
 
@@ -387,22 +394,33 @@ def _collect_io_descriptors(self, group_names_to_export: list[str] = ["policy"])
     Operations.
     """
 
-    def reset(
-        self,
-        env_ids: Sequence[int] | torch.Tensor | None = None,
-        *,
-        env_mask: wp.array | None = None,
-    ) -> dict[str, float]:
-        # Mask-first path: captured callers must provide env_mask.
-        if env_mask is None or not isinstance(env_mask, wp.array):
-            # Keep all id->mask resolution strictly outside capture.
-            if wp.get_device().is_capturing:
-                raise RuntimeError(
-                    "ObservationManager.reset requires env_mask(wp.array[bool]) during capture. "
-                    "Do not pass env_ids on captured paths."
-                )
-            env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask)
-        return self._run_steps("reset", env_mask=env_mask)
+    @captured
+    def _reset(self, env_mask: wp.array) -> dict[str, float]:
+        """Resets the observation terms of the environments selected by a boolean mask of shape (num_envs,)."""
+        # call all terms that are classes
+        for term_cfg in self._class_term_split.eager:
+            eager(term_cfg.func.reset, env_mask=env_mask)
+        for term_cfg in self._class_term_split.captured:
+            term_cfg.func.reset(env_mask=env_mask)
+        # reset terms with history
+        if any(self._group_obs_term_history_buffer.values()):
+            eager(self._reset_history, env_mask)
+        # call all modifiers/noise models that are classes
+        for mod in self._group_obs_class_instances:
+            mod.reset(env_mask=env_mask)
+
+        # nothing to log here
+        return {}
+
+    def _reset_history(self, env_mask: wp.array) -> None:
+        """Resets the history buffers of the selected environments, which are not capture-safe.
+
+        The buffers are looked up at call time because :meth:`_compute_term` can replace them, while :func:`eager`
+        replays its recorded arguments.
+        """
+        for history_buffers in self._group_obs_term_history_buffer.values():
+            for circular_buffer in history_buffers.values():
+                circular_buffer.reset(env_mask=env_mask)
 
     def compute(
         self, update_history: bool = False, return_cloned_output: bool = True
@@ -424,12 +442,18 @@ def compute(
             The observations are either concatenated into a single tensor or returned as a dictionary
             with keys corresponding to the term's name.
         """
-        return self._run_steps("compute", update_history=update_history, return_cloned_output=return_cloned_output)
-
-    def _compute_output(
-        self, update_history: bool = False, return_cloned_output: bool = True
-    ) -> dict[str, torch.Tensor | dict[str, torch.Tensor]]:
-        """Return the persistent observation buffer, or a clone of it."""
+        obs_buffer = self._compute(update_history)
+        # clone outside the recorded stage
+        return clone_obs_buffer(obs_buffer) if return_cloned_output else obs_buffer
+
+    @captured
+    def _compute(self, update_history: bool) -> dict[str, torch.Tensor | dict[str, torch.Tensor]]:
+        """Compute every term into the persistent buffers and return the persistent observation buffer."""
+        # Launch kernels for every group (writes into persistent buffers in-place).
+        for group_name, index in self._term_split.eager:
+            eager(self._compute_term, group_name, index, update_history=update_history)
+        for group_name, index in self._term_split.captured:
+            self._compute_term(group_name, index, update_history=update_history)
         # Build the obs buffer once (persistent refs to in-place-updated tensors/dicts).
         if self._obs_buffer is None:
             self._obs_buffer = {
@@ -440,8 +464,6 @@ def _compute_output(
                 )
                 for group_name in self._group_obs_term_names
             }
-        if return_cloned_output:
-            return clone_obs_buffer(self._obs_buffer)
         return self._obs_buffer
 
     def compute_group(self, group_name: str, update_history: bool = False) -> torch.Tensor | dict[str, torch.Tensor]:
@@ -497,9 +519,7 @@ def compute_group(self, group_name: str, update_history: bool = False) -> torch.
             return self._group_out_torch[group_name]
         return self._group_obs_dict[group_name]
 
-    def _compute_term(
-        self, group_name: str, index: int, update_history: bool = False, return_cloned_output: bool = True
-    ) -> None:
+    def _compute_term(self, group_name: str, index: int, update_history: bool = False) -> None:
         """Compute one observation term into its slice of the group buffer."""
         term_name = self._group_obs_term_names[group_name][index]
         term_cfg = self._group_obs_term_cfgs[group_name][index]
@@ -563,51 +583,6 @@ def _compute_term(
             else:
                 group_obs[term_name] = circular_buffer.buffer
 
-    """
-    Operations - Stage steps.
-    """
-
-    def _build_stage_steps(self, operation: str) -> list[tuple[bool, Callable[..., Any]]]:
-        if operation == "compute":
-            # Noise draws from the environment's shared random state, so terms keep their order.
-            steps = [
-                (
-                    is_warp_capturable(term_cfg.func, term_cfg.params) and term_cfg.history_length == 0,
-                    partial(self._compute_term, group_name, index),
-                )
-                for group_name, term_cfgs in self._group_obs_term_cfgs.items()
-                for index, term_cfg in enumerate(term_cfgs)
-            ]
-            return [*steps, (True, self._compute_output)]
-        if operation == "reset":
-            steps = []
-            for group_name, term_cfgs in self._group_obs_class_term_cfgs.items():
-                steps += [
-                    (is_warp_capturable(cfg.func, cfg.params), partial(self._reset_class_term, cfg))
-                    for cfg in term_cfgs
-                ]
-                # history buffers are not capture-safe
-                steps += [
-                    (False, partial(self._reset_history, group_name, term_name))
-                    for term_name in self._group_obs_term_history_buffer[group_name]
-                ]
-            steps += [(True, instance.reset) for instance in self._group_obs_class_instances]
-            return [*steps, (True, self._reset_output)]
-        return super()._build_stage_steps(operation)
-
-    @staticmethod
-    def _reset_class_term(term_cfg: ObservationTermCfg, env_mask: wp.array) -> None:
-        term_cfg.func.reset(env_mask=env_mask)
-
-    def _reset_history(self, group_name: str, term_name: str, env_mask: wp.array) -> None:
-        # computing a term before the first step may replace its buffer, so look it up on every call
-        self._group_obs_term_history_buffer[group_name][term_name].reset(env_mask=env_mask)
-
-    @staticmethod
-    def _reset_output(env_mask: wp.array) -> dict:
-        # nothing to log here
-        return {}
-
     def serialize(self) -> dict:
         """Serialize the observation term configurations for all active groups.
 
@@ -745,6 +720,8 @@ def _prepare_terms(self):  # noqa: C901
                 if group_cfg.history_length is not None:
                     term_cfg.history_length = group_cfg.history_length
                     term_cfg.flatten_history_dim = group_cfg.flatten_history_dim
+                # history buffers are not capture-safe yet
+                term_cfg.capturable = term_cfg.capturable and term_cfg.history_length == 0
                 # add term config to list
                 self._group_obs_term_names[group_name].append(term_name)
                 self._group_obs_term_cfgs[group_name].append(term_cfg)
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py
index f415ce03c907..381bad1e0386 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py
@@ -11,9 +11,8 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable, Sequence
-from functools import partial
-from typing import TYPE_CHECKING, Any
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 import torch
 import warp as wp
@@ -23,8 +22,9 @@
 
 from isaaclab_experimental.utils.warp import is_warp_capturable
 from isaaclab_experimental.utils.warp.kernels import compute_reset_scale, count_masked
+from isaaclab_experimental.utils.warp_capture import captured, eager, reset_captured_stages
 
-from .manager_base import ManagerBase, ManagerTermBase
+from .manager_base import ManagerBase, ManagerTermBase, split_resets, split_terms
 
 if TYPE_CHECKING:
     from isaaclab.envs import ManagerBasedRLEnv
@@ -70,29 +70,12 @@ def _reward_finalize(
             val = weighted * dt
             total += val
             episode_sums[term_idx, env_id] += val
+        else:
+            step_reward[env_id, term_idx] = 0.0
 
     reward_buf[env_id] = total
 
 
-@wp.kernel
-def _reward_pre_compute_reset(
-    # output
-    reward_buf: wp.array(dtype=wp.float32),
-    step_reward: wp.array(dtype=wp.float32, ndim=2),
-    term_outs: wp.array(dtype=wp.float32, ndim=2),
-):
-    """Reset per-step reward buffers.
-
-    Launched with dim = (num_envs,) to reset `reward_buf` and clear the corresponding row in `step_reward`.
-    This works even when `step_reward.shape[1] == 0` (no terms).
-    """
-    env_id = wp.tid()
-    reward_buf[env_id] = 0.0
-    for term_idx in range(term_outs.shape[0]):
-        step_reward[env_id, term_idx] = 0.0
-        term_outs[term_idx, env_id] = 0.0
-
-
 class RewardManager(ManagerBase):
     """Manager for computing reward signals for a given world.
 
@@ -130,6 +113,9 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv):
         # call the base class constructor (this will parse the terms config)
         super().__init__(cfg, env)
         self._term_name_to_term_idx = {name: i for i, name in enumerate(self._term_names)}
+        # skip terms with zero weight (kind of a micro-optimization)
+        self._term_split = split_terms([term_cfg for term_cfg in self._term_cfgs if term_cfg.weight != 0.0])
+        self._class_term_split = split_resets(self._class_term_cfgs)
 
         num_terms = len(self._term_names)
         self._num_terms = num_terms
@@ -228,71 +214,16 @@ def active_terms(self) -> list[str]:
     Operations.
     """
 
-    def reset(
-        self,
-        env_ids: Sequence[int] | torch.Tensor | None = None,
-        *,
-        env_mask: wp.array | None = None,
-    ) -> dict[str, torch.Tensor]:
+    @captured
+    def _reset(self, env_mask: wp.array) -> dict[str, torch.Tensor]:
         """Computes/reset episodic reward sums for masked envs (capturable core).
 
         Args:
-            env_ids: The specific environment indices to reset.
-                If None, all environments are considered.
             env_mask: Boolean Warp mask of shape (num_envs,) selecting reset environments.
-                If provided, takes precedence over ``env_ids``.
 
         Returns:
             A dictionary containing the information to log under the "Reward/{term_name}" key.
         """
-        # Mask-first path: captured callers must provide env_mask.
-        if env_mask is None or not isinstance(env_mask, wp.array):
-            if wp.get_device().is_capturing:
-                raise RuntimeError(
-                    "RewardManager.reset requires env_mask(wp.array[bool]) during capture. "
-                    "Do not pass env_ids on captured paths."
-                )
-            env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask)
-        return self._run_steps("reset", env_mask=env_mask)
-
-    def compute(self, dt: float) -> torch.Tensor:
-        """Computes the reward signal as a weighted sum of individual terms.
-
-        This function calls each reward term managed by the class and adds them to compute the net
-        reward signal. It also updates the episodic sums corresponding to individual reward terms.
-
-        Args:
-            dt: The time-step interval of the environment.
-
-        Returns:
-            The net reward signal of shape (num_envs,).
-        """
-        return self._run_steps("compute", dt=dt)
-
-    """
-    Operations - Stage steps.
-    """
-
-    def _build_stage_steps(self, operation: str) -> list[tuple[bool, Callable[..., Any]]]:
-        if operation == "compute":
-            term_steps = [
-                (is_warp_capturable(term_cfg.func, term_cfg.params), partial(self._compute_term, index))
-                for index, term_cfg in enumerate(self._term_cfgs)
-            ]
-            # Terms write their own output columns and the finalize step sums them in a fixed order, so
-            # running the capturable terms first leaves the reward unchanged bitwise.
-            term_steps.sort(key=lambda step: not step[0])
-            return [(True, self._reset_step_buffers), *term_steps, (True, self._finalize_step)]
-        if operation == "reset":
-            term_steps = [
-                (is_warp_capturable(term_cfg.func, term_cfg.params), partial(self._reset_term, term_cfg))
-                for term_cfg in self._class_term_cfgs
-            ]
-            return [(True, self._reset_episode_sums), *term_steps]
-        return super()._build_stage_steps(operation)
-
-    def _reset_episode_sums(self, env_mask: wp.array) -> dict[str, torch.Tensor]:
-        """Average the episodic sums of the selected environments and zero them."""
         self._episode_sum_avg_wp.zero_()
         self._reset_count_wp.zero_()
         self._reset_scale_wp.zero_()
@@ -312,38 +243,43 @@ def _reset_episode_sums(self, env_mask: wp.array) -> dict[str, torch.Tensor]:
                 inputs=[env_mask, self._reset_scale_wp, self._episode_sums_wp, self._episode_sum_avg_wp],
                 device=self.device,
             )
-        return self._reset_extras
 
-    def _reset_term(self, term_cfg: RewardTermCfg, env_mask: wp.array) -> dict[str, torch.Tensor]:
-        """Reset a class term and merge its logging values into the reset extras.
+        # reset all the reward terms; class terms may expose logging values as
+        # persistent tensor views (see :meth:`ManagerTermBase.reset`), which are
+        # merged into the manager's reset extras. Merging the same views again on
+        # subsequent resets is a no-op, so this stays capture-safe.
+        for term_cfg in self._class_term_split.eager:
+            term_extras = eager(term_cfg.func.reset, env_mask=env_mask)
+            if term_extras:
+                self._reset_extras.update(term_extras)
+        for term_cfg in self._class_term_split.captured:
+            term_extras = term_cfg.func.reset(env_mask=env_mask)
+            if term_extras:
+                self._reset_extras.update(term_extras)
 
-        Class terms may expose logging values as persistent tensor views (see :meth:`ManagerTermBase.reset`).
-        Merging the same views again on later resets is a no-op, so this stays capture-safe.
-        """
-        term_extras = term_cfg.func.reset(env_mask=env_mask)
-        if term_extras:
-            self._reset_extras.update(term_extras)
         return self._reset_extras
 
-    def _reset_step_buffers(self, dt: float) -> None:
-        """Zero the reward, step reward and term output buffers in a single kernel launch."""
-        # TODO: Investigate performance diff between two .fill_ and kernel launch
-        wp.launch(
-            kernel=_reward_pre_compute_reset,
-            dim=self.num_envs,
-            inputs=[self._reward_wp, self._step_reward_wp, self._term_outs_wp],
-            device=self.device,
-        )
+    @captured
+    def compute(self, dt: float) -> torch.Tensor:
+        """Computes the reward signal as a weighted sum of individual terms.
+
+        This function calls each reward term managed by the class and adds them to compute the net
+        reward signal. It also updates the episodic sums corresponding to individual reward terms.
 
-    def _compute_term(self, index: int, dt: float) -> None:
-        """Compute one term into its pre-zeroed output column (raw, unweighted)."""
-        term_cfg = self._term_cfgs[index]
-        # skip if weight is zero (kind of a micro-optimization)
-        if term_cfg.weight != 0.0:
+        Args:
+            dt: The time-step interval of the environment.
+
+        Returns:
+            The net reward signal of shape (num_envs,).
+        """
+        # iterate over all the reward terms (Python loop; per-term math is warp)
+        # compute each term into the persistent warp buffer (raw, unweighted)
+        for term_cfg in self._term_split.eager:
+            eager(term_cfg.func, self._env, term_cfg.out, **term_cfg.params)
+        for term_cfg in self._term_split.captured:
             term_cfg.func(self._env, term_cfg.out, **term_cfg.params)
 
-    def _finalize_step(self, dt: float) -> torch.Tensor:
-        """Update the total reward, episodic sums and step rewards in a single kernel launch."""
+        # update total reward, episodic sums and step rewards in a single kernel launch
         wp.launch(
             kernel=_reward_finalize,
             dim=self.num_envs,
@@ -384,8 +320,10 @@ def set_term_cfg(self, term_name: str, cfg: RewardTermCfg):
         self._term_cfgs[term_idx] = cfg
         # keep on-device weights in sync (call this to update weights used in compute)
         self._term_weights_tensor_view[term_idx] = float(cfg.weight)
-        # the replaced term may differ in capturability
-        self._clear_stage_steps()
+        # recorded compute graphs hold the replaced term, which may also differ in capturability
+        cfg.capturable = is_warp_capturable(cfg.func, cfg.params)
+        self._term_split = split_terms([term_cfg for term_cfg in self._term_cfgs if term_cfg.weight != 0.0])
+        reset_captured_stages(self)
 
     def get_term_cfg(self, term_name: str) -> RewardTermCfg:
         """Gets the configuration for the specified term.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/scene_entity_cfg.py b/source/isaaclab_experimental/isaaclab_experimental/managers/scene_entity_cfg.py
index f6d30489eeca..1f47d2c4facd 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/scene_entity_cfg.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/scene_entity_cfg.py
@@ -26,10 +26,10 @@ class SceneEntityCfg(_SceneEntityCfg):
     """Scene entity configuration with an optional Warp joint mask.
 
     Notes:
-    - `joint_mask` is intended for Warp kernels only.
+    - `joint_mask_wp` is intended for Warp kernels only.
     """
 
-    joint_mask: wp.array | None = None
+    joint_mask_wp: wp.array | None = None
 
     """Integer indices of selected joints — used for subset-sized gathers where a boolean mask
     cannot provide the mapping from output index k to joint index."""
@@ -38,6 +38,9 @@ class SceneEntityCfg(_SceneEntityCfg):
     """Integer indices of selected bodies — used for subset-sized body gathers."""
     body_ids_wp: wp.array | None = None
 
+    """Boolean mask of selected bodies — used for masked per-body writes."""
+    body_mask_wp: wp.array | None = None
+
     @classmethod
     def from_stable(cls, stable: _SceneEntityCfg) -> SceneEntityCfg:
         """Build a warp scene-entity cfg from a stable one.
@@ -63,11 +66,18 @@ def resolve(self, scene: InteractiveScene):
                 mask_list = [False] * entity.num_joints
                 for idx in joint_ids_list:
                     mask_list[idx] = True
-            self.joint_mask = wp.array(mask_list, dtype=wp.bool, device=scene.device)
+            self.joint_mask_wp = wp.array(mask_list, dtype=wp.bool, device=scene.device)
             self.joint_ids_wp = wp.array(joint_ids_list, dtype=wp.int32, device=scene.device)
 
-        # -- Warp body ids
+        # -- Warp body ids / mask
         if self.body_ids is not None and self.body_ids != slice(None):
-            self.body_ids_wp = wp.array(list(self.body_ids), dtype=wp.int32, device=scene.device)
+            body_ids_list = list(self.body_ids)
         elif hasattr(entity, "num_bodies"):
-            self.body_ids_wp = wp.array(list(range(entity.num_bodies)), dtype=wp.int32, device=scene.device)
+            body_ids_list = list(range(entity.num_bodies))
+        else:
+            return
+        self.body_ids_wp = wp.array(body_ids_list, dtype=wp.int32, device=scene.device)
+        if hasattr(entity, "num_bodies"):
+            selected = set(body_ids_list)
+            mask_list = [idx in selected for idx in range(entity.num_bodies)]
+            self.body_mask_wp = wp.array(mask_list, dtype=wp.bool, device=scene.device)
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py
index 32174f3eb33c..9191679aec12 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py
@@ -15,9 +15,8 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable, Sequence
-from functools import partial
-from typing import TYPE_CHECKING, Any
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 import torch
 import warp as wp
@@ -25,35 +24,14 @@
 
 from isaaclab.managers.manager_term_cfg import TerminationTermCfg
 
-from isaaclab_experimental.utils.warp import is_warp_capturable
+from isaaclab_experimental.utils.warp_capture import captured, eager
 
-from .manager_base import ManagerBase, ManagerTermBase
+from .manager_base import ManagerBase, ManagerTermBase, split_resets, split_terms
 
 if TYPE_CHECKING:
     from isaaclab.envs import ManagerBasedRLEnv
 
 
-@wp.kernel
-def _termination_pre_compute_reset(
-    # output
-    term_dones: wp.array(dtype=wp.bool, ndim=2),
-    truncated: wp.array(dtype=wp.bool),
-    terminated: wp.array(dtype=wp.bool),
-    dones: wp.array(dtype=wp.bool),
-):
-    """Reset per-step termination buffers.
-
-    Launched with dim = (num_envs,) to reset per-env flags and clear the corresponding row in `term_dones`.
-    This works even when `term_dones.shape[1] == 0` (no terms).
-    """
-    env_id = wp.tid()
-    truncated[env_id] = False
-    terminated[env_id] = False
-    dones[env_id] = False
-    for term_idx in range(term_dones.shape[1]):
-        term_dones[env_id, term_idx] = False
-
-
 @wp.kernel
 def _termination_finalize(
     # input
@@ -128,6 +106,8 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv):
         super().__init__(cfg, env)
 
         self._term_name_to_term_idx = {name: i for i, name in enumerate(self._term_names)}
+        self._term_split = split_terms(self._term_cfgs)
+        self._class_term_split = split_resets(self._class_term_cfgs)
 
         # persistent buffers (Warp)
         num_terms = len(self._term_names)
@@ -246,70 +226,16 @@ def term_dones_wp(self) -> wp.array:
     Operations.
     """
 
-    def reset(
-        self,
-        env_ids: Sequence[int] | torch.Tensor | None = None,
-        *,
-        env_mask: wp.array | None = None,
-    ) -> dict[str, torch.Tensor]:
+    @captured
+    def _reset(self, env_mask: wp.array) -> dict[str, torch.Tensor]:
         """Reset termination stats and class terms; return pre-allocated extras.
 
         Args:
-            env_ids: The specific environment indices to reset.
-                If None, all environments are considered.
             env_mask: Boolean Warp mask of shape (num_envs,) selecting reset environments.
-                If provided, takes precedence over ``env_ids``.
 
         Returns:
             A dictionary containing the information to log under the "Termination/{term_name}" key.
         """
-        # Mask-first path: captured callers must provide env_mask.
-        if env_mask is None or not isinstance(env_mask, wp.array):
-            if wp.get_device().is_capturing:
-                raise RuntimeError(
-                    "TerminationManager.reset requires env_mask(wp.array[bool]) during capture. "
-                    "Do not pass env_ids on captured paths."
-                )
-            env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask)
-        return self._run_steps("reset", env_mask=env_mask)
-
-    @property
-    def episode_termination_extras(self) -> dict[str, torch.Tensor]:
-        """Pre-allocated reset logging extras for termination terms."""
-        return self._reset_extras
-
-    def compute(self) -> torch.Tensor:
-        """Computes the termination signal as union of individual terms.
-
-        Returns:
-            The combined termination signal of shape (num_envs,).
-        """
-        return self._run_steps("compute")
-
-    """
-    Operations - Stage steps.
-    """
-
-    def _build_stage_steps(self, operation: str) -> list[tuple[bool, Callable[..., Any]]]:
-        if operation == "compute":
-            term_steps = [
-                (is_warp_capturable(term_cfg.func, term_cfg.params), partial(self._compute_term, term_cfg))
-                for term_cfg in self._term_cfgs
-            ]
-            # Terms write their own done columns and the finalize step combines them in a fixed order, so
-            # running the capturable terms first leaves the dones unchanged.
-            term_steps.sort(key=lambda step: not step[0])
-            return [(True, self._reset_step_buffers), *term_steps, (True, self._finalize_step)]
-        if operation == "reset":
-            term_steps = [
-                (is_warp_capturable(term_cfg.func, term_cfg.params), partial(self._reset_term, term_cfg))
-                for term_cfg in self._class_term_cfgs
-            ]
-            return [(True, self._average_last_episode_dones), *term_steps]
-        return super()._build_stage_steps(operation)
-
-    def _average_last_episode_dones(self, env_mask: wp.array) -> dict[str, torch.Tensor]:
-        """Average each term's last-episode done flags over all environments, as the stable manager does."""
         if len(self._term_names) > 0:
             self._term_done_avg_wp.zero_()
             wp.launch(
@@ -318,26 +244,31 @@ def _average_last_episode_dones(self, env_mask: wp.array) -> dict[str, torch.Ten
                 inputs=[self._last_episode_dones_wp, self._term_done_avg_wp],
                 device=self.device,
             )
+        for term_cfg in self._class_term_split.eager:
+            eager(term_cfg.func.reset, env_mask=env_mask)
+        for term_cfg in self._class_term_split.captured:
+            term_cfg.func.reset(env_mask=env_mask)
         return self._reset_extras
 
-    def _reset_term(self, term_cfg: TerminationTermCfg, env_mask: wp.array) -> dict[str, torch.Tensor]:
-        term_cfg.func.reset(env_mask=env_mask)
+    @property
+    def episode_termination_extras(self) -> dict[str, torch.Tensor]:
+        """Pre-allocated reset logging extras for termination terms."""
         return self._reset_extras
 
-    def _reset_step_buffers(self) -> None:
-        """Clear the done buffers in a single kernel launch."""
-        wp.launch(
-            kernel=_termination_pre_compute_reset,
-            dim=self.num_envs,
-            inputs=[self._term_dones_wp, self._truncated_wp, self._terminated_wp, self._dones_wp],
-            device=self.device,
-        )
+    @captured
+    def compute(self) -> torch.Tensor:
+        """Computes the termination signal as union of individual terms.
 
-    def _compute_term(self, term_cfg: TerminationTermCfg) -> None:
-        term_cfg.func(self._env, term_cfg.out, **term_cfg.params)
+        Returns:
+            The combined termination signal of shape (num_envs,).
+        """
+        # iterate over all the termination terms (fixed list; per-term math is Warp)
+        for term_cfg in self._term_split.eager:
+            eager(term_cfg.func, self._env, term_cfg.out, **term_cfg.params)
+        for term_cfg in self._term_split.captured:
+            term_cfg.func(self._env, term_cfg.out, **term_cfg.params)
 
-    def _finalize_step(self) -> torch.Tensor:
-        """Finalize the dones and update the last-episode term flags in a single kernel launch."""
+        # finalize dones and update last-episode term flags (single kernel launch)
         wp.launch(
             kernel=_termination_finalize,
             dim=self.num_envs,
@@ -351,6 +282,7 @@ def _finalize_step(self) -> torch.Tensor:
             ],
             device=self.device,
         )
+
         return self._dones_tensor_view
 
     def get_term(self, name: str) -> torch.Tensor:
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/__init__.pyi b/source/isaaclab_experimental/isaaclab_experimental/utils/__init__.pyi
index 00660ba9775e..f0ddccdad640 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/utils/__init__.pyi
+++ b/source/isaaclab_experimental/isaaclab_experimental/utils/__init__.pyi
@@ -4,7 +4,9 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 __all__ = [
-    "WarpGraphCache",
+    "CapturedStage",
+    "captured",
+    "eager",
     "clone_obs_buffer",
     "buffers",
     "modifiers",
@@ -13,5 +15,5 @@ __all__ = [
 ]
 
 from .torch_utils import clone_obs_buffer
-from .warp_graph_cache import WarpGraphCache
+from .warp_capture import CapturedStage, captured, eager
 from . import buffers, modifiers, noise, warp
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py
index a91c051b0d08..2d071b823a46 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py
@@ -5,7 +5,7 @@
 
 """Warp utility functions and shared kernels for isaaclab_experimental."""
 
-from .kernels import compute_reset_scale, count_masked, zero_masked_int64
+from .kernels import compute_reset_scale, count_masked
 from .utils import (
     WarpCapturable,
     is_warp_capturable,
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/kernels.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/kernels.py
index 6320240740ce..8d3f9e65d49a 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/kernels.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/kernels.py
@@ -10,14 +10,6 @@
 import warp as wp
 
 
-@wp.kernel
-def zero_masked_int64(mask: wp.array(dtype=wp.bool), values: wp.array(dtype=wp.int64)):
-    """Zero the entries of a one-dimensional int64 array selected by ``mask``."""
-    index = wp.tid()
-    if mask[index]:
-        values[index] = wp.int64(0)
-
-
 @wp.kernel
 def count_masked(
     mask: wp.array(dtype=wp.bool),
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py
index f43dd0f425f9..cfff724ce24d 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py
@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import inspect
 from collections.abc import Callable, Sequence
 from typing import Any
 
@@ -117,15 +118,16 @@ def decorator(func):
 def is_warp_capturable(func, params: dict[str, Any] | None = None) -> bool:
     """Check if a term is CUDA-graph-capturable.
 
-    Checks ``_warp_capturable`` on the term (a function, a class, or a class instance) and on its
+    Checks ``_warp_capturable`` on the term function or class (the class of an instance) and on its
     ``__wrapped__`` target. The annotation is a bool, or a predicate that decides from the term's
-    parameters (see :class:`WarpCapturable`). A class term may also set ``self._warp_capturable`` in
-    ``__init__`` to decide per instance. Returns True (capturable) if no annotation is found.
+    parameters (see :class:`WarpCapturable`). Returns True (capturable) if no annotation is found.
 
     Args:
         func: The term function, class, or class instance.
         params: The term's parameters, passed to a predicate annotation. Defaults to None (no parameters).
     """
+    if not (isinstance(func, type) or inspect.isroutine(func)):
+        func = type(func)
     for f in (func, getattr(func, "__wrapped__", None)):
         if f is not None:
             val = getattr(f, "_warp_capturable", None)
@@ -164,11 +166,10 @@ def base_height_l2(env, out, target_height, asset_cfg=..., sensor_cfg=None):
     - ``@WarpCapturable(True)`` or no decorator: capturable, returned unwrapped.
     - ``@WarpCapturable(False)`` on a function: sets ``func._warp_capturable = False`` and wraps it with a
       runtime guard that raises if ``wp.get_device().is_capturing`` is ``True``.
-    - ``@WarpCapturable(False)`` on a class term: annotates the class and guards its ``__call__``.
+    - ``@WarpCapturable(False)`` on a class term: annotates the class and guards its ``__call__``. The term's
+      ``reset`` is still recorded; decorate the ``reset`` method itself to keep it out of the recording.
     - ``@WarpCapturable(predicate)``: the term is capturable for the parameters the predicate accepts. The
       predicate receives the term's parameters as a dict, and the guard raises only for rejected ones.
-
-    A class term can also decide per instance by setting ``self._warp_capturable`` in ``__init__``.
     """
 
     def __init__(self, capturable: bool | Callable[[dict[str, Any]], bool], *, reason: str | None = None):
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_capture.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_capture.py
new file mode 100644
index 000000000000..1f0d33ef8490
--- /dev/null
+++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_capture.py
@@ -0,0 +1,298 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""CUDA graph capture of Warp frontend stages: a :func:`captured` method records its calls and replays them, and
+:func:`eager` runs a call that cannot be recorded at its place in the recording.
+
+.. code-block:: python
+
+    @captured
+    def compute(self, dt):
+        eager(term_cfg.func, self._env, term_cfg.out, **term_cfg.params)
+"""
+
+from __future__ import annotations
+
+import contextlib
+import functools
+import gc
+import os
+from collections.abc import Callable
+from typing import Any, ClassVar
+
+import torch
+import warp as wp
+
+from isaaclab.physics import PhysicsEvent, PhysicsManager
+from isaaclab.utils.version import has_kit
+
+SYNC_DEBUG_ENV_VAR = "ISAACLAB_SYNC_DEBUG"
+"""Set ``ISAACLAB_SYNC_DEBUG=1`` to run eager stage calls under ``torch.cuda.set_sync_debug_mode("error")``,
+so a hidden GPU-to-host synchronization inside a stage raises instead of stalling silently."""
+
+
+def captured(method: Callable[..., Any]) -> Callable[..., Any]:
+    """Run a method through a :class:`CapturedStage` stored on its instance.
+
+    The method body is the plain eager code of the stage, with the calls that cannot be recorded wrapped in
+    :func:`eager`. Its instance must expose ``device``. The instance keeps its stages in ``_captured_stages``,
+    keyed by the decorated method, so an override and the method it extends own separate stages.
+
+    Args:
+        method: The method to decorate.
+    """
+
+    @functools.wraps(method)
+    def wrapper(self, *args: Any, **kwargs: Any) -> Any:
+        stages = self.__dict__.setdefault("_captured_stages", {})
+        stage = stages.get(wrapper)
+        if stage is None:
+            stage = stages[wrapper] = CapturedStage(method, self.device)
+        return stage(self, *args, **kwargs)
+
+    return wrapper
+
+
+def eager(fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
+    """Call a function that cannot be recorded from inside a :class:`CapturedStage`.
+
+    Outside a recording this is a plain call. While a stage records, the call ends the current graph, runs eagerly
+    and starts the next graph, and every replay of the stage repeats it at the same place with the same arguments,
+    so the arguments must stay valid across calls like the arrays a graph reads. Replays discard its result.
+
+    Args:
+        fn: The function to call.
+        *args: Positional arguments of the call.
+        **kwargs: Keyword arguments of the call.
+
+    Returns:
+        The result of the call.
+
+    Raises:
+        RuntimeError: If the current device is capturing a graph that no :class:`CapturedStage` records.
+    """
+    recording = CapturedStage._recording
+    if recording is not None and recording.capturing:
+        return recording.split(fn, args, kwargs)
+    if wp.get_device().is_capturing:
+        name = getattr(fn, "__qualname__", type(fn).__name__)
+        raise RuntimeError(
+            f"'{name}' cannot be recorded, so it cannot run inside a graph capture that no CapturedStage started."
+        )
+    return fn(*args, **kwargs)
+
+
+def reset_captured_stages(owner: Any) -> None:
+    """Drop the graphs recorded by the :func:`captured` methods of an instance."""
+    for stage in owner.__dict__.get("_captured_stages", {}).values():
+        stage.reset()
+
+
+class CapturedStage:
+    """A stage of the Warp frontend, run eagerly or through CUDA graphs recorded from its own calls.
+
+    * The first call with a given argument layout records the stage without an eager warm-up: the recording is its
+      first execution. Later calls with the same array pointers and scalar values replay that recording and return
+      the result of the recording call, so a stage returns buffers its owner keeps.
+    * Calls that cannot be recorded go through :func:`eager`. A recording ends its graph at each of them, runs the
+      call and starts the next graph; a replay launches the graphs and repeats the calls with their recorded
+      arguments, in the recorded order.
+    * :attr:`enabled` is process-wide and set by the Warp environments. A stage runs eagerly while it is False and
+      on CPU devices.
+    * :attr:`generation` advances when the physics backend rebinds its buffers or stops (see :meth:`invalidate_on`),
+      and every stage records again on its next call.
+    * A stage called while another stage records runs inline, so its graphs and eager calls become part of that
+      recording.
+    """
+
+    enabled: ClassVar[bool] = False
+    """Whether stages record CUDA graphs. The Warp environments set it; the torch frontend never captures."""
+
+    generation: ClassVar[int] = 0
+    """Counter of physics rebinds and stops; a stage recorded under an older value records again."""
+
+    _callback_handles: ClassVar[list] = []
+
+    _recording: ClassVar[_Recording | None] = None
+    """The recording in progress, which :func:`eager` calls split."""
+
+    def __init__(self, fn: Callable[..., Any], device: wp.DeviceLike):
+        """Initialize the stage.
+
+        Args:
+            fn: Callable implementing the stage. Its GPU work must be capturable, except for :func:`eager` calls.
+            device: Warp device that runs the stage.
+        """
+        self._fn = fn
+        self._device = wp.get_device(device)
+        self._sync_debug = os.environ.get(SYNC_DEBUG_ENV_VAR, "0") == "1"
+        self._recordings: dict[tuple, tuple[list, Any]] = {}
+        self._generation = CapturedStage.generation
+
+    @property
+    def num_graphs(self) -> int:
+        """Number of graphs currently recorded for this stage."""
+        return sum(isinstance(step, wp.Graph) for steps, _ in self._recordings.values() for step in steps)
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Run the stage and return its result."""
+        if self._device.is_capturing:
+            return self._fn(*args, **kwargs)
+        if not (CapturedStage.enabled and self._device.is_cuda):
+            return self._run_eager(self._fn, args, kwargs)
+        if self._generation != CapturedStage.generation:
+            self.reset()
+            self._generation = CapturedStage.generation
+        key = _argument_key(args, kwargs)
+        recorded = self._recordings.get(key)
+        if recorded is None:
+            self._recordings[key] = recorded = self._record(args, kwargs)
+            return recorded[1]
+        steps, result = recorded
+        for step in steps:
+            if isinstance(step, wp.Graph):
+                wp.capture_launch(step)
+            else:
+                self._run_eager(*step)
+        return result
+
+    def reset(self) -> None:
+        """Drop the recordings; the stage records again on its next call."""
+        if self._recordings:
+            # graphs may still be executing; drain the device before releasing them
+            wp.synchronize_device(self._device)
+            self._recordings.clear()
+
+    @classmethod
+    def invalidate_on(cls, physics_manager: type[PhysicsManager]) -> None:
+        """Record every stage again after the physics backend rebinds its buffers or stops.
+
+        Args:
+            physics_manager: Physics manager class that dispatches the lifecycle events.
+        """
+        cls._callback_handles = [
+            physics_manager.register_callback(cls.invalidate, event)
+            for event in (PhysicsEvent.PHYSICS_READY, PhysicsEvent.STOP)
+        ]
+
+    @classmethod
+    def disable(cls) -> None:
+        """Run every stage eagerly again and stop following the physics backend."""
+        for handle in cls._callback_handles:
+            handle.deregister()
+        cls._callback_handles = []
+        cls.enabled = False
+
+    @classmethod
+    def invalidate(cls, payload: Any = None) -> None:
+        """Record every stage again on its next call."""
+        cls.generation += 1
+
+    def _record(self, args: tuple, kwargs: dict[str, Any]) -> tuple[list, Any]:
+        """Record the stage as its first run and return the recorded steps and the result."""
+        recording = _Recording(self._device, self._run_eager)
+        enclosing, CapturedStage._recording = CapturedStage._recording, recording
+        try:
+            with _paused_gc(), wp.ScopedStream(recording.capture_stream):
+                recording.begin()
+                try:
+                    result = self._fn(*args, **kwargs)
+                except BaseException:
+                    recording.abort()
+                    raise
+                recording.end()
+        finally:
+            CapturedStage._recording = enclosing
+        return recording.steps, result
+
+    def _run_eager(self, fn: Callable[..., Any], args: tuple, kwargs: dict[str, Any]) -> Any:
+        """Run eagerly, under the synchronization trap when requested."""
+        if not self._sync_debug:
+            return fn(*args, **kwargs)
+        previous = torch.cuda.get_sync_debug_mode()
+        torch.cuda.set_sync_debug_mode("error")
+        try:
+            return fn(*args, **kwargs)
+        finally:
+            torch.cuda.set_sync_debug_mode(previous)
+
+
+class _Recording:
+    """The steps of one stage recording: CUDA graphs split around :func:`eager` calls, launched as they end."""
+
+    def __init__(self, device: wp.Device, run_eager: Callable[[Callable[..., Any], tuple, dict], Any]):
+        # graphs are captured on the capture stream and launched on the stream the stage runs on
+        self.stream = wp.get_stream(device)
+        self.capture_stream = self.stream
+        self.mode = wp.CaptureMode.THREAD_LOCAL
+        if has_kit():
+            # RTX uses the legacy CUDA stream. A nonblocking stream avoids implicit synchronization.
+            self.capture_stream = wp.stream_from_torch(torch.cuda.Stream(device=device.alias))
+            self.mode = wp.CaptureMode.RELAXED
+        self.steps: list = []
+        self.capturing = False
+        self._run_eager = run_eager
+
+    def begin(self) -> None:
+        """Start the next graph."""
+        wp.capture_begin(stream=self.capture_stream, capture_mode=self.mode)
+        self.capturing = True
+
+    def end(self) -> None:
+        """End the current graph and launch it, since the recording is the first run of the stage."""
+        self.capturing = False
+        graph = wp.capture_end(stream=self.capture_stream)
+        wp.capture_launch(graph, stream=self.stream)
+        self.steps.append(graph)
+
+    def abort(self) -> None:
+        """End the current graph without launching it, after the stage raised."""
+        if self.capturing:
+            self.capturing = False
+            with contextlib.suppress(Exception):
+                wp.capture_end(stream=self.capture_stream)
+
+    def split(self, fn: Callable[..., Any], args: tuple, kwargs: dict[str, Any]) -> Any:
+        """Run a call that cannot be recorded between the current graph and the next one."""
+        self.end()
+        with wp.ScopedStream(self.stream, sync_enter=False):
+            result = self._run_eager(fn, args, kwargs)
+        self.steps.append((fn, args, kwargs))
+        self.begin()
+        return result
+
+
+@contextlib.contextmanager
+def _paused_gc():
+    """Keep collection-driven frees out of a recording, as ``NewtonQueries.capture_graph`` does.
+
+    Objects allocated during the recording stay in the youngest generation, so collecting that generation afterwards
+    frees the recording's garbage without the full collection of the whole heap.
+    """
+    was_enabled = gc.isenabled()
+    gc.disable()
+    try:
+        yield
+    finally:
+        if was_enabled:
+            gc.enable()
+            gc.collect(0)
+
+
+def _argument_key(args: tuple, kwargs: dict[str, Any]) -> tuple:
+    """Key a stage call by the pointers of its arrays and the values of its scalars."""
+    return tuple(_value_key(value) for value in args) + tuple(
+        (name, _value_key(value)) for name, value in kwargs.items()
+    )
+
+
+def _value_key(value: Any) -> Any:
+    if isinstance(value, wp.array):
+        return value.ptr
+    if isinstance(value, torch.Tensor):
+        return value.data_ptr()
+    if value is None or isinstance(value, (bool, int, float, str)):
+        return value
+    return id(value)
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py
deleted file mode 100644
index 8307158fa018..000000000000
--- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py
+++ /dev/null
@@ -1,268 +0,0 @@
-# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
-# All rights reserved.
-#
-# SPDX-License-Identifier: BSD-3-Clause
-
-"""CUDA graph capture and replay for Warp frontend stages."""
-
-from __future__ import annotations
-
-import os
-from collections.abc import Callable
-from typing import Any
-
-import torch
-import warp as wp
-from isaaclab_newton.physics import NewtonQueries
-
-from isaaclab.physics import PhysicsEvent, PhysicsManager
-from isaaclab.utils.timer import Timer
-from isaaclab.utils.version import has_kit
-
-CAPTURE_ENV_VAR = "ISAACLAB_WARP_CAPTURE"
-"""Set ``ISAACLAB_WARP_CAPTURE=0`` to run every stage eagerly (debugging and eager-vs-captured comparisons)."""
-
-SYNC_DEBUG_ENV_VAR = "ISAACLAB_SYNC_DEBUG"
-"""Set ``ISAACLAB_SYNC_DEBUG=1`` to run eager stage calls under ``torch.cuda.set_sync_debug_mode("error")``,
-so a hidden GPU-to-host synchronization inside a stage raises instead of stalling silently."""
-
-StageStep = tuple[bool, Callable[..., Any]]
-"""One step of a stage: whether its work may be recorded into a CUDA graph, and the callable that runs it."""
-
-
-class WarpGraphCache:
-    """Run Warp frontend stages eagerly or through cached CUDA graphs.
-
-    Graphs follow the physics backend's capture lifecycle:
-
-    * A stage is recorded without an eager warm-up: its first execution is the launch of the
-      newly recorded graph, so a stateful stage advances exactly once per call. Stage owners
-      allocate their buffers before the first call.
-    * Recording waits until :meth:`arm`. Environments arm the cache on each step, so stages
-      called only while constructing or resetting the environment run eagerly until then.
-    * A graph is keyed by the pointers and scalar values of the stage arguments. Passing a
-      reallocated array, or a different scalar, records the stage again.
-    * :meth:`invalidate` drops every graph and holds recording until the next :meth:`arm`.
-      :meth:`invalidate_on` does so whenever the physics backend rebinds its buffers or stops.
-    * A stage made of steps (:meth:`call_steps`) records each run of consecutive capturable steps
-      as its own graph and runs the other steps eagerly in between, keeping the step order.
-
-    A stage group is the stage name up to its first underscore. A group registered as not
-    capturable, a CPU device, or ``ISAACLAB_WARP_CAPTURE=0`` runs its stages eagerly.
-    """
-
-    def __init__(self, device: wp.DeviceLike = None):
-        """Initialize the cache.
-
-        Args:
-            device: Warp device that runs the stages. Defaults to the current Warp device.
-        """
-        self._device = wp.get_device(device)
-        self._enabled = self._device.is_cuda and os.environ.get(CAPTURE_ENV_VAR, "1") != "0"
-        self._sync_debug = os.environ.get(SYNC_DEBUG_ENV_VAR, "0") == "1"
-        self._armed = False
-        self._graphs: dict[str, tuple[tuple, wp.Graph, Any]] = {}
-        self._runs: dict[
-            int, tuple[tuple[StageStep, ...], tuple[tuple[bool, tuple[Callable[..., Any], ...]], ...]]
-        ] = {}
-        self._capturable: dict[str, bool] = {}
-        self._callback_handles = []
-
-    @property
-    def captured_stages(self) -> tuple[str, ...]:
-        """Stages currently backed by a recorded graph, sorted by name."""
-        return tuple(sorted(self._graphs))
-
-    def call(
-        self,
-        stage: str,
-        fn: Callable[..., Any],
-        /,
-        *args: Any,
-        output: Callable[[Any], Any] | None = None,
-        timer: bool = False,
-        **kwargs: Any,
-    ) -> Any:
-        """Run a single-step stage eagerly or through its recorded graph.
-
-        Args:
-            stage: Stage identifier in the form ``"GroupName_function_name"``.
-            fn: Callable implementing the stage. Its GPU work must be capturable when the group is.
-            *args: Positional arguments forwarded to :paramref:`fn`.
-            output: Transform applied to the stage result on every call, outside the graph.
-            timer: Whether to time the stage.
-            **kwargs: Keyword arguments forwarded to :paramref:`fn`.
-
-        Returns:
-            The stage result, optionally transformed by :paramref:`output`. A replayed stage returns
-            the result of the call that recorded it.
-        """
-        return self._run(stage, ((True, (fn,)),), args, kwargs, output, timer)
-
-    def call_steps(
-        self,
-        stage: str,
-        steps: tuple[StageStep, ...],
-        /,
-        *args: Any,
-        output: Callable[[Any], Any] | None = None,
-        timer: bool = False,
-        **kwargs: Any,
-    ) -> Any:
-        """Run the steps of a stage in order, recording each run of consecutive capturable steps.
-
-        Every step receives the stage arguments. A run of consecutive capturable steps is recorded into
-        one graph, named ``"<stage>[<run index>]"`` when the stage has more than one run. Steps that are
-        not capturable run eagerly at their place in the order.
-
-        Args:
-            stage: Stage identifier in the form ``"GroupName_function_name"``.
-            steps: Ordered steps of the stage. Pass the same tuple on every call.
-            *args: Positional arguments forwarded to every step.
-            output: Transform applied to the stage result on every call, outside the graphs.
-            timer: Whether to time the stage.
-            **kwargs: Keyword arguments forwarded to every step.
-
-        Returns:
-            The result of the last step, optionally transformed by :paramref:`output`. A replayed run
-            returns the result of the call that recorded it.
-        """
-        return self._run(stage, self._split_runs(steps), args, kwargs, output, timer)
-
-    def _run(
-        self,
-        stage: str,
-        runs: tuple[tuple[bool, tuple[Callable[..., Any], ...]], ...],
-        args: tuple,
-        kwargs: dict[str, Any],
-        output: Callable[[Any], Any] | None,
-        timer: bool,
-    ) -> Any:
-        """Replay, record, or eagerly run each run of a stage in order."""
-        with Timer(name=stage, msg=f"{stage} took:", enable=timer, time_unit="us"):
-            capture = self._enabled and self.is_capturable(stage.partition("_")[0])
-            key = _argument_key(args, kwargs) if capture else None
-            result = None
-            for index, (capturable, fns) in enumerate(runs):
-                name = stage if len(runs) == 1 else f"{stage}[{index}]"
-                recorded = self._graphs.get(name) if capture and capturable else None
-                if recorded is not None and recorded[0] == key:
-                    wp.capture_launch(recorded[1])
-                    result = recorded[2]
-                elif capture and capturable and self._armed:
-                    graph, result = self._record(fns, args, kwargs)
-                    self._graphs[name] = (key, graph, result)
-                    wp.capture_launch(graph)
-                else:
-                    for fn in fns:
-                        result = self._run_eager(fn, args, kwargs)
-        return output(result) if output is not None else result
-
-    def arm(self) -> None:
-        """Allow capturable stages to record on their next call."""
-        self._armed = True
-
-    def invalidate(self, group: str | None = None) -> None:
-        """Drop recorded graphs.
-
-        Args:
-            group: Stage group whose graphs are dropped; they record again on their next armed call.
-                Defaults to None, which drops every graph and holds recording until the next :meth:`arm`.
-        """
-        stages = [stage for stage in self._graphs if group is None or stage.partition("_")[0] == group]
-        if stages:
-            # Graphs may still be executing; drain the device before releasing them.
-            wp.synchronize_device(self._device)
-            for stage in stages:
-                del self._graphs[stage]
-        if group is None:
-            self._runs.clear()
-            self._armed = False
-
-    def invalidate_on(self, physics_manager: type[PhysicsManager]) -> None:
-        """Invalidate every graph when the physics backend rebinds its buffers or stops.
-
-        Args:
-            physics_manager: Physics manager class that dispatches the lifecycle events.
-        """
-        for event in (PhysicsEvent.PHYSICS_READY, PhysicsEvent.STOP):
-            self._callback_handles.append(physics_manager.register_callback(self._on_physics_event, event))
-
-    def close(self) -> None:
-        """Deregister the physics callbacks and drop every graph."""
-        for handle in self._callback_handles:
-            handle.deregister()
-        self._callback_handles.clear()
-        self.invalidate()
-
-    def register_capturability(self, group: str, capturable: bool) -> None:
-        """Register whether a stage group is safe to record.
-
-        One unsafe registration keeps the group eager.
-
-        Args:
-            group: Stage group identifier.
-            capturable: Whether the registered work is safe during CUDA graph capture.
-        """
-        self._capturable[group] = self._capturable.get(group, True) and capturable
-
-    def is_capturable(self, group: str) -> bool:
-        """Return whether a stage group is eligible for recording."""
-        return self._capturable.get(group, True)
-
-    def _on_physics_event(self, payload: Any) -> None:
-        self.invalidate()
-
-    def _split_runs(self, steps: tuple[StageStep, ...]) -> tuple[tuple[bool, tuple[Callable[..., Any], ...]], ...]:
-        """Group consecutive steps with the same capturability, memoized per steps tuple."""
-        memo = self._runs.get(id(steps))
-        if memo is None or memo[0] is not steps:
-            runs: list[tuple[bool, list[Callable[..., Any]]]] = []
-            for capturable, fn in steps:
-                if runs and runs[-1][0] == capturable:
-                    runs[-1][1].append(fn)
-                else:
-                    runs.append((capturable, [fn]))
-            memo = (steps, tuple((capturable, tuple(fns)) for capturable, fns in runs))
-            self._runs[id(steps)] = memo
-        return memo[1]
-
-    def _record(self, fns: tuple[Callable[..., Any], ...], args: tuple, kwargs: dict[str, Any]) -> tuple[wp.Graph, Any]:
-        """Record one run of steps without executing it."""
-        result = None
-
-        def capture_target() -> None:
-            nonlocal result
-            for fn in fns:
-                result = fn(*args, **kwargs)
-
-        graph = NewtonQueries.capture_graph(str(self._device), capture_target, relaxed=has_kit())
-        return graph, result
-
-    def _run_eager(self, fn: Callable[..., Any], args: tuple, kwargs: dict[str, Any]) -> Any:
-        """Run a stage eagerly, under the synchronization trap when requested."""
-        if not self._sync_debug:
-            return fn(*args, **kwargs)
-        previous = torch.cuda.get_sync_debug_mode()
-        torch.cuda.set_sync_debug_mode("error")
-        try:
-            return fn(*args, **kwargs)
-        finally:
-            torch.cuda.set_sync_debug_mode(previous)
-
-
-def _argument_key(args: tuple, kwargs: dict[str, Any]) -> tuple:
-    """Key a stage call by the pointers of its arrays and the values of its scalars."""
-    return tuple(_value_key(value) for value in args) + tuple(
-        (name, _value_key(value)) for name, value in kwargs.items()
-    )
-
-
-def _value_key(value: Any) -> Any:
-    if isinstance(value, wp.array):
-        return value.ptr
-    if isinstance(value, torch.Tensor):
-        return value.data_ptr()
-    if value is None or isinstance(value, (bool, int, float, str)):
-        return value
-    return id(value)
diff --git a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py
index e9e07bca5e9f..a3ef24ec7ac8 100644
--- a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py
+++ b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py
@@ -488,7 +488,7 @@ def prev_action(self) -> torch.Tensor:
 
 
 class MockSceneEntityCfg:
-    """Unified cfg that works for both stable (joint_ids) and experimental (joint_mask / joint_ids_wp)."""
+    """Unified cfg that works for both stable (joint_ids) and experimental (joint_mask_wp / joint_ids_wp)."""
 
     def __init__(self, name: str, joint_ids: list[int], num_joints: int, device: str):
         self.name = name
@@ -498,7 +498,7 @@ def __init__(self, name: str, joint_ids: list[int], num_joints: int, device: str
         mask = [False] * num_joints
         for idx in joint_ids:
             mask[idx] = True
-        self.joint_mask = wp.array(mask, dtype=wp.bool, device=device)
+        self.joint_mask_wp = wp.array(mask, dtype=wp.bool, device=device)
         self.joint_ids_wp = wp.array(joint_ids, dtype=wp.int32, device=device)
 
 
diff --git a/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py b/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py
index f0afa42058e3..02c8e488ad52 100644
--- a/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py
+++ b/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py
@@ -6,7 +6,7 @@
 """CUDA-graph capture safety for warp MDP terms.
 
 The Warp environments record every capturable manager stage into a CUDA graph by default
-(:class:`~isaaclab_experimental.utils.WarpGraphCache`), so a warp MDP term that is not capture-safe
+(:class:`~isaaclab_experimental.utils.CapturedStage`), so a warp MDP term that is not capture-safe
 is not an edge case — it is the normal execution path. A captured graph records launch topology and
 pointer values and replays them, so a term must launch a fixed dim, touch only buffers that
 were allocated before capture and never move, and do all of its work in kernels: host-side
diff --git a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py
index 2ad43a657324..a9b2455039e9 100644
--- a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py
+++ b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py
@@ -21,7 +21,8 @@
 pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required")
 
 import isaaclab_experimental.envs.mdp.events as warp_evt
-from isaaclab_experimental.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg
+from isaaclab_experimental.managers import EventManager, EventTermCfg, ManagerTermBase, SceneEntityCfg
+from isaaclab_experimental.utils import CapturedStage
 from parity_helpers import (
     DEVICE,
     NUM_ACTIONS,
@@ -106,9 +107,16 @@ def all_joints_cfg():
 
 
 def _make_event_term(term_type: type[ManagerTermBase], env, mode: str = "reset", **params) -> ManagerTermBase:
-    """Build a class event term the way the event manager does."""
+    """Build a class event term the way the event manager does, after resolving its asset's Warp body fields."""
     params.setdefault("asset_cfg", SceneEntityCfg("robot"))
-    return term_type(EventTermCfg(func=term_type, mode=mode, params=params), env)
+    cfg = EventTermCfg(func=term_type, mode=mode, params=params)
+    # the mock scene cannot run SceneEntityCfg.resolve, so fill the fields it would set on the stored cfg
+    asset_cfg = cfg.params["asset_cfg"]
+    num_bodies = env.scene[asset_cfg.name].num_bodies
+    body_ids = list(range(num_bodies)) if asset_cfg.body_ids in (None, slice(None)) else list(asset_cfg.body_ids)
+    asset_cfg.body_ids_wp = wp.array(body_ids, dtype=wp.int32, device=DEVICE)
+    asset_cfg.body_mask_wp = wp.array([i in body_ids for i in range(num_bodies)], dtype=wp.bool, device=DEVICE)
+    return term_type(cfg, env)
 
 
 def _all_envs() -> wp.array:
@@ -359,9 +367,49 @@ def test_com_randomization_uses_the_first_call_baseline_without_accumulating():
     assert_close(data.body_com_pos_b.torch[..., 0], torch.full((NUM_ENVS, 2), 1.75, device=DEVICE))
 
 
-@pytest.mark.parametrize("term_type", [warp_evt.randomize_rigid_body_material, warp_evt.randomize_rigid_body_mass])
-def test_mask_adapters_reject_modes_other_than_startup(term_type):
-    """The adapters convert the mask to indices on the host, which a reset or interval stage must not do."""
-    cfg = EventTermCfg(func=term_type, mode="reset", params={"asset_cfg": SceneEntityCfg("robot")})
-    with pytest.raises(ValueError, match="startup"):
-        term_type(cfg, SimpleNamespace())
+@pytest.mark.parametrize(
+    ("term_type", "params"),
+    [
+        (
+            warp_evt.randomize_rigid_body_material,
+            {
+                "static_friction_range": (0.5, 0.5),
+                "dynamic_friction_range": (0.5, 0.5),
+                "restitution_range": (0.0, 0.0),
+                "num_buckets": 1,
+            },
+        ),
+        (warp_evt.randomize_rigid_body_mass, {"mass_distribution_params": (1.0, 1.0), "operation": "scale"}),
+    ],
+    ids=["material", "mass"],
+)
+def test_mask_adapters_run_eagerly_in_reset_mode(term_type, params, monkeypatch):
+    """The adapters convert the mask to indices on the host, so reset events run them outside the recorded graph."""
+    stable_term_type = term_type.__mro__[1]
+    calls = []
+
+    def stable_init(self, cfg, env):
+        self.cfg, self._env = cfg, env
+
+    monkeypatch.setattr(stable_term_type, "__init__", stable_init)
+    monkeypatch.setattr(stable_term_type, "__call__", lambda self, env, env_ids, *args: calls.append(env_ids.tolist()))
+    monkeypatch.setattr(CapturedStage, "enabled", True)
+    env, _, _ = _make_event_env(5)
+    env.sim = SimpleNamespace(is_playing=lambda: True)
+    manager = EventManager(
+        {
+            "randomize": EventTermCfg(
+                func=term_type, mode="reset", params={**params, "asset_cfg": SimpleNamespace(name="robot")}
+            )
+        },
+        env,
+    )
+    env_mask = wp.array([i % 2 == 0 for i in range(NUM_ENVS)], dtype=wp.bool, device=DEVICE)
+    selected = [i for i in range(NUM_ENVS) if i % 2 == 0]
+
+    for _ in range(2):
+        manager.apply(
+            mode="reset", env_mask_wp=env_mask, global_env_step_count=wp.zeros(1, dtype=wp.int32, device=DEVICE)
+        )
+
+    assert calls == [selected, selected], "the adapter must run, eagerly, on every reset"
diff --git a/source/isaaclab_experimental/test/envs/test_frontend.py b/source/isaaclab_experimental/test/envs/test_frontend.py
index ff71cdb49b0d..67cf269db3a8 100644
--- a/source/isaaclab_experimental/test/envs/test_frontend.py
+++ b/source/isaaclab_experimental/test/envs/test_frontend.py
@@ -206,7 +206,7 @@ def test_from_stable_copies_all_selection_fields():
         assert getattr(warp, field) == getattr(stable, field), f"field {field!r} mismatch"
     assert isinstance(warp, WarpSceneEntityCfg)
     # Warp-only fields stay None until :meth:`resolve` runs.
-    assert warp.joint_mask is None
+    assert warp.joint_mask_wp is None
     assert warp.joint_ids_wp is None
     assert warp.body_ids_wp is None
 
diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_capture.py b/source/isaaclab_experimental/test/envs/test_warp_env_capture.py
index 3887a4292fb8..9beeae991935 100644
--- a/source/isaaclab_experimental/test/envs/test_warp_env_capture.py
+++ b/source/isaaclab_experimental/test/envs/test_warp_env_capture.py
@@ -5,12 +5,12 @@
 
 """Eager and CUDA-graph execution of the Warp environments must produce the same rollout.
 
-Each task runs the same seeded action sequence twice, once with ``ISAACLAB_WARP_CAPTURE=0`` and
-once with stage capture on. For manager-based tasks, a hard :meth:`SimulationContext.reset` halfway
-through rebuilds the Newton model and reallocates every simulation buffer; recorded graphs still
-point at the freed buffers, so the captured rollout only stays equal when the environment records
-its stages again. A variant marks one term of every manager as not capturable: its stage must record
-the other terms and run that term eagerly in its place.
+Each task runs the same seeded action sequence twice, once with capture disabled
+(:attr:`CapturedStage.enabled` False) and once with stage capture on. For manager-based tasks, a hard
+:meth:`SimulationContext.reset` halfway through rebuilds the Newton model and reallocates every simulation
+buffer; recorded graphs still point at the freed buffers, so the captured rollout only stays equal when the
+stages record again. A variant marks one term of every manager as not capturable: its stage must record the
+other terms and run that term eagerly.
 """
 
 from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg
@@ -24,8 +24,8 @@
 import pytest
 import torch
 from isaaclab_experimental.envs.frontend import WarpFrontend
+from isaaclab_experimental.utils import CapturedStage
 from isaaclab_experimental.utils.warp import WarpCapturable
-from isaaclab_experimental.utils.warp_graph_cache import CAPTURE_ENV_VAR
 
 import isaaclab.sim as sim_utils
 
@@ -33,6 +33,7 @@
 from isaaclab_tasks.utils import resolve_task_config
 
 _NUM_ENVS = 32
+_MANAGERS = ("action_manager", "observation_manager", "event_manager", "reward_manager", "termination_manager")
 _STEPS = 80
 _REBIND_STEP = 40
 
@@ -45,11 +46,34 @@
 )
 
 
+def _recorded_stages(env) -> tuple[str, ...]:
+    """``Owner.method`` of every stage currently backed by a recorded graph, sorted."""
+    owners = [env, *(getattr(env, name) for name in _MANAGERS if hasattr(env, name))]
+    return tuple(
+        sorted(
+            f"{type(owner).__name__}.{method.__name__}"
+            for owner in owners
+            for method, stage in owner.__dict__.get("_captured_stages", {}).items()
+            if stage.num_graphs
+        )
+    )
+
+
+def _stale_stages(env) -> list[str]:
+    """Stages whose graphs were recorded before the latest physics rebind."""
+    owners = [env, *(getattr(env, name) for name in _MANAGERS if hasattr(env, name))]
+    return [
+        method.__name__
+        for owner in owners
+        for method, stage in owner.__dict__.get("_captured_stages", {}).items()
+        if stage.num_graphs and stage._generation != CapturedStage.generation
+    ]
+
+
 def _rollout(
     task_id: str, capture: bool, rebind: bool, monkeypatch: pytest.MonkeyPatch, eager_terms: tuple[str, ...] = ()
 ) -> dict:
     """Run the seeded action sequence and return the rollout with the recorded stages."""
-    monkeypatch.setenv(CAPTURE_ENV_VAR, "1" if capture else "0")
     env_cfg, _ = resolve_task_config(task_id, "", overrides=("physics=newton_mjwarp",))
     env_cfg.seed = 7
     env_cfg.scene.num_envs = _NUM_ENVS
@@ -65,6 +89,8 @@ def _rollout(
         term_cfg.func = WarpCapturable(False, reason="marked eager by the test")(term_cfg.func)
     sim_utils.create_new_stage()
     env = WarpFrontend.build_env(env_cfg, task_id).unwrapped
+    # the environment enables capture when it is built
+    monkeypatch.setattr(CapturedStage, "enabled", capture)
     rollout = {"obs": [], "reward": [], "terminated": [], "truncated": []}
     try:
         env.reset()
@@ -73,9 +99,10 @@ def _rollout(
         bias = torch.rand((_NUM_ENVS, env.action_space.shape[-1]), generator=generator) * 2.0 - 1.0
         for step in range(_STEPS):
             if rebind and step == _REBIND_STEP:
-                rollout["stages_before_rebind"] = env._warp_graph_cache.captured_stages
+                rollout["stages_before_rebind"] = _recorded_stages(env)
+                generation = CapturedStage.generation
                 env.sim.reset()
-                rollout["stages_at_rebind"] = env._warp_graph_cache.captured_stages
+                rollout["rebind_advanced_generation"] = CapturedStage.generation > generation
                 env.reset()
             noise = torch.rand((_NUM_ENVS, bias.shape[1]), generator=generator) * 2.0 - 1.0
             obs, reward, terminated, truncated, _ = env.step((bias + 0.5 * noise).to(env.device))
@@ -83,7 +110,8 @@ def _rollout(
             rollout["reward"].append(reward.clone())
             rollout["terminated"].append(terminated.clone())
             rollout["truncated"].append(truncated.clone())
-        rollout["stages_at_end"] = env._warp_graph_cache.captured_stages
+        rollout["stages_at_end"] = _recorded_stages(env)
+        rollout["stale_stages"] = _stale_stages(env)
     finally:
         env.close()
     for key in ("obs", "reward", "terminated", "truncated"):
@@ -111,10 +139,11 @@ def test_captured_rollout_matches_eager(
     assert captured["stages_at_end"], "no stage was recorded; the comparison proves nothing"
     if eager_terms:
         # every stage holding an eager term still records the terms around it
-        recorded_groups = {stage.partition("_")[0] for stage in captured["stages_at_end"]}
+        recorded_groups = {stage.partition(".")[0] for stage in captured["stages_at_end"]}
         assert {"ObservationManager", "RewardManager", "TerminationManager", "EventManager"} <= recorded_groups
     if rebind:
-        assert captured["stages_at_rebind"] == (), "the rebind must drop graphs that read freed buffers"
+        assert captured["rebind_advanced_generation"], "the rebind must retire graphs that read freed buffers"
+        assert captured["stale_stages"] == [], "every stage must record again after the rebind"
         assert captured["stages_at_end"] == captured["stages_before_rebind"]
     for window in (slice(0, _REBIND_STEP), slice(_REBIND_STEP, _STEPS)):
         dones = captured["terminated"][window] | captured["truncated"][window]
diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_decimation.py b/source/isaaclab_experimental/test/envs/test_warp_env_decimation.py
index 9ef6e9a3e8ed..8ce9fc741750 100644
--- a/source/isaaclab_experimental/test/envs/test_warp_env_decimation.py
+++ b/source/isaaclab_experimental/test/envs/test_warp_env_decimation.py
@@ -82,24 +82,17 @@ def count_sim_step(*args, **kwargs):
         return sim_step(*args, **kwargs)
 
     env.sim.step = count_sim_step
-    if frontend == "warp":
-        call_steps = env._warp_graph_cache.call_steps
+    apply_action = env.action_manager.apply_action
 
-        def count_apply_action(stage, *args, **kwargs):
-            calls["apply_action"] += stage == "ActionManager_apply_action"
-            return call_steps(stage, *args, **kwargs)
+    def count_apply_action():
+        calls["apply_action"] += 1
+        apply_action()
 
-        env._warp_graph_cache.call_steps = count_apply_action
+    env.action_manager.apply_action = count_apply_action
+    if frontend == "warp":
         step_reward = env.reward_manager._step_reward_tensor_view
         term_dones = env.termination_manager._term_dones_tensor_view
     else:
-        apply_action = env.action_manager.apply_action
-
-        def count_apply_action():
-            calls["apply_action"] += 1
-            apply_action()
-
-        env.action_manager.apply_action = count_apply_action
         step_reward = env.reward_manager._step_reward
         term_dones = env.termination_manager._term_dones
     reward_names = env.reward_manager._term_names
diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_reset.py b/source/isaaclab_experimental/test/envs/test_warp_env_reset.py
index c75dea09e091..bf0ee211d098 100644
--- a/source/isaaclab_experimental/test/envs/test_warp_env_reset.py
+++ b/source/isaaclab_experimental/test/envs/test_warp_env_reset.py
@@ -14,6 +14,7 @@
 
 import torch
 from isaaclab_experimental.envs.frontend import WarpFrontend
+from isaaclab_experimental.managers import EventManager
 
 import isaaclab.sim as sim_utils
 
@@ -62,7 +63,7 @@ def test_set_term_cfg_applies_to_recorded_reset_events():
         selected = torch.arange(_NUM_ENVS, device=env.device) % 2 == 0
         # records the reset events with the configured random cart offsets
         env.reset(env_mask=selected)
-        assert "EventManager_apply_reset" in env._warp_graph_cache.captured_stages
+        assert env.event_manager._captured_stages[EventManager._apply_reset].num_graphs
 
         term_cfg = env.event_manager.get_term_cfg("reset_cart_position")
         term_cfg.params["position_range"] = (0.5, 0.5)
diff --git a/source/isaaclab_experimental/test/managers/test_event_manager.py b/source/isaaclab_experimental/test/managers/test_event_manager.py
index ae052362cdb7..c1ee9b0b13df 100644
--- a/source/isaaclab_experimental/test/managers/test_event_manager.py
+++ b/source/isaaclab_experimental/test/managers/test_event_manager.py
@@ -5,12 +5,13 @@
 
 """Tests for the Warp event manager."""
 
+from collections import Counter
 from types import SimpleNamespace
 
 import pytest
 import warp as wp
-from isaaclab_experimental.managers import EventManager, EventTermCfg
-from isaaclab_experimental.utils import WarpGraphCache
+from isaaclab_experimental.managers import EventManager, EventTermCfg, ManagerTermBase
+from isaaclab_experimental.utils import CapturedStage
 from isaaclab_experimental.utils.warp import WarpCapturable
 
 wp.init()
@@ -20,6 +21,9 @@
 NUM_ENVS = 4
 
 
+CALLS: Counter[str] = Counter()
+
+
 @WarpCapturable(False, reason="synchronizes with the host")
 def host_synchronizing_term(env, env_mask: wp.array):
     """Event term that cannot run inside a recorded stage."""
@@ -27,14 +31,35 @@ def host_synchronizing_term(env, env_mask: wp.array):
 
 def device_term(env, env_mask: wp.array):
     """Event term whose work is entirely on the device."""
+    CALLS["device_term"] += 1
+
+
+@WarpCapturable(False, reason="synchronizes with the host")
+class host_call_term(ManagerTermBase):
+    """Class event term whose call cannot be recorded but whose reset can."""
+
+    def __call__(self, env, env_mask: wp.array):
+        pass
+
+    def reset(self, env_mask: wp.array | None = None):
+        CALLS["host_call_term.reset"] += 1
+
+
+class host_reset_term(ManagerTermBase):
+    """Class event term whose reset cannot be recorded."""
+
+    def __call__(self, env, env_mask: wp.array):
+        pass
+
+    @WarpCapturable(False, reason="synchronizes with the host")
+    def reset(self, env_mask: wp.array | None = None):
+        CALLS["host_reset_term.reset"] += 1
 
 
-def test_startup_only_term_does_not_keep_the_reset_stage_eager():
+def test_startup_only_term_does_not_keep_the_reset_stage_eager(monkeypatch):
     """Startup terms run once while the environment is built, so they must not stop reset events from recording."""
-    graph_cache = WarpGraphCache(DEVICE)
-    env = SimpleNamespace(
-        num_envs=NUM_ENVS, device=DEVICE, sim=SimpleNamespace(is_playing=lambda: True), _warp_graph_cache=graph_cache
-    )
+    monkeypatch.setattr(CapturedStage, "enabled", True)
+    env = SimpleNamespace(num_envs=NUM_ENVS, device=DEVICE, sim=SimpleNamespace(is_playing=lambda: True))
     manager = EventManager(
         {
             "startup_term": EventTermCfg(func=host_synchronizing_term, mode="startup"),
@@ -42,15 +67,33 @@ def test_startup_only_term_does_not_keep_the_reset_stage_eager():
         },
         env,
     )
-    graph_cache.arm()
-
-    graph_cache.call(
-        "EventManager_apply_reset",
-        manager.apply,
-        mode="reset",
-        env_mask_wp=wp.ones(NUM_ENVS, dtype=wp.bool, device=DEVICE),
-        global_env_step_count=wp.zeros(1, dtype=wp.int32, device=DEVICE),
+    env_mask = wp.ones(NUM_ENVS, dtype=wp.bool, device=DEVICE)
+    step_count = wp.zeros(1, dtype=wp.int32, device=DEVICE)
+    CALLS.clear()
+
+    for _ in range(2):
+        manager.apply(mode="reset", env_mask_wp=env_mask, global_env_step_count=step_count)
+
+    assert CALLS["device_term"] == 1, "the second call must replay the recorded reset events"
+    assert manager._captured_stages[EventManager._apply_reset].num_graphs == 1
+
+
+def test_class_annotation_does_not_keep_the_reset_eager(monkeypatch):
+    """A class annotation covers the term's call; only an annotation on ``reset`` keeps the reset out of the graph."""
+    monkeypatch.setattr(CapturedStage, "enabled", True)
+    env = SimpleNamespace(num_envs=NUM_ENVS, device=DEVICE, sim=SimpleNamespace(is_playing=lambda: True))
+    manager = EventManager(
+        {
+            "host_call_term": EventTermCfg(func=host_call_term, mode="startup"),
+            "host_reset_term": EventTermCfg(func=host_reset_term, mode="startup"),
+        },
+        env,
     )
+    env_mask = wp.ones(NUM_ENVS, dtype=wp.bool, device=DEVICE)
+    CALLS.clear()
+
+    for _ in range(2):
+        manager.reset(env_mask=env_mask)
 
-    assert graph_cache.captured_stages == ("EventManager_apply_reset",)
-    graph_cache.close()
+    assert CALLS["host_call_term.reset"] == 1, "the second reset must replay the recorded reset"
+    assert CALLS["host_reset_term.reset"] == 2, "an annotated reset must run on every reset"
diff --git a/source/isaaclab_experimental/test/managers/test_reward_manager.py b/source/isaaclab_experimental/test/managers/test_reward_manager.py
index 5af3755aeee5..3901abb2c702 100644
--- a/source/isaaclab_experimental/test/managers/test_reward_manager.py
+++ b/source/isaaclab_experimental/test/managers/test_reward_manager.py
@@ -12,7 +12,7 @@
 import torch
 import warp as wp
 from isaaclab_experimental.managers import ManagerTermBase, RewardManager
-from isaaclab_experimental.utils import WarpGraphCache
+from isaaclab_experimental.utils import CapturedStage
 from isaaclab_experimental.utils.warp import WarpCapturable
 
 from isaaclab.managers import RewardTermCfg
@@ -32,48 +32,40 @@ def constant_reward(env, out: wp.array, value: float, sensor_cfg: str | None = N
     out.fill_(value)
 
 
+@WarpCapturable(False, reason="synchronizes with the host")
 class ConstantReward(ManagerTermBase):
-    """Class term that decides its capturability per instance."""
+    """Class term that is never capturable."""
 
-    def __init__(self, cfg: RewardTermCfg, env):
-        super().__init__(cfg, env)
-        self._warp_capturable = cfg.params["capturable"]
-
-    def __call__(self, env, out: wp.array, value: float, capturable: bool):
+    def __call__(self, env, out: wp.array, value: float):
         CALLS[value] += 1
         out.fill_(value)
 
 
-def test_capturability_is_decided_per_term_instance():
-    """The same term is recorded for one parameter set and runs eagerly for another.
+def test_capturability_is_decided_per_term_from_its_parameters(monkeypatch):
+    """The same function term is recorded for one parameter set and runs eagerly for another.
 
-    An eager instance runs its Python on every step; a recorded one only while its graph is recorded.
+    An eager term runs its Python on every step; a recorded one only while its graph is recorded.
     The guard of the predicate term raises if the instance with a sensor is ever recorded.
     """
-    graph_cache = WarpGraphCache(DEVICE)
+    monkeypatch.setattr(CapturedStage, "enabled", True)
     env = SimpleNamespace(
-        num_envs=NUM_ENVS,
-        device=DEVICE,
-        sim=SimpleNamespace(is_playing=lambda: True),
-        max_episode_length_s=1.0,
-        _warp_graph_cache=graph_cache,
+        num_envs=NUM_ENVS, device=DEVICE, sim=SimpleNamespace(is_playing=lambda: True), max_episode_length_s=1.0
     )
     cfg = {
         "function_recorded": RewardTermCfg(func=constant_reward, weight=1.0, params={"value": 1.0}),
         "function_eager": RewardTermCfg(
             func=constant_reward, weight=1.0, params={"value": 2.0, "sensor_cfg": "height_scanner"}
         ),
-        "class_recorded": RewardTermCfg(func=ConstantReward, weight=1.0, params={"value": 4.0, "capturable": True}),
-        "class_eager": RewardTermCfg(func=ConstantReward, weight=1.0, params={"value": 8.0, "capturable": False}),
+        "class_eager": RewardTermCfg(func=ConstantReward, weight=1.0, params={"value": 4.0}),
     }
     manager = RewardManager(cfg, env)
-    graph_cache.arm()
     CALLS.clear()
 
     for _ in range(3):
-        reward = graph_cache.call_steps("RewardManager_compute", manager.stage_steps("compute"), dt=0.5)
+        reward = manager.compute(dt=0.5)
     wp.synchronize()
 
-    assert CALLS == {1.0: 1, 2.0: 3, 4.0: 1, 8.0: 3}
-    assert torch.equal(reward, torch.full((NUM_ENVS,), (1.0 + 2.0 + 4.0 + 8.0) * 0.5, device=DEVICE))
-    graph_cache.close()
+    assert CALLS == {1.0: 1, 2.0: 3, 4.0: 3}
+    # each eager term splits the recording, so it holds one graph before each of the two eager terms and one after
+    assert manager._captured_stages[RewardManager.compute].num_graphs == 3
+    assert torch.equal(reward, torch.full((NUM_ENVS,), (1.0 + 2.0 + 4.0) * 0.5, device=DEVICE))
diff --git a/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py
index 66dd929894d5..62e809ed2bf1 100644
--- a/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py
+++ b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py
@@ -64,17 +64,16 @@
 
 # Sanctioned host-synchronization boundaries, by (path suffix, function name).
 SANCTIONED_BOUNDARIES = {
-    # one host predicate per step decides whether the reset pipeline runs (W7)
-    ("envs/manager_based_rl_env_warp.py", "_reset_terminated_envs"),
-    ("envs/direct_rl_env_warp.py", "step"),
+    # one host predicate per step decides whether the reset pipeline runs (W7); the direct env decides on the device
+    ("envs/manager_based_rl_env_warp.py", "step"),
     # mask to index compaction for the managers that reset by index (W11)
-    ("envs/manager_based_env_warp.py", "_reset_env_ids"),
-    # startup-only stable randomization terms take environment indices
+    ("envs/manager_based_env_warp.py", "reset"),
+    ("envs/manager_based_env_warp.py", "_reset_mask"),
+    ("envs/manager_based_rl_env_warp.py", "_reset_mask"),
+    # the stable randomization adapters take environment indices; they run outside recorded graphs
     ("envs/mdp/events.py", "_mask_to_env_ids"),
     # the camera's empty-reset predicate
     ("isaaclab/sensors/camera/camera.py", "_env_mask_has_any"),
-    # backends without a mask-native actuator reset fall back to indices; Newton and OVPhysX override it
-    ("isaaclab/actuators/actuator_control.py", "reset_native_actuators_mask"),
     # joint-limit writes mutate the solver model (event-driven, not per-step)
     ("isaaclab_newton/assets/articulation/articulation.py", "write_joint_position_limit_to_sim_mask"),
 }
diff --git a/source/isaaclab_experimental/test/utils/test_warp_capture.py b/source/isaaclab_experimental/test/utils/test_warp_capture.py
new file mode 100644
index 000000000000..0c8c3d49c8c9
--- /dev/null
+++ b/source/isaaclab_experimental/test/utils/test_warp_capture.py
@@ -0,0 +1,241 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tests for the capture of Warp frontend stages."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import warp as wp
+from isaaclab_experimental.utils.warp_capture import (
+    SYNC_DEBUG_ENV_VAR,
+    CapturedStage,
+    captured,
+    eager,
+    reset_captured_stages,
+)
+
+wp.init()
+pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required")
+
+DEVICE = "cuda:0"
+
+
+@wp.kernel
+def _affine(values: wp.array(dtype=wp.int32), scale: wp.int32, shift: wp.int32):
+    values[wp.tid()] = values[wp.tid()] * scale + shift
+
+
+class _AffineStage:
+    """Stage that applies ``x * scale + shift`` on the device and counts its Python invocations."""
+
+    def __init__(self, scale: int = 1, shift: int = 1):
+        self.scale, self.shift, self.calls = scale, shift, 0
+
+    def __call__(self, values: wp.array, *args) -> wp.array:
+        self.calls += 1
+        wp.launch(_affine, dim=values.shape[0], inputs=[values, self.scale, self.shift], device=values.device)
+        return values
+
+
+def _read(values: wp.array) -> list[int]:
+    wp.synchronize_device(values.device)
+    return values.numpy().tolist()
+
+
+@pytest.fixture(autouse=True)
+def enabled(monkeypatch):
+    monkeypatch.delenv(SYNC_DEBUG_ENV_VAR, raising=False)
+    monkeypatch.setattr(CapturedStage, "enabled", True)
+
+
+def test_first_call_records_and_runs_the_stage_once():
+    """Recording has no eager warm-up: a stateful stage advances once per call, including the first."""
+    fn = _AffineStage()
+    stage = CapturedStage(fn, DEVICE)
+    values = wp.zeros(2, dtype=wp.int32, device=DEVICE)
+
+    assert stage(values) is values
+    assert _read(values) == [1, 1]
+    assert stage(values) is values
+    assert _read(values) == [2, 2]
+    assert (fn.calls, stage.num_graphs) == (1, 1), "the second call must replay the recorded graph"
+
+
+def test_changed_argument_records_the_stage_again():
+    """A graph replays its recorded pointers and scalars, so a new array or scalar records a new graph."""
+    fn = _AffineStage()
+    stage = CapturedStage(fn, DEVICE)
+    old, new = wp.zeros(1, dtype=wp.int32, device=DEVICE), wp.zeros(1, dtype=wp.int32, device=DEVICE)
+
+    stage(old)
+    stage(new)
+    stage(new)
+    stage(new, 5)
+    assert (_read(old), _read(new), fn.calls, stage.num_graphs) == ([1], [3], 3, 3)
+
+
+def test_physics_rebind_records_the_stage_again():
+    fn = _AffineStage()
+    stage = CapturedStage(fn, DEVICE)
+    values = wp.zeros(1, dtype=wp.int32, device=DEVICE)
+    stage(values)
+
+    CapturedStage.invalidate()
+    stage(values)
+    stage(values)
+    assert (fn.calls, _read(values), stage.num_graphs) == (2, [3], 1)
+
+
+def test_stage_runs_eagerly_when_disabled_or_on_cpu(monkeypatch):
+    on_cpu = _AffineStage()
+    cpu_stage = CapturedStage(on_cpu, "cpu")
+    cpu_values = wp.zeros(1, dtype=wp.int32, device="cpu")
+    cpu_stage(cpu_values)
+    cpu_stage(cpu_values)
+
+    monkeypatch.setattr(CapturedStage, "enabled", False)
+    disabled = _AffineStage()
+    stage = CapturedStage(disabled, DEVICE)
+    values = wp.zeros(1, dtype=wp.int32, device=DEVICE)
+    stage(values)
+    stage(values)
+
+    assert (on_cpu.calls, _read(cpu_values), cpu_stage.num_graphs) == (2, [2], 0)
+    assert (disabled.calls, _read(values), stage.num_graphs) == (2, [2], 0)
+
+
+def test_eager_call_runs_between_its_graphs_on_every_call():
+    """The recording runs the graph before an eager call first, and every replay keeps that order."""
+    before, between, after = _AffineStage(2, 0), _AffineStage(1, 10), _AffineStage(3, 0)
+
+    def body(values: wp.array) -> wp.array:
+        before(values)
+        eager(between, values)
+        return after(values)
+
+    stage = CapturedStage(body, DEVICE)
+    values = wp.ones(1, dtype=wp.int32, device=DEVICE)
+
+    stage(values)
+    # the non-commuting steps only give these values in order: 1 * 2 + 10 = 12 -> 36, then 36 * 2 + 10 = 82 -> 246
+    assert _read(values) == [36]
+    stage(values)
+    assert (_read(values), between.calls, before.calls, after.calls, stage.num_graphs) == ([246], 2, 1, 1, 2)
+
+
+def test_eager_call_first_leaves_one_graph_after_it():
+    """A stage making its eager call first records all its capturable work into the single graph after it."""
+    eager_fn, first, second = _AffineStage(1, 10), _AffineStage(2, 0), _AffineStage(3, 0)
+
+    def body(values: wp.array) -> wp.array:
+        eager(eager_fn, values)
+        first(values)
+        return second(values)
+
+    stage = CapturedStage(body, DEVICE)
+    values = wp.ones(1, dtype=wp.int32, device=DEVICE)
+    stage(values)
+    stage(values)
+
+    assert _read(values) == [((1 + 10) * 6 + 10) * 6]
+    ((steps, _),) = stage._recordings.values()
+    assert [isinstance(step, wp.Graph) for step in steps] == [True, False, True]
+    # the graph before the eager call is empty; the last one multiplies by 6
+    wp.capture_launch(steps[0])
+    assert _read(values) == [456]
+    wp.capture_launch(steps[-1])
+    assert _read(values) == [456 * 6]
+
+
+class _Owner:
+    device = DEVICE
+
+    def __init__(self):
+        self.inner_fn, self.outer_fn = _AffineStage(2, 0), _AffineStage(1, 1)
+
+    @captured
+    def inner(self, values: wp.array) -> wp.array:
+        return self.inner_fn(values)
+
+    @captured
+    def outer(self, values: wp.array) -> wp.array:
+        self.outer_fn(values)
+        return self.inner(values)
+
+
+def test_decorated_methods_own_their_stages_and_nest_into_an_enclosing_recording():
+    owner = _Owner()
+    values = wp.zeros(1, dtype=wp.int32, device=DEVICE)
+
+    assert owner.outer(values) is values
+    owner.outer(values)
+    # (0 + 1) * 2 = 2 -> (2 + 1) * 2 = 6: the inner call is part of the outer graph
+    assert _read(values) == [6]
+    assert (owner.outer_fn.calls, owner.inner_fn.calls) == (1, 1)
+    stages = owner._captured_stages
+    assert (stages[_Owner.outer].num_graphs, stages[_Owner.inner].num_graphs) == (1, 0)
+
+    reset_captured_stages(owner)
+    owner.outer(values)
+    assert (owner.outer_fn.calls, _read(values)) == (2, [14])
+
+
+def test_eager_call_of_a_nested_stage_splits_the_enclosing_recording():
+    eager_fn, inner_fn, outer_fn = _AffineStage(1, 10), _AffineStage(2, 0), _AffineStage(3, 0)
+    nested = CapturedStage(lambda values: (eager(eager_fn, values), inner_fn(values)), DEVICE)
+    outer = CapturedStage(lambda values: (outer_fn(values), nested(values)), DEVICE)
+    values = wp.ones(1, dtype=wp.int32, device=DEVICE)
+
+    outer(values)
+    outer(values)
+    # 1 * 3 + 10 = 13 -> 26, then 26 * 3 + 10 = 88 -> 176: the nested eager call runs between the outer graphs
+    assert (_read(values), eager_fn.calls, inner_fn.calls, outer_fn.calls) == ([176], 2, 1, 1)
+    assert (outer.num_graphs, nested.num_graphs) == (2, 0)
+
+
+def test_eager_call_raises_inside_a_capture_no_stage_started():
+    """A call that cannot be recorded must not be captured silently into a foreign graph."""
+    fn = _AffineStage()
+    values = wp.zeros(1, dtype=wp.int32, device=DEVICE)
+
+    with pytest.raises(RuntimeError, match="no CapturedStage started"):
+        with wp.ScopedDevice(DEVICE), wp.ScopedCapture():
+            eager(fn, values)
+    assert fn.calls == 0
+
+
+def test_stage_raising_while_recording_ends_the_capture():
+    """A failed recording must not leave the device capturing; the next call records again."""
+    recorded, eager_fn = _AffineStage(2, 0), _AffineStage(1, 10)
+
+    def body(values: wp.array) -> wp.array:
+        recorded(values)
+        eager(eager_fn, values)
+        if recorded.calls == 1:
+            raise ValueError("first recording fails")
+        return values
+
+    stage = CapturedStage(body, DEVICE)
+    values = wp.ones(1, dtype=wp.int32, device=DEVICE)
+
+    with pytest.raises(ValueError, match="first recording fails"):
+        stage(values)
+    assert not wp.get_device(DEVICE).is_capturing
+    stage(values)
+    # the failed recording ran its first graph and the eager call: 1 * 2 + 10 = 12, then 12 * 2 + 10 = 34
+    assert (_read(values), stage.num_graphs) == ([34], 2)
+
+
+def test_sync_debug_trap_raises_on_a_hidden_host_sync(monkeypatch):
+    monkeypatch.setattr(CapturedStage, "enabled", False)
+    monkeypatch.setenv(SYNC_DEBUG_ENV_VAR, "1")
+    flags = torch.ones(4, dtype=torch.bool, device=DEVICE)
+    stage = CapturedStage(lambda: flags.any().item(), DEVICE)
+
+    with pytest.raises(RuntimeError, match="synchroniz"):
+        stage()
+    assert torch.cuda.get_sync_debug_mode() == 0, "the trap must be released after the stage"
diff --git a/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py b/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py
deleted file mode 100644
index 794106aa0933..000000000000
--- a/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py
+++ /dev/null
@@ -1,212 +0,0 @@
-# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
-# All rights reserved.
-#
-# SPDX-License-Identifier: BSD-3-Clause
-
-"""Tests for the WarpGraphCache capture lifecycle."""
-
-from __future__ import annotations
-
-import pytest
-import torch
-import warp as wp
-from isaaclab_experimental.utils.warp_graph_cache import CAPTURE_ENV_VAR, SYNC_DEBUG_ENV_VAR, WarpGraphCache
-
-wp.init()
-pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required")
-
-DEVICE = "cuda:0"
-
-
-@wp.kernel
-def _increment(values: wp.array(dtype=wp.int32)):
-    i = wp.tid()
-    values[i] = values[i] + 1
-
-
-class _CountingStage:
-    """Stage that increments a device counter and counts its Python invocations."""
-
-    def __init__(self):
-        self.calls = 0
-
-    def __call__(self, values: wp.array) -> wp.array:
-        self.calls += 1
-        wp.launch(_increment, dim=values.shape[0], inputs=[values], device=DEVICE)
-        return values
-
-
-def _read(values: wp.array) -> list[int]:
-    wp.synchronize_device(DEVICE)
-    return values.numpy().tolist()
-
-
-@pytest.fixture
-def cache(monkeypatch):
-    monkeypatch.delenv(CAPTURE_ENV_VAR, raising=False)
-    monkeypatch.delenv(SYNC_DEBUG_ENV_VAR, raising=False)
-    cache = WarpGraphCache(DEVICE)
-    yield cache
-    cache.close()
-
-
-def test_first_armed_call_records_and_runs_the_stage_once(cache):
-    """Recording has no eager warm-up: a stateful stage advances once per call, including the first."""
-    stage = _CountingStage()
-    values = wp.zeros(2, dtype=wp.int32, device=DEVICE)
-    cache.arm()
-
-    result = cache.call("Group_stage", stage, values)
-    assert _read(values) == [1, 1]
-    assert result is values
-
-    cache.call("Group_stage", stage, values)
-    assert _read(values) == [2, 2]
-    assert stage.calls == 1, "the second call must replay the recorded graph"
-    assert cache.captured_stages == ("Group_stage",)
-
-
-def test_stages_run_eagerly_until_armed(cache):
-    stage = _CountingStage()
-    values = wp.zeros(1, dtype=wp.int32, device=DEVICE)
-
-    cache.call("Group_stage", stage, values)
-    cache.call("Group_stage", stage, values)
-    assert (stage.calls, _read(values), cache.captured_stages) == (2, [2], ())
-
-    cache.arm()
-    cache.call("Group_stage", stage, values)
-    cache.call("Group_stage", stage, values)
-    assert (stage.calls, _read(values), cache.captured_stages) == (3, [4], ("Group_stage",))
-
-
-def test_reallocated_argument_records_the_stage_again(cache):
-    """A graph replays recorded pointers, so a new array must not reuse the old graph."""
-    stage = _CountingStage()
-    old = wp.zeros(1, dtype=wp.int32, device=DEVICE)
-    cache.arm()
-    cache.call("Group_stage", stage, old)
-
-    new = wp.zeros(1, dtype=wp.int32, device=DEVICE)
-    cache.call("Group_stage", stage, new)
-    cache.call("Group_stage", stage, new)
-
-    assert (_read(old), _read(new), stage.calls) == ([1], [2], 2)
-
-
-def test_changed_scalar_argument_records_the_stage_again(cache):
-    values = wp.zeros(1, dtype=wp.float32, device=DEVICE)
-
-    def add(values: wp.array, amount: float) -> None:
-        wp.launch(_add_scalar, dim=1, inputs=[values, amount], device=DEVICE)
-
-    cache.arm()
-    cache.call("Group_stage", add, values, 1.0)
-    cache.call("Group_stage", add, values, 10.0)
-    assert _read(values) == [11.0]
-
-
-@wp.kernel
-def _add_scalar(values: wp.array(dtype=wp.float32), amount: wp.float32):
-    values[wp.tid()] = values[wp.tid()] + amount
-
-
-def test_invalidate_drops_graphs_and_waits_for_the_next_arm(cache):
-    stage = _CountingStage()
-    values = wp.zeros(1, dtype=wp.int32, device=DEVICE)
-    cache.arm()
-    cache.call("Group_stage", stage, values)
-
-    cache.invalidate()
-    cache.call("Group_stage", stage, values)
-    assert (stage.calls, cache.captured_stages) == (2, ())
-
-    cache.arm()
-    cache.call("Group_stage", stage, values)
-    cache.call("Group_stage", stage, values)
-    assert (stage.calls, _read(values), cache.captured_stages) == (3, [4], ("Group_stage",))
-
-
-def test_group_invalidation_keeps_other_groups_recorded(cache):
-    values = wp.zeros(1, dtype=wp.int32, device=DEVICE)
-    cache.arm()
-    cache.call("First_stage", _CountingStage(), values)
-    cache.call("Second_stage", _CountingStage(), values)
-
-    cache.invalidate("First")
-
-    assert cache.captured_stages == ("Second_stage",)
-    stage = _CountingStage()
-    cache.call("First_stage", stage, values)
-    assert (stage.calls, cache.captured_stages) == (1, ("First_stage", "Second_stage"))
-
-
-def test_non_capturable_group_stays_eager(cache):
-    stage = _CountingStage()
-    values = wp.zeros(1, dtype=wp.int32, device=DEVICE)
-    cache.register_capturability("Group", True)
-    cache.register_capturability("Group", False)
-    cache.register_capturability("Group", True)
-    cache.arm()
-
-    cache.call("Group_stage", stage, values)
-    cache.call("Group_stage", stage, values)
-
-    assert (stage.calls, _read(values), cache.captured_stages) == (2, [2], ())
-
-
-def test_capture_env_var_forces_eager_execution(monkeypatch):
-    monkeypatch.setenv(CAPTURE_ENV_VAR, "0")
-    cache = WarpGraphCache(DEVICE)
-    stage = _CountingStage()
-    values = wp.zeros(1, dtype=wp.int32, device=DEVICE)
-    cache.arm()
-
-    cache.call("Group_stage", stage, values)
-    cache.call("Group_stage", stage, values)
-
-    assert (stage.calls, _read(values), cache.captured_stages) == (2, [2], ())
-
-
-def test_sync_debug_trap_raises_on_a_hidden_host_sync(monkeypatch):
-    monkeypatch.setenv(CAPTURE_ENV_VAR, "0")
-    monkeypatch.setenv(SYNC_DEBUG_ENV_VAR, "1")
-    cache = WarpGraphCache(DEVICE)
-    flags = torch.ones(4, dtype=torch.bool, device=DEVICE)
-
-    with pytest.raises(RuntimeError, match="synchroniz"):
-        cache.call("Group_stage", lambda: flags.any().item())
-    assert torch.cuda.get_sync_debug_mode() == 0, "the trap must be released after the stage"
-
-
-@wp.kernel
-def _affine(values: wp.array(dtype=wp.int32), scale: wp.int32, shift: wp.int32):
-    values[wp.tid()] = values[wp.tid()] * scale + shift
-
-
-class _AffineStep:
-    """Stage step that applies ``x * scale + shift`` on the device and counts its Python invocations."""
-
-    def __init__(self, scale: int, shift: int):
-        self.scale, self.shift, self.calls = scale, shift, 0
-
-    def __call__(self, values: wp.array) -> wp.array:
-        self.calls += 1
-        wp.launch(_affine, dim=values.shape[0], inputs=[values, self.scale, self.shift], device=DEVICE)
-        return values
-
-
-def test_steps_record_capturable_runs_and_run_the_rest_eagerly_in_order(cache):
-    """Each run of capturable steps becomes one graph; an eager step runs between them on every call."""
-    double, increment, triple = _AffineStep(2, 0), _AffineStep(1, 1), _AffineStep(3, 0)
-    steps = ((True, double), (False, increment), (True, triple))
-    values = wp.ones(1, dtype=wp.int32, device=DEVICE)
-    cache.arm()
-
-    cache.call_steps("Group_stage", steps, values)
-    cache.call_steps("Group_stage", steps, values)
-
-    # the non-commuting steps only give 57 in order: ((1 * 2 + 1) * 3 = 9) -> ((9 * 2 + 1) * 3 = 57)
-    assert _read(values) == [57]
-    assert (double.calls, increment.calls, triple.calls) == (1, 2, 1)
-    assert cache.captured_stages == ("Group_stage[0]", "Group_stage[2]")
diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst
new file mode 100644
index 000000000000..918eff1c77e5
--- /dev/null
+++ b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst
@@ -0,0 +1,4 @@
+* Changed the Cartpole, locomotion and reorient Warp direct environments to record their ``_apply_action``,
+  ``_get_dones``, ``_get_rewards``, ``_get_observations`` and ``_reset_idx`` methods as CUDA graphs with
+  :func:`~isaaclab_experimental.utils.captured`, since :class:`~isaaclab_experimental.envs.DirectRLEnvWarp` no
+  longer records them.
diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py
index 5eac6025c18d..5c7dd082ae8a 100644
--- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py
+++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py
@@ -10,6 +10,7 @@
 import warp as wp
 from isaaclab_experimental.envs import DirectRLEnvWarp
 from isaaclab_experimental.utils.warp.utils import wrap_to_pi
+from isaaclab_experimental.utils.warp_capture import captured
 
 from isaaclab.utils.seed import WarpRng
 
@@ -218,9 +219,11 @@ def _pre_physics_step(self, actions: wp.array) -> None:
             ],
         )
 
+    @captured
     def _apply_action(self) -> None:
         self.cartpole.set_joint_effort_target_mask(target=self.actions)
 
+    @captured
     def _get_observations(self) -> dict:
         wp.launch(
             get_observations,
@@ -237,6 +240,7 @@ def _get_observations(self) -> dict:
         )
         return {"policy": self.torch_obs_buf}
 
+    @captured
     def _get_rewards(self) -> None:
         wp.launch(
             compute_rewards,
@@ -257,6 +261,7 @@ def _get_rewards(self) -> None:
             ],
         )
 
+    @captured
     def _get_dones(self) -> None:
         wp.launch(
             get_dones,
@@ -273,6 +278,7 @@ def _get_dones(self) -> None:
             ],
         )
 
+    @captured
     def _reset_idx(self, mask: wp.array | None = None) -> None:
         if mask is None:
             mask = self._ALL_ENV_MASK
diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/rewards.py
index 0e7ac5ca0fcf..b9f61f1b4c05 100644
--- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/rewards.py
+++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/rewards.py
@@ -40,11 +40,11 @@ def _joint_pos_target_l2_kernel(
 def joint_pos_target_l2(env: ManagerBasedRLEnv, out, target: float, asset_cfg: SceneEntityCfg) -> None:
     """Penalize joint position deviation from a target value. Writes into ``out``."""
     asset: Articulation = env.scene[asset_cfg.name]
-    assert asset.data.joint_pos.warp.shape[1] == asset_cfg.joint_mask.shape[0]
+    assert asset.data.joint_pos.warp.shape[1] == asset_cfg.joint_mask_wp.shape[0]
     wp.launch(
         kernel=_joint_pos_target_l2_kernel,
         dim=env.num_envs,
-        inputs=[asset.data.joint_pos.warp, asset_cfg.joint_mask, out, target],
+        inputs=[asset.data.joint_pos.warp, asset_cfg.joint_mask_wp, out, target],
         device=env.device,
     )
 
diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py
index b7852ac644e0..dad1560b9cf1 100644
--- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py
+++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py
@@ -7,6 +7,7 @@
 
 import warp as wp
 from isaaclab_experimental.envs import DirectRLEnvWarp
+from isaaclab_experimental.utils.warp_capture import captured
 
 from isaaclab.envs import DirectRLEnvCfg
 from isaaclab.utils.seed import WarpRng
@@ -525,6 +526,7 @@ def _pre_physics_step(self, actions: wp.array) -> None:
             inputs=[actions, self.actions_mapped, self.joint_gears, self.action_scale],
         )
 
+    @captured
     def _apply_action(self) -> None:
         self.robot.set_joint_effort_target_mask(target=self.actions_mapped)
 
@@ -568,6 +570,7 @@ def _compute_intermediate_values(self) -> None:
             ],
         )
 
+    @captured
     def _get_observations(self) -> dict:
         wp.launch(
             observations,
@@ -594,6 +597,7 @@ def _get_observations(self) -> dict:
         )
         return {"policy": self.torch_obs_buf}
 
+    @captured
     def _get_rewards(self) -> None:
         wp.launch(
             compute_rewards,
@@ -621,6 +625,7 @@ def _get_rewards(self) -> None:
             ],
         )
 
+    @captured
     def _get_dones(self) -> None:
         self._compute_intermediate_values()
 
@@ -638,6 +643,7 @@ def _get_dones(self) -> None:
             ],
         )
 
+    @captured
     def _reset_idx(self, mask: wp.array | None = None):
         if mask is None:
             mask = self._ALL_ENV_MASK
diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
index f98a9e0e3375..e99a2d670331 100644
--- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
+++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
@@ -12,6 +12,7 @@
 import torch
 import warp as wp
 from isaaclab_experimental.envs import DirectRLEnvWarp
+from isaaclab_experimental.utils.warp_capture import captured
 
 from isaaclab.utils.seed import WarpRng
 
@@ -652,6 +653,7 @@ def _pre_physics_step(self, actions: wp.array) -> None:
         # Store actions in a persistent Warp buffer (analogous to `actions.clone()` in the Torch env).
         wp.copy(self.actions, actions)
 
+    @captured
     def _apply_action(self) -> None:
         wp.launch(
             apply_actions_to_targets,
@@ -672,6 +674,7 @@ def _apply_action(self) -> None:
         # All joints are actuated for Allegro, so default masks (None = all) are correct.
         self.hand.set_joint_position_target_mask(target=self.cur_targets)
 
+    @captured
     def _get_observations(self) -> dict:
         # if self.cfg.asymmetric_obs:
         #    self.fingertip_force_sensors = self.hand.root_physx_view.get_link_incoming_joint_force()[
@@ -687,6 +690,7 @@ def _get_observations(self) -> dict:
             raise ValueError(f"Unknown obs_type: {self.cfg.obs_type}")
         return {"policy": self.torch_obs_buf}
 
+    @captured
     def _get_rewards(self) -> None:
         # Clear reduction buffers before launching the reward kernel.
         # wp.assign(self._num_resets, 0.0)
@@ -742,6 +746,7 @@ def _get_rewards(self) -> None:
         # This avoids Torch-side index extraction and keeps the step graphable.
         self._reset_target_pose(mask=self.reset_goal_buf)
 
+    @captured
     def _get_dones(self) -> None:
         self._compute_intermediate_values()
 
@@ -765,6 +770,7 @@ def _get_dones(self) -> None:
             device=self.device,
         )
 
+    @captured
     def _reset_idx(self, mask: wp.array | None = None):
         if mask is None:
             mask = self._ALL_ENV_MASK

From 0fb73002cd2901dfba4917da31116e2912f3ecff Mon Sep 17 00:00:00 2001
From: Jichuan Hu <jichuanh@nvidia.com>
Date: Tue, 20 Oct 2026 09:52:11 -0700
Subject: [PATCH 37/48] Reach the physics backend through the simulation's
 backend manager in the Warp environments

---
 .../isaaclab_experimental/envs/direct_rl_env_warp.py        | 6 +++---
 .../isaaclab_experimental/envs/manager_based_env_warp.py    | 6 +++---
 .../isaaclab_experimental/utils/warp_capture.py             | 6 +++---
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
index b2e1a570152a..b906771136e7 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
@@ -195,8 +195,8 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
             self.scene.update(dt=self.physics_dt)
         # let the physics backend know about the env decimation so it can
         # fold the full loop into a single step() when possible
-        self.sim.physics_manager.set_decimation(self.cfg.decimation)
-        self._physics_handles_decimation = self.sim.physics_manager.handles_decimation()
+        self.sim.backend_manager.set_decimation(self.cfg.decimation)
+        self._physics_handles_decimation = self.sim.backend_manager.handles_decimation()
 
         # check if debug visualization is has been implemented by the environment
         source_code = inspect.getsource(self._set_debug_vis_impl)
@@ -241,7 +241,7 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
 
         # the task stages record CUDA graphs, and record again after the physics buffers rebind
         CapturedStage.enabled = True
-        CapturedStage.invalidate_on(self.sim.physics_manager)
+        CapturedStage.invalidate_on(self.sim.backend_manager)
 
         # setup the action and observation spaces for Gym
         self._configure_gym_env_spaces()
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
index 6dce23c172fb..932020581dc0 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
@@ -103,7 +103,7 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
 
         # the managers' stages record CUDA graphs, and record again after the physics buffers rebind
         CapturedStage.enabled = True
-        CapturedStage.invalidate_on(self.sim.physics_manager)
+        CapturedStage.invalidate_on(self.sim.backend_manager)
 
         # print useful information
         logger.info(
@@ -175,8 +175,8 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
             self.scene.update(dt=self.physics_dt)
         # let the physics backend know about the env decimation so it can
         # fold the full loop into a single step() when possible
-        self.sim.physics_manager.set_decimation(self.cfg.decimation)
-        self._physics_handles_decimation = self.sim.physics_manager.handles_decimation()
+        self.sim.backend_manager.set_decimation(self.cfg.decimation)
+        self._physics_handles_decimation = self.sim.backend_manager.handles_decimation()
 
         # TODO(jichuanh): This is a temporary solution for event_manager only, but it should be general for all managers
         # Resolve SceneEntityCfg-dependent term params once before any captured event paths.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_capture.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_capture.py
index 1f0d33ef8490..7ac964333549 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_capture.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_capture.py
@@ -166,14 +166,14 @@ def reset(self) -> None:
             self._recordings.clear()
 
     @classmethod
-    def invalidate_on(cls, physics_manager: type[PhysicsManager]) -> None:
+    def invalidate_on(cls, backend_manager: PhysicsManager) -> None:
         """Record every stage again after the physics backend rebinds its buffers or stops.
 
         Args:
-            physics_manager: Physics manager class that dispatches the lifecycle events.
+            backend_manager: Physics manager that dispatches the lifecycle events.
         """
         cls._callback_handles = [
-            physics_manager.register_callback(cls.invalidate, event)
+            backend_manager.register_callback(cls.invalidate, event)
             for event in (PhysicsEvent.PHYSICS_READY, PhysicsEvent.STOP)
         ]
 

From f6d31caca0a498d58dc9028b95faa6e4efdb9c6d Mon Sep 17 00:00:00 2001
From: Jichuan Hu <jichuanh@nvidia.com>
Date: Tue, 20 Oct 2026 09:59:05 -0700
Subject: [PATCH 38/48] Drop the history_order changelog entry that main
 already ships

---
 .../changelog.d/jichuanh-warp-scene-mask-native.fixed.rst      | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
index 7154f9e9454a..b712f5b4a8d4 100644
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
@@ -1,5 +1,2 @@
 * Fixed :meth:`~isaaclab_experimental.envs.InteractiveSceneWarp.reset` with ``env_mask`` resetting deformable
   objects and sensors in every environment, and skipping cable objects.
-* Fixed the Warp :class:`~isaaclab_experimental.managers.ObservationManager` rejecting every observation group with
-  ``Configuration for the term 'history_order' is not of type ObservationTermCfg``, which stopped all manager-based
-  Warp environments from constructing. ``history_order`` is now read as a group setting, as in the stable manager.

From ed3463deb07f46e735e6b20ed7c035c88dd123d6 Mon Sep 17 00:00:00 2001
From: Jichuan Hu <jichuanh@nvidia.com>
Date: Tue, 20 Oct 2026 10:09:38 -0700
Subject: [PATCH 39/48] Follow the hand with the Warp reorient goal marker when
 no fixed position is set

---
 .../core/reorient/reorient_warp_env.py                       | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
index 76ec933341fc..a25b9d547c07 100644
--- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
+++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
@@ -626,7 +626,10 @@ def __init__(self, cfg: AllegroHandEnvCfg, render_mode: str | None = None, **kwa
         )
         self.in_hand_pos.assign(wp.from_torch(in_hand_pos, dtype=wp.vec3f))
 
-        goal_pos = torch.tensor(self.cfg.goal_marker_position, device=self.device).repeat((self.num_envs, 1))
+        if self.cfg.goal_marker_position is None:
+            goal_pos = in_hand_pos + torch.tensor(self.cfg.goal_marker_offset, device=self.device)
+        else:
+            goal_pos = torch.tensor(self.cfg.goal_marker_position, device=self.device).repeat((self.num_envs, 1))
         self.goal_pos.assign(wp.from_torch(goal_pos, dtype=wp.vec3f))
 
         goal_rot = torch.zeros((self.num_envs, 4), device=self.device, dtype=torch.float32)

From 2030fe34401b16b7fb3bc73293ee0e725d6e0d22 Mon Sep 17 00:00:00 2001
From: Jichuan Hu <jichuanh@nvidia.com>
Date: Thu, 22 Oct 2026 00:11:43 -0700
Subject: [PATCH 40/48] Move the shared WarpRng state out of this pull request

Revert 6ac6061a137 and 9328d981380: the process-wide WarpRng state moves to its own pull request stacked on this one. The ray caster, Warp environments, MDP terms, managers, noise models and Warp tasks keep their own RNG arrays here, as on main.
---
 docs/source/concepts/warp_environments.rst    |  3 +-
 .../jichuanh-warp-scene-mask-native.added.rst |  3 -
 ...ichuanh-warp-scene-mask-native.changed.rst |  6 +-
 .../sensors/ray_caster/base_ray_caster.py     | 14 ++++-
 .../isaaclab/sensors/ray_caster/kernels.py    | 19 ++++++
 source/isaaclab/isaaclab/utils/seed.py        | 60 +-----------------
 .../isaaclab/isaaclab/utils/warp/kernels.py   | 25 --------
 source/isaaclab/test/utils/test_warp_rng.py   | 63 -------------------
 ...ichuanh-warp-scene-mask-native.changed.rst |  4 --
 ...ichuanh-warp-scene-mask-native.removed.rst |  2 -
 .../envs/direct_rl_env_warp.py                |  9 +--
 .../envs/manager_based_env_warp.py            | 39 +++++++++---
 .../isaaclab_experimental/envs/mdp/events.py  | 23 ++++---
 .../managers/command_manager.py               |  5 +-
 .../managers/event_manager.py                 | 14 +++--
 .../managers/observation_manager.py           |  6 ++
 .../utils/noise/noise_model.py                | 12 ++--
 .../test/envs/mdp/test_events_warp_parity.py  |  5 +-
 .../envs/mdp/test_observations_warp_parity.py |  7 +--
 .../test/envs/mdp/test_rewards_warp_parity.py |  7 +--
 .../envs/mdp/test_terminations_warp_parity.py |  7 +--
 .../sensors/test_newton_raycast_sensor.py     |  4 --
 ...ichuanh-warp-scene-mask-native.changed.rst |  2 -
 .../core/cartpole/cartpole_warp_env.py        | 27 +++++++-
 .../core/locomotion/locomotion_env_warp.py    | 20 ++++--
 .../core/reorient/reorient_warp_env.py        | 34 ++++++++--
 26 files changed, 184 insertions(+), 236 deletions(-)
 delete mode 100644 source/isaaclab/test/utils/test_warp_rng.py
 delete mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
 delete mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.removed.rst
 delete mode 100644 source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-scene-mask-native.changed.rst

diff --git a/docs/source/concepts/warp_environments.rst b/docs/source/concepts/warp_environments.rst
index 8dcba5296a79..1debce871a09 100644
--- a/docs/source/concepts/warp_environments.rst
+++ b/docs/source/concepts/warp_environments.rst
@@ -493,8 +493,7 @@ checks the mask to skip non-selected environments:
            return
        # ... modify state for selected envs only
 
-- RNG uses the per-env :attr:`WarpRng.state <isaaclab.utils.seed.WarpRng.state>` (``wp.uint32``) instead of
-  ``torch.rand``
+- RNG uses per-env ``env.rng_state_wp`` (``wp.uint32``) instead of ``torch.rand``
 - **Startup/prestartup** events use the torch convention ``(env, env_ids, **params)``
 - **Reset/interval** events use the warp convention ``(env, env_mask, **params)``
 
diff --git a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst
index df8793cd86bb..37303ca182c2 100644
--- a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst
+++ b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst
@@ -5,6 +5,3 @@
   :meth:`~isaaclab.utils.buffers.DelayBuffer.reset`, which select batches on the device without synchronizing the
   host. Tensor lags set through a mask are full-sized and are not range-checked.
 * Added ``env_mask`` to :meth:`~isaaclab.assets.BaseCableObject.reset`, matching the other scene assets.
-* Added :class:`~isaaclab.utils.seed.WarpRng`, the per-environment Warp random number generator state that
-  environments and sensors share, one per process. :func:`~isaaclab.utils.seed.configure_seed` reseeds it in place,
-  and environments call :meth:`~isaaclab.utils.seed.WarpRng.initialize` at construction.
diff --git a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
index 1467738a803b..ae017575fe71 100644
--- a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
+++ b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
@@ -1,7 +1,7 @@
 * Changed masked :meth:`~isaaclab.sensors.ray_caster.BaseRayCaster.reset` to resample drift in a Warp kernel from
-  the per-environment random states of :class:`~isaaclab.utils.seed.WarpRng`, so it neither synchronizes the host
-  nor allocates and can be captured in a CUDA graph. When both ``env_ids`` and ``env_mask`` are given, the mask
-  selects the resampled environments, matching :meth:`~isaaclab.sensors.SensorBase.reset`.
+  per-environment random states seeded from torch's global generator, so it neither synchronizes the host nor
+  allocates and can be captured in a CUDA graph. When both ``env_ids`` and ``env_mask`` are given, the mask selects
+  the resampled environments, matching :meth:`~isaaclab.sensors.SensorBase.reset`.
 * Changed :meth:`~isaaclab.actuators.newton.NewtonActuatorAdapter.reset` to reuse preallocated per-DOF masks, so
   masked resets of Newton-native actuator state do not allocate and can be captured in a CUDA graph.
 * **Breaking:** Changed :meth:`~isaaclab.actuators.ActuatorBase.reset` to take an ``env_mask`` argument, the
diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py
index 8c022d0df1e8..3c7e034b2167 100644
--- a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py
+++ b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py
@@ -19,7 +19,6 @@
 from ...markers import VisualizationMarkers
 from ...terrains.trimesh.utils import make_plane
 from ...utils import math as math_utils
-from ...utils.seed import WarpRng
 from ...utils.warp import ProxyArray, convert_to_warp_mesh
 from ...utils.warp.kernels import raycast_mesh_masked_kernel
 from ..sensor_base import SensorBase
@@ -118,7 +117,7 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None
         if not self._drift_sampled:
             return
         if env_mask is not None:
-            # Resample in a kernel from the shared per-environment random states with the ranges as launch
+            # Resample in a kernel from per-environment random states with the ranges as launch
             # arguments, so the masked reset neither synchronizes nor allocates and can be captured.
             # A zero range draws zero, which clears drift that was sampled before it was disabled.
             wp.launch(
@@ -129,7 +128,7 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None
                     wp.vec2f(*self.cfg.drift_range),
                     wp.vec3f(*(low for low, _ in ray_cast_range_list)),
                     wp.vec3f(*(high for _, high in ray_cast_range_list)),
-                    WarpRng.state,
+                    self._drift_rng_state,
                 ],
                 outputs=[self.drift.warp, self.ray_cast_drift.warp],
                 device=self._device,
@@ -262,6 +261,15 @@ def _initialize_rays_impl(self):
         # Drift buffers are warp-first; reset uses explicit .torch views for sampling.
         self.drift = ProxyArray(wp.zeros(self._view_count, dtype=wp.vec3f, device=self._device))
         self.ray_cast_drift = ProxyArray(wp.zeros(self._view_count, dtype=wp.vec3f, device=self._device))
+        # Per-environment random states for masked drift resampling, seeded from torch's global generator.
+        self._drift_rng_state = wp.empty(self._view_count, dtype=wp.uint32, device=self._device)
+        wp.launch(
+            ray_caster_kernels.init_rng_state_kernel,
+            dim=self._view_count,
+            inputs=[int(torch.randint(0, 2**31 - 1, (1,)))],
+            outputs=[self._drift_rng_state],
+            device=self._device,
+        )
 
         # World-frame ray buffers
         self._ray_starts_w = wp.empty((self._view_count, self.num_rays), dtype=wp.vec3f, device=self._device)
diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py b/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py
index 6fa4c65ffcc1..9d36d1a3ddd7 100644
--- a/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py
+++ b/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py
@@ -170,6 +170,25 @@ def apply_z_drift_kernel(
     ray_hits[env, ray] = wp.vec3f(hit[0], hit[1], hit[2] + ray_cast_drift[env][2])
 
 
+@wp.kernel(enable_backward=False)
+def init_rng_state_kernel(
+    # input
+    seed: wp.int32,
+    # output
+    rng_state: wp.array(dtype=wp.uint32),
+):
+    """Initialize one random number generator state per environment.
+
+    Launch with dim=(num_envs,).
+
+    Args:
+        seed: Seed shared by all environments; each environment uses its index as the offset.
+        rng_state: Random number generator states. Shape is (num_envs,).
+    """
+    env = wp.tid()
+    rng_state[env] = wp.rand_init(seed, env)
+
+
 @wp.kernel(enable_backward=False)
 def resample_drift_masked_kernel(
     # input
diff --git a/source/isaaclab/isaaclab/utils/seed.py b/source/isaaclab/isaaclab/utils/seed.py
index 6d8f13de85dd..535789305164 100644
--- a/source/isaaclab/isaaclab/utils/seed.py
+++ b/source/isaaclab/isaaclab/utils/seed.py
@@ -3,8 +3,6 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
-from __future__ import annotations
-
 import os
 import random
 
@@ -13,62 +11,6 @@
 import warp as wp
 
 
-class WarpRng:
-    """Process-wide per-environment random number generator state for Warp kernels.
-
-    Warp has no global generator: :func:`warp.rand_init` hashes a seed and an offset into a state that lives only
-    for one kernel launch. Kernels that draw random numbers per environment therefore read ``state[env]``, draw, and
-    store the advanced state back, so the next launch or CUDA-graph replay draws new numbers. This class holds that
-    array once per process, so environments, managers and sensors share it instead of each allocating a copy.
-
-    :meth:`initialize` must be called once the number of environments and the device are known; environments do this
-    at construction, so users read :attr:`state` directly. :meth:`seed` is optional; without it, the state is seeded
-    from :func:`torch.initial_seed`. Only kernels launched over environments, one thread per environment, may advance
-    the state.
-    """
-
-    state: wp.array | None = None
-    """The per-environment random number generator state, or None before :meth:`initialize`. Shape is (num_envs,)."""
-
-    _seed: int | None = None
-
-    @classmethod
-    def seed(cls, seed: int) -> None:
-        """Set the seed, and reseed the state in place if it is initialized.
-
-        Args:
-            seed: The random seed value.
-        """
-        cls._seed = seed
-        if cls.state is not None:
-            cls._reset_state()
-
-    @classmethod
-    def initialize(cls, num_envs: int, device: str) -> None:
-        """Allocate and seed a new state. Environments call it once at construction.
-
-        Args:
-            num_envs: The number of environments.
-            device: The device of the state.
-        """
-        cls.state = wp.empty(num_envs, dtype=wp.uint32, device=device)
-        cls._reset_state()
-
-    @classmethod
-    def _reset_state(cls) -> None:
-        """Seed every environment's state from the set seed, or from torch's initial seed without one."""
-        from .warp.kernels import initialize_rng_state
-
-        seed = cls._seed if cls._seed is not None else torch.initial_seed()
-        wp.launch(
-            initialize_rng_state,
-            dim=cls.state.shape[0],
-            inputs=[seed % 2**31],
-            outputs=[cls.state],
-            device=cls.state.device,
-        )
-
-
 def configure_seed(seed: int | None, torch_deterministic: bool = False) -> int:
     """Set seed across all random number generators (torch, numpy, random, warp).
 
@@ -89,7 +31,7 @@ def configure_seed(seed: int | None, torch_deterministic: bool = False) -> int:
     os.environ["PYTHONHASHSEED"] = str(seed)
     torch.cuda.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
-    WarpRng.seed(seed)
+    wp.rand_init(seed)
 
     if torch_deterministic:
         # refer to https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility
diff --git a/source/isaaclab/isaaclab/utils/warp/kernels.py b/source/isaaclab/isaaclab/utils/warp/kernels.py
index a7818691165e..3d046abe2bba 100644
--- a/source/isaaclab/isaaclab/utils/warp/kernels.py
+++ b/source/isaaclab/isaaclab/utils/warp/kernels.py
@@ -780,28 +780,3 @@ def gather_subtract_2d(
 def gather_subtract_2d_kernel(env_ids: "wp.array | torch.Tensor", item_ids: "wp.array | torch.Tensor") -> wp.Kernel:
     """Select the gather-subtract worker for the selector dtypes."""
     return _GATHER_SUBTRACT_2D_DISPATCHER.select(env_ids, item_ids)
-
-
-##
-# Random number generation.
-##
-
-
-@wp.kernel(enable_backward=False)
-def initialize_rng_state(
-    # input
-    seed: wp.int32,
-    # output
-    state: wp.array(dtype=wp.uint32),
-):
-    """Initialize one random number generator state per environment.
-
-    Launch with ``dim=(num_envs,)``. Kernels that draw from ``state[i]`` should store the advanced state back,
-    so the next launch or graph replay draws new numbers.
-
-    Args:
-        seed: Seed shared by all environments; each environment uses its index as the offset.
-        state: Random number generator states. Shape is (num_envs,).
-    """
-    env_id = wp.tid()
-    state[env_id] = wp.rand_init(seed, env_id)
diff --git a/source/isaaclab/test/utils/test_warp_rng.py b/source/isaaclab/test/utils/test_warp_rng.py
deleted file mode 100644
index 2bc3815a2b8c..000000000000
--- a/source/isaaclab/test/utils/test_warp_rng.py
+++ /dev/null
@@ -1,63 +0,0 @@
-# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
-# All rights reserved.
-#
-# SPDX-License-Identifier: BSD-3-Clause
-
-"""Tests for the process-wide per-environment Warp random number generator state."""
-
-import numpy as np
-import pytest
-import torch
-
-from isaaclab.test.utils import test_devices
-from isaaclab.utils.seed import WarpRng, configure_seed
-
-
-@pytest.fixture(autouse=True)
-def fresh_warp_rng():
-    """Start and end every test without a state or a set seed."""
-    WarpRng.state = None
-    WarpRng._seed = None
-    yield
-    WarpRng.state = None
-    WarpRng._seed = None
-
-
-@pytest.mark.parametrize("device", test_devices())
-def test_initialize_allocates_a_new_state(device):
-    """Every :meth:`WarpRng.initialize` call allocates and seeds a new state of the given size."""
-    WarpRng.initialize(4, device)
-    state = WarpRng.state
-    WarpRng.initialize(5, device)
-    assert WarpRng.state is not state
-    assert WarpRng.state.shape == (5,)
-
-
-@pytest.mark.parametrize("device", test_devices())
-def test_configure_seed_reseeds_in_place(device):
-    """configure_seed reseeds the existing state in place and reproducibly."""
-    configure_seed(7)
-    WarpRng.initialize(4, device)
-    state = WarpRng.state
-    seeded = state.numpy().copy()
-
-    state.fill_(0)
-    configure_seed(7)
-    assert WarpRng.state.ptr == state.ptr
-    np.testing.assert_array_equal(state.numpy(), seeded)
-
-    configure_seed(8)
-    assert not np.array_equal(state.numpy(), seeded)
-
-
-@pytest.mark.parametrize("device", test_devices())
-def test_default_seed_follows_torch(device):
-    """Without a set seed, the state is seeded from torch's initial seed."""
-    torch.manual_seed(5)
-    WarpRng.initialize(4, device)
-    default = WarpRng.state.numpy().copy()
-
-    WarpRng.state = None
-    WarpRng.seed(5)
-    WarpRng.initialize(4, device)
-    np.testing.assert_array_equal(WarpRng.state.numpy(), default)
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.changed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
deleted file mode 100644
index aa642663b908..000000000000
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
+++ /dev/null
@@ -1,4 +0,0 @@
-* Changed :class:`~isaaclab_experimental.envs.ManagerBasedEnvWarp` and
-  :class:`~isaaclab_experimental.envs.DirectRLEnvWarp` to initialize the process-wide
-  :class:`~isaaclab.utils.seed.WarpRng` state, which the Warp MDP terms, managers and noise models now draw from.
-  An environment created without a seed now seeds it from :func:`torch.initial_seed` instead of ``-1``.
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.removed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.removed.rst
deleted file mode 100644
index 8e486bedf7bd..000000000000
--- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.removed.rst
+++ /dev/null
@@ -1,2 +0,0 @@
-* **Breaking:** Removed ``ManagerBasedEnvWarp.rng_state_wp``. Read :attr:`~isaaclab.utils.seed.WarpRng.state`
-  instead.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
index 8ef4e8659512..2c032cea3c22 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
@@ -32,7 +32,7 @@
 from isaaclab.sim.utils import use_stage
 from isaaclab.utils import instantiate, validate
 from isaaclab.utils.noise import NoiseModel
-from isaaclab.utils.seed import WarpRng, configure_seed
+from isaaclab.utils.seed import configure_seed
 from isaaclab.utils.timer import Timer
 
 from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache
@@ -125,7 +125,7 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
 
         # set the seed for the environment
         if self.cfg.seed is not None:
-            self.cfg.seed = configure_seed(self.cfg.seed)
+            self.cfg.seed = self.seed(self.cfg.seed)
         else:
             logger.warning("Seed not set for the environment. The environment creation may not be deterministic.")
 
@@ -166,9 +166,6 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
                 # attach_stage_to_usd_context()
         logger.info(f"Scene manager: {self.scene}")
 
-        # Process-wide per-env Warp RNG state (isaaclab.utils.seed.WarpRng), shared by terms, tasks and sensors.
-        WarpRng.initialize(self.num_envs, self.device)
-
         # create event manager
         # note: this is needed here (rather than after simulation play) to allow USD-related randomization events
         #   that must happen before the simulation starts. Example: randomizing mesh scale
@@ -348,7 +345,7 @@ def reset(self, seed: int | None = None, options: dict[str, Any] | None = None)
         """
         # set the seed
         if seed is not None:
-            configure_seed(seed)
+            self.seed(seed)
 
         # reset state of scene
         self._reset_idx(self._ALL_ENV_MASK)
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
index d4a825febe43..8b6710633ae4 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
@@ -36,7 +36,7 @@
 from isaaclab.sim import SimulationContext
 from isaaclab.sim.utils import use_stage
 from isaaclab.utils import validate
-from isaaclab.utils.seed import WarpRng, configure_seed
+from isaaclab.utils.seed import configure_seed
 from isaaclab.utils.timer import Timer
 
 from isaaclab_experimental.envs.interactive_scene_warp import InteractiveSceneWarp as InteractiveScene
@@ -47,6 +47,17 @@
 logger = logging.getLogger(__name__)
 
 
+@wp.kernel
+def initialize_rng_state(
+    # input
+    seed: wp.int32,
+    # output
+    state: wp.array(dtype=wp.uint32),
+):
+    env_id = wp.tid()
+    state[env_id] = wp.rand_init(seed, wp.int32(env_id))
+
+
 class ManagerBasedEnvWarp:
     """The base environment for the manager-based workflow (experimental fork).
 
@@ -86,7 +97,7 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
 
         # set the seed for the environment
         if self.cfg.seed is not None:
-            self.cfg.seed = configure_seed(self.cfg.seed)
+            self.cfg.seed = self.seed(self.cfg.seed)
         else:
             logger.warning("Seed not set for the environment. The environment creation may not be deterministic.")
 
@@ -138,8 +149,16 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
                 # attach_stage_to_usd_context()
         logger.info(f"Scene manager: {self.scene}")
 
-        # Process-wide per-env Warp RNG state (isaaclab.utils.seed.WarpRng), shared by terms, tasks and sensors.
-        WarpRng.initialize(self.num_envs, self.device)
+        # Shared per-env Warp RNG state (accessible to all managers/terms via `env`).
+        # This is a single stream per env (no lookup) and is initialized once when `num_envs` is known.
+        self.rng_state_wp = wp.zeros((self.num_envs,), dtype=wp.uint32, device=self.device)
+        seed_val = int(self.cfg.seed) if self.cfg.seed is not None else -1
+        wp.launch(
+            kernel=initialize_rng_state,
+            dim=self.num_envs,
+            inputs=[seed_val, self.rng_state_wp],
+            device=self.device,
+        )
 
         # TODO(jichuanh): this is problematic as warp capture requires stable pointers,
         #                 using different masks for different managers/terms will cause problems.
@@ -423,10 +442,16 @@ def reset(
 
         # set the seed
         if seed is not None:
-            # also reseeds the per-env Warp RNG state in place (stable pointer for capture)
-            used_seed = configure_seed(seed)
+            used_seed = self.seed(seed)
             # keep cfg seed in sync for downstream users
             self.cfg.seed = used_seed
+            # re-initialize per-env Warp RNG state without reallocating (stable pointer for capture)
+            wp.launch(
+                kernel=initialize_rng_state,
+                dim=self.num_envs,
+                inputs=[int(used_seed), self.rng_state_wp],
+                device=self.device,
+            )
 
         # reset state of scene
         self._reset_idx(env_ids)
@@ -481,7 +506,7 @@ def reset_to(
 
         # set the seed
         if seed is not None:
-            configure_seed(seed)
+            self.seed(seed)
 
         self._reset_idx(env_ids)
 
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py
index a64af1d309ce..90e5696a6b10 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py
@@ -31,7 +31,6 @@
 
 from isaaclab.envs.mdp.events import randomize_rigid_body_mass as _StableRandomizeRigidBodyMass
 from isaaclab.envs.mdp.events import randomize_rigid_body_material as _StableRandomizeRigidBodyMaterial
-from isaaclab.utils.seed import WarpRng
 
 from isaaclab_experimental.managers import ManagerTermBase as _WarpManagerTermBase
 from isaaclab_experimental.managers import SceneEntityCfg
@@ -99,7 +98,7 @@ def randomize_rigid_body_com(
         dim=env.num_envs,
         inputs=[
             env_mask,
-            WarpRng.state,
+            env.rng_state_wp,
             asset.data.body_com_pos_b.warp,
             asset_cfg.body_ids_wp,
             fn._com_lo,
@@ -179,7 +178,7 @@ def apply_external_force_torque(
         dim=env.num_envs,
         inputs=[
             env_mask,
-            WarpRng.state,
+            env.rng_state_wp,
             apply_external_force_torque._scratch_forces,
             apply_external_force_torque._scratch_torques,
             force_range[0],
@@ -259,7 +258,7 @@ def push_by_setting_velocity(
         dim=env.num_envs,
         inputs=[
             env_mask,
-            WarpRng.state,
+            env.rng_state_wp,
             asset.data.root_vel_w.warp,
             push_by_setting_velocity._scratch_vel,
             push_by_setting_velocity._lin_lo,
@@ -378,7 +377,7 @@ def reset_root_state_uniform(
         dim=env.num_envs,
         inputs=[
             env_mask,
-            WarpRng.state,
+            env.rng_state_wp,
             asset.data.default_root_pose.warp,
             asset.data.default_root_vel.warp,
             env.env_origins_wp,
@@ -471,6 +470,11 @@ def reset_joints_by_offset(
             f"but got None for asset '{asset_cfg.name}'. "
             "Use isaaclab_experimental.managers.SceneEntityCfg and ensure joint_names are set."
         )
+    if not hasattr(env, "rng_state_wp") or env.rng_state_wp is None:
+        raise AttributeError(
+            "reset_joints_by_offset requires env.rng_state_wp to be initialized. "
+            "Use ManagerBasedEnvWarp or ManagerBasedRLEnvWarp as the base environment."
+        )
 
     wp.launch(
         kernel=_reset_joints_by_offset_kernel,
@@ -478,7 +482,7 @@ def reset_joints_by_offset(
         inputs=[
             env_mask,
             asset_cfg.joint_ids_wp,
-            WarpRng.state,
+            env.rng_state_wp,
             asset.data.default_joint_pos.warp,
             asset.data.default_joint_vel.warp,
             asset.data.joint_pos.warp,
@@ -557,6 +561,11 @@ def reset_joints_by_scale(
             f"but got None for asset '{asset_cfg.name}'. "
             "Use isaaclab_experimental.managers.SceneEntityCfg and ensure joint_names are set."
         )
+    if not hasattr(env, "rng_state_wp") or env.rng_state_wp is None:
+        raise AttributeError(
+            "reset_joints_by_scale requires env.rng_state_wp to be initialized. "
+            "Use ManagerBasedEnvWarp or ManagerBasedRLEnvWarp as the base environment."
+        )
 
     wp.launch(
         kernel=_reset_joints_by_scale_kernel,
@@ -564,7 +573,7 @@ def reset_joints_by_scale(
         inputs=[
             env_mask,
             asset_cfg.joint_ids_wp,
-            WarpRng.state,
+            env.rng_state_wp,
             asset.data.default_joint_pos.warp,
             asset.data.default_joint_vel.warp,
             asset.data.joint_pos.warp,
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py
index 2b4cd8de1828..e7d8b4b746b0 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py
@@ -18,7 +18,6 @@
 
 from isaaclab.managers.manager_term_cfg import CommandTermCfg
 from isaaclab.utils import instantiate
-from isaaclab.utils.seed import WarpRng
 
 from isaaclab_experimental.utils.warp.kernels import compute_reset_scale, count_masked
 
@@ -307,6 +306,8 @@ def _resample(self, env_mask: wp.array):
             raise TypeError(f"env_mask must be a wp.array (got {type(env_mask)}).")
         if env_mask.dtype != wp.bool or env_mask.ndim != 1:
             raise TypeError(f"env_mask must be wp.bool 1D (got dtype={env_mask.dtype}, ndim={env_mask.ndim}).")
+        if self._env.rng_state_wp is None:
+            raise RuntimeError("Environment rng_state_wp is not initialized.")
 
         # resample time-left and increment command-counter for masked envs
         wp.launch(
@@ -316,7 +317,7 @@ def _resample(self, env_mask: wp.array):
                 env_mask,
                 self.time_left_wp,
                 self.command_counter_wp,
-                WarpRng.state,
+                self._env.rng_state_wp,
                 float(self.cfg.resampling_time_range[0]),
                 float(self.cfg.resampling_time_range[1]),
             ],
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
index 580e5062d4c1..5c3832fb803b 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
@@ -30,7 +30,6 @@
 from prettytable import PrettyTable
 
 from isaaclab.managers.manager_term_cfg import EventTermCfg
-from isaaclab.utils.seed import WarpRng
 
 from .manager_base import ManagerBase
 
@@ -274,7 +273,7 @@ def reset(
                     inputs=[
                         env_mask,
                         self._interval_term_time_left_wp[i],
-                        WarpRng.state,
+                        self._env.rng_state_wp,
                         float(lower),
                         float(upper),
                     ],
@@ -332,6 +331,9 @@ def apply(
             term_cfg.func(self._env, env_mask_wp, **term_cfg.params)
 
     def _apply_interval(self, dt: float) -> None:
+        if self._env.rng_state_wp is None:
+            raise RuntimeError("EventManager._apply_interval: env.rng_state_wp is not initialized.")
+
         # iterate over all the interval terms (fixed list; captured graph-friendly)
         for i, term_cfg in enumerate(self._mode_term_cfgs["interval"]):
             lower, upper = self._interval_term_ranges[i]
@@ -361,7 +363,7 @@ def _apply_interval(self, dt: float) -> None:
                     dim=self.num_envs,
                     inputs=[
                         self._interval_term_time_left_wp[i],
-                        WarpRng.state,
+                        self._env.rng_state_wp,
                         self._scratch_term_mask_wp,
                         float(dt),
                         float(lower),
@@ -448,11 +450,13 @@ def _prepare_terms(self):
                 if term_cfg.is_global_time:
                     # allocate and seed scalar global RNG state if needed (avoid consuming env0 RNG stream)
                     if self._interval_global_rng_state_wp is None:
+                        if self._env.rng_state_wp is None:
+                            raise RuntimeError("EventManager._prepare_terms: env.rng_state_wp is not initialized.")
                         self._interval_global_rng_state_wp = wp.zeros((1,), dtype=wp.uint32, device=self.device)
                         wp.launch(
                             kernel=_seed_global_rng_from_env_rng,
                             dim=1,
-                            inputs=[WarpRng.state, self._interval_global_rng_state_wp],
+                            inputs=[self._env.rng_state_wp, self._interval_global_rng_state_wp],
                             device=self.device,
                         )
                     time_left = wp.zeros((1,), dtype=wp.float32, device=self.device)
@@ -469,7 +473,7 @@ def _prepare_terms(self):
                     wp.launch(
                         kernel=_interval_init_per_env,
                         dim=self.num_envs,
-                        inputs=[time_left, WarpRng.state, float(lower), float(upper)],
+                        inputs=[time_left, self._env.rng_state_wp, float(lower), float(upper)],
                         device=self.device,
                     )
                     self._interval_term_time_left_wp.append(time_left)
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
index 6ec29aa0e921..051aaa8a6ff6 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
@@ -780,8 +780,14 @@ def _prepare_terms(self):  # noqa: C901
                                     f" and optional parameters: {args_with_defaults}, but received: {term_params}."
                                 )
 
+                # plumb the shared per-env RNG state so Warp noise kernels can consume it
+                # (function-style NoiseCfg kernels read it off the cfg at call time)
+                if term_cfg.noise is not None and isinstance(term_cfg.noise, noise.NoiseCfg):
+                    term_cfg.noise.rng_state_wp = self._env.rng_state_wp
                 # prepare noise model classes
                 if term_cfg.noise is not None and isinstance(term_cfg.noise, noise.NoiseModelCfg):
+                    # plumb the shared per-env RNG state so Warp noise kernels can consume it
+                    term_cfg.noise.rng_state_wp = self._env.rng_state_wp
                     noise_model_cls = term_cfg.noise.class_type
                     if not issubclass(noise_model_cls, noise.NoiseModel):
                         raise TypeError(
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/noise/noise_model.py b/source/isaaclab_experimental/isaaclab_experimental/utils/noise/noise_model.py
index cbc933324830..acba926ea488 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/utils/noise/noise_model.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/utils/noise/noise_model.py
@@ -10,8 +10,10 @@
 
     noise_cfg.func(data_wp, noise_cfg) -> None
 
-Random noise kernels (gaussian, uniform) draw from the process-wide per-env Warp RNG state
-:attr:`isaaclab.utils.seed.WarpRng.state`.
+Random noise kernels (gaussian, uniform) consume the shared per-env Warp RNG state
+(``rng_state_wp``) that is set on the config at manager prep time from
+``env.rng_state_wp``.  See :func:`initialize_rng_state` in
+``isaaclab_experimental.envs.manager_based_env_warp`` for the initialization pattern.
 """
 
 from __future__ import annotations
@@ -20,8 +22,6 @@
 
 import warp as wp
 
-from isaaclab.utils.seed import WarpRng
-
 if TYPE_CHECKING:
     from . import noise_cfg
 
@@ -108,7 +108,7 @@ def uniform_noise(data: wp.array, cfg: noise_cfg.UniformNoiseCfg) -> None:
     wp.launch(
         _apply_uniform_noise,
         dim=data.shape[0],
-        inputs=[data, WarpRng.state, float(cfg.n_min), float(cfg.n_max), _OPERATION_MAP[cfg.operation]],
+        inputs=[data, cfg.rng_state_wp, float(cfg.n_min), float(cfg.n_max), _OPERATION_MAP[cfg.operation]],
         device=data.device,
     )
 
@@ -149,7 +149,7 @@ def gaussian_noise(data: wp.array, cfg: noise_cfg.GaussianNoiseCfg) -> None:
     wp.launch(
         _apply_gaussian_noise,
         dim=data.shape[0],
-        inputs=[data, WarpRng.state, float(cfg.mean), float(cfg.std), _OPERATION_MAP[cfg.operation]],
+        inputs=[data, cfg.rng_state_wp, float(cfg.mean), float(cfg.std), _OPERATION_MAP[cfg.operation]],
         device=data.device,
     )
 
diff --git a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py
index e910f2d9fc1a..e93e73267938 100644
--- a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py
+++ b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py
@@ -12,8 +12,6 @@
 import torch
 import warp as wp
 
-from isaaclab.utils.seed import WarpRng
-
 # Skip entire module if no CUDA device available
 wp.init()
 pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required")
@@ -46,7 +44,6 @@ def _clear_function_caches():
     between tests to avoid stale references from prior fixtures.
     """
     yield
-    WarpRng.state = None
     for fn in (
         warp_evt.push_by_setting_velocity,
         warp_evt.apply_external_force_torque,
@@ -105,7 +102,7 @@ class _Env:
     env.step_dt = 0.02
     env.max_episode_length_s = 10.0
     # RNG state for events (seeded deterministically)
-    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
diff --git a/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py
index 91a5f8fe138f..0f014cc2b3fd 100644
--- a/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py
+++ b/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py
@@ -12,8 +12,6 @@
 import torch
 import warp as wp
 
-from isaaclab.utils.seed import WarpRng
-
 # Skip entire module if no CUDA device available
 wp.init()
 pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required")
@@ -52,7 +50,6 @@
 @pytest.fixture(autouse=True)
 def _clear_caches():
     yield
-    WarpRng.state = None
     for fn in [warp_obs.generated_commands]:
         for attr in list(vars(fn)):
             if attr.startswith("_"):
@@ -133,7 +130,7 @@ class _Env:
     env.episode_length_buf = episode_length_buf
     env.step_dt = 0.02
     env.max_episode_length_s = 10.0
-    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
@@ -172,7 +169,7 @@ class _Env:
     env.step_dt = 0.02
     env.max_episode_length = 500
     env.max_episode_length_s = 10.0
-    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
diff --git a/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py
index 9f2d39253d4e..c81f6f73ec89 100644
--- a/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py
+++ b/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py
@@ -12,8 +12,6 @@
 import torch
 import warp as wp
 
-from isaaclab.utils.seed import WarpRng
-
 # Skip entire module if no CUDA device available
 wp.init()
 pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required")
@@ -60,7 +58,6 @@
 @pytest.fixture(autouse=True)
 def _clear_caches():
     yield
-    WarpRng.state = None
     for fn in [warp_rew.track_lin_vel_xy_exp, warp_rew.track_ang_vel_z_exp, warp_rew.undesired_contacts]:
         for attr in list(vars(fn)):
             if attr.startswith("_"):
@@ -147,7 +144,7 @@ class _Env:
     env.episode_length_buf = episode_length_buf
     env.step_dt = 0.02
     env.max_episode_length_s = 10.0
-    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
@@ -187,7 +184,7 @@ class _Env:
     env.step_dt = 0.02
     env.max_episode_length = 500
     env.max_episode_length_s = 10.0
-    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
diff --git a/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py
index e0fbee333530..9676f2b06044 100644
--- a/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py
+++ b/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py
@@ -14,8 +14,6 @@
 import torch
 import warp as wp
 
-from isaaclab.utils.seed import WarpRng
-
 # Skip entire module if no CUDA device available
 wp.init()
 pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required")
@@ -61,7 +59,6 @@
 @pytest.fixture(autouse=True)
 def _clear_caches():
     yield
-    WarpRng.state = None
     for fn in [warp_term.illegal_contact]:
         for attr in list(vars(fn)):
             if attr.startswith("_"):
@@ -142,7 +139,7 @@ class _Env:
     env.episode_length_buf = episode_length_buf
     env.step_dt = 0.02
     env.max_episode_length_s = 10.0
-    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
@@ -182,7 +179,7 @@ class _Env:
     env.step_dt = 0.02
     env.max_episode_length = 500
     env.max_episode_length_s = 10.0
-    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
diff --git a/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py b/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py
index 8cf45b728922..9c454dd72cca 100644
--- a/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py
+++ b/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py
@@ -40,7 +40,6 @@
 from isaaclab.sim import SimulationCfg
 from isaaclab.terrains import TerrainImporterCfg
 from isaaclab.utils import configclass
-from isaaclab.utils.seed import WarpRng
 
 SENSOR_HEIGHT = 2.0
 RAY_OFFSET = 0.2
@@ -223,9 +222,6 @@ def test_partial_reset_resamples_drift_of_selected_envs_only(sim):
     sensor: NewtonRaycastSensor = scene["raycast"]
     sensor.reset()
 
-    # A masked reset draws from the process-wide Warp RNG state, which an environment initializes.
-    WarpRng.initialize(scene.num_envs, sensor.device)
-
     drift = sensor.drift.torch.clone()
     ray_cast_drift = sensor.ray_cast_drift.torch.clone()
     sensor.reset(env_mask=wp.array([False, True, False], dtype=wp.bool, device=sensor.device))
diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-scene-mask-native.changed.rst b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
deleted file mode 100644
index 59b2cee59f4f..000000000000
--- a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-scene-mask-native.changed.rst
+++ /dev/null
@@ -1,2 +0,0 @@
-* Changed the Cartpole, locomotion and reorient Warp environments to draw from the process-wide
-  :class:`~isaaclab.utils.seed.WarpRng` state instead of allocating and seeding their own.
diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py
index dcafd43d57ec..2a3ccdd47f0d 100644
--- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py
+++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py
@@ -11,8 +11,6 @@
 from isaaclab_experimental.envs import DirectRLEnvWarp
 from isaaclab_experimental.utils.warp.utils import wrap_to_pi
 
-from isaaclab.utils.seed import WarpRng
-
 if TYPE_CHECKING:
     from isaaclab_tasks.core.cartpole.cartpole_direct_env_cfg import CartpoleEnvCfg
 
@@ -176,6 +174,16 @@ def reset(
         state[env_index] = rng_state
 
 
+@wp.kernel
+def initialize_state(
+    state: wp.array(dtype=wp.uint32),
+    seed: wp.int32,
+):
+    """Initialize each env's random number generator state from the seed."""
+    env_index = wp.tid()
+    state[env_index] = wp.rand_init(seed, env_index)
+
+
 class CartpoleWarpEnv(DirectRLEnvWarp):
     cfg: CartpoleEnvCfg
 
@@ -198,6 +206,19 @@ def __init__(self, cfg: CartpoleEnvCfg, render_mode: str | None = None, **kwargs
         self.observations = wp.zeros((self.num_envs), dtype=wp.vec4f, device=self.device)
         self.actions = wp.zeros((self.num_envs, self.cartpole.num_joints), dtype=wp.float32, device=self.device)
         self.rewards = wp.zeros((self.num_envs), dtype=wp.float32, device=self.device)
+        self.states = wp.zeros((self.num_envs), dtype=wp.uint32, device=self.device)
+
+        if self.cfg.seed is None:
+            self.cfg.seed = -1
+
+        wp.launch(
+            initialize_state,
+            dim=self.num_envs,
+            inputs=[
+                self.states,
+                self.cfg.seed,
+            ],
+        )
 
         # Bind torch buffers to warp buffers
         self.torch_obs_buf = wp.to_torch(self.observations)
@@ -296,6 +317,6 @@ def _reset_idx(self, mask: wp.array | None = None) -> None:
                 self.cfg.initial_pole_angle_range,
                 self.cfg.initial_pole_velocity_range,
                 mask,
-                WarpRng.state,
+                self.states,
             ],
         )
diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py
index b7852ac644e0..dfcfad6e0f55 100644
--- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py
+++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py
@@ -9,7 +9,6 @@
 from isaaclab_experimental.envs import DirectRLEnvWarp
 
 from isaaclab.envs import DirectRLEnvCfg
-from isaaclab.utils.seed import WarpRng
 from isaaclab.utils.string import resolve_matching_names_values
 
 
@@ -417,12 +416,15 @@ def survival_rate(
 
 
 @wp.kernel
-def initialize_targets(
+def initialize_state(
     env_origins: wp.array(dtype=wp.vec3f),
     targets: wp.array(dtype=wp.vec3f),
+    state: wp.array(dtype=wp.uint32),
+    seed: wp.int32,
 ):
-    """Place each env's target far along +x from the env origin."""
+    """Initialize each env's RNG state and place its target far along +x from the env origin."""
     env_index = wp.tid()
+    state[env_index] = wp.rand_init(seed, env_index)
     targets[env_index] = env_origins[env_index]
     targets[env_index] += wp.static(wp.vec3f(1000.0, 0.0, 0.0))
 
@@ -478,6 +480,7 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
         )
         self.rewards = wp.zeros((self.num_envs), dtype=wp.float32, device=self.sim.device)
         self.actions = wp.zeros((self.num_envs, self.robot.num_joints), dtype=wp.float32, device=self.sim.device)
+        self.states = wp.zeros((self.num_envs), dtype=wp.uint32, device=self.sim.device)
         self.potentials = wp.zeros(self.num_envs, dtype=wp.float32, device=self.sim.device)
         self.prev_potentials = wp.zeros_like(self.potentials)
         self.targets = wp.zeros((self.num_envs), dtype=wp.vec3f, device=self.sim.device)
@@ -493,13 +496,18 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
         self.env_origins = wp.from_torch(self.scene.env_origins, dtype=wp.vec3f)
         self.actions_mapped = wp.zeros((self.num_envs, self.robot.num_joints), dtype=wp.float32, device=self.sim.device)
 
-        # Initial targets
+        # Initial states and targets
+        if self.cfg.seed is None:
+            self.cfg.seed = -1
+
         wp.launch(
-            initialize_targets,
+            initialize_state,
             dim=self.num_envs,
             inputs=[
                 self.env_origins,
                 self.targets,
+                self.states,
+                self.cfg.seed,
             ],
         )
 
@@ -676,7 +684,7 @@ def _reset_idx(self, mask: wp.array | None = None):
                 self.soft_joint_vel_limits,
                 self.joint_pos,
                 self.joint_vel,
-                WarpRng.state,
+                self.states,
                 self.cfg.initial_joint_pos_range[0],
                 self.cfg.initial_joint_pos_range[1],
                 self.cfg.initial_joint_vel_range[0],
diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
index f98a9e0e3375..c2a34edc9da6 100644
--- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
+++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
@@ -13,12 +13,22 @@
 import warp as wp
 from isaaclab_experimental.envs import DirectRLEnvWarp
 
-from isaaclab.utils.seed import WarpRng
-
 if TYPE_CHECKING:
     from isaaclab_tasks.core.reorient.config.allegro_hand.allegro_hand_direct_env_cfg import AllegroHandEnvCfg
 
 
+@wp.kernel
+def initialize_rng_state(
+    # input
+    seed: wp.int32,
+    # output
+    state: wp.array(dtype=wp.uint32),
+):
+    """Initialize each env's random number generator state from the seed."""
+    env_id = wp.tid()
+    state[env_id] = wp.rand_init(seed, wp.int32(env_id))
+
+
 @wp.kernel
 def apply_actions_to_targets(
     # input
@@ -637,6 +647,20 @@ def __init__(self, cfg: AllegroHandEnvCfg, render_mode: str | None = None, **kwa
         self.object_pose = wp.zeros(self.num_envs, dtype=wp.transformf, device=self.device)
         self.object_vels = wp.zeros(self.num_envs, dtype=wp.spatial_vectorf, device=self.device)
 
+        # RNG state (per-env) for randomizations in reset/goal resets.
+        self.rng_state = wp.zeros(self.num_envs, dtype=wp.uint32, device=self.device)
+        if self.cfg.seed is None:
+            self.cfg.seed = -1
+        wp.launch(
+            initialize_rng_state,
+            dim=self.num_envs,
+            inputs=[
+                self.cfg.seed,
+                self.rng_state,
+            ],
+            device=self.device,
+        )
+
         # ---------------------------------------------------------------------
         # Torch views / aliases
         # ---------------------------------------------------------------------
@@ -786,7 +810,7 @@ def _reset_idx(self, mask: wp.array | None = None):
                 self.x_unit_vec,
                 self.y_unit_vec,
                 mask,
-                WarpRng.state,
+                self.rng_state,
                 self.object.data.root_link_pose_w.warp,
                 self.object.data.root_com_vel_w.warp,
             ],
@@ -806,7 +830,7 @@ def _reset_idx(self, mask: wp.array | None = None):
                 self.cfg.reset_dof_vel_noise,
                 mask,
                 self.num_hand_dofs,
-                WarpRng.state,
+                self.rng_state,
                 self.hand.data.joint_pos.warp,
                 self.hand.data.joint_vel.warp,
                 self.prev_targets,
@@ -850,7 +874,7 @@ def _reset_target_pose(self, env_ids: Sequence[int] | None = None, mask: wp.arra
                 self.y_unit_vec,
                 self.env_origins,
                 self.goal_pos,
-                WarpRng.state,
+                self.rng_state,
                 self.goal_rot,
                 self.reset_goal_buf,
                 self.goal_pos_w,

From 3ddcf7e2a9f445588a147442d53b48825883edd7 Mon Sep 17 00:00:00 2001
From: Jichuan Hu <jichuanh@nvidia.com>
Date: Thu, 22 Oct 2026 00:12:22 -0700
Subject: [PATCH 41/48] Share one per-environment Warp RNG state through
 WarpRng

Environments initialize WarpRng once at construction and configure_seed reseeds it in place. The ray caster, experimental managers, events, noise models and Warp tasks read WarpRng.state instead of keeping their own copies; rng_state_wp is removed.

Re-applies 6ac6061a137 and 9328d981380, split out of the masked-reset pull request.
---
 docs/source/concepts/warp_environments.rst    |  3 +-
 .../jichuanh-warp-rng-shared.added.rst        |  3 +
 .../jichuanh-warp-rng-shared.changed.rst      |  2 +
 .../jichuanh-warp-rng-shared.minor            |  0
 .../sensors/ray_caster/base_ray_caster.py     | 14 +----
 .../isaaclab/sensors/ray_caster/kernels.py    | 19 ------
 source/isaaclab/isaaclab/utils/seed.py        | 60 +++++++++++++++++-
 .../isaaclab/isaaclab/utils/warp/kernels.py   | 25 ++++++++
 source/isaaclab/test/utils/test_warp_rng.py   | 63 +++++++++++++++++++
 .../jichuanh-warp-rng-shared.changed.rst      |  4 ++
 .../jichuanh-warp-rng-shared.major            |  0
 .../jichuanh-warp-rng-shared.removed.rst      |  2 +
 .../envs/direct_rl_env_warp.py                |  9 ++-
 .../envs/manager_based_env_warp.py            | 39 +++---------
 .../isaaclab_experimental/envs/mdp/events.py  | 23 +++----
 .../managers/command_manager.py               |  5 +-
 .../managers/event_manager.py                 | 14 ++---
 .../managers/observation_manager.py           |  6 --
 .../utils/noise/noise_model.py                | 12 ++--
 .../test/envs/mdp/test_events_warp_parity.py  |  5 +-
 .../envs/mdp/test_observations_warp_parity.py |  7 ++-
 .../test/envs/mdp/test_rewards_warp_parity.py |  7 ++-
 .../envs/mdp/test_terminations_warp_parity.py |  7 ++-
 .../changelog.d/jichuanh-warp-rng-shared.skip |  0
 .../sensors/test_newton_raycast_sensor.py     |  4 ++
 .../jichuanh-warp-rng-shared.changed.rst      |  2 +
 .../core/cartpole/cartpole_warp_env.py        | 27 +-------
 .../core/locomotion/locomotion_env_warp.py    | 20 ++----
 .../core/reorient/reorient_warp_env.py        | 34 ++--------
 29 files changed, 235 insertions(+), 181 deletions(-)
 create mode 100644 source/isaaclab/changelog.d/jichuanh-warp-rng-shared.added.rst
 create mode 100644 source/isaaclab/changelog.d/jichuanh-warp-rng-shared.changed.rst
 create mode 100644 source/isaaclab/changelog.d/jichuanh-warp-rng-shared.minor
 create mode 100644 source/isaaclab/test/utils/test_warp_rng.py
 create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-rng-shared.changed.rst
 create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-rng-shared.major
 create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-rng-shared.removed.rst
 create mode 100644 source/isaaclab_newton/changelog.d/jichuanh-warp-rng-shared.skip
 create mode 100644 source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-rng-shared.changed.rst

diff --git a/docs/source/concepts/warp_environments.rst b/docs/source/concepts/warp_environments.rst
index 1debce871a09..8dcba5296a79 100644
--- a/docs/source/concepts/warp_environments.rst
+++ b/docs/source/concepts/warp_environments.rst
@@ -493,7 +493,8 @@ checks the mask to skip non-selected environments:
            return
        # ... modify state for selected envs only
 
-- RNG uses per-env ``env.rng_state_wp`` (``wp.uint32``) instead of ``torch.rand``
+- RNG uses the per-env :attr:`WarpRng.state <isaaclab.utils.seed.WarpRng.state>` (``wp.uint32``) instead of
+  ``torch.rand``
 - **Startup/prestartup** events use the torch convention ``(env, env_ids, **params)``
 - **Reset/interval** events use the warp convention ``(env, env_mask, **params)``
 
diff --git a/source/isaaclab/changelog.d/jichuanh-warp-rng-shared.added.rst b/source/isaaclab/changelog.d/jichuanh-warp-rng-shared.added.rst
new file mode 100644
index 000000000000..a90ea5eafdf7
--- /dev/null
+++ b/source/isaaclab/changelog.d/jichuanh-warp-rng-shared.added.rst
@@ -0,0 +1,3 @@
+* Added :class:`~isaaclab.utils.seed.WarpRng`, the per-environment Warp random number generator state that
+  environments and sensors share, one per process. :func:`~isaaclab.utils.seed.configure_seed` reseeds it in place,
+  and environments call :meth:`~isaaclab.utils.seed.WarpRng.initialize` at construction.
diff --git a/source/isaaclab/changelog.d/jichuanh-warp-rng-shared.changed.rst b/source/isaaclab/changelog.d/jichuanh-warp-rng-shared.changed.rst
new file mode 100644
index 000000000000..2fde38574e92
--- /dev/null
+++ b/source/isaaclab/changelog.d/jichuanh-warp-rng-shared.changed.rst
@@ -0,0 +1,2 @@
+* Changed masked :meth:`~isaaclab.sensors.ray_caster.BaseRayCaster.reset` to draw its drift from the shared
+  :class:`~isaaclab.utils.seed.WarpRng` state instead of random states of its own.
diff --git a/source/isaaclab/changelog.d/jichuanh-warp-rng-shared.minor b/source/isaaclab/changelog.d/jichuanh-warp-rng-shared.minor
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py
index 3c7e034b2167..8c022d0df1e8 100644
--- a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py
+++ b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py
@@ -19,6 +19,7 @@
 from ...markers import VisualizationMarkers
 from ...terrains.trimesh.utils import make_plane
 from ...utils import math as math_utils
+from ...utils.seed import WarpRng
 from ...utils.warp import ProxyArray, convert_to_warp_mesh
 from ...utils.warp.kernels import raycast_mesh_masked_kernel
 from ..sensor_base import SensorBase
@@ -117,7 +118,7 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None
         if not self._drift_sampled:
             return
         if env_mask is not None:
-            # Resample in a kernel from per-environment random states with the ranges as launch
+            # Resample in a kernel from the shared per-environment random states with the ranges as launch
             # arguments, so the masked reset neither synchronizes nor allocates and can be captured.
             # A zero range draws zero, which clears drift that was sampled before it was disabled.
             wp.launch(
@@ -128,7 +129,7 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None
                     wp.vec2f(*self.cfg.drift_range),
                     wp.vec3f(*(low for low, _ in ray_cast_range_list)),
                     wp.vec3f(*(high for _, high in ray_cast_range_list)),
-                    self._drift_rng_state,
+                    WarpRng.state,
                 ],
                 outputs=[self.drift.warp, self.ray_cast_drift.warp],
                 device=self._device,
@@ -261,15 +262,6 @@ def _initialize_rays_impl(self):
         # Drift buffers are warp-first; reset uses explicit .torch views for sampling.
         self.drift = ProxyArray(wp.zeros(self._view_count, dtype=wp.vec3f, device=self._device))
         self.ray_cast_drift = ProxyArray(wp.zeros(self._view_count, dtype=wp.vec3f, device=self._device))
-        # Per-environment random states for masked drift resampling, seeded from torch's global generator.
-        self._drift_rng_state = wp.empty(self._view_count, dtype=wp.uint32, device=self._device)
-        wp.launch(
-            ray_caster_kernels.init_rng_state_kernel,
-            dim=self._view_count,
-            inputs=[int(torch.randint(0, 2**31 - 1, (1,)))],
-            outputs=[self._drift_rng_state],
-            device=self._device,
-        )
 
         # World-frame ray buffers
         self._ray_starts_w = wp.empty((self._view_count, self.num_rays), dtype=wp.vec3f, device=self._device)
diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py b/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py
index 9d36d1a3ddd7..6fa4c65ffcc1 100644
--- a/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py
+++ b/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py
@@ -170,25 +170,6 @@ def apply_z_drift_kernel(
     ray_hits[env, ray] = wp.vec3f(hit[0], hit[1], hit[2] + ray_cast_drift[env][2])
 
 
-@wp.kernel(enable_backward=False)
-def init_rng_state_kernel(
-    # input
-    seed: wp.int32,
-    # output
-    rng_state: wp.array(dtype=wp.uint32),
-):
-    """Initialize one random number generator state per environment.
-
-    Launch with dim=(num_envs,).
-
-    Args:
-        seed: Seed shared by all environments; each environment uses its index as the offset.
-        rng_state: Random number generator states. Shape is (num_envs,).
-    """
-    env = wp.tid()
-    rng_state[env] = wp.rand_init(seed, env)
-
-
 @wp.kernel(enable_backward=False)
 def resample_drift_masked_kernel(
     # input
diff --git a/source/isaaclab/isaaclab/utils/seed.py b/source/isaaclab/isaaclab/utils/seed.py
index 535789305164..6d8f13de85dd 100644
--- a/source/isaaclab/isaaclab/utils/seed.py
+++ b/source/isaaclab/isaaclab/utils/seed.py
@@ -3,6 +3,8 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
+from __future__ import annotations
+
 import os
 import random
 
@@ -11,6 +13,62 @@
 import warp as wp
 
 
+class WarpRng:
+    """Process-wide per-environment random number generator state for Warp kernels.
+
+    Warp has no global generator: :func:`warp.rand_init` hashes a seed and an offset into a state that lives only
+    for one kernel launch. Kernels that draw random numbers per environment therefore read ``state[env]``, draw, and
+    store the advanced state back, so the next launch or CUDA-graph replay draws new numbers. This class holds that
+    array once per process, so environments, managers and sensors share it instead of each allocating a copy.
+
+    :meth:`initialize` must be called once the number of environments and the device are known; environments do this
+    at construction, so users read :attr:`state` directly. :meth:`seed` is optional; without it, the state is seeded
+    from :func:`torch.initial_seed`. Only kernels launched over environments, one thread per environment, may advance
+    the state.
+    """
+
+    state: wp.array | None = None
+    """The per-environment random number generator state, or None before :meth:`initialize`. Shape is (num_envs,)."""
+
+    _seed: int | None = None
+
+    @classmethod
+    def seed(cls, seed: int) -> None:
+        """Set the seed, and reseed the state in place if it is initialized.
+
+        Args:
+            seed: The random seed value.
+        """
+        cls._seed = seed
+        if cls.state is not None:
+            cls._reset_state()
+
+    @classmethod
+    def initialize(cls, num_envs: int, device: str) -> None:
+        """Allocate and seed a new state. Environments call it once at construction.
+
+        Args:
+            num_envs: The number of environments.
+            device: The device of the state.
+        """
+        cls.state = wp.empty(num_envs, dtype=wp.uint32, device=device)
+        cls._reset_state()
+
+    @classmethod
+    def _reset_state(cls) -> None:
+        """Seed every environment's state from the set seed, or from torch's initial seed without one."""
+        from .warp.kernels import initialize_rng_state
+
+        seed = cls._seed if cls._seed is not None else torch.initial_seed()
+        wp.launch(
+            initialize_rng_state,
+            dim=cls.state.shape[0],
+            inputs=[seed % 2**31],
+            outputs=[cls.state],
+            device=cls.state.device,
+        )
+
+
 def configure_seed(seed: int | None, torch_deterministic: bool = False) -> int:
     """Set seed across all random number generators (torch, numpy, random, warp).
 
@@ -31,7 +89,7 @@ def configure_seed(seed: int | None, torch_deterministic: bool = False) -> int:
     os.environ["PYTHONHASHSEED"] = str(seed)
     torch.cuda.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
-    wp.rand_init(seed)
+    WarpRng.seed(seed)
 
     if torch_deterministic:
         # refer to https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility
diff --git a/source/isaaclab/isaaclab/utils/warp/kernels.py b/source/isaaclab/isaaclab/utils/warp/kernels.py
index 3d046abe2bba..a7818691165e 100644
--- a/source/isaaclab/isaaclab/utils/warp/kernels.py
+++ b/source/isaaclab/isaaclab/utils/warp/kernels.py
@@ -780,3 +780,28 @@ def gather_subtract_2d(
 def gather_subtract_2d_kernel(env_ids: "wp.array | torch.Tensor", item_ids: "wp.array | torch.Tensor") -> wp.Kernel:
     """Select the gather-subtract worker for the selector dtypes."""
     return _GATHER_SUBTRACT_2D_DISPATCHER.select(env_ids, item_ids)
+
+
+##
+# Random number generation.
+##
+
+
+@wp.kernel(enable_backward=False)
+def initialize_rng_state(
+    # input
+    seed: wp.int32,
+    # output
+    state: wp.array(dtype=wp.uint32),
+):
+    """Initialize one random number generator state per environment.
+
+    Launch with ``dim=(num_envs,)``. Kernels that draw from ``state[i]`` should store the advanced state back,
+    so the next launch or graph replay draws new numbers.
+
+    Args:
+        seed: Seed shared by all environments; each environment uses its index as the offset.
+        state: Random number generator states. Shape is (num_envs,).
+    """
+    env_id = wp.tid()
+    state[env_id] = wp.rand_init(seed, env_id)
diff --git a/source/isaaclab/test/utils/test_warp_rng.py b/source/isaaclab/test/utils/test_warp_rng.py
new file mode 100644
index 000000000000..2bc3815a2b8c
--- /dev/null
+++ b/source/isaaclab/test/utils/test_warp_rng.py
@@ -0,0 +1,63 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tests for the process-wide per-environment Warp random number generator state."""
+
+import numpy as np
+import pytest
+import torch
+
+from isaaclab.test.utils import test_devices
+from isaaclab.utils.seed import WarpRng, configure_seed
+
+
+@pytest.fixture(autouse=True)
+def fresh_warp_rng():
+    """Start and end every test without a state or a set seed."""
+    WarpRng.state = None
+    WarpRng._seed = None
+    yield
+    WarpRng.state = None
+    WarpRng._seed = None
+
+
+@pytest.mark.parametrize("device", test_devices())
+def test_initialize_allocates_a_new_state(device):
+    """Every :meth:`WarpRng.initialize` call allocates and seeds a new state of the given size."""
+    WarpRng.initialize(4, device)
+    state = WarpRng.state
+    WarpRng.initialize(5, device)
+    assert WarpRng.state is not state
+    assert WarpRng.state.shape == (5,)
+
+
+@pytest.mark.parametrize("device", test_devices())
+def test_configure_seed_reseeds_in_place(device):
+    """configure_seed reseeds the existing state in place and reproducibly."""
+    configure_seed(7)
+    WarpRng.initialize(4, device)
+    state = WarpRng.state
+    seeded = state.numpy().copy()
+
+    state.fill_(0)
+    configure_seed(7)
+    assert WarpRng.state.ptr == state.ptr
+    np.testing.assert_array_equal(state.numpy(), seeded)
+
+    configure_seed(8)
+    assert not np.array_equal(state.numpy(), seeded)
+
+
+@pytest.mark.parametrize("device", test_devices())
+def test_default_seed_follows_torch(device):
+    """Without a set seed, the state is seeded from torch's initial seed."""
+    torch.manual_seed(5)
+    WarpRng.initialize(4, device)
+    default = WarpRng.state.numpy().copy()
+
+    WarpRng.state = None
+    WarpRng.seed(5)
+    WarpRng.initialize(4, device)
+    np.testing.assert_array_equal(WarpRng.state.numpy(), default)
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-rng-shared.changed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-rng-shared.changed.rst
new file mode 100644
index 000000000000..aa642663b908
--- /dev/null
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-rng-shared.changed.rst
@@ -0,0 +1,4 @@
+* Changed :class:`~isaaclab_experimental.envs.ManagerBasedEnvWarp` and
+  :class:`~isaaclab_experimental.envs.DirectRLEnvWarp` to initialize the process-wide
+  :class:`~isaaclab.utils.seed.WarpRng` state, which the Warp MDP terms, managers and noise models now draw from.
+  An environment created without a seed now seeds it from :func:`torch.initial_seed` instead of ``-1``.
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-rng-shared.major b/source/isaaclab_experimental/changelog.d/jichuanh-warp-rng-shared.major
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-rng-shared.removed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-rng-shared.removed.rst
new file mode 100644
index 000000000000..8e486bedf7bd
--- /dev/null
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-rng-shared.removed.rst
@@ -0,0 +1,2 @@
+* **Breaking:** Removed ``ManagerBasedEnvWarp.rng_state_wp``. Read :attr:`~isaaclab.utils.seed.WarpRng.state`
+  instead.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
index 2c032cea3c22..8ef4e8659512 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py
@@ -32,7 +32,7 @@
 from isaaclab.sim.utils import use_stage
 from isaaclab.utils import instantiate, validate
 from isaaclab.utils.noise import NoiseModel
-from isaaclab.utils.seed import configure_seed
+from isaaclab.utils.seed import WarpRng, configure_seed
 from isaaclab.utils.timer import Timer
 
 from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache
@@ -125,7 +125,7 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
 
         # set the seed for the environment
         if self.cfg.seed is not None:
-            self.cfg.seed = self.seed(self.cfg.seed)
+            self.cfg.seed = configure_seed(self.cfg.seed)
         else:
             logger.warning("Seed not set for the environment. The environment creation may not be deterministic.")
 
@@ -166,6 +166,9 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
                 # attach_stage_to_usd_context()
         logger.info(f"Scene manager: {self.scene}")
 
+        # Process-wide per-env Warp RNG state (isaaclab.utils.seed.WarpRng), shared by terms, tasks and sensors.
+        WarpRng.initialize(self.num_envs, self.device)
+
         # create event manager
         # note: this is needed here (rather than after simulation play) to allow USD-related randomization events
         #   that must happen before the simulation starts. Example: randomizing mesh scale
@@ -345,7 +348,7 @@ def reset(self, seed: int | None = None, options: dict[str, Any] | None = None)
         """
         # set the seed
         if seed is not None:
-            self.seed(seed)
+            configure_seed(seed)
 
         # reset state of scene
         self._reset_idx(self._ALL_ENV_MASK)
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
index 8b6710633ae4..d4a825febe43 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py
@@ -36,7 +36,7 @@
 from isaaclab.sim import SimulationContext
 from isaaclab.sim.utils import use_stage
 from isaaclab.utils import validate
-from isaaclab.utils.seed import configure_seed
+from isaaclab.utils.seed import WarpRng, configure_seed
 from isaaclab.utils.timer import Timer
 
 from isaaclab_experimental.envs.interactive_scene_warp import InteractiveSceneWarp as InteractiveScene
@@ -47,17 +47,6 @@
 logger = logging.getLogger(__name__)
 
 
-@wp.kernel
-def initialize_rng_state(
-    # input
-    seed: wp.int32,
-    # output
-    state: wp.array(dtype=wp.uint32),
-):
-    env_id = wp.tid()
-    state[env_id] = wp.rand_init(seed, wp.int32(env_id))
-
-
 class ManagerBasedEnvWarp:
     """The base environment for the manager-based workflow (experimental fork).
 
@@ -97,7 +86,7 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
 
         # set the seed for the environment
         if self.cfg.seed is not None:
-            self.cfg.seed = self.seed(self.cfg.seed)
+            self.cfg.seed = configure_seed(self.cfg.seed)
         else:
             logger.warning("Seed not set for the environment. The environment creation may not be deterministic.")
 
@@ -149,16 +138,8 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
                 # attach_stage_to_usd_context()
         logger.info(f"Scene manager: {self.scene}")
 
-        # Shared per-env Warp RNG state (accessible to all managers/terms via `env`).
-        # This is a single stream per env (no lookup) and is initialized once when `num_envs` is known.
-        self.rng_state_wp = wp.zeros((self.num_envs,), dtype=wp.uint32, device=self.device)
-        seed_val = int(self.cfg.seed) if self.cfg.seed is not None else -1
-        wp.launch(
-            kernel=initialize_rng_state,
-            dim=self.num_envs,
-            inputs=[seed_val, self.rng_state_wp],
-            device=self.device,
-        )
+        # Process-wide per-env Warp RNG state (isaaclab.utils.seed.WarpRng), shared by terms, tasks and sensors.
+        WarpRng.initialize(self.num_envs, self.device)
 
         # TODO(jichuanh): this is problematic as warp capture requires stable pointers,
         #                 using different masks for different managers/terms will cause problems.
@@ -442,16 +423,10 @@ def reset(
 
         # set the seed
         if seed is not None:
-            used_seed = self.seed(seed)
+            # also reseeds the per-env Warp RNG state in place (stable pointer for capture)
+            used_seed = configure_seed(seed)
             # keep cfg seed in sync for downstream users
             self.cfg.seed = used_seed
-            # re-initialize per-env Warp RNG state without reallocating (stable pointer for capture)
-            wp.launch(
-                kernel=initialize_rng_state,
-                dim=self.num_envs,
-                inputs=[int(used_seed), self.rng_state_wp],
-                device=self.device,
-            )
 
         # reset state of scene
         self._reset_idx(env_ids)
@@ -506,7 +481,7 @@ def reset_to(
 
         # set the seed
         if seed is not None:
-            self.seed(seed)
+            configure_seed(seed)
 
         self._reset_idx(env_ids)
 
diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py
index 90e5696a6b10..a64af1d309ce 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py
@@ -31,6 +31,7 @@
 
 from isaaclab.envs.mdp.events import randomize_rigid_body_mass as _StableRandomizeRigidBodyMass
 from isaaclab.envs.mdp.events import randomize_rigid_body_material as _StableRandomizeRigidBodyMaterial
+from isaaclab.utils.seed import WarpRng
 
 from isaaclab_experimental.managers import ManagerTermBase as _WarpManagerTermBase
 from isaaclab_experimental.managers import SceneEntityCfg
@@ -98,7 +99,7 @@ def randomize_rigid_body_com(
         dim=env.num_envs,
         inputs=[
             env_mask,
-            env.rng_state_wp,
+            WarpRng.state,
             asset.data.body_com_pos_b.warp,
             asset_cfg.body_ids_wp,
             fn._com_lo,
@@ -178,7 +179,7 @@ def apply_external_force_torque(
         dim=env.num_envs,
         inputs=[
             env_mask,
-            env.rng_state_wp,
+            WarpRng.state,
             apply_external_force_torque._scratch_forces,
             apply_external_force_torque._scratch_torques,
             force_range[0],
@@ -258,7 +259,7 @@ def push_by_setting_velocity(
         dim=env.num_envs,
         inputs=[
             env_mask,
-            env.rng_state_wp,
+            WarpRng.state,
             asset.data.root_vel_w.warp,
             push_by_setting_velocity._scratch_vel,
             push_by_setting_velocity._lin_lo,
@@ -377,7 +378,7 @@ def reset_root_state_uniform(
         dim=env.num_envs,
         inputs=[
             env_mask,
-            env.rng_state_wp,
+            WarpRng.state,
             asset.data.default_root_pose.warp,
             asset.data.default_root_vel.warp,
             env.env_origins_wp,
@@ -470,11 +471,6 @@ def reset_joints_by_offset(
             f"but got None for asset '{asset_cfg.name}'. "
             "Use isaaclab_experimental.managers.SceneEntityCfg and ensure joint_names are set."
         )
-    if not hasattr(env, "rng_state_wp") or env.rng_state_wp is None:
-        raise AttributeError(
-            "reset_joints_by_offset requires env.rng_state_wp to be initialized. "
-            "Use ManagerBasedEnvWarp or ManagerBasedRLEnvWarp as the base environment."
-        )
 
     wp.launch(
         kernel=_reset_joints_by_offset_kernel,
@@ -482,7 +478,7 @@ def reset_joints_by_offset(
         inputs=[
             env_mask,
             asset_cfg.joint_ids_wp,
-            env.rng_state_wp,
+            WarpRng.state,
             asset.data.default_joint_pos.warp,
             asset.data.default_joint_vel.warp,
             asset.data.joint_pos.warp,
@@ -561,11 +557,6 @@ def reset_joints_by_scale(
             f"but got None for asset '{asset_cfg.name}'. "
             "Use isaaclab_experimental.managers.SceneEntityCfg and ensure joint_names are set."
         )
-    if not hasattr(env, "rng_state_wp") or env.rng_state_wp is None:
-        raise AttributeError(
-            "reset_joints_by_scale requires env.rng_state_wp to be initialized. "
-            "Use ManagerBasedEnvWarp or ManagerBasedRLEnvWarp as the base environment."
-        )
 
     wp.launch(
         kernel=_reset_joints_by_scale_kernel,
@@ -573,7 +564,7 @@ def reset_joints_by_scale(
         inputs=[
             env_mask,
             asset_cfg.joint_ids_wp,
-            env.rng_state_wp,
+            WarpRng.state,
             asset.data.default_joint_pos.warp,
             asset.data.default_joint_vel.warp,
             asset.data.joint_pos.warp,
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py
index e7d8b4b746b0..2b4cd8de1828 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py
@@ -18,6 +18,7 @@
 
 from isaaclab.managers.manager_term_cfg import CommandTermCfg
 from isaaclab.utils import instantiate
+from isaaclab.utils.seed import WarpRng
 
 from isaaclab_experimental.utils.warp.kernels import compute_reset_scale, count_masked
 
@@ -306,8 +307,6 @@ def _resample(self, env_mask: wp.array):
             raise TypeError(f"env_mask must be a wp.array (got {type(env_mask)}).")
         if env_mask.dtype != wp.bool or env_mask.ndim != 1:
             raise TypeError(f"env_mask must be wp.bool 1D (got dtype={env_mask.dtype}, ndim={env_mask.ndim}).")
-        if self._env.rng_state_wp is None:
-            raise RuntimeError("Environment rng_state_wp is not initialized.")
 
         # resample time-left and increment command-counter for masked envs
         wp.launch(
@@ -317,7 +316,7 @@ def _resample(self, env_mask: wp.array):
                 env_mask,
                 self.time_left_wp,
                 self.command_counter_wp,
-                self._env.rng_state_wp,
+                WarpRng.state,
                 float(self.cfg.resampling_time_range[0]),
                 float(self.cfg.resampling_time_range[1]),
             ],
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
index 5c3832fb803b..580e5062d4c1 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py
@@ -30,6 +30,7 @@
 from prettytable import PrettyTable
 
 from isaaclab.managers.manager_term_cfg import EventTermCfg
+from isaaclab.utils.seed import WarpRng
 
 from .manager_base import ManagerBase
 
@@ -273,7 +274,7 @@ def reset(
                     inputs=[
                         env_mask,
                         self._interval_term_time_left_wp[i],
-                        self._env.rng_state_wp,
+                        WarpRng.state,
                         float(lower),
                         float(upper),
                     ],
@@ -331,9 +332,6 @@ def apply(
             term_cfg.func(self._env, env_mask_wp, **term_cfg.params)
 
     def _apply_interval(self, dt: float) -> None:
-        if self._env.rng_state_wp is None:
-            raise RuntimeError("EventManager._apply_interval: env.rng_state_wp is not initialized.")
-
         # iterate over all the interval terms (fixed list; captured graph-friendly)
         for i, term_cfg in enumerate(self._mode_term_cfgs["interval"]):
             lower, upper = self._interval_term_ranges[i]
@@ -363,7 +361,7 @@ def _apply_interval(self, dt: float) -> None:
                     dim=self.num_envs,
                     inputs=[
                         self._interval_term_time_left_wp[i],
-                        self._env.rng_state_wp,
+                        WarpRng.state,
                         self._scratch_term_mask_wp,
                         float(dt),
                         float(lower),
@@ -450,13 +448,11 @@ def _prepare_terms(self):
                 if term_cfg.is_global_time:
                     # allocate and seed scalar global RNG state if needed (avoid consuming env0 RNG stream)
                     if self._interval_global_rng_state_wp is None:
-                        if self._env.rng_state_wp is None:
-                            raise RuntimeError("EventManager._prepare_terms: env.rng_state_wp is not initialized.")
                         self._interval_global_rng_state_wp = wp.zeros((1,), dtype=wp.uint32, device=self.device)
                         wp.launch(
                             kernel=_seed_global_rng_from_env_rng,
                             dim=1,
-                            inputs=[self._env.rng_state_wp, self._interval_global_rng_state_wp],
+                            inputs=[WarpRng.state, self._interval_global_rng_state_wp],
                             device=self.device,
                         )
                     time_left = wp.zeros((1,), dtype=wp.float32, device=self.device)
@@ -473,7 +469,7 @@ def _prepare_terms(self):
                     wp.launch(
                         kernel=_interval_init_per_env,
                         dim=self.num_envs,
-                        inputs=[time_left, self._env.rng_state_wp, float(lower), float(upper)],
+                        inputs=[time_left, WarpRng.state, float(lower), float(upper)],
                         device=self.device,
                     )
                     self._interval_term_time_left_wp.append(time_left)
diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
index 051aaa8a6ff6..6ec29aa0e921 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py
@@ -780,14 +780,8 @@ def _prepare_terms(self):  # noqa: C901
                                     f" and optional parameters: {args_with_defaults}, but received: {term_params}."
                                 )
 
-                # plumb the shared per-env RNG state so Warp noise kernels can consume it
-                # (function-style NoiseCfg kernels read it off the cfg at call time)
-                if term_cfg.noise is not None and isinstance(term_cfg.noise, noise.NoiseCfg):
-                    term_cfg.noise.rng_state_wp = self._env.rng_state_wp
                 # prepare noise model classes
                 if term_cfg.noise is not None and isinstance(term_cfg.noise, noise.NoiseModelCfg):
-                    # plumb the shared per-env RNG state so Warp noise kernels can consume it
-                    term_cfg.noise.rng_state_wp = self._env.rng_state_wp
                     noise_model_cls = term_cfg.noise.class_type
                     if not issubclass(noise_model_cls, noise.NoiseModel):
                         raise TypeError(
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/noise/noise_model.py b/source/isaaclab_experimental/isaaclab_experimental/utils/noise/noise_model.py
index acba926ea488..cbc933324830 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/utils/noise/noise_model.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/utils/noise/noise_model.py
@@ -10,10 +10,8 @@
 
     noise_cfg.func(data_wp, noise_cfg) -> None
 
-Random noise kernels (gaussian, uniform) consume the shared per-env Warp RNG state
-(``rng_state_wp``) that is set on the config at manager prep time from
-``env.rng_state_wp``.  See :func:`initialize_rng_state` in
-``isaaclab_experimental.envs.manager_based_env_warp`` for the initialization pattern.
+Random noise kernels (gaussian, uniform) draw from the process-wide per-env Warp RNG state
+:attr:`isaaclab.utils.seed.WarpRng.state`.
 """
 
 from __future__ import annotations
@@ -22,6 +20,8 @@
 
 import warp as wp
 
+from isaaclab.utils.seed import WarpRng
+
 if TYPE_CHECKING:
     from . import noise_cfg
 
@@ -108,7 +108,7 @@ def uniform_noise(data: wp.array, cfg: noise_cfg.UniformNoiseCfg) -> None:
     wp.launch(
         _apply_uniform_noise,
         dim=data.shape[0],
-        inputs=[data, cfg.rng_state_wp, float(cfg.n_min), float(cfg.n_max), _OPERATION_MAP[cfg.operation]],
+        inputs=[data, WarpRng.state, float(cfg.n_min), float(cfg.n_max), _OPERATION_MAP[cfg.operation]],
         device=data.device,
     )
 
@@ -149,7 +149,7 @@ def gaussian_noise(data: wp.array, cfg: noise_cfg.GaussianNoiseCfg) -> None:
     wp.launch(
         _apply_gaussian_noise,
         dim=data.shape[0],
-        inputs=[data, cfg.rng_state_wp, float(cfg.mean), float(cfg.std), _OPERATION_MAP[cfg.operation]],
+        inputs=[data, WarpRng.state, float(cfg.mean), float(cfg.std), _OPERATION_MAP[cfg.operation]],
         device=data.device,
     )
 
diff --git a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py
index e93e73267938..e910f2d9fc1a 100644
--- a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py
+++ b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py
@@ -12,6 +12,8 @@
 import torch
 import warp as wp
 
+from isaaclab.utils.seed import WarpRng
+
 # Skip entire module if no CUDA device available
 wp.init()
 pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required")
@@ -44,6 +46,7 @@ def _clear_function_caches():
     between tests to avoid stale references from prior fixtures.
     """
     yield
+    WarpRng.state = None
     for fn in (
         warp_evt.push_by_setting_velocity,
         warp_evt.apply_external_force_torque,
@@ -102,7 +105,7 @@ class _Env:
     env.step_dt = 0.02
     env.max_episode_length_s = 10.0
     # RNG state for events (seeded deterministically)
-    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
diff --git a/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py
index 0f014cc2b3fd..91a5f8fe138f 100644
--- a/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py
+++ b/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py
@@ -12,6 +12,8 @@
 import torch
 import warp as wp
 
+from isaaclab.utils.seed import WarpRng
+
 # Skip entire module if no CUDA device available
 wp.init()
 pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required")
@@ -50,6 +52,7 @@
 @pytest.fixture(autouse=True)
 def _clear_caches():
     yield
+    WarpRng.state = None
     for fn in [warp_obs.generated_commands]:
         for attr in list(vars(fn)):
             if attr.startswith("_"):
@@ -130,7 +133,7 @@ class _Env:
     env.episode_length_buf = episode_length_buf
     env.step_dt = 0.02
     env.max_episode_length_s = 10.0
-    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
@@ -169,7 +172,7 @@ class _Env:
     env.step_dt = 0.02
     env.max_episode_length = 500
     env.max_episode_length_s = 10.0
-    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
diff --git a/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py
index c81f6f73ec89..9f2d39253d4e 100644
--- a/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py
+++ b/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py
@@ -12,6 +12,8 @@
 import torch
 import warp as wp
 
+from isaaclab.utils.seed import WarpRng
+
 # Skip entire module if no CUDA device available
 wp.init()
 pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required")
@@ -58,6 +60,7 @@
 @pytest.fixture(autouse=True)
 def _clear_caches():
     yield
+    WarpRng.state = None
     for fn in [warp_rew.track_lin_vel_xy_exp, warp_rew.track_ang_vel_z_exp, warp_rew.undesired_contacts]:
         for attr in list(vars(fn)):
             if attr.startswith("_"):
@@ -144,7 +147,7 @@ class _Env:
     env.episode_length_buf = episode_length_buf
     env.step_dt = 0.02
     env.max_episode_length_s = 10.0
-    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
@@ -184,7 +187,7 @@ class _Env:
     env.step_dt = 0.02
     env.max_episode_length = 500
     env.max_episode_length_s = 10.0
-    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
diff --git a/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py
index 9676f2b06044..e0fbee333530 100644
--- a/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py
+++ b/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py
@@ -14,6 +14,8 @@
 import torch
 import warp as wp
 
+from isaaclab.utils.seed import WarpRng
+
 # Skip entire module if no CUDA device available
 wp.init()
 pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required")
@@ -59,6 +61,7 @@
 @pytest.fixture(autouse=True)
 def _clear_caches():
     yield
+    WarpRng.state = None
     for fn in [warp_term.illegal_contact]:
         for attr in list(vars(fn)):
             if attr.startswith("_"):
@@ -139,7 +142,7 @@ class _Env:
     env.episode_length_buf = episode_length_buf
     env.step_dt = 0.02
     env.max_episode_length_s = 10.0
-    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
@@ -179,7 +182,7 @@ class _Env:
     env.step_dt = 0.02
     env.max_episode_length = 500
     env.max_episode_length_s = 10.0
-    env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
+    WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE)
     return env
 
 
diff --git a/source/isaaclab_newton/changelog.d/jichuanh-warp-rng-shared.skip b/source/isaaclab_newton/changelog.d/jichuanh-warp-rng-shared.skip
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py b/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py
index 9c454dd72cca..8cf45b728922 100644
--- a/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py
+++ b/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py
@@ -40,6 +40,7 @@
 from isaaclab.sim import SimulationCfg
 from isaaclab.terrains import TerrainImporterCfg
 from isaaclab.utils import configclass
+from isaaclab.utils.seed import WarpRng
 
 SENSOR_HEIGHT = 2.0
 RAY_OFFSET = 0.2
@@ -222,6 +223,9 @@ def test_partial_reset_resamples_drift_of_selected_envs_only(sim):
     sensor: NewtonRaycastSensor = scene["raycast"]
     sensor.reset()
 
+    # A masked reset draws from the process-wide Warp RNG state, which an environment initializes.
+    WarpRng.initialize(scene.num_envs, sensor.device)
+
     drift = sensor.drift.torch.clone()
     ray_cast_drift = sensor.ray_cast_drift.torch.clone()
     sensor.reset(env_mask=wp.array([False, True, False], dtype=wp.bool, device=sensor.device))
diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-rng-shared.changed.rst b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-rng-shared.changed.rst
new file mode 100644
index 000000000000..59b2cee59f4f
--- /dev/null
+++ b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-rng-shared.changed.rst
@@ -0,0 +1,2 @@
+* Changed the Cartpole, locomotion and reorient Warp environments to draw from the process-wide
+  :class:`~isaaclab.utils.seed.WarpRng` state instead of allocating and seeding their own.
diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py
index 2a3ccdd47f0d..dcafd43d57ec 100644
--- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py
+++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py
@@ -11,6 +11,8 @@
 from isaaclab_experimental.envs import DirectRLEnvWarp
 from isaaclab_experimental.utils.warp.utils import wrap_to_pi
 
+from isaaclab.utils.seed import WarpRng
+
 if TYPE_CHECKING:
     from isaaclab_tasks.core.cartpole.cartpole_direct_env_cfg import CartpoleEnvCfg
 
@@ -174,16 +176,6 @@ def reset(
         state[env_index] = rng_state
 
 
-@wp.kernel
-def initialize_state(
-    state: wp.array(dtype=wp.uint32),
-    seed: wp.int32,
-):
-    """Initialize each env's random number generator state from the seed."""
-    env_index = wp.tid()
-    state[env_index] = wp.rand_init(seed, env_index)
-
-
 class CartpoleWarpEnv(DirectRLEnvWarp):
     cfg: CartpoleEnvCfg
 
@@ -206,19 +198,6 @@ def __init__(self, cfg: CartpoleEnvCfg, render_mode: str | None = None, **kwargs
         self.observations = wp.zeros((self.num_envs), dtype=wp.vec4f, device=self.device)
         self.actions = wp.zeros((self.num_envs, self.cartpole.num_joints), dtype=wp.float32, device=self.device)
         self.rewards = wp.zeros((self.num_envs), dtype=wp.float32, device=self.device)
-        self.states = wp.zeros((self.num_envs), dtype=wp.uint32, device=self.device)
-
-        if self.cfg.seed is None:
-            self.cfg.seed = -1
-
-        wp.launch(
-            initialize_state,
-            dim=self.num_envs,
-            inputs=[
-                self.states,
-                self.cfg.seed,
-            ],
-        )
 
         # Bind torch buffers to warp buffers
         self.torch_obs_buf = wp.to_torch(self.observations)
@@ -317,6 +296,6 @@ def _reset_idx(self, mask: wp.array | None = None) -> None:
                 self.cfg.initial_pole_angle_range,
                 self.cfg.initial_pole_velocity_range,
                 mask,
-                self.states,
+                WarpRng.state,
             ],
         )
diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py
index dfcfad6e0f55..b7852ac644e0 100644
--- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py
+++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py
@@ -9,6 +9,7 @@
 from isaaclab_experimental.envs import DirectRLEnvWarp
 
 from isaaclab.envs import DirectRLEnvCfg
+from isaaclab.utils.seed import WarpRng
 from isaaclab.utils.string import resolve_matching_names_values
 
 
@@ -416,15 +417,12 @@ def survival_rate(
 
 
 @wp.kernel
-def initialize_state(
+def initialize_targets(
     env_origins: wp.array(dtype=wp.vec3f),
     targets: wp.array(dtype=wp.vec3f),
-    state: wp.array(dtype=wp.uint32),
-    seed: wp.int32,
 ):
-    """Initialize each env's RNG state and place its target far along +x from the env origin."""
+    """Place each env's target far along +x from the env origin."""
     env_index = wp.tid()
-    state[env_index] = wp.rand_init(seed, env_index)
     targets[env_index] = env_origins[env_index]
     targets[env_index] += wp.static(wp.vec3f(1000.0, 0.0, 0.0))
 
@@ -480,7 +478,6 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
         )
         self.rewards = wp.zeros((self.num_envs), dtype=wp.float32, device=self.sim.device)
         self.actions = wp.zeros((self.num_envs, self.robot.num_joints), dtype=wp.float32, device=self.sim.device)
-        self.states = wp.zeros((self.num_envs), dtype=wp.uint32, device=self.sim.device)
         self.potentials = wp.zeros(self.num_envs, dtype=wp.float32, device=self.sim.device)
         self.prev_potentials = wp.zeros_like(self.potentials)
         self.targets = wp.zeros((self.num_envs), dtype=wp.vec3f, device=self.sim.device)
@@ -496,18 +493,13 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
         self.env_origins = wp.from_torch(self.scene.env_origins, dtype=wp.vec3f)
         self.actions_mapped = wp.zeros((self.num_envs, self.robot.num_joints), dtype=wp.float32, device=self.sim.device)
 
-        # Initial states and targets
-        if self.cfg.seed is None:
-            self.cfg.seed = -1
-
+        # Initial targets
         wp.launch(
-            initialize_state,
+            initialize_targets,
             dim=self.num_envs,
             inputs=[
                 self.env_origins,
                 self.targets,
-                self.states,
-                self.cfg.seed,
             ],
         )
 
@@ -684,7 +676,7 @@ def _reset_idx(self, mask: wp.array | None = None):
                 self.soft_joint_vel_limits,
                 self.joint_pos,
                 self.joint_vel,
-                self.states,
+                WarpRng.state,
                 self.cfg.initial_joint_pos_range[0],
                 self.cfg.initial_joint_pos_range[1],
                 self.cfg.initial_joint_vel_range[0],
diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
index c2a34edc9da6..f98a9e0e3375 100644
--- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
+++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py
@@ -13,22 +13,12 @@
 import warp as wp
 from isaaclab_experimental.envs import DirectRLEnvWarp
 
+from isaaclab.utils.seed import WarpRng
+
 if TYPE_CHECKING:
     from isaaclab_tasks.core.reorient.config.allegro_hand.allegro_hand_direct_env_cfg import AllegroHandEnvCfg
 
 
-@wp.kernel
-def initialize_rng_state(
-    # input
-    seed: wp.int32,
-    # output
-    state: wp.array(dtype=wp.uint32),
-):
-    """Initialize each env's random number generator state from the seed."""
-    env_id = wp.tid()
-    state[env_id] = wp.rand_init(seed, wp.int32(env_id))
-
-
 @wp.kernel
 def apply_actions_to_targets(
     # input
@@ -647,20 +637,6 @@ def __init__(self, cfg: AllegroHandEnvCfg, render_mode: str | None = None, **kwa
         self.object_pose = wp.zeros(self.num_envs, dtype=wp.transformf, device=self.device)
         self.object_vels = wp.zeros(self.num_envs, dtype=wp.spatial_vectorf, device=self.device)
 
-        # RNG state (per-env) for randomizations in reset/goal resets.
-        self.rng_state = wp.zeros(self.num_envs, dtype=wp.uint32, device=self.device)
-        if self.cfg.seed is None:
-            self.cfg.seed = -1
-        wp.launch(
-            initialize_rng_state,
-            dim=self.num_envs,
-            inputs=[
-                self.cfg.seed,
-                self.rng_state,
-            ],
-            device=self.device,
-        )
-
         # ---------------------------------------------------------------------
         # Torch views / aliases
         # ---------------------------------------------------------------------
@@ -810,7 +786,7 @@ def _reset_idx(self, mask: wp.array | None = None):
                 self.x_unit_vec,
                 self.y_unit_vec,
                 mask,
-                self.rng_state,
+                WarpRng.state,
                 self.object.data.root_link_pose_w.warp,
                 self.object.data.root_com_vel_w.warp,
             ],
@@ -830,7 +806,7 @@ def _reset_idx(self, mask: wp.array | None = None):
                 self.cfg.reset_dof_vel_noise,
                 mask,
                 self.num_hand_dofs,
-                self.rng_state,
+                WarpRng.state,
                 self.hand.data.joint_pos.warp,
                 self.hand.data.joint_vel.warp,
                 self.prev_targets,
@@ -874,7 +850,7 @@ def _reset_target_pose(self, env_ids: Sequence[int] | None = None, mask: wp.arra
                 self.y_unit_vec,
                 self.env_origins,
                 self.goal_pos,
-                self.rng_state,
+                WarpRng.state,
                 self.goal_rot,
                 self.reset_goal_buf,
                 self.goal_pos_w,

From 5676dd1c78534c484a467076fe5c3fb923596bc5 Mon Sep 17 00:00:00 2001
From: Jichuan Hu <jichuanh@nvidia.com>
Date: Mon, 26 Oct 2026 02:48:16 -0700
Subject: [PATCH 42/48] Add resolve_ids next to resolve_1d_mask as the shared
 id and mask helpers

resolve_ids resolves ids or a boolean mask to ids on a device, as resolve_1d_mask does to a mask. resolve_1d_mask takes an optional scratch mask, copies a Warp mask on another device, and accepts a boolean torch mask on its device inside a graph capture. The experimental package now re-exports it instead of keeping a copy.
---
 .../changelog.d/jichuanh-id-mask.added.rst    |   2 +
 .../changelog.d/jichuanh-id-mask.changed.rst  |   4 +
 .../changelog.d/jichuanh-id-mask.minor        |   0
 source/isaaclab/isaaclab/utils/warp/utils.py  | 119 +++++++++++++-
 .../test/utils/warp/test_resolve_selection.py | 155 ++++++++++++++++++
 .../changelog.d/jichuanh-id-mask.changed.rst  |   2 +
 .../utils/warp/__init__.py                    |   3 +-
 .../isaaclab_experimental/utils/warp/utils.py |  88 ----------
 8 files changed, 279 insertions(+), 94 deletions(-)
 create mode 100644 source/isaaclab/changelog.d/jichuanh-id-mask.added.rst
 create mode 100644 source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst
 create mode 100644 source/isaaclab/changelog.d/jichuanh-id-mask.minor
 create mode 100644 source/isaaclab/test/utils/warp/test_resolve_selection.py
 create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-id-mask.changed.rst

diff --git a/source/isaaclab/changelog.d/jichuanh-id-mask.added.rst b/source/isaaclab/changelog.d/jichuanh-id-mask.added.rst
new file mode 100644
index 000000000000..935dc70f7fa4
--- /dev/null
+++ b/source/isaaclab/changelog.d/jichuanh-id-mask.added.rst
@@ -0,0 +1,2 @@
+* Added :func:`~isaaclab.utils.warp.utils.resolve_ids`, which resolves an index selection given as ids or as a
+  boolean mask to ids on a device, as :func:`~isaaclab.utils.warp.utils.resolve_1d_mask` does to a mask.
diff --git a/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst
new file mode 100644
index 000000000000..6a1b73d77c7e
--- /dev/null
+++ b/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst
@@ -0,0 +1,4 @@
+* Changed :func:`~isaaclab.utils.warp.utils.resolve_1d_mask` to make ``scratch_mask`` optional for callers that never
+  pass ids. Ids without a ``scratch_mask`` raise a :class:`ValueError`.
+* Changed :func:`~isaaclab.utils.warp.utils.resolve_1d_mask` to copy a Warp mask on another device to ``device``, and
+  to accept a boolean torch mask on ``device`` inside a CUDA graph capture.
diff --git a/source/isaaclab/changelog.d/jichuanh-id-mask.minor b/source/isaaclab/changelog.d/jichuanh-id-mask.minor
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/source/isaaclab/isaaclab/utils/warp/utils.py b/source/isaaclab/isaaclab/utils/warp/utils.py
index afd31e9c1b38..24e23707615c 100644
--- a/source/isaaclab/isaaclab/utils/warp/utils.py
+++ b/source/isaaclab/isaaclab/utils/warp/utils.py
@@ -11,6 +11,8 @@
 import torch
 import warp as wp
 
+from .proxy_array import ProxyArray
+
 ##
 # Mask resolution - ids/mask to warp boolean mask.
 ##
@@ -30,7 +32,7 @@ def resolve_1d_mask(
     ids: Sequence[int] | slice | torch.Tensor | wp.array | None = None,
     mask: wp.array | torch.Tensor | None = None,
     all_mask: wp.array,
-    scratch_mask: wp.array,
+    scratch_mask: wp.array | None = None,
     device: str,
 ) -> wp.array:
     """Resolve ids/mask into a warp boolean mask.
@@ -41,30 +43,36 @@ def resolve_1d_mask(
     Args:
         ids: Index ids. Accepts ``Sequence[int]``, ``slice``, ``torch.Tensor``,
             ``wp.array(dtype=wp.int32)``, or ``None`` (all elements).
-        mask: Direct boolean mask. ``wp.array`` is returned as-is;
+        mask: Direct boolean mask. ``wp.array`` is returned as-is, or copied when on another device;
             ``torch.Tensor`` is converted.
         all_mask: Pre-allocated all-True mask returned when both *ids* and *mask*
             are ``None``.
         scratch_mask: Pre-allocated scratch buffer populated in-place when *ids*
-            are provided. Not re-entrant (shared buffer).
+            are provided. Not re-entrant (shared buffer). Defaults to None, for callers
+            that never pass *ids*.
         device: Warp device string (e.g. ``"cuda:0"``).
 
     Returns:
         A ``wp.array(dtype=wp.bool)`` mask.
+
+    Raises:
+        ValueError: If *ids* select a subset of elements and no ``scratch_mask`` is given.
     """
     # Normalize slice(None) to None so the capture guard treats it identically to ids=None.
     if isinstance(ids, slice) and ids == slice(None):
         ids = None
 
     if wp.get_device().is_capturing:
-        if ids is not None or (mask is not None and not isinstance(mask, wp.array)):
+        if ids is not None or (mask is not None and _mask_needs_copy(mask, device)):
             raise RuntimeError(
-                "resolve_1d_mask is only capturable when mask is a wp.array or both ids and mask are None."
+                "resolve_1d_mask is only capturable when ids are None and mask is None or a boolean mask on device."
             )
 
     # --- Direct mask input ---
     if mask is not None:
         if isinstance(mask, wp.array):
+            if mask.device != wp.get_device(device):
+                mask = wp.clone(mask, device=device)
             return mask
         if isinstance(mask, torch.Tensor):
             if mask.dtype != torch.bool:
@@ -78,6 +86,9 @@ def resolve_1d_mask(
     if ids is None:
         return all_mask
 
+    if scratch_mask is None:
+        raise ValueError("resolve_1d_mask needs a scratch_mask to write ids into.")
+
     if isinstance(ids, slice):
         scratch_mask.zero_()
         wp.to_torch(scratch_mask)[ids] = True
@@ -115,6 +126,104 @@ def resolve_1d_mask(
     return scratch_mask
 
 
+def _mask_needs_copy(mask: wp.array | torch.Tensor, device: str) -> bool:
+    """Whether :func:`resolve_1d_mask` casts or copies ``mask``, which a graph capture cannot record."""
+    if isinstance(mask, wp.array):
+        return mask.device != wp.get_device(device)
+    if isinstance(mask, torch.Tensor):
+        return mask.dtype != torch.bool or str(mask.device) != device
+    return True
+
+
+##
+# Id resolution - ids/mask to ids.
+##
+
+
+def resolve_ids(
+    *,
+    ids: Sequence[int] | slice | torch.Tensor | wp.array | None = None,
+    mask: wp.array | torch.Tensor | None = None,
+    all_ids: wp.array,
+    device: str,
+    dtype: type | None = None,
+) -> wp.array | torch.Tensor:
+    """Resolve ids/mask into ids on ``device``. The mask takes precedence over the ids.
+
+    Args:
+        ids: Index ids. Accepts ``Sequence[int]``, ``slice``, ``torch.Tensor``, ``wp.array``,
+            or ``None`` (all elements). Tensors and warp arrays keep their integer dtype
+            unless ``dtype`` is given. Pass ids from ``find_*(as_proxy=True)`` as ``.warp`` or ``.torch``.
+        mask: Boolean mask. Its ids are found with :func:`torch.nonzero`, which waits for
+            the device.
+        all_ids: Pre-allocated ids of every element, returned as-is when *ids* and *mask*
+            select every element.
+        device: Warp device string (e.g. ``"cuda:0"``).
+        dtype: Warp integer dtype the ids must have, for callers whose kernels take one
+            index dtype. The ids are then a contiguous 1-D ``wp.array`` of this dtype.
+            Defaults to None: a mask resolves to an int64 ``torch.Tensor``, and ids keep
+            their type and dtype.
+
+    Returns:
+        ``all_ids``, a ``torch.Tensor`` or a ``wp.array`` of ids on ``device``.
+
+    Raises:
+        TypeError: If *ids* is a :class:`ProxyArray`, which is output-only, *mask* is
+            neither a ``wp.array`` nor a ``torch.Tensor``, or ``dtype`` is given and the
+            ids are not 32-bit or 64-bit integers.
+        ValueError: If ``dtype`` is given and the ids are not one-dimensional.
+    """
+    if mask is not None:
+        if isinstance(mask, wp.array):
+            mask = wp.to_torch(mask)
+        elif not isinstance(mask, torch.Tensor):
+            raise TypeError(f"Unsupported mask type: {type(mask)}")
+        ids = torch.nonzero(mask.to(wp.device_to_torch(device))).squeeze(-1)
+    else:
+        ids = _ids_on_device(ids, all_ids, device)
+    return ids if dtype is None else _ids_as_dtype(ids, dtype)
+
+
+def _ids_on_device(
+    ids: Sequence[int] | slice | torch.Tensor | wp.array | None, all_ids: wp.array, device: str
+) -> wp.array | torch.Tensor:
+    """Ids on ``device``, keeping the dtype of tensors and warp arrays; lists and slices become ``wp.int32``."""
+    if isinstance(ids, torch.Tensor):
+        return ids if str(ids.device) == str(wp.get_device(device)) else ids.to(wp.device_to_torch(device))
+    if ids is None:
+        return all_ids
+    if isinstance(ids, wp.array):
+        return ids if ids.device == wp.get_device(device) else ids.to(device)
+    if isinstance(ids, ProxyArray):
+        raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
+    if isinstance(ids, slice):
+        if ids == slice(None):
+            return all_ids
+        start, stop, step = ids.indices(all_ids.shape[0])
+        if step > 0:
+            # a strided view of the caller's all-ids, no copy
+            return wp.from_torch(wp.to_torch(all_ids)[start:stop:step])
+        return wp.array(list(range(start, stop, step)), dtype=wp.int32, device=device)
+    return wp.array(list(ids), dtype=wp.int32, device=device)
+
+
+def _ids_as_dtype(ids: wp.array | torch.Tensor, dtype: type) -> wp.array:
+    """``ids`` as a contiguous 1-D ``wp.array`` of ``dtype`` on their device."""
+    if dtype not in (wp.int32, wp.int64):
+        raise TypeError(f"Unsupported id dtype: {dtype}. Expected wp.int32 or wp.int64.")
+    if ids.ndim != 1:
+        raise ValueError(f"Ids must be a one-dimensional array, got {ids.ndim} dimensions.")
+    if isinstance(ids, wp.array):
+        if ids.dtype == dtype and ids.is_contiguous:
+            return ids
+        if ids.dtype not in (wp.int32, wp.int64):
+            raise TypeError(f"Warp ids must be 32-bit or 64-bit signed integers, got {ids.dtype}.")
+        ids = wp.to_torch(ids)
+    elif ids.dtype.is_floating_point or ids.dtype.is_complex or ids.dtype == torch.bool:
+        raise TypeError(f"Ids must be integers, got {ids.dtype}.")
+    return wp.from_torch(ids.to(wp.dtype_to_torch(dtype)).contiguous(), dtype=dtype)
+
+
 ##
 # Capture safety — property guard.
 ##
diff --git a/source/isaaclab/test/utils/warp/test_resolve_selection.py b/source/isaaclab/test/utils/warp/test_resolve_selection.py
new file mode 100644
index 000000000000..fc951df8f614
--- /dev/null
+++ b/source/isaaclab/test/utils/warp/test_resolve_selection.py
@@ -0,0 +1,155 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tests for resolving index selections with :func:`resolve_ids` and :func:`resolve_1d_mask`."""
+
+import pytest
+import torch
+import warp as wp
+
+from isaaclab.test.utils import test_devices
+from isaaclab.utils.warp import ProxyArray
+from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids
+
+
+def _all_ids(n: int, device: str) -> wp.array:
+    return wp.array(list(range(n)), dtype=wp.int32, device=device)
+
+
+@pytest.mark.parametrize("device", test_devices())
+def test_resolve_ids(device):
+    """Every selection resolves to device ids; a selection of every element is the shared constant itself."""
+    all_ids = _all_ids(10, device)
+    assert resolve_ids(all_ids=all_ids, device=device) is all_ids
+    assert resolve_ids(ids=slice(None), all_ids=all_ids, device=device) is all_ids
+    assert resolve_ids(ids=[3, 1], all_ids=all_ids, device=device).numpy().tolist() == [3, 1]
+    assert resolve_ids(ids=slice(2, 8, 3), all_ids=all_ids, device=device).numpy().tolist() == [2, 5]
+    assert resolve_ids(ids=slice(None, None, -4), all_ids=all_ids, device=device).numpy().tolist() == [9, 5, 1]
+    ids_torch = torch.tensor([4, 2], dtype=torch.int64, device=wp.device_to_torch(device))
+    assert resolve_ids(ids=ids_torch, all_ids=all_ids, device=device) is ids_torch
+    ids_wp = wp.array([6], dtype=wp.int32, device=device)
+    assert resolve_ids(ids=ids_wp, all_ids=all_ids, device=device) is ids_wp
+    with pytest.raises(TypeError, match="ProxyArray"):
+        resolve_ids(ids=ProxyArray(ids_wp), all_ids=all_ids, device=device)
+
+
+@pytest.mark.parametrize("device", test_devices())
+@pytest.mark.parametrize("mask_type", ["warp", "torch_bool", "torch_int"])
+def test_resolve_ids_from_a_mask(device, mask_type):
+    """A mask resolves to its ascending ids, takes precedence over ids, and honours ``dtype``."""
+    n = 4096
+    expected = (torch.rand(n, generator=torch.Generator().manual_seed(0)) < 0.5).to(wp.device_to_torch(device))
+    mask = {
+        "warp": wp.from_torch(expected, dtype=wp.bool),
+        "torch_bool": expected,
+        "torch_int": expected.to(torch.int64),
+    }[mask_type]
+    all_ids = _all_ids(n, device)
+    ids = resolve_ids(ids=[0, 1], mask=mask, all_ids=all_ids, device=device)
+    assert ids.dtype == torch.int64 and torch.equal(ids, torch.nonzero(expected).flatten())
+    ids_int32 = resolve_ids(mask=mask, all_ids=all_ids, device=device, dtype=wp.int32)
+    assert isinstance(ids_int32, wp.array) and ids_int32.dtype == wp.int32 and ids_int32.is_contiguous
+    assert torch.equal(wp.to_torch(ids_int32), torch.nonzero(expected).flatten().to(torch.int32))
+
+
+@pytest.mark.parametrize("device", test_devices())
+def test_resolve_ids_with_dtype(device):
+    """``dtype=wp.int32`` gives a contiguous one-dimensional int32 Warp array for every kind of ids."""
+    all_ids = _all_ids(10, device)
+    torch_device = wp.device_to_torch(device)
+    assert resolve_ids(all_ids=all_ids, device=device, dtype=wp.int32) is all_ids
+    for ids, expected in (
+        (torch.tensor([4, 2], dtype=torch.int64, device=torch_device), [4, 2]),
+        (torch.arange(10, device=torch_device)[::3], [0, 3, 6, 9]),
+        (wp.array([7, 1], dtype=wp.int64, device=device), [7, 1]),
+        (slice(1, 8, 3), [1, 4, 7]),
+        (slice(None, None, -4), [9, 5, 1]),
+        ([3, 8], [3, 8]),
+    ):
+        resolved = resolve_ids(ids=ids, all_ids=all_ids, device=device, dtype=wp.int32)
+        assert isinstance(resolved, wp.array) and resolved.dtype == wp.int32 and resolved.is_contiguous
+        assert resolved.numpy().tolist() == expected
+    ids_int32 = wp.array([5], dtype=wp.int32, device=device)
+    assert resolve_ids(ids=ids_int32, all_ids=all_ids, device=device, dtype=wp.int32) is ids_int32
+    with pytest.raises(TypeError):
+        resolve_ids(ids=torch.tensor([1.0], device=torch_device), all_ids=all_ids, device=device, dtype=wp.int32)
+    with pytest.raises(TypeError):
+        resolve_ids(
+            ids=wp.array([1.0], dtype=wp.float32, device=device), all_ids=all_ids, device=device, dtype=wp.int32
+        )
+    with pytest.raises(ValueError):
+        resolve_ids(
+            ids=torch.zeros((2, 2), dtype=torch.int64, device=torch_device),
+            all_ids=all_ids,
+            device=device,
+            dtype=wp.int32,
+        )
+
+
+@pytest.mark.parametrize("device", test_devices())
+@pytest.mark.parametrize("ids_type", ["list", "slice", "torch_int64", "torch_int32", "warp_int32"])
+def test_resolve_1d_mask_from_ids(device, ids_type):
+    """Ids of every type are written into the scratch mask, which is cleared first."""
+    n = 12
+    all_mask = wp.ones(n, dtype=wp.bool, device=device)
+    scratch = wp.ones(n, dtype=wp.bool, device=device)
+    selected = [1, 4, 5, 11]
+    torch_device = wp.device_to_torch(device)
+    ids = {
+        "list": selected,
+        "slice": slice(1, 12, 3),
+        "torch_int64": torch.tensor(selected, dtype=torch.int64, device=torch_device),
+        "torch_int32": torch.tensor(selected, dtype=torch.int32, device=torch_device),
+        "warp_int32": wp.array(selected, dtype=wp.int32, device=device),
+    }[ids_type]
+    expected = [False] * n
+    for i in range(n)[ids] if isinstance(ids, slice) else selected:
+        expected[i] = True
+    mask = resolve_1d_mask(ids=ids, all_mask=all_mask, scratch_mask=scratch, device=device)
+    assert mask is scratch
+    assert mask.numpy().tolist() == expected
+
+
+@pytest.mark.parametrize("device", test_devices())
+def test_resolve_1d_mask_needs_a_scratch_mask_only_for_ids(device):
+    """A mask or no selection resolves without a scratch mask; ids without one raise."""
+    all_mask = wp.ones(5, dtype=wp.bool, device=device)
+    mask = wp.array([True, False, True, False, False], dtype=wp.bool, device=device)
+    assert resolve_1d_mask(ids=[0], mask=mask, all_mask=all_mask, device=device) is mask
+    assert resolve_1d_mask(all_mask=all_mask, device=device) is all_mask
+    assert resolve_1d_mask(ids=slice(None), all_mask=all_mask, device=device) is all_mask
+    mask_torch = torch.tensor([0, 1, 1, 0, 0], device=wp.device_to_torch(device))
+    resolved = resolve_1d_mask(mask=mask_torch, all_mask=all_mask, device=device)
+    assert resolved.dtype == wp.bool and resolved.numpy().tolist() == [False, True, True, False, False]
+    with pytest.raises(ValueError, match="scratch_mask"):
+        resolve_1d_mask(ids=[1], all_mask=all_mask, device=device)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device")
+def test_resolve_1d_mask_copies_a_mask_on_another_device():
+    """A Warp mask on another device is copied to ``device``."""
+    all_mask = wp.ones(3, dtype=wp.bool, device="cuda:0")
+    mask = wp.array([True, False, True], dtype=wp.bool, device="cpu")
+    resolved = resolve_1d_mask(mask=mask, all_mask=all_mask, device="cuda:0")
+    assert str(resolved.device) == "cuda:0" and resolved.numpy().tolist() == [True, False, True]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device")
+def test_resolve_1d_mask_in_a_graph_capture():
+    """Inside a capture, a boolean mask on device resolves; ids or a mask that needs a cast or copy raise."""
+    device = "cuda:0"
+    all_mask = wp.ones(3, dtype=wp.bool, device=device)
+    scratch = wp.zeros(3, dtype=wp.bool, device=device)
+    mask_wp = wp.array([True, False, True], dtype=wp.bool, device=device)
+    mask_torch = torch.tensor([False, True, True], device=device)
+    mask_int = mask_torch.int()
+    mask_cpu = wp.array([True, False, True], dtype=wp.bool, device="cpu")
+    with wp.ScopedDevice(device), wp.ScopedCapture():
+        assert resolve_1d_mask(mask=mask_wp, all_mask=all_mask, device=device) is mask_wp
+        resolved = resolve_1d_mask(mask=mask_torch, all_mask=all_mask, device=device)
+        for selection in ({"mask": mask_int}, {"mask": mask_cpu}, {"ids": [0]}):
+            with pytest.raises(RuntimeError, match="capturable"):
+                resolve_1d_mask(**selection, all_mask=all_mask, scratch_mask=scratch, device=device)
+    assert resolved.numpy().tolist() == [False, True, True]
diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-id-mask.changed.rst
new file mode 100644
index 000000000000..ce57c9e7851f
--- /dev/null
+++ b/source/isaaclab_experimental/changelog.d/jichuanh-id-mask.changed.rst
@@ -0,0 +1,2 @@
+* Changed ``isaaclab_experimental.utils.warp.resolve_1d_mask`` to be
+  :func:`isaaclab.utils.warp.utils.resolve_1d_mask` instead of a copy of it.
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py
index 2d071b823a46..757a45e8b825 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py
@@ -5,11 +5,12 @@
 
 """Warp utility functions and shared kernels for isaaclab_experimental."""
 
+from isaaclab.utils.warp.utils import resolve_1d_mask
+
 from .kernels import compute_reset_scale, count_masked
 from .utils import (
     WarpCapturable,
     is_warp_capturable,
-    resolve_1d_mask,
     warp_capturable,
     wrap_to_pi,
     zero_masked_2d,
diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py
index a90f6b2d6fc6..6b8aafc415f8 100644
--- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py
+++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py
@@ -5,97 +5,9 @@
 
 from __future__ import annotations
 
-from collections.abc import Sequence
-
-import torch
 import warp as wp
 
 
-@wp.kernel
-def _set_mask_from_ids(
-    mask: wp.array(dtype=wp.bool),
-    ids: wp.array(dtype=wp.int32),
-):
-    """Set ``mask[ids[i]] = True`` for each thread *i*."""
-    i = wp.tid()
-    mask[ids[i]] = True
-
-
-def resolve_1d_mask(
-    *,
-    ids: Sequence[int] | slice | wp.array | torch.Tensor | None,
-    mask: wp.array | torch.Tensor | None,
-    all_mask: wp.array,
-    scratch_mask: wp.array,
-    device: str,
-) -> wp.array:
-    """Resolve ids/mask into a warp boolean mask.
-
-    Matches the contract of ``ArticulationData._resolve_1d_mask`` on dev/newton.
-    Callers must provide pre-allocated ``all_mask`` (all-True) and ``scratch_mask``
-    (reusable working buffer). No allocations happen inside this function.
-
-    Args:
-        ids: Indices to set to ``True``. ``None`` or ``slice(None)`` means all.
-        mask: Explicit boolean mask. If provided, returned directly (after
-            torch->warp normalization if needed). Takes precedence over *ids*.
-        all_mask: Pre-allocated all-True mask of shape ``(size,)``, returned
-            when both *ids* and *mask* are ``None``.
-        scratch_mask: Pre-allocated scratch mask of shape ``(size,)``, filled
-            in-place when *ids* are provided.
-        device: Warp device string.
-
-    Returns:
-        A ``wp.array(dtype=wp.bool)`` -- ``mask``, ``all_mask``, or ``scratch_mask``.
-    """
-    # Fast path: explicit mask provided.
-    if mask is not None:
-        if isinstance(mask, torch.Tensor):
-            if mask.dtype != torch.bool:
-                mask = mask.to(dtype=torch.bool)
-            if str(mask.device) != device:
-                mask = mask.to(device)
-            return wp.from_torch(mask, dtype=wp.bool)
-        return mask
-
-    # Fast path: all ids.
-    if ids is None or (isinstance(ids, slice) and ids == slice(None)):
-        return all_mask
-
-    # Normalize slice into explicit indices.
-    if isinstance(ids, slice):
-        start, stop, step = ids.indices(scratch_mask.shape[0])
-        ids = list(range(start, stop, step))
-    elif not isinstance(ids, (torch.Tensor, wp.array)):
-        ids = list(ids)
-
-    # Prepare output mask.
-    scratch_mask.fill_(False)
-
-    # Normalize ids to wp.int32 array and launch kernel.
-    if isinstance(ids, torch.Tensor):
-        if ids.numel() == 0:
-            return scratch_mask
-        if str(ids.device) != device:
-            ids = ids.to(device)
-        if ids.dtype != torch.int32:
-            ids = ids.to(dtype=torch.int32)
-        if not ids.is_contiguous():
-            ids = ids.contiguous()
-        ids_wp = wp.from_torch(ids, dtype=wp.int32)
-    elif isinstance(ids, wp.array):
-        if ids.shape[0] == 0:
-            return scratch_mask
-        ids_wp = ids
-    else:
-        if len(ids) == 0:
-            return scratch_mask
-        ids_wp = wp.array(ids, dtype=wp.int32, device=device)
-
-    wp.launch(kernel=_set_mask_from_ids, dim=ids_wp.shape[0], inputs=[scratch_mask, ids_wp], device=device)
-    return scratch_mask
-
-
 def warp_capturable(capturable: bool):
     """Annotate an MDP term's CUDA-graph capturability.
 

From 97e56623b87fdec9128e4d7105ee076590e07f9e Mon Sep 17 00:00:00 2001
From: Jichuan Hu <jichuanh@nvidia.com>
Date: Mon, 26 Oct 2026 02:48:17 -0700
Subject: [PATCH 43/48] Resolve sensor, camera and wrench composer selections
 through the shared helpers

SensorBase, Camera and WrenchComposer resolve env_ids, env_mask and body_ids through resolve_ids and resolve_1d_mask instead of their own copies. The test doubles that build these objects without their initializer get the env masks a real sensor creates.
---
 .../changelog.d/jichuanh-id-mask.changed.rst  |  7 ++++
 .../isaaclab/sensors/camera/camera.py         | 23 +----------
 .../isaaclab/isaaclab/sensors/sensor_base.py  | 13 ++-----
 .../isaaclab/utils/wrench_composer.py         | 38 +++----------------
 .../test_simulation_render_context.py         |  1 +
 .../sensors/test_sensor_update_schedule.py    |  1 +
 .../test/utils/test_wrench_composer.py        |  5 ++-
 7 files changed, 23 insertions(+), 65 deletions(-)

diff --git a/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst
index 6a1b73d77c7e..2ae1cf9e1daa 100644
--- a/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst
+++ b/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst
@@ -1,3 +1,10 @@
+* Changed the env selection helpers of :class:`~isaaclab.utils.wrench_composer.WrenchComposer`,
+  :class:`~isaaclab.sensors.SensorBase` and :class:`~isaaclab.sensors.Camera` to resolve ``env_ids`` and ``env_mask``
+  through :func:`~isaaclab.utils.warp.utils.resolve_ids` and :func:`~isaaclab.utils.warp.utils.resolve_1d_mask`.
+  Any integer sequence and slices with a negative step are accepted, ids on another device are moved to the owner's
+  device, and a :class:`~isaaclab.utils.warp.ProxyArray` is rejected with a :class:`TypeError`.
+* Changed :class:`~isaaclab.utils.wrench_composer.WrenchComposer` to resolve ``body_ids`` through
+  :func:`~isaaclab.utils.warp.utils.resolve_ids` as well, so slices other than ``slice(None)`` are accepted.
 * Changed :func:`~isaaclab.utils.warp.utils.resolve_1d_mask` to make ``scratch_mask`` optional for callers that never
   pass ids. Ids without a ``scratch_mask`` raise a :class:`ValueError`.
 * Changed :func:`~isaaclab.utils.warp.utils.resolve_1d_mask` to copy a Warp mask on another device to ``device``, and
diff --git a/source/isaaclab/isaaclab/sensors/camera/camera.py b/source/isaaclab/isaaclab/sensors/camera/camera.py
index d052dd59aa73..b1e577520bbb 100644
--- a/source/isaaclab/isaaclab/sensors/camera/camera.py
+++ b/source/isaaclab/isaaclab/sensors/camera/camera.py
@@ -26,6 +26,7 @@
     quat_from_matrix,
 )
 from ...utils.warp import ProxyArray
+from ...utils.warp.utils import resolve_ids
 from ..sensor_base import SensorBase
 from .camera_data import CameraData, RenderBufferKind
 
@@ -1028,27 +1029,7 @@ def _resolve_env_ids_wp(self, env_ids: Sequence[int] | torch.Tensor | wp.array |
         """Resolve camera indices to a Warp ``int32`` array."""
         if env_ids is None:
             return None
-        if isinstance(env_ids, wp.array):
-            if env_ids.dtype != wp.int32:
-                raise TypeError(f"Unsupported wp.array dtype for env_ids: {env_ids.dtype}. Expected wp.int32.")
-            if env_ids.ndim != 1:
-                raise ValueError("Warp camera indices must be a one-dimensional array.")
-            if str(env_ids.device) == str(self._device):
-                return env_ids
-            return env_ids.to(self._device)
-        elif isinstance(env_ids, torch.Tensor):
-            env_ids = env_ids.to(device=self._device, dtype=torch.int32).reshape(-1)
-            if not env_ids.is_contiguous():
-                env_ids = env_ids.contiguous()
-            return wp.from_torch(env_ids, dtype=wp.int32)
-        elif isinstance(env_ids, slice) and (env_ids.step or 1) > 0:
-            return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids])
-        elif isinstance(env_ids, slice):
-            # Torch cannot slice with a negative step.
-            env_ids = np.arange(self._view.count, dtype=np.int32)[env_ids]
-        else:
-            env_ids = np.asarray(env_ids, dtype=np.int32).reshape(-1)
-        return wp.array(env_ids, dtype=wp.int32, device=self._device)
+        return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self._device, dtype=wp.int32)
 
     def _env_mask_has_any(self, env_mask: wp.array) -> bool:
         """Return whether the mask selects any camera."""
diff --git a/source/isaaclab/isaaclab/sensors/sensor_base.py b/source/isaaclab/isaaclab/sensors/sensor_base.py
index 5bcec23eb91e..f0e39055540d 100644
--- a/source/isaaclab/isaaclab/sensors/sensor_base.py
+++ b/source/isaaclab/isaaclab/sensors/sensor_base.py
@@ -29,7 +29,7 @@
 from ..physics import PhysicsEvent, PhysicsManager
 from ..sim.utils.queries import get_first_matching_ancestor_prim
 from ..sim.utils.transforms import resolve_prim_pose
-from ..utils import index_fill_
+from ..utils.warp.utils import resolve_1d_mask
 from .kernels import invalidate_buffers_kernel, reset_envs_kernel, update_outdated_envs_kernel, update_timestamp_kernel
 
 if TYPE_CHECKING:
@@ -506,14 +506,9 @@ def _resolve_indices_and_mask(
         self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None
     ) -> wp.array:
         """Resolve environment indices to a warp array and mask."""
-        if (env_ids is None or env_ids == slice(None)) and env_mask is None:
-            return self._ALL_ENV_MASK
-        elif env_mask is not None:
-            return env_mask
-        else:
-            self._reset_mask.zero_()
-            index_fill_(self._reset_mask_torch, env_ids, True)
-            return self._reset_mask
+        return resolve_1d_mask(
+            ids=env_ids, mask=env_mask, all_mask=self._ALL_ENV_MASK, device=self._device, scratch_mask=self._reset_mask
+        )
 
     def _resolve_rigid_body_ancestor_expr(
         self,
diff --git a/source/isaaclab/isaaclab/utils/wrench_composer.py b/source/isaaclab/isaaclab/utils/wrench_composer.py
index 6331716cad5d..ace1570dc4dc 100644
--- a/source/isaaclab/isaaclab/utils/wrench_composer.py
+++ b/source/isaaclab/isaaclab/utils/wrench_composer.py
@@ -24,6 +24,7 @@
     set_forces_to_dual_buffers_index_kernel,
     set_forces_to_dual_buffers_mask,
 )
+from .warp.utils import resolve_ids
 
 if TYPE_CHECKING:
     from ..assets import BaseArticulation, BaseRigidObject, BaseRigidObjectCollection
@@ -721,24 +722,9 @@ def _resolve_env_ids(self, env_ids: wp.array | torch.Tensor | list | slice | Non
             Environment indices.
 
         Raises:
-            TypeError: If ``env_ids`` is an unsupported type.
+            TypeError: If ``env_ids`` is a :class:`~isaaclab.utils.warp.ProxyArray`.
         """
-        if env_ids is None:
-            return self._ALL_ENV_INDICES
-        # Check tensor types before slice comparison (tensor == slice crashes)
-        if isinstance(env_ids, torch.Tensor):
-            return env_ids
-        if isinstance(env_ids, wp.array):
-            return env_ids
-        if env_ids == slice(None):
-            return self._ALL_ENV_INDICES
-        if isinstance(env_ids, slice):
-            return wp.from_torch(wp.to_torch(self._ALL_ENV_INDICES)[env_ids])
-        if isinstance(env_ids, list):
-            return wp.array(env_ids, dtype=wp.int32, device=self.device)
-        raise TypeError(
-            f"env_ids must be None, slice(None), list, torch.Tensor, or wp.array, got {type(env_ids).__name__}"
-        )
+        return resolve_ids(ids=env_ids, all_ids=self._ALL_ENV_INDICES, device=self.device)
 
     def _resolve_body_ids(
         self, body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None
@@ -752,23 +738,9 @@ def _resolve_body_ids(
             Body indices.
 
         Raises:
-            TypeError: If ``body_ids`` is an unsupported type.
+            TypeError: If ``body_ids`` is a :class:`~isaaclab.utils.warp.ProxyArray`.
         """
-        if body_ids is None:
-            return self._ALL_BODY_INDICES
-        if isinstance(body_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if isinstance(body_ids, torch.Tensor):
-            return body_ids
-        if isinstance(body_ids, wp.array):
-            return body_ids
-        if body_ids == slice(None):
-            return self._ALL_BODY_INDICES
-        if isinstance(body_ids, Sequence):
-            return wp.array(body_ids, dtype=wp.int32, device=self.device)
-        raise TypeError(
-            f"body_ids must be None, slice(None), a sequence, torch.Tensor, or wp.array, got {type(body_ids).__name__}"
-        )
+        return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self.device)
 
     def _ensure_composed(self):
         """Compose input buffers into output buffers if dirty."""
diff --git a/source/isaaclab/test/renderers/test_simulation_render_context.py b/source/isaaclab/test/renderers/test_simulation_render_context.py
index fdcc8d4f3779..cf16bc4a2d76 100644
--- a/source/isaaclab/test/renderers/test_simulation_render_context.py
+++ b/source/isaaclab/test/renderers/test_simulation_render_context.py
@@ -317,6 +317,7 @@ def __init__(self, renderer, name, update_period=0.0):
         self._frame = ProxyArray(wp.zeros(2, dtype=wp.int64, device="cpu"))
         self._ALL_INDICES = wp.array([0, 1], dtype=wp.int32, device="cpu")
         self._ALL_ENV_MASK = wp.ones(2, dtype=wp.bool, device="cpu")
+        self._reset_mask = wp.zeros(2, dtype=wp.bool, device="cpu")
         self._create_timing_buffers()
         self._invalidate_buffers(self._ALL_ENV_MASK)
         self.pose = 0.0
diff --git a/source/isaaclab/test/sensors/test_sensor_update_schedule.py b/source/isaaclab/test/sensors/test_sensor_update_schedule.py
index 66224a3fea17..39353ecefdb3 100644
--- a/source/isaaclab/test/sensors/test_sensor_update_schedule.py
+++ b/source/isaaclab/test/sensors/test_sensor_update_schedule.py
@@ -70,6 +70,7 @@ def __init__(self, update_period: float, device: str, num_envs: int = 2):
         self._device = device
         self._num_envs = num_envs
         self._ALL_ENV_MASK = wp.ones(num_envs, dtype=wp.bool, device=device)
+        self._reset_mask = wp.zeros(num_envs, dtype=wp.bool, device=device)
         self._is_initialized = True
         self._is_visualizing = False
         self._create_timing_buffers()
diff --git a/source/isaaclab/test/utils/test_wrench_composer.py b/source/isaaclab/test/utils/test_wrench_composer.py
index 10daf0d06612..b236b052707a 100644
--- a/source/isaaclab/test/utils/test_wrench_composer.py
+++ b/source/isaaclab/test/utils/test_wrench_composer.py
@@ -477,8 +477,9 @@ def test_deprecated_writers_warn_and_keep_add_and_set_behavior() -> None:
 def test_invalid_selection_and_empty_wrench_input_are_reported() -> None:
     composer = _make_composer()
 
-    with pytest.raises(TypeError, match="env_ids must be"):
-        composer.add_forces_and_torques_index(forces=_vectors(torch.ones((1, 2, 3))), env_ids=(0,))
+    env_ids = ProxyArray(wp.array([0], dtype=wp.int32, device=composer.device))
+    with pytest.raises(TypeError, match="ProxyArray is output-only"):
+        composer.add_forces_and_torques_index(forces=_vectors(torch.ones((1, 2, 3))), env_ids=env_ids)
     for method in ("add_forces_and_torques_index", "set_forces_and_torques_mask"):
         with pytest.warns(UserWarning, match="No forces or torques"):
             getattr(composer, method)()

From 92590d00f98878892bfdc8dde0d1eb56b77cde19 Mon Sep 17 00:00:00 2001
From: Jichuan Hu <jichuanh@nvidia.com>
Date: Mon, 26 Oct 2026 02:48:18 -0700
Subject: [PATCH 44/48] Fix the deformable mask writers that raised on any env
 mask

The nodal mask writers of BaseDeformableObject, which the OVPhysX deformable object inherits, and of the PhysX deformable object called wp.nonzero on a Warp array, which raised a RuntimeError. They now resolve the mask to ids with resolve_ids and call the index writers.
---
 .../changelog.d/jichuanh-id-mask.fixed.rst    |  2 ++
 .../base_deformable_object.py                 | 24 +++++---------
 .../changelog.d/jichuanh-id-mask.fixed.rst    |  2 ++
 .../deformable_object/deformable_object.py    | 33 ++++---------------
 .../test/assets/test_deformable_object.py     | 16 +++++++++
 5 files changed, 34 insertions(+), 43 deletions(-)
 create mode 100644 source/isaaclab/changelog.d/jichuanh-id-mask.fixed.rst
 create mode 100644 source/isaaclab_physx/changelog.d/jichuanh-id-mask.fixed.rst

diff --git a/source/isaaclab/changelog.d/jichuanh-id-mask.fixed.rst b/source/isaaclab/changelog.d/jichuanh-id-mask.fixed.rst
new file mode 100644
index 000000000000..9ed262f73ec2
--- /dev/null
+++ b/source/isaaclab/changelog.d/jichuanh-id-mask.fixed.rst
@@ -0,0 +1,2 @@
+* Fixed the ``write_nodal_*_to_sim_mask`` methods of :class:`~isaaclab.assets.BaseDeformableObject`, which the
+  OVPhysX deformable object inherits, raising a :class:`RuntimeError` whenever ``env_mask`` was given.
diff --git a/source/isaaclab/isaaclab/assets/deformable_object/base_deformable_object.py b/source/isaaclab/isaaclab/assets/deformable_object/base_deformable_object.py
index 0542a460af48..4467403b79cf 100644
--- a/source/isaaclab/isaaclab/assets/deformable_object/base_deformable_object.py
+++ b/source/isaaclab/isaaclab/assets/deformable_object/base_deformable_object.py
@@ -15,6 +15,7 @@
 
 from ...utils import math as math_utils
 from ...utils.warp import ProxyArray
+from ...utils.warp.utils import resolve_ids
 from ..asset_base import AssetBase
 
 if TYPE_CHECKING:
@@ -200,6 +201,9 @@ def write_nodal_kinematic_target_to_sim_index(
 
     """
     Operations - Write to simulation (mask variants).
+
+    These convert the mask to ids, which waits for the device, and call the index variants.
+    Backends with a mask write path override them.
     """
 
     def write_nodal_state_to_sim_mask(
@@ -214,10 +218,7 @@ def write_nodal_state_to_sim_mask(
                 Shape is (num_instances, max_sim_vertices_per_body, 6).
             env_mask: Environment mask. If None, then all indices are used.
         """
-        if env_mask is not None:
-            env_ids = wp.nonzero(env_mask)
-        else:
-            env_ids = self._ALL_INDICES
+        env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device)
         self.write_nodal_state_to_sim_index(nodal_state, env_ids=env_ids, full_data=True)
 
     def write_nodal_pos_to_sim_mask(
@@ -232,10 +233,7 @@ def write_nodal_pos_to_sim_mask(
                 Shape is (num_instances, max_sim_vertices_per_body, 3).
             env_mask: Environment mask. If None, then all indices are used.
         """
-        if env_mask is not None:
-            env_ids = wp.nonzero(env_mask)
-        else:
-            env_ids = self._ALL_INDICES
+        env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device)
         self.write_nodal_pos_to_sim_index(nodal_pos, env_ids=env_ids, full_data=True)
 
     def write_nodal_velocity_to_sim_mask(
@@ -250,10 +248,7 @@ def write_nodal_velocity_to_sim_mask(
                 Shape is (num_instances, max_sim_vertices_per_body, 3).
             env_mask: Environment mask. If None, then all indices are used.
         """
-        if env_mask is not None:
-            env_ids = wp.nonzero(env_mask)
-        else:
-            env_ids = self._ALL_INDICES
+        env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device)
         self.write_nodal_velocity_to_sim_index(nodal_vel, env_ids=env_ids, full_data=True)
 
     def write_nodal_kinematic_target_to_sim_mask(
@@ -268,10 +263,7 @@ def write_nodal_kinematic_target_to_sim_mask(
                 Shape is (num_instances, max_sim_vertices_per_body, 4).
             env_mask: Environment mask. If None, then all indices are used.
         """
-        if env_mask is not None:
-            env_ids = wp.nonzero(env_mask)
-        else:
-            env_ids = self._ALL_INDICES
+        env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device)
         self.write_nodal_kinematic_target_to_sim_index(targets, env_ids=env_ids, full_data=True)
 
     """
diff --git a/source/isaaclab_physx/changelog.d/jichuanh-id-mask.fixed.rst b/source/isaaclab_physx/changelog.d/jichuanh-id-mask.fixed.rst
new file mode 100644
index 000000000000..96c018a902e7
--- /dev/null
+++ b/source/isaaclab_physx/changelog.d/jichuanh-id-mask.fixed.rst
@@ -0,0 +1,2 @@
+* Fixed the ``write_nodal_*_to_sim_mask`` methods of :class:`~isaaclab_physx.assets.DeformableObject` raising a
+  :class:`RuntimeError` whenever ``env_mask`` was given.
diff --git a/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py b/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py
index 019d50bdb1a8..d8dc3d9e7bc3 100644
--- a/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py
+++ b/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py
@@ -21,6 +21,7 @@
 from isaaclab.assets.asset_base import AssetBase
 from isaaclab.markers import VisualizationMarkers
 from isaaclab.utils.warp import ProxyArray
+from isaaclab.utils.warp.utils import resolve_ids
 
 from .deformable_object_data import DeformableObjectData
 from .kernels import (
@@ -218,10 +219,7 @@ def write_nodal_state_to_sim_mask(
                 Shape is (num_instances, max_sim_vertices_per_body, 6).
             env_mask: Environment mask. If None, then all indices are used.
         """
-        if env_mask is not None:
-            env_ids = wp.nonzero(env_mask)
-        else:
-            env_ids = self._ALL_INDICES
+        env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device)
         self.write_nodal_state_to_sim_index(nodal_state, env_ids=env_ids, full_data=True)
 
     def write_nodal_pos_to_sim_index(
@@ -288,10 +286,7 @@ def write_nodal_pos_to_sim_mask(
                 Shape is (num_instances, max_sim_vertices_per_body, 3).
             env_mask: Environment mask. If None, then all indices are used.
         """
-        if env_mask is not None:
-            env_ids = wp.nonzero(env_mask)
-        else:
-            env_ids = self._ALL_INDICES
+        env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device)
         self.write_nodal_pos_to_sim_index(nodal_pos, env_ids=env_ids, full_data=True)
 
     def write_nodal_velocity_to_sim_index(
@@ -359,10 +354,7 @@ def write_nodal_velocity_to_sim_mask(
                 Shape is (num_instances, max_sim_vertices_per_body, 3).
             env_mask: Environment mask. If None, then all indices are used.
         """
-        if env_mask is not None:
-            env_ids = wp.nonzero(env_mask)
-        else:
-            env_ids = self._ALL_INDICES
+        env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device)
         self.write_nodal_velocity_to_sim_index(nodal_vel, env_ids=env_ids, full_data=True)
 
     def write_nodal_kinematic_target_to_sim_index(
@@ -438,10 +430,7 @@ def write_nodal_kinematic_target_to_sim_mask(
                 Shape is (num_instances, max_sim_vertices_per_body, 4).
             env_mask: Environment mask. If None, then all indices are used.
         """
-        if env_mask is not None:
-            env_ids = wp.nonzero(env_mask)
-        else:
-            env_ids = self._ALL_INDICES
+        env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device)
         self.write_nodal_kinematic_target_to_sim_index(targets, env_ids=env_ids, full_data=True)
 
     """
@@ -566,17 +555,7 @@ def _resolve_env_ids(self, env_ids):
         Returns:
             A warp array of environment indices.
         """
-        if (env_ids is None) or (env_ids == slice(None)):
-            return self._ALL_INDICES
-        if isinstance(env_ids, slice):
-            return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids])
-        if isinstance(env_ids, torch.Tensor):
-            if env_ids.dtype == torch.int64:
-                env_ids = env_ids.to(torch.int32)
-            return wp.from_torch(env_ids, dtype=wp.int32)
-        if isinstance(env_ids, list):
-            return wp.array(env_ids, dtype=wp.int32, device=self.device)
-        return env_ids
+        return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device, dtype=wp.int32)
 
     def _initialize_impl(self):
         # obtain global simulation view
diff --git a/source/isaaclab_physx/test/assets/test_deformable_object.py b/source/isaaclab_physx/test/assets/test_deformable_object.py
index cfd389fa99ef..4eeb4c246f8c 100644
--- a/source/isaaclab_physx/test/assets/test_deformable_object.py
+++ b/source/isaaclab_physx/test/assets/test_deformable_object.py
@@ -244,6 +244,22 @@ def test_deformable_nodal_state_writes(deformable_scene: _DeformableScene) -> No
     torch.testing.assert_close(raw_nodal_pos.reshape(nodal_pos.shape), nodal_pos)
 
 
+@pytest.mark.isaacsim_ci
+def test_deformable_nodal_state_mask_writes(deformable_scene: _DeformableScene) -> None:
+    """Mask writes take full-size data and reach only the environments the mask selects."""
+    scene = deformable_scene
+    device = scene.device
+    env_mask = wp.array([False, True], dtype=wp.bool, device=device)
+    for deformable in (scene.volume, scene.surface):
+        initial_state = deformable.data.nodal_state_w.torch.clone()
+        state = initial_state.clone()
+        state[:, 0, :3] += torch.tensor([0.01, 0.02, 0.03], device=device)
+        state[:, 0, 3] = 0.25
+        deformable.write_nodal_state_to_sim_mask(state, env_mask=env_mask)
+        expected = torch.cat((initial_state[:1], state[1:]))
+        torch.testing.assert_close(deformable.data.nodal_state_w.torch, expected)
+
+
 @pytest.mark.isaacsim_ci
 def test_volume_kinematic_targets_hold_selected_nodes(deformable_scene: _DeformableScene) -> None:
     """Kinematic targets of environment 0 hold its nodes while environment 1 keeps moving freely."""

From f044c920104c97249838a46e6dbb60948b5ea3e3 Mon Sep 17 00:00:00 2001
From: Jichuan Hu <jichuanh@nvidia.com>
Date: Mon, 26 Oct 2026 02:48:20 -0700
Subject: [PATCH 45/48] Resolve Newton selections through the shared helpers
 and write its mask writers through masks

The Newton assets resolve ids and masks through resolve_ids and resolve_1d_mask. The RigidObjectCollection mask writers and the fixed-tendon mask setters write through the masks instead of converting them to ids, which waited for the GPU on every call; the joint-coordinate writes reuse a scratch mask; one generic masked 2-D write kernel replaces two copies.
---
 .../assets/contract/articulation_factory.py   |   1 +
 .../assets/contract/rigid_object_factory.py   |   1 +
 .../contract/test_articulation_contract.py    |   5 -
 .../contract/test_rigid_object_contract.py    |  21 ++
 .../changelog.d/jichuanh-id-mask.changed.rst  |  13 ++
 .../changelog.d/jichuanh-id-mask.fixed.rst    |   2 +
 .../assets/articulation/articulation.py       | 159 +++++++--------
 .../assets/articulation/kernels.py            |  30 +++
 .../assets/cable_object/cable_object.py       |  11 +-
 .../deformable_object/deformable_object.py    |  17 +-
 .../isaaclab_newton/assets/kernels.py         | 113 ++++++++++-
 .../assets/mpm_object/mpm_object.py           |  19 +-
 .../assets/rigid_object/rigid_object.py       |  17 +-
 .../rigid_object_collection.py                | 189 ++++++++----------
 14 files changed, 332 insertions(+), 266 deletions(-)
 create mode 100644 source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst
 create mode 100644 source/isaaclab_newton/changelog.d/jichuanh-id-mask.fixed.rst

diff --git a/source/isaaclab/test/assets/contract/articulation_factory.py b/source/isaaclab/test/assets/contract/articulation_factory.py
index 8ff52e4fdafb..6b0488bb006b 100644
--- a/source/isaaclab/test/assets/contract/articulation_factory.py
+++ b/source/isaaclab/test/assets/contract/articulation_factory.py
@@ -237,6 +237,7 @@ def _newton_articulation(N, J, B, FT, ST, device, is_fixed_base, joint_ordering,
     articulation._ALL_BODY_INDICES = indices(B, device)
     articulation._ALL_JOINT_INDICES = indices(J, device)
     articulation._ALL_ENV_MASK = wp.ones((N,), dtype=wp.bool, device=device)
+    articulation._scratch_env_mask = wp.zeros((N,), dtype=wp.bool, device=device)
     articulation._ALL_BODY_MASK = wp.ones((B,), dtype=wp.bool, device=device)
     articulation._ALL_JOINT_MASK = wp.ones((J,), dtype=wp.bool, device=device)
     articulation._resolve_and_install_ordering_maps()
diff --git a/source/isaaclab/test/assets/contract/rigid_object_factory.py b/source/isaaclab/test/assets/contract/rigid_object_factory.py
index 0bf4741ed9f5..cb108da59a1c 100644
--- a/source/isaaclab/test/assets/contract/rigid_object_factory.py
+++ b/source/isaaclab/test/assets/contract/rigid_object_factory.py
@@ -195,6 +195,7 @@ def _newton_collection(num_instances: int, num_bodies: int, device: str, monkeyp
     collection._ALL_BODY_INDICES = indices(num_bodies, device)
     collection._ALL_ENV_MASK = wp.ones((num_instances,), dtype=wp.bool, device=device)
     collection._ALL_BODY_MASK = wp.ones((num_bodies,), dtype=wp.bool, device=device)
+    collection._scratch_body_mask = wp.zeros((num_bodies,), dtype=wp.bool, device=device)
     return collection, view
 
 
diff --git a/source/isaaclab/test/assets/contract/test_articulation_contract.py b/source/isaaclab/test/assets/contract/test_articulation_contract.py
index 99abc45ae632..42878092b782 100644
--- a/source/isaaclab/test/assets/contract/test_articulation_contract.py
+++ b/source/isaaclab/test/assets/contract/test_articulation_contract.py
@@ -536,11 +536,6 @@ def read() -> ProxyArray:
 
 
 @requires("newton")
-@pytest.mark.xfail(
-    raises=pytest.fail.Exception,
-    strict=True,
-    reason="Newton fixed-tendon mask setters slice full data without validating its shape",
-)
 def test_newton_fixed_tendon_mask_setters_reject_extra_environments(monkeypatch):
     art, _ = _articulation("newton", "cpu", monkeypatch)
     with pytest.raises((AssertionError, RuntimeError)):
diff --git a/source/isaaclab/test/assets/contract/test_rigid_object_contract.py b/source/isaaclab/test/assets/contract/test_rigid_object_contract.py
index 1d294d9a8efc..514f4811ad57 100644
--- a/source/isaaclab/test/assets/contract/test_rigid_object_contract.py
+++ b/source/isaaclab/test/assets/contract/test_rigid_object_contract.py
@@ -314,6 +314,27 @@ def test_pose_and_velocity_writer_aliases_forward_to_canonical_writers(monkeypat
             torch.testing.assert_close(getattr(obj.data, getter).torch, values, atol=1e-5, rtol=1e-5)
 
 
+@pytest.mark.parametrize("quantity", ["link_pose", "com_pose", "com_velocity", "link_velocity"])
+def test_collection_mask_writers_select_bodies_by_ids(monkeypatch, backend, device, quantity):
+    """Collection mask writers given ``body_ids`` write only those bodies of the masked environments."""
+    if backend == "ovphysx":
+        pytest.skip("OVPhysX collection mask writers take no body_ids.")
+    obj, _, _, num_bodies = _create("collection", backend, monkeypatch, device)
+    obj.data.update(dt=0.01)
+    trailing = (7,) if quantity.endswith("pose") else (6,)
+    kwarg = "body_poses" if quantity.endswith("pose") else "body_velocities"
+    getter = f"body_{quantity.replace('velocity', 'vel')}_w"
+    method = getattr(obj, f"write_body_{quantity}_to_sim_mask")
+    full = _payload((_NUM_INSTANCES, num_bodies), trailing, device=device)
+    method(**{kwarg: full})
+    update = _payload((_NUM_INSTANCES, num_bodies), trailing, offset=100.0, device=device)
+    env_mask = wp.array(np.arange(_NUM_INSTANCES) == _NUM_INSTANCES - 1, dtype=wp.bool, device=str(device))
+    method(**{kwarg: update}, env_mask=env_mask, body_ids=[num_bodies - 1])
+    expected = full.clone()
+    expected[-1, -1] = update[-1, -1]
+    torch.testing.assert_close(getattr(obj.data, getter).torch, expected, atol=1e-5, rtol=1e-5)
+
+
 # backend-stored quantity -> (written quantity, per-element shape)
 _PARTIAL_WRITES = {"pose": ("link_pose", (7,)), "velocity": ("com_velocity", (6,)), "mass": ("masses", ())}
 
diff --git a/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst
new file mode 100644
index 000000000000..0c69e3841f70
--- /dev/null
+++ b/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst
@@ -0,0 +1,13 @@
+* Changed the env selection helpers of the Newton assets to resolve ``env_ids`` through
+  :func:`~isaaclab.utils.warp.utils.resolve_ids` and ``env_mask`` through
+  :func:`~isaaclab.utils.warp.utils.resolve_1d_mask`. Any integer sequence and slices with a negative step are accepted,
+  ids on another device are moved to the asset's device, and a :class:`~isaaclab.utils.warp.ProxyArray` is rejected
+  with a :class:`TypeError`.
+* Changed the body, joint and tendon selection helpers of the Newton assets to resolve ids and masks through
+  :func:`~isaaclab.utils.warp.utils.resolve_ids`, as the env selection does, so slices other than ``slice(None)`` are
+  accepted.
+* Changed the joint-coordinate writes of :class:`~isaaclab_newton.assets.Articulation` to write env ids into a
+  scratch mask of the articulation instead of allocating a mask on every call.
+* Changed the mask writers of :class:`~isaaclab_newton.assets.RigidObjectCollection` and the fixed-tendon mask setters
+  of :class:`~isaaclab_newton.assets.Articulation` to write through the masks, as the other Newton mask writers do,
+  instead of converting the masks to indices, which synchronized the host with the GPU on every call.
diff --git a/source/isaaclab_newton/changelog.d/jichuanh-id-mask.fixed.rst b/source/isaaclab_newton/changelog.d/jichuanh-id-mask.fixed.rst
new file mode 100644
index 000000000000..638c02730a49
--- /dev/null
+++ b/source/isaaclab_newton/changelog.d/jichuanh-id-mask.fixed.rst
@@ -0,0 +1,2 @@
+* Fixed the fixed-tendon mask setters of :class:`~isaaclab_newton.assets.Articulation` accepting data whose shape
+  does not match the articulation.
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
index f6b03aa5b673..5ed2cb339d61 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
@@ -12,7 +12,7 @@
 import re
 import warnings
 from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 import numpy as np
 import torch
@@ -33,6 +33,7 @@
 from isaaclab.utils.string import resolve_matching_names, resolve_matching_names_values
 from isaaclab.utils.version import get_isaac_sim_version, has_kit
 from isaaclab.utils.warp import ProxyArray
+from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids
 from isaaclab.utils.wrench_composer import WrenchComposer
 
 from isaaclab_newton.assets import kernels as shared_kernels
@@ -2298,15 +2299,11 @@ def write_joint_viscous_friction_coefficient_to_sim_mask(
     Operations - Newton Actuator Parameter Writers.
     """
 
-    @staticmethod
-    @wp.kernel(enable_backward=False)
-    def _build_env_mask_kernel(mask: wp.array(dtype=wp.bool), indices: wp.array(dtype=Any)):
-        i = wp.tid()
-        mask[wp.int32(indices[i])] = True
-
     def _env_ids_to_mask(self, env_ids: wp.array | torch.Tensor) -> wp.array:
         """Convert env_ids to a boolean Warp mask.
 
+        Needed because :func:`scatter_joint_coordinates` only takes a mask.
+
         Args:
             env_ids: Environment indices as returned by :meth:`_resolve_env_ids`, which may be a
                 warp array or a torch tensor of any integer width.
@@ -2316,11 +2313,11 @@ def _env_ids_to_mask(self, env_ids: wp.array | torch.Tensor) -> wp.array:
         """
         if env_ids is self._ALL_INDICES:
             return self._ALL_ENV_MASK
-        if isinstance(env_ids, torch.Tensor):
-            env_ids = wp.from_torch(env_ids)
-        mask = wp.zeros(self.num_instances, dtype=wp.bool, device=self.device)
-        wp.launch(self._build_env_mask_kernel, dim=env_ids.shape[0], inputs=[mask, env_ids], device=self.device)
-        return mask
+        if isinstance(env_ids, wp.array) and env_ids.dtype != wp.int32:
+            env_ids = wp.to_torch(env_ids)
+        return resolve_1d_mask(
+            ids=env_ids, all_mask=self._ALL_ENV_MASK, scratch_mask=self._scratch_env_mask, device=self.device
+        )
 
     """
     Operations - Setters.
@@ -2716,12 +2713,17 @@ def set_fixed_tendon_stiffness_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        env_ids = self._resolve_env_mask(env_mask)
-        fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask)
-        self.set_fixed_tendon_stiffness_index(
-            stiffness=self._select_full_data(stiffness, env_ids, fixed_tendon_ids),
-            fixed_tendon_ids=fixed_tendon_ids,
-            env_ids=env_ids,
+        env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK)
+        fixed_tendon_mask = self._resolve_mask(fixed_tendon_mask, self._ALL_FIXED_TENDON_MASK)
+        self.assert_shape_and_dtype_mask(stiffness, (env_mask, fixed_tendon_mask), wp.float32, "stiffness")
+        wp.launch(
+            articulation_kernels.float_data_to_buffer_with_mask
+            if isinstance(stiffness, float)
+            else shared_kernels.write_2d_data_to_buffer_with_mask,
+            dim=(self.num_instances, self.num_fixed_tendons),
+            inputs=[stiffness, env_mask, fixed_tendon_mask],
+            outputs=[self.data._fixed_tendon_stiffness],
+            device=self.device,
         )
 
     def set_fixed_tendon_damping_index(
@@ -2813,12 +2815,17 @@ def set_fixed_tendon_damping_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        env_ids = self._resolve_env_mask(env_mask)
-        fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask)
-        self.set_fixed_tendon_damping_index(
-            damping=self._select_full_data(damping, env_ids, fixed_tendon_ids),
-            fixed_tendon_ids=fixed_tendon_ids,
-            env_ids=env_ids,
+        env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK)
+        fixed_tendon_mask = self._resolve_mask(fixed_tendon_mask, self._ALL_FIXED_TENDON_MASK)
+        self.assert_shape_and_dtype_mask(damping, (env_mask, fixed_tendon_mask), wp.float32, "damping")
+        wp.launch(
+            articulation_kernels.float_data_to_buffer_with_mask
+            if isinstance(damping, float)
+            else shared_kernels.write_2d_data_to_buffer_with_mask,
+            dim=(self.num_instances, self.num_fixed_tendons),
+            inputs=[damping, env_mask, fixed_tendon_mask],
+            outputs=[self.data._fixed_tendon_damping],
+            device=self.device,
         )
 
     def set_fixed_tendon_limit_stiffness_index(
@@ -2930,12 +2937,17 @@ def set_fixed_tendon_position_limit_mask(
             fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used.
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
-        env_ids = self._resolve_env_mask(env_mask)
-        fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask)
-        self.set_fixed_tendon_position_limit_index(
-            limit=self._select_full_data(limit, env_ids, fixed_tendon_ids),
-            fixed_tendon_ids=fixed_tendon_ids,
-            env_ids=env_ids,
+        env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK)
+        fixed_tendon_mask = self._resolve_mask(fixed_tendon_mask, self._ALL_FIXED_TENDON_MASK)
+        self.assert_shape_and_dtype_mask(limit, (env_mask, fixed_tendon_mask), wp.vec2f, "limit")
+        if isinstance(limit, float):
+            raise ValueError("Fixed tendon position limits must be a tensor or array, not a float.")
+        wp.launch(
+            shared_kernels.write_2d_data_to_buffer_with_mask_kernel(wp.vec2f),
+            dim=(self.num_instances, self.num_fixed_tendons),
+            inputs=[limit, env_mask, fixed_tendon_mask],
+            outputs=[self.data._fixed_tendon_pos_limits],
+            device=self.device,
         )
 
     def set_fixed_tendon_rest_length_index(
@@ -3136,10 +3148,27 @@ def write_fixed_tendon_properties_to_sim_mask(
             fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are updated.
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
-        self.write_fixed_tendon_properties_to_sim_index(
-            fixed_tendon_ids=self._resolve_fixed_tendon_mask(fixed_tendon_mask),
-            env_ids=self._resolve_env_mask(env_mask),
+        env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK)
+        fixed_tendon_mask = self._resolve_mask(fixed_tendon_mask, self._ALL_FIXED_TENDON_MASK)
+        wp.launch(
+            articulation_kernels.copy_fixed_tendon_properties_with_mask,
+            dim=(self.num_instances, self.num_fixed_tendons),
+            inputs=[
+                env_mask,
+                fixed_tendon_mask,
+                self.data._fixed_tendon_stiffness,
+                self.data._fixed_tendon_damping,
+                self.data._fixed_tendon_pos_limits,
+            ],
+            outputs=[
+                self.data._sim_bind_fixed_tendon_stiffness,
+                self.data._sim_bind_fixed_tendon_damping,
+                self.data._sim_bind_fixed_tendon_pos_limits,
+            ],
+            device=self.device,
         )
+        # the solver keeps its own copy of the tendon properties and only re-reads them when notified
+        self._backend_manager.add_model_change(ModelFlags.TENDON_PROPERTIES)
 
     def set_spatial_tendon_stiffness_index(
         self,
@@ -3445,6 +3474,7 @@ def _clear_callbacks(self) -> None:
     def _create_buffers(self):
         self._ALL_INDICES = wp.array(np.arange(self.num_instances, dtype=np.int32), device=self.device)
         self._ALL_ENV_MASK = wp.ones((self.num_instances,), dtype=wp.bool, device=self.device)
+        self._scratch_env_mask = wp.zeros((self.num_instances,), dtype=wp.bool, device=self.device)
         self._ALL_JOINT_INDICES = wp.array(np.arange(self.num_joints, dtype=np.int32), device=self.device)
         self._ALL_JOINT_MASK = wp.ones((self.num_joints,), dtype=wp.bool, device=self.device)
         self._ALL_BODY_INDICES = wp.array(np.arange(self.num_bodies, dtype=np.int32), device=self.device)
@@ -3720,15 +3750,7 @@ def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | No
         Returns:
             Environment indices.
         """
-        if isinstance(env_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if (env_ids is None) or (env_ids == slice(None)):
-            return self._ALL_INDICES
-        if isinstance(env_ids, slice):
-            return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids])
-        if isinstance(env_ids, list):
-            return wp.array(env_ids, dtype=wp.int32, device=self.device)
-        return env_ids
+        return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device)
 
     def _resolve_joint_ids(self, joint_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor:
         """Resolve joint indices to a warp array or tensor.
@@ -3739,13 +3761,7 @@ def _resolve_joint_ids(self, joint_ids: Sequence[int] | torch.Tensor | wp.array
         Returns:
             A warp array of joint indices or a tensor of joint indices.
         """
-        if isinstance(joint_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if isinstance(joint_ids, list):
-            return wp.array(joint_ids, dtype=wp.int32, device=self.device)
-        if (joint_ids is None) or (joint_ids == slice(None)):
-            return self._ALL_JOINT_INDICES
-        return joint_ids
+        return resolve_ids(ids=joint_ids, all_ids=self._ALL_JOINT_INDICES, device=self.device)
 
     def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor:
         """Resolve body indices to a warp array or tensor.
@@ -3756,13 +3772,7 @@ def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array |
         Returns:
             A warp array of body indices or a tensor of body indices.
         """
-        if isinstance(body_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if isinstance(body_ids, list):
-            return wp.array(body_ids, dtype=wp.int32, device=self.device)
-        if (body_ids is None) or (body_ids == slice(None)):
-            return self._ALL_BODY_INDICES
-        return body_ids
+        return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self.device)
 
     def _resolve_fixed_tendon_ids(
         self, tendon_ids: Sequence[int] | torch.Tensor | wp.array | None
@@ -3775,13 +3785,7 @@ def _resolve_fixed_tendon_ids(
         Returns:
             A warp array of tendon indices or a tensor of tendon indices.
         """
-        if isinstance(tendon_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if isinstance(tendon_ids, list):
-            return wp.array(tendon_ids, dtype=wp.int32, device=self.device)
-        if (tendon_ids is None) or (tendon_ids == slice(None)):
-            return self._ALL_FIXED_TENDON_INDICES
-        return tendon_ids
+        return resolve_ids(ids=tendon_ids, all_ids=self._ALL_FIXED_TENDON_INDICES, device=self.device)
 
     def _resolve_spatial_tendon_ids(
         self, spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None
@@ -3794,27 +3798,7 @@ def _resolve_spatial_tendon_ids(
         Returns:
             A warp array of spatial tendon indices or a tensor of spatial tendon indices.
         """
-        if isinstance(spatial_tendon_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if isinstance(spatial_tendon_ids, list):
-            return wp.array(spatial_tendon_ids, dtype=wp.int32, device=self.device)
-        if (spatial_tendon_ids is None) or (spatial_tendon_ids == slice(None)):
-            return self._ALL_SPATIAL_TENDON_INDICES
-        return spatial_tendon_ids
-
-    def _resolve_env_mask(self, env_mask: wp.array | torch.Tensor | None) -> wp.array | torch.Tensor:
-        """Resolve an environment mask to environment indices."""
-        return self._ALL_INDICES if env_mask is None else self._mask_to_ids(env_mask)
-
-    def _resolve_fixed_tendon_mask(self, fixed_tendon_mask: wp.array | torch.Tensor | None) -> wp.array | torch.Tensor:
-        """Resolve a fixed tendon mask to fixed tendon indices."""
-        return self._ALL_FIXED_TENDON_INDICES if fixed_tendon_mask is None else self._mask_to_ids(fixed_tendon_mask)
-
-    @staticmethod
-    def _mask_to_ids(mask: wp.array | torch.Tensor) -> torch.Tensor:
-        """Convert a boolean mask to int32 indices."""
-        mask = wp.to_torch(mask) if isinstance(mask, wp.array) else mask
-        return torch.nonzero(mask)[:, 0].to(torch.int32)
+        return resolve_ids(ids=spatial_tendon_ids, all_ids=self._ALL_SPATIAL_TENDON_INDICES, device=self.device)
 
     @staticmethod
     def _to_torch_ids(ids: wp.array | torch.Tensor) -> torch.Tensor:
@@ -3839,12 +3823,7 @@ def _resolve_mask(self, mask: wp.array | torch.Tensor | None, full_mask: wp.arra
         Returns:
             A warp array of mask.
         """
-        if mask is None:
-            return full_mask
-
-        if isinstance(mask, torch.Tensor):
-            return wp.from_torch(mask, dtype=wp.bool)
-        return mask
+        return resolve_1d_mask(mask=mask, all_mask=full_mask, device=self.device)
 
     """
     Deprecated methods.
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py
index 60544e3a4100..7ad2bd59f057 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py
@@ -553,6 +553,36 @@ def float_data_to_buffer_with_indices_kernel(
     return _FLOAT_DATA_TO_BUFFER_WITH_INDICES_DISPATCHER.select(env_ids, joint_ids)
 
 
+@wp.kernel
+def copy_fixed_tendon_properties_with_mask(
+    env_mask: wp.array(dtype=wp.bool),
+    tendon_mask: wp.array(dtype=wp.bool),
+    stiffness: wp.array2d(dtype=wp.float32),
+    damping: wp.array2d(dtype=wp.float32),
+    pos_limits: wp.array2d(dtype=wp.vec2f),
+    sim_stiffness: wp.array2d(dtype=wp.float32),
+    sim_damping: wp.array2d(dtype=wp.float32),
+    sim_pos_limits: wp.array2d(dtype=wp.vec2f),
+):
+    """Copy the staged fixed-tendon stiffness, damping and position limits to the simulation where both masks are true.
+
+    Args:
+        env_mask: Environment mask. Shape is (num_instances,).
+        tendon_mask: Fixed-tendon mask. Shape is (num_fixed_tendons,).
+        stiffness: Staged stiffness. Shape is (num_instances, num_fixed_tendons).
+        damping: Staged damping. Shape is (num_instances, num_fixed_tendons).
+        pos_limits: Staged position limits. Shape is (num_instances, num_fixed_tendons).
+        sim_stiffness: Simulation stiffness. Shape is (num_instances, num_fixed_tendons).
+        sim_damping: Simulation damping. Shape is (num_instances, num_fixed_tendons).
+        sim_pos_limits: Simulation position limits. Shape is (num_instances, num_fixed_tendons).
+    """
+    i, j = wp.tid()
+    if env_mask[i] and tendon_mask[j]:
+        sim_stiffness[i, j] = stiffness[i, j]
+        sim_damping[i, j] = damping[i, j]
+        sim_pos_limits[i, j] = pos_limits[i, j]
+
+
 @wp.kernel
 def float_data_to_buffer_with_mask(
     in_data: wp.float32,
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/cable_object/cable_object.py b/source/isaaclab_newton/isaaclab_newton/assets/cable_object/cable_object.py
index 86bd5d8590cb..d98cef9aba62 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/cable_object/cable_object.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/cable_object/cable_object.py
@@ -18,6 +18,7 @@
 from isaaclab.assets.cable_object.base_cable_object import BaseCableObject
 from isaaclab.sim.utils.queries import has_deformable_curve_api, path_expr_to_glob, resolve_matching_prims_from_source
 from isaaclab.utils.warp import ProxyArray
+from isaaclab.utils.warp.utils import resolve_ids
 
 from .cable_object_data import CableObjectData
 from .kernels import (
@@ -207,15 +208,7 @@ def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array(dtyp
         dtype=wp.int32
     ):
         """Resolve environment indices to a Warp array."""
-        if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)):
-            return self._ALL_INDICES
-        if isinstance(env_ids, slice):
-            return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids])
-        if isinstance(env_ids, torch.Tensor):
-            return wp.from_torch(env_ids.to(device=self.device, dtype=torch.int32).contiguous(), dtype=wp.int32)
-        if isinstance(env_ids, Sequence):
-            return wp.array(list(env_ids), dtype=wp.int32, device=self.device)
-        return env_ids
+        return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device, dtype=wp.int32)
 
     def _write_segment_state_to_sim(
         self,
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/deformable_object/deformable_object.py b/source/isaaclab_newton/isaaclab_newton/assets/deformable_object/deformable_object.py
index 06e96694147e..d2016b4186ed 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/deformable_object/deformable_object.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/deformable_object/deformable_object.py
@@ -17,6 +17,7 @@
 from isaaclab.assets.deformable_object.base_deformable_object import BaseDeformableObject
 from isaaclab.markers import VisualizationMarkers
 from isaaclab.utils.warp import ProxyArray
+from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids
 
 from .deformable_object_data import DeformableObjectData
 from .kernels import (
@@ -402,23 +403,11 @@ def write_nodal_kinematic_target_to_sim_mask(
 
     def _resolve_env_ids(self, env_ids):
         """Resolve environment indices to a warp int32 array."""
-        if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)):
-            return self._ALL_INDICES
-        elif isinstance(env_ids, list):
-            return wp.array(env_ids, dtype=wp.int32, device=self.device)
-        elif isinstance(env_ids, torch.Tensor):
-            return wp.from_torch(env_ids.to(torch.int32), dtype=wp.int32)
-        return env_ids
+        return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device, dtype=wp.int32)
 
     def _resolve_mask(self, mask: wp.array | torch.Tensor | None, full_mask: wp.array) -> wp.array:
         """Resolve an environment mask to a warp bool array."""
-        if mask is None:
-            return full_mask
-        if isinstance(mask, torch.Tensor):
-            if mask.dtype != torch.bool:
-                mask = mask.to(torch.bool)
-            return wp.from_torch(mask, dtype=wp.bool)
-        return mask
+        return resolve_1d_mask(mask=mask, all_mask=full_mask, device=self.device)
 
     def _iter_particle_states(self):
         """Yield active Newton states."""
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/kernels.py b/source/isaaclab_newton/isaaclab_newton/assets/kernels.py
index eec3479f1fdc..c9f43ef3adec 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/kernels.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/kernels.py
@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import functools
 from typing import TYPE_CHECKING, Any
 
 import torch
@@ -960,6 +961,89 @@ def set_body_link_velocity_to_sim(
     body_acc_w[env_id, body_id] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
 
 
+@wp.kernel
+def set_body_com_pose_to_sim_mask(
+    data: wp.array2d(dtype=wp.transformf),
+    body_com_pos_b: wp.array2d(dtype=wp.vec3f),
+    env_mask: wp.array(dtype=wp.bool),
+    body_mask: wp.array(dtype=wp.bool),
+    body_com_pose_w: wp.array2d(dtype=wp.transformf),
+    body_link_pose_w: wp.array2d(dtype=wp.transformf),
+):
+    """Write body COM pose data, and the link poses derived from it, for the selected environments and bodies.
+
+    Args:
+        data: Input array of body COM poses. Shape is (num_envs, num_bodies).
+        body_com_pos_b: Input array of body COM positions in body frame. Shape is (num_envs, num_bodies).
+        env_mask: Input array of environment mask. Shape is (num_envs,).
+        body_mask: Input array of body mask. Shape is (num_bodies,).
+        body_com_pose_w: Output array where body COM poses are written. Shape is (num_envs, num_bodies).
+        body_link_pose_w: Output array where body link poses (derived from COM) are written.
+            Shape is (num_envs, num_bodies).
+    """
+    i, j = wp.tid()
+    if env_mask[i] and body_mask[j]:
+        body_com_pose_w[i, j] = data[i, j]
+        body_link_pose_w[i, j] = get_com_pose_in_link_frame_func(body_com_pose_w[i, j], body_com_pos_b[i, j])
+
+
+@wp.kernel
+def set_body_com_velocity_to_sim_mask(
+    data: wp.array2d(dtype=wp.spatial_vectorf),
+    env_mask: wp.array(dtype=wp.bool),
+    body_mask: wp.array(dtype=wp.bool),
+    body_com_velocity_w: wp.array2d(dtype=wp.spatial_vectorf),
+    body_acc_w: wp.array2d(dtype=wp.spatial_vectorf),
+):
+    """Write body COM velocity data, and zero the body accelerations, for the selected environments and bodies.
+
+    Args:
+        data: Input array of body COM spatial velocities. Shape is (num_envs, num_bodies).
+        env_mask: Input array of environment mask. Shape is (num_envs,).
+        body_mask: Input array of body mask. Shape is (num_bodies,).
+        body_com_velocity_w: Output array where body COM velocities are written. Shape is (num_envs, num_bodies).
+        body_acc_w: Output array where body accelerations are zeroed. Shape is (num_envs, num_bodies).
+    """
+    i, j = wp.tid()
+    if env_mask[i] and body_mask[j]:
+        body_com_velocity_w[i, j] = data[i, j]
+        body_acc_w[i, j] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
+
+
+@wp.kernel
+def set_body_link_velocity_to_sim_mask(
+    data: wp.array2d(dtype=wp.spatial_vectorf),
+    body_com_pos_b: wp.array2d(dtype=wp.vec3f),
+    body_link_pose_w: wp.array2d(dtype=wp.transformf),
+    env_mask: wp.array(dtype=wp.bool),
+    body_mask: wp.array(dtype=wp.bool),
+    body_link_velocity_w: wp.array2d(dtype=wp.spatial_vectorf),
+    body_com_velocity_w: wp.array2d(dtype=wp.spatial_vectorf),
+    body_acc_w: wp.array2d(dtype=wp.spatial_vectorf),
+):
+    """Write body link velocity data, the COM velocities derived from it, and zero the body accelerations, for the
+    selected environments and bodies.
+
+    Args:
+        data: Input array of body link spatial velocities. Shape is (num_envs, num_bodies).
+        body_com_pos_b: Input array of body COM positions in body frame. Shape is (num_envs, num_bodies).
+        body_link_pose_w: Input array of body link poses in world frame. Shape is (num_envs, num_bodies).
+        env_mask: Input array of environment mask. Shape is (num_envs,).
+        body_mask: Input array of body mask. Shape is (num_bodies,).
+        body_link_velocity_w: Output array where body link velocities are written. Shape is (num_envs, num_bodies).
+        body_com_velocity_w: Output array where body COM velocities (derived from link) are written.
+            Shape is (num_envs, num_bodies).
+        body_acc_w: Output array where body accelerations are zeroed. Shape is (num_envs, num_bodies).
+    """
+    i, j = wp.tid()
+    if env_mask[i] and body_mask[j]:
+        body_link_velocity_w[i, j] = data[i, j]
+        body_com_velocity_w[i, j] = get_link_velocity_in_com_frame_func(
+            body_link_velocity_w[i, j], body_link_pose_w[i, j], body_com_pos_b[i, j]
+        )
+        body_acc_w[i, j] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
+
+
 _SET_BODY_LINK_POSE_TO_SIM_DISPATCHER = IndexKernelDispatcher(set_body_link_pose_to_sim, ("env_ids", "body_ids"))
 _SET_BODY_COM_POSE_TO_SIM_DISPATCHER = IndexKernelDispatcher(set_body_com_pose_to_sim, ("env_ids", "body_ids"))
 _SET_BODY_COM_VELOCITY_TO_SIM_DISPATCHER = IndexKernelDispatcher(set_body_com_velocity_to_sim, ("env_ids", "body_ids"))
@@ -1034,21 +1118,18 @@ def write_2d_data_to_buffer_with_indices_kernel(
 
 
 @wp.kernel
-def write_2d_data_to_buffer_with_mask(
-    in_data: wp.array2d(dtype=wp.float32),
+def _write_2d_data_to_buffer_with_mask(
+    in_data: wp.array2d(dtype=Any),
     env_mask: wp.array(dtype=wp.bool),
     joint_mask: wp.array(dtype=wp.bool),
-    out_data: wp.array2d(dtype=wp.float32),
+    out_data: wp.array2d(dtype=Any),
 ):
-    """Write 2D float data to a buffer at specified indices.
-
-    This kernel copies float data from an input array to an output buffer at the specified
-    environment and joint/body indices.
+    """Write 2D data to a buffer where both masks are true.
 
     Args:
-        in_data: Input array containing float data. Shape is (num_instances, num_joints).
+        in_data: Input array. Shape is (num_instances, num_joints).
         env_mask: Input array of environment mask. Shape is (num_instances,).
-        joint_mask: Input array of joint/body mask. Shape is (num_instances, num_joints).
+        joint_mask: Input array of joint/body/tendon mask. Shape is (num_joints,).
         out_data: Output array where data is written. Shape is (num_instances, num_joints).
     """
     i, j = wp.tid()
@@ -1056,6 +1137,20 @@ def write_2d_data_to_buffer_with_mask(
         out_data[i, j] = in_data[i, j]
 
 
+@functools.cache
+def write_2d_data_to_buffer_with_mask_kernel(dtype: type) -> wp.Kernel:
+    """Select the masked 2-D buffer writer for the buffer's ``dtype``.
+
+    Launching the generic kernel directly resolves its types on every launch (about twice the host cost), and
+    cannot take torch input.
+    """
+    array_type = wp.array2d(dtype=dtype)
+    return wp.overload(_write_2d_data_to_buffer_with_mask, {"in_data": array_type, "out_data": array_type})
+
+
+write_2d_data_to_buffer_with_mask = write_2d_data_to_buffer_with_mask_kernel(wp.float32)
+
+
 @wp.func
 def _update_body_inertial_properties_inverse(
     env_id: int,
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py
index 3d1329cca985..8a725ca9f1ea 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py
@@ -15,6 +15,7 @@
 
 from isaaclab.assets.deformable_object.base_deformable_object import BaseDeformableObject
 from isaaclab.utils.warp import ProxyArray
+from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids
 
 from isaaclab_newton.sim.spawners.mpm.mpm import _SIMULATION_POINTS_SUFFIX
 
@@ -283,24 +284,10 @@ def _create_buffers(self):
         self._data.default_particle_state_w = self._data.default_nodal_state_w
 
     def _resolve_env_ids(self, env_ids):
-        if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)):
-            return self._ALL_INDICES
-        if isinstance(env_ids, slice):
-            return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids])
-        if isinstance(env_ids, torch.Tensor):
-            return wp.from_torch(env_ids.to(device=self.device, dtype=torch.int32), dtype=wp.int32)
-        if isinstance(env_ids, Sequence):
-            return wp.array(list(env_ids), dtype=wp.int32, device=self.device)
-        return env_ids
+        return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device, dtype=wp.int32)
 
     def _resolve_mask(self, mask: wp.array | torch.Tensor | None) -> wp.array:
-        if mask is None:
-            return self._ALL_ENV_MASK
-        if isinstance(mask, torch.Tensor):
-            if mask.dtype != torch.bool:
-                mask = mask.to(torch.bool)
-            return wp.from_torch(mask.to(device=self.device).contiguous(), dtype=wp.bool)
-        return mask
+        return resolve_1d_mask(mask=mask, all_mask=self._ALL_ENV_MASK, device=self.device)
 
     def _iter_particle_states(self):
         """Yield the Newton states whose particle arrays must receive writes."""
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py
index ae04842b9967..fc7e0c76a7c8 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py
@@ -21,6 +21,7 @@
 from isaaclab.assets.rigid_object.base_rigid_object import BaseRigidObject
 from isaaclab.sim.utils.queries import path_expr_to_glob, resolve_matching_prims_from_source
 from isaaclab.utils.warp import ProxyArray
+from isaaclab.utils.warp.utils import resolve_ids
 from isaaclab.utils.wrench_composer import WrenchComposer
 
 from isaaclab_newton.assets import kernels as shared_kernels
@@ -1122,13 +1123,7 @@ def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | No
         Returns:
             A warp array of environment indices or a tensor of environment indices.
         """
-        if (env_ids is None) or (env_ids == slice(None)):
-            return self._ALL_INDICES
-        if isinstance(env_ids, slice):
-            return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids])
-        if isinstance(env_ids, list):
-            return wp.array(env_ids, dtype=wp.int32, device=self.device)
-        return env_ids
+        return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device)
 
     def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor:
         """Resolve body indices to a warp array or tensor.
@@ -1139,13 +1134,7 @@ def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array |
         Returns:
             A warp array of body indices or a tensor of body indices.
         """
-        if isinstance(body_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if isinstance(body_ids, list):
-            return wp.array(body_ids, dtype=wp.int32, device=self.device)
-        if (body_ids is None) or (body_ids == slice(None)):
-            return self._ALL_BODY_INDICES
-        return body_ids
+        return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self.device)
 
     """
     Internal simulation callbacks.
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py
index 268e7a3e74d3..01ad44ba49d7 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py
@@ -22,6 +22,7 @@
 from isaaclab.assets.rigid_object_collection.base_rigid_object_collection import BaseRigidObjectCollection
 from isaaclab.utils import clone, validate
 from isaaclab.utils.warp import ProxyArray
+from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids
 from isaaclab.utils.wrench_composer import WrenchComposer
 
 from isaaclab_newton.assets import kernels as shared_kernels
@@ -308,16 +309,8 @@ def write_body_pose_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_ids = self._resolve_body_mask(body_mask)
-        else:
-            body_ids = self._ALL_BODY_INDICES
-        self.write_body_link_pose_to_sim_index(
-            body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward
+        self.write_body_link_pose_to_sim_mask(
+            body_poses=body_poses, body_mask=body_mask, env_mask=env_mask, skip_forward=skip_forward
         )
 
     def write_body_velocity_to_sim_index(
@@ -386,20 +379,8 @@ def write_body_velocity_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_ids = self._resolve_body_mask(body_mask)
-        else:
-            body_ids = self._ALL_BODY_INDICES
-        self.write_body_com_velocity_to_sim_index(
-            body_velocities=body_velocities,
-            env_ids=env_ids,
-            body_ids=body_ids,
-            full_data=True,
-            skip_forward=skip_forward,
+        self.write_body_com_velocity_to_sim_mask(
+            body_velocities=body_velocities, body_mask=body_mask, env_mask=env_mask, skip_forward=skip_forward
         )
 
     def write_body_link_pose_to_sim_index(
@@ -492,16 +473,25 @@ def write_body_link_pose_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_mask = self._ALL_ENV_MASK
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_ids = self._resolve_body_mask(body_mask)
-        self.write_body_link_pose_to_sim_index(
-            body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward
+        env_mask = resolve_1d_mask(mask=env_mask, all_mask=self._ALL_ENV_MASK, device=self.device)
+        # a body selection given as ids becomes a body mask on the device, so the same mask kernel applies
+        body_mask = resolve_1d_mask(
+            ids=body_ids,
+            mask=body_mask,
+            all_mask=self._ALL_BODY_MASK,
+            scratch_mask=self._scratch_body_mask,
+            device=self.device,
+        )
+        self.assert_shape_and_dtype_mask(body_poses, (env_mask, body_mask), wp.transformf, "body_poses")
+        wp.launch(
+            shared_kernels.write_2d_data_to_buffer_with_mask_kernel(wp.transformf),
+            dim=(self.num_instances, self.num_bodies),
+            inputs=[body_poses, env_mask, body_mask],
+            outputs=[self.data.body_link_pose_w],
+            device=self.device,
         )
+        if not skip_forward:
+            self.data._reset_pose(env_mask=env_mask)
 
     def write_body_com_pose_to_sim_index(
         self,
@@ -597,17 +587,25 @@ def write_body_com_pose_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_mask = self._ALL_ENV_MASK
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_ids = self._resolve_body_mask(body_mask)
-        # The index writer owns the invalidation (with from_link=False); forward skip_forward to it.
-        self.write_body_com_pose_to_sim_index(
-            body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward
+        env_mask = resolve_1d_mask(mask=env_mask, all_mask=self._ALL_ENV_MASK, device=self.device)
+        # a body selection given as ids becomes a body mask on the device, so the same mask kernel applies
+        body_mask = resolve_1d_mask(
+            ids=body_ids,
+            mask=body_mask,
+            all_mask=self._ALL_BODY_MASK,
+            scratch_mask=self._scratch_body_mask,
+            device=self.device,
+        )
+        self.assert_shape_and_dtype_mask(body_poses, (env_mask, body_mask), wp.transformf, "body_poses")
+        wp.launch(
+            shared_kernels.set_body_com_pose_to_sim_mask,
+            dim=(self.num_instances, self.num_bodies),
+            inputs=[body_poses, self.data.body_com_pos_b, env_mask, body_mask],
+            outputs=[self.data.body_com_pose_w, self.data.body_link_pose_w],
+            device=self.device,
         )
+        if not skip_forward:
+            self.data._reset_pose(env_mask=env_mask, from_link=False)
 
     def write_body_com_velocity_to_sim_index(
         self,
@@ -711,21 +709,25 @@ def write_body_com_velocity_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_mask = self._ALL_ENV_MASK
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_ids = self._resolve_body_mask(body_mask)
-        # The index writer owns the invalidation; forward skip_forward to it.
-        self.write_body_com_velocity_to_sim_index(
-            body_velocities=body_velocities,
-            env_ids=env_ids,
-            body_ids=body_ids,
-            full_data=True,
-            skip_forward=skip_forward,
+        env_mask = resolve_1d_mask(mask=env_mask, all_mask=self._ALL_ENV_MASK, device=self.device)
+        # a body selection given as ids becomes a body mask on the device, so the same mask kernel applies
+        body_mask = resolve_1d_mask(
+            ids=body_ids,
+            mask=body_mask,
+            all_mask=self._ALL_BODY_MASK,
+            scratch_mask=self._scratch_body_mask,
+            device=self.device,
         )
+        self.assert_shape_and_dtype_mask(body_velocities, (env_mask, body_mask), wp.spatial_vectorf, "body_velocities")
+        wp.launch(
+            shared_kernels.set_body_com_velocity_to_sim_mask,
+            dim=(self.num_instances, self.num_bodies),
+            inputs=[body_velocities, env_mask, body_mask],
+            outputs=[self.data.body_com_vel_w, self.data.body_com_acc_w],
+            device=self.device,
+        )
+        if not skip_forward:
+            self.data._reset_velocity(env_mask=env_mask)
 
     def write_body_link_velocity_to_sim_index(
         self,
@@ -831,21 +833,25 @@ def write_body_link_velocity_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_mask = self._ALL_ENV_MASK
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_ids = self._resolve_body_mask(body_mask)
-        # The index writer owns the invalidation (with from_com=False); forward skip_forward to it.
-        self.write_body_link_velocity_to_sim_index(
-            body_velocities=body_velocities,
-            env_ids=env_ids,
-            body_ids=body_ids,
-            full_data=True,
-            skip_forward=skip_forward,
+        env_mask = resolve_1d_mask(mask=env_mask, all_mask=self._ALL_ENV_MASK, device=self.device)
+        # a body selection given as ids becomes a body mask on the device, so the same mask kernel applies
+        body_mask = resolve_1d_mask(
+            ids=body_ids,
+            mask=body_mask,
+            all_mask=self._ALL_BODY_MASK,
+            scratch_mask=self._scratch_body_mask,
+            device=self.device,
+        )
+        self.assert_shape_and_dtype_mask(body_velocities, (env_mask, body_mask), wp.spatial_vectorf, "body_velocities")
+        wp.launch(
+            shared_kernels.set_body_link_velocity_to_sim_mask,
+            dim=(self.num_instances, self.num_bodies),
+            inputs=[body_velocities, self.data.body_com_pos_b, self.data.body_link_pose_w, env_mask, body_mask],
+            outputs=[self.data.body_link_vel_w, self.data.body_com_vel_w, self.data.body_com_acc_w],
+            device=self.device,
         )
+        if not skip_forward:
+            self.data._reset_velocity(env_mask=env_mask, from_com=False)
 
     """
     Operations - Setters.
@@ -1244,6 +1250,7 @@ def _create_buffers(self):
         )
         self._ALL_ENV_MASK = wp.ones((self.num_instances,), dtype=wp.bool, device=self.device)
         self._ALL_BODY_MASK = wp.ones((self.num_bodies,), dtype=wp.bool, device=self.device)
+        self._scratch_body_mask = wp.zeros((self.num_bodies,), dtype=wp.bool, device=self.device)
 
         # external wrench composer
         self._instantaneous_wrench_composer = WrenchComposer(self)
@@ -1285,13 +1292,7 @@ def _resolve_env_ids(self, env_ids) -> wp.array | torch.Tensor:
         Returns:
             Environment indices.
         """
-        if (env_ids is None) or (env_ids == slice(None)):
-            return self._ALL_ENV_INDICES
-        if isinstance(env_ids, slice):
-            return wp.from_torch(wp.to_torch(self._ALL_ENV_INDICES)[env_ids])
-        if isinstance(env_ids, list):
-            return wp.array(env_ids, dtype=wp.int32, device=self.device)
-        return env_ids
+        return resolve_ids(ids=env_ids, all_ids=self._ALL_ENV_INDICES, device=self.device)
 
     def _resolve_body_ids(self, body_ids) -> wp.array | torch.Tensor:
         """Resolve body indices.
@@ -1302,37 +1303,7 @@ def _resolve_body_ids(self, body_ids) -> wp.array | torch.Tensor:
         Returns:
             Body indices.
         """
-        if isinstance(body_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if isinstance(body_ids, list):
-            return wp.array(body_ids, dtype=wp.int32, device=self.device)
-        if (body_ids is None) or (body_ids == slice(None)):
-            return self._ALL_BODY_INDICES
-        if isinstance(body_ids, slice):
-            return wp.from_torch(
-                torch.arange(self.num_bodies, dtype=torch.int32, device=self.device)[body_ids], dtype=wp.int32
-            )
-        return body_ids
-
-    def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array | torch.Tensor:
-        """Resolve environment mask to indices via torch.nonzero."""
-        if env_mask is not None:
-            if isinstance(env_mask, wp.array):
-                env_mask = wp.to_torch(env_mask)
-            env_ids = torch.nonzero(env_mask)[:, 0]
-        else:
-            env_ids = self._ALL_ENV_INDICES
-        return env_ids
-
-    def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array | torch.Tensor:
-        """Resolve body mask to indices via torch.nonzero."""
-        if body_mask is not None:
-            if isinstance(body_mask, wp.array):
-                body_mask = wp.to_torch(body_mask)
-            body_ids = torch.nonzero(body_mask)[:, 0]
-        else:
-            body_ids = self._ALL_BODY_INDICES
-        return body_ids
+        return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self.device)
 
     """
     Internal simulation callbacks.

From 7b48c38e3c2417e6e458b46a82d7ef425e047065 Mon Sep 17 00:00:00 2001
From: Jichuan Hu <jichuanh@nvidia.com>
Date: Mon, 26 Oct 2026 02:48:21 -0700
Subject: [PATCH 46/48] Resolve OVPhysX selections through the shared helpers

The OVPhysX assets resolve env, body, joint and tendon selections through resolve_ids and resolve_1d_mask instead of their own copies.
---
 .../changelog.d/jichuanh-id-mask.changed.rst  |   7 ++
 .../assets/articulation/articulation.py       | 103 ++---------------
 .../deformable_object/deformable_object.py    |  18 +--
 .../assets/rigid_object/rigid_object.py       |  37 +-----
 .../rigid_object_collection.py                | 105 +++---------------
 .../test/assets/test_deformable_object.py     |  22 ++++
 6 files changed, 65 insertions(+), 227 deletions(-)
 create mode 100644 source/isaaclab_ov/changelog.d/jichuanh-id-mask.changed.rst

diff --git a/source/isaaclab_ov/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab_ov/changelog.d/jichuanh-id-mask.changed.rst
new file mode 100644
index 000000000000..3f1ff2245d66
--- /dev/null
+++ b/source/isaaclab_ov/changelog.d/jichuanh-id-mask.changed.rst
@@ -0,0 +1,7 @@
+* Changed the env selection helpers of the OVPhysX assets to resolve ``env_ids`` through
+  :func:`~isaaclab.utils.warp.utils.resolve_ids` and ``env_mask`` through
+  :func:`~isaaclab.utils.warp.utils.resolve_1d_mask`. Any integer sequence and slices with a negative step are accepted,
+  and a :class:`~isaaclab.utils.warp.ProxyArray` is rejected with a :class:`TypeError`.
+* Changed the body, joint and tendon selection helpers of the OVPhysX assets to resolve ids and masks through
+  :func:`~isaaclab.utils.warp.utils.resolve_ids` and :func:`~isaaclab.utils.warp.utils.resolve_1d_mask`, as the env
+  selection does, so slices other than ``slice(None)`` are accepted.
diff --git a/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py b/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py
index 07bc71d3f611..2162ae04b8ba 100644
--- a/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py
+++ b/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py
@@ -33,6 +33,7 @@
 from isaaclab.utils.string import resolve_matching_names
 from isaaclab.utils.warp import ProxyArray
 from isaaclab.utils.warp import kernels as warp_kernels
+from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids
 from isaaclab.utils.wrench_composer import WrenchComposer
 
 from isaaclab_ov import tensor_types as TT
@@ -4373,79 +4374,27 @@ def _log_articulation_info(self) -> None:
 
     def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor:
         """Resolve environment indices on ``self._device``."""
-        if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)):
-            return self._ALL_INDICES
-        if isinstance(env_ids, slice):
-            return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids])
-        if isinstance(env_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if isinstance(env_ids, list):
-            return wp.array(env_ids, dtype=wp.int32, device=self._device)
-        if isinstance(env_ids, torch.Tensor):
-            return env_ids.to(device=self._device)
-        if isinstance(env_ids, wp.array) and str(env_ids.device) != self._device:
-            env_ids = wp.clone(env_ids, device=self._device)
-        return env_ids
+        return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self._device)
 
     def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor:
         """Resolve body indices to a Warp signed-integer array on ``self._device``."""
-        if isinstance(body_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if body_ids is None or body_ids == slice(None):
-            return self._ALL_BODY_INDICES
-        if isinstance(body_ids, list):
-            return wp.array(body_ids, dtype=wp.int32, device=self._device)
-        if isinstance(body_ids, torch.Tensor):
-            return body_ids.to(device=self._device)
-        if isinstance(body_ids, wp.array) and str(body_ids.device) != self._device:
-            body_ids = wp.clone(body_ids, device=self._device)
-        return body_ids
+        return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self._device)
 
     def _resolve_joint_ids(self, joint_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor:
         """Resolve joint indices to a Warp signed-integer array on ``self._device``."""
-        if isinstance(joint_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if joint_ids is None or joint_ids == slice(None):
-            return self._ALL_JOINT_INDICES
-        if isinstance(joint_ids, list):
-            return wp.array(joint_ids, dtype=wp.int32, device=self._device)
-        if isinstance(joint_ids, torch.Tensor):
-            return joint_ids.to(device=self._device)
-        if isinstance(joint_ids, wp.array) and str(joint_ids.device) != self._device:
-            joint_ids = wp.clone(joint_ids, device=self._device)
-        return joint_ids
+        return resolve_ids(ids=joint_ids, all_ids=self._ALL_JOINT_INDICES, device=self._device)
 
     def _resolve_fixed_tendon_ids(
         self, tendon_ids: Sequence[int] | torch.Tensor | wp.array | None
     ) -> wp.array | torch.Tensor:
         """Resolve fixed-tendon indices to a Warp signed-integer array on ``self._device``."""
-        if isinstance(tendon_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if tendon_ids is None or tendon_ids == slice(None):
-            return self._ALL_FIXED_TENDON_INDICES
-        if isinstance(tendon_ids, list):
-            return wp.array(tendon_ids, dtype=wp.int32, device=self._device)
-        if isinstance(tendon_ids, torch.Tensor):
-            return tendon_ids.to(device=self._device)
-        if isinstance(tendon_ids, wp.array) and str(tendon_ids.device) != self._device:
-            tendon_ids = wp.clone(tendon_ids, device=self._device)
-        return tendon_ids
+        return resolve_ids(ids=tendon_ids, all_ids=self._ALL_FIXED_TENDON_INDICES, device=self._device)
 
     def _resolve_spatial_tendon_ids(
         self, tendon_ids: Sequence[int] | torch.Tensor | wp.array | None
     ) -> wp.array | torch.Tensor:
         """Resolve spatial-tendon indices to a Warp signed-integer array on ``self._device``."""
-        if isinstance(tendon_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if tendon_ids is None or tendon_ids == slice(None):
-            return self._ALL_SPATIAL_TENDON_INDICES
-        if isinstance(tendon_ids, list):
-            return wp.array(tendon_ids, dtype=wp.int32, device=self._device)
-        if isinstance(tendon_ids, torch.Tensor):
-            return tendon_ids.to(device=self._device)
-        if isinstance(tendon_ids, wp.array) and str(tendon_ids.device) != self._device:
-            tendon_ids = wp.clone(tendon_ids, device=self._device)
-        return tendon_ids
+        return resolve_ids(ids=tendon_ids, all_ids=self._ALL_SPATIAL_TENDON_INDICES, device=self._device)
 
     def _broadcast_scalar_to_2d(
         self, value: float | torch.Tensor | wp.array, shape: tuple[int, int]
@@ -4483,13 +4432,7 @@ def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array:
         no ``torch.nonzero`` conversion is needed. ``None`` returns the pre-allocated
         all-true mask.
         """
-        if env_mask is None:
-            return self._ALL_TRUE_ENV_MASK
-        if isinstance(env_mask, torch.Tensor):
-            return wp.from_torch(env_mask.to(torch.bool), dtype=wp.bool)
-        if isinstance(env_mask, wp.array) and str(env_mask.device) != self._device:
-            env_mask = wp.clone(env_mask, device=self._device)
-        return env_mask
+        return resolve_1d_mask(mask=env_mask, all_mask=self._ALL_TRUE_ENV_MASK, device=self._device)
 
     def _sim_env_ids_view(self, count: int) -> wp.array:
         """Return a cached prefix view of the simulator-index scratch buffer."""
@@ -4505,43 +4448,19 @@ def _sim_env_ids_view(self, count: int) -> wp.array:
 
     def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array:
         """Resolve a body mask to a ``wp.bool`` array on ``self._device`` (Newton-style)."""
-        if body_mask is None:
-            return self._ALL_TRUE_BODY_MASK
-        if isinstance(body_mask, torch.Tensor):
-            return wp.from_torch(body_mask.to(torch.bool), dtype=wp.bool)
-        if isinstance(body_mask, wp.array) and str(body_mask.device) != self._device:
-            body_mask = wp.clone(body_mask, device=self._device)
-        return body_mask
+        return resolve_1d_mask(mask=body_mask, all_mask=self._ALL_TRUE_BODY_MASK, device=self._device)
 
     def _resolve_joint_mask(self, joint_mask: wp.array | None) -> wp.array:
         """Resolve a joint mask to a ``wp.bool`` array on ``self._device``."""
-        if joint_mask is None:
-            return self._ALL_TRUE_JOINT_MASK
-        if isinstance(joint_mask, torch.Tensor):
-            return wp.from_torch(joint_mask.to(torch.bool), dtype=wp.bool)
-        if isinstance(joint_mask, wp.array) and str(joint_mask.device) != self._device:
-            joint_mask = wp.clone(joint_mask, device=self._device)
-        return joint_mask
+        return resolve_1d_mask(mask=joint_mask, all_mask=self._ALL_TRUE_JOINT_MASK, device=self._device)
 
     def _resolve_fixed_tendon_mask(self, tendon_mask: wp.array | None) -> wp.array:
         """Resolve a fixed-tendon mask to a ``wp.bool`` array on ``self._device``."""
-        if tendon_mask is None:
-            return self._ALL_TRUE_FIXED_TENDON_MASK
-        if isinstance(tendon_mask, torch.Tensor):
-            return wp.from_torch(tendon_mask.to(torch.bool), dtype=wp.bool)
-        if isinstance(tendon_mask, wp.array) and str(tendon_mask.device) != self._device:
-            tendon_mask = wp.clone(tendon_mask, device=self._device)
-        return tendon_mask
+        return resolve_1d_mask(mask=tendon_mask, all_mask=self._ALL_TRUE_FIXED_TENDON_MASK, device=self._device)
 
     def _resolve_spatial_tendon_mask(self, tendon_mask: wp.array | None) -> wp.array:
         """Resolve a spatial-tendon mask to a ``wp.bool`` array on ``self._device``."""
-        if tendon_mask is None:
-            return self._ALL_TRUE_SPATIAL_TENDON_MASK
-        if isinstance(tendon_mask, torch.Tensor):
-            return wp.from_torch(tendon_mask.to(torch.bool), dtype=wp.bool)
-        if isinstance(tendon_mask, wp.array) and str(tendon_mask.device) != self._device:
-            tendon_mask = wp.clone(tendon_mask, device=self._device)
-        return tendon_mask
+        return resolve_1d_mask(mask=tendon_mask, all_mask=self._ALL_TRUE_SPATIAL_TENDON_MASK, device=self._device)
 
     def _get_cpu_env_mask(self, env_mask: wp.array) -> wp.array:
         """Return a pinned-host CPU copy of :paramref:`env_mask` for a CPU-only binding write.
diff --git a/source/isaaclab_ov/isaaclab_ov/assets/deformable_object/deformable_object.py b/source/isaaclab_ov/isaaclab_ov/assets/deformable_object/deformable_object.py
index ccca84b3cd0a..ceef9f0e2049 100644
--- a/source/isaaclab_ov/isaaclab_ov/assets/deformable_object/deformable_object.py
+++ b/source/isaaclab_ov/isaaclab_ov/assets/deformable_object/deformable_object.py
@@ -22,6 +22,7 @@
 from isaaclab.assets.deformable_object.base_deformable_object import BaseDeformableObject
 from isaaclab.markers import VisualizationMarkers
 from isaaclab.utils.warp import ProxyArray
+from isaaclab.utils.warp.utils import resolve_ids
 
 from isaaclab_ov import tensor_types as TT
 from isaaclab_ov.sim.views.ovphysx_view import OvPhysxView
@@ -528,22 +529,7 @@ def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array(dtyp
         dtype=wp.int32
     ):
         """Resolve environment indices to a device-resident int32 Warp array."""
-        if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)):
-            return self._ALL_INDICES
-        if isinstance(env_ids, slice):
-            # OVPhysX consumes these IDs directly and requires contiguous DLPack tensors.
-            return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids].contiguous())
-        if isinstance(env_ids, torch.Tensor):
-            values = env_ids.to(device=self.device, dtype=torch.int32).contiguous()
-            return wp.from_torch(values, dtype=wp.int32)
-        if isinstance(env_ids, wp.array):
-            if env_ids.dtype != wp.int32:
-                values = wp.to_torch(env_ids).to(device=self.device, dtype=torch.int32).contiguous()
-                return wp.from_torch(values, dtype=wp.int32)
-            if str(env_ids.device) != self.device:
-                return wp.clone(env_ids, device=self.device)
-            return env_ids
-        return wp.array(list(env_ids), dtype=wp.int32, device=self.device)
+        return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device, dtype=wp.int32)
 
     @staticmethod
     def _arrays_overlap(first: wp.array, second: wp.array) -> bool:
diff --git a/source/isaaclab_ov/isaaclab_ov/assets/rigid_object/rigid_object.py b/source/isaaclab_ov/isaaclab_ov/assets/rigid_object/rigid_object.py
index 7e19bc75109d..a44665429701 100644
--- a/source/isaaclab_ov/isaaclab_ov/assets/rigid_object/rigid_object.py
+++ b/source/isaaclab_ov/isaaclab_ov/assets/rigid_object/rigid_object.py
@@ -22,6 +22,7 @@
 from isaaclab.sim.utils.queries import path_expr_to_glob, resolve_matching_prims_from_source
 from isaaclab.utils.string import resolve_matching_names
 from isaaclab.utils.warp import ProxyArray
+from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids
 from isaaclab.utils.wrench_composer import WrenchComposer
 
 from isaaclab_ov import tensor_types as TT
@@ -1076,17 +1077,7 @@ def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | No
         Returns:
             Environment indices on ``self._device``.
         """
-        if env_ids is None or env_ids == slice(None):
-            return self._ALL_INDICES
-        if isinstance(env_ids, slice):
-            return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids])
-        if isinstance(env_ids, list):
-            return wp.array(env_ids, dtype=wp.int32, device=self._device)
-        if isinstance(env_ids, torch.Tensor):
-            return env_ids.to(device=self._device)
-        if isinstance(env_ids, wp.array) and str(env_ids.device) != self._device:
-            env_ids = wp.clone(env_ids, device=self._device)
-        return env_ids
+        return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self._device)
 
     def _sim_env_ids_view(self, count: int) -> wp.array:
         """Return a cached prefix of the simulator-index scratch buffer."""
@@ -1109,13 +1100,7 @@ def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array |
         Returns:
             Body indices.
         """
-        if isinstance(body_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if body_ids is None or body_ids == slice(None):
-            return self._ALL_BODY_INDICES
-        if isinstance(body_ids, list):
-            return wp.array(body_ids, dtype=wp.int32, device=self._device)
-        return body_ids
+        return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self._device)
 
     def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array:
         """Resolve an environment mask to a ``wp.bool`` array.
@@ -1126,13 +1111,7 @@ def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array:
         Returns:
             A warp ``wp.bool`` array on ``self._device``.
         """
-        if env_mask is None:
-            return self._ALL_TRUE_ENV_MASK
-        if isinstance(env_mask, torch.Tensor):
-            return wp.from_torch(env_mask.to(torch.bool), dtype=wp.bool)
-        if isinstance(env_mask, wp.array) and str(env_mask.device) != self._device:
-            env_mask = wp.clone(env_mask, device=self._device)
-        return env_mask
+        return resolve_1d_mask(mask=env_mask, all_mask=self._ALL_TRUE_ENV_MASK, device=self._device)
 
     def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array:
         """Resolve a body mask to a ``wp.bool`` array.
@@ -1143,13 +1122,7 @@ def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array:
         Returns:
             A warp ``wp.bool`` array on ``self._device``.
         """
-        if body_mask is None:
-            return self._ALL_TRUE_BODY_MASK
-        if isinstance(body_mask, torch.Tensor):
-            return wp.from_torch(body_mask.to(torch.bool), dtype=wp.bool)
-        if isinstance(body_mask, wp.array) and str(body_mask.device) != self._device:
-            body_mask = wp.clone(body_mask, device=self._device)
-        return body_mask
+        return resolve_1d_mask(mask=body_mask, all_mask=self._ALL_TRUE_BODY_MASK, device=self._device)
 
     def _get_cpu_env_mask(self, env_mask: wp.array) -> wp.array:
         """Return a pinned-host CPU copy of *env_mask* for a CPU-only binding write.
diff --git a/source/isaaclab_ov/isaaclab_ov/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_ov/isaaclab_ov/assets/rigid_object_collection/rigid_object_collection.py
index c39c8c6d2f2b..1d4066c2709a 100644
--- a/source/isaaclab_ov/isaaclab_ov/assets/rigid_object_collection/rigid_object_collection.py
+++ b/source/isaaclab_ov/isaaclab_ov/assets/rigid_object_collection/rigid_object_collection.py
@@ -21,6 +21,7 @@
 from isaaclab.utils import clone, validate
 from isaaclab.utils.string import resolve_matching_names
 from isaaclab.utils.warp import ProxyArray
+from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids
 from isaaclab.utils.wrench_composer import WrenchComposer
 
 from isaaclab_ov import tensor_types as TT
@@ -439,16 +440,8 @@ def write_body_link_pose_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask
-            env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0])
-        else:
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_mask_t = wp.to_torch(body_mask) if isinstance(body_mask, wp.array) else body_mask
-            body_ids = self._resolve_body_ids(torch.nonzero(body_mask_t)[:, 0])
-        else:
-            body_ids = self._ALL_BODY_INDICES
+        env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device)
+        body_ids = resolve_ids(mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self._device)
         self.assert_shape_and_dtype(body_poses, (self._num_instances, self._num_bodies), wp.transformf, "body_poses")
         wp.launch(
             shared_kernels.set_body_link_pose_to_sim_kernel(env_ids, body_ids),
@@ -535,16 +528,8 @@ def write_body_com_pose_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask
-            env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0])
-        else:
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_mask_t = wp.to_torch(body_mask) if isinstance(body_mask, wp.array) else body_mask
-            body_ids = self._resolve_body_ids(torch.nonzero(body_mask_t)[:, 0])
-        else:
-            body_ids = self._ALL_BODY_INDICES
+        env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device)
+        body_ids = resolve_ids(mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self._device)
         self.assert_shape_and_dtype(body_poses, (self._num_instances, self._num_bodies), wp.transformf, "body_poses")
         wp.launch(
             shared_kernels.set_body_com_pose_to_sim_kernel(env_ids, body_ids),
@@ -643,16 +628,8 @@ def write_body_com_velocity_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask
-            env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0])
-        else:
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_mask_t = wp.to_torch(body_mask) if isinstance(body_mask, wp.array) else body_mask
-            body_ids = self._resolve_body_ids(torch.nonzero(body_mask_t)[:, 0])
-        else:
-            body_ids = self._ALL_BODY_INDICES
+        env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device)
+        body_ids = resolve_ids(mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self._device)
         self.assert_shape_and_dtype(
             body_velocities, (self._num_instances, self._num_bodies), wp.spatial_vectorf, "body_velocities"
         )
@@ -758,16 +735,8 @@ def write_body_link_velocity_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask
-            env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0])
-        else:
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_mask_t = wp.to_torch(body_mask) if isinstance(body_mask, wp.array) else body_mask
-            body_ids = self._resolve_body_ids(torch.nonzero(body_mask_t)[:, 0])
-        else:
-            body_ids = self._ALL_BODY_INDICES
+        env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device)
+        body_ids = resolve_ids(mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self._device)
         self.assert_shape_and_dtype(
             body_velocities, (self._num_instances, self._num_bodies), wp.spatial_vectorf, "body_velocities"
         )
@@ -869,11 +838,7 @@ def set_masses_mask(
             env_mask: Environment mask. If None, all instances are updated.
                 Shape is (num_instances,).
         """
-        if env_mask is not None:
-            env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask
-            env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0])
-        else:
-            env_ids = self._ALL_ENV_INDICES
+        env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device)
         self.assert_shape_and_dtype(masses, (self._num_instances, self._num_bodies), wp.float32, "masses")
         wp.launch(
             shared_kernels.write_2d_data_to_buffer_with_mask,
@@ -955,11 +920,7 @@ def set_coms_mask(
             env_mask: Environment mask. If None, all instances are updated.
                 Shape is (num_instances,).
         """
-        if env_mask is not None:
-            env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask
-            env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0])
-        else:
-            env_ids = self._ALL_ENV_INDICES
+        env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device)
         self.assert_shape_and_dtype(coms, (self._num_instances, self._num_bodies), wp.transformf, "coms")
         wp.launch(
             shared_kernels.write_body_com_pose_to_buffer_mask,
@@ -1051,11 +1012,7 @@ def set_inertias_mask(
             env_mask: Environment mask. If None, all instances are updated.
                 Shape is (num_instances,).
         """
-        if env_mask is not None:
-            env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask
-            env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0])
-        else:
-            env_ids = self._ALL_ENV_INDICES
+        env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device)
         self.assert_shape_and_dtype(inertias, (self._num_instances, self._num_bodies, 9), wp.float32, "inertias")
         wp.launch(
             shared_kernels.write_body_inertia_to_buffer_mask,
@@ -1388,27 +1345,11 @@ def _resolve_env_ids(self, env_ids) -> wp.array | torch.Tensor:
         resolved array onto ``self._device`` so kernel launches don't fail on a device
         mismatch.
         """
-        if env_ids is None or env_ids == slice(None):
-            return self._ALL_ENV_INDICES
-        if isinstance(env_ids, slice):
-            return wp.from_torch(wp.to_torch(self._ALL_ENV_INDICES)[env_ids])
-        if isinstance(env_ids, list):
-            return wp.array(env_ids, dtype=wp.int32, device=self._device)
-        if isinstance(env_ids, torch.Tensor):
-            return env_ids.to(device=self._device)
-        if isinstance(env_ids, wp.array) and str(env_ids.device) != self._device:
-            env_ids = wp.clone(env_ids, device=self._device)
-        return env_ids
+        return resolve_ids(ids=env_ids, all_ids=self._ALL_ENV_INDICES, device=self._device)
 
     def _resolve_body_ids(self, body_ids) -> wp.array | torch.Tensor:
         """Resolve body indices on ``self._device``."""
-        if isinstance(body_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if body_ids is None or body_ids == slice(None):
-            return self._ALL_BODY_INDICES
-        if isinstance(body_ids, list):
-            return wp.array(body_ids, dtype=wp.int32, device=self._device)
-        return body_ids
+        return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self._device)
 
     def _env_body_ids_to_view_ids(
         self, env_ids: torch.Tensor | wp.array, body_ids: torch.Tensor | wp.array, device: str = "cuda:0"
@@ -1418,6 +1359,8 @@ def _env_body_ids_to_view_ids(
         Computes ``view_id = body_id * num_instances + env_id`` for each
         (env_id, body_id) pair.  The output array is laid out column-major over
         the (env, body) grid, matching the PhysX ``root_view`` ordering.
+        The mask writers convert their masks to ids for this: a binding write mask must cover all
+        ``num_bodies * num_instances`` rows, so an env mask does not fit.
 
         Args:
             env_ids: Environment indices.
@@ -1466,13 +1409,7 @@ def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array:
         Returns:
             A ``wp.bool`` array of shape (num_instances,) on ``self._device``.
         """
-        if env_mask is None:
-            return self._ALL_TRUE_ENV_MASK
-        if isinstance(env_mask, torch.Tensor):
-            return wp.from_torch(env_mask.to(torch.bool), dtype=wp.bool)
-        if isinstance(env_mask, wp.array) and str(env_mask.device) != self._device:
-            env_mask = wp.clone(env_mask, device=self._device)
-        return env_mask
+        return resolve_1d_mask(mask=env_mask, all_mask=self._ALL_TRUE_ENV_MASK, device=self._device)
 
     def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array:
         """Resolve a body mask to a ``wp.bool`` array on ``self._device``.
@@ -1485,13 +1422,7 @@ def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array:
         Returns:
             A ``wp.bool`` array of shape (num_bodies,) on ``self._device``.
         """
-        if body_mask is None:
-            return self._ALL_TRUE_BODY_MASK
-        if isinstance(body_mask, torch.Tensor):
-            return wp.from_torch(body_mask.to(torch.bool), dtype=wp.bool)
-        if isinstance(body_mask, wp.array) and str(body_mask.device) != self._device:
-            body_mask = wp.clone(body_mask, device=self._device)
-        return body_mask
+        return resolve_1d_mask(mask=body_mask, all_mask=self._ALL_TRUE_BODY_MASK, device=self._device)
 
     def _sim_env_ids_view(self, count: int) -> wp.array:
         """Return a cached prefix of the simulator-index scratch buffer."""
diff --git a/source/isaaclab_ov/test/assets/test_deformable_object.py b/source/isaaclab_ov/test/assets/test_deformable_object.py
index bc9fc6a2e04b..5ea0b8ea84f6 100644
--- a/source/isaaclab_ov/test/assets/test_deformable_object.py
+++ b/source/isaaclab_ov/test/assets/test_deformable_object.py
@@ -386,6 +386,28 @@ def test_indexed_partial_write_preserves_retained_aliased_slice(
     torch.testing.assert_close(retained.torch, expected)
 
 
+@pytest.mark.usefixtures("host_shell")
+def test_mask_writes_update_only_masked_rows():
+    """The mask writers take full-size data and write only the environments the mask selects."""
+    asset = _make_asset_shell(deformable_type="volume", num_instances=3, num_vertices=2)
+    zeros = torch.zeros((3, 2, 3), device=asset.device)
+    asset.root_view.positions = wp.from_torch(zeros.clone(), dtype=wp.float32)
+    asset.root_view.velocities = wp.from_torch(zeros.clone(), dtype=wp.float32)
+    asset.update(0.1)
+    env_mask = wp.array([False, True, False], dtype=wp.bool, device=asset.device)
+
+    state = torch.cat((torch.full((3, 2, 3), 5.0), torch.full((3, 2, 3), -5.0)), dim=-1).to(asset.device)
+    asset.write_nodal_state_to_sim_mask(state, env_mask=env_mask)
+    expected = torch.zeros((3, 2, 6), device=asset.device)
+    expected[1] = state[1]
+    torch.testing.assert_close(asset.data.nodal_state_w.torch, expected)
+
+    positions = torch.full((3, 2, 3), 9.0, device=asset.device)
+    asset.write_nodal_pos_to_sim_mask(positions, env_mask=env_mask)
+    expected[1, :, :3] = 9.0
+    torch.testing.assert_close(asset.data.nodal_pos_w.torch, expected[..., :3])
+
+
 @pytest.mark.usefixtures("host_shell")
 def test_full_overwrite_stale_cache_does_not_read_simulator() -> None:
     asset = _make_asset_shell(deformable_type="volume", num_instances=3, num_vertices=2)

From 3d3661f29616ab829138b9dbdbbbf81f1b79e4e0 Mon Sep 17 00:00:00 2001
From: Jichuan Hu <jichuanh@nvidia.com>
Date: Mon, 26 Oct 2026 02:48:22 -0700
Subject: [PATCH 47/48] Resolve PhysX selections through the shared helpers

The PhysX assets resolve env, body, joint and tendon selections through resolve_ids and resolve_1d_mask instead of their own copies. The sensor test doubles get the env masks a real sensor creates.
---
 .../changelog.d/jichuanh-id-mask.changed.rst  |   7 +
 .../assets/articulation/articulation.py       | 282 ++++++------------
 .../assets/rigid_object/rigid_object.py       | 113 ++-----
 .../rigid_object_collection.py                | 119 ++------
 .../assets/surface_gripper/surface_gripper.py |  29 +-
 .../test/sensors/test_frame_transformer.py    |   2 +
 .../sensors/test_imu_pva_recorded_launch.py   |   2 +
 7 files changed, 154 insertions(+), 400 deletions(-)
 create mode 100644 source/isaaclab_physx/changelog.d/jichuanh-id-mask.changed.rst

diff --git a/source/isaaclab_physx/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab_physx/changelog.d/jichuanh-id-mask.changed.rst
new file mode 100644
index 000000000000..fb8da731a750
--- /dev/null
+++ b/source/isaaclab_physx/changelog.d/jichuanh-id-mask.changed.rst
@@ -0,0 +1,7 @@
+* Changed the env selection helpers of the PhysX assets to resolve ``env_ids`` through
+  :func:`~isaaclab.utils.warp.utils.resolve_ids`. Any integer sequence and slices with a negative step are accepted,
+  ids on another device are moved to the asset's device, and a :class:`~isaaclab.utils.warp.ProxyArray` is rejected
+  with a :class:`TypeError`.
+* Changed the body, joint and tendon selection helpers of the PhysX assets to resolve ids and masks through
+  :func:`~isaaclab.utils.warp.utils.resolve_ids`, as the env selection does, so slices other than ``slice(None)`` are
+  accepted.
diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py
index 3d2cb4f1d946..b1be80852be8 100644
--- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py
+++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py
@@ -27,6 +27,7 @@
 from isaaclab.utils.string import resolve_matching_names, resolve_matching_names_values
 from isaaclab.utils.version import get_isaac_sim_version, has_kit
 from isaaclab.utils.warp import ProxyArray
+from isaaclab.utils.warp.utils import resolve_ids
 from isaaclab.utils.wrench_composer import WrenchComposer
 
 from isaaclab_physx.assets import kernels as shared_kernels
@@ -590,7 +591,7 @@ def write_root_link_pose_to_sim_mask(
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
         # resolve masks
-        env_ids = self._resolve_env_mask(env_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
         # Set full data to True to ensure the the right code path is taken inside the kernel.
         self.write_root_link_pose_to_sim_index(
             root_pose=root_pose, env_ids=env_ids, full_data=True, skip_forward=skip_forward
@@ -692,7 +693,7 @@ def write_root_com_pose_to_sim_mask(
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
         # resolve masks
-        env_ids = self._resolve_env_mask(env_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
         # Set full data to True to ensure the the right code path is taken inside the kernel.
         self.write_root_com_pose_to_sim_index(
             root_pose=root_pose, env_ids=env_ids, full_data=True, skip_forward=skip_forward
@@ -866,7 +867,7 @@ def write_root_com_velocity_to_sim_mask(
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
         # resolve masks
-        env_ids = self._resolve_env_mask(env_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
         # Set full data to True to ensure the the right code path is taken inside the kernel.
         self.write_root_com_velocity_to_sim_index(
             root_velocity=root_velocity, env_ids=env_ids, full_data=True, skip_forward=skip_forward
@@ -977,7 +978,7 @@ def write_root_link_velocity_to_sim_mask(
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
         # resolve masks
-        env_ids = self._resolve_env_mask(env_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
         # Set full data to True to ensure the the right code path is taken inside the kernel.
         self.write_root_link_velocity_to_sim_index(
             root_velocity=root_velocity, env_ids=env_ids, full_data=True, skip_forward=skip_forward
@@ -1091,8 +1092,8 @@ def write_joint_state_to_sim_mask(
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
         # set into simulation
-        env_ids = self._resolve_env_mask(env_mask)
-        joint_ids = None if joint_mask is None else self._resolve_joint_mask(joint_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        joint_ids = None if joint_mask is None else self._resolve_joint_ids(joint_mask=joint_mask)
         self.write_joint_state_to_sim_index(
             position=position,
             velocity=velocity,
@@ -1196,8 +1197,8 @@ def write_joint_position_to_sim_mask(
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
         # resolve masks
-        env_ids = self._resolve_env_mask(env_mask)
-        joint_ids = None if joint_mask is None else self._resolve_joint_mask(joint_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        joint_ids = None if joint_mask is None else self._resolve_joint_ids(joint_mask=joint_mask)
         # Set full data to True to ensure the the right code path is taken inside the kernel.
         self.write_joint_position_to_sim_index(
             position=position, joint_ids=joint_ids, env_ids=env_ids, full_data=True, skip_forward=skip_forward
@@ -1301,8 +1302,8 @@ def write_joint_velocity_to_sim_mask(
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
         # resolve masks
-        env_ids = self._resolve_env_mask(env_mask)
-        joint_ids = None if joint_mask is None else self._resolve_joint_mask(joint_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        joint_ids = None if joint_mask is None else self._resolve_joint_ids(joint_mask=joint_mask)
         # Set full data to True to ensure the the right code path is taken inside the kernel.
         self.write_joint_velocity_to_sim_index(
             velocity=velocity, joint_ids=joint_ids, env_ids=env_ids, full_data=True, skip_forward=skip_forward
@@ -1412,8 +1413,8 @@ def write_joint_stiffness_to_sim_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        env_ids = self._resolve_env_mask(env_mask)
-        joint_ids = self._resolve_joint_mask(joint_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        joint_ids = self._resolve_joint_ids(joint_mask=joint_mask)
         # Set full data to True to ensure the the right code path is taken inside the kernel.
         self.write_joint_stiffness_to_sim_index(
             stiffness=stiffness, joint_ids=joint_ids, env_ids=env_ids, full_data=True
@@ -1560,8 +1561,8 @@ def write_joint_damping_to_sim_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # resolve masks
-        env_ids = self._resolve_env_mask(env_mask)
-        joint_ids = self._resolve_joint_mask(joint_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        joint_ids = self._resolve_joint_ids(joint_mask=joint_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.write_joint_damping_to_sim_index(damping=damping, joint_ids=joint_ids, env_ids=env_ids, full_data=True)
 
@@ -1672,8 +1673,8 @@ def write_joint_position_limit_to_sim_mask(
                 exceed the new limits. Defaults to True.
         """
         # resolve masks
-        env_ids = self._resolve_env_mask(env_mask)
-        joint_ids = self._resolve_joint_mask(joint_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        joint_ids = self._resolve_joint_ids(joint_mask=joint_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.write_joint_position_limit_to_sim_index(
             limits=limits,
@@ -1787,8 +1788,8 @@ def write_joint_velocity_limit_to_sim_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # resolve masks
-        env_ids = self._resolve_env_mask(env_mask)
-        joint_ids = self._resolve_joint_mask(joint_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        joint_ids = self._resolve_joint_ids(joint_mask=joint_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.write_joint_velocity_limit_to_sim_index(
             limits=limits, joint_ids=joint_ids, env_ids=env_ids, full_data=True
@@ -1897,8 +1898,8 @@ def write_joint_effort_limit_to_sim_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # resolve masks
-        env_ids = self._resolve_env_mask(env_mask)
-        joint_ids = self._resolve_joint_mask(joint_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        joint_ids = self._resolve_joint_ids(joint_mask=joint_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.write_joint_effort_limit_to_sim_index(limits=limits, joint_ids=joint_ids, env_ids=env_ids, full_data=True)
 
@@ -2005,8 +2006,8 @@ def write_joint_armature_to_sim_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # resolve masks
-        env_ids = self._resolve_env_mask(env_mask)
-        joint_ids = self._resolve_joint_mask(joint_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        joint_ids = self._resolve_joint_ids(joint_mask=joint_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.write_joint_armature_to_sim_index(armature=armature, joint_ids=joint_ids, env_ids=env_ids, full_data=True)
 
@@ -2173,8 +2174,8 @@ def write_joint_friction_coefficient_to_sim_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # resolve masks
-        env_ids = self._resolve_env_mask(env_mask)
-        joint_ids = self._resolve_joint_mask(joint_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        joint_ids = self._resolve_joint_ids(joint_mask=joint_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.write_joint_friction_coefficient_to_sim_index(
             joint_friction_coeff=joint_friction_coeff,
@@ -2283,8 +2284,8 @@ def write_joint_dynamic_friction_coefficient_to_sim_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # resolve masks
-        env_ids = self._resolve_env_mask(env_mask)
-        joint_ids = self._resolve_joint_mask(joint_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        joint_ids = self._resolve_joint_ids(joint_mask=joint_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.write_joint_dynamic_friction_coefficient_to_sim_index(
             joint_dynamic_friction_coeff=joint_dynamic_friction_coeff,
@@ -2394,8 +2395,8 @@ def write_joint_viscous_friction_coefficient_to_sim_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # resolve masks
-        env_ids = self._resolve_env_mask(env_mask)
-        joint_ids = self._resolve_joint_mask(joint_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        joint_ids = self._resolve_joint_ids(joint_mask=joint_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.write_joint_viscous_friction_coefficient_to_sim_index(
             joint_viscous_friction_coeff=joint_viscous_friction_coeff,
@@ -2496,8 +2497,8 @@ def set_masses_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        env_ids = self._resolve_env_mask(env_mask)
-        body_ids = self._resolve_body_mask(body_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        body_ids = self._resolve_body_ids(body_mask=body_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.set_masses_index(masses=masses, body_ids=body_ids, env_ids=env_ids, full_data=True)
 
@@ -2603,8 +2604,8 @@ def set_coms_mask(
         all_envs_selected = env_mask is None
         all_bodies_selected = body_mask is None
         # Resolve masks.
-        env_ids = None if all_envs_selected else self._resolve_env_mask(env_mask)
-        body_ids = None if all_bodies_selected else self._resolve_body_mask(body_mask)
+        env_ids = None if all_envs_selected else self._resolve_env_ids(env_mask=env_mask)
+        body_ids = None if all_bodies_selected else self._resolve_body_ids(body_mask=body_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.set_coms_index(coms=coms, body_ids=body_ids, env_ids=env_ids, full_data=True)
 
@@ -2700,8 +2701,8 @@ def set_inertias_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        env_ids = self._resolve_env_mask(env_mask)
-        body_ids = self._resolve_body_mask(body_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        body_ids = self._resolve_body_ids(body_mask=body_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.set_inertias_index(inertias=inertias, body_ids=body_ids, env_ids=env_ids, full_data=True)
 
@@ -2806,8 +2807,8 @@ def set_fixed_tendon_stiffness_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        env_ids = self._resolve_env_mask(env_mask)
-        fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        fixed_tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.set_fixed_tendon_stiffness_index(
             stiffness=stiffness, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True
@@ -2906,8 +2907,8 @@ def set_fixed_tendon_damping_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        env_ids = self._resolve_env_mask(env_mask)
-        fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        fixed_tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.set_fixed_tendon_damping_index(
             damping=damping, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True
@@ -3010,8 +3011,8 @@ def set_fixed_tendon_limit_stiffness_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        env_ids = self._resolve_env_mask(env_mask)
-        fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        fixed_tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.set_fixed_tendon_limit_stiffness_index(
             limit_stiffness=limit_stiffness, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True
@@ -3088,8 +3089,8 @@ def set_fixed_tendon_position_limit_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        env_ids = self._resolve_env_mask(env_mask)
-        fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        fixed_tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.set_fixed_tendon_position_limit_index(
             limit=limit, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True
@@ -3192,8 +3193,8 @@ def set_fixed_tendon_rest_length_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        env_ids = self._resolve_env_mask(env_mask)
-        fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        fixed_tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.set_fixed_tendon_rest_length_index(
             rest_length=rest_length, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True
@@ -3343,8 +3344,8 @@ def set_fixed_tendon_position_target_mask(
         """
         self.set_fixed_tendon_position_target_index(
             target=target,
-            fixed_tendon_ids=self._resolve_fixed_tendon_mask(fixed_tendon_mask),
-            env_ids=self._resolve_env_mask(env_mask),
+            fixed_tendon_ids=self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask),
+            env_ids=self._resolve_env_ids(env_mask=env_mask),
             full_data=True,
         )
 
@@ -3374,8 +3375,8 @@ def set_fixed_tendon_offset_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        env_ids = self._resolve_env_mask(env_mask)
-        fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        fixed_tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.set_fixed_tendon_offset_index(
             offset=offset, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True
@@ -3430,7 +3431,7 @@ def write_fixed_tendon_properties_to_sim_mask(
         only the selected cached values before this method pushes the complete rows.
 
         .. tip::
-            For maximum performance we recommend using the mask method. This is because in PhysX, the tensor API
+            For maximum performance we recommend using the index method. This is because in PhysX, the tensor API
             is only supporting indexing, hence masks need to be converted to indices.
 
         Args:
@@ -3439,7 +3440,7 @@ def write_fixed_tendon_properties_to_sim_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        env_ids = self._resolve_env_mask(env_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
         self.write_fixed_tendon_properties_to_sim_index(env_ids=env_ids)
 
     def set_spatial_tendon_stiffness_index(
@@ -3539,8 +3540,8 @@ def set_spatial_tendon_stiffness_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        env_ids = self._resolve_env_mask(env_mask)
-        spatial_tendon_ids = self._resolve_spatial_tendon_mask(spatial_tendon_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        spatial_tendon_ids = self._resolve_spatial_tendon_ids(spatial_tendon_mask=spatial_tendon_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.set_spatial_tendon_stiffness_index(
             stiffness=stiffness, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids, full_data=True
@@ -3639,8 +3640,8 @@ def set_spatial_tendon_damping_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        env_ids = self._resolve_env_mask(env_mask)
-        spatial_tendon_ids = self._resolve_spatial_tendon_mask(spatial_tendon_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        spatial_tendon_ids = self._resolve_spatial_tendon_ids(spatial_tendon_mask=spatial_tendon_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.set_spatial_tendon_damping_index(
             damping=damping, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids, full_data=True
@@ -3744,8 +3745,8 @@ def set_spatial_tendon_limit_stiffness_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        env_ids = self._resolve_env_mask(env_mask)
-        spatial_tendon_ids = self._resolve_spatial_tendon_mask(spatial_tendon_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        spatial_tendon_ids = self._resolve_spatial_tendon_ids(spatial_tendon_mask=spatial_tendon_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.set_spatial_tendon_limit_stiffness_index(
             limit_stiffness=limit_stiffness, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids, full_data=True
@@ -3844,8 +3845,8 @@ def set_spatial_tendon_offset_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        env_ids = self._resolve_env_mask(env_mask)
-        spatial_tendon_ids = self._resolve_spatial_tendon_mask(spatial_tendon_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        spatial_tendon_ids = self._resolve_spatial_tendon_ids(spatial_tendon_mask=spatial_tendon_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.set_spatial_tendon_offset_index(
             offset=offset, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids, full_data=True
@@ -3898,7 +3899,7 @@ def write_spatial_tendon_properties_to_sim_mask(
         this method pushes the complete rows.
 
         .. tip::
-            For maximum performance we recommend using the mask method. This is because in PhysX, the tensor API
+            For maximum performance we recommend using the index method. This is because in PhysX, the tensor API
             is only supporting indexing, hence masks need to be converted to indices.
 
         Args:
@@ -3907,7 +3908,7 @@ def write_spatial_tendon_properties_to_sim_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        env_ids = self._resolve_env_mask(env_mask)
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
         # Write spatial tendon properties to the simulation.
         self.write_spatial_tendon_properties_to_sim_index(env_ids=env_ids)
 
@@ -4463,62 +4464,23 @@ def _cpu_env_ids_view(self, count: int) -> wp.array:
             )
         return self._cpu_env_ids_views[count]
 
-    def _resolve_env_mask(self, env_mask: wp.array | None) -> torch.Tensor | wp.array:
-        """
-        Resolve environment mask to a torch tensor.
-
-        Args:
-            env_mask: Environment mask. If None, then all indices are used.
-
-        Returns:
-            A torch tensor of environment indices.
-        """
-        # resolve masks
-        if env_mask is not None:
-            if isinstance(env_mask, wp.array):
-                env_mask = wp.to_torch(env_mask)
-            env_ids = torch.nonzero(env_mask)[:, 0]
-        else:
-            env_ids = self._ALL_INDICES
-        return env_ids
-
-    def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor:
+    def _resolve_env_ids(
+        self, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_mask: wp.array | None = None
+    ) -> wp.array | torch.Tensor:
         """Resolve environment indices.
 
         Args:
             env_ids: Environment indices. If None, then all indices are used.
+            env_mask: Boolean mask; takes precedence over ``env_ids``.
 
         Returns:
             Environment indices.
         """
-        if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)):
-            return self._ALL_INDICES
-        if isinstance(env_ids, slice):
-            return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids])
-        if isinstance(env_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if isinstance(env_ids, list):
-            return wp.array(env_ids, dtype=wp.int32, device=self.device)
-        return env_ids
-
-    def _resolve_joint_mask(self, joint_mask: wp.array | None) -> torch.Tensor | wp.array:
-        """Resolve joint mask to a torch tensor.
-
-        Args:
-            joint_mask: Joint mask. If None, then all indices are used.
+        return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_INDICES, device=self.device)
 
-        Returns:
-            A torch tensor of joint indices.
-        """
-        if joint_mask is not None:
-            if isinstance(joint_mask, wp.array):
-                joint_mask = wp.to_torch(joint_mask)
-            joint_ids = torch.nonzero(joint_mask)[:, 0]
-        else:
-            joint_ids = self._ALL_JOINT_INDICES
-        return joint_ids
-
-    def _resolve_joint_ids(self, joint_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor:
+    def _resolve_joint_ids(
+        self, joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, joint_mask: wp.array | None = None
+    ) -> wp.array | torch.Tensor:
         """Resolve joint indices to a warp array or tensor.
 
         .. note::
@@ -4526,123 +4488,65 @@ def _resolve_joint_ids(self, joint_ids: Sequence[int] | torch.Tensor | wp.array
 
         Args:
             joint_ids: Joint indices. If None, then all indices are used.
+            joint_mask: Boolean mask; takes precedence over ``joint_ids``.
 
         Returns:
             A warp array of joint indices or a tensor of joint indices.
         """
-        if isinstance(joint_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if isinstance(joint_ids, list):
-            return wp.array(joint_ids, dtype=wp.int32, device=self.device)
-        if (joint_ids is None) or (joint_ids == slice(None)):
-            return self._ALL_JOINT_INDICES
-        return joint_ids
-
-    def _resolve_body_mask(self, body_mask: wp.array | None) -> torch.Tensor | wp.array:
-        """Resolve body mask to a torch tensor.
-
-        Args:
-            body_mask: Body mask. If None, then all indices are used.
-
-        Returns:
-            A torch tensor of body indices.
-        """
-        if body_mask is not None:
-            if isinstance(body_mask, wp.array):
-                body_mask = wp.to_torch(body_mask)
-            body_ids = torch.nonzero(body_mask)[:, 0]
-        else:
-            body_ids = self._ALL_BODY_INDICES
-        return body_ids
+        return resolve_ids(ids=joint_ids, mask=joint_mask, all_ids=self._ALL_JOINT_INDICES, device=self.device)
 
-    def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor:
+    def _resolve_body_ids(
+        self, body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, body_mask: wp.array | None = None
+    ) -> wp.array | torch.Tensor:
         """Resolve body indices to a warp array or tensor.
 
         Args:
             body_ids: Body indices. If None, then all indices are used.
+            body_mask: Boolean mask; takes precedence over ``body_ids``.
 
         Returns:
             A warp array of body indices or a tensor of body indices.
         """
-        if isinstance(body_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if isinstance(body_ids, list):
-            return wp.array(body_ids, dtype=wp.int32, device=self.device)
-        if (body_ids is None) or (body_ids == slice(None)):
-            return self._ALL_BODY_INDICES
-        return body_ids
-
-    def _resolve_fixed_tendon_mask(self, fixed_tendon_mask: wp.array | None) -> torch.Tensor | wp.array:
-        """Resolve fixed tendon mask to a torch tensor.
-
-        Args:
-            fixed_tendon_mask: Fixed tendon mask. If None, then all indices are used.
-
-        Returns:
-            A torch tensor of fixed tendon indices.
-        """
-        if fixed_tendon_mask is not None:
-            if isinstance(fixed_tendon_mask, wp.array):
-                fixed_tendon_mask = wp.to_torch(fixed_tendon_mask)
-            fixed_tendon_ids = torch.nonzero(fixed_tendon_mask)[:, 0]
-        else:
-            fixed_tendon_ids = self._ALL_FIXED_TENDON_INDICES
-        return fixed_tendon_ids
+        return resolve_ids(ids=body_ids, mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self.device)
 
     def _resolve_fixed_tendon_ids(
-        self, tendon_ids: Sequence[int] | torch.Tensor | wp.array | None
+        self,
+        tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None,
+        fixed_tendon_mask: wp.array | None = None,
     ) -> wp.array | torch.Tensor:
         """Resolve tendon indices to a warp array or tensor.
 
         Args:
             tendon_ids: Tendon indices. If None, then all indices are used.
+            fixed_tendon_mask: Boolean mask; takes precedence over ``tendon_ids``.
 
         Returns:
             A warp array of tendon indices or a tensor of tendon indices.
         """
-        if isinstance(tendon_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if isinstance(tendon_ids, list):
-            return wp.array(tendon_ids, dtype=wp.int32, device=self.device)
-        if (tendon_ids is None) or (tendon_ids == slice(None)):
-            return self._ALL_FIXED_TENDON_INDICES
-        return tendon_ids
-
-    def _resolve_spatial_tendon_mask(self, spatial_tendon_mask: wp.array | None) -> torch.Tensor | wp.array:
-        """Resolve spatial tendon mask to a torch tensor.
-
-        Args:
-            spatial_tendon_mask: Spatial tendon mask. If None, then all indices are used.
-
-        Returns:
-            A torch tensor of spatial tendon indices.
-        """
-        if spatial_tendon_mask is not None:
-            if isinstance(spatial_tendon_mask, wp.array):
-                spatial_tendon_mask = wp.to_torch(spatial_tendon_mask)
-            spatial_tendon_ids = torch.nonzero(spatial_tendon_mask)[:, 0]
-        else:
-            spatial_tendon_ids = self._ALL_SPATIAL_TENDON_INDICES
-        return spatial_tendon_ids
+        return resolve_ids(
+            ids=tendon_ids, mask=fixed_tendon_mask, all_ids=self._ALL_FIXED_TENDON_INDICES, device=self.device
+        )
 
     def _resolve_spatial_tendon_ids(
-        self, spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None
+        self,
+        spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None,
+        spatial_tendon_mask: wp.array | None = None,
     ) -> wp.array | torch.Tensor:
         """Resolve spatial tendon indices to a warp array or tensor.
 
         Args:
             spatial_tendon_ids: Spatial tendon indices. If None, then all indices are used.
+            spatial_tendon_mask: Boolean mask; takes precedence over ``spatial_tendon_ids``.
 
         Returns:
             A warp array of spatial tendon indices or a tensor of spatial tendon indices.
         """
-        if isinstance(spatial_tendon_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if isinstance(spatial_tendon_ids, list):
-            return wp.array(spatial_tendon_ids, dtype=wp.int32, device=self.device)
-        if (spatial_tendon_ids is None) or (spatial_tendon_ids == slice(None)):
-            return self._ALL_SPATIAL_TENDON_INDICES
-        return spatial_tendon_ids
+        return resolve_ids(
+            ids=spatial_tendon_ids,
+            mask=spatial_tendon_mask,
+            all_ids=self._ALL_SPATIAL_TENDON_INDICES,
+            device=self.device,
+        )
 
     """
     Deprecated methods.
diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py
index 47c7240c510e..48d1345349ea 100644
--- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py
+++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py
@@ -19,6 +19,7 @@
 from isaaclab.assets.rigid_object.base_rigid_object import BaseRigidObject
 from isaaclab.sim.utils.queries import path_expr_to_glob, resolve_matching_prims_from_source
 from isaaclab.utils.warp import ProxyArray
+from isaaclab.utils.warp.utils import resolve_ids
 from isaaclab.utils.wrench_composer import WrenchComposer
 
 from isaaclab_physx.assets import kernels as shared_kernels
@@ -397,10 +398,7 @@ def write_root_link_pose_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_ids = self._ALL_INDICES
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
         self.write_root_link_pose_to_sim_index(
             root_pose=root_pose, env_ids=env_ids, full_data=True, skip_forward=skip_forward
         )
@@ -491,10 +489,7 @@ def write_root_com_pose_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_ids = self._ALL_INDICES
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
         self.write_root_com_pose_to_sim_index(
             root_pose=root_pose, env_ids=env_ids, full_data=True, skip_forward=skip_forward
         )
@@ -593,10 +588,7 @@ def write_root_com_velocity_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_ids = self._ALL_INDICES
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
         self.write_root_com_velocity_to_sim_index(
             root_velocity=root_velocity, env_ids=env_ids, full_data=True, skip_forward=skip_forward
         )
@@ -699,10 +691,7 @@ def write_root_link_velocity_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_ids = self._ALL_INDICES
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
         self.write_root_link_velocity_to_sim_index(
             root_velocity=root_velocity, env_ids=env_ids, full_data=True, skip_forward=skip_forward
         )
@@ -793,14 +782,8 @@ def set_masses_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_ids = self._ALL_INDICES
-        if body_mask is not None:
-            body_ids = self._resolve_body_mask(body_mask)
-        else:
-            body_ids = self._ALL_BODY_INDICES
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        body_ids = self._resolve_body_ids(body_mask=body_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.set_masses_index(masses=masses, body_ids=body_ids, env_ids=env_ids, full_data=True)
 
@@ -885,14 +868,8 @@ def set_coms_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_ids = self._ALL_INDICES
-        if body_mask is not None:
-            body_ids = self._resolve_body_mask(body_mask)
-        else:
-            body_ids = self._ALL_BODY_INDICES
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        body_ids = self._resolve_body_ids(body_mask=body_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.set_coms_index(coms=coms, body_ids=body_ids, env_ids=env_ids, full_data=True)
 
@@ -980,14 +957,8 @@ def set_inertias_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_ids = self._ALL_INDICES
-        if body_mask is not None:
-            body_ids = self._resolve_body_mask(body_mask)
-        else:
-            body_ids = self._ALL_BODY_INDICES
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        body_ids = self._resolve_body_ids(body_mask=body_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.set_inertias_index(inertias=inertias, body_ids=body_ids, env_ids=env_ids, full_data=True)
 
@@ -1068,40 +1039,6 @@ def _process_cfg(self) -> None:
         self._data.default_root_pose = wp.array(default_root_pose, dtype=wp.transformf, device=self.device)
         self._data.default_root_vel = wp.array(default_root_vel, dtype=wp.spatial_vectorf, device=self.device)
 
-    def _resolve_env_mask(self, env_mask: wp.array | None) -> torch.Tensor | wp.array:
-        """Resolve environment mask to a torch tensor.
-
-        Args:
-            env_mask: Environment mask. If None, then all indices are used.
-
-        Returns:
-            A torch tensor of environment indices.
-        """
-        if env_mask is not None:
-            if isinstance(env_mask, wp.array):
-                env_mask = wp.to_torch(env_mask)
-            env_ids = torch.nonzero(env_mask)[:, 0]
-        else:
-            env_ids = self._ALL_INDICES
-        return env_ids
-
-    def _resolve_body_mask(self, body_mask: wp.array | None) -> torch.Tensor | wp.array:
-        """Resolve body mask to a torch tensor.
-
-        Args:
-            body_mask: Body mask. If None, then all indices are used.
-
-        Returns:
-            A torch tensor of body indices.
-        """
-        if body_mask is not None:
-            if isinstance(body_mask, wp.array):
-                body_mask = wp.to_torch(body_mask)
-            body_ids = torch.nonzero(body_mask)[:, 0]
-        else:
-            body_ids = self._ALL_BODY_INDICES
-        return body_ids
-
     def _get_cpu_env_ids(self, env_ids: wp.array | torch.Tensor, sim_env_ids: wp.array | None = None) -> wp.array:
         """Get CPU environment indices. Uses pre-allocated pinned buffer for full-index case.
 
@@ -1164,39 +1101,33 @@ def _sim_env_ids_view(self, count: int) -> wp.array:
             )
         return self._sim_env_ids_views[count]
 
-    def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor:
+    def _resolve_env_ids(
+        self, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_mask: wp.array | None = None
+    ) -> wp.array | torch.Tensor:
         """Resolve environment indices to a warp array or tensor.
 
         Args:
             env_ids: Environment indices. If None, then all indices are used.
+            env_mask: Boolean mask; takes precedence over ``env_ids``.
 
         Returns:
             A warp array of environment indices or a tensor of environment indices.
         """
-        if (env_ids is None) or (env_ids == slice(None)):
-            return self._ALL_INDICES
-        if isinstance(env_ids, slice):
-            return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids])
-        if isinstance(env_ids, list):
-            return wp.array(env_ids, dtype=wp.int32, device=self.device)
-        return env_ids
-
-    def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor:
+        return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_INDICES, device=self.device)
+
+    def _resolve_body_ids(
+        self, body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, body_mask: wp.array | None = None
+    ) -> wp.array | torch.Tensor:
         """Resolve body indices to a warp array or tensor.
 
         Args:
             body_ids: Body indices. If None, then all indices are used.
+            body_mask: Boolean mask; takes precedence over ``body_ids``.
 
         Returns:
             A warp array of body indices or a tensor of body indices.
         """
-        if isinstance(body_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if isinstance(body_ids, list):
-            return wp.array(body_ids, dtype=wp.int32, device=self.device)
-        if (body_ids is None) or (body_ids == slice(None)):
-            return self._ALL_BODY_INDICES
-        return body_ids
+        return resolve_ids(ids=body_ids, mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self.device)
 
     """
     Internal simulation callbacks.
diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py
index 47d5e1cb460b..e79295d59f04 100644
--- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py
+++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py
@@ -22,6 +22,7 @@
 from isaaclab.assets.rigid_object_collection.base_rigid_object_collection import BaseRigidObjectCollection
 from isaaclab.utils import clone, validate
 from isaaclab.utils.warp import ProxyArray
+from isaaclab.utils.warp.utils import resolve_ids
 from isaaclab.utils.wrench_composer import WrenchComposer
 
 from isaaclab_physx.assets import kernels as shared_kernels
@@ -307,14 +308,8 @@ def write_body_pose_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_ids = self._resolve_body_mask(body_mask)
-        else:
-            body_ids = self._ALL_BODY_INDICES
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        body_ids = self._resolve_body_ids(body_mask=body_mask)
         self.write_body_link_pose_to_sim_index(
             body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward
         )
@@ -391,14 +386,8 @@ def write_body_velocity_to_sim_mask(
             skip_forward: Whether to skip invalidating cached data after the write. When True, the caller
                 must invalidate stale cached data before reading it back. Defaults to False.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_ids = self._resolve_body_mask(body_mask)
-        else:
-            body_ids = self._ALL_BODY_INDICES
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        body_ids = self._resolve_body_ids(body_mask=body_mask)
         self.write_body_com_velocity_to_sim_index(
             body_velocities=body_velocities,
             env_ids=env_ids,
@@ -509,10 +498,7 @@ def write_body_link_pose_to_sim_mask(
                 must invalidate stale cached data before reading it back. Defaults to False.
             body_ids: Deprecated, use :attr:`body_mask` instead. Body indices. Defaults to None.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_ids = self._ALL_ENV_INDICES
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
         body_ids = self._resolve_mask_writer_body_ids(body_mask, body_ids)
         self.write_body_link_pose_to_sim_index(
             body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward
@@ -624,10 +610,7 @@ def write_body_com_pose_to_sim_mask(
                 must invalidate stale cached data before reading it back. Defaults to False.
             body_ids: Deprecated, use :attr:`body_mask` instead. Body indices. Defaults to None.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_ids = self._ALL_ENV_INDICES
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
         body_ids = self._resolve_mask_writer_body_ids(body_mask, body_ids)
         self.write_body_com_pose_to_sim_index(
             body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward
@@ -744,10 +727,7 @@ def write_body_com_velocity_to_sim_mask(
                 must invalidate stale cached data before reading it back. Defaults to False.
             body_ids: Deprecated, use :attr:`body_mask` instead. Body indices. Defaults to None.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_ids = self._ALL_ENV_INDICES
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
         body_ids = self._resolve_mask_writer_body_ids(body_mask, body_ids)
         self.write_body_com_velocity_to_sim_index(
             body_velocities=body_velocities,
@@ -871,10 +851,7 @@ def write_body_link_velocity_to_sim_mask(
                 must invalidate stale cached data before reading it back. Defaults to False.
             body_ids: Deprecated, use :attr:`body_mask` instead. Body indices. Defaults to None.
         """
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_ids = self._ALL_ENV_INDICES
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
         body_ids = self._resolve_mask_writer_body_ids(body_mask, body_ids)
         self.write_body_link_velocity_to_sim_index(
             body_velocities=body_velocities,
@@ -968,14 +945,8 @@ def set_masses_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_ids = self._resolve_body_mask(body_mask)
-        else:
-            body_ids = self._ALL_BODY_INDICES
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        body_ids = self._resolve_body_ids(body_mask=body_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.set_masses_index(masses=masses, body_ids=body_ids, env_ids=env_ids, full_data=True)
 
@@ -1062,14 +1033,8 @@ def set_coms_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_ids = self._resolve_body_mask(body_mask)
-        else:
-            body_ids = self._ALL_BODY_INDICES
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        body_ids = self._resolve_body_ids(body_mask=body_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.set_coms_index(coms=coms, body_ids=body_ids, env_ids=env_ids, full_data=True)
 
@@ -1157,14 +1122,8 @@ def set_inertias_mask(
             env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,).
         """
         # Resolve masks.
-        if env_mask is not None:
-            env_ids = self._resolve_env_mask(env_mask)
-        else:
-            env_ids = self._ALL_ENV_INDICES
-        if body_mask is not None:
-            body_ids = self._resolve_body_mask(body_mask)
-        else:
-            body_ids = self._ALL_BODY_INDICES
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
+        body_ids = self._resolve_body_ids(body_mask=body_mask)
         # Set full data to True to ensure the right code path is taken inside the kernel.
         self.set_inertias_index(inertias=inertias, body_ids=body_ids, env_ids=env_ids, full_data=True)
 
@@ -1311,46 +1270,20 @@ def reshape_data_to_view_3d(
     Internal helper.
     """
 
-    def _resolve_env_ids(self, env_ids) -> wp.array | torch.Tensor:
+    def _resolve_env_ids(self, env_ids=None, env_mask: wp.array | None = None) -> wp.array | torch.Tensor:
         """Resolve environment indices."""
-        if (env_ids is None) or (env_ids == slice(None)):
-            return self._ALL_ENV_INDICES
-        if isinstance(env_ids, slice):
-            return wp.from_torch(wp.to_torch(self._ALL_ENV_INDICES)[env_ids])
-        if isinstance(env_ids, list):
-            return wp.array(env_ids, dtype=wp.int32, device=self.device)
-        return env_ids
-
-    def _resolve_body_ids(self, body_ids) -> wp.array | torch.Tensor:
-        """Resolve body indices."""
-        if isinstance(body_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if isinstance(body_ids, list):
-            return wp.array(body_ids, dtype=wp.int32, device=self.device)
-        if body_ids is None or (body_ids == slice(None)):
-            return self._ALL_BODY_INDICES
-        if isinstance(body_ids, slice):
-            return wp.from_torch(
-                torch.arange(self.num_bodies, dtype=torch.int32, device=self.device)[body_ids], dtype=wp.int32
-            )
-        return body_ids
+        return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self.device)
 
-    def _resolve_env_mask(self, env_mask: wp.array | None) -> torch.Tensor | wp.array:
-        """Resolve environment mask to indices via torch.nonzero."""
-        if env_mask is not None:
-            if isinstance(env_mask, wp.array):
-                env_mask = wp.to_torch(env_mask)
-            env_ids = torch.nonzero(env_mask)[:, 0]
-        else:
-            env_ids = self._ALL_ENV_INDICES
-        return env_ids
+    def _resolve_body_ids(self, body_ids=None, body_mask: wp.array | None = None) -> wp.array | torch.Tensor:
+        """Resolve body indices."""
+        return resolve_ids(ids=body_ids, mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self.device)
 
     def _resolve_mask_writer_body_ids(
         self, body_mask: wp.array | None, body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None
     ) -> Sequence[int] | torch.Tensor | wp.array | slice:
         """Resolve the bodies of a mask writer, accepting the deprecated ``body_ids`` in place of ``body_mask``."""
         if body_ids is None:
-            return self._resolve_body_mask(body_mask)
+            return self._resolve_body_ids(body_mask=body_mask)
         if body_mask is not None:
             raise ValueError("Pass either 'body_mask' or the deprecated 'body_ids', not both.")
         warnings.warn(
@@ -1361,16 +1294,6 @@ def _resolve_mask_writer_body_ids(
         )
         return body_ids
 
-    def _resolve_body_mask(self, body_mask: wp.array | None) -> torch.Tensor | wp.array:
-        """Resolve body mask to indices via torch.nonzero."""
-        if body_mask is not None:
-            if isinstance(body_mask, wp.array):
-                body_mask = wp.to_torch(body_mask)
-            body_ids = torch.nonzero(body_mask)[:, 0]
-        else:
-            body_ids = self._ALL_BODY_INDICES
-        return body_ids
-
     def _initialize_impl(self):
         # clear body names list to prevent double counting on re-initialization
         self._body_names_list.clear()
diff --git a/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper.py b/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper.py
index e37a30f7a505..a2f5cb930c8f 100644
--- a/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper.py
+++ b/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper.py
@@ -16,6 +16,7 @@
 from isaaclab.assets import AssetBase
 from isaaclab.utils import clone
 from isaaclab.utils.version import get_isaac_sim_version, has_kit
+from isaaclab.utils.warp.utils import resolve_ids
 
 if TYPE_CHECKING:
     from isaacsim.robot.surface_gripper import GripperView
@@ -187,10 +188,7 @@ def set_grippers_command_mask(self, states: torch.Tensor | wp.array, env_mask: t
             states: A tensor/array of floats representing the gripper command (full size).
             env_mask: Boolean mask of shape (num_envs,). Defaults to None (all environments).
         """
-        if env_mask is not None:
-            env_ids = wp.from_torch(torch.argwhere(env_mask).view(-1).to(torch.int32), dtype=wp.int32)
-        else:
-            env_ids = self._ALL_INDICES
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
         self.set_grippers_command_index(states, env_ids, full_data=True)
 
     def set_grippers_command(self, states: torch.Tensor, indices: torch.Tensor | None = None) -> None:
@@ -274,10 +272,7 @@ def update_gripper_properties_mask(
             retry_interval: The retry interval. Shape (num_envs,).
             env_mask: Boolean mask of shape (num_envs,). Defaults to None (all environments).
         """
-        if env_mask is not None:
-            env_ids = wp.from_torch(torch.argwhere(env_mask).view(-1).to(torch.int32), dtype=wp.int32)
-        else:
-            env_ids = self._ALL_INDICES
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
         self.update_gripper_properties_index(
             max_grip_distance=max_grip_distance,
             coaxial_force_limit=coaxial_force_limit,
@@ -407,10 +402,7 @@ def reset_mask(self, env_mask: torch.Tensor | None = None) -> None:
         Args:
             env_mask: Boolean mask of shape (num_envs,). Defaults to None (all environments).
         """
-        if env_mask is not None:
-            env_ids = wp.from_torch(torch.argwhere(env_mask).view(-1).to(torch.int32), dtype=wp.int32)
-        else:
-            env_ids = self._ALL_INDICES
+        env_ids = self._resolve_env_ids(env_mask=env_mask)
         self.reset_index(env_ids)
 
     def reset(self, indices: torch.Tensor | None = None) -> None:
@@ -552,24 +544,17 @@ def _process_cfg(self) -> None:
     Helper functions.
     """
 
-    def _resolve_env_ids(self, env_ids) -> wp.array:
+    def _resolve_env_ids(self, env_ids=None, env_mask: wp.array | torch.Tensor | None = None) -> wp.array:
         """Resolve environment indices to a warp array.
 
         Args:
             env_ids: Environment indices. Can be None, a slice, a list, or a torch.Tensor.
+            env_mask: Boolean mask; takes precedence over ``env_ids``.
 
         Returns:
             A warp array of int32 indices.
         """
-        if env_ids is None or env_ids == slice(None):
-            return self._ALL_INDICES
-        if isinstance(env_ids, slice):
-            return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids])
-        elif isinstance(env_ids, list):
-            return wp.array(env_ids, dtype=wp.int32, device=self._device)
-        elif isinstance(env_ids, torch.Tensor):
-            return wp.from_torch(env_ids.to(torch.int32).contiguous(), dtype=wp.int32)
-        return env_ids
+        return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_INDICES, device=self._device, dtype=wp.int32)
 
     def parse_gripper_parameter(
         self, cfg_value: float | int | tuple | None, default_value: float | int | tuple | None, ndim: int = 0
diff --git a/source/isaaclab_physx/test/sensors/test_frame_transformer.py b/source/isaaclab_physx/test/sensors/test_frame_transformer.py
index 7ba9a23d3de4..50e52bfbc646 100644
--- a/source/isaaclab_physx/test/sensors/test_frame_transformer.py
+++ b/source/isaaclab_physx/test/sensors/test_frame_transformer.py
@@ -655,6 +655,8 @@ def _make_frame_transformer(use_recorded_launch: bool = True):
     sensor.cfg = SimpleNamespace(prim_path="/World/Source")
     sensor._device = device
     sensor._num_envs = 1
+    sensor._ALL_ENV_MASK = wp.ones(1, dtype=wp.bool, device=device)
+    sensor._reset_mask = wp.zeros(1, dtype=wp.bool, device=device)
     sensor._num_target_frames = 1
     sensor._frame_physx_view = transform_view
     sensor._backend_manager = SimpleNamespace(update_kinematics=lambda: None)
diff --git a/source/isaaclab_physx/test/sensors/test_imu_pva_recorded_launch.py b/source/isaaclab_physx/test/sensors/test_imu_pva_recorded_launch.py
index ea2a1b6df989..80c2956113de 100644
--- a/source/isaaclab_physx/test/sensors/test_imu_pva_recorded_launch.py
+++ b/source/isaaclab_physx/test/sensors/test_imu_pva_recorded_launch.py
@@ -122,6 +122,8 @@ def _make_sensor(sensor_type: str, use_recorded_launch: bool = True):
     sensor._device = device
     sensor._backend_manager = SimpleNamespace(update_kinematics=lambda: None)
     sensor._num_envs = 2
+    sensor._ALL_ENV_MASK = wp.ones(2, dtype=wp.bool, device=device)
+    sensor._reset_mask = wp.zeros(2, dtype=wp.bool, device=device)
     sensor._view = rigid_view
     sensor._offset_pos_b = wp.zeros(2, dtype=wp.vec3f, device=device)
     sensor._offset_quat_b = wp.array(

From 6c31123eb0a14ec8b952568689858730a7b23077 Mon Sep 17 00:00:00 2001
From: Jichuan Hu <jichuanh@nvidia.com>
Date: Mon, 26 Oct 2026 09:53:29 -0700
Subject: [PATCH 48/48] Write Newton tendon targets behind the host flag again

Writing the buffered targets every step cost 2-4% physics step time for every
articulation with tendon actuators. The only Warp caller, the Shadow Hand twin,
sets its targets outside the captured stage, so the host flag still sees them.
---
 .../jichuanh-warp-scene-mask-native.fixed.rst |  3 ---
 .../assets/articulation/articulation.py       |  7 ++---
 .../test/assets/test_articulation.py          | 26 +++----------------
 3 files changed, 8 insertions(+), 28 deletions(-)

diff --git a/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst b/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
index 8400b505e5f3..4f283403787d 100644
--- a/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
+++ b/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst
@@ -3,6 +3,3 @@
 * Fixed :meth:`~isaaclab_newton.assets.RigidObject.reset` and
   :meth:`~isaaclab_newton.assets.RigidObjectCollection.reset` ignoring ``env_mask`` and clearing the external
   wrenches of every environment.
-* Fixed fixed-tendon position targets commanded inside a captured CUDA graph reaching MuJoCo's controls only at
-  capture time. :meth:`~isaaclab_newton.assets.Articulation.write_data_to_sim` now writes the buffered targets on
-  every step when the articulation has tendon actuators, instead of behind a host-side flag.
diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
index f39e71c5360b..d1bacfb77805 100644
--- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
+++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py
@@ -450,10 +450,9 @@ def write_data_to_sim(self):
 
         # Tendon submission is solver-specific: MuJoCo drives tendons through actuator controls
         # outside the articulation view, so the manager owns how a buffered target reaches the solver.
-        # The buffer is written every step rather than behind a host-side flag, so a target
-        # commanded inside a captured CUDA graph reaches the solver on every replay.
-        if self._fixed_tendon_control is not None:
+        if self._fixed_tendon_target_dirty:
             self._fixed_tendon_control.write_data_to_sim(self._backend_manager.get_control())
+            self._fixed_tendon_target_dirty = False
 
     def update(self, dt: float):
         """Updates the simulation data.
@@ -3026,6 +3025,7 @@ def set_fixed_tendon_position_target_index(
         self._fixed_tendon_control.set_position_target_index(
             target=target, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids
         )
+        self._fixed_tendon_target_dirty = True
 
     def set_fixed_tendon_offset_index(
         self,
@@ -3077,6 +3077,7 @@ def set_fixed_tendon_position_target_mask(
         self._fixed_tendon_control.set_position_target_mask(
             target=target, fixed_tendon_mask=fixed_tendon_mask, env_mask=env_mask
         )
+        self._fixed_tendon_target_dirty = True
 
     def set_fixed_tendon_offset_mask(
         self,
diff --git a/source/isaaclab_newton/test/assets/test_articulation.py b/source/isaaclab_newton/test/assets/test_articulation.py
index b621cf310ffe..dbf4f1c23c51 100644
--- a/source/isaaclab_newton/test/assets/test_articulation.py
+++ b/source/isaaclab_newton/test/assets/test_articulation.py
@@ -20,7 +20,7 @@
 these tests come first and fail if selected after a composite-scene test.
 
 The composite scene runs on the CPU without CUDA-graph capture. On CUDA, own-scene tests cover the ordered-state
-republish and the fixed-tendon mask writes and targets recorded into a captured graph.
+republish and the fixed-tendon mask writes recorded into a captured graph.
 """
 
 from isaaclab_newton.physics import NewtonCfg, NewtonMJWarpManager
@@ -49,7 +49,6 @@
 from isaaclab_newton.controllers.ik import NewtonIKPoseObjectiveCfg
 from isaaclab_newton.envs.mdp.actions.newton_ik_actions import NewtonInverseKinematicsAction
 from isaaclab_newton.envs.mdp.actions.newton_ik_actions_cfg import NewtonInverseKinematicsActionCfg
-from isaaclab_newton.physics.mjwarp_tendon_control import resolve_fixed_tendon_actuator_columns
 from isaaclab_physx.sim.schemas import PhysxJointCfg
 from newton import JointType, ModelBuilder, ModelFlags, ShapeFlags
 from newton.selection import ArticulationView
@@ -709,33 +708,16 @@ def test_newton_ordered_state_publishes_inside_captured_cuda_graph(device: str)
 
 
 @pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA))
-def test_fixed_tendon_mask_writes_and_targets_replay_from_cuda_graph(device: str) -> None:
-    """Fixed-tendon mask setters, the mask writer, and tendon targets recorded into a CUDA graph act on every replay.
+def test_fixed_tendon_mask_writes_replay_from_cuda_graph(device: str) -> None:
+    """Fixed-tendon mask setters and the mask writer recorded into a CUDA graph act on every replay.
 
     The mask variants are the documented graph-capturable path. The inputs are filled only after capture, so the
-    replay must read them at launch time: property writes reach only the selected cells, and a tendon target
-    reaches MuJoCo's controls on every replay, not only at capture.
+    replay must read them at launch time and write only the selected cells.
     """
     with build_simulation_context(sim_cfg=newton_sim_cfg(device, use_newton_actuators=False)) as sim:
         articulation = spawn_assets({"tendon": _island_cfgs()["tendon"]})["tendon"]
         sim.reset()
         num_envs, num_tendons = articulation.num_instances, articulation.num_fixed_tendons
-        actuator_columns, _ = resolve_fixed_tendon_actuator_columns(
-            articulation.root_view, sim.backend_manager.get_model()
-        )
-
-        target = torch.zeros(num_envs, num_tendons, device=device)
-        with wp.ScopedCapture(device) as capture:
-            articulation.set_fixed_tendon_position_target_mask(target=target)
-        articulation.write_data_to_sim()
-        for value in (0.3, 0.6):
-            target.fill_(value)
-            wp.capture_launch(capture.graph)
-            articulation.write_data_to_sim()
-            ctrl = wp.to_torch(articulation.root_view.get_attribute("mujoco.ctrl", sim.backend_manager.get_control()))
-            tendon_ctrl = ctrl[:, 0, actuator_columns]
-            torch.testing.assert_close(tendon_ctrl, torch.full_like(tendon_ctrl, value))
-
         env_mask = wp.array([False, True], dtype=wp.bool, device=device)
         tendon_mask = wp.array([False, True], dtype=wp.bool, device=device)
         selected = torch.zeros(num_envs, num_tendons, 1, dtype=torch.bool, device=device)