[Warp] Part 6/12: Run commands and Rough velocity tasks on the Warp managers and apply term changes under capture

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

@@ -14,16 +14,12 @@ The following classes are part of the public :mod:`isaaclab_experimental.envs` A
    :nosignatures:
 
    DirectRLEnvWarp
-   InteractiveSceneWarp
    ManagerBasedEnvWarp
    ManagerBasedRLEnvWarp
 
 .. autoclass:: DirectRLEnvWarp
    :show-inheritance:
 
-.. autoclass:: InteractiveSceneWarp
-   :show-inheritance:
-
 .. autoclass:: ManagerBasedEnvWarp
    :show-inheritance:
 

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

@@ -8,8 +8,10 @@ Warp Environments
    The warp environment infrastructure lives in ``isaaclab_experimental`` and
    ``isaaclab_tasks_experimental``. It's an experimental feature.
 
-Experimental Warp environments compute observations, rewards, resets, and actions
-in Warp kernels. CUDA graph capture reduces Python overhead during execution.
+Experimental Warp environments compute MDP terms (observations, rewards, terminations, events, commands,
+and actions) in Warp kernels and record the capturable terms of each manager stage into CUDA graphs, which
+reduces Python overhead during execution. Scene writes, actuator models, the stable curriculum and recorder
+managers, and terms marked non-capturable still run eagerly.
 
 Here, **stable** means the default Torch environments in ``isaaclab`` and
 ``isaaclab_tasks``. Pass ``--frontend warp`` to use their experimental Warp counterparts.
@@ -94,19 +96,28 @@ MDP term for its warp twin). Select the Newton solver explicitly with
 - ``Isaac-Velocity-Flat-G1``
 - ``Isaac-Velocity-Flat-H1``
 - ``Isaac-Velocity-Flat-UnitreeGo2``
+- ``Isaac-Velocity-Rough-AnymalD``
+- ``Isaac-Velocity-Rough-Cassie``
+- ``Isaac-Velocity-Rough-G1``
+- ``Isaac-Velocity-Rough-H1``
+- ``Isaac-Velocity-Rough-UnitreeGo2``
 
 The following contributed tasks also have full twin coverage:
 
 - ``IsaacContrib-Velocity-Flat-AnymalB``
 - ``IsaacContrib-Velocity-Flat-AnymalC``
 - ``IsaacContrib-Velocity-Flat-UnitreeA1``
 - ``IsaacContrib-Velocity-Flat-UnitreeGo1``
+- ``IsaacContrib-Velocity-Rough-AnymalB``
+- ``IsaacContrib-Velocity-Rough-AnymalC``
+- ``IsaacContrib-Velocity-Rough-UnitreeA1``
+- ``IsaacContrib-Velocity-Rough-UnitreeGo1``
 
 A missing twin is a hard error listing the affected terms, so a partially
 covered task fails at build time rather than silently changing behavior.
-Rough-terrain velocity configurations that use ``height_scan`` cannot currently be
-adapted: that observation term has no Warp twin. Check the terms used by a task rather
-than assuming that every terrain configuration has the same limitation.
+The Rough tasks keep the stable terrain curriculum. Their ``height_scan``
+observation, and ``base_height_l2`` when it reads the height scanner, read a ray
+caster whose refresh is decided on the host, so these terms run eagerly.
 
 
 Performance Comparison
@@ -251,6 +262,10 @@ specific to warp envs; for Newton physics features and related guides see
 - Some scene-side operations (asset write, actuator models, certain sensor types) still go
   through torch. They participate in the step but are not yet captured into the graph; they
   set the lower bound on observed step time.
+- A term marked ``@WarpCapturable(False)``, or rejected by its parameter predicate, runs eagerly
+  between the recorded graphs of its stage. Observation, reward, termination, command and action
+  stages run their eager terms before the recorded ones; event stages keep the configured order, so
+  an eager event term splits its stage into several graphs.
 - Sensors that depend on the Kit RTX renderer (camera-based observations) cannot be combined
   with the warp env path — they need Kit, which the warp runtime does not initialise.
 
@@ -377,11 +392,12 @@ Project Structure
 Warp-specific implementations that diverge from the torch-based managers and env classes live in the ``_experimental`` packages:
 
 - ``isaaclab_experimental`` — warp managers, base env classes, warp MDP terms
