From ce4db9c4c08e711c184065d5324b784b30cdde12 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 30 Sep 2026 02:03:58 -0700 Subject: [PATCH 01/50] 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/50] 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/50] 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/50] 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/50] 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/50] 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/50] 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 2bb8e9e3068049308e52c71ec6943bd7b9896f69 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 1 Oct 2026 00:37:32 -0700 Subject: [PATCH 08/50] 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 09/50] 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 10/50] 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 11/50] 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 12/50] 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 13/50] 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 14/50] 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 15/50] 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 16/50] 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 17/50] 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 18/50] 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 19/50] 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 20/50] 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 21/50] 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 f6d31caca0a498d58dc9028b95faa6e4efdb9c6d Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 20 Oct 2026 09:59:05 -0700 Subject: [PATCH 22/50] 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 23/50] 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 5676dd1c78534c484a467076fe5c3fb923596bc5 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Mon, 26 Oct 2026 02:48:16 -0700 Subject: [PATCH 24/50] 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 25/50] 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 26/50] 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 27/50] 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 28/50] 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 29/50] 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 30/50] 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 31/50] 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 e390289b4e203470ac3c41f628e85b20d3c2fce1 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 12:08:26 -0700 Subject: [PATCH 32/50] 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 2048c626fc09043944ac1a17e44949ff871c5be1 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 12:16:17 -0700 Subject: [PATCH 33/50] 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 38721eac3a64f993b54bedbd51b6662dbd5f650d Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 27 Oct 2026 15:15:36 -0700 Subject: [PATCH 34/50] 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 35/50] 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 36/50] 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 37/50] 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 38/50] 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 39/50] 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 40/50] 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 41/50] 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 42/50] 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 43/50] 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 44/50] 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 45/50] 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 46/50] 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 47/50] 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 dc2aba6ae591033ea7eeee4e195012c71aa101c5 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Thu, 29 Oct 2026 02:03:59 -0700 Subject: [PATCH 48/50] 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 49/50] 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 50/50] 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)