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-id-mask.added.rst b/source/isaaclab/changelog.d/jichuanh-id-mask.added.rst new file mode 100644 index 000000000000..03136d5b7fdd --- /dev/null +++ b/source/isaaclab/changelog.d/jichuanh-id-mask.added.rst @@ -0,0 +1,4 @@ +* Added :func:`~isaaclab.utils.warp.memory.resolve_ids`, which resolves an index selection given as ids or as a + boolean mask to ids on a device, as :func:`~isaaclab.utils.warp.memory.resolve_1d_mask` does to a mask. +* Added :func:`~isaaclab.utils.warp.memory.write_2d_with_mask`, which writes 2-D data or one scalar into a Warp array + where a row mask and a column mask are both true, launching over the mask lengths. 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..2d36907fce7c --- /dev/null +++ b/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst @@ -0,0 +1,11 @@ +* Changed the env selection helpers of :class:`~isaaclab.utils.wrench_composer.WrenchComposer` and + :class:`~isaaclab.sensors.Camera` to resolve ``env_ids`` and ``env_mask`` through + :func:`~isaaclab.utils.warp.memory.resolve_ids` and :func:`~isaaclab.utils.warp.memory.resolve_1d_mask`. + Any integer sequence and slices with a negative step are accepted, ids on another device are moved to the owner's + device, and a :class:`~isaaclab.utils.warp.ProxyArray` is rejected with a :class:`TypeError`. The + :class:`~isaaclab.sensors.Camera` now rejects floating-point ``env_ids`` with a :class:`TypeError` and + multi-dimensional ``env_ids`` with a :class:`ValueError` instead of casting and flattening them. +* Changed :class:`~isaaclab.utils.wrench_composer.WrenchComposer` to resolve ``body_ids`` through + :func:`~isaaclab.utils.warp.memory.resolve_ids` as well, so slices other than ``slice(None)`` are accepted. +* Changed :func:`~isaaclab.utils.warp.memory.resolve_1d_mask` to make ``scratch_mask`` optional for callers that never + pass ids. Their masks are returned as given. Ids without a ``scratch_mask`` raise a :class:`ValueError`. 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/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/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..c52b4d0c252e --- /dev/null +++ b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst @@ -0,0 +1,10 @@ +* Added ``env_mask`` to :meth:`~isaaclab.actuators.ActuatorCollection.reset` and + :meth:`~isaaclab.actuators.ActuatorControl.reset_native_actuators`. Isaac Lab actuator models and backend-native + actuator state reset only the masked environments from the mask itself, without converting it to indices. +* Added ``batch_mask`` to :meth:`~isaaclab.utils.buffers.DelayBuffer.set_time_lag` and + :meth:`~isaaclab.utils.buffers.DelayBuffer.reset`, which select batches on the device without synchronizing the + host. Tensor lags set through a mask are full-sized and are not range-checked. +* Added ``env_mask`` to :meth:`~isaaclab.assets.BaseCableObject.reset`, matching the other scene assets. +* Added ``env_mask`` to :meth:`~isaaclab.scene.InteractiveScene.reset`, which forwards the mask to every asset and + sensor so only the masked environments are reset, without synchronizing the host. Surface grippers still reset by + index. 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..cf37839edc0b --- /dev/null +++ b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst @@ -0,0 +1,18 @@ +* Changed masked :meth:`~isaaclab.sensors.ray_caster.BaseRayCaster.reset` to resample drift in a Warp kernel from + per-environment random states seeded from torch's global generator, so it neither synchronizes the host nor + allocates and can be captured in a CUDA graph. When both ``env_ids`` and ``env_mask`` are given, the mask selects + the resampled environments, matching :meth:`~isaaclab.sensors.SensorBase.reset`. +* Changed :meth:`~isaaclab.actuators.newton.NewtonActuatorAdapter.reset` to reuse preallocated per-DOF masks, so + masked resets of Newton-native actuator state do not allocate and can be captured in a CUDA graph. +* **Breaking:** Changed :meth:`~isaaclab.actuators.ActuatorBase.reset` to take an ``env_mask`` argument, the + boolean mask of the environments to reset that masked resets pass instead of indices. Custom actuator models must + accept ``env_mask`` 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. +* 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/changelog.d/jichuanh-warp-scene-mask-native.major b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.major new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab/isaaclab/actuators/actuator_base.py b/source/isaaclab/isaaclab/actuators/actuator_base.py index 8b031f6fb534..a75221b05574 100644 --- a/source/isaaclab/isaaclab/actuators/actuator_base.py +++ b/source/isaaclab/isaaclab/actuators/actuator_base.py @@ -279,11 +279,17 @@ def joint_indices(self) -> slice | torch.Tensor: """ @abstractmethod - def reset(self, env_ids: Sequence[int]): + def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: torch.Tensor | None = None): """Reset the internals within the group. + .. caution:: + If both ``env_ids`` and ``env_mask`` are provided, ``env_mask`` takes precedence. + Args: - env_ids: List of environment IDs to reset. + env_ids: Environment indices or a slice selecting the environments to reset. None resets all + environments. + env_mask: Boolean mask of the environments to reset. Shape is (num_envs,). Implementations reset from it + without synchronizing the host, for example through :func:`~isaaclab.utils.array.index_fill_`. """ raise NotImplementedError diff --git a/source/isaaclab/isaaclab/actuators/actuator_collection.py b/source/isaaclab/isaaclab/actuators/actuator_collection.py index 1204dc3a94ab..c0eebf6caedb 100644 --- a/source/isaaclab/isaaclab/actuators/actuator_collection.py +++ b/source/isaaclab/isaaclab/actuators/actuator_collection.py @@ -180,18 +180,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 receive the mask as a zero-copy boolean tensor through the ``env_mask`` + argument of :meth:`~isaaclab.actuators.ActuatorBase.reset`. Backend-native actuator state resets from + the mask directly. + + .. caution:: + If both ``env_ids`` and ``env_mask`` are provided, ``env_mask`` takes precedence. + Args: env_ids: Environment indices to reset. Defaults to all environments. + env_mask: Environment mask. Shape is (num_instances,). Defaults to None. """ group_env_ids = self._control._normalize_index_sequence(env_ids) + # 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) - self._control.reset_native_actuators(slice(None) if group_env_ids is None else group_env_ids) + 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: """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 35e84e956302..64f7d4b37261 100644 --- a/source/isaaclab/isaaclab/actuators/actuator_control.py +++ b/source/isaaclab/isaaclab/actuators/actuator_control.py @@ -307,11 +307,15 @@ def submit_commands(self, collection: ActuatorCollection) -> None: """ raise NotImplementedError - def reset_native_actuators(self, env_ids: Sequence[int] | slice) -> None: + def reset_native_actuators(self, env_ids: Sequence[int] | slice, env_mask: wp.array | None = None) -> None: """Reset backend-native actuator state. + .. caution:: + If both ``env_ids`` and ``env_mask`` are provided, ``env_mask`` takes precedence. + Args: env_ids: Environments to reset. + env_mask: Environment mask. Shape is (num_instances,). Defaults to None. """ diff --git a/source/isaaclab/isaaclab/actuators/actuator_net.py b/source/isaaclab/isaaclab/actuators/actuator_net.py index c5dca36a609d..5d06e4bfb500 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] | 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(): - 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] | slice | 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 bc20ea721fb7..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): + 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]): + def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: torch.Tensor | None = None): pass def compute( @@ -451,18 +451,25 @@ class DelayedPDActuator(IdealPDActuator): def __init__(self, cfg: DelayedPDActuatorCfg, *args, **kwargs): super().__init__(cfg, *args, **kwargs) + # masked resets apply the sampled delays without range-checking them on the host + if not 0 <= cfg.min_delay <= cfg.max_delay: + raise ValueError(f"Expected 0 <= min_delay <= max_delay, received {cfg.min_delay} and {cfg.max_delay}.") # instantiate the delay buffers self.positions_delay_buffer = DelayBuffer(cfg.max_delay, self._num_envs, device=self._device) self.velocities_delay_buffer = DelayBuffer(cfg.max_delay, self._num_envs, device=self._device) self.efforts_delay_buffer = DelayBuffer(cfg.max_delay, self._num_envs, device=self._device) self._delay_buffers = (self.positions_delay_buffer, self.velocities_delay_buffer, self.efforts_delay_buffer) - def reset(self, env_ids: Sequence[int]): - super().reset(env_ids) + def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: torch.Tensor | None = None): + super().reset(env_ids, env_mask) # number of environments (since env_ids can be a slice) if env_ids is None: env_ids = slice(None) - num_envs = len(range(self._num_envs)[env_ids]) if isinstance(env_ids, slice) else len(env_ids) + if env_mask is not None: + # a mask samples every environment and the delay buffers keep the masked ones on the device + num_envs = self._num_envs + else: + num_envs = len(range(self._num_envs)[env_ids]) if isinstance(env_ids, slice) else len(env_ids) # set a new random delay for environments in env_ids time_lags = torch.randint( low=self.cfg.min_delay, @@ -473,8 +480,8 @@ def reset(self, env_ids: Sequence[int]): ) # set delays and reset buffers for delay_buffer in self._delay_buffers: - delay_buffer.set_time_lag(time_lags, env_ids) - delay_buffer.reset(env_ids) + delay_buffer.set_time_lag(time_lags, env_ids, batch_mask=env_mask) + delay_buffer.reset(env_ids, batch_mask=env_mask) def compute( self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor diff --git a/source/isaaclab/isaaclab/actuators/newton/adapter.py b/source/isaaclab/isaaclab/actuators/newton/adapter.py index 07600f67fdfa..c08525370d5e 100644 --- a/source/isaaclab/isaaclab/actuators/newton/adapter.py +++ b/source/isaaclab/isaaclab/actuators/newton/adapter.py @@ -32,12 +32,8 @@ from newton.selection import ArticulationView from ...utils import index_fill_ -from .kernels import ( - build_implicit_dof_mask, - build_per_dof_env_mask_kernel, - set_mask_kernel, - zero_at_indices_kernel, -) +from ...utils.warp.memory import resolve_1d_mask +from .kernels import build_implicit_dof_mask, build_per_dof_env_mask_kernel, zero_at_indices_kernel if TYPE_CHECKING: from .. import ActuatorCollection @@ -106,6 +102,12 @@ def __init__( self._states_a = [act.state() for act in actuators] self._states_b = [act.state() for act in actuators] + # Newton actuator states reset by a per-DOF mask, which every partial reset expands from the env mask. + # Allocated once and fully rewritten on each reset, so a reset does not allocate. + self._reset_dof_masks = [wp.zeros(act.indices.shape[0], dtype=wp.bool, device=device) for act in actuators] + # Env masks for resolving env ids in :meth:`reset`. + self._all_env_mask = wp.ones(num_envs, dtype=wp.bool, device=device) + self._scratch_env_mask = wp.zeros(num_envs, dtype=wp.bool, device=device) # Pre-clamp computed effort buffer. Each Newton actuator scatter-adds # its raw drive output to ``sim_control.joint_computed_f`` when @@ -172,7 +174,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: @@ -180,6 +186,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``), @@ -188,7 +196,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) @@ -196,23 +204,17 @@ 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) + if env_mask is None: + num_envs = len(range(self._num_envs)[env_ids]) if isinstance(env_ids, slice) else len(env_ids) + if num_envs == 0: + return + env_mask = resolve_1d_mask( + ids=env_ids, + all_mask=self._all_env_mask, + scratch_mask=self._scratch_env_mask, ) - 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/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), 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/assets/articulation/base_articulation.py b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py index 215f216005e0..dddb3c4c2b60 100644 --- a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py +++ b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py @@ -2091,7 +2091,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: @@ -2109,7 +2109,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). """ @@ -2119,7 +2120,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: @@ -2137,7 +2138,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/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/isaaclab/assets/deformable_object/base_deformable_object.py b/source/isaaclab/isaaclab/assets/deformable_object/base_deformable_object.py index 0542a460af48..2144ccbad474 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.memory 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/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/isaaclab/sensors/camera/camera.py b/source/isaaclab/isaaclab/sensors/camera/camera.py index d052dd59aa73..56c0e2f6dd2c 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.memory 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/ray_caster/base_ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py index f8c5aba6af6c..99c70dbb29a9 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py @@ -119,12 +119,27 @@ 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 env_mask is not None: + # 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=[ + 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)), + 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: 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 @@ -251,6 +266,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 866b9eabf0ef..d4188f3edd8e 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py @@ -151,6 +151,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/utils/buffers/delay_buffer.py b/source/isaaclab/isaaclab/utils/buffers/delay_buffer.py index 17398b58869e..9ab8c652994f 100644 --- a/source/isaaclab/isaaclab/utils/buffers/delay_buffer.py +++ b/source/isaaclab/isaaclab/utils/buffers/delay_buffer.py @@ -122,18 +122,29 @@ def time_lags(self) -> torch.Tensor: Operations. """ - def set_time_lag(self, time_lag: int | torch.Tensor, batch_ids: Sequence[int] | 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),). + * 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 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. @@ -150,6 +161,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 batch_mask is not None: + time_lag = time_lag.to(device=self.device, dtype=self._time_lags.dtype) + 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 else: @@ -163,20 +178,37 @@ 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] | 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. 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) - if self._hold_prob is not None: + index_fill_(self._num_pushes, indices if batch_mask is None else batch_mask, 0) + if self._hold_prob is None: + return + 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, + self.history_length + 1, + (self.batch_size,), + dtype=self._time_lags.dtype, + device=self.device, + ) + torch.where(batch_mask, 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/isaaclab/utils/warp/memory.py b/source/isaaclab/isaaclab/utils/warp/memory.py index 79bd61115a17..a55fa0618fb1 100644 --- a/source/isaaclab/isaaclab/utils/warp/memory.py +++ b/source/isaaclab/isaaclab/utils/warp/memory.py @@ -20,6 +20,7 @@ import warp as wp from .index_kernel import selector_dtype +from .proxy_array import ProxyArray _CONCAT_TILE_SIZE = 32 @@ -290,6 +291,42 @@ def fill( return fill +def write_2d_with_mask( + data: "float | torch.Tensor | wp.array", + 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. + + 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 ``(row_mask.shape[0], column_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. + 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"``). + """ + if isinstance(data, (int, float)): + kernel = make_masked_fill(out.dtype, 2) + 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=device, + ) + + @cache def _make_populate_mask(dtype) -> wp.Kernel: @wp.kernel(module="unique") @@ -316,7 +353,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, ) -> wp.array: """Resolve a selection using caller-owned boolean storage on its device. @@ -331,19 +368,26 @@ def resolve_1d_mask( preallocated device selectors avoid this synchronization. Integer widths and strides are preserved. Masks must have the destination shape and boolean dtype. The returned scratch storage is reused, so this is not reentrant. + Callers that never select by ids may omit ``scratch_mask``; their masks are then + returned as given, and the callers launch over the mask lengths. + + Raises: + ValueError: If ids or a slice select a subset of elements and no ``scratch_mask`` is given. """ - device = scratch_mask.device - if scratch_mask.ndim != 1 or scratch_mask.dtype != wp.bool: - raise ValueError("scratch_mask must be a rank-1 boolean Warp array") - if all_mask.shape != scratch_mask.shape or all_mask.dtype != wp.bool or all_mask.device != device: - raise ValueError("all_mask must match scratch_mask shape, dtype and device") + device = all_mask.device + if scratch_mask is not None: + if scratch_mask.ndim != 1 or scratch_mask.dtype != wp.bool: + raise ValueError("scratch_mask must be a rank-1 boolean Warp array") + if all_mask.shape != scratch_mask.shape or all_mask.dtype != wp.bool or all_mask.device != scratch_mask.device: + raise ValueError("all_mask must match scratch_mask shape, dtype and device") temporary_input = False selection = mask if mask is not None else ids if selection is None or (isinstance(selection, slice) and selection == slice(None)): return all_mask if isinstance(selection, slice): - start, stop, step = selection.indices(scratch_mask.shape[0]) + start, stop, step = selection.indices(all_mask.shape[0]) + scratch_mask = _require_scratch_mask(scratch_mask) scratch_mask.zero_() count = len(range(start, stop, step)) if count: @@ -381,11 +425,12 @@ def resolve_1d_mask( if selection.ndim != 1: raise ValueError("Selectors must be rank-1 arrays") if mask is not None or selection.dtype == wp.bool: - if selection.dtype != wp.bool or selection.shape != scratch_mask.shape: + if selection.dtype != wp.bool or (scratch_mask is not None and selection.shape != scratch_mask.shape): raise ValueError("Boolean masks must match the destination shape and dtype") return selection if selection.dtype not in (wp.int32, wp.int64): raise TypeError("Index arrays must use signed int32 or int64 values") + scratch_mask = _require_scratch_mask(scratch_mask) scratch_mask.zero_() if selection.size: wp.launch( @@ -395,3 +440,94 @@ def resolve_1d_mask( # Eager uploads must outlive scatter even on imported nonblocking CUDA streams. wp.synchronize_stream(wp.get_stream(device)) return scratch_mask + + +def _require_scratch_mask(scratch_mask: wp.array | None) -> wp.array: + """The scratch mask :func:`resolve_1d_mask` writes ids into.""" + if scratch_mask is None: + raise ValueError("resolve_1d_mask needs a scratch_mask to write ids into.") + return scratch_mask + + +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) diff --git a/source/isaaclab/isaaclab/utils/wrench_composer.py b/source/isaaclab/isaaclab/utils/wrench_composer.py index 6331716cad5d..f9ac3263b3be 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.memory 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/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) 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..5c2d5a6cd38c --- /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=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/assets/contract/articulation_factory.py b/source/isaaclab/test/assets/contract/articulation_factory.py index e47337eccf1c..24fd46c4e8e8 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 77a9b0c74817..5a6182f57015 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,20 @@ 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 + 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") + + 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}) @@ -535,18 +539,6 @@ def read() -> ProxyArray: method(**{kwarg: wp.from_torch(extra_env.contiguous(), dtype=dtype)}) -@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)): - 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 1d294d9a8efc..cd800bd91696 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,30 @@ 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. + 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) + 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) for value in (bad, wp.from_torch(bad.contiguous(), dtype=wp_dtype)): 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/test/utils/test_delay_buffer.py b/source/isaaclab/test/utils/test_delay_buffer.py index d2e46d19d13d..6055fc0ffde0 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, 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(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)) + + @pytest.mark.parametrize( ("time_lag", "batch_ids"), [ diff --git a/source/isaaclab/test/utils/test_wrench_composer.py b/source/isaaclab/test/utils/test_wrench_composer.py index 10daf0d06612..fccb5f3c8773 100644 --- a/source/isaaclab/test/utils/test_wrench_composer.py +++ b/source/isaaclab/test/utils/test_wrench_composer.py @@ -474,11 +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() - 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,)) 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 new file mode 100644 index 000000000000..6bc5c013e0cc --- /dev/null +++ b/source/isaaclab/test/utils/warp/test_resolve_selection.py @@ -0,0 +1,111 @@ +# 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 DeviceScope, test_devices +from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.memory 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()) +@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) + 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_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"]) +def test_resolve_ids_from_a_mask(device, mask_type): + """A mask resolves to its ascending ids, takes precedence over ids, and honours ``dtype``.""" + 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 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 ids_int32.numpy().tolist() == [1, 2, 4, 7] + + +@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) is mask + assert resolve_1d_mask(all_mask=all_mask) is all_mask + assert resolve_1d_mask(ids=slice(None), all_mask=all_mask) is all_mask + 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) + 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) + + +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +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) + 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) is scratch + assert scratch.numpy().tolist() == [True, False, True] 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..4fb308d76a86 --- /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.memory 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_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..e04a0303897f --- /dev/null +++ b/source/isaaclab_contrib/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst @@ -0,0 +1,2 @@ +* 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 e2a914a17efa..894c87fdd9b4 100644 --- a/source/isaaclab_contrib/isaaclab_contrib/assets/multirotor/multirotor.py +++ b/source/isaaclab_contrib/isaaclab_contrib/assets/multirotor/multirotor.py @@ -41,12 +41,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) @@ -244,26 +251,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 5d195eeffe7a..6faf717f5d07 100644 --- a/source/isaaclab_contrib/test/assets/test_multirotor.py +++ b/source/isaaclab_contrib/test/assets/test_multirotor.py @@ -18,12 +18,15 @@ 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 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 +159,30 @@ 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 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]: 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 f792b994f67b..000000000000 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py +++ /dev/null @@ -1,44 +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. - - 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) - for articulation in self._articulations.values(): - articulation.reset(env_ids, env_mask=env_mask) - for deformable_object in self._deformable_objects.values(): - deformable_object.reset(env_ids) - 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 (no env_mask support) - for sensor in self._sensors.values(): - sensor.reset(env_ids) 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 f326d39d48c8..fda31c129f73 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,6 +33,7 @@ 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 @@ -40,7 +41,6 @@ from isaaclab.utils.timer import Timer from isaaclab.utils.warp.memory import resolve_1d_mask -from isaaclab_experimental.envs.interactive_scene_warp import InteractiveSceneWarp as InteractiveScene from isaaclab_experimental.utils.manager_call_switch import ManagerCallMode, ManagerCallSwitch # import logger 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_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..9614b51dd445 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst @@ -0,0 +1,14 @@ +* Changed the env selection helpers of the Newton assets to resolve ``env_ids`` through + :func:`~isaaclab.utils.warp.memory.resolve_ids` and ``env_mask`` through + :func:`~isaaclab.utils.warp.memory.resolve_1d_mask`. Any integer sequence and slices with a negative step are accepted, + ids on another device are moved to the 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.memory.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. 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/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/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.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/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/articulation/actuator_control.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/actuator_control.py index 76ddb3a83c0a..a0ad89c7d11b 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/actuator_control.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/actuator_control.py @@ -232,9 +232,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 self._articulation._backend_manager.backend.actuators is not None: - self._articulation._backend_manager.backend.actuators.reset(env_ids) + self._articulation._backend_manager.backend.actuators.reset(env_ids, env_mask) def _bam_actuators(self) -> list[tuple[Actuator, list[str]]]: """Return the Newton BAM actuators that drive this articulation, each with the BAM groups it drives. diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index 4596d05b610e..ce7646a3d3be 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 @@ -34,6 +34,7 @@ from isaaclab.utils.version import get_isaac_sim_version, has_kit from isaaclab.utils.warp import ProxyArray, memory from isaaclab.utils.warp import math as warp_math +from isaaclab.utils.warp.memory 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 @@ -381,7 +382,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) @@ -2291,15 +2292,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. @@ -2309,11 +2306,7 @@ 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 + return resolve_1d_mask(ids=env_ids, all_mask=self._ALL_ENV_MASK, scratch_mask=self._scratch_env_mask) """ Operations - Setters. @@ -2674,8 +2667,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). @@ -2683,13 +2676,10 @@ 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") + write_2d_with_mask(stiffness, env_mask, fixed_tendon_mask, self.data._fixed_tendon_stiffness, self.device) def set_fixed_tendon_damping_index( self, @@ -2748,8 +2738,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). @@ -2757,13 +2747,10 @@ 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") + 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, @@ -2865,8 +2852,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) @@ -2874,13 +2861,12 @@ 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.") + 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, @@ -3080,10 +3066,31 @@ 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=(env_mask.shape[0], fixed_tendon_mask.shape[0]), + 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; + # 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, @@ -3389,6 +3396,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) @@ -3666,15 +3674,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. @@ -3685,13 +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. """ - 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. @@ -3702,13 +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. """ - 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 @@ -3721,13 +3709,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 @@ -3740,27 +3722,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: @@ -3785,12 +3747,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) """ 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 600caa545870..d1b7e749132c 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py @@ -498,6 +498,36 @@ def write_joint_friction_param_to_buffer_kernel( return _WRITE_JOINT_FRICTION_PARAM_TO_BUFFER_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 update_default_joint_values( source: wp.array(dtype=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 2bc5ab64aaa5..f494d1996973 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.memory import resolve_ids from .cable_object_data import CableObjectData from .kernels import ( @@ -58,13 +59,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.""" @@ -207,15 +209,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 40d15f7461aa..7b8059e8662f 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.memory 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) 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 f2d98ce53633..86cb7645f167 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/kernels.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/kernels.py @@ -706,6 +706,93 @@ 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 (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 (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]: + 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 [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 [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]: + 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 [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 [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 [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]: + 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")) 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 bc340f5b7131..a010df1ca9f8 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.memory 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) 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 ddffc620823e..6e7e8e65560b 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, memory +from isaaclab.utils.warp.memory import resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_newton.assets import kernels as shared_kernels @@ -129,8 +130,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. @@ -1099,13 +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. """ - 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. @@ -1116,13 +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. """ - 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 8cf53a9c4b35..57528d8b3a58 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, memory +from isaaclab.utils.warp.memory 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 @@ -175,8 +176,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. @@ -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,18 @@ 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) + # 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, ) + 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) + if not skip_forward: + self.data._reset_pose(env_mask=env_mask) def write_body_com_pose_to_sim_index( self, @@ -597,17 +580,24 @@ 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) + # 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, ) + 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=(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, + ) + 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 +701,24 @@ 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) + # 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, ) + 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=(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, + ) + if not skip_forward: + self.data._reset_velocity(env_mask=env_mask) def write_body_link_velocity_to_sim_index( self, @@ -831,21 +824,24 @@ 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) + # 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, + ) + 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=(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, ) + if not skip_forward: + self.data._reset_velocity(env_mask=env_mask, from_com=False) """ Operations - Setters. @@ -1228,6 +1224,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) @@ -1269,13 +1266,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. @@ -1286,37 +1277,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. diff --git a/source/isaaclab_newton/test/assets/test_articulation.py b/source/isaaclab_newton/test/assets/test_articulation.py index a3baf7022d26..1f41f2c04927 100644 --- a/source/isaaclab_newton/test/assets/test_articulation.py +++ b/source/isaaclab_newton/test/assets/test_articulation.py @@ -16,11 +16,11 @@ world for its own duration. Configurations that fail initialization, nested-root IK, the rebind scenario, which swaps the live Newton -state, and the CUDA-graph check build their own small scenes. Only one simulation context can be alive, so +state, and the CUDA-graph checks build their own small scenes. Only one simulation context can be alive, so 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, one own-scene test covers the -ordered-state republish recorded into a captured graph. +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 recorded into a captured graph. """ from isaaclab_newton.physics import NewtonCfg, NewtonManager @@ -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. @@ -707,6 +729,56 @@ 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_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 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 + 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) + selected[1, 1] = True + stiffness = torch.zeros(num_envs, num_tendons, device=device) + damping = torch.zeros_like(stiffness) + limits = torch.zeros(num_envs, 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) + + ## # Composite scene. ## @@ -1254,9 +1326,9 @@ def one_step_response() -> tuple[torch.Tensor, torch.Tensor]: actuator_reset = actuator.reset reset_env_ids = [] - def record_actuator_reset(env_ids: torch.Tensor | None = None) -> None: + def record_actuator_reset(env_ids: torch.Tensor | None = None, env_mask: torch.Tensor | None = None) -> 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_newton/test/assets/test_newton_actuators_newton.py b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py index b106a373b13f..6596c01e8ee8 100644 --- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py +++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py @@ -383,35 +383,96 @@ def _run( return _Run(sim=sim, articulations=articulations, results=results) -def _record_lab_state_reset(run: _Run) -> dict[str, tuple[torch.Tensor, torch.Tensor]]: - """Reset one Lab-path environment of the delayed island and record each group's effort and PD demand. +_RESET_SELECTORS = ("env_ids", "env_mask") +"""Partial-reset selectors: an index tensor and a Warp mask.""" + + +def _reset_env_0(articulation: Articulation, selector: str) -> None: + """Reset environment 0 alone through the given selector.""" + if selector == "env_ids": + articulation.reset(env_ids=torch.tensor([0], device=articulation.device, dtype=torch.long)) + else: + env_mask = torch.arange(articulation.num_instances, device=articulation.device) == 0 + articulation.reset(env_mask=wp.from_torch(env_mask)) + + +def _record_lab_state_reset(run: _Run) -> dict[str, dict[str, tuple[torch.Tensor, torch.Tensor]]]: + """Reset one Lab-path environment of the delayed island per selector and record each group's effort and PD demand. Reset environments must accept fresh commands while the remaining environments retain their delay history. """ articulation = run.articulations["delayed"] 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([0], device=articulation.device)) - commands.set_position_index(value=new_target) - articulation.write_data_to_sim() expected_target = old_target.clone() expected_target[0] = new_target[0] recorded = {} - for name, actuator in articulation.actuators.items(): - joints = actuator.joint_indices - demand = ( - actuator.stiffness * (expected_target[:, joints] - articulation.data.joint_pos.torch[:, joints]) - - actuator.damping * articulation.data.joint_vel.torch[:, joints] - ) - recorded[name] = (actuator.computed_effort.clone(), demand) + for selector in _RESET_SELECTORS: + # 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() + _reset_env_0(articulation, selector) + commands.set_position_index(value=new_target) + articulation.write_data_to_sim() + recorded[selector] = {} + for name, actuator in articulation.actuators.items(): + joints = actuator.joint_indices + demand = ( + actuator.stiffness * (expected_target[:, joints] - articulation.data.joint_pos.torch[:, joints]) + - actuator.damping * articulation.data.joint_vel.torch[:, joints] + ) + recorded[selector][name] = (actuator.computed_effort.clone(), demand) return recorded +def _warm_delayed_island(run: _Run) -> list: + """Step the delayed island until its delay buffers hold history and return its delayed actuators' state pairs. + + Only the reset articulation's own delayed actuators are returned. The adapter is model-wide, so the other + islands' actuators share its state buffers; their state is not part of this articulation's contract. + """ + articulation = run.articulations["delayed"] + for _ in range(3): + articulation.write_data_to_sim() + run.sim.step() + articulation.update(DT) + adapter = run.sim.backend_manager.backend.actuators + assert adapter is not None + own_actuators = [] + for group_name in articulation.actuators._native_group_names: + group_actuators = articulation.actuators[group_name] + own_actuators.extend(group_actuators if isinstance(group_actuators, tuple) else (group_actuators,)) + stateful_pairs = [ + (act, st) + for act, st in zip(adapter.actuators, adapter._states_a) + if any(act is own for own in own_actuators) and st is not None and st.delay_state is not None + ] + assert 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() + assert (pushes_before > 0).all(), "expected non-zero num_pushes for all DOFs after warmup" + return stateful_pairs + + +def _assert_delay_state_cleared_for_env_0_only(run: _Run, stateful_pairs: list) -> None: + """Assert ``num_pushes`` is zero for environment 0's DOFs and positive for every other environment's DOFs.""" + # Map each entry of ``act.indices`` to its env via the adapter's per-env DOF count. The adapter is + # model-wide (includes free-joint DOFs on floating-base articulations), so ``adapter.num_joints`` is the + # stride. + adapter = run.sim.backend_manager.backend.actuators + 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 == 0: + assert int(pushes_after[i]) == 0, f"DOF {i} (env {env}) should be reset to 0, got {pushes_after[i]}" + else: + assert int(pushes_after[i]) > 0, f"DOF {i} (env {env}) was NOT in the reset selection but is 0" + + @pytest.fixture(scope="module", params=test_devices(DeviceScope.CUDA)) def device(request: pytest.FixtureRequest) -> str: """Simulation device of the actuator scenes.""" @@ -757,6 +818,31 @@ def test_sync_torque_telemetry_keeps_user_order_effort_buffers_unmapped() -> Non np.testing.assert_allclose(applied.numpy(), np.asarray([[100.0, 200.0, 300.0]], dtype=np.float32)) +# --------------------------------------------------------------------------- +# Own scene; runs before the module-scoped scenes exist, since only one simulation context can be alive +# --------------------------------------------------------------------------- + + +def test_newton_masked_reset_replays_from_cuda_graph_without_allocation(device: str) -> None: + """A masked reset records into a CUDA graph without allocating and, on replay, clears only the masked env. + + The scene holds only the delayed island: the adapter resets every actuator of the model, and Newton's neural + actuator state reset cannot be recorded into a graph. + """ + sim_cfg = newton_sim_cfg(device, dt=DT, use_newton_actuators=True, solver_cfg=_solver_cfg()) + with build_simulation_context(sim_cfg=sim_cfg) as sim: + run = _run(sim, {"delayed": _islands()["delayed"]}, use_newton_actuators=True) + articulation = run.articulations["delayed"] + env_mask = wp.from_torch(torch.arange(articulation.num_instances, device=device) == 0) + # Allocating inside the capture raises while the memory pool is disabled. + with wp.ScopedMempool(device, False), wp.ScopedCapture(device) as capture: + articulation.reset(env_mask=env_mask) + # Repopulate the delay buffers, so the replay rather than the capture must clear environment 0. + stateful_pairs = _warm_delayed_island(run) + wp.capture_launch(capture.graph) + _assert_delay_state_cleared_for_env_0_only(run, stateful_pairs) + + # --------------------------------------------------------------------------- # Equivalence tests with different actuator types # --------------------------------------------------------------------------- @@ -945,46 +1031,14 @@ def gains(name: str) -> torch.Tensor: def test_newton_state_reset_isolated_to_reset_env(newton_run: _Run) -> None: - """Newton: ``num_pushes`` zeroes for env 0's DOFs only after reset of [0]. - - Only the reset articulation's own delayed actuators are checked. The adapter is model-wide, so the other - islands' actuators share its state buffers; their state is not part of this articulation's contract. - """ - articulation = newton_run.articulations["delayed"] - adapter = newton_run.sim.backend_manager.backend.actuators - assert adapter is not None - own_actuators = [] - for group_name in articulation.actuators._native_group_names: - group_actuators = articulation.actuators[group_name] - own_actuators.extend(group_actuators if isinstance(group_actuators, tuple) else (group_actuators,)) - stateful_pairs = [ - (act, st) - for act, st in zip(adapter.actuators, adapter._states_a) - if any(act is own for own in own_actuators) and st is not None and st.delay_state is not None - ] - assert 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() - assert (pushes_before > 0).all(), "expected non-zero num_pushes for all DOFs after warmup" - - articulation.reset(env_ids=torch.tensor([0], device=articulation.device, dtype=torch.long)) - - # Map each entry of ``act.indices`` to its env via the adapter's per-env DOF count. The adapter is - # model-wide (includes free-joint DOFs on floating-base articulations), so ``adapter.num_joints`` is the - # stride. - 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 == 0: - assert int(pushes_after[i]) == 0, f"DOF {i} (env {env}) should be reset to 0, got {pushes_after[i]}" - else: - assert int(pushes_after[i]) > 0, f"DOF {i} (env {env}) was NOT in reset env_ids but num_pushes is 0" + """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"], "env_ids") + _assert_delay_state_cleared_for_env_0_only(newton_run, stateful_pairs) -def test_lab_state_reset_isolated_to_reset_env(lab_run: dict) -> None: +@pytest.mark.parametrize("selector", _RESET_SELECTORS) +def test_lab_state_reset_isolated_to_reset_env(lab_run: dict, selector: str) -> None: """Reset environments accept fresh commands while the remaining environments retain their history.""" - for computed_effort, demand in lab_run["state_reset"].values(): + for computed_effort, demand in lab_run["state_reset"][selector].values(): torch.testing.assert_close(computed_effort, demand) diff --git a/source/isaaclab_newton/test/assets/test_rigid_object.py b/source/isaaclab_newton/test/assets/test_rigid_object.py index 4b80f67a9cc4..f78d09582f9f 100644 --- a/source/isaaclab_newton/test/assets/test_rigid_object.py +++ b/source/isaaclab_newton/test/assets/test_rigid_object.py @@ -11,9 +11,9 @@ Every environment holds a rigid-object cube and a three-cube collection, all locally spawned shapes. The cube also stands beside the collection as an unrelated rigid body, so the collection view must select exactly its configured bodies. Scene gravity is off; a test that needs gravity applies it to every world for its own duration. -Configurations that fail initialization, VBD state publication, and the single-instance cases build their own -scenes. Only one simulation context can be alive, so those tests come first and fail if selected after a -shared-scene test. +Configurations that fail initialization, VBD state publication, the single-instance cases, and the CUDA-graph +check of the collection mask writers build their own scenes. Only one simulation context can be alive, so those +tests come first and fail if selected after a shared-scene test. """ from isaaclab_newton.physics import NewtonCfg @@ -247,6 +247,46 @@ def test_collection_single_instance_initialization(device: str) -> None: torch.testing.assert_close(object_collection.data.body_link_pose_w.torch, pose) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +def test_collection_mask_writes_replay_from_cuda_graph_on_selected_cells(device: str) -> None: + """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. + """ + with build_simulation_context(sim_cfg=newton_sim_cfg(device)) as sim: + collection = spawn_assets({"collection": _collection_cfg(2)})["collection"] + sim.reset() + sim.step() + collection.update(sim.cfg.dt) + + env_mask = wp.array([False, True], dtype=wp.bool, device=device) + body_mask = wp.array([True, False], dtype=wp.bool, device=device) + selected = torch.zeros(collection.num_instances, collection.num_bodies, 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(collection.data, buffer_name).torch.clone() + target = torch.zeros_like(before) + with wp.ScopedCapture(device) as capture: + getattr(collection, 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_like(target[..., :3]) + 2.0 + else: + target.copy_(torch.rand_like(target) + 1.0) + wp.capture_launch(capture.graph) + + written = getattr(collection.data, buffer_name).torch + torch.testing.assert_close(written, torch.where(selected, target, before), msg=writer) + + ## # Shared scene. ## @@ -342,7 +382,8 @@ def test_external_force_on_single_body(scene: _Scene) -> None: """A permanent wrench reaches the solver in the local and the global frame, and a reset clears it. Cube 0 receives a force equal to its weight and hovers while cube 1 falls. Cube 0 then receives a force set - and added 1 m off its center along y and must turn about x while cube 1 keeps falling. + and added 1 m off its center along y and must turn about x while cube 1 keeps falling. A masked reset clears + the wrench of the selected environment only. """ cube = scene.cube composer = cube.permanent_wrench_composer @@ -371,6 +412,12 @@ def test_external_force_on_single_body(scene: _Scene) -> None: assert cube.data.root_ang_vel_b.torch[0, 0].abs() > 0.1 assert cube.data.root_pos_w.torch[1, 2] < 1.0 + ones = torch.ones_like(zeros) + composer.set_forces_and_torques_index(forces=ones, torques=ones) + cube.reset(env_mask=wp.from_torch(torch.arange(cube.num_instances, device=scene.device) == 0)) + assert torch.count_nonzero(composer.out_force_b.torch[0]) == 0 + assert torch.count_nonzero(composer.out_force_b.torch[1:]) == composer.out_force_b.torch[1:].numel() + @pytest.mark.isaacsim_ci def test_rigid_body_set_material_properties(scene: _Scene) -> None: @@ -861,7 +908,8 @@ def test_collection_external_force_on_single_body(scene: _Scene) -> None: """A permanent wrench reaches the solver in the local and the global frame, and a reset clears it. Every other cube receives a force equal to its weight and hovers while the rest fall. Those cubes then - receive a force set and added 1 m off their centers along y and must turn about x. + receive a force set and added 1 m off their centers along y and must turn about x. A masked reset clears the + wrenches of the selected environment only. """ collection = scene.collection composer = collection.permanent_wrench_composer @@ -890,6 +938,12 @@ def test_collection_external_force_on_single_body(scene: _Scene) -> None: assert torch.all(collection.data.body_com_ang_vel_b.torch[:, 0::2, 0] > 0.1) assert torch.all(height[:, 1::2] < 1.0) + ones = torch.ones_like(zeros) + composer.set_forces_and_torques_index(forces=ones, torques=ones) + collection.reset(env_mask=wp.from_torch(torch.arange(collection.num_instances, device=scene.device) == 0)) + assert torch.count_nonzero(composer.out_force_b.torch[0]) == 0 + assert torch.count_nonzero(composer.out_force_b.torch[1:]) == composer.out_force_b.torch[1:].numel() + @pytest.mark.isaacsim_ci def test_collection_gravity_vec_w_tracks_model_gravity(scene: _Scene) -> None: 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 4216a21eb42d..1075d4332d54 100644 --- a/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py +++ b/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py @@ -211,6 +211,47 @@ def test_legacy_multi_mesh_tracks_ad_hoc_regex_target(sim): torch.testing.assert_close(sensor.data.pos_w.torch, pose[:, :3]) +@pytest.mark.parametrize("sim", [pytest.param(False, id="eager")], indirect=True) +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)} + scene = InteractiveScene(scene_cfg) + sim.reset() + sensor: NewtonRaycastSensor = scene["raycast"] + sensor.reset() + + 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]) + + # 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() + + # 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): """Read the same box through a static target and a tracked target.""" 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..aba8c1cc69a4 --- /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.memory.resolve_ids` and ``env_mask`` through + :func:`~isaaclab.utils.warp.memory.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.memory.resolve_ids` and :func:`~isaaclab.utils.warp.memory.resolve_1d_mask`, as the env + selection does, so slices other than ``slice(None)`` are accepted. 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..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,9 +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) + self._actuator_runtime.reset(env_ids, env_mask) def stage_user_command( self, diff --git a/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py b/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py index 005d86a6624a..a8dd4c860afb 100644 --- a/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py +++ b/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py @@ -34,6 +34,7 @@ from isaaclab.utils.warp import ProxyArray, memory from isaaclab.utils.warp import kernels as warp_kernels from isaaclab.utils.warp import math as warp_math +from isaaclab.utils.warp.memory import resolve_1d_mask, resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_ov import tensor_types as TT @@ -220,7 +221,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) @@ -4365,79 +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``.""" - 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] @@ -4459,13 +4408,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) def _sim_env_ids_view(self, count: int) -> wp.array: """Return a cached prefix view of the simulator-index scratch buffer.""" @@ -4481,43 +4424,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) 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) 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) 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) 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..ee4ca2dd0d24 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.memory 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 cea8ee574f06..f022446b3b7d 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, memory +from isaaclab.utils.warp.memory import resolve_1d_mask, resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_ov import tensor_types as TT @@ -1061,17 +1062,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.""" @@ -1094,13 +1085,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. @@ -1111,13 +1096,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) def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array: """Resolve a body mask to a ``wp.bool`` array. @@ -1128,13 +1107,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) 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 6e60ee458e7c..deff20d9afbd 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, memory +from isaaclab.utils.warp.memory 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 = 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), @@ -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 = 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), @@ -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 = 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" ) @@ -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 = 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" ) @@ -863,11 +832,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 = 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( memory.make_masked_copy(wp.float32, 2), @@ -943,11 +908,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 = 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( memory.make_masked_copy(wp.transformf, 2), @@ -1033,11 +994,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 = 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( memory.make_masked_copy(wp.float32, 3), @@ -1363,34 +1320,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. """ - 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, 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``.""" - 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, 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" @@ -1400,6 +1341,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. @@ -1448,13 +1391,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) def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array: """Resolve a body mask to a ``wp.bool`` array on ``self._device``. @@ -1467,13 +1404,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) 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_articulation.py b/source/isaaclab_ov/test/assets/test_articulation.py index 66021b7b501f..607f62fffc2a 100644 --- a/source/isaaclab_ov/test/assets/test_articulation.py +++ b/source/isaaclab_ov/test/assets/test_articulation.py @@ -1636,6 +1636,13 @@ def __getitem__(self, name): # The delay state is flattened environment-major over the group's joints. assert stateful_pairs[0].delay_state.num_pushes.numpy().tolist() == [0] * num_joints + [1] * num_joints + scene.step(articulation) + 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] * num_joints + [0] * num_joints + env = Env(articulation) asset_cfg = SceneEntityCfg("robot") event_params = { diff --git a/source/isaaclab_ov/test/assets/test_deformable_object.py b/source/isaaclab_ov/test/assets/test_deformable_object.py index 7a2f90460fae..38e32c3f9041 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) 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..1483b1fd1d87 --- /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.memory.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.memory.resolve_ids`, as the env selection does, so slices other than ``slice(None)`` are + accepted. 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/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.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/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) diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 4603a29cd657..0476c34ba7dd 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -28,6 +28,7 @@ from isaaclab.utils.version import get_isaac_sim_version, has_kit from isaaclab.utils.warp import ProxyArray, memory from isaaclab.utils.warp import math as warp_math +from isaaclab.utils.warp.memory import resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_physx.assets import kernels as shared_kernels @@ -232,7 +233,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) @@ -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 @@ -691,7 +692,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 @@ -864,7 +865,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 @@ -974,7 +975,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 @@ -1088,8 +1089,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, @@ -1193,8 +1194,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 @@ -1297,8 +1298,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 @@ -1383,8 +1384,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 @@ -1506,8 +1507,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) @@ -1617,8 +1618,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, @@ -1707,8 +1708,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 @@ -1792,8 +1793,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) @@ -1875,8 +1876,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) @@ -2042,8 +2043,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, @@ -2151,8 +2152,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, @@ -2261,8 +2262,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, @@ -2363,8 +2364,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) @@ -2470,8 +2471,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) @@ -2567,8 +2568,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) @@ -2651,8 +2652,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 @@ -2729,8 +2730,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 @@ -2815,8 +2816,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 @@ -2825,7 +2826,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, @@ -2869,7 +2870,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: @@ -2893,8 +2894,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 @@ -2977,8 +2978,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 @@ -3106,8 +3107,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, ) @@ -3137,8 +3138,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 @@ -3193,7 +3194,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: @@ -3202,7 +3203,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( @@ -3280,8 +3281,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 @@ -3358,8 +3359,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 @@ -3445,8 +3446,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 @@ -3523,8 +3524,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 @@ -3577,7 +3578,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: @@ -3586,7 +3587,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) @@ -4142,62 +4143,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:: @@ -4205,123 +4167,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/deformable_object/deformable_object.py b/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py index 019d50bdb1a8..b832033e79ea 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.memory 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/isaaclab_physx/assets/rigid_object/rigid_object.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py index 29cf3fde174f..fdd0d4d8e2f8 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, memory +from isaaclab.utils.warp.memory import resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_physx.assets import kernels as shared_kernels @@ -137,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. @@ -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 ) @@ -781,14 +770,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) @@ -861,14 +844,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) @@ -944,14 +921,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) @@ -1032,40 +1003,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. @@ -1128,39 +1065,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 26950ffebb02..a6ee40228835 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, memory +from isaaclab.utils.warp.memory import resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_physx.assets import kernels as shared_kernels @@ -174,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. @@ -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, @@ -958,14 +935,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) @@ -1042,14 +1013,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) @@ -1127,14 +1092,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) @@ -1281,46 +1240,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( @@ -1331,16 +1264,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..52201c763c8a 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.memory 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/assets/test_articulation.py b/source/isaaclab_physx/test/assets/test_articulation.py index 48344bd748c8..b33ecfa87ef9 100644 --- a/source/isaaclab_physx/test/assets/test_articulation.py +++ b/source/isaaclab_physx/test/assets/test_articulation.py @@ -1086,9 +1086,9 @@ def test_floating_articulation_root_writes_and_wrenches(articulation_scene: _Art actuator_reset = actuator.reset reset_env_ids = [] - def record_actuator_reset(env_ids=None) -> None: + def record_actuator_reset(env_ids=None, env_mask=None) -> 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) composers = (floating.instantaneous_wrench_composer, floating.permanent_wrench_composer) 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.""" 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 3e84d69abcec..066e1ffb6e5a 100644 --- a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py +++ b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py @@ -296,9 +296,16 @@ def test_network_actuators_drive_haa_joints(rollouts: _Rollouts, robot: str) -> assert not torch.allclose(trace["applied_effort"][0][:, haa_ids], trace["applied_effort"][-1][:, haa_ids]) -def test_reset_clears_actuator_state_of_selected_envs(rollouts: _Rollouts) -> None: - """Resetting environment 0 clears the delay history of its DoFs only.""" +@pytest.mark.parametrize("selector", ["env_ids", "env_mask"]) +def test_reset_clears_actuator_state_of_selected_envs(rollouts: _Rollouts, selector: str) -> None: + """Resetting environment 0 by indices or by mask clears the delay history of its DoFs only.""" articulation = rollouts.robots["delayed"] + # Refill the delay history that a previous reset cleared. + sim = SimulationContext.instance() + for _ in range(3): + articulation.write_data_to_sim() + sim.step() + articulation.update(DT) adapter = articulation.newton_actuator_adapter delay_states = [ (actuator, state.delay_state) @@ -308,7 +315,10 @@ def test_reset_clears_actuator_state_of_selected_envs(rollouts: _Rollouts) -> No assert delay_states assert all((delay_state.num_pushes.numpy() > 0).all() for _, delay_state in delay_states) - articulation.reset(env_ids=torch.tensor([0], device=articulation.device)) + if selector == "env_ids": + articulation.reset(env_ids=torch.tensor([0], device=articulation.device)) + else: + articulation.reset(env_mask=wp.from_torch(torch.arange(NUM_ENVS, device=articulation.device) == 0)) for actuator, delay_state in delay_states: env_of_dof = actuator.indices.numpy() // adapter.num_joints pushes = delay_state.num_pushes.numpy() diff --git a/source/isaaclab_physx/test/assets/test_rigid_object.py b/source/isaaclab_physx/test/assets/test_rigid_object.py index 764e54ff941d..bdeb573cd776 100644 --- a/source/isaaclab_physx/test/assets/test_rigid_object.py +++ b/source/isaaclab_physx/test/assets/test_rigid_object.py @@ -644,17 +644,23 @@ def test_rigid_object_wrench_delivery_and_reset(scene: _Scene) -> None: > 0.99 ) - # A partial reset clears only the selected environment; a full reset deactivates the composers. + # A partial reset by indices or by mask clears only the selected environment; a full reset deactivates the + # composers. composers = (cube_object.instantaneous_wrench_composer, cube_object.permanent_wrench_composer) ones = torch.ones((_NUM_ENVS, 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) - 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() + 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(): + 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 cube_object.reset() for composer in composers: assert not composer.active @@ -834,17 +840,22 @@ def test_collection_wrench_delivery_and_reset(scene: _Scene) -> None: ) assert torch.all(heading[..., 1] > 0.99) - # A partial reset clears the external wrenches of the selected environments only. + # A partial reset by indices or by mask clears the external wrenches of the selected environments only. composers = (collection.instantaneous_wrench_composer, collection.permanent_wrench_composer) ones = torch.ones((_NUM_ENVS, _NUM_CUBES, 3), device=device) - collection.permanent_wrench_composer.set_forces_and_torques_index(forces=ones, torques=ones) - collection.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones) - 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() + 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(): + collection.permanent_wrench_composer.set_forces_and_torques_index(forces=ones, torques=ones) + collection.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones) + 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 collection.reset()