Resolve id and mask selections through shared helpers

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

@@ -0,0 +1,2 @@
+* Added :func:`~isaaclab.utils.warp.utils.resolve_ids`, which resolves an index selection given as ids or as a
+  boolean mask to ids on a device, as :func:`~isaaclab.utils.warp.utils.resolve_1d_mask` does to a mask.

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

@@ -0,0 +1,11 @@
+* Changed the env selection helpers of :class:`~isaaclab.utils.wrench_composer.WrenchComposer`,
+  :class:`~isaaclab.sensors.SensorBase` and :class:`~isaaclab.sensors.Camera` to resolve ``env_ids`` and ``env_mask``
+  through :func:`~isaaclab.utils.warp.utils.resolve_ids` and :func:`~isaaclab.utils.warp.utils.resolve_1d_mask`.
+  Any integer sequence and slices with a negative step are accepted, ids on another device are moved to the owner's
+  device, and a :class:`~isaaclab.utils.warp.ProxyArray` is rejected with a :class:`TypeError`.
+* Changed :class:`~isaaclab.utils.wrench_composer.WrenchComposer` to resolve ``body_ids`` through
+  :func:`~isaaclab.utils.warp.utils.resolve_ids` as well, so slices other than ``slice(None)`` are accepted.
+* Changed :func:`~isaaclab.utils.warp.utils.resolve_1d_mask` to make ``scratch_mask`` optional for callers that never
+  pass ids. Ids without a ``scratch_mask`` raise a :class:`ValueError`.
+* Changed :func:`~isaaclab.utils.warp.utils.resolve_1d_mask` to copy a Warp mask on another device to ``device``, and
+  to accept a boolean torch mask on ``device`` inside a CUDA graph capture.

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

@@ -0,0 +1,2 @@
+* Fixed the ``write_nodal_*_to_sim_mask`` methods of :class:`~isaaclab.assets.BaseDeformableObject`, which the
+  OVPhysX deformable object inherits, raising a :class:`RuntimeError` whenever ``env_mask`` was given.

source/isaaclab/isaaclab/assets/deformable_object/base_deformable_object.py

@@ -15,6 +15,7 @@
 
 from ...utils import math as math_utils
 from ...utils.warp import ProxyArray
+from ...utils.warp.utils import resolve_ids
 from ..asset_base import AssetBase
 
 if TYPE_CHECKING:
