From ce4db9c4c08e711c184065d5324b784b30cdde12 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 30 Sep 2026 02:03:58 -0700 Subject: [PATCH 01/80] 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 Date: Wed, 30 Sep 2026 02:05:55 -0700 Subject: [PATCH 02/80] 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 Date: Wed, 30 Sep 2026 02:11:37 -0700 Subject: [PATCH 03/80] 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 Date: Wed, 30 Sep 2026 02:23:05 -0700 Subject: [PATCH 04/80] 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 Date: Wed, 30 Sep 2026 03:34:56 -0700 Subject: [PATCH 05/80] 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 Date: Wed, 30 Sep 2026 03:43:37 -0700 Subject: [PATCH 06/80] 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 Date: Wed, 30 Sep 2026 03:58:52 -0700 Subject: [PATCH 07/80] 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 Date: Wed, 30 Sep 2026 02:54:05 -0700 Subject: [PATCH 08/80] 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 Date: Wed, 30 Sep 2026 03:01:33 -0700 Subject: [PATCH 09/80] 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 Date: Wed, 30 Sep 2026 03:12:38 -0700 Subject: [PATCH 10/80] 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 ` + 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 Date: Wed, 30 Sep 2026 03:28:37 -0700 Subject: [PATCH 11/80] 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 Date: Wed, 30 Sep 2026 03:46:19 -0700 Subject: [PATCH 12/80] 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 `, 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 ` 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 ` - replaying recorded event stages with the replaced configuration. +* Fixed :meth:`EventManager.set_term_cfg ` and + :meth:`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 ``"[]"`` 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 Date: Wed, 30 Sep 2026 03:47:31 -0700 Subject: [PATCH 13/80] 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 Date: Wed, 30 Sep 2026 04:05:46 -0700 Subject: [PATCH 14/80] 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 `, 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 Date: Wed, 30 Sep 2026 11:01:48 -0700 Subject: [PATCH 15/80] 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 Date: Wed, 30 Sep 2026 15:32:12 -0700 Subject: [PATCH 16/80] 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 Date: Wed, 30 Sep 2026 15:41:24 -0700 Subject: [PATCH 17/80] 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 Date: Wed, 30 Sep 2026 15:44:08 -0700 Subject: [PATCH 18/80] 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 1e520d6cc6e92ae4acf34ebc07c31896dc8d5644 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 30 Sep 2026 04:09:01 -0700 Subject: [PATCH 19/80] Run command terms through the Warp command manager Warp environments now build the experimental CommandManager and record its compute and reset stages. The frontend swaps command term class_types for their Warp twins (commands join WARP_MANAGED_GROUPS); the term configs stay the stable ones. Add Warp twins of UniformVelocityCommand, UniformPoseCommand and NullCommand. The pose twin keeps the per-episode sticky success flag that the pose_command_success termination writes, checks both the position and the orientation thresholds, and logs Metrics/success_rate like the stable term. The null twin never resamples, so its time_left stays finite. The debug visualization of the stable velocity and pose commands moves into private mixins shared with the twins, keeping marker_pos_offset and environment_ids. Warp reward and observation terms read the command through CommandManager.get_command_wp on every call instead of caching the first environment's buffer in a function attribute. --- .../jichuanh-warp-frontend-managers.skip | 0 .../isaaclab/envs/mdp/commands/_debug_vis.py | 132 ++++++++ .../envs/mdp/commands/pose_command.py | 42 +-- .../envs/mdp/commands/velocity_command.py | 76 +---- .../jichuanh-warp-frontend-managers.added.rst | 6 + .../jichuanh-warp-frontend-managers.fixed.rst | 2 + .../jichuanh-warp-frontend-managers.minor | 0 .../isaaclab_experimental/envs/frontend.py | 25 +- .../envs/manager_based_rl_env_warp.py | 28 +- .../envs/mdp/__init__.pyi | 3 + .../envs/mdp/commands/__init__.py | 14 + .../envs/mdp/commands/__init__.pyi | 14 + .../envs/mdp/commands/null_command.py | 78 +++++ .../envs/mdp/commands/pose_command.py | 278 ++++++++++++++++ .../envs/mdp/commands/velocity_command.py | 308 ++++++++++++++++++ .../envs/mdp/observations.py | 15 +- .../isaaclab_experimental/envs/mdp/rewards.py | 25 +- .../managers/command_manager.py | 27 +- .../test/envs/mdp/parity_helpers.py | 6 + .../test/envs/mdp/test_capture_safety.py | 22 +- .../envs/mdp/test_commands_warp_parity.py | 192 +++++++++++ .../envs/mdp/test_observations_warp_parity.py | 11 +- .../test/envs/mdp/test_rewards_warp_parity.py | 2 +- .../test/envs/test_warp_env_commands.py | 50 +++ .../jichuanh-warp-frontend-managers.fixed.rst | 2 + .../core/reach/mdp/rewards.py | 28 +- .../core/velocity/mdp/rewards.py | 62 ++-- 27 files changed, 1204 insertions(+), 244 deletions(-) create mode 100644 source/isaaclab/changelog.d/jichuanh-warp-frontend-managers.skip create mode 100644 source/isaaclab/isaaclab/envs/mdp/commands/_debug_vis.py create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.minor create mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.py create mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.pyi create mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/null_command.py create mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py create mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py create mode 100644 source/isaaclab_experimental/test/envs/mdp/test_commands_warp_parity.py create mode 100644 source/isaaclab_experimental/test/envs/test_warp_env_commands.py create mode 100644 source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst diff --git a/source/isaaclab/changelog.d/jichuanh-warp-frontend-managers.skip b/source/isaaclab/changelog.d/jichuanh-warp-frontend-managers.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/_debug_vis.py b/source/isaaclab/isaaclab/envs/mdp/commands/_debug_vis.py new file mode 100644 index 000000000000..dfdacd68e10d --- /dev/null +++ b/source/isaaclab/isaaclab/envs/mdp/commands/_debug_vis.py @@ -0,0 +1,132 @@ +# 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 + +"""Debug-visualization mixins shared by the command terms and their Warp twins.""" + +from __future__ import annotations + +import torch + +import isaaclab.utils.math as math_utils +from isaaclab.markers import VisualizationMarkers + + +class _VelocityCommandDebugVis: + """Velocity command markers: the commanded and the current base velocity as arrows. + + Reads ``robot``, ``cfg`` and the Torch command buffer ``vel_command_b`` from the command term. + """ + + def _set_debug_vis_impl(self, debug_vis: bool): + # set visibility of markers + # note: parent only deals with callbacks. not their visibility + if debug_vis: + # create markers if necessary for the first time + if not hasattr(self, "goal_vel_visualizer"): + # -- goal + self.goal_vel_visualizer = VisualizationMarkers(self.cfg.goal_vel_visualizer_cfg) + # -- current + self.current_vel_visualizer = VisualizationMarkers(self.cfg.current_vel_visualizer_cfg) + # set their visibility to true + self.goal_vel_visualizer.set_visibility(True) + self.current_vel_visualizer.set_visibility(True) + else: + if hasattr(self, "goal_vel_visualizer"): + self.goal_vel_visualizer.set_visibility(False) + self.current_vel_visualizer.set_visibility(False) + + def _debug_vis_callback(self, event): + # check if robot is initialized + # note: this is needed in-case the robot is de-initialized. we can't access the data + if not self.robot.is_initialized: + return + # get marker location + # -- base state + base_pos_w = self.robot.data.root_pos_w.torch.clone() + offset = self.cfg.marker_pos_offset + base_pos_w[:, 0] += offset[0] + base_pos_w[:, 1] += offset[1] + base_pos_w[:, 2] += offset[2] + # -- resolve the scales and quaternions + vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.vel_command_b[:, :2]) + vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow( + self.robot.data.root_lin_vel_b.torch[:, :2] + ) + # display markers + environment_ids = self._env.scene._ALL_INDICES + self.goal_vel_visualizer.visualize( + base_pos_w, + vel_des_arrow_quat, + vel_des_arrow_scale, + environment_ids=environment_ids, + ) + self.current_vel_visualizer.visualize( + base_pos_w, + vel_arrow_quat, + vel_arrow_scale, + environment_ids=environment_ids, + ) + + def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Converts the XY base velocity command to arrow direction rotation.""" + # obtain default scale of the marker + default_scale = self.goal_vel_visualizer.cfg.markers["arrow"].scale + # arrow-scale + arrow_scale = torch.tensor(default_scale, device=self.device).repeat(xy_velocity.shape[0], 1) + arrow_scale[:, 0] *= torch.linalg.norm(xy_velocity, dim=1) * 3.0 + # arrow-direction + heading_angle = torch.atan2(xy_velocity[:, 1], xy_velocity[:, 0]) + zeros = torch.zeros_like(heading_angle) + arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle) + # convert everything back from base to world frame + base_quat_w = self.robot.data.root_quat_w.torch + arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat) + + return arrow_scale, arrow_quat + + +class _PoseCommandDebugVis: + """Pose command markers: the goal pose and the current pose of the commanded body as frames. + + Reads ``robot``, ``body_idx``, ``cfg`` and the Torch world-frame command ``pose_command_w`` from the + command term. + """ + + def _set_debug_vis_impl(self, debug_vis: bool): + # create markers if necessary for the first time + if debug_vis: + if not hasattr(self, "goal_pose_visualizer"): + # -- goal pose + self.goal_pose_visualizer = VisualizationMarkers(self.cfg.goal_pose_visualizer_cfg) + # -- current body pose + self.current_pose_visualizer = VisualizationMarkers(self.cfg.current_pose_visualizer_cfg) + # set their visibility to true + self.goal_pose_visualizer.set_visibility(True) + self.current_pose_visualizer.set_visibility(True) + else: + if hasattr(self, "goal_pose_visualizer"): + self.goal_pose_visualizer.set_visibility(False) + self.current_pose_visualizer.set_visibility(False) + + def _debug_vis_callback(self, event): + # check if robot is initialized + # note: this is needed in-case the robot is de-initialized. we can't access the data + if not self.robot.is_initialized: + return + # update the markers + environment_ids = self._env.scene._ALL_INDICES + # -- goal pose + self.goal_pose_visualizer.visualize( + self.pose_command_w[:, :3], + self.pose_command_w[:, 3:], + environment_ids=environment_ids, + ) + # -- current body pose + body_link_pose_w = self.robot.data.body_link_pose_w.torch[:, self.body_idx] + self.current_pose_visualizer.visualize( + body_link_pose_w[:, :3], + body_link_pose_w[:, 3:7], + environment_ids=environment_ids, + ) diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py index 89ef3e5a5258..62d4952b0f18 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py @@ -14,17 +14,18 @@ from isaaclab.assets import Articulation from isaaclab.managers import CommandTerm -from isaaclab.markers import VisualizationMarkers from isaaclab.utils import index_fill_ from isaaclab.utils.leapp import POSE7_ELEMENT_NAMES from isaaclab.utils.math import combine_frame_transforms, compute_pose_error, quat_from_euler_xyz, quat_unique +from ._debug_vis import _PoseCommandDebugVis + if TYPE_CHECKING: from ... import ManagerBasedEnv from .commands_cfg import UniformPoseCommandCfg -class UniformPoseCommand(CommandTerm): +class UniformPoseCommand(_PoseCommandDebugVis, CommandTerm): """Command generator for generating pose commands uniformly. The command generator generates poses by sampling positions uniformly within specified @@ -181,40 +182,3 @@ def _resample_command(self, env_ids: Sequence[int]): def _update_command(self): pass - - def _set_debug_vis_impl(self, debug_vis: bool): - # create markers if necessary for the first time - if debug_vis: - if not hasattr(self, "goal_pose_visualizer"): - # -- goal pose - self.goal_pose_visualizer = VisualizationMarkers(self.cfg.goal_pose_visualizer_cfg) - # -- current body pose - self.current_pose_visualizer = VisualizationMarkers(self.cfg.current_pose_visualizer_cfg) - # set their visibility to true - self.goal_pose_visualizer.set_visibility(True) - self.current_pose_visualizer.set_visibility(True) - else: - if hasattr(self, "goal_pose_visualizer"): - self.goal_pose_visualizer.set_visibility(False) - self.current_pose_visualizer.set_visibility(False) - - def _debug_vis_callback(self, event): - # check if robot is initialized - # note: this is needed in-case the robot is de-initialized. we can't access the data - if not self.robot.is_initialized: - return - # update the markers - environment_ids = self._env.scene._ALL_INDICES - # -- goal pose - self.goal_pose_visualizer.visualize( - self.pose_command_w[:, :3], - self.pose_command_w[:, 3:], - environment_ids=environment_ids, - ) - # -- current body pose - body_link_pose_w = self.robot.data.body_link_pose_w.torch[:, self.body_idx] - self.current_pose_visualizer.visualize( - body_link_pose_w[:, :3], - body_link_pose_w[:, 3:7], - environment_ids=environment_ids, - ) diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py index f580e87d3052..e8545e8f1acd 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py @@ -16,9 +16,10 @@ import isaaclab.utils.math as math_utils from isaaclab.assets import Articulation from isaaclab.managers import CommandTerm -from isaaclab.markers import VisualizationMarkers from isaaclab.utils import index_fill_ +from ._debug_vis import _VelocityCommandDebugVis + if TYPE_CHECKING: from ... import ManagerBasedEnv from .commands_cfg import NormalVelocityCommandCfg, UniformVelocityCommandCfg @@ -26,7 +27,7 @@ logger = logging.getLogger(__name__) -class UniformVelocityCommand(CommandTerm): +class UniformVelocityCommand(_VelocityCommandDebugVis, CommandTerm): r"""Command generator that generates a velocity command in SE(2) from uniform distribution. The command comprises of a linear velocity in x and y direction and an angular velocity around @@ -191,77 +192,6 @@ def _update_command(self): # Enforce standing (i.e., zero velocity command) for standing envs index_fill_(self.vel_command_b, self.is_standing_env, 0.0) - def _set_debug_vis_impl(self, debug_vis: bool): - # set visibility of markers - # note: parent only deals with callbacks. not their visibility - if debug_vis: - # create markers if necessary for the first time - if not hasattr(self, "goal_vel_visualizer"): - # -- goal - self.goal_vel_visualizer = VisualizationMarkers(self.cfg.goal_vel_visualizer_cfg) - # -- current - self.current_vel_visualizer = VisualizationMarkers(self.cfg.current_vel_visualizer_cfg) - # set their visibility to true - self.goal_vel_visualizer.set_visibility(True) - self.current_vel_visualizer.set_visibility(True) - else: - if hasattr(self, "goal_vel_visualizer"): - self.goal_vel_visualizer.set_visibility(False) - self.current_vel_visualizer.set_visibility(False) - - def _debug_vis_callback(self, event): - # check if robot is initialized - # note: this is needed in-case the robot is de-initialized. we can't access the data - if not self.robot.is_initialized: - return - # get marker location - # -- base state - base_pos_w = self.robot.data.root_pos_w.torch.clone() - offset = self.cfg.marker_pos_offset - base_pos_w[:, 0] += offset[0] - base_pos_w[:, 1] += offset[1] - base_pos_w[:, 2] += offset[2] - # -- resolve the scales and quaternions - vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2]) - vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow( - self.robot.data.root_lin_vel_b.torch[:, :2] - ) - # display markers - environment_ids = self._env.scene._ALL_INDICES - self.goal_vel_visualizer.visualize( - base_pos_w, - vel_des_arrow_quat, - vel_des_arrow_scale, - environment_ids=environment_ids, - ) - self.current_vel_visualizer.visualize( - base_pos_w, - vel_arrow_quat, - vel_arrow_scale, - environment_ids=environment_ids, - ) - - """ - Internal helpers. - """ - - def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Converts the XY base velocity command to arrow direction rotation.""" - # obtain default scale of the marker - default_scale = self.goal_vel_visualizer.cfg.markers["arrow"].scale - # arrow-scale - arrow_scale = torch.tensor(default_scale, device=self.device).repeat(xy_velocity.shape[0], 1) - arrow_scale[:, 0] *= torch.linalg.norm(xy_velocity, dim=1) * 3.0 - # arrow-direction - heading_angle = torch.atan2(xy_velocity[:, 1], xy_velocity[:, 0]) - zeros = torch.zeros_like(heading_angle) - arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle) - # convert everything back from base to world frame - base_quat_w = self.robot.data.root_quat_w.torch - arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat) - - return arrow_scale, arrow_quat - class NormalVelocityCommand(UniformVelocityCommand): """Command generator that generates a velocity command in SE(2) from a normal distribution. diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst new file mode 100644 index 000000000000..6ffd6c18ad87 --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst @@ -0,0 +1,6 @@ +* Added Warp twins of the ``UniformVelocityCommand``, ``UniformPoseCommand`` and ``NullCommand`` command + terms. ``--frontend warp`` now runs command terms on the Warp command manager and records its compute and + reset stages. The pose command logs ``Metrics/success_rate`` from its position and orientation thresholds, + like the stable term. +* Added :meth:`~isaaclab_experimental.managers.CommandManager.get_command_wp` to read a command as a Warp + array. diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst new file mode 100644 index 000000000000..e6e1100aff3e --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst @@ -0,0 +1,2 @@ +* Fixed a second Warp velocity environment in the same process failing to build with a ``RecursionError``, and + Warp observation and reward terms reading the command buffer of the first environment created in the process. diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.minor b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.minor new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index cf2183c9310c..1d3eb394da88 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -92,13 +92,14 @@ class WarpFrontend: # are adapted (SceneEntityCfg promotion + MDP twin swap). The event manager # is warp-first too — it invokes term funcs with a Warp env-mask, so a stable # event func (which expects torch ``env_ids``) breaks at runtime; its funcs - # must be swapped to warp twins. The curriculum, recorder and command - # managers run on the stable (torch) implementation, so their terms are left - # untouched. A stable term left on a warp manager would break, so a missing - # twin in these groups is a hard error; a stable term on a stable manager is - # correct, so those groups are skipped. + # must be swapped to warp twins. Command terms are classes whose + # ``class_type`` swaps to a warp twin; their cfgs stay stable. The curriculum + # and recorder managers run on the stable (torch) implementation, so their + # terms are left untouched. A stable term left on a warp manager would break, + # so a missing twin in these groups is a hard error; a stable term on a + # stable manager is correct, so those groups are skipped. WARP_MANAGED_GROUPS: ClassVar[frozenset[str]] = frozenset( - {"observations", "rewards", "terminations", "actions", "events"} + {"observations", "rewards", "terminations", "actions", "events", "commands"} ) # ------------------------------------------------------------------ @@ -669,10 +670,10 @@ def _walk_terms(node: Any, path: tuple[str, ...] = ()) -> Iterator[tuple[tuple[s """Yield ``(path, term)`` for every MDP term cfg in the cfg tree. A "term" is a :class:`ManagerTermBaseCfg` (observation/reward/ - termination/event/curriculum) *or* an :class:`ActionTermCfg` — the - latter is a separate base that is **not** a ``ManagerTermBaseCfg`` - subclass, yet carries a swappable ``class_type``, so it must be matched - explicitly. + termination/event/curriculum) *or* an :class:`ActionTermCfg` / + :class:`CommandTermCfg` — separate bases that are **not** + ``ManagerTermBaseCfg`` subclasses, yet carry a swappable ``class_type``, + so they must be matched explicitly. Behavior at each node: @@ -700,9 +701,9 @@ def _walk_terms(node: Any, path: tuple[str, ...] = ()) -> Iterator[tuple[tuple[s cfg layouts (extra observation groups, new nesting, etc.) are picked up automatically as long as their terms subclass one of the term base cfgs. """ - from isaaclab.managers.manager_term_cfg import ActionTermCfg, ManagerTermBaseCfg + from isaaclab.managers.manager_term_cfg import ActionTermCfg, CommandTermCfg, ManagerTermBaseCfg - if isinstance(node, (ManagerTermBaseCfg, ActionTermCfg)): + if isinstance(node, (ManagerTermBaseCfg, ActionTermCfg, CommandTermCfg)): yield path, node # a term: yield and stop; never descend into params/func return if isinstance(node, Mapping): 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..9ea31d051f55 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 @@ -24,10 +24,10 @@ from isaaclab.envs.common import VecEnvStepReturn from isaaclab.envs.manager_based_rl_env_cfg import ManagerBasedRLEnvCfg -from isaaclab.managers import CommandManager, CurriculumManager +from isaaclab.managers import CurriculumManager from isaaclab.utils.timer import Timer -from isaaclab_experimental.managers import RewardManager, TerminationManager +from isaaclab_experimental.managers import CommandManager, RewardManager, TerminationManager from isaaclab_experimental.utils.torch_utils import clone_obs_buffer from isaaclab_experimental.utils.warp import zero_masked_int64 @@ -153,7 +153,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 implementation) + # -- command manager self.command_manager = CommandManager(self.cfg.commands, self) print("[INFO] Command Manager: ", self.command_manager) @@ -171,12 +171,8 @@ 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) - ) + # the stable curriculum manager resets by environment index + self._resets_by_index = self._resets_by_index or bool(self.curriculum_manager.active_terms) # setup the action and observation spaces for Gym self._configure_gym_env_spaces() @@ -318,7 +314,12 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: self._reset_terminated_envs() # -- update command - self.command_manager.compute(dt=float(self.step_dt)) + self._warp_graph_cache.call_steps( + "CommandManager_compute", + self.command_manager.stage_steps("compute"), + dt=float(self.step_dt), + timer=DEBUG_TIMER_STEP, + ) # -- step interval events if "interval" in self.event_manager.available_modes: @@ -533,7 +534,12 @@ def _reset_mask(self, env_mask: wp.array, env_ids: torch.Tensor | None = None) - curriculum_info = self.curriculum_manager.reset(env_ids=env_ids) # -- command + event + termination managers - command_info = self.command_manager.reset(env_ids=env_ids) + command_info = self._warp_graph_cache.call_steps( + "CommandManager_reset", + self.command_manager.stage_steps("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 ) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/__init__.pyi b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/__init__.pyi index a6dadcfce5bc..3ff2c2555e63 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/__init__.pyi +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/__init__.pyi @@ -20,6 +20,9 @@ from .actions import ( # noqa: F401 JointPositionActionCfg, ) +# Experimental Warp-first command terms. Their configurations stay the stable ones. +from .commands import NullCommand, UniformPoseCommand, UniformVelocityCommand # noqa: F401 + # Override stable terms with experimental Warp-first implementations. These leaf # modules are import-clean (no eager backend imports), so re-exporting them here # is safe. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.py new file mode 100644 index 000000000000..1e60915d6d74 --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.py @@ -0,0 +1,14 @@ +# 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 + +"""Experimental command terms (Warp-first). + +Provides Warp-first twins of the stable :mod:`isaaclab.envs.mdp.commands` terms. Their configurations are the +stable command configurations: the Warp frontend points their ``class_type`` at these twins. +""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.pyi b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.pyi new file mode 100644 index 000000000000..e0a69a9ab9c7 --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.pyi @@ -0,0 +1,14 @@ +# 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 + +__all__ = [ + "NullCommand", + "UniformPoseCommand", + "UniformVelocityCommand", +] + +from .null_command import NullCommand +from .pose_command import UniformPoseCommand +from .velocity_command import UniformVelocityCommand diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/null_command.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/null_command.py new file mode 100644 index 000000000000..970cc10c38ae --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/null_command.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 + +"""Warp-first command generator that does nothing.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch +import warp as wp + +from isaaclab_experimental.managers import CommandTerm + +if TYPE_CHECKING: + from isaaclab.envs.mdp.commands.commands_cfg import NullCommandCfg + + +class NullCommand(CommandTerm): + """Command generator that does nothing. + + Warp-first twin of :class:`isaaclab.envs.mdp.commands.NullCommand`. It does not generate any commands and + is used for environments that do not require any commands. Like the stable term, it never resamples, so + its infinite resampling time is never drawn. + """ + + cfg: NullCommandCfg + """Configuration for the command generator.""" + + def __str__(self) -> str: + msg = "NullCommand:\n" + msg += "\tCommand dimension: N/A\n" + msg += f"\tResampling time range: {self.cfg.resampling_time_range}" + return msg + + """ + Properties + """ + + @property + def command(self): + """Null command. + + Raises: + RuntimeError: No command is generated. Always raises this error. + """ + raise RuntimeError("NullCommandTerm does not generate any commands.") + + """ + Operations. + """ + + def reset( + self, + env_ids: Sequence[int] | torch.Tensor | None = None, + *, + env_mask: wp.array | None = None, + ) -> dict[str, torch.Tensor]: + return self.reset_extras + + def compute(self, dt: float): + pass + + """ + Implementation specific functions. + """ + + def _update_metrics(self): + pass + + def _resample_command(self, env_mask: wp.array): + pass + + def _update_command(self): + pass diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py new file mode 100644 index 000000000000..51742f531b25 --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py @@ -0,0 +1,278 @@ +# 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-first pose command generator.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch +import warp as wp + +from isaaclab.envs.mdp.commands._debug_vis import _PoseCommandDebugVis +from isaaclab.utils.leapp import POSE7_ELEMENT_NAMES + +from isaaclab_experimental.managers import CommandTerm + +if TYPE_CHECKING: + from isaaclab.assets import Articulation + from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.envs.mdp.commands.commands_cfg import UniformPoseCommandCfg + + +@wp.kernel +def _resample_pose_command( + env_mask: wp.array(dtype=wp.bool), + rng_state: wp.array(dtype=wp.uint32), + pos_x: wp.vec2f, + pos_y: wp.vec2f, + pos_z: wp.vec2f, + roll: wp.vec2f, + pitch: wp.vec2f, + yaw: wp.vec2f, + make_quat_unique: bool, + command: wp.array(dtype=wp.float32, ndim=2), +): + """Sample a new root-frame pose command: uniform position [m] and roll-pitch-yaw [rad] as (x, y, z, w).""" + env_id = wp.tid() + if env_mask[env_id]: + state = rng_state[env_id] + command[env_id, 0] = wp.randf(state, pos_x[0], pos_x[1]) + command[env_id, 1] = wp.randf(state, pos_y[0], pos_y[1]) + command[env_id, 2] = wp.randf(state, pos_z[0], pos_z[1]) + # extrinsic XYZ Euler angles, matching ``quat_from_euler_xyz`` + q_roll = wp.quat_from_axis_angle(wp.vec3f(1.0, 0.0, 0.0), wp.randf(state, roll[0], roll[1])) + q_pitch = wp.quat_from_axis_angle(wp.vec3f(0.0, 1.0, 0.0), wp.randf(state, pitch[0], pitch[1])) + q_yaw = wp.quat_from_axis_angle(wp.vec3f(0.0, 0.0, 1.0), wp.randf(state, yaw[0], yaw[1])) + quat = q_yaw * q_pitch * q_roll + if make_quat_unique and quat[3] < 0.0: + quat = -quat + command[env_id, 3] = quat[0] + command[env_id, 4] = quat[1] + command[env_id, 5] = quat[2] + command[env_id, 6] = quat[3] + rng_state[env_id] = state + + +@wp.kernel +def _update_pose_metrics( + command_b: wp.array(dtype=wp.float32, ndim=2), + root_pos_w: wp.array(dtype=wp.vec3f), + root_quat_w: wp.array(dtype=wp.quatf), + body_pos_w: wp.array(dtype=wp.vec3f, ndim=2), + body_quat_w: wp.array(dtype=wp.quatf, ndim=2), + body_idx: int, + check_position: bool, + position_threshold: float, + check_orientation: bool, + orientation_threshold: float, + command_w: wp.array(dtype=wp.float32, ndim=2), + position_error: wp.array(dtype=wp.float32), + orientation_error: wp.array(dtype=wp.float32), + succeeded: wp.array(dtype=wp.bool), +): + """Express the command in the world frame, compute the tracking errors and update the sticky success flag. + + The orientation error is the axis-angle magnitude of the body-to-goal rotation, ``2 atan2(|xyz|, |w|)``, as in + :func:`~isaaclab.utils.math.compute_pose_error`. Each configured threshold is an accept predicate + (``error < threshold``); with no threshold configured, the success flag is left untouched. + """ + env_id = wp.tid() + root_quat = root_quat_w[env_id] + goal_pos_b = wp.vec3f(command_b[env_id, 0], command_b[env_id, 1], command_b[env_id, 2]) + goal_quat_b = wp.quatf(command_b[env_id, 3], command_b[env_id, 4], command_b[env_id, 5], command_b[env_id, 6]) + goal_pos_w = root_pos_w[env_id] + wp.quat_rotate(root_quat, goal_pos_b) + goal_quat_w = root_quat * goal_quat_b + command_w[env_id, 0] = goal_pos_w[0] + command_w[env_id, 1] = goal_pos_w[1] + command_w[env_id, 2] = goal_pos_w[2] + command_w[env_id, 3] = goal_quat_w[0] + command_w[env_id, 4] = goal_quat_w[1] + command_w[env_id, 5] = goal_quat_w[2] + command_w[env_id, 6] = goal_quat_w[3] + + pos_error = wp.length(body_pos_w[env_id, body_idx] - goal_pos_w) + quat_error = body_quat_w[env_id, body_idx] * wp.quat_inverse(goal_quat_w) + rot_error = 2.0 * wp.atan2(wp.length(wp.vec3f(quat_error[0], quat_error[1], quat_error[2])), wp.abs(quat_error[3])) + position_error[env_id] = pos_error + orientation_error[env_id] = rot_error + + if check_position or check_orientation: + success = bool(True) + if check_position: + success = success and (pos_error < position_threshold) + if check_orientation: + success = success and (rot_error < orientation_threshold) + if success: + succeeded[env_id] = True + + +@wp.kernel +def _finalize_pose_success( + env_mask: wp.array(dtype=wp.bool), + succeeded: wp.array(dtype=wp.bool), + success_rate: wp.array(dtype=wp.float32), +): + """Move the episode's sticky success flag into the success-rate metric and clear it for the next episode.""" + env_id = wp.tid() + if env_mask[env_id]: + success_rate[env_id] = wp.where(succeeded[env_id], 1.0, 0.0) + succeeded[env_id] = False + + +class UniformPoseCommand(_PoseCommandDebugVis, CommandTerm): + """Command generator for generating pose commands uniformly. + + Warp-first twin of :class:`isaaclab.envs.mdp.commands.UniformPoseCommand`. Positions are sampled uniformly + within the configured ranges and orientations from uniform roll-pitch-yaw angles, in the robot's root frame. + Commands are sampled from the per-environment Warp random state, so they follow the stable distribution + but not its random sequence. + + When a success threshold is configured, the per-episode "ever successful" flag is kept in :attr:`_succeeded` + (shared with the ``pose_command_success`` termination) and logged under ``Metrics/success_rate``. + + The command buffers are Warp arrays; :attr:`pose_command_b`, :attr:`pose_command_w` and :attr:`_succeeded` + are their zero-copy Torch views. + """ + + cfg: UniformPoseCommandCfg + """Configuration for the command generator.""" + + def __init__(self, cfg: UniformPoseCommandCfg, env: ManagerBasedRLEnv): + """Initialize the command generator class. + + Args: + cfg: The configuration parameters for the command generator. + env: The environment object. + """ + super().__init__(cfg, env) + + # extract the robot and body index for which the command is generated + self.robot: Articulation = env.scene[cfg.asset_name] + self.body_idx = self.robot.find_bodies(cfg.body_name)[0][0] + + # -- commands: (x, y, z, qx, qy, qz, qw) in root frame and in world frame + self._pose_command_b_wp = wp.zeros((self.num_envs, 7), dtype=wp.float32, device=self.device) + self._pose_command_w_wp = wp.zeros_like(self._pose_command_b_wp) + self.pose_command_b = wp.to_torch(self._pose_command_b_wp) + self.pose_command_w = wp.to_torch(self._pose_command_w_wp) + self.pose_command_b[:, 6] = 1.0 + # -- metrics + self.metrics["position_error"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self.metrics["orientation_error"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + # -- per-episode sticky success bit (only used when at least one success threshold is set) + self._track_success = ( + cfg.position_success_threshold is not None or cfg.orientation_success_threshold is not None + ) + self._succeeded_wp = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) + if self._track_success: + self._succeeded = wp.to_torch(self._succeeded_wp) + self.metrics["success_rate"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + + # adds (optional) cmd kind and element names for leapp export + self.cfg.cmd_kind = self.cfg.cmd_kind or "command/body/pose" + self.cfg.element_names = self.cfg.element_names or POSE7_ELEMENT_NAMES + + def __str__(self) -> str: + msg = "UniformPoseCommand:\n" + msg += f"\tCommand dimension: {tuple(self.command.shape[1:])}\n" + msg += f"\tResampling time range: {self.cfg.resampling_time_range}\n" + return msg + + """ + Properties + """ + + @property + def command(self) -> wp.array: + """The desired pose command. Shape is (num_envs, 7). + + The first three elements correspond to the position, followed by the quaternion orientation in (x, y, z, w). + """ + return self._pose_command_b_wp + + """ + Operations. + """ + + def reset( + self, + env_ids: Sequence[int] | torch.Tensor | None = None, + *, + env_mask: wp.array | None = None, + ) -> dict[str, torch.Tensor]: + """Log the ending episode's metrics, then reset the command of the selected environments. + + 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 the environments to reset. + Takes precedence over ``env_ids``. + + Returns: + Persistent views of the metric means over the reset environments. + """ + env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + if self._track_success: + wp.launch( + kernel=_finalize_pose_success, + dim=self.num_envs, + inputs=[env_mask, self._succeeded_wp, self.metrics["success_rate"]], + device=self.device, + ) + return super().reset(env_mask=env_mask) + + """ + Implementation specific functions. + """ + + def _update_metrics(self): + position_threshold = self.cfg.position_success_threshold + orientation_threshold = self.cfg.orientation_success_threshold + wp.launch( + kernel=_update_pose_metrics, + dim=self.num_envs, + inputs=[ + self._pose_command_b_wp, + self.robot.data.root_pos_w.warp, + self.robot.data.root_quat_w.warp, + self.robot.data.body_pos_w.warp, + self.robot.data.body_quat_w.warp, + self.body_idx, + position_threshold is not None, + float(position_threshold or 0.0), + orientation_threshold is not None, + float(orientation_threshold or 0.0), + self._pose_command_w_wp, + self.metrics["position_error"], + self.metrics["orientation_error"], + self._succeeded_wp, + ], + device=self.device, + ) + + def _resample_command(self, env_mask: wp.array): + ranges = self.cfg.ranges + wp.launch( + kernel=_resample_pose_command, + dim=self.num_envs, + inputs=[ + env_mask, + self._env.rng_state_wp, + wp.vec2f(*ranges.pos_x), + wp.vec2f(*ranges.pos_y), + wp.vec2f(*ranges.pos_z), + wp.vec2f(*ranges.roll), + wp.vec2f(*ranges.pitch), + wp.vec2f(*ranges.yaw), + self.cfg.make_quat_unique, + self._pose_command_b_wp, + ], + device=self.device, + ) + + def _update_command(self): + pass diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py new file mode 100644 index 000000000000..55703a72afd5 --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py @@ -0,0 +1,308 @@ +# 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-first velocity command generator.""" + +from __future__ import annotations + +import logging +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch +import warp as wp +from isaaclab_newton.kernels.state_kernels import body_ang_vel_from_root, body_lin_vel_from_root + +from isaaclab.envs.mdp.commands._debug_vis import _VelocityCommandDebugVis + +from isaaclab_experimental.managers import CommandTerm +from isaaclab_experimental.utils.warp import wrap_to_pi + +if TYPE_CHECKING: + from isaaclab.assets import Articulation + from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.envs.mdp.commands.commands_cfg import UniformVelocityCommandCfg + +logger = logging.getLogger(__name__) + + +@wp.kernel +def _accumulate_velocity_errors( + command: wp.array(dtype=wp.float32, ndim=2), + root_pose_w: wp.array(dtype=wp.transformf), + root_vel_w: wp.array(dtype=wp.spatial_vectorf), + error_xy_sum: wp.array(dtype=wp.float32), + error_yaw_sum: wp.array(dtype=wp.float32), + step_count: wp.array(dtype=wp.float32), +): + """Add this step's planar velocity error [m/s] and yaw-rate error [rad/s] to the episode sums.""" + env_id = wp.tid() + lin_vel_b = body_lin_vel_from_root(root_pose_w[env_id], root_vel_w[env_id]) + ang_vel_b = body_ang_vel_from_root(root_pose_w[env_id], root_vel_w[env_id]) + dx = command[env_id, 0] - lin_vel_b[0] + dy = command[env_id, 1] - lin_vel_b[1] + error_xy_sum[env_id] += wp.sqrt(dx * dx + dy * dy) + error_yaw_sum[env_id] += wp.abs(command[env_id, 2] - ang_vel_b[2]) + step_count[env_id] += 1.0 + + +@wp.kernel +def _finalize_velocity_metrics( + env_mask: wp.array(dtype=wp.bool), + error_xy_sum: wp.array(dtype=wp.float32), + error_yaw_sum: wp.array(dtype=wp.float32), + step_count: wp.array(dtype=wp.float32), + xy_threshold: float, + yaw_threshold: float, + error_xy: wp.array(dtype=wp.float32), + error_yaw: wp.array(dtype=wp.float32), + success_rate: wp.array(dtype=wp.float32), +): + """Turn the ending episode's error sums into means and a success flag, then clear the sums.""" + env_id = wp.tid() + if env_mask[env_id]: + denominator = wp.max(step_count[env_id], 1.0) + mean_error_xy = error_xy_sum[env_id] / denominator + mean_error_yaw = error_yaw_sum[env_id] / denominator + error_xy[env_id] = mean_error_xy + error_yaw[env_id] = mean_error_yaw + success_rate[env_id] = wp.where((mean_error_xy < xy_threshold) and (mean_error_yaw < yaw_threshold), 1.0, 0.0) + error_xy_sum[env_id] = 0.0 + error_yaw_sum[env_id] = 0.0 + step_count[env_id] = 0.0 + + +@wp.kernel +def _resample_velocity_command( + env_mask: wp.array(dtype=wp.bool), + rng_state: wp.array(dtype=wp.uint32), + lin_vel_x: wp.vec2f, + lin_vel_y: wp.vec2f, + ang_vel_z: wp.vec2f, + heading: wp.vec2f, + heading_command: bool, + rel_heading_envs: float, + rel_standing_envs: float, + command: wp.array(dtype=wp.float32, ndim=2), + heading_target: wp.array(dtype=wp.float32), + is_heading_env: wp.array(dtype=wp.bool), + is_standing_env: wp.array(dtype=wp.bool), +): + """Sample a new base velocity command, heading target and heading/standing roles for the masked envs.""" + env_id = wp.tid() + if env_mask[env_id]: + state = rng_state[env_id] + command[env_id, 0] = wp.randf(state, lin_vel_x[0], lin_vel_x[1]) + command[env_id, 1] = wp.randf(state, lin_vel_y[0], lin_vel_y[1]) + command[env_id, 2] = wp.randf(state, ang_vel_z[0], ang_vel_z[1]) + if heading_command: + heading_target[env_id] = wp.randf(state, heading[0], heading[1]) + is_heading_env[env_id] = wp.randf(state, 0.0, 1.0) <= rel_heading_envs + is_standing_env[env_id] = wp.randf(state, 0.0, 1.0) <= rel_standing_envs + rng_state[env_id] = state + + +@wp.kernel +def _update_velocity_command( + root_pose_w: wp.array(dtype=wp.transformf), + heading_target: wp.array(dtype=wp.float32), + is_heading_env: wp.array(dtype=wp.bool), + is_standing_env: wp.array(dtype=wp.bool), + heading_command: bool, + heading_control_stiffness: float, + ang_vel_z: wp.vec2f, + command: wp.array(dtype=wp.float32, ndim=2), +): + """Track the heading target with a clipped yaw-rate command and zero the command of standing envs.""" + env_id = wp.tid() + if heading_command and is_heading_env[env_id]: + # heading of the base x-axis in the world frame + forward_w = wp.quat_rotate(wp.transform_get_rotation(root_pose_w[env_id]), wp.vec3f(1.0, 0.0, 0.0)) + heading_error = wrap_to_pi(heading_target[env_id] - wp.atan2(forward_w[1], forward_w[0])) + command[env_id, 2] = wp.clamp(heading_control_stiffness * heading_error, ang_vel_z[0], ang_vel_z[1]) + if is_standing_env[env_id]: + command[env_id, 0] = 0.0 + command[env_id, 1] = 0.0 + command[env_id, 2] = 0.0 + + +class UniformVelocityCommand(_VelocityCommandDebugVis, CommandTerm): + r"""Command generator that generates a velocity command in SE(2) from uniform distribution. + + Warp-first twin of :class:`isaaclab.envs.mdp.commands.UniformVelocityCommand`. The command comprises a + linear velocity in x and y direction and an angular velocity around the z-axis, in the robot's base frame. + Commands are sampled from the per-environment Warp random state, so they follow the stable distribution + but not its random sequence. + + The command buffer is a Warp array; :attr:`vel_command_b` is its zero-copy Torch view. + """ + + cfg: UniformVelocityCommandCfg + """The configuration of the command generator.""" + + def __init__(self, cfg: UniformVelocityCommandCfg, env: ManagerBasedRLEnv): + """Initialize the command generator. + + Args: + cfg: The configuration of the command generator. + env: The environment. + + Raises: + ValueError: If the heading command is active but the heading range is not provided. + """ + super().__init__(cfg, env) + + # check configuration + if self.cfg.heading_command and self.cfg.ranges.heading is None: + raise ValueError( + "The velocity command has heading commands active (heading_command=True) but the `ranges.heading`" + " parameter is set to None." + ) + if self.cfg.ranges.heading and not self.cfg.heading_command: + logger.warning( + f"The velocity command has the 'ranges.heading' attribute set to '{self.cfg.ranges.heading}'" + " but the heading command is not active. Consider setting the flag for the heading command to True." + ) + + # obtain the robot asset + self.robot: Articulation = env.scene[cfg.asset_name] + + # -- command: x vel, y vel, yaw vel + self._vel_command_b_wp = wp.zeros((self.num_envs, 3), dtype=wp.float32, device=self.device) + self.vel_command_b = wp.to_torch(self._vel_command_b_wp) + self._heading_target_wp = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self._is_heading_env_wp = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) + self._is_standing_env_wp = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) + # -- metrics: finalized per-episode means/rates, written at reset() and read by the base class + self.metrics["error_vel_xy"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self.metrics["error_vel_yaw"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self.metrics["success_rate"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + # -- per-episode running sums (cleared at episode reset) + self._error_xy_sum = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self._error_yaw_sum = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self._step_count = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + + # adds (optional) cmd kind and element names for leapp export + self.cfg.cmd_kind = self.cfg.cmd_kind or "command/body/velocity" + self.cfg.element_names = self.cfg.element_names or ["lin_vel_x", "lin_vel_y", "ang_vel_z"] + + def __str__(self) -> str: + """Return a string representation of the command generator.""" + msg = "UniformVelocityCommand:\n" + msg += f"\tCommand dimension: {tuple(self.command.shape[1:])}\n" + msg += f"\tResampling time range: {self.cfg.resampling_time_range}\n" + msg += f"\tHeading command: {self.cfg.heading_command}\n" + if self.cfg.heading_command: + msg += f"\tHeading probability: {self.cfg.rel_heading_envs}\n" + msg += f"\tStanding probability: {self.cfg.rel_standing_envs}" + return msg + + """ + Properties + """ + + @property + def command(self) -> wp.array: + """The desired base velocity command in the base frame [m/s, m/s, rad/s]. Shape is (num_envs, 3).""" + return self._vel_command_b_wp + + """ + Operations. + """ + + def reset( + self, + env_ids: Sequence[int] | torch.Tensor | None = None, + *, + env_mask: wp.array | None = None, + ) -> dict[str, torch.Tensor]: + """Finalize the ending episode's metrics, then reset the command of the selected environments. + + 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 the environments to reset. + Takes precedence over ``env_ids``. + + Returns: + Persistent views of the metric means over the reset environments. + """ + env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + wp.launch( + kernel=_finalize_velocity_metrics, + dim=self.num_envs, + inputs=[ + env_mask, + self._error_xy_sum, + self._error_yaw_sum, + self._step_count, + self.cfg.vel_xy_success_threshold, + self.cfg.vel_yaw_success_threshold, + self.metrics["error_vel_xy"], + self.metrics["error_vel_yaw"], + self.metrics["success_rate"], + ], + device=self.device, + ) + return super().reset(env_mask=env_mask) + + """ + Implementation specific functions. + """ + + def _update_metrics(self): + # the per-episode mean is finalized in reset(), independent of episode length + wp.launch( + kernel=_accumulate_velocity_errors, + dim=self.num_envs, + inputs=[ + self._vel_command_b_wp, + self.robot.data.root_link_pose_w.warp, + self.robot.data.root_com_vel_w.warp, + self._error_xy_sum, + self._error_yaw_sum, + self._step_count, + ], + device=self.device, + ) + + def _resample_command(self, env_mask: wp.array): + ranges = self.cfg.ranges + wp.launch( + kernel=_resample_velocity_command, + dim=self.num_envs, + inputs=[ + env_mask, + self._env.rng_state_wp, + wp.vec2f(*ranges.lin_vel_x), + wp.vec2f(*ranges.lin_vel_y), + wp.vec2f(*ranges.ang_vel_z), + wp.vec2f(*(ranges.heading or (0.0, 0.0))), + self.cfg.heading_command, + self.cfg.rel_heading_envs, + self.cfg.rel_standing_envs, + self._vel_command_b_wp, + self._heading_target_wp, + self._is_heading_env_wp, + self._is_standing_env_wp, + ], + device=self.device, + ) + + def _update_command(self): + wp.launch( + kernel=_update_velocity_command, + dim=self.num_envs, + inputs=[ + self.robot.data.root_link_pose_w.warp, + self._heading_target_wp, + self._is_heading_env_wp, + self._is_standing_env_wp, + self.cfg.heading_command, + self.cfg.heading_control_stiffness, + wp.vec2f(*self.cfg.ranges.ang_vel_z), + self._vel_command_b_wp, + ], + device=self.device, + ) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py index 4bc49cc1e2ae..d76f448dd21e 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py @@ -353,20 +353,9 @@ def generated_commands(env: ManagerBasedEnv, out, command_name: str) -> None: """The generated command from the command manager. Writes into ``out``. Warp-first override of :func:`isaaclab.envs.mdp.observations.generated_commands`. - Uses ``wp.from_torch`` to create a zero-copy warp view of the command tensor on first call. + Reads the command manager's Warp command array on every call. """ - # TODO(warp-migration): Cross-manager access (observation → command). Replace with direct - # warp getter once all managers are guaranteed to be warp-native. - fn = generated_commands - if not getattr(fn, "_is_warmed_up", False) or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - if isinstance(cmd, wp.array): - fn._cmd_wp = cmd - else: - fn._cmd_wp = wp.from_torch(cmd) - fn._cmd_name = command_name - fn._is_warmed_up = True - wp.copy(out, fn._cmd_wp) + wp.copy(out, env.command_manager.get_command_wp(command_name)) """ diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py index 880697b562f2..fb5a92637ac9 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py @@ -509,24 +509,13 @@ def track_lin_vel_xy_exp( Warp-first override of :func:`isaaclab.envs.mdp.rewards.track_lin_vel_xy_exp`. """ asset: Articulation = env.scene[asset_cfg.name] - # cache the warp view of the command tensor on first call (zero-copy) - # TODO(warp-migration): Cross-manager access (reward → command). Replace with direct - # warp getter once all managers are guaranteed to be warp-native. - if not getattr(track_lin_vel_xy_exp, "_is_warmed_up", False) or track_lin_vel_xy_exp._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - if isinstance(cmd, wp.array): - track_lin_vel_xy_exp._cmd_wp = cmd - else: - track_lin_vel_xy_exp._cmd_wp = wp.from_torch(cmd) - track_lin_vel_xy_exp._cmd_name = command_name - track_lin_vel_xy_exp._is_warmed_up = True wp.launch( kernel=_track_lin_vel_xy_exp_kernel, dim=env.num_envs, inputs=[ asset.data.root_link_pose_w.warp, asset.data.root_com_vel_w.warp, - track_lin_vel_xy_exp._cmd_wp, + env.command_manager.get_command_wp(command_name), 1.0 / (std * std), out, ], @@ -560,23 +549,13 @@ def track_ang_vel_z_exp( Warp-first override of :func:`isaaclab.envs.mdp.rewards.track_ang_vel_z_exp`. """ asset: Articulation = env.scene[asset_cfg.name] - # TODO(warp-migration): Cross-manager access (reward → command). Replace with direct - # warp getter once all managers are guaranteed to be warp-native. - if not getattr(track_ang_vel_z_exp, "_is_warmed_up", False) or track_ang_vel_z_exp._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - if isinstance(cmd, wp.array): - track_ang_vel_z_exp._cmd_wp = cmd - else: - track_ang_vel_z_exp._cmd_wp = wp.from_torch(cmd) - track_ang_vel_z_exp._cmd_name = command_name - track_ang_vel_z_exp._is_warmed_up = True wp.launch( kernel=_track_ang_vel_z_exp_kernel, dim=env.num_envs, inputs=[ asset.data.root_link_pose_w.warp, asset.data.root_com_vel_w.warp, - track_ang_vel_z_exp._cmd_wp, + env.command_manager.get_command_wp(command_name), 2, 1.0 / (std * std), out, diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py index e7d8b4b746b0..0d5f350b7116 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py @@ -19,6 +19,7 @@ from isaaclab.managers.manager_term_cfg import CommandTermCfg from isaaclab.utils import instantiate +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 @@ -402,7 +403,11 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv): self._reset_extras: dict[str, torch.Tensor] = {} for term_name, term in self._terms.items(): for metric_name, metric_value in term.reset_extras.items(): - self._reset_extras[f"Metrics/{term_name}/{metric_name}"] = metric_value + # success rates share the unified ``Metrics/success_rate`` key across tasks, as in the stable terms + if metric_name == "success_rate": + self._reset_extras["Metrics/success_rate"] = metric_value + else: + self._reset_extras[f"Metrics/{term_name}/{metric_name}"] = metric_value def __str__(self) -> str: """Returns: A string representation for the command manager.""" @@ -546,6 +551,23 @@ def get_command(self, name: str) -> torch.Tensor: return wp.to_torch(command) return command + def get_command_wp(self, name: str) -> wp.array: + """Returns the command for the specified command term as a Warp array. + + Warp terms read the command through this accessor on every call rather than caching it, so a + recorded stage always refers to the buffers of the current environment. + + Args: + name: The name of the command term. + + Returns: + The command array of the specified command term. Shape is (num_envs, command_dim). + """ + command = self._terms[name].command + if isinstance(command, torch.Tensor): + return wp.from_torch(command) + return command + def get_term(self, name: str) -> CommandTerm: """Returns the command term with the specified name. @@ -585,5 +607,8 @@ def _prepare_terms(self): raise TypeError(f"Returned object for the term '{term_name}' is not of type CommandType.") # pre-build reset extras once for capture-friendly reset logging term._prepare_reset_extras() + # a command term that is not capturable keeps the command stages eager + if not is_warp_capturable(term): + self._all_terms_capturable = False # add class to dict self._terms[term_name] = term diff --git a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py index e9e07bca5e9f..6afeab79dcb3 100644 --- a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py +++ b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py @@ -547,6 +547,9 @@ def __init__(self, command_tensor: torch.Tensor, cmd_term: MockCommandTerm): def get_command(self, name: str) -> torch.Tensor: return self._cmd + def get_command_wp(self, name: str) -> wp.array: + return wp.from_torch(self._cmd) + def get_term(self, name: str): return self._term @@ -647,6 +650,9 @@ def __init__(self, term): def get_command(self, name: str) -> torch.Tensor: return self._term.pose_command_b + def get_command_wp(self, name: str) -> wp.array: + return wp.from_torch(self._term.pose_command_b) + def get_term(self, name: str): return self._term 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..ad77fbe3db51 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py @@ -429,12 +429,28 @@ def _ids(module: str, *names: str) -> list[str]: + _ids(f"{_TE}.velocity.mdp.terminations", "terrain_out_of_bounds"), _MIRROR, ), + # command terms are classes driven through compute() and reset(), which this harness does not model + **dict.fromkeys( + _ids(f"{_SH}.commands.pose_command", "UniformPoseCommand") + + _ids(f"{_SH}.commands.velocity_command", "UniformVelocityCommand"), + "covered by the capture-mutate-replay checks in test_commands_warp_parity.py", + ), + **dict.fromkeys(_ids(f"{_SH}.commands.null_command", "NullCommand"), "launches no kernels"), } -# Every modality a manager can run. ``actions`` is included because ``ActionManager`` is -# graph-captured like the rest, so an unsafe action term would otherwise bypass every check. -_MDP_LEAF_MODULES = ("rewards", "terminations", "observations", "events", "actions.joint_actions") +# Every modality a manager can run. ``actions`` and ``commands`` are included because their managers +# are graph-captured like the rest, so an unsafe action or command term would otherwise bypass every check. +_MDP_LEAF_MODULES = ( + "rewards", + "terminations", + "observations", + "events", + "actions.joint_actions", + "commands.null_command", + "commands.pose_command", + "commands.velocity_command", +) def _warp_mdp_modules() -> list: diff --git a/source/isaaclab_experimental/test/envs/mdp/test_commands_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_commands_warp_parity.py new file mode 100644 index 000000000000..9ac58e265bb4 --- /dev/null +++ b/source/isaaclab_experimental/test/envs/mdp/test_commands_warp_parity.py @@ -0,0 +1,192 @@ +# 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 + +"""Parity tests for the warp-first command terms. + +The command manager runs captured, so each command term is recorded once, has its inputs overwritten +in place, and is replayed; the result is compared against the stable term evaluated on the new data. +""" + +from __future__ import annotations + +import math +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +import warp as wp + +# Skip entire module if no CUDA device available +wp.init() +pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required") + +from isaaclab_experimental.envs.mdp.commands import NullCommand, UniformPoseCommand, UniformVelocityCommand +from isaaclab_experimental.utils.warp import resolve_1d_mask +from parity_helpers import ( + DEVICE, + NUM_BODIES, + NUM_ENVS, + MockArticulation, + MockArticulationData, + assert_close, + assert_equal, + copy_np_to_wp, + make_pose_command_term, + mutate_body_data, + quat_rotate_inv_np, +) + +import isaaclab.envs.mdp.commands.velocity_command as stable_velocity_command +from isaaclab.envs.mdp.commands.commands_cfg import NullCommandCfg, UniformPoseCommandCfg, UniformVelocityCommandCfg + +_DT = 0.02 + + +def _command_env(robot) -> SimpleNamespace: + """The environment surface a command term reads.""" + all_mask = wp.ones(NUM_ENVS, dtype=wp.bool, device=DEVICE) + scratch_mask = wp.zeros(NUM_ENVS, dtype=wp.bool, device=DEVICE) + return SimpleNamespace( + scene={"robot": robot}, + num_envs=NUM_ENVS, + device=DEVICE, + rng_state_wp=wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 7, device=DEVICE), + resolve_env_mask=lambda env_ids=None, env_mask=None: resolve_1d_mask( + ids=env_ids, mask=env_mask, all_mask=all_mask, scratch_mask=scratch_mask, device=DEVICE + ), + sim=SimpleNamespace(vis_marker_registry=SimpleNamespace(clear_debug_vis_callback=lambda term: None)), + ) + + +def _make_term(term_class, cfg, robot): + """Construct a command term the way :class:`CommandManager` does.""" + term = term_class(cfg, _command_env(robot)) + term._prepare_reset_extras() + return term + + +def _capture_compute(term) -> wp.Graph: + """Record one :meth:`compute` without resampling; the resampling timer is pushed past the test.""" + term.time_left_wp.fill_(1.0e6) + term.compute(_DT) # warm-up outside the capture + with wp.ScopedCapture() as capture: + term.compute(_DT) + return capture.graph + + +def test_pose_command_matches_stable_after_capture_and_mutation(): + """World-frame goal, tracking errors and the sticky success flag follow the stable term on replay.""" + art_data = MockArticulationData(num_bodies=NUM_BODIES) + robot = MockArticulation(art_data, num_bodies=NUM_BODIES) + stable = make_pose_command_term(robot) + stable.metrics = {} + cfg = UniformPoseCommandCfg( + asset_name="robot", + body_name="body", + resampling_time_range=(1.0, 1.0), + position_success_threshold=stable.cfg.position_success_threshold, + orientation_success_threshold=stable.cfg.orientation_success_threshold, + ranges=UniformPoseCommandCfg.Ranges( + pos_x=(0.0, 0.0), pos_y=(0.0, 0.0), pos_z=(0.0, 0.0), roll=(0.0, 0.0), pitch=(0.0, 0.0), yaw=(0.0, 0.0) + ), + ) + term = _make_term(UniformPoseCommand, cfg, robot) + graph = _capture_compute(term) + + # both the articulation state and the command move, so a stale view of either shows up on replay + mutate_body_data(art_data) + fresh = make_pose_command_term(robot, seed=505) + stable.pose_command_b[:] = fresh.pose_command_b + term.pose_command_b[:] = fresh.pose_command_b + # cleared so the replay has to set the sticky flags again + stable._succeeded.zero_() + term._succeeded.zero_() + wp.synchronize() + wp.capture_launch(graph) + wp.synchronize() + stable._update_metrics() + + assert_close(term.pose_command_w, stable.pose_command_w) + assert_close(wp.to_torch(term.metrics["position_error"]), stable.metrics["position_error"]) + assert_close(wp.to_torch(term.metrics["orientation_error"]), stable.metrics["orientation_error"], atol=1e-4) + assert_equal(term._succeeded, stable._succeeded) + assert stable._succeeded.any() and not stable._succeeded.all(), "thresholds must split the environments" + + # the episode success rate is the mean sticky flag over the reset envs, which then start clean + selected = torch.arange(NUM_ENVS, device=DEVICE) % 2 == 0 + success_rate = term.reset(env_mask=wp.from_torch(selected))["success_rate"] + wp.synchronize() + assert success_rate.item() == pytest.approx(stable._succeeded[selected].float().mean().item()) + assert not term._succeeded[selected].any() + assert_equal(term._succeeded[~selected], stable._succeeded[~selected]) + + +def test_velocity_command_matches_stable_after_capture_and_mutation(): + """Heading control, standing envs and tracking-error sums follow the stable term on replay.""" + art_data = MockArticulationData() + robot = MockArticulation(art_data) + cfg = UniformVelocityCommandCfg( + asset_name="robot", + resampling_time_range=(1.0, 1.0), + heading_command=True, + heading_control_stiffness=0.7, + rel_standing_envs=0.3, + rel_heading_envs=0.6, + ranges=UniformVelocityCommandCfg.Ranges( + lin_vel_x=(-1.0, 1.0), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-0.8, 0.8), heading=(-math.pi, math.pi) + ), + ) + term = _make_term(UniformVelocityCommand, cfg, robot) + term.reset() # samples commands, heading targets and roles + graph = _capture_compute(term) + stable_error_sums = (term._error_xy_sum.numpy().copy(), term._error_yaw_sum.numpy().copy()) + + # yaw-only root rotations, so the stable term's heading is the yaw by construction + rng = np.random.RandomState(606) + yaw = rng.uniform(-math.pi, math.pi, NUM_ENVS).astype(np.float32) + quat = np.zeros((NUM_ENVS, 4), dtype=np.float32) + quat[:, 2] = np.sin(yaw / 2.0) + quat[:, 3] = np.cos(yaw / 2.0) + lin_vel_w = rng.randn(NUM_ENVS, 3).astype(np.float32) + ang_vel_w = rng.randn(NUM_ENVS, 3).astype(np.float32) + pose = np.concatenate([art_data.root_pos_w.warp.numpy(), quat], axis=1) + copy_np_to_wp(art_data.root_link_pose_w, pose) + copy_np_to_wp(art_data.root_quat_w, quat) + copy_np_to_wp(art_data.root_com_vel_w, np.concatenate([lin_vel_w, ang_vel_w], axis=1)) + copy_np_to_wp(art_data.root_lin_vel_b, quat_rotate_inv_np(quat, lin_vel_w)) + copy_np_to_wp(art_data.root_ang_vel_b, quat_rotate_inv_np(quat, ang_vel_w)) + art_data.heading_w = SimpleNamespace(torch=torch.tensor(yaw, device=DEVICE)) + + stable = object.__new__(stable_velocity_command.UniformVelocityCommand) + stable.cfg = cfg + stable.robot = robot + stable.vel_command_b = term.vel_command_b.clone() + stable.heading_target = wp.to_torch(term._heading_target_wp).clone() + stable.is_heading_env = wp.to_torch(term._is_heading_env_wp).clone() + stable.is_standing_env = wp.to_torch(term._is_standing_env_wp).clone() + stable._error_xy_sum = torch.tensor(stable_error_sums[0], device=DEVICE) + stable._error_yaw_sum = torch.tensor(stable_error_sums[1], device=DEVICE) + stable._step_count = torch.zeros(NUM_ENVS, device=DEVICE) + assert stable.is_heading_env.any() and stable.is_standing_env.any() and not stable.is_standing_env.all() + + wp.capture_launch(graph) + wp.synchronize() + stable._update_metrics() + stable._update_command() + + assert_close(term.vel_command_b, stable.vel_command_b) + assert_close(wp.to_torch(term._error_xy_sum), stable._error_xy_sum) + assert_close(wp.to_torch(term._error_yaw_sum), stable._error_yaw_sum) + + +def test_null_command_leaves_resampling_time_finite(): + """The null command never draws from its infinite resampling range.""" + term = _make_term(NullCommand, NullCommandCfg(), robot=None) + + assert term.reset() == {} + term.compute(_DT) + + assert np.isfinite(term.time_left_wp.numpy()).all() 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..e5c24c108cbb 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 @@ -47,15 +47,6 @@ # ============================================================================ -@pytest.fixture(autouse=True) -def _clear_caches(): - yield - for fn in [warp_obs.generated_commands]: - for attr in list(vars(fn)): - if attr.startswith("_"): - delattr(fn, attr) - - @pytest.fixture() def art_data(): return MockArticulationData(NUM_ENVS, NUM_JOINTS, DEVICE) @@ -313,7 +304,7 @@ def test_generated_commands(self, warp_env_bodies, stable_env_bodies, cmd_tensor actual = run_warp_obs(warp_obs.generated_commands, warp_env_bodies, (NUM_ENVS, CMD_DIM), command_name="vel") assert_close(actual, expected) - # Mutate the command tensor in-place: the cached zero-copy view must pick it up on replay. + # Mutate the command tensor in-place: the replay must read the updated command. def mutate(): cmd_tensor[:] = torch.randn_like(cmd_tensor) 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..7dc943b58fb2 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 @@ -56,7 +56,7 @@ @pytest.fixture(autouse=True) def _clear_caches(): yield - for fn in [warp_rew.track_lin_vel_xy_exp, warp_rew.track_ang_vel_z_exp, warp_rew.undesired_contacts]: + for fn in [warp_rew.undesired_contacts]: for attr in list(vars(fn)): if attr.startswith("_"): delattr(fn, attr) diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_commands.py b/source/isaaclab_experimental/test/envs/test_warp_env_commands.py new file mode 100644 index 000000000000..cfa8514d9186 --- /dev/null +++ b/source/isaaclab_experimental/test/envs/test_warp_env_commands.py @@ -0,0 +1,50 @@ +# 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 + +"""Commands of Warp environments built one after another in one process.""" + +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_experimental.envs.mdp as warp_mdp +import torch +import warp as wp +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 + +_TASK = "Isaac-Velocity-Flat-UnitreeGo2" +_NUM_ENVS = 8 + + +def test_each_velocity_env_reads_its_own_commands(): + """A second Warp velocity environment builds, and its command-reading terms read its own commands. + + The terms read the command through the environment's command manager on every call, so nothing on the term + functions refers to an earlier environment. + """ + for seed in (7, 8): + env_cfg, _ = resolve_task_config(_TASK, "", overrides=("physics=newton_mjwarp",)) + env_cfg.seed = seed + env_cfg.scene.num_envs = _NUM_ENVS + sim_utils.create_new_stage() + env = WarpFrontend.build_env(env_cfg, _TASK).unwrapped + try: + env.reset() + env.step(torch.zeros((_NUM_ENVS, env.action_space.shape[-1]), device=env.device)) + out = wp.zeros((_NUM_ENVS, 3), dtype=wp.float32, device=env.device) + + warp_mdp.generated_commands(env, out, command_name="base_velocity") + + assert torch.equal(wp.to_torch(out), env.command_manager.get_command("base_velocity")) + finally: + env.close() diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst new file mode 100644 index 000000000000..e3f997611b9b --- /dev/null +++ b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst @@ -0,0 +1,2 @@ +* Fixed the Warp velocity and reach reward terms reading the command buffer of the first environment + created in the process instead of their own environment's. diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/rewards.py index cdab28c55625..065005a49e4c 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/rewards.py @@ -6,7 +6,6 @@ """Warp-first reward terms for the reach task. All functions follow the ``func(env, out, **params) -> None`` signature. -Command tensors are cached as zero-copy warp views on first call. """ from __future__ import annotations @@ -51,12 +50,7 @@ def _position_command_error_kernel( def position_command_error(env: ManagerBasedRLEnv, out, command_name: str, asset_cfg: SceneEntityCfg) -> None: """Penalize tracking of the position error using L2-norm.""" - fn = position_command_error asset: Articulation = env.scene[asset_cfg.name] - if not hasattr(fn, "_cmd_wp") or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name wp.launch( kernel=_position_command_error_kernel, dim=env.num_envs, @@ -64,7 +58,7 @@ def position_command_error(env: ManagerBasedRLEnv, out, command_name: str, asset asset.data.root_pos_w.warp, asset.data.root_quat_w.warp, asset.data.body_pos_w.warp, - fn._cmd_wp, + env.command_manager.get_command_wp(command_name), asset_cfg.body_ids[0], out, ], @@ -103,12 +97,7 @@ def position_command_error_tanh( env: ManagerBasedRLEnv, out, std: float, command_name: str, asset_cfg: SceneEntityCfg ) -> None: """Reward tracking of the position using the tanh kernel.""" - fn = position_command_error_tanh asset: Articulation = env.scene[asset_cfg.name] - if not hasattr(fn, "_cmd_wp") or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name wp.launch( kernel=_position_command_error_tanh_kernel, dim=env.num_envs, @@ -116,7 +105,7 @@ def position_command_error_tanh( asset.data.root_pos_w.warp, asset.data.root_quat_w.warp, asset.data.body_pos_w.warp, - fn._cmd_wp, + env.command_manager.get_command_wp(command_name), asset_cfg.body_ids[0], 1.0 / std, out, @@ -156,15 +145,16 @@ def _orientation_command_error_kernel( def orientation_command_error(env: ManagerBasedRLEnv, out, command_name: str, asset_cfg: SceneEntityCfg) -> None: """Penalize tracking orientation error using shortest path.""" - fn = orientation_command_error asset: Articulation = env.scene[asset_cfg.name] - if not hasattr(fn, "_cmd_wp") or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name wp.launch( kernel=_orientation_command_error_kernel, dim=env.num_envs, - inputs=[asset.data.root_quat_w.warp, asset.data.body_quat_w.warp, fn._cmd_wp, asset_cfg.body_ids[0], out], + inputs=[ + asset.data.root_quat_w.warp, + asset.data.body_quat_w.warp, + env.command_manager.get_command_wp(command_name), + asset_cfg.body_ids[0], + out, + ], device=env.device, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/rewards.py index 4861b58e5e33..3cfb380a7939 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/rewards.py @@ -6,8 +6,6 @@ """Warp-first reward functions for the velocity locomotion environment. All functions follow the ``func(env, out, **params) -> None`` signature. -Cross-manager torch tensors (contact sensor, commands) are cached as zero-copy -warp views on first call via ``wp.from_torch``. """ from __future__ import annotations @@ -51,14 +49,7 @@ def _feet_air_time_kernel( def feet_air_time(env: ManagerBasedRLEnv, out, command_name: str, sensor_cfg: SceneEntityCfg, threshold: float) -> None: """Reward long steps taken by the feet using L2-kernel.""" - fn = feet_air_time contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] - # Cache command bridge (persistent pointer) - if not getattr(fn, "_is_warmed_up", False) or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name - fn._is_warmed_up = True # Newton contact sensor returns persistent wp.arrays — use directly, no wp.from_torch needed first_contact = contact_sensor.compute_first_contact(env.step_dt) wp.launch( @@ -68,7 +59,7 @@ def feet_air_time(env: ManagerBasedRLEnv, out, command_name: str, sensor_cfg: Sc contact_sensor.data.last_air_time.warp, first_contact, sensor_cfg.body_ids_wp, - fn._cmd_wp, + env.command_manager.get_command_wp(command_name), threshold, out, ], @@ -116,13 +107,7 @@ def _feet_air_time_positive_biped_kernel( def feet_air_time_positive_biped(env, out, command_name: str, threshold: float, sensor_cfg: SceneEntityCfg) -> None: """Reward long steps taken by the feet for bipeds.""" - fn = feet_air_time_positive_biped contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] - if not getattr(fn, "_is_warmed_up", False) or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name - fn._is_warmed_up = True wp.launch( kernel=_feet_air_time_positive_biped_kernel, dim=env.num_envs, @@ -130,7 +115,7 @@ def feet_air_time_positive_biped(env, out, command_name: str, threshold: float, contact_sensor.data.current_air_time.warp, contact_sensor.data.current_contact_time.warp, sensor_cfg.body_ids_wp, - fn._cmd_wp, + env.command_manager.get_command_wp(command_name), threshold, out, ], @@ -180,7 +165,7 @@ def feet_air_time_variance( data.last_air_time, data.last_contact_time, sensor_cfg.body_ids_wp, - wp.from_torch(env.command_manager.get_command(command_name)), + env.command_manager.get_command_wp(command_name), max_time, out, ], @@ -278,17 +263,17 @@ def track_lin_vel_xy_yaw_frame_exp( env, out, std: float, command_name: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> None: """Reward tracking of linear velocity commands (xy axes) in the gravity aligned robot frame.""" - fn = track_lin_vel_xy_yaw_frame_exp asset = env.scene[asset_cfg.name] - if not getattr(fn, "_is_warmed_up", False) or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name - fn._is_warmed_up = True wp.launch( kernel=_track_lin_vel_xy_yaw_frame_exp_kernel, dim=env.num_envs, - inputs=[asset.data.root_quat_w.warp, asset.data.root_lin_vel_w.warp, fn._cmd_wp, 1.0 / (std * std), out], + inputs=[ + asset.data.root_quat_w.warp, + asset.data.root_lin_vel_w.warp, + env.command_manager.get_command_wp(command_name), + 1.0 / (std * std), + out, + ], device=env.device, ) @@ -315,17 +300,16 @@ def track_ang_vel_z_world_exp( env, out, command_name: str, std: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> None: """Reward tracking of angular velocity commands (yaw) in world frame.""" - fn = track_ang_vel_z_world_exp asset = env.scene[asset_cfg.name] - if not getattr(fn, "_is_warmed_up", False) or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name - fn._is_warmed_up = True wp.launch( kernel=_track_ang_vel_z_world_exp_kernel, dim=env.num_envs, - inputs=[asset.data.root_ang_vel_w.warp, fn._cmd_wp, 1.0 / (std * std), out], + inputs=[ + asset.data.root_ang_vel_w.warp, + env.command_manager.get_command_wp(command_name), + 1.0 / (std * std), + out, + ], device=env.device, ) @@ -360,16 +344,16 @@ def stand_still_joint_deviation_l1( env, out, command_name: str, command_threshold: float = 0.06, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> None: """Penalize offsets from the default joint positions when the command is very small.""" - fn = stand_still_joint_deviation_l1 asset = env.scene[asset_cfg.name] - if not getattr(fn, "_is_warmed_up", False) or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name - fn._is_warmed_up = True wp.launch( kernel=_stand_still_joint_deviation_l1_kernel, dim=env.num_envs, - inputs=[asset.data.joint_pos.warp, asset.data.default_joint_pos.warp, fn._cmd_wp, command_threshold, out], + inputs=[ + asset.data.joint_pos.warp, + asset.data.default_joint_pos.warp, + env.command_manager.get_command_wp(command_name), + command_threshold, + out, + ], device=env.device, ) From 1454ebd92ad0df87bfbf5f0d5042ab8be7b69b9b Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 30 Sep 2026 04:09:40 -0700 Subject: [PATCH 20/80] Add a Warp height_scan twin for the rough velocity tasks The Warp observation manager infers the output size of a ray-caster term from the sensor's ray count, so height_scan needs no IO descriptor. The Newton ray caster refreshes through host-side timestamps and scene transforms, so height_scan is marked non-capturable and runs eagerly. terrain_out_of_bounds reads the terrain bounds from the environment on every call instead of caching the first environment's in function attributes. The rough velocity tasks keep the stable terrain curriculum and now adapt under --frontend warp. --- .../jichuanh-warp-frontend-managers.added.rst | 2 + .../envs/mdp/observations.py | 34 +++++++++++- .../managers/observation_manager.py | 12 +++- .../test/envs/mdp/test_capture_safety.py | 55 ++++++++++++++++--- .../test/envs/test_frontend_cfg_conversion.py | 9 +++ .../test/utils/test_mask_boundary_gate.py | 1 + .../jichuanh-warp-frontend-managers.fixed.rst | 2 + .../core/velocity/mdp/terminations.py | 33 +++++------ 8 files changed, 118 insertions(+), 30 deletions(-) diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst index 6ffd6c18ad87..eee8d6a68be3 100644 --- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst @@ -4,3 +4,5 @@ like the stable term. * Added :meth:`~isaaclab_experimental.managers.CommandManager.get_command_wp` to read a command as a Warp array. +* Added a Warp twin of the ``height_scan`` observation, so the Rough velocity tasks run with + ``--frontend warp``. It reads a ray caster refreshed on the host, so it runs eagerly. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py index d76f448dd21e..2a34fd9b3f62 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py @@ -12,7 +12,8 @@ where ``out`` is a pre-allocated Warp array with float32 dtype and shape ``(num_envs, D)``. Output dimension ``D`` is inferred from decorator metadata: ``axes`` for root-state terms, -``out_dim`` for body/command/action/time terms, or ``joint_ids`` count for joint terms. +``out_dim`` for body/command/action/time terms, the ray count for ray-caster terms, or ``joint_ids`` +count for joint terms. """ from __future__ import annotations @@ -36,10 +37,12 @@ record_shape, ) from isaaclab_experimental.managers import SceneEntityCfg +from isaaclab_experimental.utils.warp import WarpCapturable if TYPE_CHECKING: from isaaclab.assets import Articulation from isaaclab.envs import ManagerBasedEnv + from isaaclab.sensors import RayCaster # --------------------------------------------------------------------------- @@ -400,3 +403,32 @@ def body_incoming_wrench(env: ManagerBasedEnv, out, sensor_cfg: SceneEntityCfg) inputs=[sensor.data.force.warp, sensor.data.torque.warp, sensor_cfg.body_ids_wp, out], device=env.device, ) + + +@wp.kernel +def _height_scan_kernel( + pos_w: wp.array(dtype=wp.vec3f), + ray_hits_w: wp.array(dtype=wp.vec3f, ndim=2), + offset: float, + out: wp.array(dtype=wp.float32, ndim=2), +): + """Height of the sensor frame above each ray hit [m], minus the offset. Launch with dim=(num_envs, num_rays).""" + env_id, ray_id = wp.tid() + out[env_id, ray_id] = pos_w[env_id][2] - ray_hits_w[env_id, ray_id][2] - offset + + +@WarpCapturable(False, reason="Ray-caster reads refresh the sensor through host-side timestamps and scene transforms.") +def height_scan(env: ManagerBasedEnv, out, sensor_cfg: SceneEntityCfg, offset: float = 0.5) -> None: + """Height scan from the given sensor w.r.t. the sensor's frame [m]. Writes into ``out``. + + Warp-first override of :func:`isaaclab.envs.mdp.observations.height_scan`. The provided offset + (Defaults to 0.5) is subtracted from the returned values. Missed rays keep their infinite hit + height, as in the stable term. + """ + sensor: RayCaster = env.scene.sensors[sensor_cfg.name] + wp.launch( + kernel=_height_scan_kernel, + dim=(env.num_envs, sensor.num_rays), + inputs=[sensor.data.pos_w.warp, sensor.data.ray_hits_w.warp, offset, out], + device=env.device, + ) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py index 80d89688c413..2407d7633b3e 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py @@ -922,7 +922,8 @@ def _infer_term_dim_scalar(self, term_cfg: ObservationTermCfg) -> int: 1. ``out_dim`` on the function's ``@generic_io_descriptor_warp`` decorator. 2. ``axes`` on the decorator (e.g. ``axes=["X","Y","Z"]`` → dim 3). 3. Explicit ``term_dim`` / ``out_dim`` / ``obs_dim`` in ``term_cfg.params`` (legacy). - 4. ``asset_cfg.joint_ids`` count (joint-based observations). + 4. Number of rays of the ray-caster sensor in ``sensor_cfg`` (one value per ray). + 5. ``asset_cfg.joint_ids`` count (joint-based observations). """ # --- 1-2. Decorator metadata (preferred) --- func = term_cfg.func @@ -946,7 +947,14 @@ def _infer_term_dim_scalar(self, term_cfg: ObservationTermCfg) -> int: if k in term_cfg.params: return int(term_cfg.params[k]) - # --- 3. Joint-based fallback via asset_cfg --- + # --- 4. Ray-based fallback via sensor_cfg --- + sensor_cfg = term_cfg.params.get("sensor_cfg") + if sensor_cfg is not None: + num_rays = getattr(self._env.scene.sensors[sensor_cfg.name], "num_rays", None) + if num_rays is not None: + return int(num_rays) + + # --- 5. Joint-based fallback via asset_cfg --- asset_cfg = term_cfg.params.get("asset_cfg") if asset_cfg is None: raise ValueError( 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 ad77fbe3db51..5de0308c0138 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py @@ -46,6 +46,8 @@ import isaaclab_tasks_experimental.core as tasks_experimental_core import isaaclab_tasks_experimental.core.locomotion.mdp.rewards as warp_loco_rew import isaaclab_tasks_experimental.core.velocity.mdp.rewards as warp_velocity_rew +import isaaclab_tasks_experimental.core.velocity.mdp.terminations as warp_velocity_term +from isaaclab_experimental.utils.warp import is_warp_capturable from parity_helpers import ( DEVICE, NUM_BODIES, @@ -62,6 +64,7 @@ copy_np_to_wp, make_pose_command_term, mutate_body_data, + mutate_root_state, proxy_array, ) @@ -70,6 +73,7 @@ from isaaclab.managers.manager_term_cfg import RewardTermCfg, TerminationTermCfg import isaaclab_tasks.core.velocity.mdp.rewards as stable_velocity_rew +import isaaclab_tasks.core.velocity.mdp.terminations as stable_velocity_term @dataclasses.dataclass(frozen=True) @@ -291,6 +295,33 @@ def mutate(): ) +def _build_terrain_out_of_bounds() -> CaptureCase: + """A termination term reading the terrain size and the root positions. + + The term first runs in an environment on a plane, then in one on generated terrain: bounds held + across calls (rather than read from the environment) keep the plane's "never out of bounds". + """ + + def make_env(terrain_type: str) -> SimpleNamespace: + art_data = MockArticulationData() + terrain_generator = SimpleNamespace(size=(2.0, 3.0), num_rows=2, num_cols=2, border_width=0.5) + scene = MockScene({"robot": MockArticulation(art_data)}, env_origins=None) + scene.cfg = SimpleNamespace(terrain=SimpleNamespace(terrain_type=terrain_type)) + scene.terrain = SimpleNamespace(cfg=SimpleNamespace(terrain_generator=terrain_generator)) + return SimpleNamespace(scene=scene, num_envs=NUM_ENVS, device=DEVICE, art_data=art_data) + + warp_velocity_term.terrain_out_of_bounds(make_env("plane"), wp.zeros(NUM_ENVS, dtype=wp.bool, device=DEVICE)) + env = make_env("generator") + params = {"distance_buffer": 0.5} + + def mutate() -> None: + mutate_root_state(np.random.RandomState(707), env.art_data) + + return CaptureCase( + warp_velocity_term.terrain_out_of_bounds, stable_velocity_term.terrain_out_of_bounds, env, env, params, mutate + ) + + _SHARED_REW = "isaaclab_experimental.envs.mdp.rewards" _SHARED_TERM = "isaaclab_experimental.envs.mdp.terminations" _LOCO_REW = "isaaclab_tasks_experimental.core.locomotion.mdp.rewards" @@ -308,6 +339,12 @@ def mutate(): "reward", _build_feet_air_time_variance, ), + CaptureSpec( + "isaaclab_tasks_experimental.core.velocity.mdp.terminations", + "terrain_out_of_bounds", + "termination", + _build_terrain_out_of_bounds, + ), ] @@ -425,8 +462,7 @@ def _ids(module: str, *names: str) -> list[str]: "stand_still_joint_deviation_l1", "track_ang_vel_z_world_exp", "track_lin_vel_xy_yaw_frame_exp", - ) - + _ids(f"{_TE}.velocity.mdp.terminations", "terrain_out_of_bounds"), + ), _MIRROR, ), # command terms are classes driven through compute() and reset(), which this harness does not model @@ -467,32 +503,35 @@ def _warp_mdp_modules() -> list: return modules -def _discover_warp_mdp_terms() -> set[str]: - """Return every public warp MDP term as a ``":"`` identity. +def _discover_warp_mdp_terms() -> dict[str, object]: + """Return every public warp MDP term, keyed by its ``":"`` identity. Qualified rather than bare: the same term name legitimately appears in more than one task mirror (``survival_success_rate`` is twinned by both cartpole and locomotion), and keying by name alone would let a spec for one of them mark the other as declared. """ - terms: set[str] = set() + terms: dict[str, object] = {} for module in _warp_mdp_modules(): for name, obj in vars(module).items(): if name.startswith("_") or not (inspect.isfunction(obj) or inspect.isclass(obj)): continue if getattr(obj, "__module__", "") == module.__name__: - terms.add(f"{module.__name__}:{name}") + terms[f"{module.__name__}:{name}"] = obj return terms def test_every_warp_mdp_term_is_declared(): """Every warp MDP term is either exercised here or listed as an unaudited pre-existing term. + A term annotated ``@WarpCapturable(False)`` is declared by that annotation: it never runs captured. The declarations must also stay truthful: a term moved into the harness loses its unaudited row, and rows for terms that no longer exist are dropped, so the backlog cannot rot. """ specified = {spec.qualified for spec in CAPTURE_SPECS} - discovered = _discover_warp_mdp_terms() - undeclared = sorted(discovered - (specified | set(CAPTURE_UNAUDITED))) + terms = _discover_warp_mdp_terms() + discovered = set(terms) + eager = {name for name, term in terms.items() if not is_warp_capturable(term)} + undeclared = sorted(discovered - (specified | set(CAPTURE_UNAUDITED) | eager)) assert not undeclared, ( "warp MDP terms with no capture declaration: " 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..2547f8a9c3be 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py +++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py @@ -58,10 +58,19 @@ "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", "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", } ) 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..02fb21c776c3 100644 --- a/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py +++ b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py @@ -138,6 +138,7 @@ def test_mask_first_code_has_no_unsanctioned_host_syncs(): ("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"), + ("isaaclab_experimental/envs/mdp/observations.py", "height_scan"), } NON_CAPTURABLE_SCAN_ROOTS = [ diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst index e3f997611b9b..30c9c6c3b491 100644 --- a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst +++ b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst @@ -1,2 +1,4 @@ * Fixed the Warp velocity and reach reward terms reading the command buffer of the first environment created in the process instead of their own environment's. +* Fixed the Warp ``terrain_out_of_bounds`` termination using the terrain bounds of the first environment + created in the process. diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/terminations.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/terminations.py index 53f701700097..7c4291092ccb 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/terminations.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/terminations.py @@ -36,32 +36,27 @@ def _terrain_out_of_bounds_kernel( def terrain_out_of_bounds( env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), distance_buffer: float = 3.0 ) -> None: - """Terminate when the actor moves too close to the edge of the terrain.""" - fn = terrain_out_of_bounds - if not getattr(fn, "_is_warmed_up", False): - terrain_type = env.scene.cfg.terrain.terrain_type - if terrain_type == "plane": - fn._is_plane = True - elif terrain_type == "generator": - fn._is_plane = False - terrain_gen_cfg = env.scene.terrain.cfg.terrain_generator - grid_width, grid_length = terrain_gen_cfg.size - n_rows, n_cols = terrain_gen_cfg.num_rows, terrain_gen_cfg.num_cols - border_width = terrain_gen_cfg.border_width - fn._half_width = 0.5 * (n_rows * grid_width + 2 * border_width) - fn._half_height = 0.5 * (n_cols * grid_length + 2 * border_width) - else: - raise ValueError("Received unsupported terrain type, must be either 'plane' or 'generator'.") - fn._is_warmed_up = True + """Terminate when the actor moves too close to the edge of the terrain. - if fn._is_plane: + The terrain bounds are read from the environment on every call, so each environment uses its own terrain. + """ + terrain_type = env.scene.cfg.terrain.terrain_type + if terrain_type == "plane": + # we have infinite terrain because it is a plane out.zero_() return + if terrain_type != "generator": + raise ValueError("Received unsupported terrain type, must be either 'plane' or 'generator'.") + terrain_gen_cfg = env.scene.terrain.cfg.terrain_generator + grid_width, grid_length = terrain_gen_cfg.size + border_width = terrain_gen_cfg.border_width + half_width = 0.5 * (terrain_gen_cfg.num_rows * grid_width + 2 * border_width) + half_height = 0.5 * (terrain_gen_cfg.num_cols * grid_length + 2 * border_width) asset: Articulation = env.scene[asset_cfg.name] wp.launch( kernel=_terrain_out_of_bounds_kernel, dim=env.num_envs, - inputs=[asset.data.root_pos_w.warp, fn._half_width, fn._half_height, distance_buffer, out], + inputs=[asset.data.root_pos_w.warp, half_width, half_height, distance_buffer, out], device=env.device, ) From a26abb310aa37d8c9e69b46324546524145ed378 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 30 Sep 2026 04:10:19 -0700 Subject: [PATCH 21/80] Keep external Warp terms on the Warp frontend A term defined outside the mirrored packages has no twin. The frontend now uses it unchanged when it is a Warp term: a class term that subclasses the Warp ManagerTermBase, or a function term that declares its capture behavior with WarpCapturable. Any other such term is reported as not a Warp term instead of being replaced by an unrelated built-in twin of the same name. --- .../jichuanh-warp-frontend-managers.added.rst | 3 + .../jichuanh-warp-frontend-managers.fixed.rst | 2 + .../isaaclab_experimental/envs/frontend.py | 58 +++++++-- .../test/envs/test_frontend_external_terms.py | 116 ++++++++++++++++++ 4 files changed, 167 insertions(+), 12 deletions(-) create mode 100644 source/isaaclab_experimental/test/envs/test_frontend_external_terms.py diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst index eee8d6a68be3..005d258fef1f 100644 --- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst @@ -6,3 +6,6 @@ array. * Added a Warp twin of the ``height_scan`` observation, so the Rough velocity tasks run with ``--frontend warp``. It reads a ray caster refreshed on the host, so it runs eagerly. +* Added support for Warp terms defined outside Isaac Lab on the Warp frontend: class terms that subclass + :class:`~isaaclab_experimental.managers.ManagerTermBase` and function terms decorated with + :class:`~isaaclab_experimental.utils.warp.WarpCapturable` are used as configured. diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst index e6e1100aff3e..21bc3b032cc6 100644 --- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst @@ -1,2 +1,4 @@ * Fixed a second Warp velocity environment in the same process failing to build with a ``RecursionError``, and Warp observation and reward terms reading the command buffer of the first environment created in the process. +* Fixed the Warp frontend replacing a Torch term defined outside Isaac Lab with an unrelated built-in Warp + term of the same name. Such a term is now reported as not being a Warp term. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index 1d3eb394da88..565caafbef69 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -302,7 +302,7 @@ def _swap_mdp(cls, cfg: Any, label: str) -> None: searched: set[str] = set() # module names, for the error message swapped = 0 - missing: list[tuple[str, str, str]] = [] # (location, attr, symbol) + missing: list[str] = [] for path, term in cls._walk_terms(cfg): if not path or path[0] not in cls.WARP_MANAGED_GROUPS: continue # term runs on a stable manager; leave it stable @@ -317,15 +317,21 @@ def _swap_mdp(cls, cfg: Any, label: str) -> None: searched.update(m.__name__ for m in module_cache[origin]) twin = cls._resolve_warp_twin(stable.__name__, module_cache[origin]) if twin is None: - missing.append((location, attr, stable.__name__)) # collect every miss; report once below + # collect every miss; report once below + if cls._mirror_module(origin) is None: + missing.append(f"{location}.{attr}: {stable.__name__!r} from {origin!r} is not a warp term") + else: + missing.append(f"{location}.{attr}: no warp twin for {stable.__name__!r}") continue setattr(term, attr, twin) swapped += 1 if missing: - lines = "\n ".join(f"{loc}.{attr}: no warp twin for {sym!r}" for loc, attr, sym in missing) + lines = "\n ".join(missing) raise FrontendIncompatibleError( - f"warp env {label!r}: missing warp MDP twins (searched {sorted(searched)}):\n {lines}" + f"warp env {label!r}: missing warp MDP twins (searched {sorted(searched)}):\n {lines}\n" + "Terms defined outside the mirrored packages are used as-is when they are warp terms: subclass" + " isaaclab_experimental.managers.ManagerTermBase, or declare a function term with @WarpCapturable." ) logger.info("frontend.warp: swapped %d MDP symbol(s) to warp twins", swapped) @@ -631,16 +637,20 @@ def _twin_modules(cls, symbol_module: str, cfg_route_modules: list[ModuleType]) from another task family and, for core-defined symbols, the shared :mod:`isaaclab_experimental.envs.mdp` twins (the mirror of :mod:`isaaclab.envs.mdp`). + + A symbol defined outside the mirrored packages has no twin, so none are consulted: a + same-named warp term would be unrelated to it. """ + mirrored = cls._mirror_module(symbol_module) + if mirrored is None: + return [] modules = list(cfg_route_modules) # Fallback: the mirror of the package that defines the symbol. - mirrored = cls._mirror_module(symbol_module) - if mirrored is not None: - target = cls._nearest_mdp_module(mirrored) - if target is not None: - module = cls._import_twin_module(target) - if module not in modules: - modules.append(module) + target = cls._nearest_mdp_module(mirrored) + if target is not None: + module = cls._import_twin_module(target) + if module not in modules: + modules.append(module) return modules @classmethod @@ -657,14 +667,38 @@ def _resolve_warp_twin(cls, name: str, modules: Iterable[ModuleType]) -> Any | N @classmethod def _is_swap_candidate(cls, value: Any) -> bool: - """Heuristic: callable or class whose origin is *not already* warp.""" + """Whether a term callable or class must be swapped for its warp twin. + + Terms from the warp packages are already warp (idempotent). A term defined outside the mirrored + packages, e.g. in an external project, is kept when it declares itself a warp term (see + :meth:`_is_native_warp_term`). + """ if not callable(value): return False origin = getattr(value, "__module__", "") or "" if origin.startswith(cls.WARP_ROOT_PREFIXES): return False # already a warp twin (idempotent) + if cls._mirror_module(origin) is None and cls._is_native_warp_term(value): + return False # a warp term from another package return True + @staticmethod + def _is_native_warp_term(value: Any) -> bool: + """Whether a term defined outside the warp packages is a warp term. + + A class term is one when it subclasses the warp :class:`~isaaclab_experimental.managers.ManagerTermBase` + (warp action and command terms included). A function term is one when it declares its capture + behavior with :class:`~isaaclab_experimental.utils.warp.WarpCapturable`. + """ + from isaaclab.utils import string_to_callable + + from isaaclab_experimental.managers import ManagerTermBase + + term = string_to_callable(value) if isinstance(value, str) else value + if isinstance(term, type): + return issubclass(term, ManagerTermBase) + return any(hasattr(f, "_warp_capturable") for f in (term, getattr(term, "__wrapped__", None)) if f is not None) + @staticmethod def _walk_terms(node: Any, path: tuple[str, ...] = ()) -> Iterator[tuple[tuple[str, ...], Any]]: """Yield ``(path, term)`` for every MDP term cfg in the cfg tree. diff --git a/source/isaaclab_experimental/test/envs/test_frontend_external_terms.py b/source/isaaclab_experimental/test/envs/test_frontend_external_terms.py new file mode 100644 index 000000000000..6d64721f8745 --- /dev/null +++ b/source/isaaclab_experimental/test/envs/test_frontend_external_terms.py @@ -0,0 +1,116 @@ +# 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 terms from an external project run on the Warp frontend unchanged. + +The project lives outside the packages the frontend mirrors, so its terms have no twins. Its Warp terms +must be used as-is, even when one shares its name with a built-in term, while a Torch term from the same +project must still be rejected instead of being replaced by a same-named built-in twin. +""" + +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 importlib +import textwrap + +import pytest +import torch +from isaaclab_experimental.envs.frontend import WarpFrontend + +import isaaclab.sim as sim_utils +from isaaclab.managers import RewardTermCfg +from isaaclab.utils import configclass + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.utils import resolve_task_config + +_NUM_ENVS = 4 + +_EXTERNAL_MDP = textwrap.dedent(''' + import torch + import warp as wp + + from isaaclab_experimental.managers import ManagerTermBase + from isaaclab_experimental.utils.warp import WarpCapturable + + + @wp.kernel + def _fill(value: float, out: wp.array(dtype=wp.float32)): + out[wp.tid()] = value + + + @WarpCapturable(True) + def joint_vel_l2(env, out) -> None: + """A Warp reward sharing its name with the built-in ``joint_vel_l2``.""" + wp.launch(_fill, dim=env.num_envs, inputs=[2.0, out], device=env.device) + + + class constant_bonus(ManagerTermBase): + """A Warp class reward.""" + + def __call__(self, env, out, value: float) -> None: + wp.launch(_fill, dim=env.num_envs, inputs=[value, out], device=env.device) + + + def joint_vel_l1(env) -> torch.Tensor: + """A Torch reward sharing its name with the built-in ``joint_vel_l1``.""" + return torch.zeros(env.num_envs, device=env.device) +''') + + +@pytest.fixture +def external_mdp(tmp_path, monkeypatch): + """An external project package on the import path.""" + package = tmp_path / "my_warp_project" + package.mkdir() + (package / "__init__.py").write_text("") + (package / "mdp.py").write_text(_EXTERNAL_MDP) + monkeypatch.syspath_prepend(str(tmp_path)) + return importlib.import_module("my_warp_project.mdp") + + +def _cartpole_cfg(rewards): + env_cfg, _ = resolve_task_config("Isaac-Cartpole", "", overrides=("physics=newton_mjwarp",)) + env_cfg.seed = 7 + env_cfg.scene.num_envs = _NUM_ENVS + env_cfg.rewards = rewards + return env_cfg + + +def test_external_warp_terms_run_as_declared(external_mdp): + @configclass + class RewardsCfg: + function_term = RewardTermCfg(func=external_mdp.joint_vel_l2, weight=1.0) + class_term = RewardTermCfg(func=external_mdp.constant_bonus, weight=0.5, params={"value": 3.0}) + + env_cfg = _cartpole_cfg(RewardsCfg()) + sim_utils.create_new_stage() + env = WarpFrontend.build_env(env_cfg, "Isaac-Cartpole").unwrapped + try: + env.reset() + actions = torch.zeros((_NUM_ENVS, env.action_space.shape[-1]), device=env.device) + for _ in range(3): + _, reward, _, _, _ = env.step(actions) + + assert env.reward_manager.get_term_cfg("function_term").func is external_mdp.joint_vel_l2 + expected = (1.0 * 2.0 + 0.5 * 3.0) * env.step_dt + assert torch.allclose(reward, torch.full_like(reward, expected)) + finally: + env.close() + + +def test_external_torch_term_is_rejected(external_mdp): + @configclass + class RewardsCfg: + torch_term = RewardTermCfg(func=external_mdp.joint_vel_l1, weight=1.0) + + reason = WarpFrontend.check_compatibility(_cartpole_cfg(RewardsCfg())) + + assert reason is not None and "'joint_vel_l1' from 'my_warp_project.mdp' is not a warp term" in reason From 565740d75dde42b72af18de9909914aa74e10b78 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 30 Sep 2026 04:24:50 -0700 Subject: [PATCH 22/80] Run base_height_l2 eagerly when it reads a height scanner Reading a Newton ray caster refreshes it on the host and launches its own query graph, which fails inside a recorded reward stage: Go2 Rough stopped at its first captured step with CUDA error 900. base_height_l2 now declares itself capturable only without a sensor, so the reward stage records its other terms and runs this one eagerly; Go2 Flat keeps it recorded. The height-scanner case of the capture harness becomes an eager parity test, since that configuration is never recorded. --- .../jichuanh-warp-frontend-managers.fixed.rst | 2 ++ .../isaaclab_experimental/envs/mdp/rewards.py | 5 +++++ .../test/envs/mdp/test_capture_safety.py | 11 ++++++++++- .../test/utils/test_mask_boundary_gate.py | 1 + 4 files changed, 18 insertions(+), 1 deletion(-) diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst index 21bc3b032cc6..747ec18193a2 100644 --- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst @@ -1,4 +1,6 @@ * Fixed a second Warp velocity environment in the same process failing to build with a ``RecursionError``, and Warp observation and reward terms reading the command buffer of the first environment created in the process. +* Fixed the Warp ``base_height_l2`` reward failing inside a recorded reward stage when it reads a + ray-caster height scanner. It now runs eagerly in that configuration. * Fixed the Warp frontend replacing a Torch term defined outside Isaac Lab with an unrelated built-in Warp term of the same name. Such a term is now reported as not being a Warp term. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py index fb5a92637ac9..60852cc5f48b 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py @@ -25,6 +25,7 @@ ) from isaaclab_experimental.managers import ManagerTermBase, SceneEntityCfg +from isaaclab_experimental.utils.warp import WarpCapturable if TYPE_CHECKING: from isaaclab.assets import Articulation @@ -226,6 +227,10 @@ def _base_height_l2_kernel( out[i] = error * error +@WarpCapturable( + lambda params: params.get("sensor_cfg") is None, + reason="Ray-caster reads refresh the sensor through host-side timestamps and scene transforms.", +) def base_height_l2( env: ManagerBasedRLEnv, out: wp.array(dtype=wp.float32), 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 5de0308c0138..538ab7259733 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py @@ -332,7 +332,6 @@ def mutate() -> None: CaptureSpec(_LOCO_REW, "terminated_penalty", "reward", _build_terminated_penalty), CaptureSpec(_LOCO_REW, "survival_success_rate", "reward", _build_survival_success_rate, expect_nonzero=False), CaptureSpec(_SHARED_REW, "base_height_l2", "reward", partial(_build_base_height_l2, terrain=False)), - CaptureSpec(_SHARED_REW, "base_height_l2", "reward", partial(_build_base_height_l2, terrain=True)), CaptureSpec( "isaaclab_tasks_experimental.core.velocity.mdp.rewards", "feet_air_time_variance", @@ -576,6 +575,16 @@ def test_term_is_capture_safe(spec: CaptureSpec): assert expected.any(), "mutated inputs produced a degenerate expectation; the replay proves little" +def test_base_height_l2_with_height_scanner_matches_stable(): + """With a height scanner the term runs eagerly (reading the sensor refreshes it on the host).""" + case = _build_base_height_l2(terrain=True) + out = wp.zeros((NUM_ENVS,), dtype=wp.float32, device=DEVICE) + + case.warp_fn(case.warp_env, out, **case.params) + + assert_close(wp.to_torch(out), case.stable_fn(case.stable_env, **case.params)) + + def test_feet_air_time_variance_requires_two_feet(): case = _build_feet_air_time_variance(body_ids=[0]) with pytest.raises(RuntimeError, match="at least two bodies"): 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 02fb21c776c3..516b56b105dc 100644 --- a/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py +++ b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py @@ -139,6 +139,7 @@ def test_mask_first_code_has_no_unsanctioned_host_syncs(): ("isaaclab_experimental/envs/mdp/events.py", "randomize_rigid_body_mass"), ("isaaclab_experimental/envs/mdp/events.py", "randomize_rigid_body_material"), ("isaaclab_experimental/envs/mdp/observations.py", "height_scan"), + ("isaaclab_experimental/envs/mdp/rewards.py", "base_height_l2"), } NON_CAPTURABLE_SCAN_ROOTS = [ From 9443fe0100c8be95366da90c9037b9f221a64682 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 30 Sep 2026 04:50:57 -0700 Subject: [PATCH 23/80] Record Warp stages again when a term configuration changes A recorded stage keeps the scalars a term read while it recorded, so a configuration change reaches a captured stage only through a setter that drops the recorded graphs. The Warp termination and command managers gain set_term_cfg and get_term_cfg that do so, matching the reward and event managers. The stable modify_term_cfg curriculum writes into a manager's term configuration directly. After writing a value it now hands the term configuration back to the manager's set_term_cfg when the manager has one, which is a no-op for the stable managers and records the Warp stages again. Capture-then-mutate tests cover a reward weight raised from zero, a replaced termination threshold, and a command range changed by modify_term_cfg. --- ...ichuanh-warp-frontend-managers.changed.rst | 3 + .../jichuanh-warp-frontend-managers.skip | 0 .../isaaclab/isaaclab/envs/mdp/curriculums.py | 20 +++- .../jichuanh-warp-frontend-managers.added.rst | 5 + .../managers/command_manager.py | 38 +++++++ .../managers/termination_manager.py | 43 ++++++++ .../test/managers/test_command_manager.py | 100 ++++++++++++++++++ .../test/managers/test_reward_manager.py | 27 +++++ .../test/managers/test_termination_manager.py | 38 +++++++ 9 files changed, 273 insertions(+), 1 deletion(-) create mode 100644 source/isaaclab/changelog.d/jichuanh-warp-frontend-managers.changed.rst delete mode 100644 source/isaaclab/changelog.d/jichuanh-warp-frontend-managers.skip create mode 100644 source/isaaclab_experimental/test/managers/test_command_manager.py diff --git a/source/isaaclab/changelog.d/jichuanh-warp-frontend-managers.changed.rst b/source/isaaclab/changelog.d/jichuanh-warp-frontend-managers.changed.rst new file mode 100644 index 000000000000..486f2e3b24bd --- /dev/null +++ b/source/isaaclab/changelog.d/jichuanh-warp-frontend-managers.changed.rst @@ -0,0 +1,3 @@ +* 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. diff --git a/source/isaaclab/changelog.d/jichuanh-warp-frontend-managers.skip b/source/isaaclab/changelog.d/jichuanh-warp-frontend-managers.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab/isaaclab/envs/mdp/curriculums.py b/source/isaaclab/isaaclab/envs/mdp/curriculums.py index 65d2e858a8f4..5568ec2ab9dd 100644 --- a/source/isaaclab/isaaclab/envs/mdp/curriculums.py +++ b/source/isaaclab/isaaclab/envs/mdp/curriculums.py @@ -263,7 +263,9 @@ class modify_term_cfg(modify_env_param): The same applies to other managers, such as "observations", "commands", "rewards", and "terminations". Internally, it replaces the first occurrence of "s." in the address with "_manager.cfg.", - thus transforming the simplified address into a full manager path. + thus transforming the simplified address into a full manager path. After writing a new value, it hands the + modified term configuration back to the manager's ``set_term_cfg`` when the manager has one, so managers that + keep values read from the configuration refresh them. Usage: .. code-block:: python @@ -287,9 +289,25 @@ def override_value(env, env_ids, data, value, num_steps): def __init__(self, cfg, env): # initialize the parent super().__init__(cfg, env) + # the simplified address starts with the manager's configuration group and the term name + group, _, term_path = self._address.partition(".") + self._manager_name = f"{group[:-1]}_manager" + self._term_name = re.split(r"[.\[]", term_path, maxsplit=1)[0] # overwrite the simplified address with the full manager path self._address = self._address.replace("s.", "_manager.cfg.", 1) + def _process_accessors(self, root: ManagerBasedRLEnv, path: str) -> tuple[callable, callable]: + get_value, set_value = super()._process_accessors(root, path) + manager = getattr(root, self._manager_name) + if not hasattr(manager, "set_term_cfg"): + return get_value, set_value + + def set_value_and_term_cfg(val): + set_value(val) + manager.set_term_cfg(self._term_name, manager.get_term_cfg(self._term_name)) + + return get_value, set_value_and_term_cfg + class DifficultyScheduler(ManagerTermBase): """Adaptive difficulty scheduler for curriculum learning. diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst index 005d258fef1f..3f2183dfaf59 100644 --- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst @@ -9,3 +9,8 @@ * Added support for Warp terms defined outside Isaac Lab on the Warp frontend: class terms that subclass :class:`~isaaclab_experimental.managers.ManagerTermBase` and function terms decorated with :class:`~isaaclab_experimental.utils.warp.WarpCapturable` are used as configured. +* Added ``set_term_cfg`` and ``get_term_cfg`` to the Warp + :class:`~isaaclab_experimental.managers.TerminationManager` and + :class:`~isaaclab_experimental.managers.CommandManager`. Setting a term configuration records the manager's + stages again, so a curriculum that changes a termination parameter or a command range through + ``modify_term_cfg`` takes effect under CUDA graph capture. diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py index 0d5f350b7116..0996895f6e18 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py @@ -568,6 +568,44 @@ def get_command_wp(self, name: str) -> wp.array: return wp.from_torch(command) return command + def set_term_cfg(self, name: str, cfg: CommandTermCfg): + """Sets the configuration of the specified command term. + + A command term reads its configuration, such as its sampling ranges, when a recorded command stage records, + so the command stages record again on their next call. + + Args: + name: The name of the command term. + cfg: The configuration for the command term. + + Raises: + ValueError: If the term name is not found. + """ + if name not in self._terms: + raise ValueError(f"Command term '{name}' not found.") + self._terms[name].cfg = cfg + if isinstance(self.cfg, dict): + self.cfg[name] = cfg + else: + setattr(self.cfg, name, cfg) + self._clear_stage_steps() + + def get_term_cfg(self, name: str) -> CommandTermCfg: + """Gets the configuration of the specified command term. + + Args: + name: The name of the command term. + + Returns: + The configuration of the command term. + + Raises: + ValueError: If the term name is not found. + """ + if name not in self._terms: + raise ValueError(f"Command term '{name}' not found.") + return self._terms[name].cfg + def get_term(self, name: str) -> CommandTerm: """Returns the command term with the specified name. diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py index 32174f3eb33c..1c3b61c473c1 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py @@ -353,6 +353,49 @@ def _finalize_step(self) -> torch.Tensor: ) return self._dones_tensor_view + """ + Operations - Term settings. + """ + + def set_term_cfg(self, term_name: str, cfg: TerminationTermCfg): + """Sets the configuration of the specified term into the manager. + + Recorded termination stages read the term's configuration when they record, so they record again on their + next call. + + Args: + term_name: The name of the termination term. + cfg: The configuration for the termination term. + + Raises: + ValueError: If the term name is not found. + """ + if term_name not in self._term_names: + raise ValueError(f"Termination term '{term_name}' not found.") + term_idx = self._term_name_to_term_idx[term_name] + # the term keeps writing into its column of the done buffer + cfg.out = self._term_out_views_wp[term_idx] + self._term_cfgs[term_idx] = cfg + self._class_term_cfgs = [term_cfg for term_cfg in self._term_cfgs if isinstance(term_cfg.func, ManagerTermBase)] + wp.to_torch(self._term_is_time_out_wp)[term_idx] = bool(cfg.time_out) + self._clear_stage_steps() + + def get_term_cfg(self, term_name: str) -> TerminationTermCfg: + """Gets the configuration for the specified term. + + Args: + term_name: The name of the termination term. + + Returns: + The configuration of the termination term. + + Raises: + ValueError: If the term name is not found. + """ + if term_name not in self._term_names: + raise ValueError(f"Termination term '{term_name}' not found.") + return self._term_cfgs[self._term_name_to_term_idx[term_name]] + def get_term(self, name: str) -> torch.Tensor: """Returns the termination term value at current step with the specified name. diff --git a/source/isaaclab_experimental/test/managers/test_command_manager.py b/source/isaaclab_experimental/test/managers/test_command_manager.py new file mode 100644 index 000000000000..2de666a71497 --- /dev/null +++ b/source/isaaclab_experimental/test/managers/test_command_manager.py @@ -0,0 +1,100 @@ +# 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 command manager.""" + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +import warp as wp +from isaaclab_experimental.envs.mdp.commands import UniformVelocityCommand +from isaaclab_experimental.managers import CommandManager +from isaaclab_experimental.utils import WarpGraphCache +from isaaclab_experimental.utils.warp import resolve_1d_mask + +import isaaclab.envs.mdp as mdp +from isaaclab.managers import CurriculumTermCfg +from isaaclab.utils import configclass + +wp.init() +pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required") + +DEVICE = "cuda:0" +NUM_ENVS = 4 +DT = 0.02 + + +@configclass +class CommandsCfg: + base_velocity = mdp.UniformVelocityCommandCfg( + class_type=UniformVelocityCommand, + asset_name="robot", + # shorter than a step, so every step resamples + resampling_time_range=(DT / 2, DT / 2), + heading_command=False, + rel_standing_envs=0.0, + ranges=mdp.UniformVelocityCommandCfg.Ranges(lin_vel_x=(0.0, 0.0), lin_vel_y=(0.0, 0.0), ang_vel_z=(0.0, 0.0)), + ) + + +def _override(env, env_ids, data, value): + return value + + +def _env(graph_cache: WarpGraphCache) -> SimpleNamespace: + """The environment surface the command manager and the velocity command read.""" + pose = np.zeros((NUM_ENVS, 7), dtype=np.float32) + pose[:, 6] = 1.0 + robot = SimpleNamespace( + data=SimpleNamespace( + root_link_pose_w=SimpleNamespace(warp=wp.array(pose, dtype=wp.transformf, device=DEVICE)), + root_com_vel_w=SimpleNamespace(warp=wp.zeros(NUM_ENVS, dtype=wp.spatial_vectorf, device=DEVICE)), + ) + ) + all_mask = wp.ones(NUM_ENVS, dtype=wp.bool, device=DEVICE) + scratch_mask = wp.zeros(NUM_ENVS, dtype=wp.bool, device=DEVICE) + return SimpleNamespace( + num_envs=NUM_ENVS, + device=DEVICE, + scene={"robot": robot}, + rng_state_wp=wp.array(np.arange(NUM_ENVS, dtype=np.uint32), device=DEVICE), + resolve_env_mask=lambda env_ids=None, env_mask=None: resolve_1d_mask( + ids=env_ids, mask=env_mask, all_mask=all_mask, scratch_mask=scratch_mask, device=DEVICE + ), + sim=SimpleNamespace( + is_playing=lambda: True, + vis_marker_registry=SimpleNamespace(clear_debug_vis_callback=lambda term: None), + ), + _warp_graph_cache=graph_cache, + ) + + +def test_modify_term_cfg_applies_to_the_recorded_command_stage(): + """A curriculum that changes a command range after the command stage recorded changes the next commands. + + The range is a kernel argument read while the stage records, so a replay without recording again keeps + sampling from the old range. + """ + graph_cache = WarpGraphCache(DEVICE) + env = _env(graph_cache) + env.command_manager = CommandManager(CommandsCfg(), env) + params = { + "address": "commands.base_velocity.ranges.lin_vel_x", + "modify_fn": _override, + "modify_params": {"value": (1.5, 1.5)}, + } + curriculum = mdp.modify_term_cfg(CurriculumTermCfg(func=mdp.modify_term_cfg, params=params), env) + graph_cache.arm() + graph_cache.call_steps("CommandManager_compute", env.command_manager.stage_steps("compute"), dt=DT) + + curriculum(env, None, **params) + graph_cache.call_steps("CommandManager_compute", env.command_manager.stage_steps("compute"), dt=DT) + wp.synchronize() + + command = env.command_manager.get_command("base_velocity") + assert torch.equal(command[:, 0], torch.full((NUM_ENVS,), 1.5, device=DEVICE)) + graph_cache.close() diff --git a/source/isaaclab_experimental/test/managers/test_reward_manager.py b/source/isaaclab_experimental/test/managers/test_reward_manager.py index 5af3755aeee5..3b13737ec7a7 100644 --- a/source/isaaclab_experimental/test/managers/test_reward_manager.py +++ b/source/isaaclab_experimental/test/managers/test_reward_manager.py @@ -77,3 +77,30 @@ def test_capturability_is_decided_per_term_instance(): 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() + + +def test_set_term_cfg_applies_to_the_recorded_reward_stage(): + """Raising a weight from zero after the stage recorded adds the term, as ``modify_reward_weight`` expects. + + A zero-weight term is skipped while the stage records, so a replay without recording again never runs it. + """ + 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, + ) + manager = RewardManager({"late": RewardTermCfg(func=constant_reward, weight=0.0, params={"value": 3.0})}, env) + graph_cache.arm() + graph_cache.call_steps("RewardManager_compute", manager.stage_steps("compute"), dt=0.5) + + term_cfg = manager.get_term_cfg("late") + term_cfg.weight = 2.0 + manager.set_term_cfg("late", term_cfg) + reward = graph_cache.call_steps("RewardManager_compute", manager.stage_steps("compute"), dt=0.5) + wp.synchronize() + + assert torch.equal(reward, torch.full((NUM_ENVS,), 3.0 * 2.0 * 0.5, device=DEVICE)) + graph_cache.close() diff --git a/source/isaaclab_experimental/test/managers/test_termination_manager.py b/source/isaaclab_experimental/test/managers/test_termination_manager.py index 8d60dbad3329..7dcdb8c29f40 100644 --- a/source/isaaclab_experimental/test/managers/test_termination_manager.py +++ b/source/isaaclab_experimental/test/managers/test_termination_manager.py @@ -8,8 +8,10 @@ from types import SimpleNamespace import pytest +import torch import warp as wp from isaaclab_experimental.managers import TerminationManager +from isaaclab_experimental.utils import WarpGraphCache from isaaclab.managers import TerminationTermCfg @@ -37,3 +39,39 @@ def test_episode_termination_metric_averages_over_all_environments(): extras = manager.reset(env_mask=env_mask) assert float(extras["Episode_Termination/fell"]) == pytest.approx(0.5) + + +@wp.kernel +def _above_kernel(values: wp.array(dtype=wp.float32), threshold: float, out: wp.array(dtype=wp.bool)): + i = wp.tid() + out[i] = values[i] > threshold + + +def value_above(env, out: wp.array, threshold: float): + """Warp-first termination term that compares the environment's values with a threshold.""" + wp.launch(_above_kernel, dim=env.num_envs, inputs=[env.values, threshold, out], device=env.device) + + +@pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required") +def test_set_term_cfg_applies_to_the_recorded_termination_stage(): + """A threshold replaced after the stage recorded is used on the next call, not the recorded one.""" + graph_cache = WarpGraphCache("cuda:0") + env = SimpleNamespace( + num_envs=4, + device="cuda:0", + sim=SimpleNamespace(is_playing=lambda: True), + values=wp.array([0.1, 0.3, 0.5, 0.7], dtype=wp.float32, device="cuda:0"), + _warp_graph_cache=graph_cache, + ) + manager = TerminationManager({"above": TerminationTermCfg(func=value_above, params={"threshold": 0.6})}, env) + graph_cache.arm() + graph_cache.call_steps("TerminationManager_compute", manager.stage_steps("compute")) + + term_cfg = manager.get_term_cfg("above") + term_cfg.params["threshold"] = 0.2 + manager.set_term_cfg("above", term_cfg) + dones = graph_cache.call_steps("TerminationManager_compute", manager.stage_steps("compute")) + wp.synchronize() + + assert torch.equal(dones.cpu(), torch.tensor([False, True, True, True])) + graph_cache.close() From 1252efc58e72e42d6d9a87fb00f9778ebd509729 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 30 Sep 2026 04:25:01 -0700 Subject: [PATCH 24/80] Document the Warp capture contract and rough task support The migration guide describes the stable-config frontend, the command terms, the author contract for recorded stages and Warp terms from other packages. The environment overview lists the rough velocity tasks and the stages that stay eager. --- .../concepts/warp_environment_migration.rst | 119 +++++++++++++++--- docs/source/concepts/warp_environments.rst | 26 +++- 2 files changed, 124 insertions(+), 21 deletions(-) diff --git a/docs/source/concepts/warp_environment_migration.rst b/docs/source/concepts/warp_environment_migration.rst index f1c3c5fc9067..7eb5e3d8b099 100644 --- a/docs/source/concepts/warp_environment_migration.rst +++ b/docs/source/concepts/warp_environment_migration.rst @@ -42,11 +42,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 @@ -55,15 +56,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 @@ -136,7 +138,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. @@ -172,15 +177,89 @@ asset. Both stages use warp kernels with pre-allocated ``_raw_actions`` and ``_p buffers. +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 +``env.rng_state_wp``; the command manager records ``compute`` and ``reset`` like the other stages. +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 stage of a Warp environment (for example ``RewardManager_compute`` or +``CommandManager_reset``) is recorded into a CUDA graph on its first call after the first +``step`` and replayed afterwards. 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 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 reward, event, + termination and command managers' ``set_term_cfg`` records their stages again, and the + ``modify_reward_weight`` and ``modify_term_cfg`` curricula go through it. 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; the observation manager has no + ``set_term_cfg``, so observation settings are fixed once their stage records. +- **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=...)``. It runs + eagerly between the recorded parts of its stage and 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=...)``; a class term + may instead set ``self._warp_capturable`` in ``__init__``. ``height_scan``, ``base_height_l2`` with a + height scanner and the rigid-body randomization events are examples. + +``ISAACLAB_WARP_CAPTURE=0`` 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 @@ -227,6 +306,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. Available Warp MDP Terms @@ -238,7 +320,7 @@ Available Warp MDP Terms * - Category - Available Terms - * - Observations (11) + * - Observations (12) - | ``base_pos_z`` | ``base_lin_vel`` | ``base_ang_vel`` @@ -250,6 +332,7 @@ Available Warp MDP Terms | ``joint_vel_rel`` | ``last_action`` | ``generated_commands`` + | ``height_scan`` * - Rewards (16) - | ``is_alive`` | ``is_terminated`` @@ -282,6 +365,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. diff --git a/docs/source/concepts/warp_environments.rst b/docs/source/concepts/warp_environments.rst index 59e56aeddb1d..6c86acb51eae 100644 --- a/docs/source/concepts/warp_environments.rst +++ b/docs/source/concepts/warp_environments.rst @@ -9,8 +9,10 @@ Warp Experimental Environments ``isaaclab_tasks_experimental``. It's an experimental feature. The experimental extensions introduce **warp-first** environment infrastructure with CUDA graph capture -support. All environment-side computation (observations, rewards, resets, actions) runs as pure Warp -kernels, eliminating Python overhead and enabling CUDA graph capture for maximum throughput. +support. ``--frontend warp`` runs stable task configurations on Warp managers with Newton physics: MDP +terms (observations, rewards, terminations, events, commands, actions) run as Warp kernels, and the +capturable terms of each manager stage are recorded into CUDA graphs. Scene writes, actuator models, the +stable curriculum and recorder managers, and terms marked non-capturable still run eagerly. Throughout this page, **stable** refers to the standard Torch-based implementations in the ``isaaclab`` and ``isaaclab_tasks`` packages (the default runtime, used when ``--frontend warp`` @@ -68,6 +70,11 @@ 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: @@ -75,12 +82,16 @@ The following contributed tasks also have full twin coverage: - ``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. Quick Start @@ -255,6 +266,11 @@ specific to warp envs; for Newton physics maturity and specialist 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 runs of its stage. Observation and event terms, which draw from the shared + per-environment random state, and action terms keep their configured order, so an eager term there + splits the stage into several graphs. A non-capturable command term keeps the whole command stage + eager. - 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. From cb244640b304573ffcbe1cc8ec11a90e5cf2ff41 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 30 Sep 2026 05:24:15 -0700 Subject: [PATCH 25/80] Record Warp observation stages again when a term configuration changes A recorded observation stage keeps the noise parameters its terms read while it recorded, so an observation noise curriculum written through modify_term_cfg, like the noise ADR terms of the Lift task, changed nothing once the stage was captured. The Warp observation manager gains set_term_cfg and get_term_cfg, naming a term "/". Setting a term prepares it again and drops the recorded observation graphs. The per-term preparation moves out of _prepare_terms so both paths share it; an existing noise model keeps its state. modify_term_cfg hands observation terms back under that name. A capture-then-mutate test changes a noise bias through modify_term_cfg after the observation stage recorded; it fails without the setter. --- .../concepts/warp_environment_migration.rst | 12 +- ...ichuanh-warp-frontend-managers.changed.rst | 2 + .../isaaclab/isaaclab/envs/mdp/curriculums.py | 4 +- .../jichuanh-warp-frontend-managers.added.rst | 3 + .../jichuanh-warp-frontend-managers.fixed.rst | 2 + .../managers/observation_manager.py | 305 +++++++++++------- .../test/managers/test_observation_manager.py | 66 +++- 7 files changed, 273 insertions(+), 121 deletions(-) diff --git a/docs/source/concepts/warp_environment_migration.rst b/docs/source/concepts/warp_environment_migration.rst index 7eb5e3d8b099..b22cba3ce2b8 100644 --- a/docs/source/concepts/warp_environment_migration.rst +++ b/docs/source/concepts/warp_environment_migration.rst @@ -214,12 +214,12 @@ part of it: 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 reward, event, - termination and command managers' ``set_term_cfg`` records their stages again, and the - ``modify_reward_weight`` and ``modify_term_cfg`` curricula go through it. 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; the observation manager has no - ``set_term_cfg``, so observation settings are fixed once their stage records. + or ``env.common_step_counter``, is frozen into the graph when the stage records. The managers' + ``set_term_cfg`` records their stages again (observation terms are named ``"/"``), and the + ``modify_reward_weight`` and ``modify_term_cfg`` curricula go through it, so an observation noise + curriculum applies to recorded observations. 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=...)``. It runs eagerly between the recorded parts of its stage and raises if it is ever called while a stage diff --git a/source/isaaclab/changelog.d/jichuanh-warp-frontend-managers.changed.rst b/source/isaaclab/changelog.d/jichuanh-warp-frontend-managers.changed.rst index 486f2e3b24bd..1e57e3b1fe92 100644 --- a/source/isaaclab/changelog.d/jichuanh-warp-frontend-managers.changed.rst +++ b/source/isaaclab/changelog.d/jichuanh-warp-frontend-managers.changed.rst @@ -1,3 +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 ``"/"`` name. diff --git a/source/isaaclab/isaaclab/envs/mdp/curriculums.py b/source/isaaclab/isaaclab/envs/mdp/curriculums.py index 5568ec2ab9dd..f3381159463e 100644 --- a/source/isaaclab/isaaclab/envs/mdp/curriculums.py +++ b/source/isaaclab/isaaclab/envs/mdp/curriculums.py @@ -292,7 +292,9 @@ def __init__(self, cfg, env): # the simplified address starts with the manager's configuration group and the term name group, _, term_path = self._address.partition(".") self._manager_name = f"{group[:-1]}_manager" - self._term_name = re.split(r"[.\[]", term_path, maxsplit=1)[0] + names = re.split(r"[.\[]", term_path) + # observation terms are nested in observation groups and named "/" + self._term_name = "/".join(names[:2]) if group == "observations" else names[0] # overwrite the simplified address with the full manager path self._address = self._address.replace("s.", "_manager.cfg.", 1) diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst index 3f2183dfaf59..b5f53bcb44ba 100644 --- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst @@ -14,3 +14,6 @@ :class:`~isaaclab_experimental.managers.CommandManager`. Setting a term configuration records the manager's stages again, so a curriculum that changes a termination parameter or a command range through ``modify_term_cfg`` takes effect under CUDA graph capture. +* Added ``set_term_cfg`` and ``get_term_cfg`` to the Warp :class:`~isaaclab_experimental.managers.ObservationManager`. + Observation terms are named ``"/"``. Setting a term configuration records the observation stages + again. diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst index 747ec18193a2..14820a067666 100644 --- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.fixed.rst @@ -4,3 +4,5 @@ ray-caster height scanner. It now runs eagerly in that configuration. * Fixed the Warp frontend replacing a Torch term defined outside Isaac Lab with an unrelated built-in Warp term of the same name. Such a term is now reported as not being a Warp term. +* Fixed observation noise curricula that change a noise configuration through ``modify_term_cfg`` having no effect + on the observations of a Warp environment under CUDA graph capture. diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py index 2407d7633b3e..a48964cb3866 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py @@ -563,6 +563,63 @@ def _compute_term( else: group_obs[term_name] = circular_buffer.buffer + """ + Operations - Term settings. + """ + + def set_term_cfg(self, term_name: str, cfg: ObservationTermCfg): + """Sets the configuration of the specified term into the manager. + + Recorded observation stages read the term's configuration, such as its noise, when they record, so they + record again on their next call. The term keeps writing into its slice of the group's observation. + + Args: + term_name: The name of the observation term, in the form ``"/"``. + cfg: The configuration for the observation term. + + Raises: + ValueError: If the term name is not found, or if the configuration changes the term's output width. + """ + group_name, index = self._find_term(term_name) + group_cfg = self.cfg[group_name] if isinstance(self.cfg, dict) else getattr(self.cfg, group_name) + old_cfg = self._group_obs_term_cfgs[group_name][index] + term_dim = self._prepare_term(group_name, self._group_obs_term_names[group_name][index], cfg, group_cfg) + if term_dim != old_cfg._term_dim: + raise ValueError( + f"Observation term '{term_name}' must keep its output width {old_cfg._term_dim}, got {term_dim}." + ) + cfg.out_wp, cfg.out_torch = old_cfg.out_wp, old_cfg.out_torch + self._group_obs_term_cfgs[group_name][index] = cfg + if cfg is not old_cfg and isinstance(cfg.func, ManagerTermBase): + cfg.func.reset() + self._group_obs_class_term_cfgs[group_name] = [ + term_cfg for term_cfg in self._group_obs_term_cfgs[group_name] if isinstance(term_cfg.func, ManagerTermBase) + ] + self._group_obs_class_instances = self._collect_class_instances() + self._clear_stage_steps() + + def get_term_cfg(self, term_name: str) -> ObservationTermCfg: + """Gets the configuration for the specified term. + + Args: + term_name: The name of the observation term, in the form ``"/"``. + + Returns: + The configuration of the observation term. + + Raises: + ValueError: If the term name is not found. + """ + group_name, index = self._find_term(term_name) + return self._group_obs_term_cfgs[group_name][index] + + def _find_term(self, term_name: str) -> tuple[str, int]: + """Return the group and the index within the group of a ``"/"`` name.""" + group_name, _, name = term_name.partition("/") + if name not in self._group_obs_term_names.get(group_name, ()): + raise ValueError(f"Observation term '{term_name}' not found.") + return group_name, self._group_obs_term_names[group_name].index(name) + """ Operations - Stage steps. """ @@ -724,124 +781,14 @@ def _prepare_terms(self): # noqa: C901 f"Configuration for the term '{term_name}' is not of type ObservationTermCfg." f" Received: '{type(term_cfg)}'." ) - # resolve common terms in the config - # Warp-first signature is (env, out, **params) - self._resolve_common_term_cfg(f"{group_name}/{term_name}", term_cfg, min_argc=2) - - # check noise settings - if not group_cfg.enable_corruption: - term_cfg.noise = None - elif term_cfg.noise is not None and not type(term_cfg.noise).__module__.startswith( - "isaaclab_experimental." - ): - raise TypeError( - f"Observation term '{term_name}' uses noise cfg" - f" '{type(term_cfg.noise).__module__}.{type(term_cfg.noise).__name__}', which the Warp" - " observation manager would silently ignore. Use the warp-native noise cfgs from" - " isaaclab_experimental.utils.noise (the warp frontend converts stable cfgs" - " automatically)." - ) - # check group history params and override terms - 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 + term_dim = self._prepare_term(group_name, term_name, term_cfg, group_cfg) # 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) - - # infer dimensions (Warp-first: terms write to out; we infer dim from resolved scene info) - term_dim = self._infer_term_dim_scalar(term_cfg) - # Cache the "raw" term output dimension (before history reshaping) for Warp buffer allocation. - # This matches the tensor shape produced directly by the term into `out`: (num_envs, term_dim). - term_cfg._term_dim = int(term_dim) group_term_cfgs.append(term_cfg) - group_term_raw_dims.append(int(term_dim)) + group_term_raw_dims.append(term_dim) obs_dims = (self._env.num_envs, term_dim) - # if scale is set, check if single float or tuple - if term_cfg.scale is not None: - if not isinstance(term_cfg.scale, (float, int, tuple)): - raise TypeError( - f"Scale for observation term '{term_name}' in group '{group_name}'" - f" is not of type float, int or tuple. Received: '{type(term_cfg.scale)}'." - ) - if isinstance(term_cfg.scale, tuple) and len(term_cfg.scale) != obs_dims[1]: - raise ValueError( - f"Scale for observation term '{term_name}' in group '{group_name}'" - f" does not match the dimensions of the observation. Expected: {obs_dims[1]}" - f" but received: {len(term_cfg.scale)}." - ) - - scale_vals = ( - term_cfg.scale if isinstance(term_cfg.scale, tuple) else [float(term_cfg.scale)] * obs_dims[1] - ) - term_cfg.scale_wp = wp.array(scale_vals, dtype=wp.float32, device=self._env.device) - - # prepare modifiers for each observation - if term_cfg.modifiers is not None: - # initialize list of modifiers for term - for mod_cfg in term_cfg.modifiers: - # check if class modifier and initialize with observation size when adding - if isinstance(mod_cfg, modifiers.ModifierCfg): - # to list of modifiers - if inspect.isclass(mod_cfg.func): - if not issubclass(mod_cfg.func, modifiers.ModifierBase): - raise TypeError( - f"Modifier function '{mod_cfg.func}' for observation term '{term_name}'" - f" is not a subclass of 'ModifierBase'. Received: '{type(mod_cfg.func)}'." - ) - mod_cfg.func = mod_cfg.func(cfg=mod_cfg, data_dim=obs_dims, device=self._env.device) - - # add to list of class modifiers - self._group_obs_class_instances.append(mod_cfg.func) - else: - raise TypeError( - f"Modifier configuration '{mod_cfg}' of observation term '{term_name}' is not of" - f" required type ModifierCfg, Received: '{type(mod_cfg)}'" - ) - - # check if function is callable - if not callable(mod_cfg.func): - raise AttributeError( - f"Modifier '{mod_cfg}' of observation term '{term_name}' is not callable." - f" Received: {mod_cfg.func}" - ) - - # check if term's arguments are matched by params - term_params = list(mod_cfg.params.keys()) - args = inspect.signature(mod_cfg.func).parameters - args_with_defaults = [arg for arg in args if args[arg].default is not inspect.Parameter.empty] - args_without_defaults = [arg for arg in args if args[arg].default is inspect.Parameter.empty] - args = args_without_defaults + args_with_defaults - # ignore first argument for data - if len(args) > 1: - if set(args[1:]) != set(term_params + args_with_defaults): - raise ValueError( - f"Modifier '{mod_cfg}' of observation term '{term_name}' expects" - f" mandatory parameters: {args_without_defaults[1:]}" - 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( - f"Class type for observation term '{term_name}' NoiseModelCfg" - f" is not a subclass of 'NoiseModel'. Received: '{type(noise_model_cls)}'." - ) - # initialize func to be the noise model class instance - term_cfg.noise.func = instantiate( - term_cfg.noise, num_envs=self._env.num_envs, device=self._env.device - ) - self._group_obs_class_instances.append(term_cfg.noise.func) - # create history buffers and calculate history term dimensions if term_cfg.history_length > 0: group_entry_history_buffer[term_name] = CircularBuffer( @@ -915,6 +862,140 @@ def _prepare_terms(self): # noqa: C901 # add history buffers for each group self._group_obs_term_history_buffer[group_name] = group_entry_history_buffer + self._group_obs_class_instances = self._collect_class_instances() + + def _prepare_term( + self, group_name: str, term_name: str, term_cfg: ObservationTermCfg, group_cfg: ObservationGroupCfg + ) -> int: + """Resolve a term configuration and prepare its scale, modifiers and noise. + + Args: + group_name: Name of the term's group. + term_name: Name of the term within its group. + term_cfg: The term configuration. + group_cfg: The group configuration. + + Returns: + The width of the term's output. + """ + # resolve common terms in the config + # Warp-first signature is (env, out, **params) + self._resolve_common_term_cfg(f"{group_name}/{term_name}", term_cfg, min_argc=2) + + # check noise settings + if not group_cfg.enable_corruption: + term_cfg.noise = None + elif term_cfg.noise is not None and not type(term_cfg.noise).__module__.startswith("isaaclab_experimental."): + raise TypeError( + f"Observation term '{term_name}' uses noise cfg" + f" '{type(term_cfg.noise).__module__}.{type(term_cfg.noise).__name__}', which the Warp" + " observation manager would silently ignore. Use the warp-native noise cfgs from" + " isaaclab_experimental.utils.noise (the warp frontend converts stable cfgs" + " automatically)." + ) + # check group history params and override terms + 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 + + # infer dimensions (Warp-first: terms write to out; we infer dim from resolved scene info) + term_dim = int(self._infer_term_dim_scalar(term_cfg)) + # Cache the "raw" term output dimension (before history reshaping) for Warp buffer allocation. + # This matches the tensor shape produced directly by the term into `out`: (num_envs, term_dim). + term_cfg._term_dim = term_dim + obs_dims = (self._env.num_envs, term_dim) + + # if scale is set, check if single float or tuple + if term_cfg.scale is not None: + if not isinstance(term_cfg.scale, (float, int, tuple)): + raise TypeError( + f"Scale for observation term '{term_name}' in group '{group_name}'" + f" is not of type float, int or tuple. Received: '{type(term_cfg.scale)}'." + ) + if isinstance(term_cfg.scale, tuple) and len(term_cfg.scale) != obs_dims[1]: + raise ValueError( + f"Scale for observation term '{term_name}' in group '{group_name}'" + f" does not match the dimensions of the observation. Expected: {obs_dims[1]}" + f" but received: {len(term_cfg.scale)}." + ) + + scale_vals = term_cfg.scale if isinstance(term_cfg.scale, tuple) else [float(term_cfg.scale)] * obs_dims[1] + term_cfg.scale_wp = wp.array(scale_vals, dtype=wp.float32, device=self._env.device) + + # prepare modifiers for each observation + if term_cfg.modifiers is not None: + # initialize list of modifiers for term + for mod_cfg in term_cfg.modifiers: + # check if class modifier and initialize with observation size when adding + if isinstance(mod_cfg, modifiers.ModifierCfg): + # to list of modifiers + if inspect.isclass(mod_cfg.func): + if not issubclass(mod_cfg.func, modifiers.ModifierBase): + raise TypeError( + f"Modifier function '{mod_cfg.func}' for observation term '{term_name}'" + f" is not a subclass of 'ModifierBase'. Received: '{type(mod_cfg.func)}'." + ) + mod_cfg.func = mod_cfg.func(cfg=mod_cfg, data_dim=obs_dims, device=self._env.device) + else: + raise TypeError( + f"Modifier configuration '{mod_cfg}' of observation term '{term_name}' is not of" + f" required type ModifierCfg, Received: '{type(mod_cfg)}'" + ) + + # check if function is callable + if not callable(mod_cfg.func): + raise AttributeError( + f"Modifier '{mod_cfg}' of observation term '{term_name}' is not callable." + f" Received: {mod_cfg.func}" + ) + + # check if term's arguments are matched by params + term_params = list(mod_cfg.params.keys()) + args = inspect.signature(mod_cfg.func).parameters + args_with_defaults = [arg for arg in args if args[arg].default is not inspect.Parameter.empty] + args_without_defaults = [arg for arg in args if args[arg].default is inspect.Parameter.empty] + args = args_without_defaults + args_with_defaults + # ignore first argument for data + if len(args) > 1: + if set(args[1:]) != set(term_params + args_with_defaults): + raise ValueError( + f"Modifier '{mod_cfg}' of observation term '{term_name}' expects" + f" mandatory parameters: {args_without_defaults[1:]}" + 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( + f"Class type for observation term '{term_name}' NoiseModelCfg" + f" is not a subclass of 'NoiseModel'. Received: '{type(noise_model_cls)}'." + ) + # initialize func to be the noise model class instance, keeping an existing one and its state + if not isinstance(term_cfg.noise.func, noise.NoiseModel): + term_cfg.noise.func = instantiate(term_cfg.noise, num_envs=self._env.num_envs, device=self._env.device) + + return term_dim + + def _collect_class_instances(self) -> list[modifiers.ModifierBase | noise.NoiseModel]: + """Return the class modifiers and noise models of every term, which are reset with the manager.""" + instances = [] + for term_cfgs in self._group_obs_term_cfgs.values(): + for term_cfg in term_cfgs: + for mod_cfg in term_cfg.modifiers or (): + if isinstance(mod_cfg.func, modifiers.ModifierBase): + instances.append(mod_cfg.func) + if isinstance(term_cfg.noise, noise.NoiseModelCfg): + instances.append(term_cfg.noise.func) + return instances + def _infer_term_dim_scalar(self, term_cfg: ObservationTermCfg) -> int: """Infer observation output dimension (D,) using decorator metadata, scene info, or manager state. diff --git a/source/isaaclab_experimental/test/managers/test_observation_manager.py b/source/isaaclab_experimental/test/managers/test_observation_manager.py index 3094a9863da0..dd17149e8421 100644 --- a/source/isaaclab_experimental/test/managers/test_observation_manager.py +++ b/source/isaaclab_experimental/test/managers/test_observation_manager.py @@ -3,14 +3,19 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Tests for the Warp observation manager configuration parsing.""" +"""Tests for the Warp observation manager.""" from types import SimpleNamespace +import pytest +import torch import warp as wp from isaaclab_experimental.managers import ObservationManager +from isaaclab_experimental.utils import WarpGraphCache +from isaaclab_experimental.utils.noise import ConstantNoiseCfg -from isaaclab.managers import ObservationGroupCfg, ObservationTermCfg +import isaaclab.envs.mdp as mdp +from isaaclab.managers import CurriculumTermCfg, ObservationGroupCfg, ObservationTermCfg from isaaclab.utils import configclass @@ -35,3 +40,60 @@ def test_group_settings_are_not_parsed_as_terms(): assert manager.active_terms["policy"] == ["first", "second"] assert manager.group_obs_dim["policy"] == (5,) + + +@configclass +class NoisyPolicyCfg(ObservationGroupCfg): + """Group with one term corrupted by a constant bias.""" + + value = ObservationTermCfg(func=constant_obs, params={"out_dim": 2}, noise=ConstantNoiseCfg(bias=0.0)) + + def __post_init__(self): + self.enable_corruption = True + + +@configclass +class ObservationsCfg: + policy: NoisyPolicyCfg = NoisyPolicyCfg() + + +def _override(env, env_ids, data, value): + return value + + +@pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required") +def test_modify_term_cfg_applies_noise_to_the_recorded_observation_stage(): + """An observation noise curriculum changes the observations of a recorded stage, as noise ADR expects. + + The noise parameters are kernel arguments read while the stage records, so a replay without recording again + keeps applying the old noise. + """ + graph_cache = WarpGraphCache("cuda:0") + env = SimpleNamespace( + num_envs=4, + device="cuda:0", + sim=SimpleNamespace(is_playing=lambda: True), + rng_state_wp=wp.zeros(4, dtype=wp.uint32, device="cuda:0"), + _warp_graph_cache=graph_cache, + ) + env.observation_manager = ObservationManager(ObservationsCfg(), env) + params = { + "address": "observations.policy.value.noise.bias", + "modify_fn": _override, + "modify_params": {"value": 0.5}, + } + curriculum = mdp.modify_term_cfg(CurriculumTermCfg(func=mdp.modify_term_cfg, params=params), env) + graph_cache.arm() + graph_cache.call_steps( + "ObservationManager_compute", env.observation_manager.stage_steps("compute"), return_cloned_output=False + ) + assert graph_cache.captured_stages == ("ObservationManager_compute",) + + curriculum(env, None, **params) + obs = graph_cache.call_steps( + "ObservationManager_compute", env.observation_manager.stage_steps("compute"), return_cloned_output=False + ) + wp.synchronize() + + assert torch.equal(obs["policy"], torch.full((4, 2), 1.5, device="cuda:0")) + graph_cache.close() From 3f906ebe8022b97134b5f4fcca5b94e055b0c93a Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 30 Sep 2026 15:13:49 -0700 Subject: [PATCH 26/80] Record Warp stages again only when a term configuration changes Curricula such as modify_reward_weight set their term on every reset past their threshold, and set_term_cfg dropped the manager's graphs on every call, so the reward stage was recorded again on every reset step. The reward stage now checks zero weights on the device: while recording, each term runs under a conditional graph node on an int32 flag kept beside the weights, so any weight change, including to and from zero, is a device write. Without conditional node support (CUDA < 12.4) the record-time check stays and a zero crossing records again. The reward, event, termination, command and observation setters compare the term against the values taken when the steps were built, since curricula change the stored configuration in place, and drop the graphs only when what the recorded steps read differs: the function, parameters and capturability, an event term's trigger settings, and the whole configuration of command and observation terms. An unchanged event or observation term is no longer resolved again, which reallocated its Warp buffers. --- .../concepts/warp_environment_migration.rst | 12 +-- .../jichuanh-warp-frontend-managers.added.rst | 9 +- ...ichuanh-warp-frontend-managers.changed.rst | 6 ++ .../managers/command_manager.py | 18 ++-- .../managers/event_manager.py | 26 +++++- .../managers/manager_base.py | 53 +++++++++++- .../managers/observation_manager.py | 40 ++++++--- .../managers/reward_manager.py | 37 +++++++- .../managers/termination_manager.py | 13 ++- .../test/envs/test_warp_env_curriculum.py | 85 +++++++++++++++++++ .../test/managers/test_command_manager.py | 20 ++++- .../test/managers/test_event_manager.py | 47 ++++++++++ .../test/managers/test_observation_manager.py | 22 +++-- .../test/managers/test_reward_manager.py | 85 +++++++++++++------ 14 files changed, 393 insertions(+), 80 deletions(-) create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.changed.rst create mode 100644 source/isaaclab_experimental/test/envs/test_warp_env_curriculum.py diff --git a/docs/source/concepts/warp_environment_migration.rst b/docs/source/concepts/warp_environment_migration.rst index b22cba3ce2b8..acfec1681ac8 100644 --- a/docs/source/concepts/warp_environment_migration.rst +++ b/docs/source/concepts/warp_environment_migration.rst @@ -215,11 +215,13 @@ part of it: 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 (observation terms are named ``"/"``), and the - ``modify_reward_weight`` and ``modify_term_cfg`` curricula go through it, so an observation noise - curriculum applies to recorded observations. 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. + ``set_term_cfg`` records their stages again when it changes a term's configuration (observation terms + are named ``"/"``), 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, and reward weights are read on the device, so a weight change + applies without recording again (before CUDA 12.4, a weight set to or from zero records again). 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=...)``. It runs eagerly between the recorded parts of its stage and raises if it is ever called while a stage diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst index b5f53bcb44ba..4b911286d641 100644 --- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.added.rst @@ -11,9 +11,8 @@ :class:`~isaaclab_experimental.utils.warp.WarpCapturable` are used as configured. * Added ``set_term_cfg`` and ``get_term_cfg`` to the Warp :class:`~isaaclab_experimental.managers.TerminationManager` and - :class:`~isaaclab_experimental.managers.CommandManager`. Setting a term configuration records the manager's - stages again, so a curriculum that changes a termination parameter or a command range through - ``modify_term_cfg`` takes effect under CUDA graph capture. + :class:`~isaaclab_experimental.managers.CommandManager`. A curriculum that changes a termination parameter or + a command range through ``modify_term_cfg`` takes effect under CUDA graph capture. * Added ``set_term_cfg`` and ``get_term_cfg`` to the Warp :class:`~isaaclab_experimental.managers.ObservationManager`. - Observation terms are named ``"/"``. Setting a term configuration records the observation stages - again. + Observation terms are named ``"/"``. Setting a changed term configuration records the observation + stages again. diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.changed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.changed.rst new file mode 100644 index 000000000000..7f2cd0f3b24e --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.changed.rst @@ -0,0 +1,6 @@ +* Changed ``set_term_cfg`` of the Warp reward, event, termination and command managers to record the manager's + stages again only when the term's function, parameters or capturability change, or an event term's trigger + settings. Setting an unchanged configuration, as curricula do on every reset, keeps the recorded stages. +* Changed the Warp reward manager to check zero term weights inside the recorded reward stage, so a weight set + to or from zero applies without recording the stage again. This needs CUDA 12.4 or newer; with older CUDA + versions a weight crossing zero records the stage again. diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py index 0996895f6e18..7aae4740907f 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py @@ -9,8 +9,8 @@ import inspect from abc import abstractmethod -from collections.abc import Sequence -from typing import TYPE_CHECKING +from collections.abc import Iterable, Sequence +from typing import TYPE_CHECKING, Any import torch import warp as wp @@ -22,7 +22,7 @@ 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 +from .manager_base import ManagerBase, ManagerTermBase, _config_key # import omni.kit.app @@ -572,7 +572,7 @@ def set_term_cfg(self, name: str, cfg: CommandTermCfg): """Sets the configuration of the specified command term. A command term reads its configuration, such as its sampling ranges, when a recorded command stage records, - so the command stages record again on their next call. + so the command stages record again on their next call when the configuration changes. Args: name: The name of the command term. @@ -583,12 +583,14 @@ def set_term_cfg(self, name: str, cfg: CommandTermCfg): """ if name not in self._terms: raise ValueError(f"Command term '{name}' not found.") + changed = self._term_cfg_changed(name, cfg) self._terms[name].cfg = cfg if isinstance(self.cfg, dict): self.cfg[name] = cfg else: setattr(self.cfg, name, cfg) - self._clear_stage_steps() + if changed: + self._clear_stage_steps() def get_term_cfg(self, name: str) -> CommandTermCfg: """Gets the configuration of the specified command term. @@ -621,6 +623,12 @@ def get_term(self, name: str) -> CommandTerm: Helper functions. """ + def _named_term_cfgs(self) -> Iterable[tuple[str, CommandTermCfg]]: + return ((name, term.cfg) for name, term in self._terms.items()) + + def _term_signature(self, term_cfg: CommandTermCfg) -> Any: + return _config_key(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/event_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py index a961e5b2fb16..75b9c86b93ee 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py @@ -23,7 +23,7 @@ import inspect import logging -from collections.abc import Callable, Sequence +from collections.abc import Callable, Iterable, Sequence from functools import partial from typing import Any @@ -229,7 +229,8 @@ def available_modes(self) -> list[str]: def set_term_cfg(self, term_name: str, cfg: EventTermCfg): """Set the configuration of an existing event term. - Recorded event stages read the replaced configuration, so they are recorded again on their next call. + Recorded event stages record again on their next call when the replacement changes the term's function, + parameters, capturability or trigger settings. Args: term_name: Name of the event term. @@ -249,7 +250,10 @@ def set_term_cfg(self, term_name: str, cfg: EventTermCfg): 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) + # an unchanged configuration was resolved already; resolving it again reallocates its Warp selections + changed = self._term_cfg_changed(term_name, cfg) + if changed: + 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 @@ -259,7 +263,8 @@ 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) ] - self._clear_stage_steps() + if changed: + self._clear_stage_steps() return raise ValueError(f"Event term '{term_name}' not found.") @@ -458,6 +463,19 @@ def _resample_interval(self, index: int, env_mask: wp.array) -> None: def _reset_output(env_mask: wp.array) -> dict[str, float]: return {} + def _named_term_cfgs(self) -> Iterable[tuple[str, EventTermCfg]]: + for mode, names in self._mode_term_names.items(): + yield from zip(names, self._mode_term_cfgs[mode]) + + def _term_signature(self, term_cfg: EventTermCfg) -> Any: + # trigger settings are kernel arguments and decide which reset steps exist + trigger = ( + term_cfg.interval_range_s, + term_cfg.min_step_count_between_reset, + term_cfg.resample_interval_on_reset, + ) + return super()._term_signature(term_cfg), trigger + 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 95c7d45f1c0c..b4c1f2b940d6 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py @@ -17,12 +17,15 @@ import contextlib import copy +import dataclasses import inspect import logging from abc import ABC, abstractmethod -from collections.abc import Callable, Sequence +from collections.abc import Callable, Iterable, Sequence from typing import TYPE_CHECKING, Any +import numpy as np +import torch import warp as wp import isaaclab.utils.string as string_utils @@ -37,8 +40,6 @@ if TYPE_CHECKING: - import torch - from isaaclab.envs import ManagerBasedEnv # import logger @@ -197,6 +198,8 @@ def __init__(self, cfg: object, env: ManagerBasedEnv): 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]], ...]] = {} + # what the built steps read from each term configuration, compared when a term configuration is set + self._term_signatures: dict[str, Any] | None = None self._all_terms_capturable = True # parse config to create terms information @@ -302,6 +305,8 @@ def stage_steps(self, operation: str) -> tuple[tuple[bool, Callable[..., Any]], """ steps = self._stage_steps.get(operation) if steps is None: + if self._term_signatures is None: + self._term_signatures = {name: self._term_signature(cfg) for name, cfg in self._named_term_cfgs()} steps = self._stage_steps[operation] = tuple(self._build_stage_steps(operation)) return steps @@ -333,10 +338,31 @@ def _run_steps(self, operation: str, *args: Any, **kwargs: Any) -> Any: def _clear_stage_steps(self) -> None: """Rebuild the steps on next use and drop the graphs recorded from them.""" self._stage_steps.clear() + self._term_signatures = None graph_cache = getattr(self._env, "_warp_graph_cache", None) if graph_cache is not None: graph_cache.invalidate(type(self).__name__) + def _named_term_cfgs(self) -> Iterable[tuple[str, Any]]: + """Name and configuration of each term whose configuration can be set.""" + return () + + def _term_signature(self, term_cfg: ManagerTermBaseCfg) -> Any: + """What the steps read from a term configuration when they are built or recorded. + + Values the steps read from device arrays on every call are left out, since a replayed graph sees + their changes. + """ + return (term_cfg.func, _config_key(term_cfg.params), is_warp_capturable(term_cfg.func, term_cfg.params)) + + def _term_cfg_changed(self, term_name: str, term_cfg: Any) -> bool: + """Whether the built steps read a different configuration of the term. + + Compares against the values taken when the steps were built, since curricula change a stored + configuration in place before setting it again. + """ + return self._term_signatures is None or self._term_signatures.get(term_name) != self._term_signature(term_cfg) + @abstractmethod def _prepare_terms(self): """Prepare terms information from the configuration object.""" @@ -509,3 +535,24 @@ def _process_term_cfg_at_play(self, term_name: str, term_cfg: ManagerTermBaseCfg if inspect.isclass(term_cfg.func): logger.info(f"Initializing term '{term_name}' with class '{term_cfg.func.__name__}'.") term_cfg.func = term_cfg.func(cfg=term_cfg, env=self._env) + + +def _config_key(value: Any) -> Any: + """Comparable copy of a configuration value, taken by value. + + Arrays are keyed by the memory a recorded graph reads, configuration objects by their attributes + (including those set when they are resolved, such as the Warp joint mask), and other objects by identity. + """ + if isinstance(value, wp.array): + return ("wp.array", value.ptr, value.shape, value.strides, value.dtype) + if isinstance(value, torch.Tensor): + return ("tensor", value.data_ptr(), tuple(value.shape), value.stride(), value.dtype) + if isinstance(value, np.ndarray): + return ("ndarray", value.shape, value.dtype.str, value.tobytes()) + if isinstance(value, dict): + return ("dict", tuple((key, _config_key(item)) for key, item in value.items())) + if isinstance(value, (list, tuple)): + return (type(value), tuple(_config_key(item) for item in value)) + if dataclasses.is_dataclass(value) and not isinstance(value, type): + return (type(value), _config_key(vars(value))) + return value diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py index a48964cb3866..229934a8b242 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py @@ -47,7 +47,7 @@ from __future__ import annotations import inspect -from collections.abc import Callable, Sequence +from collections.abc import Callable, Iterable, Sequence from functools import partial from typing import TYPE_CHECKING, Any @@ -65,7 +65,7 @@ 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 +from .manager_base import ManagerBase, ManagerTermBase, _config_key if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv @@ -571,7 +571,8 @@ def set_term_cfg(self, term_name: str, cfg: ObservationTermCfg): """Sets the configuration of the specified term into the manager. Recorded observation stages read the term's configuration, such as its noise, when they record, so they - record again on their next call. The term keeps writing into its slice of the group's observation. + record again on their next call when the configuration changes. The term keeps writing into its slice of + the group's observation. Args: term_name: The name of the observation term, in the form ``"/"``. @@ -581,22 +582,26 @@ def set_term_cfg(self, term_name: str, cfg: ObservationTermCfg): ValueError: If the term name is not found, or if the configuration changes the term's output width. """ group_name, index = self._find_term(term_name) - group_cfg = self.cfg[group_name] if isinstance(self.cfg, dict) else getattr(self.cfg, group_name) old_cfg = self._group_obs_term_cfgs[group_name][index] - term_dim = self._prepare_term(group_name, self._group_obs_term_names[group_name][index], cfg, group_cfg) - if term_dim != old_cfg._term_dim: - raise ValueError( - f"Observation term '{term_name}' must keep its output width {old_cfg._term_dim}, got {term_dim}." - ) - cfg.out_wp, cfg.out_torch = old_cfg.out_wp, old_cfg.out_torch + # an unchanged configuration was prepared already; preparing it again reallocates its Warp buffers + changed = self._term_cfg_changed(term_name, cfg) + if changed: + group_cfg = self.cfg[group_name] if isinstance(self.cfg, dict) else getattr(self.cfg, group_name) + term_dim = self._prepare_term(group_name, self._group_obs_term_names[group_name][index], cfg, group_cfg) + if term_dim != old_cfg._term_dim: + raise ValueError( + f"Observation term '{term_name}' must keep its output width {old_cfg._term_dim}, got {term_dim}." + ) + cfg.out_wp, cfg.out_torch = old_cfg.out_wp, old_cfg.out_torch + if cfg is not old_cfg and isinstance(cfg.func, ManagerTermBase): + cfg.func.reset() self._group_obs_term_cfgs[group_name][index] = cfg - if cfg is not old_cfg and isinstance(cfg.func, ManagerTermBase): - cfg.func.reset() self._group_obs_class_term_cfgs[group_name] = [ term_cfg for term_cfg in self._group_obs_term_cfgs[group_name] if isinstance(term_cfg.func, ManagerTermBase) ] self._group_obs_class_instances = self._collect_class_instances() - self._clear_stage_steps() + if changed: + self._clear_stage_steps() def get_term_cfg(self, term_name: str) -> ObservationTermCfg: """Gets the configuration for the specified term. @@ -692,6 +697,15 @@ def serialize(self) -> dict: Helper functions. """ + def _named_term_cfgs(self) -> Iterable[tuple[str, ObservationTermCfg]]: + for group_name, term_names in self._group_obs_term_names.items(): + for term_name, term_cfg in zip(term_names, self._group_obs_term_cfgs[group_name]): + yield f"{group_name}/{term_name}", term_cfg + + def _term_signature(self, term_cfg: ObservationTermCfg) -> Any: + # noise, clip, scale, modifiers and history are read while the stage records + return _config_key(term_cfg), is_warp_capturable(term_cfg.func, term_cfg.params) + def _prepare_terms(self): # noqa: C901 """Prepares a list of observation terms functions.""" # create buffers to store information for each observation group diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py index f415ce03c907..00a6b851f1b4 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py @@ -11,7 +11,7 @@ from __future__ import annotations -from collections.abc import Callable, Sequence +from collections.abc import Callable, Iterable, Sequence from functools import partial from typing import TYPE_CHECKING, Any @@ -179,6 +179,13 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv): self._term_weights_wp = wp.array( [float(term_cfg.weight) for term_cfg in self._term_cfgs], dtype=wp.float32, device=self.device ) + # per-term flags whether the weight is nonzero, read by the conditional nodes of recorded graphs + self._term_active_wp = wp.array( + [int(term_cfg.weight != 0.0) for term_cfg in self._term_cfgs], dtype=wp.int32, device=self.device + ) + self._term_active_views_wp = [self._term_active_wp[index : index + 1] for index in range(num_terms)] + # without conditional nodes, a recorded graph keeps or skips a term by its weight at record time + self._conditional_graphs = wp.get_device(self.device).is_cuda and wp.is_conditional_graph_supported() # persistent reset-time logging buffers (warp buffers) self._episode_sum_avg_wp = wp.zeros((num_terms,), dtype=wp.float32, device=self.device) @@ -189,6 +196,7 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv): self._reward_tensor_view = wp.to_torch(self._reward_wp) self._step_reward_tensor_view = wp.to_torch(self._step_reward_wp) self._term_weights_tensor_view = wp.to_torch(self._term_weights_wp) + self._term_active_tensor_view = wp.to_torch(self._term_active_wp) self._episode_sum_avg_tensor_view = wp.to_torch(self._episode_sum_avg_wp) self._reset_extras = { "Episode_Reward/" + term_name: self._episode_sum_avg_tensor_view[term_idx] @@ -338,8 +346,14 @@ def _reset_step_buffers(self, dt: float) -> None: 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] + if self._conditional_graphs and wp.get_device(self.device).is_capturing: + # the graph checks the weight on every launch, so a weight set later needs no new recording + wp.capture_if( + self._term_active_views_wp[index], + on_true=partial(term_cfg.func, self._env, term_cfg.out, **term_cfg.params), + ) # skip if weight is zero (kind of a micro-optimization) - if term_cfg.weight != 0.0: + elif term_cfg.weight != 0.0: term_cfg.func(self._env, term_cfg.out, **term_cfg.params) def _finalize_step(self, dt: float) -> torch.Tensor: @@ -367,6 +381,9 @@ def _finalize_step(self, dt: float) -> torch.Tensor: def set_term_cfg(self, term_name: str, cfg: RewardTermCfg): """Sets the configuration of the specified term into the manager. + Recorded reward stages read the weight on every call. They record again on their next call only + when the term's function, parameters or capturability change. + Args: term_name: The name of the reward term. cfg: The configuration for the reward term. @@ -381,11 +398,14 @@ def set_term_cfg(self, term_name: str, cfg: RewardTermCfg): # set the configuration (preserve the pre-allocated output view) term_idx = self._term_names.index(term_name) cfg.out = self._term_out_views_wp[term_idx] + changed = self._term_cfg_changed(term_name, cfg) self._term_cfgs[term_idx] = cfg + self._class_term_cfgs = [term_cfg for term_cfg in self._term_cfgs if isinstance(term_cfg.func, ManagerTermBase)] # 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() + self._term_active_tensor_view[term_idx] = int(cfg.weight != 0.0) + if changed: + self._clear_stage_steps() def get_term_cfg(self, term_name: str) -> RewardTermCfg: """Gets the configuration for the specified term. @@ -425,6 +445,15 @@ def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequenc Helper functions. """ + def _named_term_cfgs(self) -> Iterable[tuple[str, RewardTermCfg]]: + return zip(self._term_names, self._term_cfgs) + + def _term_signature(self, term_cfg: RewardTermCfg) -> Any: + signature = super()._term_signature(term_cfg) + if self._conditional_graphs: + return signature + return signature, term_cfg.weight == 0.0 + def _prepare_terms(self): # check if config is dict already if isinstance(self.cfg, dict): diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py index 1c3b61c473c1..d06f56114071 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py @@ -15,7 +15,7 @@ from __future__ import annotations -from collections.abc import Callable, Sequence +from collections.abc import Callable, Iterable, Sequence from functools import partial from typing import TYPE_CHECKING, Any @@ -360,8 +360,8 @@ def _finalize_step(self) -> torch.Tensor: def set_term_cfg(self, term_name: str, cfg: TerminationTermCfg): """Sets the configuration of the specified term into the manager. - Recorded termination stages read the term's configuration when they record, so they record again on their - next call. + Recorded termination stages read whether the term is a time-out on every call. They record again on their + next call only when the term's function, parameters or capturability change. Args: term_name: The name of the termination term. @@ -375,10 +375,12 @@ def set_term_cfg(self, term_name: str, cfg: TerminationTermCfg): term_idx = self._term_name_to_term_idx[term_name] # the term keeps writing into its column of the done buffer cfg.out = self._term_out_views_wp[term_idx] + changed = self._term_cfg_changed(term_name, cfg) self._term_cfgs[term_idx] = cfg self._class_term_cfgs = [term_cfg for term_cfg in self._term_cfgs if isinstance(term_cfg.func, ManagerTermBase)] wp.to_torch(self._term_is_time_out_wp)[term_idx] = bool(cfg.time_out) - self._clear_stage_steps() + if changed: + self._clear_stage_steps() def get_term_cfg(self, term_name: str) -> TerminationTermCfg: """Gets the configuration for the specified term. @@ -415,6 +417,9 @@ def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequenc Helper functions. """ + def _named_term_cfgs(self) -> Iterable[tuple[str, TerminationTermCfg]]: + return zip(self._term_names, self._term_cfgs) + def _prepare_terms(self): # check if config is dict already if isinstance(self.cfg, dict): diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_curriculum.py b/source/isaaclab_experimental/test/envs/test_warp_env_curriculum.py new file mode 100644 index 000000000000..410f5ced6e01 --- /dev/null +++ b/source/isaaclab_experimental/test/envs/test_warp_env_curriculum.py @@ -0,0 +1,85 @@ +# 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 + +"""Curricula of Warp 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()))) + +from collections import Counter + +import torch +from isaaclab_experimental.envs.frontend import WarpFrontend +from isaaclab_experimental.utils import WarpGraphCache + +import isaaclab.sim as sim_utils + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.utils import resolve_task_config + +_TASK = "Isaac-Reach-Franka" +_NUM_ENVS = 16 + + +def test_unchanged_term_settings_keep_the_recorded_stages(monkeypatch): + """Curricula set their terms on every reset; a stage records again only when a setting it read changes. + + Past its threshold, Reach's ``modify_reward_weight`` raises two penalties on every reset (episodes of 15 + steps), and the recorded reward stage reads weights on the device. Setting every term of every manager to + its own configuration records nothing either. + """ + recordings = Counter() + call_steps = WarpGraphCache.call_steps + + def counting_call_steps(self, stage, *args, **kwargs): + recorded = set(self.captured_stages) + result = call_steps(self, stage, *args, **kwargs) + recordings[stage] += len(set(self.captured_stages) - recorded) + return result + + monkeypatch.setattr(WarpGraphCache, "call_steps", counting_call_steps) + env_cfg, _ = resolve_task_config(_TASK, "", overrides=("physics=newton_mjwarp",)) + env_cfg.scene.num_envs = _NUM_ENVS + env_cfg.seed = 3 + env_cfg.episode_length_s = 0.5 + sim_utils.create_new_stage() + env = WarpFrontend.build_env(env_cfg, _TASK).unwrapped + try: + env.reset() + env.common_step_counter = 5000 + actions = torch.zeros((_NUM_ENVS, env.action_space.shape[-1]), device=env.device) + + def step(num_steps: int) -> int: + reset_steps = 0 + for _ in range(num_steps): + _, _, terminated, truncated, _ = env.step(actions) + reset_steps += int(bool((terminated | truncated).any())) + return reset_steps + + assert step(60) >= 3 + assert env.reward_manager.get_term_cfg("action_rate").weight == -0.005 + assert recordings["RewardManager_compute"] == 1 + + recorded = dict(recordings) + term_names = { + env.reward_manager: env.reward_manager.active_terms, + env.termination_manager: env.termination_manager.active_terms, + env.command_manager: env.command_manager.active_terms, + env.event_manager: [name for names in env.event_manager.active_terms.values() for name in names], + env.observation_manager: [ + f"{group}/{name}" for group, names in env.observation_manager.active_terms.items() for name in names + ], + } + for manager, names in term_names.items(): + for name in names: + manager.set_term_cfg(name, manager.get_term_cfg(name)) + assert step(20) >= 1 + assert dict(recordings) == recorded + finally: + env.close() diff --git a/source/isaaclab_experimental/test/managers/test_command_manager.py b/source/isaaclab_experimental/test/managers/test_command_manager.py index 2de666a71497..3a1296e21ba1 100644 --- a/source/isaaclab_experimental/test/managers/test_command_manager.py +++ b/source/isaaclab_experimental/test/managers/test_command_manager.py @@ -28,10 +28,18 @@ DT = 0.02 +class CountingVelocityCommand(UniformVelocityCommand): + """Velocity command counting the runs of its Python; a recorded stage runs it once per recording.""" + + def compute(self, dt: float): + self._env.recordings += 1 + super().compute(dt) + + @configclass class CommandsCfg: base_velocity = mdp.UniformVelocityCommandCfg( - class_type=UniformVelocityCommand, + class_type=CountingVelocityCommand, asset_name="robot", # shorter than a step, so every step resamples resampling_time_range=(DT / 2, DT / 2), @@ -60,6 +68,7 @@ def _env(graph_cache: WarpGraphCache) -> SimpleNamespace: return SimpleNamespace( num_envs=NUM_ENVS, device=DEVICE, + recordings=0, scene={"robot": robot}, rng_state_wp=wp.array(np.arange(NUM_ENVS, dtype=np.uint32), device=DEVICE), resolve_env_mask=lambda env_ids=None, env_mask=None: resolve_1d_mask( @@ -77,7 +86,8 @@ def test_modify_term_cfg_applies_to_the_recorded_command_stage(): """A curriculum that changes a command range after the command stage recorded changes the next commands. The range is a kernel argument read while the stage records, so a replay without recording again keeps - sampling from the old range. + sampling from the old range. Writing the same range again, as a curriculum does on every reset past its + threshold, records nothing new. """ graph_cache = WarpGraphCache(DEVICE) env = _env(graph_cache) @@ -91,10 +101,12 @@ def test_modify_term_cfg_applies_to_the_recorded_command_stage(): graph_cache.arm() graph_cache.call_steps("CommandManager_compute", env.command_manager.stage_steps("compute"), dt=DT) - curriculum(env, None, **params) - graph_cache.call_steps("CommandManager_compute", env.command_manager.stage_steps("compute"), dt=DT) + for _ in range(3): + curriculum(env, None, **params) + graph_cache.call_steps("CommandManager_compute", env.command_manager.stage_steps("compute"), dt=DT) wp.synchronize() command = env.command_manager.get_command("base_velocity") assert torch.equal(command[:, 0], torch.full((NUM_ENVS,), 1.5, device=DEVICE)) + assert env.recordings == 2 graph_cache.close() diff --git a/source/isaaclab_experimental/test/managers/test_event_manager.py b/source/isaaclab_experimental/test/managers/test_event_manager.py index ae052362cdb7..89718b8050f1 100644 --- a/source/isaaclab_experimental/test/managers/test_event_manager.py +++ b/source/isaaclab_experimental/test/managers/test_event_manager.py @@ -29,6 +29,12 @@ def device_term(env, env_mask: wp.array): """Event term whose work is entirely on the device.""" +def select(env, env_mask: wp.array): + """Event term that stores which environments it was applied to.""" + env.recordings += 1 + wp.copy(env.selected, env_mask) + + 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) @@ -54,3 +60,44 @@ def test_startup_only_term_does_not_keep_the_reset_stage_eager(): assert graph_cache.captured_stages == ("EventManager_apply_reset",) graph_cache.close() + + +def test_set_term_cfg_records_the_reset_events_again_only_when_the_trigger_changes(): + """A curriculum sets its event term on every reset; only a changed trigger spacing records the stage again. + + The spacing is a kernel argument of the recorded stage, so the changed spacing applies only after a new + recording. Every recording runs the term's Python once. + """ + graph_cache = WarpGraphCache(DEVICE) + env = SimpleNamespace( + num_envs=NUM_ENVS, + device=DEVICE, + sim=SimpleNamespace(is_playing=lambda: True), + selected=wp.zeros(NUM_ENVS, dtype=wp.bool, device=DEVICE), + recordings=0, + _warp_graph_cache=graph_cache, + ) + manager = EventManager({"select": EventTermCfg(func=select, mode="reset")}, env) + env_mask = wp.ones(NUM_ENVS, dtype=wp.bool, device=DEVICE) + step_count = wp.zeros(1, dtype=wp.int32, device=DEVICE) + term_cfg = manager.get_term_cfg("select") + graph_cache.arm() + + def apply_reset(step: int, min_step_count_between_reset: int) -> list[bool]: + term_cfg.min_step_count_between_reset = min_step_count_between_reset + manager.set_term_cfg("select", term_cfg) + step_count.fill_(step) + graph_cache.call_steps( + "EventManager_apply_reset", + manager.stage_steps("apply_reset"), + env_mask_wp=env_mask, + global_env_step_count=step_count, + ) + return env.selected.numpy().tolist() + + assert apply_reset(step=5, min_step_count_between_reset=0) == [True] * NUM_ENVS + assert apply_reset(step=6, min_step_count_between_reset=0) == [True] * NUM_ENVS + assert env.recordings == 1 + assert apply_reset(step=7, min_step_count_between_reset=10) == [False] * NUM_ENVS + assert env.recordings == 2 + graph_cache.close() diff --git a/source/isaaclab_experimental/test/managers/test_observation_manager.py b/source/isaaclab_experimental/test/managers/test_observation_manager.py index dd17149e8421..1b0e67700122 100644 --- a/source/isaaclab_experimental/test/managers/test_observation_manager.py +++ b/source/isaaclab_experimental/test/managers/test_observation_manager.py @@ -42,11 +42,17 @@ def test_group_settings_are_not_parsed_as_terms(): assert manager.group_obs_dim["policy"] == (5,) +def counted_constant_obs(env, out: wp.array, out_dim: int): + """``constant_obs`` counting the runs of its Python; a recorded stage runs it once per recording.""" + env.recordings += 1 + out.fill_(1.0) + + @configclass class NoisyPolicyCfg(ObservationGroupCfg): """Group with one term corrupted by a constant bias.""" - value = ObservationTermCfg(func=constant_obs, params={"out_dim": 2}, noise=ConstantNoiseCfg(bias=0.0)) + value = ObservationTermCfg(func=counted_constant_obs, params={"out_dim": 2}, noise=ConstantNoiseCfg(bias=0.0)) def __post_init__(self): self.enable_corruption = True @@ -66,7 +72,8 @@ def test_modify_term_cfg_applies_noise_to_the_recorded_observation_stage(): """An observation noise curriculum changes the observations of a recorded stage, as noise ADR expects. The noise parameters are kernel arguments read while the stage records, so a replay without recording again - keeps applying the old noise. + keeps applying the old noise. Writing the same bias again, as a curriculum does on every reset past its + threshold, records nothing new. """ graph_cache = WarpGraphCache("cuda:0") env = SimpleNamespace( @@ -74,6 +81,7 @@ def test_modify_term_cfg_applies_noise_to_the_recorded_observation_stage(): device="cuda:0", sim=SimpleNamespace(is_playing=lambda: True), rng_state_wp=wp.zeros(4, dtype=wp.uint32, device="cuda:0"), + recordings=0, _warp_graph_cache=graph_cache, ) env.observation_manager = ObservationManager(ObservationsCfg(), env) @@ -89,11 +97,13 @@ def test_modify_term_cfg_applies_noise_to_the_recorded_observation_stage(): ) assert graph_cache.captured_stages == ("ObservationManager_compute",) - curriculum(env, None, **params) - obs = graph_cache.call_steps( - "ObservationManager_compute", env.observation_manager.stage_steps("compute"), return_cloned_output=False - ) + for _ in range(3): + curriculum(env, None, **params) + obs = graph_cache.call_steps( + "ObservationManager_compute", env.observation_manager.stage_steps("compute"), return_cloned_output=False + ) wp.synchronize() assert torch.equal(obs["policy"], torch.full((4, 2), 1.5, device="cuda:0")) + assert env.recordings == 2 graph_cache.close() diff --git a/source/isaaclab_experimental/test/managers/test_reward_manager.py b/source/isaaclab_experimental/test/managers/test_reward_manager.py index 3b13737ec7a7..e1f9fcfaf468 100644 --- a/source/isaaclab_experimental/test/managers/test_reward_manager.py +++ b/source/isaaclab_experimental/test/managers/test_reward_manager.py @@ -44,6 +44,17 @@ def __call__(self, env, out: wp.array, value: float, capturable: bool): out.fill_(value) +def _env(graph_cache: WarpGraphCache) -> SimpleNamespace: + """The environment surface the reward manager reads.""" + return SimpleNamespace( + num_envs=NUM_ENVS, + device=DEVICE, + sim=SimpleNamespace(is_playing=lambda: True), + max_episode_length_s=1.0, + _warp_graph_cache=graph_cache, + ) + + def test_capturability_is_decided_per_term_instance(): """The same term is recorded for one parameter set and runs eagerly for another. @@ -51,13 +62,6 @@ def test_capturability_is_decided_per_term_instance(): 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( @@ -66,7 +70,7 @@ def test_capturability_is_decided_per_term_instance(): "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) + manager = RewardManager(cfg, _env(graph_cache)) graph_cache.arm() CALLS.clear() @@ -79,28 +83,55 @@ def test_capturability_is_decided_per_term_instance(): graph_cache.close() -def test_set_term_cfg_applies_to_the_recorded_reward_stage(): - """Raising a weight from zero after the stage recorded adds the term, as ``modify_reward_weight`` expects. - - A zero-weight term is skipped while the stage records, so a replay without recording again never runs it. +@pytest.mark.parametrize( + "conditional_graphs, expected_recordings", + [(True, [1, 1, 1, 1, 2, 2]), (False, [1, 2, 2, 2, 3, 4])], + ids=["conditional_nodes", "no_conditional_nodes"], +) +def test_set_term_cfg_records_the_reward_stage_again_only_when_its_graph_changes( + monkeypatch, conditional_graphs, expected_recordings +): + """``modify_reward_weight`` sets its term on every reset past its threshold, changing the weight in place. + + The replayed reward matches the eager reward after every setting. Weight changes, including to and from zero, + need no new recording when the graph checks the weight on the device; without conditional graph nodes a zero + crossing records again. A parameter change records again once. Every recording runs the always-on term's + Python once, so its call count is the number of recordings. """ + if conditional_graphs and not wp.is_conditional_graph_supported(): + pytest.skip("Conditional graph nodes require CUDA 12.4+.") + if not conditional_graphs: + monkeypatch.setattr(wp, "is_conditional_graph_supported", lambda: False) 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, + manager = RewardManager( + { + "always": RewardTermCfg(func=constant_reward, weight=1.0, params={"value": 1.0}), + "late": RewardTermCfg(func=constant_reward, weight=0.0, params={"value": 3.0}), + }, + _env(graph_cache), ) - manager = RewardManager({"late": RewardTermCfg(func=constant_reward, weight=0.0, params={"value": 3.0})}, env) - graph_cache.arm() - graph_cache.call_steps("RewardManager_compute", manager.stage_steps("compute"), dt=0.5) + settings = [(0.0, 3.0), (2.0, 3.0), (2.0, 3.0), (4.0, 3.0), (4.0, 5.0), (0.0, 5.0)] + late = manager.get_term_cfg("late") + + def run(compute) -> tuple[list[torch.Tensor], list[int]]: + rewards, calls = [], [] + for weight, value in settings: + late.weight = weight + late.params["value"] = value + manager.set_term_cfg("late", late) + rewards.append(compute().clone()) + calls.append(CALLS[1.0]) + return rewards, calls - term_cfg = manager.get_term_cfg("late") - term_cfg.weight = 2.0 - manager.set_term_cfg("late", term_cfg) - reward = graph_cache.call_steps("RewardManager_compute", manager.stage_steps("compute"), dt=0.5) - wp.synchronize() + graph_cache.arm() + CALLS.clear() + replayed, recordings = run( + lambda: graph_cache.call_steps("RewardManager_compute", manager.stage_steps("compute"), dt=0.5) + ) + eager, _ = run(lambda: manager.compute(dt=0.5)) - assert torch.equal(reward, torch.full((NUM_ENVS,), 3.0 * 2.0 * 0.5, device=DEVICE)) + assert recordings == expected_recordings + for (weight, value), replayed_reward, eager_reward in zip(settings, replayed, eager): + assert torch.equal(replayed_reward, eager_reward) + assert torch.equal(replayed_reward, torch.full((NUM_ENVS,), (1.0 + weight * value) * 0.5, device=DEVICE)) graph_cache.close() From 2bb8e9e3068049308e52c71ec6943bd7b9896f69 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 1 Oct 2026 00:37:32 -0700 Subject: [PATCH 27/80] 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 Date: Thu, 1 Oct 2026 00:46:22 -0700 Subject: [PATCH 28/80] 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 Date: Thu, 1 Oct 2026 00:46:59 -0700 Subject: [PATCH 29/80] 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 Date: Thu, 1 Oct 2026 00:48:22 -0700 Subject: [PATCH 30/80] 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 Date: Thu, 1 Oct 2026 00:48:56 -0700 Subject: [PATCH 31/80] 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 Date: Thu, 1 Oct 2026 00:54:23 -0700 Subject: [PATCH 32/80] 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 Date: Thu, 1 Oct 2026 00:55:06 -0700 Subject: [PATCH 33/80] 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 Date: Thu, 1 Oct 2026 11:47:31 -0700 Subject: [PATCH 34/80] 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 Date: Thu, 1 Oct 2026 11:47:31 -0700 Subject: [PATCH 35/80] 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 Date: Thu, 1 Oct 2026 11:47:31 -0700 Subject: [PATCH 36/80] 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 Date: Thu, 1 Oct 2026 11:50:15 -0700 Subject: [PATCH 37/80] 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 ` (``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 Date: Thu, 1 Oct 2026 11:54:09 -0700 Subject: [PATCH 38/80] 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 Date: Thu, 1 Oct 2026 12:23:28 -0700 Subject: [PATCH 39/80] 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 Date: Thu, 1 Oct 2026 12:23:29 -0700 Subject: [PATCH 40/80] 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 Date: Fri, 2 Oct 2026 15:57:44 -0700 Subject: [PATCH 41/80] 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 ` 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 `, 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 ` 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 ` 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 ` + 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 ``"[]"`` 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 bd5a5a05d25a9949e5bf614505511454557053d9 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Sun, 4 Oct 2026 02:29:03 -0700 Subject: [PATCH 42/80] Record Warp stages again only when a set term configuration changes Port Part 3's change detection onto the per-call capture. Each manager takes a signature of every term configuration when its terms are prepared (and again when scene entities resolve at play); set_term_cfg compares against it and calls reset_captured_stages only when the recorded stages would read something new: the function, parameters (by value) or capturability, an event term's trigger settings, an observation or command term's whole configuration, or whether a reward weight is zero. Curricula that set an unchanged configuration on every reset, such as modify_reward_weight past its threshold, keep the recorded stages. Zero weights stay in Part 2's style: the reward stage leaves zero-weight terms out and re-splits by weight, so only a weight turning zero or nonzero records again; other weights are read on the device. This drops Part 3's wp.capture_if path and its CUDA 12.4 fallback. --- ...ichuanh-warp-frontend-managers.changed.rst | 9 +- .../managers/command_manager.py | 20 +++-- .../managers/event_manager.py | 26 +++++- .../managers/manager_base.py | 61 +++++++++++++- .../managers/observation_manager.py | 45 ++++++---- .../managers/reward_manager.py | 24 ++++-- .../managers/termination_manager.py | 15 ++-- .../test/envs/test_warp_env_curriculum.py | 83 +++++++++++++++++++ .../test/managers/test_command_manager.py | 1 + .../test/managers/test_event_manager.py | 39 +++++++++ .../test/managers/test_observation_manager.py | 1 + .../test/managers/test_reward_manager.py | 43 ++++++++++ 12 files changed, 325 insertions(+), 42 deletions(-) create mode 100644 source/isaaclab_experimental/test/envs/test_warp_env_curriculum.py diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.changed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.changed.rst index 7f2cd0f3b24e..8e1bd09d94a4 100644 --- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.changed.rst +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-managers.changed.rst @@ -1,6 +1,5 @@ * Changed ``set_term_cfg`` of the Warp reward, event, termination and command managers to record the manager's - stages again only when the term's function, parameters or capturability change, or an event term's trigger - settings. Setting an unchanged configuration, as curricula do on every reset, keeps the recorded stages. -* Changed the Warp reward manager to check zero term weights inside the recorded reward stage, so a weight set - to or from zero applies without recording the stage again. This needs CUDA 12.4 or newer; with older CUDA - versions a weight crossing zero records the stage again. + stages again only when the term's function, parameters or capturability change, an event term's trigger + settings change, or a reward term's weight turns zero or nonzero. Setting an unchanged configuration, as + curricula do on every reset, keeps the recorded stages, and other weight changes apply to the recorded reward + stage, which reads the weights on the device. diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py index ddcbb2a6077d..408cf97fec86 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py @@ -9,8 +9,8 @@ import inspect from abc import abstractmethod -from collections.abc import Sequence -from typing import TYPE_CHECKING +from collections.abc import Iterable, Sequence +from typing import TYPE_CHECKING, Any import torch import warp as wp @@ -22,9 +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 isaaclab_experimental.utils.warp_capture import captured, eager -from .manager_base import ManagerBase, ManagerTermBase, split_terms +from .manager_base import ManagerBase, ManagerTermBase, _config_key, split_terms # import omni.kit.app @@ -559,7 +559,7 @@ def set_term_cfg(self, name: str, cfg: CommandTermCfg): """Sets the configuration of the specified command term. A command term reads its configuration, such as its sampling ranges, when a command stage records, so the - command stages record again on their next call. + command stages record again on their next call when the configuration changes. Args: name: The name of the command term. @@ -570,12 +570,14 @@ def set_term_cfg(self, name: str, cfg: CommandTermCfg): """ if name not in self._terms: raise ValueError(f"Command term '{name}' not found.") + changed = self._term_cfg_changed(name, cfg) self._terms[name].cfg = cfg if isinstance(self.cfg, dict): self.cfg[name] = cfg else: setattr(self.cfg, name, cfg) - reset_captured_stages(self) + if changed: + self._record_term_again(name, cfg) def get_term_cfg(self, name: str) -> CommandTermCfg: """Gets the configuration of the specified command term. @@ -608,6 +610,12 @@ def get_term(self, name: str) -> CommandTerm: Helper functions. """ + def _named_term_cfgs(self) -> Iterable[tuple[str, CommandTermCfg]]: + return ((name, term.cfg) for name, term in self._terms.items()) + + def _term_signature(self, term_cfg: CommandTermCfg) -> Any: + return _config_key(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/event_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py index 7b959c9b4521..aed79e88d38e 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py @@ -23,7 +23,8 @@ import inspect import logging -from collections.abc import Sequence +from collections.abc import Iterable, Sequence +from typing import Any import torch import warp as wp @@ -33,7 +34,7 @@ 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 isaaclab_experimental.utils.warp_capture import captured, eager from .manager_base import ManagerBase, ManagerTermBase, split_resets @@ -232,6 +233,9 @@ def set_term_cfg(self, term_name: str, cfg: EventTermCfg): 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. + Recorded event stages record again on their next call only when the term's function, parameters, + capturability or trigger settings change. + Args: term_name: The name of the event term. cfg: The configuration for the event term. @@ -242,6 +246,7 @@ def set_term_cfg(self, term_name: str, cfg: EventTermCfg): term_found = False for mode, terms in self._mode_term_names.items(): if term_name in terms: + changed = self._term_cfg_changed(term_name, cfg) cfg.capturable = is_warp_capturable(cfg.func, cfg.params) self._mode_term_cfgs[mode][terms.index(term_name)] = cfg term_found = True @@ -250,7 +255,8 @@ def set_term_cfg(self, term_name: str, cfg: EventTermCfg): 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) + if changed: + self._record_term_again(term_name, cfg) def get_term_cfg(self, term_name: str) -> EventTermCfg: for mode, terms in self._mode_term_names.items(): @@ -405,6 +411,20 @@ def _apply_reset(self, env_mask_wp: wp.array, global_env_step_count_wp: wp.array else: eager(term_cfg.func, self._env, self._scratch_term_mask_wp, **term_cfg.params) + def _named_term_cfgs(self) -> Iterable[tuple[str, EventTermCfg]]: + for mode, names in self._mode_term_names.items(): + yield from zip(names, self._mode_term_cfgs[mode]) + + def _term_signature(self, term_cfg: EventTermCfg) -> Any: + # the trigger settings are kernel arguments and decide which interval timers reset + trigger = ( + term_cfg.interval_range_s, + term_cfg.is_global_time, + term_cfg.min_step_count_between_reset, + term_cfg.resample_interval_on_reset, + ) + return super()._term_signature(term_cfg), trigger + 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 = { diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py index 36f5cd4a033c..191c835bd17d 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py @@ -17,12 +17,15 @@ import contextlib import copy +import dataclasses import inspect import logging from abc import ABC, abstractmethod -from collections.abc import Callable, Sequence +from collections.abc import Callable, Iterable, Sequence from typing import TYPE_CHECKING, Any, NamedTuple +import numpy as np +import torch import warp as wp import isaaclab.utils.string as string_utils @@ -39,8 +42,6 @@ if TYPE_CHECKING: - import torch - from isaaclab.envs import ManagerBasedEnv # import logger @@ -232,9 +233,12 @@ def __init__(self, cfg: object, env: ManagerBasedEnv): # self._resolve_terms_handle = None self._resolve_terms_handle = None + # what the recorded stages read from each term configuration, compared when a term configuration is set + self._term_signatures: dict[str, Any] = {} # parse config to create terms information if self.cfg: self._prepare_terms() + self._take_term_signatures() def __del__(self): """Delete the manager.""" @@ -349,6 +353,35 @@ 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 {} + def _named_term_cfgs(self) -> Iterable[tuple[str, Any]]: + """Name and configuration of each term whose configuration can be set.""" + return () + + def _term_signature(self, term_cfg: ManagerTermBaseCfg) -> Any: + """What the recorded stages read from a term configuration when they record. + + Values the stages read from device arrays on every call are left out, since a replayed graph sees their + changes. + """ + return term_cfg.func, _config_key(term_cfg.params), is_warp_capturable(term_cfg.func, term_cfg.params) + + def _take_term_signatures(self) -> None: + """Take the signature of every term configuration, as the stages recorded from now on read them.""" + self._term_signatures = {name: self._term_signature(cfg) for name, cfg in self._named_term_cfgs()} + + def _term_cfg_changed(self, term_name: str, term_cfg: Any) -> bool: + """Whether the recorded stages read a different configuration of the term. + + Compares against the signature taken when the term was prepared or last changed, since curricula change a + stored configuration in place before setting it again. + """ + return self._term_signatures.get(term_name) != self._term_signature(term_cfg) + + def _record_term_again(self, term_name: str, term_cfg: Any) -> None: + """Take the signature of a changed term configuration and drop the recordings that read the previous one.""" + self._term_signatures[term_name] = self._term_signature(term_cfg) + reset_captured_stages(self) + """ Internal callbacks. """ @@ -379,6 +412,7 @@ def _resolve_terms_callback(self, event): # set the flag self._is_scene_entities_resolved = True # class terms were just created + self._take_term_signatures() reset_captured_stages(self) """ @@ -509,3 +543,24 @@ def _process_term_cfg_at_play(self, term_name: str, term_cfg: ManagerTermBaseCfg if inspect.isclass(term_cfg.func): logger.info(f"Initializing term '{term_name}' with class '{term_cfg.func.__name__}'.") term_cfg.func = term_cfg.func(cfg=term_cfg, env=self._env) + + +def _config_key(value: Any) -> Any: + """Comparable copy of a configuration value, taken by value. + + Arrays are keyed by the memory a recorded graph reads, configuration objects by their attributes + (including those set when they are resolved, such as the Warp joint mask), and other objects by identity. + """ + if isinstance(value, wp.array): + return ("wp.array", value.ptr, value.shape, value.strides, value.dtype) + if isinstance(value, torch.Tensor): + return ("tensor", value.data_ptr(), tuple(value.shape), value.stride(), value.dtype) + if isinstance(value, np.ndarray): + return ("ndarray", value.shape, value.dtype.str, value.tobytes()) + if isinstance(value, dict): + return ("dict", tuple((key, _config_key(item)) for key, item in value.items())) + if isinstance(value, (list, tuple)): + return (type(value), tuple(_config_key(item) for item in value)) + if dataclasses.is_dataclass(value) and not isinstance(value, type): + return (type(value), _config_key(vars(value))) + return value diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py index d40fac9287cc..e76a3d7cdda2 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py @@ -47,7 +47,7 @@ from __future__ import annotations import inspect -from collections.abc import Sequence +from collections.abc import Iterable, Sequence from typing import TYPE_CHECKING, Any import numpy as np @@ -62,9 +62,10 @@ 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_capture import captured, eager, reset_captured_stages +from isaaclab_experimental.utils.warp import is_warp_capturable +from isaaclab_experimental.utils.warp_capture import captured, eager -from .manager_base import ManagerBase, ManagerTermBase, split_resets, split_terms +from .manager_base import ManagerBase, ManagerTermBase, _config_key, split_resets, split_terms if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv @@ -584,7 +585,8 @@ def set_term_cfg(self, term_name: str, cfg: ObservationTermCfg): """Sets the configuration of the specified term into the manager. Recorded observation stages read the term's configuration, such as its noise, when they record, so they - record again on their next call. The term keeps writing into its slice of the group's observation. + record again on their next call when the configuration changes. The term keeps writing into its slice of + the group's observation. Args: term_name: The name of the observation term, in the form ``"/"``. @@ -595,23 +597,27 @@ def set_term_cfg(self, term_name: str, cfg: ObservationTermCfg): """ group_name, index = self._find_term(term_name) old_cfg = self._group_obs_term_cfgs[group_name][index] - group_cfg = self.cfg[group_name] if isinstance(self.cfg, dict) else getattr(self.cfg, group_name) - term_dim = self._prepare_term(group_name, self._group_obs_term_names[group_name][index], cfg, group_cfg) - if term_dim != old_cfg._term_dim: - raise ValueError( - f"Observation term '{term_name}' must keep its output width {old_cfg._term_dim}, got {term_dim}." - ) - cfg.out_wp, cfg.out_torch = old_cfg.out_wp, old_cfg.out_torch - if cfg is not old_cfg and isinstance(cfg.func, ManagerTermBase): - cfg.func.reset() + # an unchanged configuration was prepared already; preparing it again reallocates its Warp buffers + changed = self._term_cfg_changed(term_name, cfg) + if changed: + group_cfg = self.cfg[group_name] if isinstance(self.cfg, dict) else getattr(self.cfg, group_name) + term_dim = self._prepare_term(group_name, self._group_obs_term_names[group_name][index], cfg, group_cfg) + if term_dim != old_cfg._term_dim: + raise ValueError( + f"Observation term '{term_name}' must keep its output width {old_cfg._term_dim}, got {term_dim}." + ) + cfg.out_wp, cfg.out_torch = old_cfg.out_wp, old_cfg.out_torch + if cfg is not old_cfg and isinstance(cfg.func, ManagerTermBase): + cfg.func.reset() self._group_obs_term_cfgs[group_name][index] = cfg self._group_obs_class_term_cfgs[group_name] = [ term_cfg for term_cfg in self._group_obs_term_cfgs[group_name] if isinstance(term_cfg.func, ManagerTermBase) ] self._group_obs_class_instances = self._collect_class_instances() - # recorded stages hold the replaced term, which may also differ in capturability self._split_terms() - reset_captured_stages(self) + # recorded stages hold the replaced term, which may also differ in capturability + if changed: + self._record_term_again(term_name, cfg) def get_term_cfg(self, term_name: str) -> ObservationTermCfg: """Gets the configuration for the specified term. @@ -662,6 +668,15 @@ def serialize(self) -> dict: Helper functions. """ + def _named_term_cfgs(self) -> Iterable[tuple[str, ObservationTermCfg]]: + for group_name, term_names in self._group_obs_term_names.items(): + for term_name, term_cfg in zip(term_names, self._group_obs_term_cfgs[group_name]): + yield f"{group_name}/{term_name}", term_cfg + + def _term_signature(self, term_cfg: ObservationTermCfg) -> Any: + # noise, clip, scale, modifiers and history are read while the stage records + return _config_key(term_cfg), is_warp_capturable(term_cfg.func, term_cfg.params) + def _split_terms(self) -> None: """Split the terms, as (group, index) pairs, and the class term resets by whether they can be recorded.""" terms = [ diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py index 381bad1e0386..0ae0f5264cfd 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py @@ -11,8 +11,8 @@ from __future__ import annotations -from collections.abc import Sequence -from typing import TYPE_CHECKING +from collections.abc import Iterable, Sequence +from typing import TYPE_CHECKING, Any import torch import warp as wp @@ -22,7 +22,7 @@ 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 isaaclab_experimental.utils.warp_capture import captured, eager from .manager_base import ManagerBase, ManagerTermBase, split_resets, split_terms @@ -303,6 +303,9 @@ def compute(self, dt: float) -> torch.Tensor: def set_term_cfg(self, term_name: str, cfg: RewardTermCfg): """Sets the configuration of the specified term into the manager. + Recorded reward stages read the weights on the device on every call. They record again on their next call + only when the term's function, parameters or capturability change, or its weight turns zero or nonzero. + Args: term_name: The name of the reward term. cfg: The configuration for the reward term. @@ -317,13 +320,17 @@ def set_term_cfg(self, term_name: str, cfg: RewardTermCfg): # set the configuration (preserve the pre-allocated output view) term_idx = self._term_names.index(term_name) cfg.out = self._term_out_views_wp[term_idx] + changed = self._term_cfg_changed(term_name, cfg) 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) - # 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) + self._class_term_cfgs = [term_cfg for term_cfg in self._term_cfgs if isinstance(term_cfg.func, ManagerTermBase)] + self._class_term_split = split_resets(self._class_term_cfgs) + # recorded stages hold the replaced term, which may also differ in capturability or in whether it runs + if changed: + self._record_term_again(term_name, cfg) def get_term_cfg(self, term_name: str) -> RewardTermCfg: """Gets the configuration for the specified term. @@ -363,6 +370,13 @@ def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequenc Helper functions. """ + def _named_term_cfgs(self) -> Iterable[tuple[str, RewardTermCfg]]: + return zip(self._term_names, self._term_cfgs) + + def _term_signature(self, term_cfg: RewardTermCfg) -> Any: + # a zero weight leaves the term out of the recorded stage + return super()._term_signature(term_cfg), term_cfg.weight == 0.0 + def _prepare_terms(self): # check if config is dict already if isinstance(self.cfg, dict): diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py index d16c474c30fa..30063cc43b8b 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py @@ -15,7 +15,7 @@ from __future__ import annotations -from collections.abc import Sequence +from collections.abc import Iterable, Sequence from typing import TYPE_CHECKING import torch @@ -25,7 +25,7 @@ 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, reset_captured_stages +from isaaclab_experimental.utils.warp_capture import captured, eager from .manager_base import ManagerBase, ManagerTermBase, split_resets, split_terms @@ -294,7 +294,7 @@ def set_term_cfg(self, term_name: str, cfg: TerminationTermCfg): """Sets the configuration of the specified term into the manager. Recorded termination stages read whether the term is a time-out on every call. They record again on their - next call, since they hold the replaced term. + next call only when the term's function, parameters or capturability change. Args: term_name: The name of the termination term. @@ -308,14 +308,16 @@ def set_term_cfg(self, term_name: str, cfg: TerminationTermCfg): term_idx = self._term_name_to_term_idx[term_name] # the term keeps writing into its column of the done buffer cfg.out = self._term_out_views_wp[term_idx] + changed = self._term_cfg_changed(term_name, cfg) cfg.capturable = is_warp_capturable(cfg.func, cfg.params) self._term_cfgs[term_idx] = cfg wp.to_torch(self._term_is_time_out_wp)[term_idx] = bool(cfg.time_out) - # recorded stages hold the replaced term, which may also differ in capturability self._class_term_cfgs = [term_cfg for term_cfg in self._term_cfgs if isinstance(term_cfg.func, ManagerTermBase)] self._term_split = split_terms(self._term_cfgs) self._class_term_split = split_resets(self._class_term_cfgs) - reset_captured_stages(self) + # recorded stages hold the replaced term, which may also differ in capturability + if changed: + self._record_term_again(term_name, cfg) def get_term_cfg(self, term_name: str) -> TerminationTermCfg: """Gets the configuration for the specified term. @@ -352,6 +354,9 @@ def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequenc Helper functions. """ + def _named_term_cfgs(self) -> Iterable[tuple[str, TerminationTermCfg]]: + return zip(self._term_names, self._term_cfgs) + def _prepare_terms(self): # check if config is dict already if isinstance(self.cfg, dict): diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_curriculum.py b/source/isaaclab_experimental/test/envs/test_warp_env_curriculum.py new file mode 100644 index 000000000000..3afc247e688d --- /dev/null +++ b/source/isaaclab_experimental/test/envs/test_warp_env_curriculum.py @@ -0,0 +1,83 @@ +# 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 + +"""Curricula of Warp 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()))) + +from collections import Counter + +import torch +from isaaclab_experimental.envs.frontend import WarpFrontend +from isaaclab_experimental.utils import CapturedStage + +import isaaclab.sim as sim_utils + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.utils import resolve_task_config + +_TASK = "Isaac-Reach-Franka" +_NUM_ENVS = 16 + + +def test_unchanged_term_settings_keep_the_recorded_stages(monkeypatch): + """Curricula set their terms on every reset; a stage records again only when a setting it read changes. + + Past its threshold, Reach's ``modify_reward_weight`` raises two penalties on every reset (episodes of 15 + steps), and the recorded reward stage reads weights on the device. Setting every term of every manager to + its own configuration records nothing either. + """ + recordings = Counter() + record = CapturedStage._record + + def counting_record(self, args, kwargs): + recordings[self._fn.__qualname__] += 1 + return record(self, args, kwargs) + + monkeypatch.setattr(CapturedStage, "_record", counting_record) + env_cfg, _ = resolve_task_config(_TASK, "", overrides=("physics=newton_mjwarp",)) + env_cfg.scene.num_envs = _NUM_ENVS + env_cfg.seed = 3 + env_cfg.episode_length_s = 0.5 + sim_utils.create_new_stage() + env = WarpFrontend.build_env(env_cfg, _TASK).unwrapped + try: + env.reset() + env.common_step_counter = 5000 + actions = torch.zeros((_NUM_ENVS, env.action_space.shape[-1]), device=env.device) + + def step(num_steps: int) -> int: + reset_steps = 0 + for _ in range(num_steps): + _, _, terminated, truncated, _ = env.step(actions) + reset_steps += int(bool((terminated | truncated).any())) + return reset_steps + + assert step(60) >= 3 + assert env.reward_manager.get_term_cfg("action_rate").weight == -0.005 + assert recordings["RewardManager.compute"] == 1 + + recorded = dict(recordings) + term_names = { + env.reward_manager: env.reward_manager.active_terms, + env.termination_manager: env.termination_manager.active_terms, + env.command_manager: env.command_manager.active_terms, + env.event_manager: [name for names in env.event_manager.active_terms.values() for name in names], + env.observation_manager: [ + f"{group}/{name}" for group, names in env.observation_manager.active_terms.items() for name in names + ], + } + for manager, names in term_names.items(): + for name in names: + manager.set_term_cfg(name, manager.get_term_cfg(name)) + assert step(20) >= 1 + assert dict(recordings) == recorded + finally: + env.close() diff --git a/source/isaaclab_experimental/test/managers/test_command_manager.py b/source/isaaclab_experimental/test/managers/test_command_manager.py index 88c40d13a241..304ac03d25b9 100644 --- a/source/isaaclab_experimental/test/managers/test_command_manager.py +++ b/source/isaaclab_experimental/test/managers/test_command_manager.py @@ -107,3 +107,4 @@ def test_modify_term_cfg_applies_to_the_recorded_command_stage(monkeypatch): command = env.command_manager.get_command("base_velocity") assert torch.equal(command[:, 0], torch.full((NUM_ENVS,), 1.5, device=DEVICE)) + assert env.recordings == 2 diff --git a/source/isaaclab_experimental/test/managers/test_event_manager.py b/source/isaaclab_experimental/test/managers/test_event_manager.py index c1ee9b0b13df..af8173ff6c04 100644 --- a/source/isaaclab_experimental/test/managers/test_event_manager.py +++ b/source/isaaclab_experimental/test/managers/test_event_manager.py @@ -34,6 +34,12 @@ def device_term(env, env_mask: wp.array): CALLS["device_term"] += 1 +def select(env, env_mask: wp.array): + """Event term that stores which environments it was applied to.""" + CALLS["select"] += 1 + wp.copy(env.selected, env_mask) + + @WarpCapturable(False, reason="synchronizes with the host") class host_call_term(ManagerTermBase): """Class event term whose call cannot be recorded but whose reset can.""" @@ -97,3 +103,36 @@ def test_class_annotation_does_not_keep_the_reset_eager(monkeypatch): 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" + + +def test_set_term_cfg_records_the_reset_events_again_only_when_the_trigger_changes(monkeypatch): + """A curriculum sets its event term on every reset; only a changed trigger spacing records the stage again. + + The spacing is a kernel argument of the recorded stage, so the changed spacing applies only after a new + recording. Every recording runs the term's Python once. + """ + monkeypatch.setattr(CapturedStage, "enabled", True) + env = SimpleNamespace( + num_envs=NUM_ENVS, + device=DEVICE, + sim=SimpleNamespace(is_playing=lambda: True), + selected=wp.zeros(NUM_ENVS, dtype=wp.bool, device=DEVICE), + ) + manager = EventManager({"select": EventTermCfg(func=select, mode="reset")}, env) + env_mask = wp.ones(NUM_ENVS, dtype=wp.bool, device=DEVICE) + step_count = wp.zeros(1, dtype=wp.int32, device=DEVICE) + term_cfg = manager.get_term_cfg("select") + CALLS.clear() + + def apply_reset(step: int, min_step_count_between_reset: int) -> list[bool]: + term_cfg.min_step_count_between_reset = min_step_count_between_reset + manager.set_term_cfg("select", term_cfg) + step_count.fill_(step) + manager.apply(mode="reset", env_mask_wp=env_mask, global_env_step_count=step_count) + return env.selected.numpy().tolist() + + assert apply_reset(step=5, min_step_count_between_reset=0) == [True] * NUM_ENVS + assert apply_reset(step=6, min_step_count_between_reset=0) == [True] * NUM_ENVS + assert CALLS["select"] == 1 + assert apply_reset(step=7, min_step_count_between_reset=10) == [False] * NUM_ENVS + assert CALLS["select"] == 2 diff --git a/source/isaaclab_experimental/test/managers/test_observation_manager.py b/source/isaaclab_experimental/test/managers/test_observation_manager.py index 07d58a8ddfcf..2d0256a7199f 100644 --- a/source/isaaclab_experimental/test/managers/test_observation_manager.py +++ b/source/isaaclab_experimental/test/managers/test_observation_manager.py @@ -92,3 +92,4 @@ def test_modify_term_cfg_applies_noise_to_the_recorded_observation_stage(monkeyp wp.synchronize() assert torch.equal(obs["policy"], torch.full((4, 2), 1.5, device="cuda:0")) + assert env.recordings == 2 diff --git a/source/isaaclab_experimental/test/managers/test_reward_manager.py b/source/isaaclab_experimental/test/managers/test_reward_manager.py index 3901abb2c702..8c50a634751d 100644 --- a/source/isaaclab_experimental/test/managers/test_reward_manager.py +++ b/source/isaaclab_experimental/test/managers/test_reward_manager.py @@ -69,3 +69,46 @@ def test_capturability_is_decided_per_term_from_its_parameters(monkeypatch): # 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)) + + +def test_set_term_cfg_records_the_reward_stage_again_only_when_it_runs_a_different_term(monkeypatch): + """``modify_reward_weight`` sets its term on every reset past its threshold, changing the weight in place. + + The replayed reward matches the eager reward after every setting. A weight change needs no new recording, since + the stage reads the weights on the device; a weight turning zero or nonzero, or a parameter change, records the + stage again once. Every recording runs the always-on term's Python once, so its call count is the number of + recordings. + """ + env = SimpleNamespace( + num_envs=NUM_ENVS, device=DEVICE, sim=SimpleNamespace(is_playing=lambda: True), max_episode_length_s=1.0 + ) + manager = RewardManager( + { + "always": RewardTermCfg(func=constant_reward, weight=1.0, params={"value": 1.0}), + "late": RewardTermCfg(func=constant_reward, weight=0.0, params={"value": 3.0}), + }, + env, + ) + settings = [(0.0, 3.0), (2.0, 3.0), (2.0, 3.0), (4.0, 3.0), (4.0, 5.0), (0.0, 5.0)] + late = manager.get_term_cfg("late") + + def run() -> tuple[list[torch.Tensor], list[int]]: + rewards, calls = [], [] + for weight, value in settings: + late.weight = weight + late.params["value"] = value + manager.set_term_cfg("late", late) + rewards.append(manager.compute(dt=0.5).clone()) + calls.append(CALLS[1.0]) + return rewards, calls + + monkeypatch.setattr(CapturedStage, "enabled", True) + CALLS.clear() + replayed, recordings = run() + monkeypatch.setattr(CapturedStage, "enabled", False) + eager, _ = run() + + assert recordings == [1, 2, 2, 2, 3, 4] + for (weight, value), replayed_reward, eager_reward in zip(settings, replayed, eager): + assert torch.equal(replayed_reward, eager_reward) + assert torch.equal(replayed_reward, torch.full((NUM_ENVS,), (1.0 + weight * value) * 0.5, device=DEVICE)) From 76fb382019b5ef3a69853bb38df94354521de21b Mon Sep 17 00:00:00 2001 From: jichuanh Date: Sun, 4 Oct 2026 02:30:15 -0700 Subject: [PATCH 43/80] Document the Warp capture contract with the per-call capture names Describe stages as @captured manager methods recorded on their first call, the eager() placement of non-capturable terms, CapturedStage.enabled instead of the removed ISAACLAB_WARP_CAPTURE, WarpRng.state for command resampling, and zero reward weights recording the reward stage again. Drop the instance-level _warp_capturable option, which Part 2's per-term capturability does not read. --- .../concepts/warp_environment_migration.rst | 49 ++++++++++--------- docs/source/concepts/warp_environments.rst | 7 ++- 2 files changed, 30 insertions(+), 26 deletions(-) diff --git a/docs/source/concepts/warp_environment_migration.rst b/docs/source/concepts/warp_environment_migration.rst index bc6fa35b539b..d4b31fee0cd4 100644 --- a/docs/source/concepts/warp_environment_migration.rst +++ b/docs/source/concepts/warp_environment_migration.rst @@ -183,7 +183,8 @@ 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 -``env.rng_state_wp``; the command manager records ``compute`` and ``reset`` like the other stages. +:attr:`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. @@ -191,12 +192,12 @@ caching it, so a term never holds another environment's buffer. Capture Safety ^^^^^^^^^^^^^^ -Each manager stage of a Warp environment (for example ``RewardManager_compute`` or -``CommandManager_reset``) is recorded into a CUDA graph on its first call after the first -``step`` and replayed afterwards. 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 term runs once, while the stage records. +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: @@ -216,22 +217,26 @@ part of it: 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 a term's configuration (observation terms - are named ``"/"``), 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, and reward weights are read on the device, so a weight change - applies without recording again (before CUDA 12.4, a weight set to or from zero records again). 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. + ``set_term_cfg`` records their stages again when it changes what they read from a term's configuration + (observation terms are named ``"/"``), 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=...)``. It runs - eagerly between the recorded parts of its stage and 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=...)``; a class term - may instead set ``self._warp_capturable`` in ``__init__``. ``height_scan``, ``base_height_l2`` with a - height scanner and the rigid-body randomization events are examples. - -``ISAACLAB_WARP_CAPTURE=0`` runs every stage eagerly, and ``ISAACLAB_SYNC_DEBUG=1`` makes a hidden + 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 ` 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. diff --git a/docs/source/concepts/warp_environments.rst b/docs/source/concepts/warp_environments.rst index 6c86acb51eae..9a18c78735dc 100644 --- a/docs/source/concepts/warp_environments.rst +++ b/docs/source/concepts/warp_environments.rst @@ -267,10 +267,9 @@ specific to warp envs; for Newton physics maturity and specialist guides see 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 runs of its stage. Observation and event terms, which draw from the shared - per-environment random state, and action terms keep their configured order, so an eager term there - splits the stage into several graphs. A non-capturable command term keeps the whole command stage - eager. + 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. From 0fb73002cd2901dfba4917da31116e2912f3ecff Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 20 Oct 2026 09:52:11 -0700 Subject: [PATCH 44/80] 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 Date: Tue, 20 Oct 2026 09:59:05 -0700 Subject: [PATCH 45/80] 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 2030fe34401b16b7fb3bc73293ee0e725d6e0d22 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Thu, 22 Oct 2026 00:11:43 -0700 Subject: [PATCH 46/80] 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 ` (``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 Date: Thu, 22 Oct 2026 00:12:22 -0700 Subject: [PATCH 47/80] 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 ` (``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 Date: Mon, 26 Oct 2026 02:48:16 -0700 Subject: [PATCH 48/80] 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 Date: Mon, 26 Oct 2026 02:48:17 -0700 Subject: [PATCH 49/80] 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 Date: Mon, 26 Oct 2026 02:48:18 -0700 Subject: [PATCH 50/80] 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 Date: Mon, 26 Oct 2026 02:48:20 -0700 Subject: [PATCH 51/80] 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 Date: Mon, 26 Oct 2026 02:48:21 -0700 Subject: [PATCH 52/80] 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 Date: Mon, 26 Oct 2026 02:48:22 -0700 Subject: [PATCH 53/80] 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 Date: Mon, 26 Oct 2026 09:53:29 -0700 Subject: [PATCH 54/80] 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) From 44021e5e3ce1f8b75e9a79da83ef64519e39f479 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Mon, 26 Oct 2026 13:59:29 -0700 Subject: [PATCH 55/80] Reset only the masked environments from the stable scene InteractiveScene.reset takes env_mask and forwards it to every asset and sensor, so the Warp environments use the stable scene and InteractiveSceneWarp is removed. --- .../isaaclab_experimental.envs.rst | 4 -- .../jichuanh-warp-scene-mask-native.added.rst | 3 ++ .../isaaclab/scene/interactive_scene.py | 27 ++++++---- .../jichuanh-warp-scene-mask-native.fixed.rst | 2 - .../jichuanh-warp-scene-mask-native.major | 0 ...ichuanh-warp-scene-mask-native.removed.rst | 2 + .../isaaclab_experimental/envs/__init__.pyi | 2 - .../isaaclab_experimental/envs/frontend.py | 11 ++--- .../envs/interactive_scene_warp.py | 49 ------------------- .../envs/manager_based_env_warp.py | 2 +- .../test/envs/test_frontend.py | 4 +- .../scene/test_interactive_scene_mask.py} | 7 ++- 12 files changed, 32 insertions(+), 81 deletions(-) delete mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.major create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.removed.rst delete mode 100644 source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py rename source/{isaaclab_experimental/test/envs/test_interactive_scene_warp.py => isaaclab_newton/test/scene/test_interactive_scene_mask.py} (90%) diff --git a/docs/source/api/lab_experimental/isaaclab_experimental.envs.rst b/docs/source/api/lab_experimental/isaaclab_experimental.envs.rst index 09e295daa8e5..300b482a45ca 100644 --- a/docs/source/api/lab_experimental/isaaclab_experimental.envs.rst +++ b/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: 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..c52b4d0c252e 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 ``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. diff --git a/source/isaaclab/isaaclab/scene/interactive_scene.py b/source/isaaclab/isaaclab/scene/interactive_scene.py index 2cb70929c1bb..d3c4388a9ee5 100644 --- a/source/isaaclab/isaaclab/scene/interactive_scene.py +++ b/source/isaaclab/isaaclab/scene/interactive_scene.py @@ -415,29 +415,38 @@ def state(self) -> dict[str, dict[str, dict[str, torch.Tensor]]]: Operations. """ - def reset(self, env_ids: Sequence[int] | None = None): + def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None): """Resets the scene entities. + .. 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 mask of the environments to reset. Shape is (num_envs,). Defaults to None. """ + # entities written before masked resets existed only accept indices + mask_kwargs = {} if env_mask is None else {"env_mask": env_mask} # -- assets for articulation in self._articulations.values(): - articulation.reset(env_ids) + articulation.reset(env_ids, **mask_kwargs) for cable_object in self._cable_objects.values(): - cable_object.reset(env_ids) + cable_object.reset(env_ids, **mask_kwargs) for deformable_object in self._deformable_objects.values(): - deformable_object.reset(env_ids) + deformable_object.reset(env_ids, **mask_kwargs) for rigid_object in self._rigid_objects.values(): - rigid_object.reset(env_ids) - for surface_gripper in self._surface_grippers.values(): - surface_gripper.reset(env_ids) + rigid_object.reset(env_ids, **mask_kwargs) + if self._surface_grippers: + # surface grippers reset by index only + gripper_ids = env_ids if env_mask is None else wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) + for surface_gripper in self._surface_grippers.values(): + surface_gripper.reset(gripper_ids) for rigid_object_collection in self._rigid_object_collections.values(): - rigid_object_collection.reset(env_ids) + rigid_object_collection.reset(env_ids, **mask_kwargs) # -- sensors for sensor in self._sensors.values(): - sensor.reset(env_ids) + sensor.reset(env_ids, **mask_kwargs) def write_data_to_sim(self): """Writes the data of the scene entities to the simulation.""" 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 deleted file mode 100644 index b712f5b4a8d4..000000000000 --- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst +++ /dev/null @@ -1,2 +0,0 @@ -* Fixed :meth:`~isaaclab_experimental.envs.InteractiveSceneWarp.reset` with ``env_mask`` resetting deformable - objects and sensors in every environment, and skipping cable objects. diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.major b/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.major new file mode 100644 index 000000000000..e69de29bb2d1 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..60116ae09cf4 --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.removed.rst @@ -0,0 +1,2 @@ +* **Breaking:** Removed ``isaaclab_experimental.envs.InteractiveSceneWarp``. Use + :class:`~isaaclab.scene.InteractiveScene`, whose :meth:`~isaaclab.scene.InteractiveScene.reset` takes ``env_mask``. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/__init__.pyi b/source/isaaclab_experimental/isaaclab_experimental/envs/__init__.pyi index ef3d01d6b4ec..109fa5f8d06c 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/__init__.pyi +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/__init__.pyi @@ -6,13 +6,11 @@ __all__ = [ "mdp", "DirectRLEnvWarp", - "InteractiveSceneWarp", "ManagerBasedEnvWarp", "ManagerBasedRLEnvWarp", ] from . import mdp from .direct_rl_env_warp import DirectRLEnvWarp -from .interactive_scene_warp import InteractiveSceneWarp from .manager_based_env_warp import ManagerBasedEnvWarp from .manager_based_rl_env_warp import ManagerBasedRLEnvWarp diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index cf2183c9310c..99836cdc7f2b 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -46,8 +46,6 @@ from isaaclab.envs import DirectMARLEnvCfg, DirectRLEnvCfg, ManagerBasedRLEnvCfg from isaaclab.managers.scene_entity_cfg import SceneEntityCfg as _StableSceneEntityCfg -from isaaclab_experimental.envs.interactive_scene_warp import InteractiveSceneWarp - logger = logging.getLogger(__name__) @@ -140,8 +138,8 @@ def _build_manager_env(cls, env_cfg: Any, **construct_kwargs: Any) -> gym.Env: def _build_direct_env(cls, env_cfg: Any, task_id: str, **construct_kwargs: Any) -> gym.Env: """Construct a direct warp env. - Direct task data stays stable, while the frontend selects the Warp scene - implementation and a hand-written Warp env class. The class is + Direct task data stays stable, while the frontend selects a + hand-written Warp env class. The class is resolved by name from the mirrored experimental package (or an explicit ``warp_entry_point`` override); see :meth:`_resolve_direct_warp_class`. If name resolution finds nothing, the task itself must be a warp-native @@ -157,11 +155,8 @@ def _build_direct_env(cls, env_cfg: Any, task_id: str, **construct_kwargs: Any) if env_class is None: # No warp twin by name: the task itself must be a warp-native registration. cls._assert_direct_warp_registration(task_id) - else: - cls._require_newton_physics(env_cfg, type(env_cfg).__name__) - env_cfg.scene.class_type = InteractiveSceneWarp - if env_class is None: return gym.make(task_id, cfg=env_cfg, **construct_kwargs) + cls._require_newton_physics(env_cfg, type(env_cfg).__name__) return env_class(cfg=env_cfg, **construct_kwargs) # ------------------------------------------------------------------ diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py deleted file mode 100644 index 5fa1a2eda227..000000000000 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py +++ /dev/null @@ -1,49 +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 - -"""Warp-native interactive scene with env_mask support for reset.""" - -from __future__ import annotations - -from collections.abc import Sequence - -import warp as wp - -from isaaclab.scene import InteractiveScene - - -class InteractiveSceneWarp(InteractiveScene): - """Interactive scene with warp-native env_mask support for reset. - - Extends :class:`InteractiveScene` to accept a boolean warp mask for selective resets, - avoiding the need to convert between env_ids and masks. - """ - - 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 - 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, 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) - for rigid_object_collection in self._rigid_object_collections.values(): - rigid_object_collection.reset(env_ids, env_mask=env_mask) - # -- sensors - for sensor in self._sensors.values(): - sensor.reset(env_ids, env_mask=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..697913544212 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 @@ -33,13 +33,13 @@ export_articulations_data, export_scene_data, ) +from isaaclab.scene import InteractiveScene 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.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.utils.warp import resolve_1d_mask diff --git a/source/isaaclab_experimental/test/envs/test_frontend.py b/source/isaaclab_experimental/test/envs/test_frontend.py index ff71cdb49b0d..a9a64dc4f1bb 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend.py +++ b/source/isaaclab_experimental/test/envs/test_frontend.py @@ -154,12 +154,12 @@ def fake_env(*, cfg: Any, **kwargs: Any) -> Any: assert calls == [(cfg, {"render_mode": "rgb_array"})] -def test_warp_direct_build_selects_warp_scene_from_cfg(): +def test_warp_direct_build_keeps_the_stable_scene(): cfg = types.SimpleNamespace(scene=types.SimpleNamespace(class_type=None)) expected_env = object() def fake_env(*, cfg: Any, **kwargs: Any) -> Any: - assert cfg.scene.class_type is fe.InteractiveSceneWarp + assert cfg.scene.class_type is None assert kwargs == {"render_mode": "rgb_array"} return expected_env diff --git a/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py b/source/isaaclab_newton/test/scene/test_interactive_scene_mask.py similarity index 90% rename from source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py rename to source/isaaclab_newton/test/scene/test_interactive_scene_mask.py index 440cc82a29a9..7cc9bbf153bb 100644 --- a/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py +++ b/source/isaaclab_newton/test/scene/test_interactive_scene_mask.py @@ -3,16 +3,15 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Tests for masked resets of the Warp interactive scene on Newton.""" +"""Tests for masked resets of the 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.scene import InteractiveScene, InteractiveSceneCfg from isaaclab.sensors.imu import ImuCfg from isaaclab.sim import SimulationCfg from isaaclab.terrains import TerrainImporterCfg @@ -45,7 +44,7 @@ def test_masked_reset_keeps_sensor_state_of_unselected_envs(): 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)) + scene = InteractiveScene(ImuSceneCfg(num_envs=2)) sim.reset() scene.reset() imu = scene["imu"] From c2d540beeb801627e507ec4793b7386187da6b7c Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 12:08:09 -0700 Subject: [PATCH 56/80] Check the Warp command and external-term contracts without env builds The second-environment command check moves into the generated_commands parity test on mock environments, replacing two Go2 builds. The external Warp terms are checked on the adapted cfg instead of a Cartpole build. The curriculum test drops its reward-weight phase, which the reward manager unit test already covers. --- .../envs/mdp/test_observations_warp_parity.py | 11 +++- .../test/envs/test_frontend_external_terms.py | 34 +++---------- .../test/envs/test_warp_env_commands.py | 50 ------------------- .../test/envs/test_warp_env_curriculum.py | 11 ++-- 4 files changed, 19 insertions(+), 87 deletions(-) delete mode 100644 source/isaaclab_experimental/test/envs/test_warp_env_commands.py 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 e8605133ae21..03ecdea228b3 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 @@ -7,6 +7,8 @@ from __future__ import annotations +import copy + import numpy as np import pytest import torch @@ -307,7 +309,7 @@ def test_last_action(self, warp_env, stable_env, art_data): class TestObservationParityNewTerms: """Verify observation Warp kernels for newly migrated terms match stable torch implementations.""" - def test_generated_commands(self, warp_env_bodies, stable_env_bodies, cmd_tensor): + def test_generated_commands(self, warp_env_bodies, stable_env_bodies, cmd_tensor, cmd_term): expected = stable_obs.generated_commands(stable_env_bodies, command_name="vel") actual = run_warp_obs(warp_obs.generated_commands, warp_env_bodies, (NUM_ENVS, CMD_DIM), command_name="vel") assert_close(actual, expected) @@ -320,3 +322,10 @@ def mutate(): warp_obs.generated_commands, warp_env_bodies, mutate, shape=(NUM_ENVS, CMD_DIM), command_name="vel" ) assert_close(actual_cap, stable_obs.generated_commands(stable_env_bodies, command_name="vel")) + + # A second environment reads its own command, not the first environment's. + other_env = copy.copy(warp_env_bodies) + other_cmd = torch.randn_like(cmd_tensor) + other_env.command_manager = MockCommandManager(other_cmd, cmd_term) + actual_other = run_warp_obs(warp_obs.generated_commands, other_env, (NUM_ENVS, CMD_DIM), command_name="vel") + assert_close(actual_other, other_cmd) diff --git a/source/isaaclab_experimental/test/envs/test_frontend_external_terms.py b/source/isaaclab_experimental/test/envs/test_frontend_external_terms.py index 6d64721f8745..3a551fed26b6 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend_external_terms.py +++ b/source/isaaclab_experimental/test/envs/test_frontend_external_terms.py @@ -3,36 +3,25 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Warp terms from an external project run on the Warp frontend unchanged. +"""Warp terms from an external project pass through the Warp frontend unchanged. The project lives outside the packages the frontend mirrors, so its terms have no twins. Its Warp terms must be used as-is, even when one shares its name with a built-in term, while a Torch term from the same project must still be rejected instead of being replaced by a same-named built-in twin. """ -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 importlib import textwrap import pytest -import torch from isaaclab_experimental.envs.frontend import WarpFrontend -import isaaclab.sim as sim_utils from isaaclab.managers import RewardTermCfg from isaaclab.utils import configclass import isaaclab_tasks # noqa: F401 from isaaclab_tasks.utils import resolve_task_config -_NUM_ENVS = 4 - _EXTERNAL_MDP = textwrap.dedent(''' import torch import warp as wp @@ -78,32 +67,21 @@ def external_mdp(tmp_path, monkeypatch): def _cartpole_cfg(rewards): env_cfg, _ = resolve_task_config("Isaac-Cartpole", "", overrides=("physics=newton_mjwarp",)) - env_cfg.seed = 7 - env_cfg.scene.num_envs = _NUM_ENVS env_cfg.rewards = rewards return env_cfg -def test_external_warp_terms_run_as_declared(external_mdp): +def test_external_warp_terms_are_kept_as_declared(external_mdp): @configclass class RewardsCfg: function_term = RewardTermCfg(func=external_mdp.joint_vel_l2, weight=1.0) class_term = RewardTermCfg(func=external_mdp.constant_bonus, weight=0.5, params={"value": 3.0}) env_cfg = _cartpole_cfg(RewardsCfg()) - sim_utils.create_new_stage() - env = WarpFrontend.build_env(env_cfg, "Isaac-Cartpole").unwrapped - try: - env.reset() - actions = torch.zeros((_NUM_ENVS, env.action_space.shape[-1]), device=env.device) - for _ in range(3): - _, reward, _, _, _ = env.step(actions) - - assert env.reward_manager.get_term_cfg("function_term").func is external_mdp.joint_vel_l2 - expected = (1.0 * 2.0 + 0.5 * 3.0) * env.step_dt - assert torch.allclose(reward, torch.full_like(reward, expected)) - finally: - env.close() + WarpFrontend.adapt_cfg(env_cfg) + + assert env_cfg.rewards.function_term.func is external_mdp.joint_vel_l2 + assert env_cfg.rewards.class_term.func is external_mdp.constant_bonus def test_external_torch_term_is_rejected(external_mdp): diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_commands.py b/source/isaaclab_experimental/test/envs/test_warp_env_commands.py deleted file mode 100644 index cfa8514d9186..000000000000 --- a/source/isaaclab_experimental/test/envs/test_warp_env_commands.py +++ /dev/null @@ -1,50 +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 - -"""Commands of Warp environments built one after another in one process.""" - -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_experimental.envs.mdp as warp_mdp -import torch -import warp as wp -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 - -_TASK = "Isaac-Velocity-Flat-UnitreeGo2" -_NUM_ENVS = 8 - - -def test_each_velocity_env_reads_its_own_commands(): - """A second Warp velocity environment builds, and its command-reading terms read its own commands. - - The terms read the command through the environment's command manager on every call, so nothing on the term - functions refers to an earlier environment. - """ - for seed in (7, 8): - env_cfg, _ = resolve_task_config(_TASK, "", overrides=("physics=newton_mjwarp",)) - env_cfg.seed = seed - env_cfg.scene.num_envs = _NUM_ENVS - sim_utils.create_new_stage() - env = WarpFrontend.build_env(env_cfg, _TASK).unwrapped - try: - env.reset() - env.step(torch.zeros((_NUM_ENVS, env.action_space.shape[-1]), device=env.device)) - out = wp.zeros((_NUM_ENVS, 3), dtype=wp.float32, device=env.device) - - warp_mdp.generated_commands(env, out, command_name="base_velocity") - - assert torch.equal(wp.to_torch(out), env.command_manager.get_command("base_velocity")) - finally: - env.close() diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_curriculum.py b/source/isaaclab_experimental/test/envs/test_warp_env_curriculum.py index 3afc247e688d..dc521b4a9185 100644 --- a/source/isaaclab_experimental/test/envs/test_warp_env_curriculum.py +++ b/source/isaaclab_experimental/test/envs/test_warp_env_curriculum.py @@ -30,9 +30,8 @@ def test_unchanged_term_settings_keep_the_recorded_stages(monkeypatch): """Curricula set their terms on every reset; a stage records again only when a setting it read changes. - Past its threshold, Reach's ``modify_reward_weight`` raises two penalties on every reset (episodes of 15 - steps), and the recorded reward stage reads weights on the device. Setting every term of every manager to - its own configuration records nothing either. + Setting every term of every manager to its own configuration records nothing. Episodes last 15 steps, so + each 20-step run resets. """ recordings = Counter() record = CapturedStage._record @@ -50,7 +49,6 @@ def counting_record(self, args, kwargs): env = WarpFrontend.build_env(env_cfg, _TASK).unwrapped try: env.reset() - env.common_step_counter = 5000 actions = torch.zeros((_NUM_ENVS, env.action_space.shape[-1]), device=env.device) def step(num_steps: int) -> int: @@ -60,10 +58,7 @@ def step(num_steps: int) -> int: reset_steps += int(bool((terminated | truncated).any())) return reset_steps - assert step(60) >= 3 - assert env.reward_manager.get_term_cfg("action_rate").weight == -0.005 - assert recordings["RewardManager.compute"] == 1 - + assert step(20) >= 1 recorded = dict(recordings) term_names = { env.reward_manager: env.reward_manager.active_terms, From e390289b4e203470ac3c41f628e85b20d3c2fce1 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 12:08:26 -0700 Subject: [PATCH 57/80] Trim duplicate id-resolution and mask-writer tests Assert the no-copy slice views once in the resolve_ids test and keep one wiring line per backend resolver in the asset contracts. Merge the two resolve_ids tests into one dtype-parametrized table, use a fixed small mask, and scope the CUDA-only resolve_1d_mask tests to the CUDA devices. Drop the Newton early return in the articulation writer contract so the fixed-tendon mask setters are checked for extra environments with torch and warp data, and remove the narrower standalone Newton test. Fold the collection body_ids mask-writer check into the writer read-back test and remove the duplicate ProxyArray check from the wrench composer test. --- .../contract/test_articulation_contract.py | 29 +---- .../contract/test_rigid_object_contract.py | 46 +++----- .../test/utils/test_wrench_composer.py | 5 +- .../test/utils/warp/test_resolve_selection.py | 109 ++++++++---------- 4 files changed, 73 insertions(+), 116 deletions(-) diff --git a/source/isaaclab/test/assets/contract/test_articulation_contract.py b/source/isaaclab/test/assets/contract/test_articulation_contract.py index 35f43d65898c..abfb6d1eff57 100644 --- a/source/isaaclab/test/assets/contract/test_articulation_contract.py +++ b/source/isaaclab/test/assets/contract/test_articulation_contract.py @@ -139,20 +139,12 @@ def test_finders(monkeypatch, backend): finder("nonexistent_xyz") -def test_resolve_ids_accept_tensor_views(monkeypatch, backend): - """Selector resolution honors tensor views, and environment slices alias the cached indices without copies.""" - art, _ = _articulation(backend, "cpu", monkeypatch, num_instances=4, num_joints=4, num_bodies=4) - ids = torch.arange(4, dtype=torch.int32) - for resolve in (art._resolve_env_ids, art._resolve_joint_ids, art._resolve_body_ids): - assert resolve(ids).shape[0] == 4 - assert resolve(ids[:2]).shape[0] == 2 - - cached = wp.to_torch(art._ALL_INDICES) - for selection in (slice(None), slice(1, None, 2), slice(0, 0)): - resolved = wp.to_torch(art._resolve_env_ids(selection)) - torch.testing.assert_close(resolved, cached[selection]) - assert resolved.data_ptr() == cached[selection].data_ptr() - assert resolved.stride() == cached[selection].stride() +def test_resolvers_select_from_their_own_cached_indices(monkeypatch, backend): + """Each selector resolver returns its own cached indices for a selection of every element.""" + art, _ = _articulation(backend, "cpu", monkeypatch) + assert art._resolve_env_ids(slice(None)) is art._ALL_INDICES + assert art._resolve_joint_ids(slice(None)) is art._ALL_JOINT_INDICES + assert art._resolve_body_ids(slice(None)) is art._ALL_BODY_INDICES # --------------------------------------------------------------------------- @@ -526,8 +518,6 @@ def read() -> ProxyArray: method(**{kwarg: 1.0}) _assert_reads_back(read(), torch.ones_like(expected), f"{writer} scalar") - if backend == "newton" and selection == "mask" and is_tendon: - return # pinned by test_newton_fixed_tendon_mask_setters_reject_extra_environments extra_env = _payload((N + 1, *shape[1:]), trailing, dtype, device) with pytest.raises((AssertionError, RuntimeError)): method(**{kwarg: extra_env}) @@ -535,13 +525,6 @@ def read() -> ProxyArray: method(**{kwarg: wp.from_torch(extra_env.contiguous(), dtype=dtype)}) -@requires("newton") -def test_newton_fixed_tendon_mask_setters_reject_extra_environments(monkeypatch): - art, _ = _articulation("newton", "cpu", monkeypatch) - with pytest.raises((AssertionError, RuntimeError)): - art.set_fixed_tendon_stiffness_mask(stiffness=torch.ones((N + 1, FT))) - - # quantity -> (writer, keyword, trailing size). Joint partial writes are covered by the ordering contract. _PARTIAL_WRITES = { "root_pose": ("write_root_link_pose_to_sim", "root_pose", 7), 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 514f4811ad57..1a6c85274bce 100644 --- a/source/isaaclab/test/assets/contract/test_rigid_object_contract.py +++ b/source/isaaclab/test/assets/contract/test_rigid_object_contract.py @@ -97,17 +97,10 @@ def test_counts_names_finders_and_index_resolution(monkeypatch, backend, asset): if asset == "collection": with pytest.warns(DeprecationWarning): assert obj.find_objects(".*", as_proxy=True)[0] is proxy - assert obj._resolve_body_ids(torch.arange(3, dtype=torch.int32)[:2]).shape[0] == 2 - # Environment selections resolve to strided views of the cached indices without copying. - env_ids = torch.arange(4, dtype=torch.int32) - assert obj._resolve_env_ids(env_ids[:2]).shape[0] == 2 - cached = wp.to_torch(obj._ALL_INDICES if asset == "object" else obj._ALL_ENV_INDICES) - for selection in (slice(None), slice(1, None, 2), slice(0, 0)): - resolved = wp.to_torch(obj._resolve_env_ids(selection)) - torch.testing.assert_close(resolved, cached[selection]) - assert resolved.data_ptr() == cached[selection].data_ptr() - assert resolved.stride() == cached[selection].stride() + # Each resolver returns its own cached indices for a selection of every element. + assert obj._resolve_env_ids(slice(None)) is (obj._ALL_INDICES if asset == "object" else obj._ALL_ENV_INDICES) + assert obj._resolve_body_ids(slice(None)) is obj._ALL_BODY_INDICES def _expected_properties(top: str) -> dict[str, type]: @@ -295,6 +288,18 @@ def test_writer_reads_back_full_and_partial_writes(monkeypatch, backend, device, expected[last] = update[last] torch.testing.assert_close(getattr(obj.data, getter).torch, expected, atol=1e-5, rtol=1e-5) + if ( + asset == "collection" + and selection == "mask" + and quantity.endswith(("pose", "velocity")) + and backend != "ovphysx" + ): + # Collection mask writers given ``body_ids`` write only those bodies of the masked environments. + update = _payload(shape, trailing, offset=200.0, device=device) + method(**{kwarg: update}, env_mask=selectors["env_mask"], body_ids=[num_bodies - 1]) + expected[-1, -1] = update[-1, -1] + torch.testing.assert_close(getattr(obj.data, getter).torch, expected, atol=1e-5, rtol=1e-5) + # Data with an extra environment is rejected for torch and warp inputs. bad = _payload((_NUM_INSTANCES + 1, *shape[1:]), trailing, device=device) for value in (bad, wp.from_torch(bad.contiguous(), dtype=wp_dtype)): @@ -314,27 +319,6 @@ 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/test/utils/test_wrench_composer.py b/source/isaaclab/test/utils/test_wrench_composer.py index b236b052707a..fccb5f3c8773 100644 --- a/source/isaaclab/test/utils/test_wrench_composer.py +++ b/source/isaaclab/test/utils/test_wrench_composer.py @@ -474,12 +474,9 @@ def test_deprecated_writers_warn_and_keep_add_and_set_behavior() -> None: _assert_vectors(composer.out_torque_b, torch.zeros((2, 2, 3))) -def test_invalid_selection_and_empty_wrench_input_are_reported() -> None: +def test_empty_wrench_input_is_reported() -> None: composer = _make_composer() - 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)() diff --git a/source/isaaclab/test/utils/warp/test_resolve_selection.py b/source/isaaclab/test/utils/warp/test_resolve_selection.py index fc951df8f614..a84a09fd68da 100644 --- a/source/isaaclab/test/utils/warp/test_resolve_selection.py +++ b/source/isaaclab/test/utils/warp/test_resolve_selection.py @@ -9,7 +9,7 @@ import torch import warp as wp -from isaaclab.test.utils import test_devices +from isaaclab.test.utils import DeviceScope, test_devices from isaaclab.utils.warp import ProxyArray from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids @@ -19,34 +19,62 @@ def _all_ids(n: int, device: str) -> wp.array: @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.""" +@pytest.mark.parametrize("dtype", [None, wp.int32]) +def test_resolve_ids(device, dtype): + """Every selection resolves to device ids; a selection of every element is the shared constant itself. + + Without ``dtype``, tensors and Warp arrays pass through and positive-step slices view the shared ids without a + copy. ``dtype=wp.int32`` gives a contiguous one-dimensional int32 Warp array for every kind of ids. + """ 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 + torch_device = wp.device_to_torch(device) + ids_int32 = wp.array([5], dtype=wp.int32, device=device) + for ids in (None, slice(None)): + assert resolve_ids(ids=ids, all_ids=all_ids, device=device, dtype=dtype) is all_ids + assert resolve_ids(ids=ids_int32, all_ids=all_ids, device=device, dtype=dtype) is ids_int32 + for ids, expected in ( + ([3, 8], [3, 8]), + (slice(2, 8, 3), [2, 5]), + (slice(None, None, -4), [9, 5, 1]), + (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]), + ): + resolved = resolve_ids(ids=ids, all_ids=all_ids, device=device, dtype=dtype) + if dtype is None: + assert resolved is ids or isinstance(ids, (list, slice)) + else: + assert isinstance(resolved, wp.array) and resolved.dtype == dtype and resolved.is_contiguous + assert (wp.to_torch(resolved) if isinstance(resolved, wp.array) else resolved).tolist() == expected with pytest.raises(TypeError, match="ProxyArray"): - resolve_ids(ids=ProxyArray(ids_wp), all_ids=all_ids, device=device) + resolve_ids(ids=ProxyArray(ids_int32), all_ids=all_ids, device=device, dtype=dtype) + if dtype is None: + cached = wp.to_torch(all_ids) + for selection in (slice(1, None, 2), slice(0, 0)): + resolved = wp.to_torch(resolve_ids(ids=selection, all_ids=all_ids, device=device)) + assert resolved.data_ptr() == cached[selection].data_ptr() + assert resolved.stride() == cached[selection].stride() + else: + for ids, error in ( + (torch.tensor([1.0], device=torch_device), TypeError), + (wp.array([1.0], dtype=wp.float32, device=device), TypeError), + (torch.zeros((2, 2), dtype=torch.int64, device=torch_device), ValueError), + ): + with pytest.raises(error): + resolve_ids(ids=ids, all_ids=all_ids, device=device, dtype=dtype) @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)) + expected = torch.tensor([0, 1, 1, 0, 1, 0, 0, 1], dtype=torch.bool, device=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) + all_ids = _all_ids(expected.shape[0], 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) @@ -54,40 +82,6 @@ def test_resolve_ids_from_a_mask(device, mask_type): 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): @@ -127,19 +121,18 @@ def test_resolve_1d_mask_needs_a_scratch_mask_only_for_ids(device): 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(): +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +def test_resolve_1d_mask_copies_a_mask_on_another_device(device): """A Warp mask on another device is copied to ``device``.""" - all_mask = wp.ones(3, dtype=wp.bool, device="cuda:0") + all_mask = wp.ones(3, dtype=wp.bool, device=device) 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] + resolved = resolve_1d_mask(mask=mask, all_mask=all_mask, device=device) + assert str(resolved.device) == device 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(): +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +def test_resolve_1d_mask_in_a_graph_capture(device): """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) From 1d17a83b6ba76748ed513ce5c9cbe15d716f1eff Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 12:13:48 -0700 Subject: [PATCH 58/80] Trim redundant builds and layout checks from the capture tests Merge the two manager Cartpole capture arms, flip compute_final_obs within one rollout instead of building disabled twins, reach both reset-log copy sites with three resets instead of a 120-step rollout, and move the set_term_cfg check to an EventManager unit test. Drop the non-capturable inventory and the assertions on the internal graph split layout; the value assertions keep the behaviour. --- .../test/envs/test_warp_env_capture.py | 22 ++------ .../test/envs/test_warp_env_final_obs.py | 35 +++++++------ .../test/envs/test_warp_env_log.py | 42 +++++----------- .../test/envs/test_warp_env_reset.py | 27 ---------- .../test/managers/test_event_manager.py | 45 +++++++++++++++++ .../test/managers/test_reward_manager.py | 2 - .../test/utils/test_mask_boundary_gate.py | 50 ------------------- .../test/utils/test_warp_capture.py | 11 +--- 8 files changed, 85 insertions(+), 149 deletions(-) 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 9beeae991935..61f3e895d328 100644 --- a/source/isaaclab_experimental/test/envs/test_warp_env_capture.py +++ b/source/isaaclab_experimental/test/envs/test_warp_env_capture.py @@ -9,8 +9,8 @@ (: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. +stages record again. The manager-based task also 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 @@ -59,17 +59,6 @@ def _recorded_stages(env) -> tuple[str, ...]: ) -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: @@ -111,7 +100,6 @@ def _rollout( rollout["terminated"].append(terminated.clone()) rollout["truncated"].append(truncated.clone()) rollout["stages_at_end"] = _recorded_stages(env) - rollout["stale_stages"] = _stale_stages(env) finally: env.close() for key in ("obs", "reward", "terminated", "truncated"): @@ -123,11 +111,10 @@ def _rollout( @pytest.mark.parametrize( ("task_id", "rebind", "eager_terms"), [ - ("Isaac-Cartpole", True, ()), - ("Isaac-Cartpole", False, _CARTPOLE_EAGER_TERMS), + ("Isaac-Cartpole", True, _CARTPOLE_EAGER_TERMS), ("Isaac-Cartpole-Direct", False, ()), ], - ids=["manager-cartpole", "manager-cartpole-eager-terms", "direct-cartpole"], + ids=["manager-cartpole", "direct-cartpole"], ) def test_captured_rollout_matches_eager( task_id: str, rebind: bool, eager_terms: tuple[str, ...], monkeypatch: pytest.MonkeyPatch @@ -143,7 +130,6 @@ def test_captured_rollout_matches_eager( assert {"ObservationManager", "RewardManager", "TerminationManager", "EventManager"} <= recorded_groups if rebind: 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_final_obs.py b/source/isaaclab_experimental/test/envs/test_warp_env_final_obs.py index 07c2e4413244..eb7a372a7b82 100644 --- a/source/isaaclab_experimental/test/envs/test_warp_env_final_obs.py +++ b/source/isaaclab_experimental/test/envs/test_warp_env_final_obs.py @@ -30,17 +30,19 @@ _NUM_ENVS = 16 _STEPS = 120 +# the environments read ``compute_final_obs`` on every step; each window holds resets +_ENABLE_STEP = 75 _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): +def _cartpole_cfg(task_id: str): """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 + env_cfg.compute_final_obs = False 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) @@ -54,9 +56,12 @@ def _cartpole_cfg(task_id: str, compute_final_obs: bool): 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) +def _rollout(task_id: str, frontend: str) -> dict: + """Push every cart with a constant force and record the observations of every step. + + The terminal observation is enabled from step :data:`_ENABLE_STEP` on. + """ + env_cfg = _cartpole_cfg(task_id) sim_utils.create_new_stage() if frontend == "warp": env = WarpFrontend.build_env(env_cfg, task_id).unwrapped @@ -68,6 +73,8 @@ def _rollout(task_id: str, frontend: str, compute_final_obs: bool) -> dict: # 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): + if step == _ENABLE_STEP: + env.cfg.compute_final_obs = True obs, _, terminated, truncated, extras = env.step(push) dones = terminated | truncated rollout["obs"].append(obs["policy"].clone()) @@ -86,17 +93,17 @@ def _rollout(task_id: str, frontend: str, compute_final_obs: bool) -> dict: @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] +def rollouts() -> dict[tuple[str, str], dict]: + arms = [(task_id, frontend) for task_id in _TASK_IDS for frontend in ("stable", "warp")] 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)] + rollout = rollouts[(task_id, "warp")] cart = rollout["cart_index"] - reset_steps = torch.nonzero(rollout["dones"].any(dim=1)).squeeze(-1).tolist() + resets = rollout["dones"].any(dim=1) + reset_steps = [step for step in range(_ENABLE_STEP, _STEPS) if resets[step]] assert reset_steps, "no environment terminated; the terminal observation is not exercised" assert sorted(rollout["final_obs"]) == reset_steps for step in reset_steps: @@ -110,9 +117,9 @@ def test_final_obs_is_the_observation_before_reset(task_id: str, rollouts: dict) @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"]) + rollout = rollouts[(task_id, "warp")] + assert rollout["dones"][:_ENABLE_STEP].any() + assert not any(rollout["final_obs_key"][:_ENABLE_STEP]) @pytest.mark.parametrize( @@ -121,7 +128,7 @@ def test_final_obs_is_absent_when_disabled(task_id: str, rollouts: dict): [("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)] + stable, warp = rollouts[(task_id, "stable")], rollouts[(task_id, "warp")] 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"]) 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 fd5a88e0549a..be45328d7eba 100644 --- a/source/isaaclab_experimental/test/envs/test_warp_env_log.py +++ b/source/isaaclab_experimental/test/envs/test_warp_env_log.py @@ -29,46 +29,30 @@ from isaaclab_tasks.utils import resolve_task_config _NUM_ENVS = 16 -_STEPS = 120 -def _make_env(task_id: str): +def _env_cfg(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 - + return env_cfg -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 _make_env(task_id: str): + return WarpFrontend.build_env(_env_cfg(task_id), task_id).unwrapped -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 +def _make_base_env(task_id: str): + env_cfg = _env_cfg(task_id) WarpFrontend.adapt_cfg(env_cfg) - sim_utils.create_new_stage() - env = ManagerBasedEnvWarp(cfg=env_cfg) + return ManagerBasedEnvWarp(cfg=env_cfg) + + +# the base and the RL environment each copy the log in their own ``_reset_mask`` +@pytest.mark.parametrize("make_env", [_make_base_env, _make_env], ids=["base", "rl"]) +def test_reset_logs_keep_their_values(make_env, monkeypatch: pytest.MonkeyPatch): + env = make_env("Isaac-Cartpole") logged = [] try: # a manager logs a persistent buffer that it refreshes on every reset 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 bf0ee211d098..c9c78933e822 100644 --- a/source/isaaclab_experimental/test/envs/test_warp_env_reset.py +++ b/source/isaaclab_experimental/test/envs/test_warp_env_reset.py @@ -14,7 +14,6 @@ import torch from isaaclab_experimental.envs.frontend import WarpFrontend -from isaaclab_experimental.managers import EventManager import isaaclab.sim as sim_utils @@ -52,29 +51,3 @@ 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 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) - 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 index c1ee9b0b13df..04cd4c98ffe0 100644 --- a/source/isaaclab_experimental/test/managers/test_event_manager.py +++ b/source/isaaclab_experimental/test/managers/test_event_manager.py @@ -34,6 +34,19 @@ def device_term(env, env_mask: wp.array): CALLS["device_term"] += 1 +@wp.kernel +def _write_value(env_mask: wp.array(dtype=wp.bool), out: wp.array(dtype=wp.float32), value: wp.float32): + env_id = wp.tid() + if env_mask[env_id]: + out[env_id] = value + + +def write_value(env, env_mask: wp.array, value: float): + """Event term that writes its configured value into the selected environments.""" + CALLS["write_value"] += 1 + wp.launch(_write_value, dim=env.num_envs, inputs=[env_mask, env.out, value], device=env.device) + + @WarpCapturable(False, reason="synchronizes with the host") class host_call_term(ManagerTermBase): """Class event term whose call cannot be recorded but whose reset can.""" @@ -97,3 +110,35 @@ def test_class_annotation_does_not_keep_the_reset_eager(monkeypatch): 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" + + +def test_set_term_cfg_applies_new_params_to_the_recorded_reset_events(monkeypatch): + """A replaced term configuration must not replay the scalars recorded from the old one.""" + monkeypatch.setattr(CapturedStage, "enabled", True) + env = SimpleNamespace( + num_envs=NUM_ENVS, + device=DEVICE, + sim=SimpleNamespace(is_playing=lambda: True), + out=wp.zeros(NUM_ENVS, dtype=wp.float32, device=DEVICE), + ) + manager = EventManager( + {"write_value": EventTermCfg(func=write_value, mode="reset", params={"value": 1.0})}, + env, + ) + env_mask = wp.ones(NUM_ENVS, dtype=wp.bool, device=DEVICE) + step_count = wp.zeros(1, dtype=wp.int32, device=DEVICE) + CALLS.clear() + + def apply_reset() -> list[float]: + manager.apply(mode="reset", env_mask_wp=env_mask, global_env_step_count=step_count) + return env.out.numpy().tolist() + + assert apply_reset() == [1.0] * NUM_ENVS + assert apply_reset() == [1.0] * NUM_ENVS + assert CALLS["write_value"] == 1, "the second reset must replay the recorded reset events" + + term_cfg = manager.get_term_cfg("write_value") + term_cfg.params["value"] = 2.0 + manager.set_term_cfg("write_value", term_cfg) + + assert apply_reset() == [2.0] * NUM_ENVS diff --git a/source/isaaclab_experimental/test/managers/test_reward_manager.py b/source/isaaclab_experimental/test/managers/test_reward_manager.py index 3901abb2c702..324feed831db 100644 --- a/source/isaaclab_experimental/test/managers/test_reward_manager.py +++ b/source/isaaclab_experimental/test/managers/test_reward_manager.py @@ -66,6 +66,4 @@ def test_capturability_is_decided_per_term_from_its_parameters(monkeypatch): wp.synchronize() 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 62e809ed2bf1..929e88d2e5a9 100644 --- a/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py +++ b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py @@ -14,8 +14,6 @@ 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 @@ -130,51 +128,3 @@ def test_mask_first_code_has_no_unsanctioned_host_syncs(): 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)`` 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"), - ("isaaclab_experimental/envs/mdp/events.py", "randomize_rigid_body_material"), -} - -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)`` or a parameter predicate.""" - 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 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)`` 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")): - 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)}" - ) diff --git a/source/isaaclab_experimental/test/utils/test_warp_capture.py b/source/isaaclab_experimental/test/utils/test_warp_capture.py index 0c8c3d49c8c9..7b1cff0c4c40 100644 --- a/source/isaaclab_experimental/test/utils/test_warp_capture.py +++ b/source/isaaclab_experimental/test/utils/test_warp_capture.py @@ -127,8 +127,8 @@ def body(values: wp.array) -> wp.array: 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.""" +def test_eager_call_first_runs_before_the_recorded_work(): + """A stage making its eager call first runs it before its recorded work, also when it replays the recording.""" eager_fn, first, second = _AffineStage(1, 10), _AffineStage(2, 0), _AffineStage(3, 0) def body(values: wp.array) -> wp.array: @@ -142,13 +142,6 @@ def body(values: wp.array) -> wp.array: 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: From 2048c626fc09043944ac1a17e44949ff871c5be1 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 12:16:17 -0700 Subject: [PATCH 59/80] Trim redundant setup from the masked-reset tests Replace the Newton scene test, which built a simulation and stepped it to check that a masked scene reset reaches one sensor, with a stub-scene unit test of InteractiveScene.reset: every asset and sensor receives the mask, surface grippers receive the masked ids, and a reset by ids passes no mask. Drop the eager masked cases that the CUDA-graph replay tests already cover for the Newton actuator state and the ray caster drift, and remove the Warp observation manager history_order test, whose fix is already on main and is unrelated to masked resets. --- .../test/scene/test_interactive_scene.py | 36 +++++++++++ .../test/managers/test_observation_manager.py | 37 ----------- .../assets/test_newton_actuators_newton.py | 5 +- .../test/scene/test_interactive_scene_mask.py | 62 ------------------- .../sensors/test_newton_raycast_sensor.py | 12 +--- 5 files changed, 39 insertions(+), 113 deletions(-) delete mode 100644 source/isaaclab_experimental/test/managers/test_observation_manager.py delete mode 100644 source/isaaclab_newton/test/scene/test_interactive_scene_mask.py diff --git a/source/isaaclab/test/scene/test_interactive_scene.py b/source/isaaclab/test/scene/test_interactive_scene.py index 40946843a879..8797dbe5069f 100644 --- a/source/isaaclab/test/scene/test_interactive_scene.py +++ b/source/isaaclab/test/scene/test_interactive_scene.py @@ -12,6 +12,7 @@ import numpy as np import pytest import torch +import warp as wp from isaaclab_physx.sim.schemas import PhysxRigidBodyCfg import isaaclab.sim as sim_utils @@ -184,6 +185,41 @@ def test_relative_deformable_state(): torch.testing.assert_close(written_state["env_ids"], env_ids) +def test_masked_reset_reaches_every_entity(): + """A masked reset passes the mask to every asset and sensor and the selected ids to surface grippers. + + A reset by ids passes no mask, so entities that accept only ids keep working. + """ + groups = ( + "_articulations", + "_cable_objects", + "_deformable_objects", + "_rigid_objects", + "_surface_grippers", + "_rigid_object_collections", + "_sensors", + ) + calls = {} + + def entity(group): + return SimpleNamespace(reset=lambda env_ids=None, **kwargs: calls.update({group: (env_ids, kwargs)})) + + scene = SimpleNamespace(**{group: {"entity": entity(group)} for group in groups}) + env_mask = wp.array([True, False, True], dtype=wp.bool, device="cpu") + + InteractiveScene.reset(scene, env_mask=env_mask) + assert calls.keys() == set(groups) + for group, (env_ids, kwargs) in calls.items(): + if group == "_surface_grippers": + assert env_ids.tolist() == [0, 2] and kwargs == {} + else: + assert env_ids is None and kwargs.keys() == {"env_mask"} and kwargs["env_mask"] is env_mask + + calls.clear() + InteractiveScene.reset(scene, env_ids=[1]) + assert calls == dict.fromkeys(groups, ([1], {})) + + def test_scene_publishes_plan_before_replicate(monkeypatch: pytest.MonkeyPatch): """A cfg-driven scene publishes the exact plan it forwards to replication.""" import isaaclab.cloner.replicate_session as replicate_session_module diff --git a/source/isaaclab_experimental/test/managers/test_observation_manager.py b/source/isaaclab_experimental/test/managers/test_observation_manager.py deleted file mode 100644 index 3094a9863da0..000000000000 --- a/source/isaaclab_experimental/test/managers/test_observation_manager.py +++ /dev/null @@ -1,37 +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 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,) 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 d0e287c7095c..6596c01e8ee8 100644 --- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py +++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py @@ -1030,11 +1030,10 @@ def gains(name: str) -> torch.Tensor: # --------------------------------------------------------------------------- -@pytest.mark.parametrize("selector", _RESET_SELECTORS) -def test_newton_state_reset_isolated_to_reset_env(newton_run: _Run, selector: str) -> None: +def test_newton_state_reset_isolated_to_reset_env(newton_run: _Run) -> None: """Newton: ``num_pushes`` zeroes for env 0's DOFs only after a partial reset of env 0.""" stateful_pairs = _warm_delayed_island(newton_run) - _reset_env_0(newton_run.articulations["delayed"], selector) + _reset_env_0(newton_run.articulations["delayed"], "env_ids") _assert_delay_state_cleared_for_env_0_only(newton_run, stateful_pairs) diff --git a/source/isaaclab_newton/test/scene/test_interactive_scene_mask.py b/source/isaaclab_newton/test/scene/test_interactive_scene_mask.py deleted file mode 100644 index 7cc9bbf153bb..000000000000 --- a/source/isaaclab_newton/test/scene/test_interactive_scene_mask.py +++ /dev/null @@ -1,62 +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 masked resets of the interactive scene on Newton.""" - -import torch -import warp as wp -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -import isaaclab.sim as sim_utils -from isaaclab.assets import RigidObjectCfg -from isaaclab.scene import InteractiveScene, 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 = InteractiveScene(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/test/sensors/test_newton_raycast_sensor.py b/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py index 9c454dd72cca..bba5a231f66c 100644 --- a/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py +++ b/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py @@ -213,7 +213,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 only, also on graph replay.""" + """Slice resets and replayed masked resets resample the drift of the selected environments only.""" 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)} @@ -222,16 +222,6 @@ def test_partial_reset_resamples_drift_of_selected_envs_only(sim): 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]) From 6318c479fae8cb91ac7858ce318b38bce4c7827f Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 12:18:05 -0700 Subject: [PATCH 60/80] Drop unused WarpRng setup from the parity tests No observation, reward or termination term reads WarpRng.state, so the three parity tests no longer set or clear it; the event parity test, whose terms draw from it, keeps it. Run the WarpRng tests on the CUDA devices only, since WarpRng has no device-specific path. --- source/isaaclab/test/utils/test_warp_rng.py | 8 ++++---- .../test/envs/mdp/test_observations_warp_parity.py | 5 ----- .../test/envs/mdp/test_rewards_warp_parity.py | 5 ----- .../test/envs/mdp/test_terminations_warp_parity.py | 5 ----- 4 files changed, 4 insertions(+), 19 deletions(-) diff --git a/source/isaaclab/test/utils/test_warp_rng.py b/source/isaaclab/test/utils/test_warp_rng.py index 2bc3815a2b8c..e408854d2e9a 100644 --- a/source/isaaclab/test/utils/test_warp_rng.py +++ b/source/isaaclab/test/utils/test_warp_rng.py @@ -9,7 +9,7 @@ import pytest import torch -from isaaclab.test.utils import test_devices +from isaaclab.test.utils import DeviceScope, test_devices from isaaclab.utils.seed import WarpRng, configure_seed @@ -23,7 +23,7 @@ def fresh_warp_rng(): WarpRng._seed = None -@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) 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) @@ -33,7 +33,7 @@ def test_initialize_allocates_a_new_state(device): assert WarpRng.state.shape == (5,) -@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) def test_configure_seed_reseeds_in_place(device): """configure_seed reseeds the existing state in place and reproducibly.""" configure_seed(7) @@ -50,7 +50,7 @@ def test_configure_seed_reseeds_in_place(device): assert not np.array_equal(state.numpy(), seeded) -@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) def test_default_seed_follows_torch(device): """Without a set seed, the state is seeded from torch's initial seed.""" torch.manual_seed(5) 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..4e709473dc4c 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,6 @@ 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) return env @@ -172,7 +168,6 @@ 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) 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..29d5a30d3c53 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,6 @@ 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) return env @@ -187,7 +183,6 @@ 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) 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..0cace559d50a 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,6 @@ 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) return env @@ -182,7 +178,6 @@ 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) return env From e74ad95c070667e4e895eccb584cf54a4efc71d2 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 12:56:51 -0700 Subject: [PATCH 61/80] Raise a clear error on a masked ray-caster reset without WarpRng state --- .../changelog.d/jichuanh-warp-rng-shared.changed.rst | 3 ++- .../isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py | 5 +++++ .../test/sensors/test_newton_raycast_sensor.py | 3 +++ 3 files changed, 10 insertions(+), 1 deletion(-) diff --git a/source/isaaclab/changelog.d/jichuanh-warp-rng-shared.changed.rst b/source/isaaclab/changelog.d/jichuanh-warp-rng-shared.changed.rst index 2fde38574e92..4cc96b750f2c 100644 --- a/source/isaaclab/changelog.d/jichuanh-warp-rng-shared.changed.rst +++ b/source/isaaclab/changelog.d/jichuanh-warp-rng-shared.changed.rst @@ -1,2 +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. + :class:`~isaaclab.utils.seed.WarpRng` state instead of random states of its own. It raises a ``RuntimeError`` + when that state is not initialized. 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 5a8882487c0f..276b282e194d 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py @@ -124,6 +124,11 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None # 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. + if WarpRng.state is None: + raise RuntimeError( + "A masked ray-caster reset draws drift from WarpRng.state; call WarpRng.initialize(num_envs," + " device) first (the Warp environments do this at construction)." + ) wp.launch( ray_caster_kernels.resample_drift_masked_kernel, dim=self._view_count, 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..4b032471c96d 100644 --- a/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py +++ b/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py @@ -224,6 +224,9 @@ def test_partial_reset_resamples_drift_of_selected_envs_only(sim): sensor.reset() # A masked reset draws from the process-wide Warp RNG state, which an environment initializes. + WarpRng.state = None + with pytest.raises(RuntimeError, match="WarpRng.initialize"): + sensor.reset(env_mask=wp.array([False, True, False], dtype=wp.bool, device=sensor.device)) WarpRng.initialize(scene.num_envs, sensor.device) drift = sensor.drift.torch.clone() From 38721eac3a64f993b54bedbd51b6662dbd5f650d Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 15:15:36 -0700 Subject: [PATCH 62/80] Resample masked ray-caster drift from the device mask --- .../isaaclab/sensors/ray_caster/base_ray_caster.py | 2 +- .../test/sensors/test_newton_raycast_sensor.py | 12 ++++++++++-- 2 files changed, 11 insertions(+), 3 deletions(-) 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 3741c39d7f29..99c70dbb29a9 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py @@ -127,7 +127,7 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None ray_caster_kernels.resample_drift_masked_kernel, dim=self._view_count, inputs=[ - env_mask, + reset_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)), 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 bba5a231f66c..1075d4332d54 100644 --- a/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py +++ b/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py @@ -212,8 +212,8 @@ 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): - """Slice resets and replayed masked resets resample the drift of the selected environments only.""" +def test_partial_reset_resamples_drift_of_selected_envs_only(sim, monkeypatch): + """Slice, host-mask and replayed masked resets resample the drift of the selected environments only.""" 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)} @@ -243,6 +243,14 @@ def test_partial_reset_resamples_drift_of_selected_envs_only(sim): 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() + # A host mask reaches the drift kernel on the sensor device. Strict launches reject a host array that + # GPUs with heterogeneous memory management would otherwise read silently. + monkeypatch.setattr(wp.config, "launch_array_access_mode", wp.config.LaunchArrayAccessMode.STRICT) + drift = sensor.drift.torch.clone() + sensor.reset(env_mask=wp.array([False, True, False], dtype=wp.bool, device="cpu")) + torch.testing.assert_close(sensor.drift.torch[[0, 2]], drift[[0, 2]]) + assert not torch.allclose(sensor.drift.torch[1], drift[1]) + @configclass class MultiMeshTargetSceneCfg(RaycastTestSceneCfg): From 0e9e57d8d59f4877ef676a69f246aee60a2e8658 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 15:15:43 -0700 Subject: [PATCH 63/80] Pass the reset mask to actuator models only when one is given --- .../jichuanh-warp-scene-mask-native.changed.rst | 5 +++-- .../isaaclab/actuators/actuator_collection.py | 5 +++-- .../test/actuators/test_actuator_collection.py | 17 +++++++++++++++++ 3 files changed, 23 insertions(+), 4 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 ae017575fe71..cf37839edc0b 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 @@ -6,8 +6,9 @@ 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`. + accept ``env_mask`` to support masked resets and reset from it without synchronizing the host, for example through + :func:`~isaaclab.utils.array.index_fill_` or :func:`torch.where`. Resets without a mask still call + ``reset(env_ids)``. * 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_collection.py b/source/isaaclab/isaaclab/actuators/actuator_collection.py index cd9366b6095a..504216c02ded 100644 --- a/source/isaaclab/isaaclab/actuators/actuator_collection.py +++ b/source/isaaclab/isaaclab/actuators/actuator_collection.py @@ -195,11 +195,12 @@ 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. """ 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) + # actuator models written before masked resets existed only accept indices + mask_kwargs = {} if env_mask is None else {"env_mask": 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, env_mask=group_env_mask) + actuator.reset(group_env_ids, **mask_kwargs) 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: diff --git a/source/isaaclab/test/actuators/test_actuator_collection.py b/source/isaaclab/test/actuators/test_actuator_collection.py index 831477f44a17..7f6db0e662bb 100644 --- a/source/isaaclab/test/actuators/test_actuator_collection.py +++ b/source/isaaclab/test/actuators/test_actuator_collection.py @@ -47,6 +47,13 @@ def compute(self, control_action, joint_pos, joint_vel): return super().compute(control_action, joint_pos, joint_vel) +class IndexResetActuator(ImplicitActuator): + """Custom actuator with the index-only :meth:`reset` signature from before masked resets.""" + + def reset(self, env_ids): + self.reset_env_ids = env_ids + + def _ideal_cfg(joints: list[str], *, stiffness: float, damping: float, effort_limit: float): return IdealPDActuatorCfg( joint_names_expr=joints, @@ -824,6 +831,16 @@ def test_custom_singleton_compute_receives_original_selector(): assert collection["all"].observed_joint_indices == slice(None) +def test_index_reset_reaches_custom_actuator_without_mask_argument(): + cfg = _implicit_cfg() + cfg.class_type = IndexResetActuator + collection = ActuatorCollection({"all": cfg}, FakeActuatorControl()) + + collection.reset([1]) + + assert collection["all"].reset_env_ids == [1] + + def test_multi_group_explicit_outputs_match_pd_formula(): joint_names = [f"joint_{index}" for index in range(4)] control = FakeActuatorControl(joint_names=joint_names) From 59f9cf5283fb2e2bba1d944713ad1634411cf42f Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 15:17:25 -0700 Subject: [PATCH 64/80] Accept env_mask in the multirotor reset --- .../jichuanh-warp-scene-mask-native.fixed.rst | 4 ++-- .../assets/multirotor/multirotor.py | 12 +++++++++--- .../test/assets/test_multirotor.py | 15 +++++++++++++++ 3 files changed, 26 insertions(+), 5 deletions(-) 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 index 177eb50253d4..e04a0303897f 100644 --- 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 @@ -1,2 +1,2 @@ -* Fixed :class:`~isaaclab_contrib.assets.Multirotor` resets with only ``env_mask`` resetting the thrusters of every - environment. +* Fixed :meth:`~isaaclab_contrib.assets.Multirotor.reset` rejecting the ``env_mask`` argument of the articulation + reset. A masked reset now resets the thrusters and thrust targets of the masked environments only. diff --git a/source/isaaclab_contrib/isaaclab_contrib/assets/multirotor/multirotor.py b/source/isaaclab_contrib/isaaclab_contrib/assets/multirotor/multirotor.py index f65ded646726..ac92b910fcd0 100644 --- a/source/isaaclab_contrib/isaaclab_contrib/assets/multirotor/multirotor.py +++ b/source/isaaclab_contrib/isaaclab_contrib/assets/multirotor/multirotor.py @@ -252,26 +252,32 @@ def set_thrust_target( # set targets self._data.thrust_target[env_ids, thruster_ids] = target - def reset(self, env_ids: Sequence[int] | None = None): + def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None): """Reset the multirotor to default state. This method resets both the base articulation state (pose, velocities) and multirotor-specific state (thruster targets) to their default values as specified in the configuration. + .. 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). Can be a sequence of integers or None. + env_mask: Environment mask. Shape is (num_instances,). Defaults to None. Note: The default thruster state is set via the :attr:`MultirotorCfg.init_state.rps` configuration parameter. """ # call parent reset - super().reset(env_ids) + super().reset(env_ids, env_mask=env_mask) # reset multirotor-specific data - if env_ids is None: + if env_mask is not None: + env_ids = wp.to_torch(env_mask) + elif env_ids is None: env_ids = self._ALL_INDICES elif not isinstance(env_ids, (torch.Tensor, slice)): env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) diff --git a/source/isaaclab_contrib/test/assets/test_multirotor.py b/source/isaaclab_contrib/test/assets/test_multirotor.py index d0bf04b2baf4..6faf717f5d07 100644 --- a/source/isaaclab_contrib/test/assets/test_multirotor.py +++ b/source/isaaclab_contrib/test/assets/test_multirotor.py @@ -19,6 +19,7 @@ import pytest import torch import warp as wp +from isaaclab_physx.assets.articulation import Articulation from isaaclab_physx.physics import PhysxCfg import isaaclab.sim.utils.prims as prim_utils @@ -168,6 +169,20 @@ def test_thruster_collection_resets_masked_envs_unit(): assert [env_ids.tolist() for env_ids in reset_env_ids] == [[0, 2]] +def test_reset_with_only_mask_restores_masked_thrust_targets_unit(monkeypatch): + """A mask-only reset, as the scene forwards it, reaches the articulation and restores the masked envs only.""" + parent_masks = [] + monkeypatch.setattr(Articulation, "reset", lambda self, env_ids=None, env_mask=None: parent_masks.append(env_mask)) + m = Multirotor.__new__(Multirotor) + m._data = types.SimpleNamespace(thrust_target=torch.ones(4, 2), default_thruster_rps=torch.zeros(4, 2)) + env_mask = wp.array([True, False, True, False], dtype=wp.bool, device="cpu") + + m.reset(env_mask=env_mask) + + assert parent_masks == [env_mask] + torch.testing.assert_close(m._data.thrust_target, torch.tensor([[0.0, 0.0], [1.0, 1.0], [0.0, 0.0], [1.0, 1.0]])) + + def generate_multirotor( multirotor_cfg: MultirotorCfg, num_multirotors: int, device: str ) -> tuple[Multirotor, torch.Tensor]: From d8d4911965621f66ec094badd1a8026ef4b08219 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 15:17:35 -0700 Subject: [PATCH 65/80] Match the base reset selector type in the actuator overrides --- source/isaaclab/isaaclab/actuators/actuator_net.py | 4 ++-- source/isaaclab/isaaclab/actuators/actuator_pd.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/source/isaaclab/isaaclab/actuators/actuator_net.py b/source/isaaclab/isaaclab/actuators/actuator_net.py index 9a026fe07d5a..5d06e4bfb500 100644 --- a/source/isaaclab/isaaclab/actuators/actuator_net.py +++ b/source/isaaclab/isaaclab/actuators/actuator_net.py @@ -68,7 +68,7 @@ def __init__(self, cfg: ActuatorNetLSTMCfg, *args, **kwargs): Operations. """ - def reset(self, env_ids: Sequence[int] | None = None, env_mask: torch.Tensor | None = None): + def reset(self, env_ids: Sequence[int] | slice | 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(): @@ -139,7 +139,7 @@ def __init__(self, cfg: ActuatorNetMLPCfg, *args, **kwargs): Operations. """ - def reset(self, env_ids: Sequence[int] | None = None, env_mask: torch.Tensor | None = None): + def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: torch.Tensor | None = None): # reset the history for the specified environments selection = env_ids if env_mask is None else env_mask index_fill_(self._joint_pos_error_history, selection, 0.0) diff --git a/source/isaaclab/isaaclab/actuators/actuator_pd.py b/source/isaaclab/isaaclab/actuators/actuator_pd.py index c2d2e5a2c192..9531e354bddb 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, env_mask: torch.Tensor | None = None): + def reset(self, env_ids: Sequence[int] | slice | 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] | None = None, env_mask: torch.Tensor | None = None): + def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: torch.Tensor | None = None): pass def compute( From 13ec11cc33c4c04e36fb1f5985556aa41456bfb1 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 15:18:49 -0700 Subject: [PATCH 66/80] Accept every id form main accepted in resolve_1d_mask Warp ids of either integer width on any device are moved to the target device as contiguous int32, slices with a negative step select their elements, and a mask on "cuda" counts as on the current CUDA device. The Newton articulation no longer round-trips int64 ids through torch. --- .../changelog.d/jichuanh-id-mask.changed.rst | 5 ++-- source/isaaclab/isaaclab/utils/warp/utils.py | 23 ++++++++++------- .../contract/test_rigid_object_contract.py | 13 +++++++--- .../test/utils/warp/test_resolve_selection.py | 25 ++++++++++++++++--- .../assets/articulation/articulation.py | 2 -- 5 files changed, 47 insertions(+), 21 deletions(-) diff --git a/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst index 2ae1cf9e1daa..b9f819e8f3be 100644 --- a/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst +++ b/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst @@ -7,5 +7,6 @@ :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. +* Changed :func:`~isaaclab.utils.warp.utils.resolve_1d_mask` to copy a Warp mask or Warp ids on another device to + ``device``, to accept 64-bit Warp ids and slices with a negative step, and to accept a boolean torch mask on + ``device``, including ``"cuda"`` for the current CUDA device, inside a CUDA graph capture. diff --git a/source/isaaclab/isaaclab/utils/warp/utils.py b/source/isaaclab/isaaclab/utils/warp/utils.py index 24e23707615c..a4d637692cc2 100644 --- a/source/isaaclab/isaaclab/utils/warp/utils.py +++ b/source/isaaclab/isaaclab/utils/warp/utils.py @@ -38,11 +38,12 @@ def resolve_1d_mask( """Resolve ids/mask into a warp boolean mask. Callers provide pre-allocated ``all_mask`` (all-True) and ``scratch_mask`` (reusable - work buffer) so this function never allocates. + work buffer), so nothing is allocated for ``None``, a slice, a boolean mask on ``device`` + or contiguous ``wp.int32`` ids on ``device``. Other selections are copied to ``device`` first. Args: ids: Index ids. Accepts ``Sequence[int]``, ``slice``, ``torch.Tensor``, - ``wp.array(dtype=wp.int32)``, or ``None`` (all elements). + ``wp.array`` of 32-bit or 64-bit integers, or ``None`` (all elements). 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* @@ -77,8 +78,8 @@ def resolve_1d_mask( 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) + if str(mask.device) != str(wp.get_device(device)): + mask = mask.to(wp.device_to_torch(device)) return wp.from_torch(mask, dtype=wp.bool) raise TypeError(f"Unsupported mask type: {type(mask)}") @@ -90,8 +91,12 @@ def resolve_1d_mask( raise ValueError("resolve_1d_mask needs a scratch_mask to write ids into.") if isinstance(ids, slice): + # a mask has no order, so a negative step selects what its reversed, positive-step range selects + selected = range(*ids.indices(scratch_mask.shape[0])) + if selected.step < 0: + selected = selected[::-1] scratch_mask.zero_() - wp.to_torch(scratch_mask)[ids] = True + wp.to_torch(scratch_mask)[selected.start : selected.stop : selected.step] = True return scratch_mask # --- Normalize to concrete type --- @@ -114,9 +119,9 @@ def resolve_1d_mask( elif isinstance(ids, wp.array): if ids.shape[0] == 0: return scratch_mask - if ids.dtype != wp.int32: - raise TypeError(f"Unsupported wp.array dtype for ids: {ids.dtype}. Expected wp.int32 index array.") - ids_wp = ids + if ids.device != wp.get_device(device): + ids = ids.to(device) + ids_wp = _ids_as_dtype(ids, wp.int32) else: if len(ids) == 0: return scratch_mask @@ -131,7 +136,7 @@ def _mask_needs_copy(mask: wp.array | torch.Tensor, device: str) -> bool: 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 mask.dtype != torch.bool or str(mask.device) != str(wp.get_device(device)) return True 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 1a6c85274bce..49fae35e4b0d 100644 --- a/source/isaaclab/test/assets/contract/test_rigid_object_contract.py +++ b/source/isaaclab/test/assets/contract/test_rigid_object_contract.py @@ -295,10 +295,15 @@ def test_writer_reads_back_full_and_partial_writes(monkeypatch, backend, device, and backend != "ovphysx" ): # Collection mask writers given ``body_ids`` write only those bodies of the masked environments. - update = _payload(shape, trailing, offset=200.0, device=device) - method(**{kwarg: update}, env_mask=selectors["env_mask"], body_ids=[num_bodies - 1]) - expected[-1, -1] = update[-1, -1] - torch.testing.assert_close(getattr(obj.data, getter).torch, expected, atol=1e-5, rtol=1e-5) + last_body = num_bodies - 1 + for offset, body_ids in enumerate( + ([last_body], wp.array([last_body], dtype=wp.int64, device=device), slice(None, last_body - 1, -2)), + start=2, + ): + update = _payload(shape, trailing, offset=100.0 * offset, device=device) + method(**{kwarg: update}, env_mask=selectors["env_mask"], body_ids=body_ids) + expected[-1, -1] = update[-1, -1] + torch.testing.assert_close(getattr(obj.data, getter).torch, expected, atol=1e-5, rtol=1e-5) # Data with an extra environment is rejected for torch and warp inputs. bad = _payload((_NUM_INSTANCES + 1, *shape[1:]), trailing, device=device) diff --git a/source/isaaclab/test/utils/warp/test_resolve_selection.py b/source/isaaclab/test/utils/warp/test_resolve_selection.py index a84a09fd68da..347f888da5ad 100644 --- a/source/isaaclab/test/utils/warp/test_resolve_selection.py +++ b/source/isaaclab/test/utils/warp/test_resolve_selection.py @@ -83,7 +83,9 @@ def test_resolve_ids_from_a_mask(device, mask_type): @pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("ids_type", ["list", "slice", "torch_int64", "torch_int32", "warp_int32"]) +@pytest.mark.parametrize( + "ids_type", ["list", "slice", "reverse_slice", "torch_int64", "torch_int32", "warp_int32", "warp_int64"] +) 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 @@ -94,9 +96,11 @@ def test_resolve_1d_mask_from_ids(device, ids_type): ids = { "list": selected, "slice": slice(1, 12, 3), + "reverse_slice": slice(None, None, -2), "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), + "warp_int64": wp.array(selected, dtype=wp.int64, device=device), }[ids_type] expected = [False] * n for i in range(n)[ids] if isinstance(ids, slice) else selected: @@ -122,17 +126,29 @@ def test_resolve_1d_mask_needs_a_scratch_mask_only_for_ids(device): @pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -def test_resolve_1d_mask_copies_a_mask_on_another_device(device): - """A Warp mask on another device is copied to ``device``.""" +def test_resolve_1d_mask_copies_a_selection_on_another_device(device, monkeypatch): + """A Warp mask or Warp ids on another device are copied to ``device``. + + The strict launch mode rejects a kernel argument on another device, which the default mode may read through + host-memory addressing or fault on. + """ + monkeypatch.setattr(wp.config, "launch_array_access_mode", wp.config.LaunchArrayAccessMode.STRICT) all_mask = wp.ones(3, dtype=wp.bool, device=device) + scratch = wp.zeros(3, dtype=wp.bool, device=device) mask = wp.array([True, False, True], dtype=wp.bool, device="cpu") resolved = resolve_1d_mask(mask=mask, all_mask=all_mask, device=device) assert str(resolved.device) == device and resolved.numpy().tolist() == [True, False, True] + ids = wp.array([0, 2], dtype=wp.int32, device="cpu") + assert resolve_1d_mask(ids=ids, all_mask=all_mask, scratch_mask=scratch, device=device) is scratch + assert scratch.numpy().tolist() == [True, False, True] @pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) def test_resolve_1d_mask_in_a_graph_capture(device): - """Inside a capture, a boolean mask on device resolves; ids or a mask that needs a cast or copy raise.""" + """Inside a capture, a boolean mask on device resolves; ids or a mask that needs a cast or copy raise. + + ``"cuda"`` names the current CUDA device, so a mask on it needs no copy. + """ 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) @@ -142,6 +158,7 @@ def test_resolve_1d_mask_in_a_graph_capture(device): 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) + assert resolve_1d_mask(mask=mask_torch, all_mask=all_mask, device="cuda").ptr == resolved.ptr 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) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index efe2d22a2aef..d8f2b17f2821 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -2311,8 +2311,6 @@ 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, 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 ) From 784b0a2d253a615717598ffde810131609fb3c74 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 15:20:06 -0700 Subject: [PATCH 67/80] Refresh the solver on every replay of the Newton tendon mask writer Inside a CUDA graph capture the writer records the solver's tendon property refresh instead of queuing it on the host, where only the first replay's physics step would apply it. --- .../changelog.d/jichuanh-id-mask.changed.rst | 3 ++- .../assets/articulation/articulation.py | 8 +++++-- .../test/assets/test_articulation.py | 22 +++++++++++++++++++ 3 files changed, 30 insertions(+), 3 deletions(-) diff --git a/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst index 0c69e3841f70..e0e59800bd6d 100644 --- a/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst +++ b/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst @@ -10,4 +10,5 @@ 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. + instead of converting the masks to indices, which synchronized the host with the GPU on every call. Recorded into + a CUDA graph, ``write_fixed_tendon_properties_to_sim_mask`` refreshes the solver's tendon properties on every replay. diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index d8f2b17f2821..48c3a20e7862 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -3163,8 +3163,12 @@ def write_fixed_tendon_properties_to_sim_mask( ], 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) + # the solver keeps its own copy of the tendon properties and only re-reads them when notified; + # a capture records the refresh, so every replay reaches the solver + if wp.get_device(self.device).is_capturing: + self._backend_manager.backend.solver.notify_model_changed(ModelFlags.TENDON_PROPERTIES) + else: + self._backend_manager.add_model_change(ModelFlags.TENDON_PROPERTIES) def set_spatial_tendon_stiffness_index( self, diff --git a/source/isaaclab_newton/test/assets/test_articulation.py b/source/isaaclab_newton/test/assets/test_articulation.py index 20e39199a7d8..c226f6fbcd40 100644 --- a/source/isaaclab_newton/test/assets/test_articulation.py +++ b/source/isaaclab_newton/test/assets/test_articulation.py @@ -1696,6 +1696,28 @@ def test_hand_with_tendons_targets_only_given_envs_and_properties_reach_solver(s articulation.write_fixed_tendon_properties_to_sim_index() +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +def test_fixed_tendon_mask_writer_reaches_solver_on_every_replay(device: str) -> None: + """Each replay of a CUDA graph holding the fixed-tendon mask writer refreshes the solver's tendon properties. + + The solver keeps its own copy of them, so a refresh queued only while capturing would leave the physics on the + values of the first replay. + """ + 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() + stiffness = torch.zeros(articulation.num_instances, articulation.num_fixed_tendons, device=device) + with wp.ScopedCapture(device) as capture: + articulation.set_fixed_tendon_stiffness_mask(stiffness=stiffness) + articulation.write_fixed_tendon_properties_to_sim_mask() + solver_model = sim.backend_manager.backend.solver.mjw_model + for value in (12.0, 34.0): + stiffness.fill_(value) + wp.capture_launch(capture.graph) + sim.step() + np.testing.assert_allclose(solver_model.tendon_stiffness.numpy(), value) + + @pytest.mark.isaacsim_ci @pytest.mark.parametrize("island", ["fixed", "floating", "ordered"], ids=["fixed", "floating", "floating_reordered"]) def test_get_gravity_compensation_forces_matches_jacobian_gravity(scene: _Scene, island: str) -> None: From 6a4f173813b36725f70d935127e572cbc1eb7157 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 15:22:27 -0700 Subject: [PATCH 68/80] Tidy the mask writers, their docs and the selection tests - Give the Newton collection mask writers one env and body mask resolver. - Route the OVPhysX collection mask writers through the class id resolvers, which now take a mask as the PhysX ones do. - Use the dedicated-implementation tip on the Newton tendon mask setters and document the units of the Newton body mask kernels. - Note the camera's stricter env_ids in the changelog. - Check resolved ids against literal ids, and drop the integer mask rows, which are outside the boolean mask contract. --- .../changelog.d/jichuanh-id-mask.changed.rst | 4 +- .../test/utils/warp/test_resolve_selection.py | 18 ++--- .../assets/articulation/articulation.py | 12 ++-- .../isaaclab_newton/assets/kernels.py | 30 +++++---- .../rigid_object_collection.py | 67 +++++++++---------- .../rigid_object_collection.py | 30 ++++----- 6 files changed, 79 insertions(+), 82 deletions(-) diff --git a/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst index b9f819e8f3be..a94a9dba2c3a 100644 --- a/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst +++ b/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst @@ -2,7 +2,9 @@ :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`. + 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.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 diff --git a/source/isaaclab/test/utils/warp/test_resolve_selection.py b/source/isaaclab/test/utils/warp/test_resolve_selection.py index 347f888da5ad..610b90cf72d0 100644 --- a/source/isaaclab/test/utils/warp/test_resolve_selection.py +++ b/source/isaaclab/test/utils/warp/test_resolve_selection.py @@ -65,21 +65,17 @@ def test_resolve_ids(device, dtype): @pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("mask_type", ["warp", "torch_bool", "torch_int"]) +@pytest.mark.parametrize("mask_type", ["warp", "torch"]) def test_resolve_ids_from_a_mask(device, mask_type): """A mask resolves to its ascending ids, takes precedence over ids, and honours ``dtype``.""" - expected = torch.tensor([0, 1, 1, 0, 1, 0, 0, 1], dtype=torch.bool, device=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(expected.shape[0], device) + mask_bool = torch.tensor([0, 1, 1, 0, 1, 0, 0, 1], dtype=torch.bool, device=wp.device_to_torch(device)) + mask = wp.from_torch(mask_bool, dtype=wp.bool) if mask_type == "warp" else mask_bool + all_ids = _all_ids(mask_bool.shape[0], 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()) + assert ids.dtype == torch.int64 and ids.tolist() == [1, 2, 4, 7] 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)) + assert ids_int32.numpy().tolist() == [1, 2, 4, 7] @pytest.mark.parametrize("device", test_devices()) @@ -118,7 +114,7 @@ def test_resolve_1d_mask_needs_a_scratch_mask_only_for_ids(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)) + mask_torch = torch.tensor([False, True, True, False, False], 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"): diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index 48c3a20e7862..879e42a4be21 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -2700,8 +2700,8 @@ def set_fixed_tendon_stiffness_mask( This method expects full data. .. tip:: - 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. + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. Args: stiffness: Fixed tendon stiffness. Shape is (num_instances, num_fixed_tendons). @@ -2802,8 +2802,8 @@ def set_fixed_tendon_damping_mask( This method expects full data. .. tip:: - 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. + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. Args: damping: Fixed tendon damping. Shape is (num_instances, num_fixed_tendons). @@ -2924,8 +2924,8 @@ def set_fixed_tendon_position_limit_mask( This method expects full data. .. tip:: - 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. + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. Args: limit: Fixed tendon position limits ``[lower, upper]`` [m]. Shape is (num_instances, num_fixed_tendons) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/kernels.py b/source/isaaclab_newton/isaaclab_newton/assets/kernels.py index c9f43ef3adec..3a232ea7ea1f 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/kernels.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/kernels.py @@ -973,13 +973,14 @@ def set_body_com_pose_to_sim_mask( """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). + data: Input array of body COM poses (position [m] and quaternion). Shape is (num_envs, num_bodies). + body_com_pos_b: Input array of body COM positions in body frame [m]. 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. + body_com_pose_w: Output array where body COM poses (position [m] and quaternion) are written. Shape is (num_envs, num_bodies). + body_link_pose_w: Output array where body link poses derived from COM are written (position [m] and + quaternion). Shape is (num_envs, num_bodies). """ i, j = wp.tid() if env_mask[i] and body_mask[j]: @@ -998,11 +999,12 @@ def set_body_com_velocity_to_sim_mask( """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). + data: Input array of body COM spatial velocities [m/s, rad/s]. 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). + body_com_velocity_w: Output array where body COM velocities [m/s, rad/s] are written. + Shape is (num_envs, num_bodies). + body_acc_w: Output array where body accelerations [m/s^2, rad/s^2] are zeroed. Shape is (num_envs, num_bodies). """ i, j = wp.tid() if env_mask[i] and body_mask[j]: @@ -1025,15 +1027,17 @@ def set_body_link_velocity_to_sim_mask( 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). + data: Input array of body link spatial velocities [m/s, rad/s]. Shape is (num_envs, num_bodies). + body_com_pos_b: Input array of body COM positions in body frame [m]. Shape is (num_envs, num_bodies). + body_link_pose_w: Input array of body link poses in world frame (position [m] and quaternion). + 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. + body_link_velocity_w: Output array where body link velocities [m/s, rad/s] are written. + Shape is (num_envs, num_bodies). + body_com_velocity_w: Output array where body COM velocities [m/s, rad/s], 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). + body_acc_w: Output array where body accelerations [m/s^2, rad/s^2] are zeroed. Shape is (num_envs, num_bodies). """ i, j = wp.tid() if env_mask[i] and body_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 ad916dbfc09e..311683ba4446 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 @@ -473,15 +473,7 @@ 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. """ - 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, - ) + env_mask, body_mask = self._resolve_env_and_body_masks(env_mask, body_mask, body_ids) 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), @@ -587,15 +579,7 @@ 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. """ - 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, - ) + env_mask, body_mask = self._resolve_env_and_body_masks(env_mask, body_mask, body_ids) 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, @@ -709,15 +693,7 @@ 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. """ - 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, - ) + env_mask, body_mask = self._resolve_env_and_body_masks(env_mask, body_mask, body_ids) 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, @@ -833,15 +809,7 @@ 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. """ - 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, - ) + env_mask, body_mask = self._resolve_env_and_body_masks(env_mask, body_mask, body_ids) 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, @@ -1305,6 +1273,33 @@ def _resolve_body_ids(self, body_ids) -> wp.array | torch.Tensor: """ return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self.device) + def _resolve_env_and_body_masks( + self, + env_mask: wp.array | None, + body_mask: wp.array | None, + body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None, + ) -> tuple[wp.array, wp.array]: + """Resolve the environment and body selections of a mask writer to masks. + + Args: + env_mask: Environment mask. If None, then all environments are used. + body_mask: Body mask. If None, then ``body_ids`` or all bodies are used. + body_ids: Body indices, used when ``body_mask`` is None. If None, then all bodies are used. + + Returns: + The environment mask and the body mask. + """ + 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, + ) + return env_mask, body_mask + """ Internal simulation callbacks. """ 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 1d4066c2709a..d4fac59b9a4e 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 @@ -440,8 +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. """ - 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) + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) 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), @@ -528,8 +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. """ - 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) + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) 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), @@ -628,8 +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. """ - 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) + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) self.assert_shape_and_dtype( body_velocities, (self._num_instances, self._num_bodies), wp.spatial_vectorf, "body_velocities" ) @@ -735,8 +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. """ - 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) + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) self.assert_shape_and_dtype( body_velocities, (self._num_instances, self._num_bodies), wp.spatial_vectorf, "body_velocities" ) @@ -838,7 +838,7 @@ def set_masses_mask( env_mask: Environment mask. If None, all instances are updated. Shape is (num_instances,). """ - env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device) + env_ids = self._resolve_env_ids(env_mask=env_mask) 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, @@ -920,7 +920,7 @@ def set_coms_mask( env_mask: Environment mask. If None, all instances are updated. Shape is (num_instances,). """ - env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device) + env_ids = self._resolve_env_ids(env_mask=env_mask) 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, @@ -1012,7 +1012,7 @@ def set_inertias_mask( env_mask: Environment mask. If None, all instances are updated. Shape is (num_instances,). """ - env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device) + env_ids = self._resolve_env_ids(env_mask=env_mask) 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, @@ -1338,18 +1338,18 @@ def _binding_write( 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 on ``self._device``. Tests sometimes hand us indices on CPU even when the sim runs on GPU; we move the resolved array onto ``self._device`` so kernel launches don't fail on a device mismatch. """ - return resolve_ids(ids=env_ids, all_ids=self._ALL_ENV_INDICES, device=self._device) + return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device) - def _resolve_body_ids(self, body_ids) -> wp.array | torch.Tensor: + def _resolve_body_ids(self, body_ids=None, body_mask: wp.array | None = None) -> wp.array | torch.Tensor: """Resolve body indices on ``self._device``.""" - return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self._device) + return resolve_ids(ids=body_ids, mask=body_mask, 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" From bc9a8bbf78cf8ed43669f75fe48b1855d73065b5 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 15:24:36 -0700 Subject: [PATCH 69/80] Reject Newton masks whose length differs from the asset dimension The tendon and collection mask writers launch over the full dimension, so a short mask was read past its end. The host-side check follows the asset's shape-check switch and needs no device sync, so it stays capturable. --- .../test/assets/contract/test_articulation_contract.py | 6 ++++++ .../test/assets/contract/test_rigid_object_contract.py | 6 ++++++ .../changelog.d/jichuanh-id-mask.changed.rst | 3 +++ .../isaaclab_newton/assets/articulation/articulation.py | 7 ++++++- .../rigid_object_collection/rigid_object_collection.py | 5 +++++ 5 files changed, 26 insertions(+), 1 deletion(-) diff --git a/source/isaaclab/test/assets/contract/test_articulation_contract.py b/source/isaaclab/test/assets/contract/test_articulation_contract.py index abfb6d1eff57..935791ea2a0d 100644 --- a/source/isaaclab/test/assets/contract/test_articulation_contract.py +++ b/source/isaaclab/test/assets/contract/test_articulation_contract.py @@ -524,6 +524,12 @@ def read() -> ProxyArray: with pytest.raises((AssertionError, RuntimeError)): method(**{kwarg: wp.from_torch(extra_env.contiguous(), dtype=dtype)}) + if backend == "newton" and selection == "mask" and not writer.endswith("_target"): + # A mask shorter than the asset dimension is rejected, even with data of the mask's length. + short_env = _payload((N - 1, *shape[1:]), trailing, dtype, device) + with pytest.raises(AssertionError, match="Mask length"): + method(**{kwarg: short_env, "env_mask": _mask(N - 1, [0], device)}) + # quantity -> (writer, keyword, trailing size). Joint partial writes are covered by the ordering contract. _PARTIAL_WRITES = { 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 49fae35e4b0d..b8ba83e6146b 100644 --- a/source/isaaclab/test/assets/contract/test_rigid_object_contract.py +++ b/source/isaaclab/test/assets/contract/test_rigid_object_contract.py @@ -304,6 +304,12 @@ def test_writer_reads_back_full_and_partial_writes(monkeypatch, backend, device, method(**{kwarg: update}, env_mask=selectors["env_mask"], body_ids=body_ids) expected[-1, -1] = update[-1, -1] torch.testing.assert_close(getattr(obj.data, getter).torch, expected, atol=1e-5, rtol=1e-5) + if backend == "newton": + # A mask shorter than the asset dimension is rejected, even with data of the mask's length. + short_env = _payload((_NUM_INSTANCES - 1, num_bodies), trailing, device=device) + short_mask = wp.ones(_NUM_INSTANCES - 1, dtype=wp.bool, device=device) + with pytest.raises(AssertionError, match="Mask length"): + method(**{kwarg: short_env, "env_mask": short_mask}) # Data with an extra environment is rejected for torch and warp inputs. bad = _payload((_NUM_INSTANCES + 1, *shape[1:]), trailing, device=device) diff --git a/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst index e0e59800bd6d..258dfee34a3e 100644 --- a/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst +++ b/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst @@ -12,3 +12,6 @@ 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. Recorded into a CUDA graph, ``write_fixed_tendon_properties_to_sim_mask`` refreshes the solver's tendon properties on every replay. +* Changed the mask writers of :class:`~isaaclab_newton.assets.Articulation` and the pose and velocity mask writers of + :class:`~isaaclab_newton.assets.RigidObjectCollection` to reject a mask whose length differs from the asset + dimension it selects, while shape checks are enabled. diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index 879e42a4be21..1cd940064aaa 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -3823,7 +3823,12 @@ def _resolve_mask(self, mask: wp.array | torch.Tensor | None, full_mask: wp.arra Returns: A warp array of mask. """ - return resolve_1d_mask(mask=mask, all_mask=full_mask, device=self.device) + mask = resolve_1d_mask(mask=mask, all_mask=full_mask, device=self.device) + if self._check_shapes: + assert mask.shape[0] == full_mask.shape[0], ( + f"{type(self).__name__}: Mask length mismatch: {mask.shape[0]} != {full_mask.shape[0]}" + ) + return mask """ Deprecated methods. 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 311683ba4446..1aa24cb050a8 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 @@ -1298,6 +1298,11 @@ def _resolve_env_and_body_masks( scratch_mask=self._scratch_body_mask, device=self.device, ) + if self._check_shapes: + for mask, full_mask in ((env_mask, self._ALL_ENV_MASK), (body_mask, self._ALL_BODY_MASK)): + assert mask.shape[0] == full_mask.shape[0], ( + f"{type(self).__name__}: Mask length mismatch: {mask.shape[0]} != {full_mask.shape[0]}" + ) return env_mask, body_mask """ From 981e3a418ee61e87a2b39f33b2419df03e8677d0 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 15:25:36 -0700 Subject: [PATCH 70/80] Accept strided torch limits in the Newton tendon position limit mask setter A vec2f kernel argument needs contiguous components, so a torch view with a strided last dimension is made contiguous before the launch, as the index setter's torch assignment accepted it. --- .../test/assets/contract/test_articulation_contract.py | 6 ++++++ .../isaaclab_newton/assets/articulation/articulation.py | 3 +++ 2 files changed, 9 insertions(+) diff --git a/source/isaaclab/test/assets/contract/test_articulation_contract.py b/source/isaaclab/test/assets/contract/test_articulation_contract.py index 935791ea2a0d..11e647cd55bb 100644 --- a/source/isaaclab/test/assets/contract/test_articulation_contract.py +++ b/source/isaaclab/test/assets/contract/test_articulation_contract.py @@ -518,6 +518,12 @@ def read() -> ProxyArray: method(**{kwarg: 1.0}) _assert_reads_back(read(), torch.ones_like(expected), f"{writer} scalar") + if writer == "set_fixed_tendon_position_limit": + # Limit pairs whose components are strided in memory write as contiguous ones do. + expected = _payload(shape, trailing, dtype, device, offset=200.0) + method(**{kwarg: expected.repeat_interleave(2, dim=-1)[..., ::2]}) + _assert_reads_back(read(), expected, f"{writer} strided") + extra_env = _payload((N + 1, *shape[1:]), trailing, dtype, device) with pytest.raises((AssertionError, RuntimeError)): method(**{kwarg: extra_env}) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index 1cd940064aaa..8c8729f51ca9 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -2938,6 +2938,9 @@ def set_fixed_tendon_position_limit_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.") + if isinstance(limit, torch.Tensor): + # a vec2f kernel argument needs contiguous components; a contiguous tensor is returned as is + limit = limit.contiguous() wp.launch( shared_kernels.write_2d_data_to_buffer_with_mask_kernel(wp.vec2f), dim=(self.num_instances, self.num_fixed_tendons), From cff8f86b83678c357d4c978b881136b15abd1cd6 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 15:27:27 -0700 Subject: [PATCH 71/80] Run the tendon replay test before the module's shared scene opens The test builds its own simulation context, which fails once the module-scoped scene holds one. --- .../test/assets/test_articulation.py | 44 +++++++++---------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/source/isaaclab_newton/test/assets/test_articulation.py b/source/isaaclab_newton/test/assets/test_articulation.py index c226f6fbcd40..344967e0599b 100644 --- a/source/isaaclab_newton/test/assets/test_articulation.py +++ b/source/isaaclab_newton/test/assets/test_articulation.py @@ -663,6 +663,28 @@ def test_hard_reset_recreates_ordered_state_and_actuators(device: str) -> None: assert not torch.equal(articulation.data.joint_pos.torch, position) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +def test_fixed_tendon_mask_writer_reaches_solver_on_every_replay(device: str) -> None: + """Each replay of a CUDA graph holding the fixed-tendon mask writer refreshes the solver's tendon properties. + + The solver keeps its own copy of them, so a refresh queued only while capturing would leave the physics on the + values of the first replay. + """ + 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() + stiffness = torch.zeros(articulation.num_instances, articulation.num_fixed_tendons, device=device) + with wp.ScopedCapture(device) as capture: + articulation.set_fixed_tendon_stiffness_mask(stiffness=stiffness) + articulation.write_fixed_tendon_properties_to_sim_mask() + solver_model = sim.backend_manager.backend.solver.mjw_model + for value in (12.0, 34.0): + stiffness.fill_(value) + wp.capture_launch(capture.graph) + sim.step() + np.testing.assert_allclose(solver_model.tendon_stiffness.numpy(), value) + + @pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) def test_newton_ordered_state_publishes_inside_captured_cuda_graph(device: str) -> None: """Republish the user-order state shadows from inside a captured CUDA graph. @@ -1696,28 +1718,6 @@ def test_hand_with_tendons_targets_only_given_envs_and_properties_reach_solver(s articulation.write_fixed_tendon_properties_to_sim_index() -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -def test_fixed_tendon_mask_writer_reaches_solver_on_every_replay(device: str) -> None: - """Each replay of a CUDA graph holding the fixed-tendon mask writer refreshes the solver's tendon properties. - - The solver keeps its own copy of them, so a refresh queued only while capturing would leave the physics on the - values of the first replay. - """ - 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() - stiffness = torch.zeros(articulation.num_instances, articulation.num_fixed_tendons, device=device) - with wp.ScopedCapture(device) as capture: - articulation.set_fixed_tendon_stiffness_mask(stiffness=stiffness) - articulation.write_fixed_tendon_properties_to_sim_mask() - solver_model = sim.backend_manager.backend.solver.mjw_model - for value in (12.0, 34.0): - stiffness.fill_(value) - wp.capture_launch(capture.graph) - sim.step() - np.testing.assert_allclose(solver_model.tendon_stiffness.numpy(), value) - - @pytest.mark.isaacsim_ci @pytest.mark.parametrize("island", ["fixed", "floating", "ordered"], ids=["fixed", "floating", "floating_reordered"]) def test_get_gravity_compensation_forces_matches_jacobian_gravity(scene: _Scene, island: str) -> None: From bbb5104d959194e3196f6cf4c481e9f865474895 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 16:20:19 -0700 Subject: [PATCH 72/80] Revert "Reject Newton masks whose length differs from the asset dimension" This reverts commit bc9a8bbf78cf8ed43669f75fe48b1855d73065b5. --- .../test/assets/contract/test_articulation_contract.py | 6 ------ .../test/assets/contract/test_rigid_object_contract.py | 6 ------ .../changelog.d/jichuanh-id-mask.changed.rst | 3 --- .../isaaclab_newton/assets/articulation/articulation.py | 7 +------ .../rigid_object_collection/rigid_object_collection.py | 5 ----- 5 files changed, 1 insertion(+), 26 deletions(-) diff --git a/source/isaaclab/test/assets/contract/test_articulation_contract.py b/source/isaaclab/test/assets/contract/test_articulation_contract.py index 11e647cd55bb..be7339eca266 100644 --- a/source/isaaclab/test/assets/contract/test_articulation_contract.py +++ b/source/isaaclab/test/assets/contract/test_articulation_contract.py @@ -530,12 +530,6 @@ def read() -> ProxyArray: with pytest.raises((AssertionError, RuntimeError)): method(**{kwarg: wp.from_torch(extra_env.contiguous(), dtype=dtype)}) - if backend == "newton" and selection == "mask" and not writer.endswith("_target"): - # A mask shorter than the asset dimension is rejected, even with data of the mask's length. - short_env = _payload((N - 1, *shape[1:]), trailing, dtype, device) - with pytest.raises(AssertionError, match="Mask length"): - method(**{kwarg: short_env, "env_mask": _mask(N - 1, [0], device)}) - # quantity -> (writer, keyword, trailing size). Joint partial writes are covered by the ordering contract. _PARTIAL_WRITES = { 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 b8ba83e6146b..49fae35e4b0d 100644 --- a/source/isaaclab/test/assets/contract/test_rigid_object_contract.py +++ b/source/isaaclab/test/assets/contract/test_rigid_object_contract.py @@ -304,12 +304,6 @@ def test_writer_reads_back_full_and_partial_writes(monkeypatch, backend, device, method(**{kwarg: update}, env_mask=selectors["env_mask"], body_ids=body_ids) expected[-1, -1] = update[-1, -1] torch.testing.assert_close(getattr(obj.data, getter).torch, expected, atol=1e-5, rtol=1e-5) - if backend == "newton": - # A mask shorter than the asset dimension is rejected, even with data of the mask's length. - short_env = _payload((_NUM_INSTANCES - 1, num_bodies), trailing, device=device) - short_mask = wp.ones(_NUM_INSTANCES - 1, dtype=wp.bool, device=device) - with pytest.raises(AssertionError, match="Mask length"): - method(**{kwarg: short_env, "env_mask": short_mask}) # Data with an extra environment is rejected for torch and warp inputs. bad = _payload((_NUM_INSTANCES + 1, *shape[1:]), trailing, device=device) diff --git a/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst index 258dfee34a3e..e0e59800bd6d 100644 --- a/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst +++ b/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst @@ -12,6 +12,3 @@ 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. Recorded into a CUDA graph, ``write_fixed_tendon_properties_to_sim_mask`` refreshes the solver's tendon properties on every replay. -* Changed the mask writers of :class:`~isaaclab_newton.assets.Articulation` and the pose and velocity mask writers of - :class:`~isaaclab_newton.assets.RigidObjectCollection` to reject a mask whose length differs from the asset - dimension it selects, while shape checks are enabled. diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index 8c8729f51ca9..b56e17176357 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -3826,12 +3826,7 @@ def _resolve_mask(self, mask: wp.array | torch.Tensor | None, full_mask: wp.arra Returns: A warp array of mask. """ - mask = resolve_1d_mask(mask=mask, all_mask=full_mask, device=self.device) - if self._check_shapes: - assert mask.shape[0] == full_mask.shape[0], ( - f"{type(self).__name__}: Mask length mismatch: {mask.shape[0]} != {full_mask.shape[0]}" - ) - 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/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py index 1aa24cb050a8..311683ba4446 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 @@ -1298,11 +1298,6 @@ def _resolve_env_and_body_masks( scratch_mask=self._scratch_body_mask, device=self.device, ) - if self._check_shapes: - for mask, full_mask in ((env_mask, self._ALL_ENV_MASK), (body_mask, self._ALL_BODY_MASK)): - assert mask.shape[0] == full_mask.shape[0], ( - f"{type(self).__name__}: Mask length mismatch: {mask.shape[0]} != {full_mask.shape[0]}" - ) return env_mask, body_mask """ From 029ee47a41e08fcf92ddb27dd2dc97e8a51e19be Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 16:20:32 -0700 Subject: [PATCH 73/80] Resolve the Newton collection writer masks inline again Each mask writer stays self-contained, as the articulation writers are. --- .../rigid_object_collection.py | 67 ++++++++++--------- 1 file changed, 36 insertions(+), 31 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 311683ba4446..ad916dbfc09e 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 @@ -473,7 +473,15 @@ 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. """ - env_mask, body_mask = self._resolve_env_and_body_masks(env_mask, body_mask, body_ids) + 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), @@ -579,7 +587,15 @@ 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. """ - env_mask, body_mask = self._resolve_env_and_body_masks(env_mask, body_mask, body_ids) + 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, @@ -693,7 +709,15 @@ 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. """ - env_mask, body_mask = self._resolve_env_and_body_masks(env_mask, body_mask, body_ids) + 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, @@ -809,7 +833,15 @@ 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. """ - env_mask, body_mask = self._resolve_env_and_body_masks(env_mask, body_mask, body_ids) + 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, @@ -1273,33 +1305,6 @@ def _resolve_body_ids(self, body_ids) -> wp.array | torch.Tensor: """ return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self.device) - def _resolve_env_and_body_masks( - self, - env_mask: wp.array | None, - body_mask: wp.array | None, - body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None, - ) -> tuple[wp.array, wp.array]: - """Resolve the environment and body selections of a mask writer to masks. - - Args: - env_mask: Environment mask. If None, then all environments are used. - body_mask: Body mask. If None, then ``body_ids`` or all bodies are used. - body_ids: Body indices, used when ``body_mask`` is None. If None, then all bodies are used. - - Returns: - The environment mask and the body mask. - """ - 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, - ) - return env_mask, body_mask - """ Internal simulation callbacks. """ From 551e508f73d5a2cfdb6ed9610af310136e79ceac Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 16:22:34 -0700 Subject: [PATCH 74/80] Launch the rewritten Newton mask writers over the mask lengths The fixed-tendon and collection body mask writers launched over the asset dimensions, so a mask shorter than the asset was read past its end. They now launch over the masks, as the other Newton mask writers do, and a short mask writes only the rows it covers. --- .../test/assets/contract/test_articulation_contract.py | 8 ++++++++ .../test/assets/contract/test_rigid_object_contract.py | 7 +++++++ .../isaaclab_newton/assets/articulation/articulation.py | 8 ++++---- .../rigid_object_collection/rigid_object_collection.py | 8 ++++---- 4 files changed, 23 insertions(+), 8 deletions(-) diff --git a/source/isaaclab/test/assets/contract/test_articulation_contract.py b/source/isaaclab/test/assets/contract/test_articulation_contract.py index be7339eca266..3dea2cce14d2 100644 --- a/source/isaaclab/test/assets/contract/test_articulation_contract.py +++ b/source/isaaclab/test/assets/contract/test_articulation_contract.py @@ -524,6 +524,14 @@ def read() -> ProxyArray: method(**{kwarg: expected.repeat_interleave(2, dim=-1)[..., ::2]}) _assert_reads_back(read(), expected, f"{writer} strided") + if backend == "newton" and selection == "mask" and not writer.endswith("_target"): + # A short mask writes only the rows it covers, even when the memory past its end selects more rows. + expected = read().torch.clone() + values = _payload(shape, trailing, dtype, device, offset=300.0) + method(**{kwarg: values[: N - 1], "env_mask": _mask(N, range(N), device)[: N - 1]}) + expected[: N - 1] = values[: N - 1] + _assert_reads_back(read(), expected, f"{writer} short mask") + extra_env = _payload((N + 1, *shape[1:]), trailing, dtype, device) with pytest.raises((AssertionError, RuntimeError)): method(**{kwarg: extra_env}) 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 49fae35e4b0d..cd800bd91696 100644 --- a/source/isaaclab/test/assets/contract/test_rigid_object_contract.py +++ b/source/isaaclab/test/assets/contract/test_rigid_object_contract.py @@ -304,6 +304,13 @@ def test_writer_reads_back_full_and_partial_writes(monkeypatch, backend, device, method(**{kwarg: update}, env_mask=selectors["env_mask"], body_ids=body_ids) expected[-1, -1] = update[-1, -1] torch.testing.assert_close(getattr(obj.data, getter).torch, expected, atol=1e-5, rtol=1e-5) + if backend == "newton": + # A short mask writes only the rows it covers, even when the memory past its end selects more rows. + update = _payload(shape, trailing, offset=500.0, device=device) + short_mask = wp.ones(_NUM_INSTANCES, dtype=wp.bool, device=device)[: _NUM_INSTANCES - 1] + method(**{kwarg: update[: _NUM_INSTANCES - 1]}, env_mask=short_mask) + expected[: _NUM_INSTANCES - 1] = update[: _NUM_INSTANCES - 1] + torch.testing.assert_close(getattr(obj.data, getter).torch, expected, atol=1e-5, rtol=1e-5) # Data with an extra environment is rejected for torch and warp inputs. bad = _payload((_NUM_INSTANCES + 1, *shape[1:]), trailing, device=device) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index b56e17176357..259c36e79bbf 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -2716,7 +2716,7 @@ def set_fixed_tendon_stiffness_mask( 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), + dim=(env_mask.shape[0], fixed_tendon_mask.shape[0]), inputs=[stiffness, env_mask, fixed_tendon_mask], outputs=[self.data._fixed_tendon_stiffness], device=self.device, @@ -2818,7 +2818,7 @@ def set_fixed_tendon_damping_mask( 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), + dim=(env_mask.shape[0], fixed_tendon_mask.shape[0]), inputs=[damping, env_mask, fixed_tendon_mask], outputs=[self.data._fixed_tendon_damping], device=self.device, @@ -2943,7 +2943,7 @@ def set_fixed_tendon_position_limit_mask( limit = limit.contiguous() wp.launch( shared_kernels.write_2d_data_to_buffer_with_mask_kernel(wp.vec2f), - dim=(self.num_instances, self.num_fixed_tendons), + dim=(env_mask.shape[0], fixed_tendon_mask.shape[0]), inputs=[limit, env_mask, fixed_tendon_mask], outputs=[self.data._fixed_tendon_pos_limits], device=self.device, @@ -3151,7 +3151,7 @@ def write_fixed_tendon_properties_to_sim_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), + dim=(env_mask.shape[0], fixed_tendon_mask.shape[0]), inputs=[ env_mask, fixed_tendon_mask, 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 ad916dbfc09e..2e53f72d7d86 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 @@ -485,7 +485,7 @@ def write_body_link_pose_to_sim_mask( 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), + dim=(env_mask.shape[0], body_mask.shape[0]), inputs=[body_poses, env_mask, body_mask], outputs=[self.data.body_link_pose_w], device=self.device, @@ -599,7 +599,7 @@ def write_body_com_pose_to_sim_mask( 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), + dim=(env_mask.shape[0], body_mask.shape[0]), 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, @@ -721,7 +721,7 @@ def write_body_com_velocity_to_sim_mask( 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), + dim=(env_mask.shape[0], body_mask.shape[0]), inputs=[body_velocities, env_mask, body_mask], outputs=[self.data.body_com_vel_w, self.data.body_com_acc_w], device=self.device, @@ -845,7 +845,7 @@ def write_body_link_velocity_to_sim_mask( 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), + dim=(env_mask.shape[0], body_mask.shape[0]), 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, From 86ce1c95fde0e188b3398fc1a4f9ee4af6ebc4bd Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 16:39:04 -0700 Subject: [PATCH 75/80] Write the rewritten Newton mask writers through a core masked write write_2d_with_mask launches the masked 2-D write, or the scalar fill, over the mask lengths. The Newton fixed-tendon setters and the collection link pose writer use it, and the Newton float32 masked writer is back to main's version. The core masked writer matches the one the PhysX tendon setters move to the core kernels. --- .../changelog.d/jichuanh-id-mask.added.rst | 2 + .../isaaclab/isaaclab/utils/warp/kernels.py | 59 +++++++++++++++++++ source/isaaclab/isaaclab/utils/warp/utils.py | 42 +++++++++++++ .../test/utils/warp/test_write_with_mask.py | 38 ++++++++++++ .../assets/articulation/articulation.py | 33 ++--------- .../isaaclab_newton/assets/kernels.py | 30 +++------- .../rigid_object_collection.py | 10 +--- 7 files changed, 156 insertions(+), 58 deletions(-) create mode 100644 source/isaaclab/test/utils/warp/test_write_with_mask.py diff --git a/source/isaaclab/changelog.d/jichuanh-id-mask.added.rst b/source/isaaclab/changelog.d/jichuanh-id-mask.added.rst index 935dc70f7fa4..9dfc5d1942ba 100644 --- a/source/isaaclab/changelog.d/jichuanh-id-mask.added.rst +++ b/source/isaaclab/changelog.d/jichuanh-id-mask.added.rst @@ -1,2 +1,4 @@ * 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. +* Added :func:`~isaaclab.utils.warp.utils.write_2d_with_mask`, which writes 2-D data or one scalar into a Warp array + where an environment mask and an item mask are both true, launching over the mask lengths. diff --git a/source/isaaclab/isaaclab/utils/warp/kernels.py b/source/isaaclab/isaaclab/utils/warp/kernels.py index 3d046abe2bba..bc9b63d11339 100644 --- a/source/isaaclab/isaaclab/utils/warp/kernels.py +++ b/source/isaaclab/isaaclab/utils/warp/kernels.py @@ -5,6 +5,7 @@ """Custom kernels for warp.""" +import functools from typing import Any import torch @@ -780,3 +781,61 @@ 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) + + +## +# Masked writes. +## + + +@wp.kernel +def write_2d_data_to_buffer_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), +): + """Write 2D data to a buffer where both the environment and the item mask are set. + + Args: + in_data: Input array containing data for every environment and item. Shape is (num_envs, num_items). + env_mask: Input array of environment mask. Shape is (num_envs,). + item_mask: Input array of item mask. Shape is (num_items,). + out_data: Output array where data is 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] + + +@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 dtype.""" + array = wp.array2d(dtype=dtype) + return wp.overload(write_2d_data_to_buffer_with_mask, {"in_data": array, "out_data": array}) + + +@wp.kernel +def fill_2d_buffer_with_mask( + value: Any, + env_mask: wp.array(dtype=wp.bool), + item_mask: wp.array(dtype=wp.bool), + out_data: wp.array2d(dtype=Any), +): + """Write one value to a buffer where both the environment and the item mask are set. + + Args: + value: Value written to every selected entry. + env_mask: Input array of environment mask. Shape is (num_envs,). + item_mask: Input array of item mask. Shape is (num_items,). + out_data: Output array where the value is written. Shape is (num_envs, num_items). + """ + i, j = wp.tid() + if env_mask[i] and item_mask[j]: + out_data[i, j] = value + + +@functools.cache +def fill_2d_buffer_with_mask_kernel(dtype: type) -> wp.Kernel: + """Select the masked 2-D buffer fill for the buffer dtype.""" + return wp.overload(fill_2d_buffer_with_mask, {"value": dtype, "out_data": wp.array2d(dtype=dtype)}) diff --git a/source/isaaclab/isaaclab/utils/warp/utils.py b/source/isaaclab/isaaclab/utils/warp/utils.py index a4d637692cc2..cb1197803a0d 100644 --- a/source/isaaclab/isaaclab/utils/warp/utils.py +++ b/source/isaaclab/isaaclab/utils/warp/utils.py @@ -11,6 +11,7 @@ import torch import warp as wp +from .kernels import fill_2d_buffer_with_mask_kernel, write_2d_data_to_buffer_with_mask_kernel from .proxy_array import ProxyArray ## @@ -140,6 +141,47 @@ def _mask_needs_copy(mask: wp.array | torch.Tensor, device: str) -> bool: return True +## +# Masked writes. +## + + +def write_2d_with_mask( + data: float | torch.Tensor | wp.array, + env_mask: wp.array, + item_mask: wp.array, + out: wp.array, + device: str, +) -> None: + """Write ``data`` into ``out`` where both masks are true, launching over the mask lengths. + + A mask shorter than ``out`` writes only the rows or columns it covers. Callers resolve the masks and check the + data shape. + + Args: + data: Values of shape ``(env_mask.shape[0], item_mask.shape[0])`` in the dtype of ``out``, as a ``wp.array`` + or a ``torch.Tensor`` with the dtype's trailing dimensions, or one scalar for every selected entry. + env_mask: Boolean mask over the first dimension of ``out``. + item_mask: Boolean mask over the second dimension of ``out``. + out: Output array. + device: Warp device string (e.g. ``"cuda:0"``). + """ + if isinstance(data, (int, float)): + kernel = fill_2d_buffer_with_mask_kernel(out.dtype) + else: + if isinstance(data, torch.Tensor) and data.ndim > 2: + # a vector-typed kernel argument needs contiguous components + data = data.contiguous() + kernel = write_2d_data_to_buffer_with_mask_kernel(out.dtype) + wp.launch( + kernel, + dim=(env_mask.shape[0], item_mask.shape[0]), + inputs=[data, env_mask, item_mask], + outputs=[out], + device=device, + ) + + ## # Id resolution - ids/mask to ids. ## diff --git a/source/isaaclab/test/utils/warp/test_write_with_mask.py b/source/isaaclab/test/utils/warp/test_write_with_mask.py new file mode 100644 index 000000000000..7909f05676f3 --- /dev/null +++ b/source/isaaclab/test/utils/warp/test_write_with_mask.py @@ -0,0 +1,38 @@ +# 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 writing through masks with :func:`write_2d_with_mask`.""" + +import pytest +import torch +import warp as wp + +from isaaclab.test.utils import test_devices +from isaaclab.utils.warp.utils import write_2d_with_mask + + +@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("data_type", ["scalar", "torch", "warp", "torch_strided_vec2f"]) +def test_write_2d_with_mask(device, data_type): + """Only entries selected by both masks are written, and a short mask writes only the rows it covers. + + The environment mask and the data are views whose memory continues with a selected third row, so a launch over the + full output would write that row too. + """ + torch_device = wp.device_to_torch(device) + env_mask = wp.array([True, False, True], dtype=wp.bool, device=device)[:2] + item_mask = wp.array([False, True], dtype=wp.bool, device=device) + if data_type == "torch_strided_vec2f": + values = torch.arange(1, 13, dtype=torch.float32, device=torch_device).reshape(3, 2, 2) + data = values.repeat_interleave(2, dim=-1)[..., ::2][:2] + out = wp.zeros((3, 2), dtype=wp.vec2f, device=device) + else: + values = torch.arange(1, 7, dtype=torch.float32, device=torch_device).reshape(3, 2) + data = {"scalar": 7.0, "torch": values[:2], "warp": wp.from_torch(values)[:2]}[data_type] + out = wp.zeros((3, 2), dtype=wp.float32, device=device) + write_2d_with_mask(data, env_mask, item_mask, out, device) + expected = torch.zeros_like(values) + expected[0, 1] = 7.0 if data_type == "scalar" else values[0, 1] + assert torch.equal(wp.to_torch(out), expected) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index 259c36e79bbf..5f88d6b82bb4 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -33,7 +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.warp.utils import resolve_1d_mask, resolve_ids, write_2d_with_mask from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_newton.assets import kernels as shared_kernels @@ -2712,15 +2712,7 @@ def set_fixed_tendon_stiffness_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) 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=(env_mask.shape[0], fixed_tendon_mask.shape[0]), - inputs=[stiffness, env_mask, fixed_tendon_mask], - outputs=[self.data._fixed_tendon_stiffness], - device=self.device, - ) + write_2d_with_mask(stiffness, env_mask, fixed_tendon_mask, self.data._fixed_tendon_stiffness, self.device) def set_fixed_tendon_damping_index( self, @@ -2814,15 +2806,7 @@ def set_fixed_tendon_damping_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) 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=(env_mask.shape[0], fixed_tendon_mask.shape[0]), - inputs=[damping, env_mask, fixed_tendon_mask], - outputs=[self.data._fixed_tendon_damping], - device=self.device, - ) + write_2d_with_mask(damping, env_mask, fixed_tendon_mask, self.data._fixed_tendon_damping, self.device) def set_fixed_tendon_limit_stiffness_index( self, @@ -2938,16 +2922,7 @@ def set_fixed_tendon_position_limit_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.") - if isinstance(limit, torch.Tensor): - # a vec2f kernel argument needs contiguous components; a contiguous tensor is returned as is - limit = limit.contiguous() - wp.launch( - shared_kernels.write_2d_data_to_buffer_with_mask_kernel(wp.vec2f), - dim=(env_mask.shape[0], fixed_tendon_mask.shape[0]), - inputs=[limit, env_mask, fixed_tendon_mask], - outputs=[self.data._fixed_tendon_pos_limits], - device=self.device, - ) + write_2d_with_mask(limit, env_mask, fixed_tendon_mask, self.data._fixed_tendon_pos_limits, self.device) def set_fixed_tendon_rest_length_index( self, diff --git a/source/isaaclab_newton/isaaclab_newton/assets/kernels.py b/source/isaaclab_newton/isaaclab_newton/assets/kernels.py index 3a232ea7ea1f..5733b4d55888 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/kernels.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/kernels.py @@ -5,7 +5,6 @@ from __future__ import annotations -import functools from typing import TYPE_CHECKING, Any import torch @@ -1122,18 +1121,21 @@ def write_2d_data_to_buffer_with_indices_kernel( @wp.kernel -def _write_2d_data_to_buffer_with_mask( - in_data: wp.array2d(dtype=Any), +def write_2d_data_to_buffer_with_mask( + in_data: wp.array2d(dtype=wp.float32), env_mask: wp.array(dtype=wp.bool), joint_mask: wp.array(dtype=wp.bool), - out_data: wp.array2d(dtype=Any), + out_data: wp.array2d(dtype=wp.float32), ): - """Write 2D data to a buffer where both masks are true. + """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. Args: - in_data: Input array. Shape is (num_instances, num_joints). + in_data: Input array containing float data. Shape is (num_instances, num_joints). env_mask: Input array of environment mask. Shape is (num_instances,). - joint_mask: Input array of joint/body/tendon mask. Shape is (num_joints,). + 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). """ i, j = wp.tid() @@ -1141,20 +1143,6 @@ 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/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py index 2e53f72d7d86..7b37424d4a94 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,7 +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.warp.utils import resolve_1d_mask, resolve_ids, write_2d_with_mask from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_newton.assets import kernels as shared_kernels @@ -483,13 +483,7 @@ def write_body_link_pose_to_sim_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=(env_mask.shape[0], body_mask.shape[0]), - inputs=[body_poses, env_mask, body_mask], - outputs=[self.data.body_link_pose_w], - device=self.device, - ) + write_2d_with_mask(body_poses, env_mask, body_mask, self.data.body_link_pose_w, self.device) if not skip_forward: self.data._reset_pose(env_mask=env_mask) From 82587c8c6014646f16813b82783e95c44d329023 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 18:21:40 -0700 Subject: [PATCH 76/80] Raise the modifiers deprecation warning at the caller of the Warp observation manager --- .../isaaclab_experimental/managers/observation_manager.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py index 65ca4e42da08..b359b265d9f1 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py @@ -776,6 +776,8 @@ def _prepare_terms(self): # noqa: C901 f"Configuration for the term '{term_name}' is not of type ObservationTermCfg." f" Received: '{type(term_cfg)}'." ) + # migrated here so the deprecation warning points at the code that built the manager + term_cfg._migrate_modifiers() term_dim = self._prepare_term(group_name, term_name, term_cfg, group_cfg) # add term config to list self._group_obs_term_names[group_name].append(term_name) From ccfe282d5dcbbf071001d0d14b943d3fc2308ff5 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Wed, 28 Oct 2026 16:23:20 -0700 Subject: [PATCH 77/80] Import the root-to-body velocity helpers from the shared Warp math Main removed isaaclab_newton.kernels.state_kernels; body_lin_vel_from_root and body_ang_vel_from_root now live in isaaclab.utils.warp.math. --- .../isaaclab_experimental/envs/mdp/commands/velocity_command.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py index 130ab134bd41..cbd522c36dd5 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py @@ -13,10 +13,10 @@ import torch import warp as wp -from isaaclab_newton.kernels.state_kernels import body_ang_vel_from_root, body_lin_vel_from_root from isaaclab.envs.mdp.commands._debug_vis import _VelocityCommandDebugVis from isaaclab.utils.seed import WarpRng +from isaaclab.utils.warp.math import body_ang_vel_from_root, body_lin_vel_from_root from isaaclab_experimental.managers import CommandTerm from isaaclab_experimental.utils.warp import wrap_to_pi From dc2aba6ae591033ea7eeee4e195012c71aa101c5 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Thu, 29 Oct 2026 02:03:59 -0700 Subject: [PATCH 78/80] Align resolve_ids and write_2d_with_mask with the shared memory helpers Ids resolve on the device of all_ids and keep their dtype; callers that need contiguous int32 ids pass them through memory.resolve_1d_indices, as main's indexed writers do. write_2d_with_mask launches on the device of its output. The deformable mask tests check the backend writes. --- .../base_deformable_object.py | 8 +-- .../isaaclab/sensors/camera/camera.py | 4 +- source/isaaclab/isaaclab/utils/warp/memory.py | 62 +++++-------------- .../isaaclab/utils/wrench_composer.py | 4 +- .../test/utils/warp/test_resolve_selection.py | 46 +++++--------- .../test/utils/warp/test_write_with_mask.py | 2 +- .../assets/articulation/articulation.py | 16 ++--- .../assets/cable_object/cable_object.py | 4 +- .../deformable_object/deformable_object.py | 4 +- .../assets/mpm_object/mpm_object.py | 4 +- .../assets/rigid_object/rigid_object.py | 4 +- .../rigid_object_collection.py | 6 +- .../assets/articulation/articulation.py | 10 +-- .../deformable_object/deformable_object.py | 4 +- .../assets/rigid_object/rigid_object.py | 4 +- .../rigid_object_collection.py | 4 +- .../test/assets/test_deformable_object.py | 5 ++ .../assets/articulation/articulation.py | 17 ++--- .../deformable_object/deformable_object.py | 12 ++-- .../assets/rigid_object/rigid_object.py | 4 +- .../rigid_object_collection.py | 4 +- .../assets/surface_gripper/surface_gripper.py | 4 +- .../test/assets/test_deformable_object.py | 6 ++ 23 files changed, 97 insertions(+), 141 deletions(-) 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 2144ccbad474..42ce490c5a67 100644 --- a/source/isaaclab/isaaclab/assets/deformable_object/base_deformable_object.py +++ b/source/isaaclab/isaaclab/assets/deformable_object/base_deformable_object.py @@ -218,7 +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. """ - env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES) self.write_nodal_state_to_sim_index(nodal_state, env_ids=env_ids, full_data=True) def write_nodal_pos_to_sim_mask( @@ -233,7 +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. """ - env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES) self.write_nodal_pos_to_sim_index(nodal_pos, env_ids=env_ids, full_data=True) def write_nodal_velocity_to_sim_mask( @@ -248,7 +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. """ - env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES) self.write_nodal_velocity_to_sim_index(nodal_vel, env_ids=env_ids, full_data=True) def write_nodal_kinematic_target_to_sim_mask( @@ -263,7 +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. """ - env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES) self.write_nodal_kinematic_target_to_sim_index(targets, env_ids=env_ids, full_data=True) """ diff --git a/source/isaaclab/isaaclab/sensors/camera/camera.py b/source/isaaclab/isaaclab/sensors/camera/camera.py index 56c0e2f6dd2c..03d6f12eea4b 100644 --- a/source/isaaclab/isaaclab/sensors/camera/camera.py +++ b/source/isaaclab/isaaclab/sensors/camera/camera.py @@ -26,7 +26,7 @@ quat_from_matrix, ) from ...utils.warp import ProxyArray -from ...utils.warp.memory import resolve_ids +from ...utils.warp.memory import resolve_1d_indices, resolve_ids from ..sensor_base import SensorBase from .camera_data import CameraData, RenderBufferKind @@ -1029,7 +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 - return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self._device, dtype=wp.int32) + return resolve_1d_indices(resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES)) def _env_mask_has_any(self, env_mask: wp.array) -> bool: """Return whether the mask selects any camera.""" diff --git a/source/isaaclab/isaaclab/utils/warp/memory.py b/source/isaaclab/isaaclab/utils/warp/memory.py index a55fa0618fb1..cfebcfde0b64 100644 --- a/source/isaaclab/isaaclab/utils/warp/memory.py +++ b/source/isaaclab/isaaclab/utils/warp/memory.py @@ -296,7 +296,6 @@ def write_2d_with_mask( row_mask: wp.array, column_mask: wp.array, out: wp.array, - device: str, ) -> None: """Write ``data`` into ``out`` where both masks are true, launching over the mask lengths. @@ -308,8 +307,7 @@ def write_2d_with_mask( or a ``torch.Tensor`` with the dtype's trailing dimensions, or one scalar for every selected entry. row_mask: Boolean mask over the first dimension of ``out``. column_mask: Boolean mask over the second dimension of ``out``. - out: Output array. - device: Warp device string (e.g. ``"cuda:0"``). + out: Output array. The launch runs on its device. """ if isinstance(data, (int, float)): kernel = make_masked_fill(out.dtype, 2) @@ -323,7 +321,7 @@ def write_2d_with_mask( dim=(row_mask.shape[0], column_mask.shape[0]), inputs=[data, row_mask, column_mask], outputs=[out], - device=device, + device=out.device, ) @@ -454,55 +452,40 @@ 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. + """Resolve ids/mask into ids on the device of ``all_ids``. The mask takes precedence over the ids. + + Pass the result to :func:`resolve_1d_indices` where contiguous int32 ids are needed. 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``. + or ``None`` (all elements). Tensors and warp arrays keep their integer dtype. 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``. + ``all_ids``, a ``torch.Tensor`` or a ``wp.array`` of ids on the device of ``all_ids``. + A mask resolves to an int64 ``torch.Tensor``; lists and slices resolve to ``wp.int32``. 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. + TypeError: If *ids* is a :class:`ProxyArray`, which is output-only, or *mask* is + neither a ``wp.array`` nor a ``torch.Tensor``. """ + device = all_ids.device 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``.""" + return torch.nonzero(mask.to(wp.device_to_torch(device))).squeeze(-1) if isinstance(ids, torch.Tensor): - return ids if str(ids.device) == str(wp.get_device(device)) else ids.to(wp.device_to_torch(device)) + return ids if str(ids.device) == str(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) + return ids if ids.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): @@ -514,20 +497,3 @@ def _ids_on_device( 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) diff --git a/source/isaaclab/isaaclab/utils/wrench_composer.py b/source/isaaclab/isaaclab/utils/wrench_composer.py index f9ac3263b3be..49e95b91ab5b 100644 --- a/source/isaaclab/isaaclab/utils/wrench_composer.py +++ b/source/isaaclab/isaaclab/utils/wrench_composer.py @@ -724,7 +724,7 @@ def _resolve_env_ids(self, env_ids: wp.array | torch.Tensor | list | slice | Non Raises: TypeError: If ``env_ids`` is a :class:`~isaaclab.utils.warp.ProxyArray`. """ - return resolve_ids(ids=env_ids, all_ids=self._ALL_ENV_INDICES, device=self.device) + return resolve_ids(ids=env_ids, all_ids=self._ALL_ENV_INDICES) def _resolve_body_ids( self, body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None @@ -740,7 +740,7 @@ def _resolve_body_ids( Raises: TypeError: If ``body_ids`` is a :class:`~isaaclab.utils.warp.ProxyArray`. """ - return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self.device) + return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES) def _ensure_composed(self): """Compose input buffers into output buffers if dirty.""" diff --git a/source/isaaclab/test/utils/warp/test_resolve_selection.py b/source/isaaclab/test/utils/warp/test_resolve_selection.py index 6bc5c013e0cc..804d0879b0a2 100644 --- a/source/isaaclab/test/utils/warp/test_resolve_selection.py +++ b/source/isaaclab/test/utils/warp/test_resolve_selection.py @@ -11,7 +11,7 @@ from isaaclab.test.utils import DeviceScope, test_devices from isaaclab.utils.warp import ProxyArray -from isaaclab.utils.warp.memory import resolve_1d_mask, resolve_ids +from isaaclab.utils.warp.memory import resolve_1d_indices, resolve_1d_mask, resolve_ids def _all_ids(n: int, device: str) -> wp.array: @@ -19,19 +19,17 @@ def _all_ids(n: int, device: str) -> wp.array: @pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("dtype", [None, wp.int32]) -def test_resolve_ids(device, dtype): +def test_resolve_ids(device): """Every selection resolves to device ids; a selection of every element is the shared constant itself. - Without ``dtype``, tensors and Warp arrays pass through and positive-step slices view the shared ids without a - copy. ``dtype=wp.int32`` gives a contiguous one-dimensional int32 Warp array for every kind of ids. + Tensors and Warp arrays pass through, and positive-step slices view the shared ids without a copy. """ all_ids = _all_ids(10, device) torch_device = wp.device_to_torch(device) ids_int32 = wp.array([5], dtype=wp.int32, device=device) for ids in (None, slice(None)): - assert resolve_ids(ids=ids, all_ids=all_ids, device=device, dtype=dtype) is all_ids - assert resolve_ids(ids=ids_int32, all_ids=all_ids, device=device, dtype=dtype) is ids_int32 + assert resolve_ids(ids=ids, all_ids=all_ids) is all_ids + assert resolve_ids(ids=ids_int32, all_ids=all_ids) is ids_int32 for ids, expected in ( ([3, 8], [3, 8]), (slice(2, 8, 3), [2, 5]), @@ -40,40 +38,28 @@ def test_resolve_ids(device, dtype): (torch.arange(10, device=torch_device)[::3], [0, 3, 6, 9]), (wp.array([7, 1], dtype=wp.int64, device=device), [7, 1]), ): - resolved = resolve_ids(ids=ids, all_ids=all_ids, device=device, dtype=dtype) - if dtype is None: - assert resolved is ids or isinstance(ids, (list, slice)) - else: - assert isinstance(resolved, wp.array) and resolved.dtype == dtype and resolved.is_contiguous + resolved = resolve_ids(ids=ids, all_ids=all_ids) + assert resolved is ids or isinstance(ids, (list, slice)) assert (wp.to_torch(resolved) if isinstance(resolved, wp.array) else resolved).tolist() == expected with pytest.raises(TypeError, match="ProxyArray"): - resolve_ids(ids=ProxyArray(ids_int32), all_ids=all_ids, device=device, dtype=dtype) - if dtype is None: - cached = wp.to_torch(all_ids) - for selection in (slice(1, None, 2), slice(0, 0)): - resolved = wp.to_torch(resolve_ids(ids=selection, all_ids=all_ids, device=device)) - assert resolved.data_ptr() == cached[selection].data_ptr() - assert resolved.stride() == cached[selection].stride() - else: - for ids, error in ( - (torch.tensor([1.0], device=torch_device), TypeError), - (wp.array([1.0], dtype=wp.float32, device=device), TypeError), - (torch.zeros((2, 2), dtype=torch.int64, device=torch_device), ValueError), - ): - with pytest.raises(error): - resolve_ids(ids=ids, all_ids=all_ids, device=device, dtype=dtype) + resolve_ids(ids=ProxyArray(ids_int32), all_ids=all_ids) + cached = wp.to_torch(all_ids) + for selection in (slice(1, None, 2), slice(0, 0)): + resolved = wp.to_torch(resolve_ids(ids=selection, all_ids=all_ids)) + assert resolved.data_ptr() == cached[selection].data_ptr() + assert resolved.stride() == cached[selection].stride() @pytest.mark.parametrize("device", test_devices()) @pytest.mark.parametrize("mask_type", ["warp", "torch"]) def test_resolve_ids_from_a_mask(device, mask_type): - """A mask resolves to its ascending ids, takes precedence over ids, and honours ``dtype``.""" + """A mask resolves to its ascending ids and takes precedence over ids; :func:`resolve_1d_indices` casts to int32.""" mask_bool = torch.tensor([0, 1, 1, 0, 1, 0, 0, 1], dtype=torch.bool, device=wp.device_to_torch(device)) mask = wp.from_torch(mask_bool, dtype=wp.bool) if mask_type == "warp" else mask_bool all_ids = _all_ids(mask_bool.shape[0], device) - ids = resolve_ids(ids=[0, 1], mask=mask, all_ids=all_ids, device=device) + ids = resolve_ids(ids=[0, 1], mask=mask, all_ids=all_ids) assert ids.dtype == torch.int64 and ids.tolist() == [1, 2, 4, 7] - ids_int32 = resolve_ids(mask=mask, all_ids=all_ids, device=device, dtype=wp.int32) + ids_int32 = resolve_1d_indices(resolve_ids(mask=mask, all_ids=all_ids)) assert isinstance(ids_int32, wp.array) and ids_int32.dtype == wp.int32 and ids_int32.is_contiguous assert ids_int32.numpy().tolist() == [1, 2, 4, 7] diff --git a/source/isaaclab/test/utils/warp/test_write_with_mask.py b/source/isaaclab/test/utils/warp/test_write_with_mask.py index 4fb308d76a86..086e4e8f5143 100644 --- a/source/isaaclab/test/utils/warp/test_write_with_mask.py +++ b/source/isaaclab/test/utils/warp/test_write_with_mask.py @@ -32,7 +32,7 @@ def test_write_2d_with_mask(device, data_type): values = torch.arange(1, 7, dtype=torch.float32, device=torch_device).reshape(3, 2) data = {"scalar": 7.0, "torch": values[:2], "warp": wp.from_torch(values)[:2]}[data_type] out = wp.zeros((3, 2), dtype=wp.float32, device=device) - write_2d_with_mask(data, env_mask, item_mask, out, device) + write_2d_with_mask(data, env_mask, item_mask, out) expected = torch.zeros_like(values) expected[0, 1] = 7.0 if data_type == "scalar" else values[0, 1] assert torch.equal(wp.to_torch(out), expected) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index e7417a7c1413..142fe68ac5e4 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -2679,7 +2679,7 @@ def set_fixed_tendon_stiffness_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) self.assert_shape_and_dtype_mask(stiffness, (env_mask, fixed_tendon_mask), wp.float32, "stiffness") - write_2d_with_mask(stiffness, env_mask, fixed_tendon_mask, self.data._fixed_tendon_stiffness, self.device) + write_2d_with_mask(stiffness, env_mask, fixed_tendon_mask, self.data._fixed_tendon_stiffness) def set_fixed_tendon_damping_index( self, @@ -2750,7 +2750,7 @@ def set_fixed_tendon_damping_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) self.assert_shape_and_dtype_mask(damping, (env_mask, fixed_tendon_mask), wp.float32, "damping") - write_2d_with_mask(damping, env_mask, fixed_tendon_mask, self.data._fixed_tendon_damping, self.device) + write_2d_with_mask(damping, env_mask, fixed_tendon_mask, self.data._fixed_tendon_damping) def set_fixed_tendon_limit_stiffness_index( self, @@ -2866,7 +2866,7 @@ def set_fixed_tendon_position_limit_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.") - write_2d_with_mask(limit, env_mask, fixed_tendon_mask, self.data._fixed_tendon_pos_limits, self.device) + write_2d_with_mask(limit, env_mask, fixed_tendon_mask, self.data._fixed_tendon_pos_limits) def set_fixed_tendon_rest_length_index( self, @@ -3674,7 +3674,7 @@ def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | No Returns: Environment indices. """ - return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device) + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES) 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. @@ -3685,7 +3685,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. """ - return resolve_ids(ids=joint_ids, all_ids=self._ALL_JOINT_INDICES, device=self.device) + return resolve_ids(ids=joint_ids, all_ids=self._ALL_JOINT_INDICES) 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. @@ -3696,7 +3696,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. """ - return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self.device) + return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES) def _resolve_fixed_tendon_ids( self, tendon_ids: Sequence[int] | torch.Tensor | wp.array | None @@ -3709,7 +3709,7 @@ def _resolve_fixed_tendon_ids( Returns: A warp array of tendon indices or a tensor of tendon indices. """ - return resolve_ids(ids=tendon_ids, all_ids=self._ALL_FIXED_TENDON_INDICES, device=self.device) + return resolve_ids(ids=tendon_ids, all_ids=self._ALL_FIXED_TENDON_INDICES) def _resolve_spatial_tendon_ids( self, spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None @@ -3722,7 +3722,7 @@ def _resolve_spatial_tendon_ids( Returns: A warp array of spatial tendon indices or a tensor of spatial tendon indices. """ - return resolve_ids(ids=spatial_tendon_ids, all_ids=self._ALL_SPATIAL_TENDON_INDICES, device=self.device) + return resolve_ids(ids=spatial_tendon_ids, all_ids=self._ALL_SPATIAL_TENDON_INDICES) @staticmethod def _to_torch_ids(ids: wp.array | torch.Tensor) -> torch.Tensor: 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 651f39febab7..f8ae4d3f3a72 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,7 +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.memory import resolve_ids +from isaaclab.utils.warp.memory import resolve_1d_indices, resolve_ids from .cable_object_data import CableObjectData from .kernels import ( @@ -208,7 +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.""" - return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device, dtype=wp.int32) + return resolve_1d_indices(resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES)) 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 7b8059e8662f..565d19b8d530 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,7 +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.memory import resolve_1d_mask, resolve_ids +from isaaclab.utils.warp.memory import resolve_1d_indices, resolve_1d_mask, resolve_ids from .deformable_object_data import DeformableObjectData from .kernels import ( @@ -403,7 +403,7 @@ def write_nodal_kinematic_target_to_sim_mask( def _resolve_env_ids(self, env_ids): """Resolve environment indices to a warp int32 array.""" - return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device, dtype=wp.int32) + return resolve_1d_indices(resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES)) 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.""" 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 a010df1ca9f8..4375abd9d994 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,7 +15,7 @@ from isaaclab.assets.deformable_object.base_deformable_object import BaseDeformableObject from isaaclab.utils.warp import ProxyArray -from isaaclab.utils.warp.memory import resolve_1d_mask, resolve_ids +from isaaclab.utils.warp.memory import resolve_1d_indices, resolve_1d_mask, resolve_ids from isaaclab_newton.sim.spawners.mpm.mpm import _SIMULATION_POINTS_SUFFIX @@ -284,7 +284,7 @@ def _create_buffers(self): self._data.default_particle_state_w = self._data.default_nodal_state_w def _resolve_env_ids(self, env_ids): - return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device, dtype=wp.int32) + return resolve_1d_indices(resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES)) def _resolve_mask(self, mask: wp.array | torch.Tensor | None) -> wp.array: return resolve_1d_mask(mask=mask, all_mask=self._ALL_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 be77ed145934..f1e12d9bcc9e 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 @@ -1100,7 +1100,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. """ - return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device) + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES) 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. @@ -1111,7 +1111,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. """ - return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self.device) + return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES) """ 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 f73fd483de87..d97037f49966 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 @@ -482,7 +482,7 @@ def write_body_link_pose_to_sim_mask( scratch_mask=self._scratch_body_mask, ) self.assert_shape_and_dtype_mask(body_poses, (env_mask, body_mask), wp.transformf, "body_poses") - write_2d_with_mask(body_poses, env_mask, body_mask, self.data.body_link_pose_w, self.device) + write_2d_with_mask(body_poses, env_mask, body_mask, self.data.body_link_pose_w) if not skip_forward: self.data._reset_pose(env_mask=env_mask) @@ -1266,7 +1266,7 @@ def _resolve_env_ids(self, env_ids) -> wp.array | torch.Tensor: Returns: Environment indices. """ - return resolve_ids(ids=env_ids, all_ids=self._ALL_ENV_INDICES, device=self.device) + return resolve_ids(ids=env_ids, all_ids=self._ALL_ENV_INDICES) def _resolve_body_ids(self, body_ids) -> wp.array | torch.Tensor: """Resolve body indices. @@ -1277,7 +1277,7 @@ def _resolve_body_ids(self, body_ids) -> wp.array | torch.Tensor: Returns: Body indices. """ - return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self.device) + return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES) """ Internal simulation callbacks. diff --git a/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py b/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py index 53b81f633fbd..0ad05a179a43 100644 --- a/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py +++ b/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py @@ -4366,27 +4366,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``.""" - return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self._device) + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES) 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``.""" - return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self._device) + return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES) 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``.""" - return resolve_ids(ids=joint_ids, all_ids=self._ALL_JOINT_INDICES, device=self._device) + return resolve_ids(ids=joint_ids, all_ids=self._ALL_JOINT_INDICES) 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``.""" - return resolve_ids(ids=tendon_ids, all_ids=self._ALL_FIXED_TENDON_INDICES, device=self._device) + return resolve_ids(ids=tendon_ids, all_ids=self._ALL_FIXED_TENDON_INDICES) 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``.""" - return resolve_ids(ids=tendon_ids, all_ids=self._ALL_SPATIAL_TENDON_INDICES, device=self._device) + return resolve_ids(ids=tendon_ids, all_ids=self._ALL_SPATIAL_TENDON_INDICES) def _broadcast_scalar_to_2d( self, value: float | torch.Tensor | wp.array, shape: tuple[int, int] 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 ee4ca2dd0d24..bdfc5b34e062 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,7 +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.memory import resolve_ids +from isaaclab.utils.warp.memory import resolve_1d_indices, resolve_ids from isaaclab_ov import tensor_types as TT from isaaclab_ov.sim.views.ovphysx_view import OvPhysxView @@ -529,7 +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.""" - return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device, dtype=wp.int32) + return resolve_1d_indices(resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES)) @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 f022446b3b7d..8bbca14a2757 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 @@ -1062,7 +1062,7 @@ def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | No Returns: Environment indices on ``self._device``. """ - return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self._device) + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES) def _sim_env_ids_view(self, count: int) -> wp.array: """Return a cached prefix of the simulator-index scratch buffer.""" @@ -1085,7 +1085,7 @@ def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | Returns: Body indices. """ - return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self._device) + return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES) def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array: """Resolve an environment mask to a ``wp.bool`` array. 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 deff20d9afbd..7d371fd8e78a 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 @@ -1327,11 +1327,11 @@ def _resolve_env_ids(self, env_ids=None, env_mask: wp.array | None = None) -> wp resolved array onto ``self._device`` so kernel launches don't fail on a device mismatch. """ - return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device) + return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_ENV_INDICES) def _resolve_body_ids(self, body_ids=None, body_mask: wp.array | None = None) -> wp.array | torch.Tensor: """Resolve body indices on ``self._device``.""" - return resolve_ids(ids=body_ids, mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self._device) + return resolve_ids(ids=body_ids, mask=body_mask, all_ids=self._ALL_BODY_INDICES) def _env_body_ids_to_view_ids( self, env_ids: torch.Tensor | wp.array, body_ids: torch.Tensor | wp.array, device: str = "cuda:0" diff --git a/source/isaaclab_ov/test/assets/test_deformable_object.py b/source/isaaclab_ov/test/assets/test_deformable_object.py index 38e32c3f9041..6863791ae774 100644 --- a/source/isaaclab_ov/test/assets/test_deformable_object.py +++ b/source/isaaclab_ov/test/assets/test_deformable_object.py @@ -401,11 +401,16 @@ def test_mask_writes_update_only_masked_rows(): expected = torch.zeros((3, 2, 6), device=asset.device) expected[1] = state[1] torch.testing.assert_close(asset.data.nodal_state_w.torch, expected) + view = asset.root_view + assert view.last_indices.numpy().tolist() == [1] + assert (view.position_write_count, view.velocity_write_count) == (1, 1) 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]) + assert view.last_indices.numpy().tolist() == [1] + assert (view.position_write_count, view.velocity_write_count) == (2, 1) @pytest.mark.usefixtures("host_shell") diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index e6eb9f1e1ab6..efa766f47b25 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -4155,7 +4155,7 @@ def _resolve_env_ids( Returns: Environment indices. """ - return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) + return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_INDICES) def _resolve_joint_ids( self, joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, joint_mask: wp.array | None = None @@ -4172,7 +4172,7 @@ def _resolve_joint_ids( Returns: A warp array of joint indices or a tensor of joint indices. """ - return resolve_ids(ids=joint_ids, mask=joint_mask, all_ids=self._ALL_JOINT_INDICES, device=self.device) + return resolve_ids(ids=joint_ids, mask=joint_mask, all_ids=self._ALL_JOINT_INDICES) def _resolve_body_ids( self, body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, body_mask: wp.array | None = None @@ -4186,7 +4186,7 @@ def _resolve_body_ids( Returns: A warp array of body indices or a tensor of body indices. """ - return resolve_ids(ids=body_ids, mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self.device) + return resolve_ids(ids=body_ids, mask=body_mask, all_ids=self._ALL_BODY_INDICES) def _resolve_fixed_tendon_ids( self, @@ -4202,9 +4202,7 @@ def _resolve_fixed_tendon_ids( Returns: A warp array of tendon indices or a tensor of tendon indices. """ - return resolve_ids( - ids=tendon_ids, mask=fixed_tendon_mask, all_ids=self._ALL_FIXED_TENDON_INDICES, device=self.device - ) + return resolve_ids(ids=tendon_ids, mask=fixed_tendon_mask, all_ids=self._ALL_FIXED_TENDON_INDICES) def _resolve_spatial_tendon_ids( self, @@ -4220,12 +4218,7 @@ def _resolve_spatial_tendon_ids( Returns: A warp array of spatial tendon indices or a tensor of spatial tendon indices. """ - return resolve_ids( - ids=spatial_tendon_ids, - mask=spatial_tendon_mask, - all_ids=self._ALL_SPATIAL_TENDON_INDICES, - device=self.device, - ) + return resolve_ids(ids=spatial_tendon_ids, mask=spatial_tendon_mask, all_ids=self._ALL_SPATIAL_TENDON_INDICES) """ Deprecated methods. 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 b832033e79ea..2d1b9fb6a741 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,7 +21,7 @@ from isaaclab.assets.asset_base import AssetBase from isaaclab.markers import VisualizationMarkers from isaaclab.utils.warp import ProxyArray -from isaaclab.utils.warp.memory import resolve_ids +from isaaclab.utils.warp.memory import resolve_1d_indices, resolve_ids from .deformable_object_data import DeformableObjectData from .kernels import ( @@ -219,7 +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. """ - env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES) self.write_nodal_state_to_sim_index(nodal_state, env_ids=env_ids, full_data=True) def write_nodal_pos_to_sim_index( @@ -286,7 +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. """ - env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES) self.write_nodal_pos_to_sim_index(nodal_pos, env_ids=env_ids, full_data=True) def write_nodal_velocity_to_sim_index( @@ -354,7 +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. """ - env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES) self.write_nodal_velocity_to_sim_index(nodal_vel, env_ids=env_ids, full_data=True) def write_nodal_kinematic_target_to_sim_index( @@ -430,7 +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. """ - env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES) self.write_nodal_kinematic_target_to_sim_index(targets, env_ids=env_ids, full_data=True) """ @@ -555,7 +555,7 @@ def _resolve_env_ids(self, env_ids): Returns: A warp array of environment indices. """ - return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device, dtype=wp.int32) + return resolve_1d_indices(resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES)) def _initialize_impl(self): # obtain global simulation view 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 c1564550794a..835c744b2eb1 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 @@ -1077,7 +1077,7 @@ def _resolve_env_ids( Returns: A warp array of environment indices or a tensor of environment indices. """ - return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) + return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_INDICES) def _resolve_body_ids( self, body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, body_mask: wp.array | None = None @@ -1091,7 +1091,7 @@ def _resolve_body_ids( Returns: A warp array of body indices or a tensor of body indices. """ - return resolve_ids(ids=body_ids, mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self.device) + return resolve_ids(ids=body_ids, mask=body_mask, all_ids=self._ALL_BODY_INDICES) """ 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 8ea4d343cba5..2b92fd10c9d2 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 @@ -1242,11 +1242,11 @@ def reshape_data_to_view_3d( def _resolve_env_ids(self, env_ids=None, env_mask: wp.array | None = None) -> wp.array | torch.Tensor: """Resolve environment indices.""" - return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self.device) + return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_ENV_INDICES) 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) + return resolve_ids(ids=body_ids, mask=body_mask, all_ids=self._ALL_BODY_INDICES) def _resolve_mask_writer_body_ids( self, body_mask: wp.array | None, body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None 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 52201c763c8a..9c5adcda5c62 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,7 +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.memory import resolve_ids +from isaaclab.utils.warp.memory import resolve_1d_indices, resolve_ids if TYPE_CHECKING: from isaacsim.robot.surface_gripper import GripperView @@ -554,7 +554,7 @@ def _resolve_env_ids(self, env_ids=None, env_mask: wp.array | torch.Tensor | Non Returns: A warp array of int32 indices. """ - return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_INDICES, device=self._device, dtype=wp.int32) + return resolve_1d_indices(resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_INDICES)) 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/assets/test_deformable_object.py b/source/isaaclab_physx/test/assets/test_deformable_object.py index 4eeb4c246f8c..4cb8b89fc562 100644 --- a/source/isaaclab_physx/test/assets/test_deformable_object.py +++ b/source/isaaclab_physx/test/assets/test_deformable_object.py @@ -258,6 +258,12 @@ def test_deformable_nodal_state_mask_writes(deformable_scene: _DeformableScene) 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) + for getter, component in ( + (deformable.root_view.get_simulation_nodal_positions, slice(0, 3)), + (deformable.root_view.get_simulation_nodal_velocities, slice(3, 6)), + ): + raw = wp.to_torch(getter()).to(device).reshape(expected[..., component].shape) + torch.testing.assert_close(raw, expected[..., component]) @pytest.mark.isaacsim_ci From 92a413efa88e1c238da2c2430b42e46a8be155f5 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Thu, 29 Oct 2026 02:20:13 -0700 Subject: [PATCH 79/80] Write strided vector components through scalar views instead of copying them --- source/isaaclab/isaaclab/utils/warp/memory.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/source/isaaclab/isaaclab/utils/warp/memory.py b/source/isaaclab/isaaclab/utils/warp/memory.py index cfebcfde0b64..fe7878eb4c32 100644 --- a/source/isaaclab/isaaclab/utils/warp/memory.py +++ b/source/isaaclab/isaaclab/utils/warp/memory.py @@ -309,20 +309,17 @@ def write_2d_with_mask( column_mask: Boolean mask over the second dimension of ``out``. out: Output array. The launch runs on its device. """ + dim = (row_mask.shape[0], column_mask.shape[0]) if isinstance(data, (int, float)): kernel = make_masked_fill(out.dtype, 2) + elif isinstance(data, torch.Tensor) and data.ndim > 2 and data.stride(-1) != 1: + # a vector-typed argument needs packed components; copy strided ones through scalar views instead + scalar = out.dtype._wp_scalar_type_ + kernel, dim = make_masked_copy(scalar, 3), (*dim, data.shape[-1]) + data, out = wp.from_torch(data, dtype=scalar), out.view(scalar) else: - if isinstance(data, torch.Tensor) and data.ndim > 2: - # a vector-typed kernel argument needs contiguous components - data = data.contiguous() kernel = make_masked_copy(out.dtype, 2) - wp.launch( - kernel, - dim=(row_mask.shape[0], column_mask.shape[0]), - inputs=[data, row_mask, column_mask], - outputs=[out], - device=out.device, - ) + wp.launch(kernel, dim=dim, inputs=[data, row_mask, column_mask], outputs=[out], device=out.device) @cache From efee1b9ffddd250313247835f7b9af4c53c46959 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Thu, 29 Oct 2026 02:44:14 -0700 Subject: [PATCH 80/80] Update camera poses on reset through the resolved mask --- source/isaaclab/changelog.d/jichuanh-id-mask.fixed.rst | 2 ++ source/isaaclab/isaaclab/sensors/camera/camera.py | 2 +- source/isaaclab/test/sensors/test_camera.py | 8 +++++++- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/source/isaaclab/changelog.d/jichuanh-id-mask.fixed.rst b/source/isaaclab/changelog.d/jichuanh-id-mask.fixed.rst index 9ed262f73ec2..1675a56aa24e 100644 --- a/source/isaaclab/changelog.d/jichuanh-id-mask.fixed.rst +++ b/source/isaaclab/changelog.d/jichuanh-id-mask.fixed.rst @@ -1,2 +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. diff --git a/source/isaaclab/isaaclab/sensors/camera/camera.py b/source/isaaclab/isaaclab/sensors/camera/camera.py index 03d6f12eea4b..d7acf0c1f94a 100644 --- a/source/isaaclab/isaaclab/sensors/camera/camera.py +++ b/source/isaaclab/isaaclab/sensors/camera/camera.py @@ -665,7 +665,7 @@ def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: wp.array # reset the data # note: this recomputation is useful if one performs events such as randomizations on the camera poses. if env_mask is not None: - self._update_poses(env_mask=env_mask, frame_op=2) + self._update_poses(env_mask=reset_mask, frame_op=2) elif env_ids is None: self._update_poses(frame_op=2) else: diff --git a/source/isaaclab/test/sensors/test_camera.py b/source/isaaclab/test/sensors/test_camera.py index e0036daa966f..a8353a31577d 100644 --- a/source/isaaclab/test/sensors/test_camera.py +++ b/source/isaaclab/test/sensors/test_camera.py @@ -521,7 +521,7 @@ def setup_camera_device(device): @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) -def test_camera_multi_regex_init(setup_camera_device, device): +def test_camera_multi_regex_init(setup_camera_device, device, monkeypatch): """Test multi-camera initialization with regex prim paths and content validation.""" sim, camera_cfg, dt = setup_camera_device @@ -540,6 +540,12 @@ def test_camera_multi_regex_init(setup_camera_device, device): camera.reset(slice(1, None, 2)) torch.testing.assert_close(camera.frame.torch[::2], torch.full((5,), 7, device=device, dtype=torch.int64)) torch.testing.assert_close(camera.frame.torch[1::2], torch.zeros(4, device=device, dtype=torch.int64)) + if device.startswith("cuda"): + # a mask on another device resets the same rows; strict launches reject a CPU array on CUDA + monkeypatch.setattr(wp.config, "launch_array_access_mode", wp.config.LaunchArrayAccessMode.STRICT) + camera.frame.torch.fill_(7) + camera.reset(env_mask=wp.array([i == 0 for i in range(num_cameras)], dtype=wp.bool, device="cpu")) + assert camera.frame.torch.tolist() == [0] + [7] * (num_cameras - 1) assert camera._sensor_prims[1].GetPath().pathString == "/World/Origin_1/CameraSensor" assert isinstance(camera._sensor_prims[0], UsdGeom.Camera)