-- ``isaaclab_tasks_experimental`` — warp task configs and task-specific MDP terms
+- ``isaaclab_tasks_experimental`` — task-specific warp MDP terms and direct warp envs
 
-Any new warp implementation that differs from the torch-based managers or env classes belongs in these packages.
-Warp task configs reference Newton physics directly (no ``PresetCfg``) since the warp path
-is Newton-only.
+Task configurations are not duplicated: ``--frontend warp`` builds the environment from the stable task
+configuration with Newton physics selected. The experimental packages mirror the stable module tree, and each
+stable term is replaced by the same-named warp term of the mirrored ``mdp`` package (its *twin*). A term
+without a twin is a build-time error.
 
 
 Writing Warp MDP Terms
@@ -390,15 +406,16 @@ Writing Warp MDP Terms
 Imports
 """""""
 
-Warp task configs import from the experimental packages:
+Warp terms import the warp managers' helpers from the experimental packages:
 
 .. code-block:: python
 
-   # Warp
-   from isaaclab_experimental.managers import ObservationTermCfg, RewardTermCfg, SceneEntityCfg
-   import isaaclab_experimental.envs.mdp as mdp
+   from isaaclab_experimental.managers import ManagerTermBase, SceneEntityCfg
+   from isaaclab_experimental.utils.warp import WarpCapturable
 
-The term config classes have the same interface — only the import path changes.
+Term configurations keep the stable config classes (for example
+:class:`~isaaclab.managers.RewardTermCfg` and the command configurations in :mod:`isaaclab.envs.mdp`);
+the frontend only replaces their ``func`` or ``class_type``.
 
 
 Common Pattern
@@ -471,7 +488,10 @@ array. This is resolved via a fallback chain (see
 
 3. **Legacy params**: ``term_dim``, ``out_dim``, or ``obs_dim`` keys in ``term_cfg.params``.
 
-4. **Asset config fallback**: Count of ``asset_cfg.joint_ids`` (or ``joint_ids_wp``) for
+4. **Ray-caster fallback**: Number of rays of the sensor in ``sensor_cfg``, for terms that
+   write one value per ray (e.g. ``height_scan``).
+
+5. **Asset config fallback**: Count of ``asset_cfg.joint_ids`` (or ``joint_ids_wp``) for
    joint-level terms.
 
 
@@ -493,7 +513,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)``
 
@@ -509,15 +530,96 @@ buffers.
 
 .. _warp-env-capture-safety:
 
+Command Terms
+"""""""""""""
+
+Command terms subclass :class:`isaaclab_experimental.managers.CommandTerm`. They keep their command and
+metrics in persistent ``wp.array`` buffers and resample the environments selected by a mask with
+:attr:`WarpRng.state <isaaclab.utils.seed.WarpRng.state>`; the command manager records ``compute`` and
+``reset`` like the other managers.
+Warp terms read a command with ``env.command_manager.get_command_wp(name)`` on every call instead of
+caching it, so a term never holds another environment's buffer.
+
+
 Capture Safety
 """"""""""""""
 