@@ -200,6 +201,9 @@ def write_nodal_kinematic_target_to_sim_index(
 
     """
     Operations - Write to simulation (mask variants).
+
+    These convert the mask to ids, which waits for the device, and call the index variants.
+    Backends with a mask write path override them.
     """
 
     def write_nodal_state_to_sim_mask(
@@ -214,10 +218,7 @@ def write_nodal_state_to_sim_mask(
                 Shape is (num_instances, max_sim_vertices_per_body, 6).
             env_mask: Environment mask. If None, then all indices are used.
         """
-        if env_mask is not None:
-            env_ids = wp.nonzero(env_mask)
-        else:
-            env_ids = self._ALL_INDICES
+        env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device)
         self.write_nodal_state_to_sim_index(nodal_state, env_ids=env_ids, full_data=True)
 
     def write_nodal_pos_to_sim_mask(
@@ -232,10 +233,7 @@ def write_nodal_pos_to_sim_mask(
                 Shape is (num_instances, max_sim_vertices_per_body, 3).
             env_mask: Environment mask. If None, then all indices are used.
         """
-        if env_mask is not None:
-            env_ids = wp.nonzero(env_mask)
-        else:
-            env_ids = self._ALL_INDICES
+        env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device)
         self.write_nodal_pos_to_sim_index(nodal_pos, env_ids=env_ids, full_data=True)
 
     def write_nodal_velocity_to_sim_mask(
@@ -250,10 +248,7 @@ def write_nodal_velocity_to_sim_mask(
                 Shape is (num_instances, max_sim_vertices_per_body, 3).
             env_mask: Environment mask. If None, then all indices are used.
         """
-        if env_mask is not None:
-            env_ids = wp.nonzero(env_mask)
-        else:
-            env_ids = self._ALL_INDICES
+        env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device)
         self.write_nodal_velocity_to_sim_index(nodal_vel, env_ids=env_ids, full_data=True)
 
     def write_nodal_kinematic_target_to_sim_mask(
@@ -268,10 +263,7 @@ def write_nodal_kinematic_target_to_sim_mask(
                 Shape is (num_instances, max_sim_vertices_per_body, 4).
             env_mask: Environment mask. If None, then all indices are used.
         """
-        if env_mask is not None:
-            env_ids = wp.nonzero(env_mask)
-        else:
-            env_ids = self._ALL_INDICES
+        env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device)
         self.write_nodal_kinematic_target_to_sim_index(targets, env_ids=env_ids, full_data=True)
 
     """

source/isaaclab/isaaclab/sensors/camera/camera.py

@@ -26,6 +26,7 @@
     quat_from_matrix,
 )
 from ...utils.warp import ProxyArray
+from ...utils.warp.utils import resolve_ids
 from ..sensor_base import SensorBase
 from .camera_data import CameraData, RenderBufferKind
 
@@ -1028,27 +1029,7 @@ def _resolve_env_ids_wp(self, env_ids: Sequence[int] | torch.Tensor | wp.array |
         """Resolve camera indices to a Warp ``int32`` array."""
         if env_ids is None:
             return None
-        if isinstance(env_ids, wp.array):
-            if env_ids.dtype != wp.int32:
-                raise TypeError(f"Unsupported wp.array dtype for env_ids: {env_ids.dtype}. Expected wp.int32.")
-            if env_ids.ndim != 1:
-                raise ValueError("Warp camera indices must be a one-dimensional array.")
-            if str(env_ids.device) == str(self._device):
-                return env_ids
-            return env_ids.to(self._device)
-        elif isinstance(env_ids, torch.Tensor):
-            env_ids = env_ids.to(device=self._device, dtype=torch.int32).reshape(-1)
-            if not env_ids.is_contiguous():
-                env_ids = env_ids.contiguous()
-            return wp.from_torch(env_ids, dtype=wp.int32)
-        elif isinstance(env_ids, slice) and (env_ids.step or 1) > 0:
-            return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids])
-        elif isinstance(env_ids, slice):
-            # Torch cannot slice with a negative step.
-            env_ids = np.arange(self._view.count, dtype=np.int32)[env_ids]
-        else:
-            env_ids = np.asarray(env_ids, dtype=np.int32).reshape(-1)
-        return wp.array(env_ids, dtype=wp.int32, device=self._device)
+        return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self._device, dtype=wp.int32)
 
     def _env_mask_has_any(self, env_mask: wp.array) -> bool:
         """Return whether the mask selects any camera."""

source/isaaclab/isaaclab/sensors/sensor_base.py

@@ -29,7 +29,7 @@
 from ..physics import PhysicsEvent, PhysicsManager
 from ..sim.utils.queries import get_first_matching_ancestor_prim
 from ..sim.utils.transforms import resolve_prim_pose
-from ..utils import index_fill_
+from ..utils.warp.utils import resolve_1d_mask
 from .kernels import invalidate_buffers_kernel, reset_envs_kernel, update_outdated_envs_kernel, update_timestamp_kernel
 
 if TYPE_CHECKING:
@@ -506,14 +506,9 @@ def _resolve_indices_and_mask(
         self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None
     ) -> wp.array:
         """Resolve environment indices to a warp array and mask."""
-        if (env_ids is None or env_ids == slice(None)) and env_mask is None:
-            return self._ALL_ENV_MASK
-        elif env_mask is not None:
-            return env_mask
-        else:
-            self._reset_mask.zero_()
-            index_fill_(self._reset_mask_torch, env_ids, True)
-            return self._reset_mask
+        return resolve_1d_mask(
+            ids=env_ids, mask=env_mask, all_mask=self._ALL_ENV_MASK, device=self._device, scratch_mask=self._reset_mask
+        )
 
     def _resolve_rigid_body_ancestor_expr(
         self,

source/isaaclab/isaaclab/utils/warp/utils.py

@@ -11,6 +11,8 @@
 import torch
 import warp as wp
 
+from .proxy_array import ProxyArray
+
 ##
 # Mask resolution - ids/mask to warp boolean mask.
 ##
@@ -30,7 +32,7 @@ def resolve_1d_mask(
     ids: Sequence[int] | slice | torch.Tensor | wp.array | None = None,
     mask: wp.array | torch.Tensor | None = None,
     all_mask: wp.array,
-    scratch_mask: wp.array,
+    scratch_mask: wp.array | None = None,
     device: str,
 ) -> wp.array:
     """Resolve ids/mask into a warp boolean mask.
@@ -41,30 +43,36 @@ def resolve_1d_mask(
     Args:
         ids: Index ids. Accepts ``Sequence[int]``, ``slice``, ``torch.Tensor``,
             ``wp.array(dtype=wp.int32)``, or ``None`` (all elements).
-        mask: Direct boolean mask. ``wp.array`` is returned as-is;
+        mask: Direct boolean mask. ``wp.array`` is returned as-is, or copied when on another device;
             ``torch.Tensor`` is converted.
         all_mask: Pre-allocated all-True mask returned when both *ids* and *mask*
             are ``None``.
         scratch_mask: Pre-allocated scratch buffer populated in-place when *ids*
-            are provided. Not re-entrant (shared buffer).
+            are provided. Not re-entrant (shared buffer). Defaults to None, for callers
+            that never pass *ids*.
         device: Warp device string (e.g. ``"cuda:0"``).
 
     Returns:
         A ``wp.array(dtype=wp.bool)`` mask.
+
+    Raises:
+        ValueError: If *ids* select a subset of elements and no ``scratch_mask`` is given.
     """
     # Normalize slice(None) to None so the capture guard treats it identically to ids=None.
     if isinstance(ids, slice) and ids == slice(None):
         ids = None
 
     if wp.get_device().is_capturing:
-        if ids is not None or (mask is not None and not isinstance(mask, wp.array)):
+        if ids is not None or (mask is not None and _mask_needs_copy(mask, device)):
             raise RuntimeError(
-                "resolve_1d_mask is only capturable when mask is a wp.array or both ids and mask are None."
+                "resolve_1d_mask is only capturable when ids are None and mask is None or a boolean mask on device."
             )
 
     # --- Direct mask input ---
     if mask is not None:
         if isinstance(mask, wp.array):
+            if mask.device != wp.get_device(device):
+                mask = wp.clone(mask, device=device)
             return mask
         if isinstance(mask, torch.Tensor):
             if mask.dtype != torch.bool:
@@ -78,6 +86,9 @@ def resolve_1d_mask(
     if ids is None:
         return all_mask
 
+    if scratch_mask is None:
+        raise ValueError("resolve_1d_mask needs a scratch_mask to write ids into.")
+
     if isinstance(ids, slice):
         scratch_mask.zero_()
         wp.to_torch(scratch_mask)[ids] = True
@@ -115,6 +126,104 @@ def resolve_1d_mask(
     return scratch_mask
 
 
+def _mask_needs_copy(mask: wp.array | torch.Tensor, device: str) -> bool:
+    """Whether :func:`resolve_1d_mask` casts or copies ``mask``, which a graph capture cannot record."""
+    if isinstance(mask, wp.array):
+        return mask.device != wp.get_device(device)
+    if isinstance(mask, torch.Tensor):
+        return mask.dtype != torch.bool or str(mask.device) != device
+    return True
+
+
+##
+# Id resolution - ids/mask to ids.
+##
+
+
+def resolve_ids(
+    *,
+    ids: Sequence[int] | slice | torch.Tensor | wp.array | None = None,
+    mask: wp.array | torch.Tensor | None = None,
+    all_ids: wp.array,
+    device: str,
+    dtype: type | None = None,
+) -> wp.array | torch.Tensor:
+    """Resolve ids/mask into ids on ``device``. The mask takes precedence over the ids.
+
+    Args:
+        ids: Index ids. Accepts ``Sequence[int]``, ``slice``, ``torch.Tensor``, ``wp.array``,
+            or ``None`` (all elements). Tensors and warp arrays keep their integer dtype
+            unless ``dtype`` is given. Pass ids from ``find_*(as_proxy=True)`` as ``.warp`` or ``.torch``.
+        mask: Boolean mask. Its ids are found with :func:`torch.nonzero`, which waits for
+            the device.
+        all_ids: Pre-allocated ids of every element, returned as-is when *ids* and *mask*
+            select every element.
+        device: Warp device string (e.g. ``"cuda:0"``).
+        dtype: Warp integer dtype the ids must have, for callers whose kernels take one
+            index dtype. The ids are then a contiguous 1-D ``wp.array`` of this dtype.
+            Defaults to None: a mask resolves to an int64 ``torch.Tensor``, and ids keep
+            their type and dtype.
+
+    Returns:
+        ``all_ids``, a ``torch.Tensor`` or a ``wp.array`` of ids on ``device``.
+
+    Raises:
+        TypeError: If *ids* is a :class:`ProxyArray`, which is output-only, *mask* is
+            neither a ``wp.array`` nor a ``torch.Tensor``, or ``dtype`` is given and the
+            ids are not 32-bit or 64-bit integers.
+        ValueError: If ``dtype`` is given and the ids are not one-dimensional.
+    """
+    if mask is not None:
+        if isinstance(mask, wp.array):
+            mask = wp.to_torch(mask)
+        elif not isinstance(mask, torch.Tensor):
+            raise TypeError(f"Unsupported mask type: {type(mask)}")
+        ids = torch.nonzero(mask.to(wp.device_to_torch(device))).squeeze(-1)
+    else:
+        ids = _ids_on_device(ids, all_ids, device)
+    return ids if dtype is None else _ids_as_dtype(ids, dtype)
+
+
+def _ids_on_device(
+    ids: Sequence[int] | slice | torch.Tensor | wp.array | None, all_ids: wp.array, device: str
+) -> wp.array | torch.Tensor:
+    """Ids on ``device``, keeping the dtype of tensors and warp arrays; lists and slices become ``wp.int32``."""
+    if isinstance(ids, torch.Tensor):
+        return ids if str(ids.device) == str(wp.get_device(device)) else ids.to(wp.device_to_torch(device))
+    if ids is None:
+        return all_ids
+    if isinstance(ids, wp.array):
+        return ids if ids.device == wp.get_device(device) else ids.to(device)
+    if isinstance(ids, ProxyArray):
+        raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
+    if isinstance(ids, slice):
+        if ids == slice(None):
+            return all_ids
+        start, stop, step = ids.indices(all_ids.shape[0])
+        if step > 0:
+            # a strided view of the caller's all-ids, no copy
+            return wp.from_torch(wp.to_torch(all_ids)[start:stop:step])
+        return wp.array(list(range(start, stop, step)), dtype=wp.int32, device=device)
+    return wp.array(list(ids), dtype=wp.int32, device=device)
+
+
+def _ids_as_dtype(ids: wp.array | torch.Tensor, dtype: type) -> wp.array:
+    """``ids`` as a contiguous 1-D ``wp.array`` of ``dtype`` on their device."""
+    if dtype not in (wp.int32, wp.int64):
+        raise TypeError(f"Unsupported id dtype: {dtype}. Expected wp.int32 or wp.int64.")
+    if ids.ndim != 1:
+        raise ValueError(f"Ids must be a one-dimensional array, got {ids.ndim} dimensions.")
+    if isinstance(ids, wp.array):
+        if ids.dtype == dtype and ids.is_contiguous:
+            return ids
+        if ids.dtype not in (wp.int32, wp.int64):
+            raise TypeError(f"Warp ids must be 32-bit or 64-bit signed integers, got {ids.dtype}.")
+        ids = wp.to_torch(ids)
+    elif ids.dtype.is_floating_point or ids.dtype.is_complex or ids.dtype == torch.bool:
+        raise TypeError(f"Ids must be integers, got {ids.dtype}.")
+    return wp.from_torch(ids.to(wp.dtype_to_torch(dtype)).contiguous(), dtype=dtype)
+
+
 ##
 # Capture safety — property guard.
 ##

source/isaaclab/isaaclab/utils/wrench_composer.py

@@ -24,6 +24,7 @@
     set_forces_to_dual_buffers_index_kernel,
     set_forces_to_dual_buffers_mask,
 )
+from .warp.utils import resolve_ids
 
 if TYPE_CHECKING:
     from ..assets import BaseArticulation, BaseRigidObject, BaseRigidObjectCollection
@@ -721,24 +722,9 @@ def _resolve_env_ids(self, env_ids: wp.array | torch.Tensor | list | slice | Non
             Environment indices.
 
         Raises:
-            TypeError: If ``env_ids`` is an unsupported type.
+            TypeError: If ``env_ids`` is a :class:`~isaaclab.utils.warp.ProxyArray`.
         """
-        if env_ids is None:
-            return self._ALL_ENV_INDICES
-        # Check tensor types before slice comparison (tensor == slice crashes)
-        if isinstance(env_ids, torch.Tensor):
-            return env_ids
-        if isinstance(env_ids, wp.array):
-            return env_ids
-        if env_ids == slice(None):
-            return self._ALL_ENV_INDICES
-        if isinstance(env_ids, slice):
-            return wp.from_torch(wp.to_torch(self._ALL_ENV_INDICES)[env_ids])
-        if isinstance(env_ids, list):
-            return wp.array(env_ids, dtype=wp.int32, device=self.device)
-        raise TypeError(
-            f"env_ids must be None, slice(None), list, torch.Tensor, or wp.array, got {type(env_ids).__name__}"
-        )
+        return resolve_ids(ids=env_ids, all_ids=self._ALL_ENV_INDICES, device=self.device)
 
     def _resolve_body_ids(
         self, body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None
@@ -752,23 +738,9 @@ def _resolve_body_ids(
             Body indices.
 
         Raises:
-            TypeError: If ``body_ids`` is an unsupported type.
+            TypeError: If ``body_ids`` is a :class:`~isaaclab.utils.warp.ProxyArray`.
         """
-        if body_ids is None:
-            return self._ALL_BODY_INDICES
-        if isinstance(body_ids, ProxyArray):
-            raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.")
-        if isinstance(body_ids, torch.Tensor):
-            return body_ids
-        if isinstance(body_ids, wp.array):
-            return body_ids
-        if body_ids == slice(None):
-            return self._ALL_BODY_INDICES
-        if isinstance(body_ids, Sequence):
-            return wp.array(body_ids, dtype=wp.int32, device=self.device)
-        raise TypeError(
-            f"body_ids must be None, slice(None), a sequence, torch.Tensor, or wp.array, got {type(body_ids).__name__}"
-        )
+        return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self.device)
 
     def _ensure_composed(self):
         """Compose input buffers into output buffers if dirty."""

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

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