[Warp] Part 5/8: Capture Warp stages per term without warm-up and step like the stable environments

docs/source/api/lab_experimental/isaaclab_experimental.utils.rst

@@ -13,15 +13,13 @@ The following classes are part of the public :mod:`isaaclab_experimental.utils`
 .. autosummary::
    :nosignatures:
 
-   ManagerCallMode
-   ManagerCallSwitch
-   WarpGraphCache
+   CapturedStage
+   captured
+   eager
 
-.. autoclass:: ManagerCallMode
+.. autoclass:: CapturedStage
    :show-inheritance:
 
-.. autoclass:: ManagerCallSwitch
-   :show-inheritance:
+.. autofunction:: captured
 
-.. autoclass:: WarpGraphCache
-   :show-inheritance:
+.. autofunction:: eager

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)``
 

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.

source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst

@@ -0,0 +1,11 @@
+* 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
+  to accept a boolean torch mask on ``device`` inside a CUDA graph capture.

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.

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.

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.

source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst

@@ -0,0 +1,7 @@
+* 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.

source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst

@@ -0,0 +1,17 @@
+* 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.
+* **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.
+* 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.

source/isaaclab/isaaclab/actuators/actuator_base.py

@@ -279,11 +279,17 @@ def joint_indices(self) -> slice | torch.Tensor:
     """
 
     @abstractmethod
-    def reset(self, env_ids: Sequence[int]):
+    def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: torch.Tensor | None = None):
         """Reset the internals within the group.
 
+        .. caution::
+            If both ``env_ids`` and ``env_mask`` are provided, ``env_mask`` takes precedence.
+
         Args:
-            env_ids: List of environment IDs to reset.
+            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
 

source/isaaclab/isaaclab/actuators/actuator_collection.py

@@ -180,18 +180,27 @@ 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 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.
+
         Args:
             env_ids: Environment indices to reset. Defaults to all environments.
+            env_mask: Environment mask. Shape is (num_instances,). Defaults to None.
         """
         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)
-        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.

source/isaaclab/isaaclab/actuators/actuator_control.py

@@ -307,11 +307,15 @@ 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.
 
+        .. caution::
+            If both ``env_ids`` and ``env_mask`` are provided, ``env_mask`` takes precedence.
+
         Args:
             env_ids: Environments to reset.
+            env_mask: Environment mask. Shape is (num_instances,). Defaults to None.
         """
 
 

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

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(
@@ -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]):
-        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)
-        num_envs = len(range(self._num_envs)[env_ids]) if isinstance(env_ids, slice) else len(env_ids)
+        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:
+            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,
@@ -473,8 +480,8 @@ def reset(self, env_ids: Sequence[int]):
         )
         # 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

source/isaaclab/isaaclab/actuators/newton/adapter.py

@@ -32,12 +32,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
@@ -106,6 +102,12 @@ def __init__(
 
         self._states_a = [act.state() for act in actuators]
         self._states_b = [act.state() for act in actuators]
+        # 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)
 
         # Pre-clamp computed effort buffer. Each Newton actuator scatter-adds
         # its raw drive output to ``sim_control.joint_computed_f`` when
@@ -172,14 +174,20 @@ 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:
             env_ids: Environment indices to reset. ``None`` (or
                 ``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``),
@@ -188,31 +196,26 @@ 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)
                 if sb is not 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)
+        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 = resolve_1d_mask(
+                ids=env_ids,
+                all_mask=self._all_env_mask,
+                scratch_mask=self._scratch_env_mask,
+                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)
+        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],

source/isaaclab/isaaclab/actuators/newton/kernels.py

@@ -6,15 +6,14 @@
 """Shared Warp kernels for the Newton actuator fast path."""
 
 from collections.abc import Sequence
-from typing import Any
 
 import torch
 import warp as wp
 
 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),

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