-When writing terms that run inside the captured step loop, keep in mind:
+Each manager method decorated with :func:`~isaaclab_experimental.utils.captured` is a stage (for example
+``RewardManager.compute``, or the ``_reset`` that ``CommandManager.reset`` calls with the reset mask): its
+first call records it into CUDA graphs and later calls replay them. That first call is the recording
+itself; there is no eager warm-up. A stage records again when it receives a different array or scalar
+argument, and every graph is dropped when the physics backend rebinds its buffers or stops. Only device
+work is replayed: Python in a recorded term runs once, while the stage records.
+
+Keeping a term safe to record is the author's contract; a synchronization check catches only
+part of it:
+
+- **Allocate in** ``__init__`` — class terms allocate every buffer before their first call and
+  function terms allocate nothing. A buffer allocated while recording belongs to the graph.
+- **Write into persistent buffers** — outputs, metrics and state are ``wp.array`` buffers that are
+  overwritten in place, never replaced. Values returned for logging, for example from ``reset``,
+  are persistent views over such buffers.
+- **Select environments with** ``env_mask`` — kernels launch over ``num_envs`` and test the mask.
+  Converting a mask to indices (``nonzero``) or reading device values on the host (``.item()``,
+  ``.cpu()``, ``.numpy()``) synchronizes the device and cannot be recorded. Only the environment
+  base classes do it, once per step, to decide whether any environment resets.
+- **Read asset data when called** — pass sim-bound (Tier 1) buffers such as ``root_link_pose_w``
+  and ``root_com_vel_w`` to the kernel and compute derived quantities inline. Lazily computed
+  properties refresh through a host-side timestamp that does not replay, and a view cached across
+  a rebind points at freed memory.
+- **Change settings through the managers** — a scalar read in Python, such as a configuration value
+  or ``env.common_step_counter``, is frozen into the graph when the stage records. The managers'
+  ``set_term_cfg`` records their stages again when it changes what they read from a term's configuration
+  (observation terms are named ``"<group>/<term>"``), and the ``modify_reward_weight`` and
+  ``modify_term_cfg`` curricula go through it, so an observation noise curriculum applies to recorded
+  observations. Setting an unchanged configuration keeps the recorded stages. Reward weights are read on
+  the device, so a weight change applies without recording again, unless the weight turns zero or nonzero,
+  which takes the term out of the stage or puts it back. A value that changes every step belongs in a
+  device buffer that the kernel reads. Writing a term configuration directly, bypassing ``set_term_cfg``,
+  is not supported while stages are recorded.
+- **Mark host-dependent terms** — a term that relies on host-side work, such as a sensor whose
+  refresh is decided on the host, is decorated with ``@WarpCapturable(False, reason=...)``. Its stage
+  calls it through :func:`~isaaclab_experimental.utils.eager`, between recorded graphs, on every call,
+  and it raises if it is ever called while a stage records. When only some configurations need the host,
+  pass a predicate over the term's parameters, e.g.
+  ``@WarpCapturable(lambda params: params.get("sensor_cfg") is None, reason=...)``. On a class term the
+  annotation covers its call; decorate its ``reset`` method to keep the reset out of the recording. A
+  command term's class annotation covers both its ``compute`` and its ``reset``. ``height_scan``,
+  ``base_height_l2`` with a height scanner and the rigid-body randomization events are examples.
+
+Setting :attr:`CapturedStage.enabled <isaaclab_experimental.utils.CapturedStage.enabled>` to ``False``
+after the environment is built runs every stage eagerly, and ``ISAACLAB_SYNC_DEBUG=1`` makes a hidden
+device-to-host synchronization inside an eager stage raise. Comparing an eager and a captured
+rollout of the same seeded task is the end-to-end check.
+
+Terms from Other Packages
+"""""""""""""""""""""""""
+
+Terms defined outside the mirrored packages, for example in an external project, have no twins.
+The frontend uses them unchanged when they are warp terms: a class term that subclasses
+:class:`isaaclab_experimental.managers.ManagerTermBase` (warp action and command terms included),
+or a function term that declares its capture behavior with ``@WarpCapturable(True)`` or
+``@WarpCapturable(False, reason=...)``. Any other term from such a package is rejected, even when a
+built-in warp term has the same name.
+
+.. code-block:: python
+
+   import warp as wp
+
+   from isaaclab_experimental.utils.warp import WarpCapturable
+
+
+   @wp.kernel
+   def _upright_kernel(pose: wp.array(dtype=wp.transformf), out: wp.array(dtype=wp.float32)):
+       i = wp.tid()
+       up = wp.quat_rotate(wp.transform_get_rotation(pose[i]), wp.vec3f(0.0, 0.0, 1.0))
+       out[i] = up[2]
+
 
-- **No ``wp.to_torch``** or torch arithmetic — stay in warp throughout
-- **No lazy-evaluated properties** — use sim-bound (Tier 1) data directly; if a derived
-  quantity is needed, compute it inline in the kernel
-- **No dynamic allocation** — all buffers must be pre-allocated in ``__init__``
+   @WarpCapturable(True)
+   def upright_bonus(env, out, asset_cfg) -> None:
+       asset = env.scene[asset_cfg.name]
+       wp.launch(_upright_kernel, dim=env.num_envs, inputs=[asset.data.root_link_pose_w.warp, out])
 
 
 Parity Testing
@@ -564,6 +666,9 @@ This test should always be run, even for terms without a torch counterpart.
    torch.testing.assert_close(wp.to_torch(out_captured), wp.to_torch(out_uncaptured))
 
 See ``source/isaaclab_experimental/test/envs/mdp/`` for complete parity test examples.
+``test_capture_safety.py`` discovers every warp MDP term: a new capturable term needs a ``CAPTURE_SPECS``
+entry there, which records the term, overwrites its inputs in place, replays it, and compares the result
+with the stable term on the new data.
 
 
 .. _warp-env-mdp-terms:
@@ -577,7 +682,7 @@ Available Warp MDP Terms
 
    * - Category
      - Available Terms
-   * - Observations (11)
+   * - Observations (12)
      - | ``base_pos_z``
        | ``base_lin_vel``
        | ``base_ang_vel``
@@ -589,6 +694,7 @@ Available Warp MDP Terms
        | ``joint_vel_rel``
        | ``last_action``
        | ``generated_commands``
+       | ``height_scan``
    * - Rewards (16)
      - | ``is_alive``
        | ``is_terminated``
@@ -621,6 +727,10 @@ Available Warp MDP Terms
    * - Actions (2)
      - | ``JointPositionAction``
        | ``JointEffortAction``
+   * - Commands (3)
+     - | ``UniformVelocityCommand``
+       | ``UniformPoseCommand``
+       | ``NullCommand``
 
 Terms not listed here remain in torch only. When using an env that requires unlisted terms,
 those terms must be implemented in warp first.

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

@@ -0,0 +1,4 @@
+* Added :func:`~isaaclab.utils.warp.memory.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.memory.resolve_1d_mask` does to a mask.
+* Added :func:`~isaaclab.utils.warp.memory.write_2d_with_mask`, which writes 2-D data or one scalar into a Warp array
+  where a row mask and a column mask are both true, launching over the mask lengths.

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` and
+  :class:`~isaaclab.sensors.Camera` to resolve ``env_ids`` and ``env_mask`` through
+  :func:`~isaaclab.utils.warp.memory.resolve_ids` and :func:`~isaaclab.utils.warp.memory.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`. The
+  :class:`~isaaclab.sensors.Camera` now rejects floating-point ``env_ids`` with a :class:`TypeError` and
+  multi-dimensional ``env_ids`` with a :class:`ValueError` instead of casting and flattening them.
+* Changed :class:`~isaaclab.utils.wrench_composer.WrenchComposer` to resolve ``body_ids`` through
+  :func:`~isaaclab.utils.warp.memory.resolve_ids` as well, so slices other than ``slice(None)`` are accepted.
+* Changed :func:`~isaaclab.utils.warp.memory.resolve_1d_mask` to make ``scratch_mask`` optional for callers that never
+  pass ids. Their masks are returned as given. Ids without a ``scratch_mask`` raise a :class:`ValueError`.

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

@@ -0,0 +1,4 @@
+* 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.
+* Fixed :meth:`~isaaclab.sensors.Camera.reset` with an ``env_mask`` on another device than the camera: the camera
+  poses are now updated through the mask resolved on the camera's device, like the timestamps and buffers.

source/isaaclab/changelog.d/jichuanh-warp-frontend-managers.changed.rst

@@ -0,0 +1,5 @@
+* Changed :class:`~isaaclab.envs.mdp.curriculums.modify_term_cfg` to hand the modified term configuration back to
+  the manager's ``set_term_cfg`` after writing a value, so managers that keep values read from the configuration,
+  such as the recorded stages of the Warp managers, use the new value.
+* Changed :class:`~isaaclab.envs.mdp.curriculums.modify_term_cfg` to hand an observation term back to the
+  observation manager's ``set_term_cfg`` under its ``"<group>/<term>"`` name.

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,3 @@
+* 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. It raises a ``RuntimeError``
+  when that state is not initialized.

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

@@ -0,0 +1,10 @@
+* 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.
+* Added ``env_mask`` to :meth:`~isaaclab.scene.InteractiveScene.reset`, which forwards the mask to every asset and
+  sensor so only the masked environments are reset, without synchronizing the host. Surface grippers still reset by
+  index.