From 5676dd1c78534c484a467076fe5c3fb923596bc5 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Mon, 26 Oct 2026 02:48:16 -0700 Subject: [PATCH 1/6] Add resolve_ids next to resolve_1d_mask as the shared id and mask helpers resolve_ids resolves ids or a boolean mask to ids on a device, as resolve_1d_mask does to a mask. resolve_1d_mask takes an optional scratch mask, copies a Warp mask on another device, and accepts a boolean torch mask on its device inside a graph capture. The experimental package now re-exports it instead of keeping a copy. --- .../changelog.d/jichuanh-id-mask.added.rst | 2 + .../changelog.d/jichuanh-id-mask.changed.rst | 4 + .../changelog.d/jichuanh-id-mask.minor | 0 source/isaaclab/isaaclab/utils/warp/utils.py | 119 +++++++++++++- .../test/utils/warp/test_resolve_selection.py | 155 ++++++++++++++++++ .../changelog.d/jichuanh-id-mask.changed.rst | 2 + .../utils/warp/__init__.py | 3 +- .../isaaclab_experimental/utils/warp/utils.py | 88 ---------- 8 files changed, 279 insertions(+), 94 deletions(-) create mode 100644 source/isaaclab/changelog.d/jichuanh-id-mask.added.rst create mode 100644 source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst create mode 100644 source/isaaclab/changelog.d/jichuanh-id-mask.minor create mode 100644 source/isaaclab/test/utils/warp/test_resolve_selection.py create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-id-mask.changed.rst diff --git a/source/isaaclab/changelog.d/jichuanh-id-mask.added.rst b/source/isaaclab/changelog.d/jichuanh-id-mask.added.rst new file mode 100644 index 000000000000..935dc70f7fa4 --- /dev/null +++ b/source/isaaclab/changelog.d/jichuanh-id-mask.added.rst @@ -0,0 +1,2 @@ +* Added :func:`~isaaclab.utils.warp.utils.resolve_ids`, which resolves an index selection given as ids or as a + boolean mask to ids on a device, as :func:`~isaaclab.utils.warp.utils.resolve_1d_mask` does to a mask. diff --git a/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst new file mode 100644 index 000000000000..6a1b73d77c7e --- /dev/null +++ b/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst @@ -0,0 +1,4 @@ +* Changed :func:`~isaaclab.utils.warp.utils.resolve_1d_mask` to make ``scratch_mask`` optional for callers that never + pass ids. Ids without a ``scratch_mask`` raise a :class:`ValueError`. +* Changed :func:`~isaaclab.utils.warp.utils.resolve_1d_mask` to copy a Warp mask on another device to ``device``, and + to accept a boolean torch mask on ``device`` inside a CUDA graph capture. diff --git a/source/isaaclab/changelog.d/jichuanh-id-mask.minor b/source/isaaclab/changelog.d/jichuanh-id-mask.minor new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab/isaaclab/utils/warp/utils.py b/source/isaaclab/isaaclab/utils/warp/utils.py index afd31e9c1b38..24e23707615c 100644 --- a/source/isaaclab/isaaclab/utils/warp/utils.py +++ b/source/isaaclab/isaaclab/utils/warp/utils.py @@ -11,6 +11,8 @@ import torch import warp as wp +from .proxy_array import ProxyArray + ## # Mask resolution - ids/mask to warp boolean mask. ## @@ -30,7 +32,7 @@ def resolve_1d_mask( ids: Sequence[int] | slice | torch.Tensor | wp.array | None = None, mask: wp.array | torch.Tensor | None = None, all_mask: wp.array, - scratch_mask: wp.array, + scratch_mask: wp.array | None = None, device: str, ) -> wp.array: """Resolve ids/mask into a warp boolean mask. @@ -41,30 +43,36 @@ def resolve_1d_mask( Args: ids: Index ids. Accepts ``Sequence[int]``, ``slice``, ``torch.Tensor``, ``wp.array(dtype=wp.int32)``, or ``None`` (all elements). - mask: Direct boolean mask. ``wp.array`` is returned as-is; + mask: Direct boolean mask. ``wp.array`` is returned as-is, or copied when on another device; ``torch.Tensor`` is converted. all_mask: Pre-allocated all-True mask returned when both *ids* and *mask* are ``None``. scratch_mask: Pre-allocated scratch buffer populated in-place when *ids* - are provided. Not re-entrant (shared buffer). + are provided. Not re-entrant (shared buffer). Defaults to None, for callers + that never pass *ids*. device: Warp device string (e.g. ``"cuda:0"``). Returns: A ``wp.array(dtype=wp.bool)`` mask. + + Raises: + ValueError: If *ids* select a subset of elements and no ``scratch_mask`` is given. """ # Normalize slice(None) to None so the capture guard treats it identically to ids=None. if isinstance(ids, slice) and ids == slice(None): ids = None if wp.get_device().is_capturing: - if ids is not None or (mask is not None and not isinstance(mask, wp.array)): + if ids is not None or (mask is not None and _mask_needs_copy(mask, device)): raise RuntimeError( - "resolve_1d_mask is only capturable when mask is a wp.array or both ids and mask are None." + "resolve_1d_mask is only capturable when ids are None and mask is None or a boolean mask on device." ) # --- Direct mask input --- if mask is not None: if isinstance(mask, wp.array): + if mask.device != wp.get_device(device): + mask = wp.clone(mask, device=device) return mask if isinstance(mask, torch.Tensor): if mask.dtype != torch.bool: @@ -78,6 +86,9 @@ def resolve_1d_mask( if ids is None: return all_mask + if scratch_mask is None: + raise ValueError("resolve_1d_mask needs a scratch_mask to write ids into.") + if isinstance(ids, slice): scratch_mask.zero_() wp.to_torch(scratch_mask)[ids] = True @@ -115,6 +126,104 @@ def resolve_1d_mask( return scratch_mask +def _mask_needs_copy(mask: wp.array | torch.Tensor, device: str) -> bool: + """Whether :func:`resolve_1d_mask` casts or copies ``mask``, which a graph capture cannot record.""" + if isinstance(mask, wp.array): + return mask.device != wp.get_device(device) + if isinstance(mask, torch.Tensor): + return mask.dtype != torch.bool or str(mask.device) != device + return True + + +## +# Id resolution - ids/mask to ids. +## + + +def resolve_ids( + *, + ids: Sequence[int] | slice | torch.Tensor | wp.array | None = None, + mask: wp.array | torch.Tensor | None = None, + all_ids: wp.array, + device: str, + dtype: type | None = None, +) -> wp.array | torch.Tensor: + """Resolve ids/mask into ids on ``device``. The mask takes precedence over the ids. + + Args: + ids: Index ids. Accepts ``Sequence[int]``, ``slice``, ``torch.Tensor``, ``wp.array``, + or ``None`` (all elements). Tensors and warp arrays keep their integer dtype + unless ``dtype`` is given. Pass ids from ``find_*(as_proxy=True)`` as ``.warp`` or ``.torch``. + mask: Boolean mask. Its ids are found with :func:`torch.nonzero`, which waits for + the device. + all_ids: Pre-allocated ids of every element, returned as-is when *ids* and *mask* + select every element. + device: Warp device string (e.g. ``"cuda:0"``). + dtype: Warp integer dtype the ids must have, for callers whose kernels take one + index dtype. The ids are then a contiguous 1-D ``wp.array`` of this dtype. + Defaults to None: a mask resolves to an int64 ``torch.Tensor``, and ids keep + their type and dtype. + + Returns: + ``all_ids``, a ``torch.Tensor`` or a ``wp.array`` of ids on ``device``. + + Raises: + TypeError: If *ids* is a :class:`ProxyArray`, which is output-only, *mask* is + neither a ``wp.array`` nor a ``torch.Tensor``, or ``dtype`` is given and the + ids are not 32-bit or 64-bit integers. + ValueError: If ``dtype`` is given and the ids are not one-dimensional. + """ + if mask is not None: + if isinstance(mask, wp.array): + mask = wp.to_torch(mask) + elif not isinstance(mask, torch.Tensor): + raise TypeError(f"Unsupported mask type: {type(mask)}") + ids = torch.nonzero(mask.to(wp.device_to_torch(device))).squeeze(-1) + else: + ids = _ids_on_device(ids, all_ids, device) + return ids if dtype is None else _ids_as_dtype(ids, dtype) + + +def _ids_on_device( + ids: Sequence[int] | slice | torch.Tensor | wp.array | None, all_ids: wp.array, device: str +) -> wp.array | torch.Tensor: + """Ids on ``device``, keeping the dtype of tensors and warp arrays; lists and slices become ``wp.int32``.""" + if isinstance(ids, torch.Tensor): + return ids if str(ids.device) == str(wp.get_device(device)) else ids.to(wp.device_to_torch(device)) + if ids is None: + return all_ids + if isinstance(ids, wp.array): + return ids if ids.device == wp.get_device(device) else ids.to(device) + if isinstance(ids, ProxyArray): + raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") + if isinstance(ids, slice): + if ids == slice(None): + return all_ids + start, stop, step = ids.indices(all_ids.shape[0]) + if step > 0: + # a strided view of the caller's all-ids, no copy + return wp.from_torch(wp.to_torch(all_ids)[start:stop:step]) + return wp.array(list(range(start, stop, step)), dtype=wp.int32, device=device) + return wp.array(list(ids), dtype=wp.int32, device=device) + + +def _ids_as_dtype(ids: wp.array | torch.Tensor, dtype: type) -> wp.array: + """``ids`` as a contiguous 1-D ``wp.array`` of ``dtype`` on their device.""" + if dtype not in (wp.int32, wp.int64): + raise TypeError(f"Unsupported id dtype: {dtype}. Expected wp.int32 or wp.int64.") + if ids.ndim != 1: + raise ValueError(f"Ids must be a one-dimensional array, got {ids.ndim} dimensions.") + if isinstance(ids, wp.array): + if ids.dtype == dtype and ids.is_contiguous: + return ids + if ids.dtype not in (wp.int32, wp.int64): + raise TypeError(f"Warp ids must be 32-bit or 64-bit signed integers, got {ids.dtype}.") + ids = wp.to_torch(ids) + elif ids.dtype.is_floating_point or ids.dtype.is_complex or ids.dtype == torch.bool: + raise TypeError(f"Ids must be integers, got {ids.dtype}.") + return wp.from_torch(ids.to(wp.dtype_to_torch(dtype)).contiguous(), dtype=dtype) + + ## # Capture safety — property guard. ## diff --git a/source/isaaclab/test/utils/warp/test_resolve_selection.py b/source/isaaclab/test/utils/warp/test_resolve_selection.py new file mode 100644 index 000000000000..fc951df8f614 --- /dev/null +++ b/source/isaaclab/test/utils/warp/test_resolve_selection.py @@ -0,0 +1,155 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Tests for resolving index selections with :func:`resolve_ids` and :func:`resolve_1d_mask`.""" + +import pytest +import torch +import warp as wp + +from isaaclab.test.utils import test_devices +from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids + + +def _all_ids(n: int, device: str) -> wp.array: + return wp.array(list(range(n)), dtype=wp.int32, device=device) + + +@pytest.mark.parametrize("device", test_devices()) +def test_resolve_ids(device): + """Every selection resolves to device ids; a selection of every element is the shared constant itself.""" + all_ids = _all_ids(10, device) + assert resolve_ids(all_ids=all_ids, device=device) is all_ids + assert resolve_ids(ids=slice(None), all_ids=all_ids, device=device) is all_ids + assert resolve_ids(ids=[3, 1], all_ids=all_ids, device=device).numpy().tolist() == [3, 1] + assert resolve_ids(ids=slice(2, 8, 3), all_ids=all_ids, device=device).numpy().tolist() == [2, 5] + assert resolve_ids(ids=slice(None, None, -4), all_ids=all_ids, device=device).numpy().tolist() == [9, 5, 1] + ids_torch = torch.tensor([4, 2], dtype=torch.int64, device=wp.device_to_torch(device)) + assert resolve_ids(ids=ids_torch, all_ids=all_ids, device=device) is ids_torch + ids_wp = wp.array([6], dtype=wp.int32, device=device) + assert resolve_ids(ids=ids_wp, all_ids=all_ids, device=device) is ids_wp + with pytest.raises(TypeError, match="ProxyArray"): + resolve_ids(ids=ProxyArray(ids_wp), all_ids=all_ids, device=device) + + +@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("mask_type", ["warp", "torch_bool", "torch_int"]) +def test_resolve_ids_from_a_mask(device, mask_type): + """A mask resolves to its ascending ids, takes precedence over ids, and honours ``dtype``.""" + n = 4096 + expected = (torch.rand(n, generator=torch.Generator().manual_seed(0)) < 0.5).to(wp.device_to_torch(device)) + mask = { + "warp": wp.from_torch(expected, dtype=wp.bool), + "torch_bool": expected, + "torch_int": expected.to(torch.int64), + }[mask_type] + all_ids = _all_ids(n, device) + ids = resolve_ids(ids=[0, 1], mask=mask, all_ids=all_ids, device=device) + assert ids.dtype == torch.int64 and torch.equal(ids, torch.nonzero(expected).flatten()) + ids_int32 = resolve_ids(mask=mask, all_ids=all_ids, device=device, dtype=wp.int32) + assert isinstance(ids_int32, wp.array) and ids_int32.dtype == wp.int32 and ids_int32.is_contiguous + assert torch.equal(wp.to_torch(ids_int32), torch.nonzero(expected).flatten().to(torch.int32)) + + +@pytest.mark.parametrize("device", test_devices()) +def test_resolve_ids_with_dtype(device): + """``dtype=wp.int32`` gives a contiguous one-dimensional int32 Warp array for every kind of ids.""" + all_ids = _all_ids(10, device) + torch_device = wp.device_to_torch(device) + assert resolve_ids(all_ids=all_ids, device=device, dtype=wp.int32) is all_ids + for ids, expected in ( + (torch.tensor([4, 2], dtype=torch.int64, device=torch_device), [4, 2]), + (torch.arange(10, device=torch_device)[::3], [0, 3, 6, 9]), + (wp.array([7, 1], dtype=wp.int64, device=device), [7, 1]), + (slice(1, 8, 3), [1, 4, 7]), + (slice(None, None, -4), [9, 5, 1]), + ([3, 8], [3, 8]), + ): + resolved = resolve_ids(ids=ids, all_ids=all_ids, device=device, dtype=wp.int32) + assert isinstance(resolved, wp.array) and resolved.dtype == wp.int32 and resolved.is_contiguous + assert resolved.numpy().tolist() == expected + ids_int32 = wp.array([5], dtype=wp.int32, device=device) + assert resolve_ids(ids=ids_int32, all_ids=all_ids, device=device, dtype=wp.int32) is ids_int32 + with pytest.raises(TypeError): + resolve_ids(ids=torch.tensor([1.0], device=torch_device), all_ids=all_ids, device=device, dtype=wp.int32) + with pytest.raises(TypeError): + resolve_ids( + ids=wp.array([1.0], dtype=wp.float32, device=device), all_ids=all_ids, device=device, dtype=wp.int32 + ) + with pytest.raises(ValueError): + resolve_ids( + ids=torch.zeros((2, 2), dtype=torch.int64, device=torch_device), + all_ids=all_ids, + device=device, + dtype=wp.int32, + ) + + +@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("ids_type", ["list", "slice", "torch_int64", "torch_int32", "warp_int32"]) +def test_resolve_1d_mask_from_ids(device, ids_type): + """Ids of every type are written into the scratch mask, which is cleared first.""" + n = 12 + all_mask = wp.ones(n, dtype=wp.bool, device=device) + scratch = wp.ones(n, dtype=wp.bool, device=device) + selected = [1, 4, 5, 11] + torch_device = wp.device_to_torch(device) + ids = { + "list": selected, + "slice": slice(1, 12, 3), + "torch_int64": torch.tensor(selected, dtype=torch.int64, device=torch_device), + "torch_int32": torch.tensor(selected, dtype=torch.int32, device=torch_device), + "warp_int32": wp.array(selected, dtype=wp.int32, device=device), + }[ids_type] + expected = [False] * n + for i in range(n)[ids] if isinstance(ids, slice) else selected: + expected[i] = True + mask = resolve_1d_mask(ids=ids, all_mask=all_mask, scratch_mask=scratch, device=device) + assert mask is scratch + assert mask.numpy().tolist() == expected + + +@pytest.mark.parametrize("device", test_devices()) +def test_resolve_1d_mask_needs_a_scratch_mask_only_for_ids(device): + """A mask or no selection resolves without a scratch mask; ids without one raise.""" + all_mask = wp.ones(5, dtype=wp.bool, device=device) + mask = wp.array([True, False, True, False, False], dtype=wp.bool, device=device) + assert resolve_1d_mask(ids=[0], mask=mask, all_mask=all_mask, device=device) is mask + assert resolve_1d_mask(all_mask=all_mask, device=device) is all_mask + assert resolve_1d_mask(ids=slice(None), all_mask=all_mask, device=device) is all_mask + mask_torch = torch.tensor([0, 1, 1, 0, 0], device=wp.device_to_torch(device)) + resolved = resolve_1d_mask(mask=mask_torch, all_mask=all_mask, device=device) + assert resolved.dtype == wp.bool and resolved.numpy().tolist() == [False, True, True, False, False] + with pytest.raises(ValueError, match="scratch_mask"): + resolve_1d_mask(ids=[1], all_mask=all_mask, device=device) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_resolve_1d_mask_copies_a_mask_on_another_device(): + """A Warp mask on another device is copied to ``device``.""" + all_mask = wp.ones(3, dtype=wp.bool, device="cuda:0") + mask = wp.array([True, False, True], dtype=wp.bool, device="cpu") + resolved = resolve_1d_mask(mask=mask, all_mask=all_mask, device="cuda:0") + assert str(resolved.device) == "cuda:0" and resolved.numpy().tolist() == [True, False, True] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_resolve_1d_mask_in_a_graph_capture(): + """Inside a capture, a boolean mask on device resolves; ids or a mask that needs a cast or copy raise.""" + device = "cuda:0" + all_mask = wp.ones(3, dtype=wp.bool, device=device) + scratch = wp.zeros(3, dtype=wp.bool, device=device) + mask_wp = wp.array([True, False, True], dtype=wp.bool, device=device) + mask_torch = torch.tensor([False, True, True], device=device) + mask_int = mask_torch.int() + mask_cpu = wp.array([True, False, True], dtype=wp.bool, device="cpu") + with wp.ScopedDevice(device), wp.ScopedCapture(): + assert resolve_1d_mask(mask=mask_wp, all_mask=all_mask, device=device) is mask_wp + resolved = resolve_1d_mask(mask=mask_torch, all_mask=all_mask, device=device) + for selection in ({"mask": mask_int}, {"mask": mask_cpu}, {"ids": [0]}): + with pytest.raises(RuntimeError, match="capturable"): + resolve_1d_mask(**selection, all_mask=all_mask, scratch_mask=scratch, device=device) + assert resolved.numpy().tolist() == [False, True, True] diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-id-mask.changed.rst new file mode 100644 index 000000000000..ce57c9e7851f --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/jichuanh-id-mask.changed.rst @@ -0,0 +1,2 @@ +* Changed ``isaaclab_experimental.utils.warp.resolve_1d_mask`` to be + :func:`isaaclab.utils.warp.utils.resolve_1d_mask` instead of a copy of it. diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py index 2d071b823a46..757a45e8b825 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py @@ -5,11 +5,12 @@ """Warp utility functions and shared kernels for isaaclab_experimental.""" +from isaaclab.utils.warp.utils import resolve_1d_mask + from .kernels import compute_reset_scale, count_masked from .utils import ( WarpCapturable, is_warp_capturable, - resolve_1d_mask, warp_capturable, wrap_to_pi, zero_masked_2d, diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py index a90f6b2d6fc6..6b8aafc415f8 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py @@ -5,97 +5,9 @@ from __future__ import annotations -from collections.abc import Sequence - -import torch import warp as wp -@wp.kernel -def _set_mask_from_ids( - mask: wp.array(dtype=wp.bool), - ids: wp.array(dtype=wp.int32), -): - """Set ``mask[ids[i]] = True`` for each thread *i*.""" - i = wp.tid() - mask[ids[i]] = True - - -def resolve_1d_mask( - *, - ids: Sequence[int] | slice | wp.array | torch.Tensor | None, - mask: wp.array | torch.Tensor | None, - all_mask: wp.array, - scratch_mask: wp.array, - device: str, -) -> wp.array: - """Resolve ids/mask into a warp boolean mask. - - Matches the contract of ``ArticulationData._resolve_1d_mask`` on dev/newton. - Callers must provide pre-allocated ``all_mask`` (all-True) and ``scratch_mask`` - (reusable working buffer). No allocations happen inside this function. - - Args: - ids: Indices to set to ``True``. ``None`` or ``slice(None)`` means all. - mask: Explicit boolean mask. If provided, returned directly (after - torch->warp normalization if needed). Takes precedence over *ids*. - all_mask: Pre-allocated all-True mask of shape ``(size,)``, returned - when both *ids* and *mask* are ``None``. - scratch_mask: Pre-allocated scratch mask of shape ``(size,)``, filled - in-place when *ids* are provided. - device: Warp device string. - - Returns: - A ``wp.array(dtype=wp.bool)`` -- ``mask``, ``all_mask``, or ``scratch_mask``. - """ - # Fast path: explicit mask provided. - if mask is not None: - if isinstance(mask, torch.Tensor): - if mask.dtype != torch.bool: - mask = mask.to(dtype=torch.bool) - if str(mask.device) != device: - mask = mask.to(device) - return wp.from_torch(mask, dtype=wp.bool) - return mask - - # Fast path: all ids. - if ids is None or (isinstance(ids, slice) and ids == slice(None)): - return all_mask - - # Normalize slice into explicit indices. - if isinstance(ids, slice): - start, stop, step = ids.indices(scratch_mask.shape[0]) - ids = list(range(start, stop, step)) - elif not isinstance(ids, (torch.Tensor, wp.array)): - ids = list(ids) - - # Prepare output mask. - scratch_mask.fill_(False) - - # Normalize ids to wp.int32 array and launch kernel. - if isinstance(ids, torch.Tensor): - if ids.numel() == 0: - return scratch_mask - if str(ids.device) != device: - ids = ids.to(device) - if ids.dtype != torch.int32: - ids = ids.to(dtype=torch.int32) - if not ids.is_contiguous(): - ids = ids.contiguous() - ids_wp = wp.from_torch(ids, dtype=wp.int32) - elif isinstance(ids, wp.array): - if ids.shape[0] == 0: - return scratch_mask - ids_wp = ids - else: - if len(ids) == 0: - return scratch_mask - ids_wp = wp.array(ids, dtype=wp.int32, device=device) - - wp.launch(kernel=_set_mask_from_ids, dim=ids_wp.shape[0], inputs=[scratch_mask, ids_wp], device=device) - return scratch_mask - - def warp_capturable(capturable: bool): """Annotate an MDP term's CUDA-graph capturability. From 97e56623b87fdec9128e4d7105ee076590e07f9e Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Mon, 26 Oct 2026 02:48:17 -0700 Subject: [PATCH 2/6] Resolve sensor, camera and wrench composer selections through the shared helpers SensorBase, Camera and WrenchComposer resolve env_ids, env_mask and body_ids through resolve_ids and resolve_1d_mask instead of their own copies. The test doubles that build these objects without their initializer get the env masks a real sensor creates. --- .../changelog.d/jichuanh-id-mask.changed.rst | 7 ++++ .../isaaclab/sensors/camera/camera.py | 23 +---------- .../isaaclab/isaaclab/sensors/sensor_base.py | 13 ++----- .../isaaclab/utils/wrench_composer.py | 38 +++---------------- .../test_simulation_render_context.py | 1 + .../sensors/test_sensor_update_schedule.py | 1 + .../test/utils/test_wrench_composer.py | 5 ++- 7 files changed, 23 insertions(+), 65 deletions(-) diff --git a/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst index 6a1b73d77c7e..2ae1cf9e1daa 100644 --- a/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst +++ b/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst @@ -1,3 +1,10 @@ +* Changed the env selection helpers of :class:`~isaaclab.utils.wrench_composer.WrenchComposer`, + :class:`~isaaclab.sensors.SensorBase` and :class:`~isaaclab.sensors.Camera` to resolve ``env_ids`` and ``env_mask`` + through :func:`~isaaclab.utils.warp.utils.resolve_ids` and :func:`~isaaclab.utils.warp.utils.resolve_1d_mask`. + Any integer sequence and slices with a negative step are accepted, ids on another device are moved to the owner's + device, and a :class:`~isaaclab.utils.warp.ProxyArray` is rejected with a :class:`TypeError`. +* Changed :class:`~isaaclab.utils.wrench_composer.WrenchComposer` to resolve ``body_ids`` through + :func:`~isaaclab.utils.warp.utils.resolve_ids` as well, so slices other than ``slice(None)`` are accepted. * Changed :func:`~isaaclab.utils.warp.utils.resolve_1d_mask` to make ``scratch_mask`` optional for callers that never pass ids. Ids without a ``scratch_mask`` raise a :class:`ValueError`. * Changed :func:`~isaaclab.utils.warp.utils.resolve_1d_mask` to copy a Warp mask on another device to ``device``, and diff --git a/source/isaaclab/isaaclab/sensors/camera/camera.py b/source/isaaclab/isaaclab/sensors/camera/camera.py index d052dd59aa73..b1e577520bbb 100644 --- a/source/isaaclab/isaaclab/sensors/camera/camera.py +++ b/source/isaaclab/isaaclab/sensors/camera/camera.py @@ -26,6 +26,7 @@ quat_from_matrix, ) from ...utils.warp import ProxyArray +from ...utils.warp.utils import resolve_ids from ..sensor_base import SensorBase from .camera_data import CameraData, RenderBufferKind @@ -1028,27 +1029,7 @@ def _resolve_env_ids_wp(self, env_ids: Sequence[int] | torch.Tensor | wp.array | """Resolve camera indices to a Warp ``int32`` array.""" if env_ids is None: return None - if isinstance(env_ids, wp.array): - if env_ids.dtype != wp.int32: - raise TypeError(f"Unsupported wp.array dtype for env_ids: {env_ids.dtype}. Expected wp.int32.") - if env_ids.ndim != 1: - raise ValueError("Warp camera indices must be a one-dimensional array.") - if str(env_ids.device) == str(self._device): - return env_ids - return env_ids.to(self._device) - elif isinstance(env_ids, torch.Tensor): - env_ids = env_ids.to(device=self._device, dtype=torch.int32).reshape(-1) - if not env_ids.is_contiguous(): - env_ids = env_ids.contiguous() - return wp.from_torch(env_ids, dtype=wp.int32) - elif isinstance(env_ids, slice) and (env_ids.step or 1) > 0: - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - elif isinstance(env_ids, slice): - # Torch cannot slice with a negative step. - env_ids = np.arange(self._view.count, dtype=np.int32)[env_ids] - else: - env_ids = np.asarray(env_ids, dtype=np.int32).reshape(-1) - return wp.array(env_ids, dtype=wp.int32, device=self._device) + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self._device, dtype=wp.int32) def _env_mask_has_any(self, env_mask: wp.array) -> bool: """Return whether the mask selects any camera.""" diff --git a/source/isaaclab/isaaclab/sensors/sensor_base.py b/source/isaaclab/isaaclab/sensors/sensor_base.py index 5bcec23eb91e..f0e39055540d 100644 --- a/source/isaaclab/isaaclab/sensors/sensor_base.py +++ b/source/isaaclab/isaaclab/sensors/sensor_base.py @@ -29,7 +29,7 @@ from ..physics import PhysicsEvent, PhysicsManager from ..sim.utils.queries import get_first_matching_ancestor_prim from ..sim.utils.transforms import resolve_prim_pose -from ..utils import index_fill_ +from ..utils.warp.utils import resolve_1d_mask from .kernels import invalidate_buffers_kernel, reset_envs_kernel, update_outdated_envs_kernel, update_timestamp_kernel if TYPE_CHECKING: @@ -506,14 +506,9 @@ def _resolve_indices_and_mask( self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None ) -> wp.array: """Resolve environment indices to a warp array and mask.""" - if (env_ids is None or env_ids == slice(None)) and env_mask is None: - return self._ALL_ENV_MASK - elif env_mask is not None: - return env_mask - else: - self._reset_mask.zero_() - index_fill_(self._reset_mask_torch, env_ids, True) - return self._reset_mask + return resolve_1d_mask( + ids=env_ids, mask=env_mask, all_mask=self._ALL_ENV_MASK, device=self._device, scratch_mask=self._reset_mask + ) def _resolve_rigid_body_ancestor_expr( self, diff --git a/source/isaaclab/isaaclab/utils/wrench_composer.py b/source/isaaclab/isaaclab/utils/wrench_composer.py index 6331716cad5d..ace1570dc4dc 100644 --- a/source/isaaclab/isaaclab/utils/wrench_composer.py +++ b/source/isaaclab/isaaclab/utils/wrench_composer.py @@ -24,6 +24,7 @@ set_forces_to_dual_buffers_index_kernel, set_forces_to_dual_buffers_mask, ) +from .warp.utils import resolve_ids if TYPE_CHECKING: from ..assets import BaseArticulation, BaseRigidObject, BaseRigidObjectCollection @@ -721,24 +722,9 @@ def _resolve_env_ids(self, env_ids: wp.array | torch.Tensor | list | slice | Non Environment indices. Raises: - TypeError: If ``env_ids`` is an unsupported type. + TypeError: If ``env_ids`` is a :class:`~isaaclab.utils.warp.ProxyArray`. """ - if env_ids is None: - return self._ALL_ENV_INDICES - # Check tensor types before slice comparison (tensor == slice crashes) - if isinstance(env_ids, torch.Tensor): - return env_ids - if isinstance(env_ids, wp.array): - return env_ids - if env_ids == slice(None): - return self._ALL_ENV_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_ENV_INDICES)[env_ids]) - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self.device) - raise TypeError( - f"env_ids must be None, slice(None), list, torch.Tensor, or wp.array, got {type(env_ids).__name__}" - ) + return resolve_ids(ids=env_ids, all_ids=self._ALL_ENV_INDICES, device=self.device) def _resolve_body_ids( self, body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None @@ -752,23 +738,9 @@ def _resolve_body_ids( Body indices. Raises: - TypeError: If ``body_ids`` is an unsupported type. + TypeError: If ``body_ids`` is a :class:`~isaaclab.utils.warp.ProxyArray`. """ - if body_ids is None: - return self._ALL_BODY_INDICES - if isinstance(body_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(body_ids, torch.Tensor): - return body_ids - if isinstance(body_ids, wp.array): - return body_ids - if body_ids == slice(None): - return self._ALL_BODY_INDICES - if isinstance(body_ids, Sequence): - return wp.array(body_ids, dtype=wp.int32, device=self.device) - raise TypeError( - f"body_ids must be None, slice(None), a sequence, torch.Tensor, or wp.array, got {type(body_ids).__name__}" - ) + return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self.device) def _ensure_composed(self): """Compose input buffers into output buffers if dirty.""" diff --git a/source/isaaclab/test/renderers/test_simulation_render_context.py b/source/isaaclab/test/renderers/test_simulation_render_context.py index fdcc8d4f3779..cf16bc4a2d76 100644 --- a/source/isaaclab/test/renderers/test_simulation_render_context.py +++ b/source/isaaclab/test/renderers/test_simulation_render_context.py @@ -317,6 +317,7 @@ def __init__(self, renderer, name, update_period=0.0): self._frame = ProxyArray(wp.zeros(2, dtype=wp.int64, device="cpu")) self._ALL_INDICES = wp.array([0, 1], dtype=wp.int32, device="cpu") self._ALL_ENV_MASK = wp.ones(2, dtype=wp.bool, device="cpu") + self._reset_mask = wp.zeros(2, dtype=wp.bool, device="cpu") self._create_timing_buffers() self._invalidate_buffers(self._ALL_ENV_MASK) self.pose = 0.0 diff --git a/source/isaaclab/test/sensors/test_sensor_update_schedule.py b/source/isaaclab/test/sensors/test_sensor_update_schedule.py index 66224a3fea17..39353ecefdb3 100644 --- a/source/isaaclab/test/sensors/test_sensor_update_schedule.py +++ b/source/isaaclab/test/sensors/test_sensor_update_schedule.py @@ -70,6 +70,7 @@ def __init__(self, update_period: float, device: str, num_envs: int = 2): self._device = device self._num_envs = num_envs self._ALL_ENV_MASK = wp.ones(num_envs, dtype=wp.bool, device=device) + self._reset_mask = wp.zeros(num_envs, dtype=wp.bool, device=device) self._is_initialized = True self._is_visualizing = False self._create_timing_buffers() diff --git a/source/isaaclab/test/utils/test_wrench_composer.py b/source/isaaclab/test/utils/test_wrench_composer.py index 10daf0d06612..b236b052707a 100644 --- a/source/isaaclab/test/utils/test_wrench_composer.py +++ b/source/isaaclab/test/utils/test_wrench_composer.py @@ -477,8 +477,9 @@ def test_deprecated_writers_warn_and_keep_add_and_set_behavior() -> None: def test_invalid_selection_and_empty_wrench_input_are_reported() -> None: composer = _make_composer() - with pytest.raises(TypeError, match="env_ids must be"): - composer.add_forces_and_torques_index(forces=_vectors(torch.ones((1, 2, 3))), env_ids=(0,)) + env_ids = ProxyArray(wp.array([0], dtype=wp.int32, device=composer.device)) + with pytest.raises(TypeError, match="ProxyArray is output-only"): + composer.add_forces_and_torques_index(forces=_vectors(torch.ones((1, 2, 3))), env_ids=env_ids) for method in ("add_forces_and_torques_index", "set_forces_and_torques_mask"): with pytest.warns(UserWarning, match="No forces or torques"): getattr(composer, method)() From 92590d00f98878892bfdc8dde0d1eb56b77cde19 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Mon, 26 Oct 2026 02:48:18 -0700 Subject: [PATCH 3/6] Fix the deformable mask writers that raised on any env mask The nodal mask writers of BaseDeformableObject, which the OVPhysX deformable object inherits, and of the PhysX deformable object called wp.nonzero on a Warp array, which raised a RuntimeError. They now resolve the mask to ids with resolve_ids and call the index writers. --- .../changelog.d/jichuanh-id-mask.fixed.rst | 2 ++ .../base_deformable_object.py | 24 +++++--------- .../changelog.d/jichuanh-id-mask.fixed.rst | 2 ++ .../deformable_object/deformable_object.py | 33 ++++--------------- .../test/assets/test_deformable_object.py | 16 +++++++++ 5 files changed, 34 insertions(+), 43 deletions(-) create mode 100644 source/isaaclab/changelog.d/jichuanh-id-mask.fixed.rst create mode 100644 source/isaaclab_physx/changelog.d/jichuanh-id-mask.fixed.rst diff --git a/source/isaaclab/changelog.d/jichuanh-id-mask.fixed.rst b/source/isaaclab/changelog.d/jichuanh-id-mask.fixed.rst new file mode 100644 index 000000000000..9ed262f73ec2 --- /dev/null +++ b/source/isaaclab/changelog.d/jichuanh-id-mask.fixed.rst @@ -0,0 +1,2 @@ +* Fixed the ``write_nodal_*_to_sim_mask`` methods of :class:`~isaaclab.assets.BaseDeformableObject`, which the + OVPhysX deformable object inherits, raising a :class:`RuntimeError` whenever ``env_mask`` was given. diff --git a/source/isaaclab/isaaclab/assets/deformable_object/base_deformable_object.py b/source/isaaclab/isaaclab/assets/deformable_object/base_deformable_object.py index 0542a460af48..4467403b79cf 100644 --- a/source/isaaclab/isaaclab/assets/deformable_object/base_deformable_object.py +++ b/source/isaaclab/isaaclab/assets/deformable_object/base_deformable_object.py @@ -15,6 +15,7 @@ from ...utils import math as math_utils from ...utils.warp import ProxyArray +from ...utils.warp.utils import resolve_ids from ..asset_base import AssetBase if TYPE_CHECKING: @@ -200,6 +201,9 @@ def write_nodal_kinematic_target_to_sim_index( """ Operations - Write to simulation (mask variants). + + These convert the mask to ids, which waits for the device, and call the index variants. + Backends with a mask write path override them. """ def write_nodal_state_to_sim_mask( @@ -214,10 +218,7 @@ def write_nodal_state_to_sim_mask( Shape is (num_instances, max_sim_vertices_per_body, 6). env_mask: Environment mask. If None, then all indices are used. """ - if env_mask is not None: - env_ids = wp.nonzero(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) self.write_nodal_state_to_sim_index(nodal_state, env_ids=env_ids, full_data=True) def write_nodal_pos_to_sim_mask( @@ -232,10 +233,7 @@ def write_nodal_pos_to_sim_mask( Shape is (num_instances, max_sim_vertices_per_body, 3). env_mask: Environment mask. If None, then all indices are used. """ - if env_mask is not None: - env_ids = wp.nonzero(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) self.write_nodal_pos_to_sim_index(nodal_pos, env_ids=env_ids, full_data=True) def write_nodal_velocity_to_sim_mask( @@ -250,10 +248,7 @@ def write_nodal_velocity_to_sim_mask( Shape is (num_instances, max_sim_vertices_per_body, 3). env_mask: Environment mask. If None, then all indices are used. """ - if env_mask is not None: - env_ids = wp.nonzero(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) self.write_nodal_velocity_to_sim_index(nodal_vel, env_ids=env_ids, full_data=True) def write_nodal_kinematic_target_to_sim_mask( @@ -268,10 +263,7 @@ def write_nodal_kinematic_target_to_sim_mask( Shape is (num_instances, max_sim_vertices_per_body, 4). env_mask: Environment mask. If None, then all indices are used. """ - if env_mask is not None: - env_ids = wp.nonzero(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) self.write_nodal_kinematic_target_to_sim_index(targets, env_ids=env_ids, full_data=True) """ diff --git a/source/isaaclab_physx/changelog.d/jichuanh-id-mask.fixed.rst b/source/isaaclab_physx/changelog.d/jichuanh-id-mask.fixed.rst new file mode 100644 index 000000000000..96c018a902e7 --- /dev/null +++ b/source/isaaclab_physx/changelog.d/jichuanh-id-mask.fixed.rst @@ -0,0 +1,2 @@ +* Fixed the ``write_nodal_*_to_sim_mask`` methods of :class:`~isaaclab_physx.assets.DeformableObject` raising a + :class:`RuntimeError` whenever ``env_mask`` was given. diff --git a/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py b/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py index 019d50bdb1a8..d8dc3d9e7bc3 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py @@ -21,6 +21,7 @@ from isaaclab.assets.asset_base import AssetBase from isaaclab.markers import VisualizationMarkers from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_ids from .deformable_object_data import DeformableObjectData from .kernels import ( @@ -218,10 +219,7 @@ def write_nodal_state_to_sim_mask( Shape is (num_instances, max_sim_vertices_per_body, 6). env_mask: Environment mask. If None, then all indices are used. """ - if env_mask is not None: - env_ids = wp.nonzero(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) self.write_nodal_state_to_sim_index(nodal_state, env_ids=env_ids, full_data=True) def write_nodal_pos_to_sim_index( @@ -288,10 +286,7 @@ def write_nodal_pos_to_sim_mask( Shape is (num_instances, max_sim_vertices_per_body, 3). env_mask: Environment mask. If None, then all indices are used. """ - if env_mask is not None: - env_ids = wp.nonzero(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) self.write_nodal_pos_to_sim_index(nodal_pos, env_ids=env_ids, full_data=True) def write_nodal_velocity_to_sim_index( @@ -359,10 +354,7 @@ def write_nodal_velocity_to_sim_mask( Shape is (num_instances, max_sim_vertices_per_body, 3). env_mask: Environment mask. If None, then all indices are used. """ - if env_mask is not None: - env_ids = wp.nonzero(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) self.write_nodal_velocity_to_sim_index(nodal_vel, env_ids=env_ids, full_data=True) def write_nodal_kinematic_target_to_sim_index( @@ -438,10 +430,7 @@ def write_nodal_kinematic_target_to_sim_mask( Shape is (num_instances, max_sim_vertices_per_body, 4). env_mask: Environment mask. If None, then all indices are used. """ - if env_mask is not None: - env_ids = wp.nonzero(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) self.write_nodal_kinematic_target_to_sim_index(targets, env_ids=env_ids, full_data=True) """ @@ -566,17 +555,7 @@ def _resolve_env_ids(self, env_ids): Returns: A warp array of environment indices. """ - if (env_ids is None) or (env_ids == slice(None)): - return self._ALL_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - if isinstance(env_ids, torch.Tensor): - if env_ids.dtype == torch.int64: - env_ids = env_ids.to(torch.int32) - return wp.from_torch(env_ids, dtype=wp.int32) - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self.device) - return env_ids + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device, dtype=wp.int32) def _initialize_impl(self): # obtain global simulation view diff --git a/source/isaaclab_physx/test/assets/test_deformable_object.py b/source/isaaclab_physx/test/assets/test_deformable_object.py index cfd389fa99ef..4eeb4c246f8c 100644 --- a/source/isaaclab_physx/test/assets/test_deformable_object.py +++ b/source/isaaclab_physx/test/assets/test_deformable_object.py @@ -244,6 +244,22 @@ def test_deformable_nodal_state_writes(deformable_scene: _DeformableScene) -> No torch.testing.assert_close(raw_nodal_pos.reshape(nodal_pos.shape), nodal_pos) +@pytest.mark.isaacsim_ci +def test_deformable_nodal_state_mask_writes(deformable_scene: _DeformableScene) -> None: + """Mask writes take full-size data and reach only the environments the mask selects.""" + scene = deformable_scene + device = scene.device + env_mask = wp.array([False, True], dtype=wp.bool, device=device) + for deformable in (scene.volume, scene.surface): + initial_state = deformable.data.nodal_state_w.torch.clone() + state = initial_state.clone() + state[:, 0, :3] += torch.tensor([0.01, 0.02, 0.03], device=device) + state[:, 0, 3] = 0.25 + deformable.write_nodal_state_to_sim_mask(state, env_mask=env_mask) + expected = torch.cat((initial_state[:1], state[1:])) + torch.testing.assert_close(deformable.data.nodal_state_w.torch, expected) + + @pytest.mark.isaacsim_ci def test_volume_kinematic_targets_hold_selected_nodes(deformable_scene: _DeformableScene) -> None: """Kinematic targets of environment 0 hold its nodes while environment 1 keeps moving freely.""" From f044c920104c97249838a46e6dbb60948b5ea3e3 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Mon, 26 Oct 2026 02:48:20 -0700 Subject: [PATCH 4/6] Resolve Newton selections through the shared helpers and write its mask writers through masks The Newton assets resolve ids and masks through resolve_ids and resolve_1d_mask. The RigidObjectCollection mask writers and the fixed-tendon mask setters write through the masks instead of converting them to ids, which waited for the GPU on every call; the joint-coordinate writes reuse a scratch mask; one generic masked 2-D write kernel replaces two copies. --- .../assets/contract/articulation_factory.py | 1 + .../assets/contract/rigid_object_factory.py | 1 + .../contract/test_articulation_contract.py | 5 - .../contract/test_rigid_object_contract.py | 21 ++ .../changelog.d/jichuanh-id-mask.changed.rst | 13 ++ .../changelog.d/jichuanh-id-mask.fixed.rst | 2 + .../assets/articulation/articulation.py | 159 +++++++-------- .../assets/articulation/kernels.py | 30 +++ .../assets/cable_object/cable_object.py | 11 +- .../deformable_object/deformable_object.py | 17 +- .../isaaclab_newton/assets/kernels.py | 113 ++++++++++- .../assets/mpm_object/mpm_object.py | 19 +- .../assets/rigid_object/rigid_object.py | 17 +- .../rigid_object_collection.py | 189 ++++++++---------- 14 files changed, 332 insertions(+), 266 deletions(-) create mode 100644 source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst create mode 100644 source/isaaclab_newton/changelog.d/jichuanh-id-mask.fixed.rst diff --git a/source/isaaclab/test/assets/contract/articulation_factory.py b/source/isaaclab/test/assets/contract/articulation_factory.py index 8ff52e4fdafb..6b0488bb006b 100644 --- a/source/isaaclab/test/assets/contract/articulation_factory.py +++ b/source/isaaclab/test/assets/contract/articulation_factory.py @@ -237,6 +237,7 @@ def _newton_articulation(N, J, B, FT, ST, device, is_fixed_base, joint_ordering, articulation._ALL_BODY_INDICES = indices(B, device) articulation._ALL_JOINT_INDICES = indices(J, device) articulation._ALL_ENV_MASK = wp.ones((N,), dtype=wp.bool, device=device) + articulation._scratch_env_mask = wp.zeros((N,), dtype=wp.bool, device=device) articulation._ALL_BODY_MASK = wp.ones((B,), dtype=wp.bool, device=device) articulation._ALL_JOINT_MASK = wp.ones((J,), dtype=wp.bool, device=device) articulation._resolve_and_install_ordering_maps() diff --git a/source/isaaclab/test/assets/contract/rigid_object_factory.py b/source/isaaclab/test/assets/contract/rigid_object_factory.py index 0bf4741ed9f5..cb108da59a1c 100644 --- a/source/isaaclab/test/assets/contract/rigid_object_factory.py +++ b/source/isaaclab/test/assets/contract/rigid_object_factory.py @@ -195,6 +195,7 @@ def _newton_collection(num_instances: int, num_bodies: int, device: str, monkeyp collection._ALL_BODY_INDICES = indices(num_bodies, device) collection._ALL_ENV_MASK = wp.ones((num_instances,), dtype=wp.bool, device=device) collection._ALL_BODY_MASK = wp.ones((num_bodies,), dtype=wp.bool, device=device) + collection._scratch_body_mask = wp.zeros((num_bodies,), dtype=wp.bool, device=device) return collection, view diff --git a/source/isaaclab/test/assets/contract/test_articulation_contract.py b/source/isaaclab/test/assets/contract/test_articulation_contract.py index 99abc45ae632..42878092b782 100644 --- a/source/isaaclab/test/assets/contract/test_articulation_contract.py +++ b/source/isaaclab/test/assets/contract/test_articulation_contract.py @@ -536,11 +536,6 @@ def read() -> ProxyArray: @requires("newton") -@pytest.mark.xfail( - raises=pytest.fail.Exception, - strict=True, - reason="Newton fixed-tendon mask setters slice full data without validating its shape", -) def test_newton_fixed_tendon_mask_setters_reject_extra_environments(monkeypatch): art, _ = _articulation("newton", "cpu", monkeypatch) with pytest.raises((AssertionError, RuntimeError)): diff --git a/source/isaaclab/test/assets/contract/test_rigid_object_contract.py b/source/isaaclab/test/assets/contract/test_rigid_object_contract.py index 1d294d9a8efc..514f4811ad57 100644 --- a/source/isaaclab/test/assets/contract/test_rigid_object_contract.py +++ b/source/isaaclab/test/assets/contract/test_rigid_object_contract.py @@ -314,6 +314,27 @@ def test_pose_and_velocity_writer_aliases_forward_to_canonical_writers(monkeypat torch.testing.assert_close(getattr(obj.data, getter).torch, values, atol=1e-5, rtol=1e-5) +@pytest.mark.parametrize("quantity", ["link_pose", "com_pose", "com_velocity", "link_velocity"]) +def test_collection_mask_writers_select_bodies_by_ids(monkeypatch, backend, device, quantity): + """Collection mask writers given ``body_ids`` write only those bodies of the masked environments.""" + if backend == "ovphysx": + pytest.skip("OVPhysX collection mask writers take no body_ids.") + obj, _, _, num_bodies = _create("collection", backend, monkeypatch, device) + obj.data.update(dt=0.01) + trailing = (7,) if quantity.endswith("pose") else (6,) + kwarg = "body_poses" if quantity.endswith("pose") else "body_velocities" + getter = f"body_{quantity.replace('velocity', 'vel')}_w" + method = getattr(obj, f"write_body_{quantity}_to_sim_mask") + full = _payload((_NUM_INSTANCES, num_bodies), trailing, device=device) + method(**{kwarg: full}) + update = _payload((_NUM_INSTANCES, num_bodies), trailing, offset=100.0, device=device) + env_mask = wp.array(np.arange(_NUM_INSTANCES) == _NUM_INSTANCES - 1, dtype=wp.bool, device=str(device)) + method(**{kwarg: update}, env_mask=env_mask, body_ids=[num_bodies - 1]) + expected = full.clone() + expected[-1, -1] = update[-1, -1] + torch.testing.assert_close(getattr(obj.data, getter).torch, expected, atol=1e-5, rtol=1e-5) + + # backend-stored quantity -> (written quantity, per-element shape) _PARTIAL_WRITES = {"pose": ("link_pose", (7,)), "velocity": ("com_velocity", (6,)), "mass": ("masses", ())} diff --git a/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst new file mode 100644 index 000000000000..0c69e3841f70 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst @@ -0,0 +1,13 @@ +* Changed the env selection helpers of the Newton assets to resolve ``env_ids`` through + :func:`~isaaclab.utils.warp.utils.resolve_ids` and ``env_mask`` through + :func:`~isaaclab.utils.warp.utils.resolve_1d_mask`. Any integer sequence and slices with a negative step are accepted, + ids on another device are moved to the asset's device, and a :class:`~isaaclab.utils.warp.ProxyArray` is rejected + with a :class:`TypeError`. +* Changed the body, joint and tendon selection helpers of the Newton assets to resolve ids and masks through + :func:`~isaaclab.utils.warp.utils.resolve_ids`, as the env selection does, so slices other than ``slice(None)`` are + accepted. +* Changed the joint-coordinate writes of :class:`~isaaclab_newton.assets.Articulation` to write env ids into a + scratch mask of the articulation instead of allocating a mask on every call. +* Changed the mask writers of :class:`~isaaclab_newton.assets.RigidObjectCollection` and the fixed-tendon mask setters + of :class:`~isaaclab_newton.assets.Articulation` to write through the masks, as the other Newton mask writers do, + instead of converting the masks to indices, which synchronized the host with the GPU on every call. diff --git a/source/isaaclab_newton/changelog.d/jichuanh-id-mask.fixed.rst b/source/isaaclab_newton/changelog.d/jichuanh-id-mask.fixed.rst new file mode 100644 index 000000000000..638c02730a49 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/jichuanh-id-mask.fixed.rst @@ -0,0 +1,2 @@ +* Fixed the fixed-tendon mask setters of :class:`~isaaclab_newton.assets.Articulation` accepting data whose shape + does not match the articulation. diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index f6b03aa5b673..5ed2cb339d61 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -12,7 +12,7 @@ import re import warnings from collections.abc import Sequence -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING import numpy as np import torch @@ -33,6 +33,7 @@ from isaaclab.utils.string import resolve_matching_names, resolve_matching_names_values from isaaclab.utils.version import get_isaac_sim_version, has_kit from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_newton.assets import kernels as shared_kernels @@ -2298,15 +2299,11 @@ def write_joint_viscous_friction_coefficient_to_sim_mask( Operations - Newton Actuator Parameter Writers. """ - @staticmethod - @wp.kernel(enable_backward=False) - def _build_env_mask_kernel(mask: wp.array(dtype=wp.bool), indices: wp.array(dtype=Any)): - i = wp.tid() - mask[wp.int32(indices[i])] = True - def _env_ids_to_mask(self, env_ids: wp.array | torch.Tensor) -> wp.array: """Convert env_ids to a boolean Warp mask. + Needed because :func:`scatter_joint_coordinates` only takes a mask. + Args: env_ids: Environment indices as returned by :meth:`_resolve_env_ids`, which may be a warp array or a torch tensor of any integer width. @@ -2316,11 +2313,11 @@ def _env_ids_to_mask(self, env_ids: wp.array | torch.Tensor) -> wp.array: """ if env_ids is self._ALL_INDICES: return self._ALL_ENV_MASK - if isinstance(env_ids, torch.Tensor): - env_ids = wp.from_torch(env_ids) - mask = wp.zeros(self.num_instances, dtype=wp.bool, device=self.device) - wp.launch(self._build_env_mask_kernel, dim=env_ids.shape[0], inputs=[mask, env_ids], device=self.device) - return mask + if isinstance(env_ids, wp.array) and env_ids.dtype != wp.int32: + env_ids = wp.to_torch(env_ids) + return resolve_1d_mask( + ids=env_ids, all_mask=self._ALL_ENV_MASK, scratch_mask=self._scratch_env_mask, device=self.device + ) """ Operations - Setters. @@ -2716,12 +2713,17 @@ def set_fixed_tendon_stiffness_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask) - self.set_fixed_tendon_stiffness_index( - stiffness=self._select_full_data(stiffness, env_ids, fixed_tendon_ids), - fixed_tendon_ids=fixed_tendon_ids, - env_ids=env_ids, + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + fixed_tendon_mask = self._resolve_mask(fixed_tendon_mask, self._ALL_FIXED_TENDON_MASK) + self.assert_shape_and_dtype_mask(stiffness, (env_mask, fixed_tendon_mask), wp.float32, "stiffness") + wp.launch( + articulation_kernels.float_data_to_buffer_with_mask + if isinstance(stiffness, float) + else shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(self.num_instances, self.num_fixed_tendons), + inputs=[stiffness, env_mask, fixed_tendon_mask], + outputs=[self.data._fixed_tendon_stiffness], + device=self.device, ) def set_fixed_tendon_damping_index( @@ -2813,12 +2815,17 @@ def set_fixed_tendon_damping_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask) - self.set_fixed_tendon_damping_index( - damping=self._select_full_data(damping, env_ids, fixed_tendon_ids), - fixed_tendon_ids=fixed_tendon_ids, - env_ids=env_ids, + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + fixed_tendon_mask = self._resolve_mask(fixed_tendon_mask, self._ALL_FIXED_TENDON_MASK) + self.assert_shape_and_dtype_mask(damping, (env_mask, fixed_tendon_mask), wp.float32, "damping") + wp.launch( + articulation_kernels.float_data_to_buffer_with_mask + if isinstance(damping, float) + else shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(self.num_instances, self.num_fixed_tendons), + inputs=[damping, env_mask, fixed_tendon_mask], + outputs=[self.data._fixed_tendon_damping], + device=self.device, ) def set_fixed_tendon_limit_stiffness_index( @@ -2930,12 +2937,17 @@ def set_fixed_tendon_position_limit_mask( fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - env_ids = self._resolve_env_mask(env_mask) - fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask) - self.set_fixed_tendon_position_limit_index( - limit=self._select_full_data(limit, env_ids, fixed_tendon_ids), - fixed_tendon_ids=fixed_tendon_ids, - env_ids=env_ids, + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + fixed_tendon_mask = self._resolve_mask(fixed_tendon_mask, self._ALL_FIXED_TENDON_MASK) + self.assert_shape_and_dtype_mask(limit, (env_mask, fixed_tendon_mask), wp.vec2f, "limit") + if isinstance(limit, float): + raise ValueError("Fixed tendon position limits must be a tensor or array, not a float.") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask_kernel(wp.vec2f), + dim=(self.num_instances, self.num_fixed_tendons), + inputs=[limit, env_mask, fixed_tendon_mask], + outputs=[self.data._fixed_tendon_pos_limits], + device=self.device, ) def set_fixed_tendon_rest_length_index( @@ -3136,10 +3148,27 @@ def write_fixed_tendon_properties_to_sim_mask( fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are updated. env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - self.write_fixed_tendon_properties_to_sim_index( - fixed_tendon_ids=self._resolve_fixed_tendon_mask(fixed_tendon_mask), - env_ids=self._resolve_env_mask(env_mask), + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + fixed_tendon_mask = self._resolve_mask(fixed_tendon_mask, self._ALL_FIXED_TENDON_MASK) + wp.launch( + articulation_kernels.copy_fixed_tendon_properties_with_mask, + dim=(self.num_instances, self.num_fixed_tendons), + inputs=[ + env_mask, + fixed_tendon_mask, + self.data._fixed_tendon_stiffness, + self.data._fixed_tendon_damping, + self.data._fixed_tendon_pos_limits, + ], + outputs=[ + self.data._sim_bind_fixed_tendon_stiffness, + self.data._sim_bind_fixed_tendon_damping, + self.data._sim_bind_fixed_tendon_pos_limits, + ], + device=self.device, ) + # the solver keeps its own copy of the tendon properties and only re-reads them when notified + self._backend_manager.add_model_change(ModelFlags.TENDON_PROPERTIES) def set_spatial_tendon_stiffness_index( self, @@ -3445,6 +3474,7 @@ def _clear_callbacks(self) -> None: def _create_buffers(self): self._ALL_INDICES = wp.array(np.arange(self.num_instances, dtype=np.int32), device=self.device) self._ALL_ENV_MASK = wp.ones((self.num_instances,), dtype=wp.bool, device=self.device) + self._scratch_env_mask = wp.zeros((self.num_instances,), dtype=wp.bool, device=self.device) self._ALL_JOINT_INDICES = wp.array(np.arange(self.num_joints, dtype=np.int32), device=self.device) self._ALL_JOINT_MASK = wp.ones((self.num_joints,), dtype=wp.bool, device=self.device) self._ALL_BODY_INDICES = wp.array(np.arange(self.num_bodies, dtype=np.int32), device=self.device) @@ -3720,15 +3750,7 @@ def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | No Returns: Environment indices. """ - if isinstance(env_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if (env_ids is None) or (env_ids == slice(None)): - return self._ALL_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self.device) - return env_ids + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device) def _resolve_joint_ids(self, joint_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: """Resolve joint indices to a warp array or tensor. @@ -3739,13 +3761,7 @@ def _resolve_joint_ids(self, joint_ids: Sequence[int] | torch.Tensor | wp.array Returns: A warp array of joint indices or a tensor of joint indices. """ - if isinstance(joint_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(joint_ids, list): - return wp.array(joint_ids, dtype=wp.int32, device=self.device) - if (joint_ids is None) or (joint_ids == slice(None)): - return self._ALL_JOINT_INDICES - return joint_ids + return resolve_ids(ids=joint_ids, all_ids=self._ALL_JOINT_INDICES, device=self.device) def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: """Resolve body indices to a warp array or tensor. @@ -3756,13 +3772,7 @@ def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | Returns: A warp array of body indices or a tensor of body indices. """ - if isinstance(body_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(body_ids, list): - return wp.array(body_ids, dtype=wp.int32, device=self.device) - if (body_ids is None) or (body_ids == slice(None)): - return self._ALL_BODY_INDICES - return body_ids + return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self.device) def _resolve_fixed_tendon_ids( self, tendon_ids: Sequence[int] | torch.Tensor | wp.array | None @@ -3775,13 +3785,7 @@ def _resolve_fixed_tendon_ids( Returns: A warp array of tendon indices or a tensor of tendon indices. """ - if isinstance(tendon_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(tendon_ids, list): - return wp.array(tendon_ids, dtype=wp.int32, device=self.device) - if (tendon_ids is None) or (tendon_ids == slice(None)): - return self._ALL_FIXED_TENDON_INDICES - return tendon_ids + return resolve_ids(ids=tendon_ids, all_ids=self._ALL_FIXED_TENDON_INDICES, device=self.device) def _resolve_spatial_tendon_ids( self, spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None @@ -3794,27 +3798,7 @@ def _resolve_spatial_tendon_ids( Returns: A warp array of spatial tendon indices or a tensor of spatial tendon indices. """ - if isinstance(spatial_tendon_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(spatial_tendon_ids, list): - return wp.array(spatial_tendon_ids, dtype=wp.int32, device=self.device) - if (spatial_tendon_ids is None) or (spatial_tendon_ids == slice(None)): - return self._ALL_SPATIAL_TENDON_INDICES - return spatial_tendon_ids - - def _resolve_env_mask(self, env_mask: wp.array | torch.Tensor | None) -> wp.array | torch.Tensor: - """Resolve an environment mask to environment indices.""" - return self._ALL_INDICES if env_mask is None else self._mask_to_ids(env_mask) - - def _resolve_fixed_tendon_mask(self, fixed_tendon_mask: wp.array | torch.Tensor | None) -> wp.array | torch.Tensor: - """Resolve a fixed tendon mask to fixed tendon indices.""" - return self._ALL_FIXED_TENDON_INDICES if fixed_tendon_mask is None else self._mask_to_ids(fixed_tendon_mask) - - @staticmethod - def _mask_to_ids(mask: wp.array | torch.Tensor) -> torch.Tensor: - """Convert a boolean mask to int32 indices.""" - mask = wp.to_torch(mask) if isinstance(mask, wp.array) else mask - return torch.nonzero(mask)[:, 0].to(torch.int32) + return resolve_ids(ids=spatial_tendon_ids, all_ids=self._ALL_SPATIAL_TENDON_INDICES, device=self.device) @staticmethod def _to_torch_ids(ids: wp.array | torch.Tensor) -> torch.Tensor: @@ -3839,12 +3823,7 @@ def _resolve_mask(self, mask: wp.array | torch.Tensor | None, full_mask: wp.arra Returns: A warp array of mask. """ - if mask is None: - return full_mask - - if isinstance(mask, torch.Tensor): - return wp.from_torch(mask, dtype=wp.bool) - return mask + return resolve_1d_mask(mask=mask, all_mask=full_mask, device=self.device) """ Deprecated methods. diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py index 60544e3a4100..7ad2bd59f057 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py @@ -553,6 +553,36 @@ def float_data_to_buffer_with_indices_kernel( return _FLOAT_DATA_TO_BUFFER_WITH_INDICES_DISPATCHER.select(env_ids, joint_ids) +@wp.kernel +def copy_fixed_tendon_properties_with_mask( + env_mask: wp.array(dtype=wp.bool), + tendon_mask: wp.array(dtype=wp.bool), + stiffness: wp.array2d(dtype=wp.float32), + damping: wp.array2d(dtype=wp.float32), + pos_limits: wp.array2d(dtype=wp.vec2f), + sim_stiffness: wp.array2d(dtype=wp.float32), + sim_damping: wp.array2d(dtype=wp.float32), + sim_pos_limits: wp.array2d(dtype=wp.vec2f), +): + """Copy the staged fixed-tendon stiffness, damping and position limits to the simulation where both masks are true. + + Args: + env_mask: Environment mask. Shape is (num_instances,). + tendon_mask: Fixed-tendon mask. Shape is (num_fixed_tendons,). + stiffness: Staged stiffness. Shape is (num_instances, num_fixed_tendons). + damping: Staged damping. Shape is (num_instances, num_fixed_tendons). + pos_limits: Staged position limits. Shape is (num_instances, num_fixed_tendons). + sim_stiffness: Simulation stiffness. Shape is (num_instances, num_fixed_tendons). + sim_damping: Simulation damping. Shape is (num_instances, num_fixed_tendons). + sim_pos_limits: Simulation position limits. Shape is (num_instances, num_fixed_tendons). + """ + i, j = wp.tid() + if env_mask[i] and tendon_mask[j]: + sim_stiffness[i, j] = stiffness[i, j] + sim_damping[i, j] = damping[i, j] + sim_pos_limits[i, j] = pos_limits[i, j] + + @wp.kernel def float_data_to_buffer_with_mask( in_data: wp.float32, diff --git a/source/isaaclab_newton/isaaclab_newton/assets/cable_object/cable_object.py b/source/isaaclab_newton/isaaclab_newton/assets/cable_object/cable_object.py index 86bd5d8590cb..d98cef9aba62 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/cable_object/cable_object.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/cable_object/cable_object.py @@ -18,6 +18,7 @@ from isaaclab.assets.cable_object.base_cable_object import BaseCableObject from isaaclab.sim.utils.queries import has_deformable_curve_api, path_expr_to_glob, resolve_matching_prims_from_source from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_ids from .cable_object_data import CableObjectData from .kernels import ( @@ -207,15 +208,7 @@ def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array(dtyp dtype=wp.int32 ): """Resolve environment indices to a Warp array.""" - if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)): - return self._ALL_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - if isinstance(env_ids, torch.Tensor): - return wp.from_torch(env_ids.to(device=self.device, dtype=torch.int32).contiguous(), dtype=wp.int32) - if isinstance(env_ids, Sequence): - return wp.array(list(env_ids), dtype=wp.int32, device=self.device) - return env_ids + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device, dtype=wp.int32) def _write_segment_state_to_sim( self, diff --git a/source/isaaclab_newton/isaaclab_newton/assets/deformable_object/deformable_object.py b/source/isaaclab_newton/isaaclab_newton/assets/deformable_object/deformable_object.py index 06e96694147e..d2016b4186ed 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/deformable_object/deformable_object.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/deformable_object/deformable_object.py @@ -17,6 +17,7 @@ from isaaclab.assets.deformable_object.base_deformable_object import BaseDeformableObject from isaaclab.markers import VisualizationMarkers from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids from .deformable_object_data import DeformableObjectData from .kernels import ( @@ -402,23 +403,11 @@ def write_nodal_kinematic_target_to_sim_mask( def _resolve_env_ids(self, env_ids): """Resolve environment indices to a warp int32 array.""" - if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)): - return self._ALL_INDICES - elif isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self.device) - elif isinstance(env_ids, torch.Tensor): - return wp.from_torch(env_ids.to(torch.int32), dtype=wp.int32) - return env_ids + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device, dtype=wp.int32) def _resolve_mask(self, mask: wp.array | torch.Tensor | None, full_mask: wp.array) -> wp.array: """Resolve an environment mask to a warp bool array.""" - if mask is None: - return full_mask - if isinstance(mask, torch.Tensor): - if mask.dtype != torch.bool: - mask = mask.to(torch.bool) - return wp.from_torch(mask, dtype=wp.bool) - return mask + return resolve_1d_mask(mask=mask, all_mask=full_mask, device=self.device) def _iter_particle_states(self): """Yield active Newton states.""" diff --git a/source/isaaclab_newton/isaaclab_newton/assets/kernels.py b/source/isaaclab_newton/isaaclab_newton/assets/kernels.py index eec3479f1fdc..c9f43ef3adec 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/kernels.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/kernels.py @@ -5,6 +5,7 @@ from __future__ import annotations +import functools from typing import TYPE_CHECKING, Any import torch @@ -960,6 +961,89 @@ def set_body_link_velocity_to_sim( body_acc_w[env_id, body_id] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) +@wp.kernel +def set_body_com_pose_to_sim_mask( + data: wp.array2d(dtype=wp.transformf), + body_com_pos_b: wp.array2d(dtype=wp.vec3f), + env_mask: wp.array(dtype=wp.bool), + body_mask: wp.array(dtype=wp.bool), + body_com_pose_w: wp.array2d(dtype=wp.transformf), + body_link_pose_w: wp.array2d(dtype=wp.transformf), +): + """Write body COM pose data, and the link poses derived from it, for the selected environments and bodies. + + Args: + data: Input array of body COM poses. Shape is (num_envs, num_bodies). + body_com_pos_b: Input array of body COM positions in body frame. Shape is (num_envs, num_bodies). + env_mask: Input array of environment mask. Shape is (num_envs,). + body_mask: Input array of body mask. Shape is (num_bodies,). + body_com_pose_w: Output array where body COM poses are written. Shape is (num_envs, num_bodies). + body_link_pose_w: Output array where body link poses (derived from COM) are written. + Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + if env_mask[i] and body_mask[j]: + body_com_pose_w[i, j] = data[i, j] + body_link_pose_w[i, j] = get_com_pose_in_link_frame_func(body_com_pose_w[i, j], body_com_pos_b[i, j]) + + +@wp.kernel +def set_body_com_velocity_to_sim_mask( + data: wp.array2d(dtype=wp.spatial_vectorf), + env_mask: wp.array(dtype=wp.bool), + body_mask: wp.array(dtype=wp.bool), + body_com_velocity_w: wp.array2d(dtype=wp.spatial_vectorf), + body_acc_w: wp.array2d(dtype=wp.spatial_vectorf), +): + """Write body COM velocity data, and zero the body accelerations, for the selected environments and bodies. + + Args: + data: Input array of body COM spatial velocities. Shape is (num_envs, num_bodies). + env_mask: Input array of environment mask. Shape is (num_envs,). + body_mask: Input array of body mask. Shape is (num_bodies,). + body_com_velocity_w: Output array where body COM velocities are written. Shape is (num_envs, num_bodies). + body_acc_w: Output array where body accelerations are zeroed. Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + if env_mask[i] and body_mask[j]: + body_com_velocity_w[i, j] = data[i, j] + body_acc_w[i, j] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + + +@wp.kernel +def set_body_link_velocity_to_sim_mask( + data: wp.array2d(dtype=wp.spatial_vectorf), + body_com_pos_b: wp.array2d(dtype=wp.vec3f), + body_link_pose_w: wp.array2d(dtype=wp.transformf), + env_mask: wp.array(dtype=wp.bool), + body_mask: wp.array(dtype=wp.bool), + body_link_velocity_w: wp.array2d(dtype=wp.spatial_vectorf), + body_com_velocity_w: wp.array2d(dtype=wp.spatial_vectorf), + body_acc_w: wp.array2d(dtype=wp.spatial_vectorf), +): + """Write body link velocity data, the COM velocities derived from it, and zero the body accelerations, for the + selected environments and bodies. + + Args: + data: Input array of body link spatial velocities. Shape is (num_envs, num_bodies). + body_com_pos_b: Input array of body COM positions in body frame. Shape is (num_envs, num_bodies). + body_link_pose_w: Input array of body link poses in world frame. Shape is (num_envs, num_bodies). + env_mask: Input array of environment mask. Shape is (num_envs,). + body_mask: Input array of body mask. Shape is (num_bodies,). + body_link_velocity_w: Output array where body link velocities are written. Shape is (num_envs, num_bodies). + body_com_velocity_w: Output array where body COM velocities (derived from link) are written. + Shape is (num_envs, num_bodies). + body_acc_w: Output array where body accelerations are zeroed. Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + if env_mask[i] and body_mask[j]: + body_link_velocity_w[i, j] = data[i, j] + body_com_velocity_w[i, j] = get_link_velocity_in_com_frame_func( + body_link_velocity_w[i, j], body_link_pose_w[i, j], body_com_pos_b[i, j] + ) + body_acc_w[i, j] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + + _SET_BODY_LINK_POSE_TO_SIM_DISPATCHER = IndexKernelDispatcher(set_body_link_pose_to_sim, ("env_ids", "body_ids")) _SET_BODY_COM_POSE_TO_SIM_DISPATCHER = IndexKernelDispatcher(set_body_com_pose_to_sim, ("env_ids", "body_ids")) _SET_BODY_COM_VELOCITY_TO_SIM_DISPATCHER = IndexKernelDispatcher(set_body_com_velocity_to_sim, ("env_ids", "body_ids")) @@ -1034,21 +1118,18 @@ def write_2d_data_to_buffer_with_indices_kernel( @wp.kernel -def write_2d_data_to_buffer_with_mask( - in_data: wp.array2d(dtype=wp.float32), +def _write_2d_data_to_buffer_with_mask( + in_data: wp.array2d(dtype=Any), env_mask: wp.array(dtype=wp.bool), joint_mask: wp.array(dtype=wp.bool), - out_data: wp.array2d(dtype=wp.float32), + out_data: wp.array2d(dtype=Any), ): - """Write 2D float data to a buffer at specified indices. - - This kernel copies float data from an input array to an output buffer at the specified - environment and joint/body indices. + """Write 2D data to a buffer where both masks are true. Args: - in_data: Input array containing float data. Shape is (num_instances, num_joints). + in_data: Input array. Shape is (num_instances, num_joints). env_mask: Input array of environment mask. Shape is (num_instances,). - joint_mask: Input array of joint/body mask. Shape is (num_instances, num_joints). + joint_mask: Input array of joint/body/tendon mask. Shape is (num_joints,). out_data: Output array where data is written. Shape is (num_instances, num_joints). """ i, j = wp.tid() @@ -1056,6 +1137,20 @@ def write_2d_data_to_buffer_with_mask( out_data[i, j] = in_data[i, j] +@functools.cache +def write_2d_data_to_buffer_with_mask_kernel(dtype: type) -> wp.Kernel: + """Select the masked 2-D buffer writer for the buffer's ``dtype``. + + Launching the generic kernel directly resolves its types on every launch (about twice the host cost), and + cannot take torch input. + """ + array_type = wp.array2d(dtype=dtype) + return wp.overload(_write_2d_data_to_buffer_with_mask, {"in_data": array_type, "out_data": array_type}) + + +write_2d_data_to_buffer_with_mask = write_2d_data_to_buffer_with_mask_kernel(wp.float32) + + @wp.func def _update_body_inertial_properties_inverse( env_id: int, diff --git a/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py index 3d1329cca985..8a725ca9f1ea 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py @@ -15,6 +15,7 @@ from isaaclab.assets.deformable_object.base_deformable_object import BaseDeformableObject from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids from isaaclab_newton.sim.spawners.mpm.mpm import _SIMULATION_POINTS_SUFFIX @@ -283,24 +284,10 @@ def _create_buffers(self): self._data.default_particle_state_w = self._data.default_nodal_state_w def _resolve_env_ids(self, env_ids): - if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)): - return self._ALL_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - if isinstance(env_ids, torch.Tensor): - return wp.from_torch(env_ids.to(device=self.device, dtype=torch.int32), dtype=wp.int32) - if isinstance(env_ids, Sequence): - return wp.array(list(env_ids), dtype=wp.int32, device=self.device) - return env_ids + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device, dtype=wp.int32) def _resolve_mask(self, mask: wp.array | torch.Tensor | None) -> wp.array: - if mask is None: - return self._ALL_ENV_MASK - if isinstance(mask, torch.Tensor): - if mask.dtype != torch.bool: - mask = mask.to(torch.bool) - return wp.from_torch(mask.to(device=self.device).contiguous(), dtype=wp.bool) - return mask + return resolve_1d_mask(mask=mask, all_mask=self._ALL_ENV_MASK, device=self.device) def _iter_particle_states(self): """Yield the Newton states whose particle arrays must receive writes.""" diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py index ae04842b9967..fc7e0c76a7c8 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py @@ -21,6 +21,7 @@ from isaaclab.assets.rigid_object.base_rigid_object import BaseRigidObject from isaaclab.sim.utils.queries import path_expr_to_glob, resolve_matching_prims_from_source from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_newton.assets import kernels as shared_kernels @@ -1122,13 +1123,7 @@ def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | No Returns: A warp array of environment indices or a tensor of environment indices. """ - if (env_ids is None) or (env_ids == slice(None)): - return self._ALL_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self.device) - return env_ids + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device) def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: """Resolve body indices to a warp array or tensor. @@ -1139,13 +1134,7 @@ def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | Returns: A warp array of body indices or a tensor of body indices. """ - if isinstance(body_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(body_ids, list): - return wp.array(body_ids, dtype=wp.int32, device=self.device) - if (body_ids is None) or (body_ids == slice(None)): - return self._ALL_BODY_INDICES - return body_ids + return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self.device) """ Internal simulation callbacks. diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py index 268e7a3e74d3..01ad44ba49d7 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py @@ -22,6 +22,7 @@ from isaaclab.assets.rigid_object_collection.base_rigid_object_collection import BaseRigidObjectCollection from isaaclab.utils import clone, validate from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_newton.assets import kernels as shared_kernels @@ -308,16 +309,8 @@ def write_body_pose_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - else: - body_ids = self._ALL_BODY_INDICES - self.write_body_link_pose_to_sim_index( - body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward + self.write_body_link_pose_to_sim_mask( + body_poses=body_poses, body_mask=body_mask, env_mask=env_mask, skip_forward=skip_forward ) def write_body_velocity_to_sim_index( @@ -386,20 +379,8 @@ def write_body_velocity_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - else: - body_ids = self._ALL_BODY_INDICES - self.write_body_com_velocity_to_sim_index( - body_velocities=body_velocities, - env_ids=env_ids, - body_ids=body_ids, - full_data=True, - skip_forward=skip_forward, + self.write_body_com_velocity_to_sim_mask( + body_velocities=body_velocities, body_mask=body_mask, env_mask=env_mask, skip_forward=skip_forward ) def write_body_link_pose_to_sim_index( @@ -492,16 +473,25 @@ def write_body_link_pose_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_mask = self._ALL_ENV_MASK - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - self.write_body_link_pose_to_sim_index( - body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward + env_mask = resolve_1d_mask(mask=env_mask, all_mask=self._ALL_ENV_MASK, device=self.device) + # a body selection given as ids becomes a body mask on the device, so the same mask kernel applies + body_mask = resolve_1d_mask( + ids=body_ids, + mask=body_mask, + all_mask=self._ALL_BODY_MASK, + scratch_mask=self._scratch_body_mask, + device=self.device, + ) + self.assert_shape_and_dtype_mask(body_poses, (env_mask, body_mask), wp.transformf, "body_poses") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask_kernel(wp.transformf), + dim=(self.num_instances, self.num_bodies), + inputs=[body_poses, env_mask, body_mask], + outputs=[self.data.body_link_pose_w], + device=self.device, ) + if not skip_forward: + self.data._reset_pose(env_mask=env_mask) def write_body_com_pose_to_sim_index( self, @@ -597,17 +587,25 @@ def write_body_com_pose_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_mask = self._ALL_ENV_MASK - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - # The index writer owns the invalidation (with from_link=False); forward skip_forward to it. - self.write_body_com_pose_to_sim_index( - body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward + env_mask = resolve_1d_mask(mask=env_mask, all_mask=self._ALL_ENV_MASK, device=self.device) + # a body selection given as ids becomes a body mask on the device, so the same mask kernel applies + body_mask = resolve_1d_mask( + ids=body_ids, + mask=body_mask, + all_mask=self._ALL_BODY_MASK, + scratch_mask=self._scratch_body_mask, + device=self.device, + ) + self.assert_shape_and_dtype_mask(body_poses, (env_mask, body_mask), wp.transformf, "body_poses") + wp.launch( + shared_kernels.set_body_com_pose_to_sim_mask, + dim=(self.num_instances, self.num_bodies), + inputs=[body_poses, self.data.body_com_pos_b, env_mask, body_mask], + outputs=[self.data.body_com_pose_w, self.data.body_link_pose_w], + device=self.device, ) + if not skip_forward: + self.data._reset_pose(env_mask=env_mask, from_link=False) def write_body_com_velocity_to_sim_index( self, @@ -711,21 +709,25 @@ def write_body_com_velocity_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_mask = self._ALL_ENV_MASK - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - # The index writer owns the invalidation; forward skip_forward to it. - self.write_body_com_velocity_to_sim_index( - body_velocities=body_velocities, - env_ids=env_ids, - body_ids=body_ids, - full_data=True, - skip_forward=skip_forward, + env_mask = resolve_1d_mask(mask=env_mask, all_mask=self._ALL_ENV_MASK, device=self.device) + # a body selection given as ids becomes a body mask on the device, so the same mask kernel applies + body_mask = resolve_1d_mask( + ids=body_ids, + mask=body_mask, + all_mask=self._ALL_BODY_MASK, + scratch_mask=self._scratch_body_mask, + device=self.device, ) + self.assert_shape_and_dtype_mask(body_velocities, (env_mask, body_mask), wp.spatial_vectorf, "body_velocities") + wp.launch( + shared_kernels.set_body_com_velocity_to_sim_mask, + dim=(self.num_instances, self.num_bodies), + inputs=[body_velocities, env_mask, body_mask], + outputs=[self.data.body_com_vel_w, self.data.body_com_acc_w], + device=self.device, + ) + if not skip_forward: + self.data._reset_velocity(env_mask=env_mask) def write_body_link_velocity_to_sim_index( self, @@ -831,21 +833,25 @@ def write_body_link_velocity_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_mask = self._ALL_ENV_MASK - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - # The index writer owns the invalidation (with from_com=False); forward skip_forward to it. - self.write_body_link_velocity_to_sim_index( - body_velocities=body_velocities, - env_ids=env_ids, - body_ids=body_ids, - full_data=True, - skip_forward=skip_forward, + env_mask = resolve_1d_mask(mask=env_mask, all_mask=self._ALL_ENV_MASK, device=self.device) + # a body selection given as ids becomes a body mask on the device, so the same mask kernel applies + body_mask = resolve_1d_mask( + ids=body_ids, + mask=body_mask, + all_mask=self._ALL_BODY_MASK, + scratch_mask=self._scratch_body_mask, + device=self.device, + ) + self.assert_shape_and_dtype_mask(body_velocities, (env_mask, body_mask), wp.spatial_vectorf, "body_velocities") + wp.launch( + shared_kernels.set_body_link_velocity_to_sim_mask, + dim=(self.num_instances, self.num_bodies), + inputs=[body_velocities, self.data.body_com_pos_b, self.data.body_link_pose_w, env_mask, body_mask], + outputs=[self.data.body_link_vel_w, self.data.body_com_vel_w, self.data.body_com_acc_w], + device=self.device, ) + if not skip_forward: + self.data._reset_velocity(env_mask=env_mask, from_com=False) """ Operations - Setters. @@ -1244,6 +1250,7 @@ def _create_buffers(self): ) self._ALL_ENV_MASK = wp.ones((self.num_instances,), dtype=wp.bool, device=self.device) self._ALL_BODY_MASK = wp.ones((self.num_bodies,), dtype=wp.bool, device=self.device) + self._scratch_body_mask = wp.zeros((self.num_bodies,), dtype=wp.bool, device=self.device) # external wrench composer self._instantaneous_wrench_composer = WrenchComposer(self) @@ -1285,13 +1292,7 @@ def _resolve_env_ids(self, env_ids) -> wp.array | torch.Tensor: Returns: Environment indices. """ - if (env_ids is None) or (env_ids == slice(None)): - return self._ALL_ENV_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_ENV_INDICES)[env_ids]) - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self.device) - return env_ids + return resolve_ids(ids=env_ids, all_ids=self._ALL_ENV_INDICES, device=self.device) def _resolve_body_ids(self, body_ids) -> wp.array | torch.Tensor: """Resolve body indices. @@ -1302,37 +1303,7 @@ def _resolve_body_ids(self, body_ids) -> wp.array | torch.Tensor: Returns: Body indices. """ - if isinstance(body_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(body_ids, list): - return wp.array(body_ids, dtype=wp.int32, device=self.device) - if (body_ids is None) or (body_ids == slice(None)): - return self._ALL_BODY_INDICES - if isinstance(body_ids, slice): - return wp.from_torch( - torch.arange(self.num_bodies, dtype=torch.int32, device=self.device)[body_ids], dtype=wp.int32 - ) - return body_ids - - def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array | torch.Tensor: - """Resolve environment mask to indices via torch.nonzero.""" - if env_mask is not None: - if isinstance(env_mask, wp.array): - env_mask = wp.to_torch(env_mask) - env_ids = torch.nonzero(env_mask)[:, 0] - else: - env_ids = self._ALL_ENV_INDICES - return env_ids - - def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array | torch.Tensor: - """Resolve body mask to indices via torch.nonzero.""" - if body_mask is not None: - if isinstance(body_mask, wp.array): - body_mask = wp.to_torch(body_mask) - body_ids = torch.nonzero(body_mask)[:, 0] - else: - body_ids = self._ALL_BODY_INDICES - return body_ids + return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self.device) """ Internal simulation callbacks. From 7b48c38e3c2417e6e458b46a82d7ef425e047065 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Mon, 26 Oct 2026 02:48:21 -0700 Subject: [PATCH 5/6] Resolve OVPhysX selections through the shared helpers The OVPhysX assets resolve env, body, joint and tendon selections through resolve_ids and resolve_1d_mask instead of their own copies. --- .../changelog.d/jichuanh-id-mask.changed.rst | 7 ++ .../assets/articulation/articulation.py | 103 ++--------------- .../deformable_object/deformable_object.py | 18 +-- .../assets/rigid_object/rigid_object.py | 37 +----- .../rigid_object_collection.py | 105 +++--------------- .../test/assets/test_deformable_object.py | 22 ++++ 6 files changed, 65 insertions(+), 227 deletions(-) create mode 100644 source/isaaclab_ov/changelog.d/jichuanh-id-mask.changed.rst diff --git a/source/isaaclab_ov/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab_ov/changelog.d/jichuanh-id-mask.changed.rst new file mode 100644 index 000000000000..3f1ff2245d66 --- /dev/null +++ b/source/isaaclab_ov/changelog.d/jichuanh-id-mask.changed.rst @@ -0,0 +1,7 @@ +* Changed the env selection helpers of the OVPhysX assets to resolve ``env_ids`` through + :func:`~isaaclab.utils.warp.utils.resolve_ids` and ``env_mask`` through + :func:`~isaaclab.utils.warp.utils.resolve_1d_mask`. Any integer sequence and slices with a negative step are accepted, + and a :class:`~isaaclab.utils.warp.ProxyArray` is rejected with a :class:`TypeError`. +* Changed the body, joint and tendon selection helpers of the OVPhysX assets to resolve ids and masks through + :func:`~isaaclab.utils.warp.utils.resolve_ids` and :func:`~isaaclab.utils.warp.utils.resolve_1d_mask`, as the env + selection does, so slices other than ``slice(None)`` are accepted. diff --git a/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py b/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py index 07bc71d3f611..2162ae04b8ba 100644 --- a/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py +++ b/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py @@ -33,6 +33,7 @@ from isaaclab.utils.string import resolve_matching_names from isaaclab.utils.warp import ProxyArray from isaaclab.utils.warp import kernels as warp_kernels +from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_ov import tensor_types as TT @@ -4373,79 +4374,27 @@ def _log_articulation_info(self) -> None: def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: """Resolve environment indices on ``self._device``.""" - if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)): - return self._ALL_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - if isinstance(env_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self._device) - if isinstance(env_ids, torch.Tensor): - return env_ids.to(device=self._device) - if isinstance(env_ids, wp.array) and str(env_ids.device) != self._device: - env_ids = wp.clone(env_ids, device=self._device) - return env_ids + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self._device) def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: """Resolve body indices to a Warp signed-integer array on ``self._device``.""" - if isinstance(body_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if body_ids is None or body_ids == slice(None): - return self._ALL_BODY_INDICES - if isinstance(body_ids, list): - return wp.array(body_ids, dtype=wp.int32, device=self._device) - if isinstance(body_ids, torch.Tensor): - return body_ids.to(device=self._device) - if isinstance(body_ids, wp.array) and str(body_ids.device) != self._device: - body_ids = wp.clone(body_ids, device=self._device) - return body_ids + return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self._device) def _resolve_joint_ids(self, joint_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: """Resolve joint indices to a Warp signed-integer array on ``self._device``.""" - if isinstance(joint_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if joint_ids is None or joint_ids == slice(None): - return self._ALL_JOINT_INDICES - if isinstance(joint_ids, list): - return wp.array(joint_ids, dtype=wp.int32, device=self._device) - if isinstance(joint_ids, torch.Tensor): - return joint_ids.to(device=self._device) - if isinstance(joint_ids, wp.array) and str(joint_ids.device) != self._device: - joint_ids = wp.clone(joint_ids, device=self._device) - return joint_ids + return resolve_ids(ids=joint_ids, all_ids=self._ALL_JOINT_INDICES, device=self._device) def _resolve_fixed_tendon_ids( self, tendon_ids: Sequence[int] | torch.Tensor | wp.array | None ) -> wp.array | torch.Tensor: """Resolve fixed-tendon indices to a Warp signed-integer array on ``self._device``.""" - if isinstance(tendon_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if tendon_ids is None or tendon_ids == slice(None): - return self._ALL_FIXED_TENDON_INDICES - if isinstance(tendon_ids, list): - return wp.array(tendon_ids, dtype=wp.int32, device=self._device) - if isinstance(tendon_ids, torch.Tensor): - return tendon_ids.to(device=self._device) - if isinstance(tendon_ids, wp.array) and str(tendon_ids.device) != self._device: - tendon_ids = wp.clone(tendon_ids, device=self._device) - return tendon_ids + return resolve_ids(ids=tendon_ids, all_ids=self._ALL_FIXED_TENDON_INDICES, device=self._device) def _resolve_spatial_tendon_ids( self, tendon_ids: Sequence[int] | torch.Tensor | wp.array | None ) -> wp.array | torch.Tensor: """Resolve spatial-tendon indices to a Warp signed-integer array on ``self._device``.""" - if isinstance(tendon_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if tendon_ids is None or tendon_ids == slice(None): - return self._ALL_SPATIAL_TENDON_INDICES - if isinstance(tendon_ids, list): - return wp.array(tendon_ids, dtype=wp.int32, device=self._device) - if isinstance(tendon_ids, torch.Tensor): - return tendon_ids.to(device=self._device) - if isinstance(tendon_ids, wp.array) and str(tendon_ids.device) != self._device: - tendon_ids = wp.clone(tendon_ids, device=self._device) - return tendon_ids + return resolve_ids(ids=tendon_ids, all_ids=self._ALL_SPATIAL_TENDON_INDICES, device=self._device) def _broadcast_scalar_to_2d( self, value: float | torch.Tensor | wp.array, shape: tuple[int, int] @@ -4483,13 +4432,7 @@ def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array: no ``torch.nonzero`` conversion is needed. ``None`` returns the pre-allocated all-true mask. """ - if env_mask is None: - return self._ALL_TRUE_ENV_MASK - if isinstance(env_mask, torch.Tensor): - return wp.from_torch(env_mask.to(torch.bool), dtype=wp.bool) - if isinstance(env_mask, wp.array) and str(env_mask.device) != self._device: - env_mask = wp.clone(env_mask, device=self._device) - return env_mask + return resolve_1d_mask(mask=env_mask, all_mask=self._ALL_TRUE_ENV_MASK, device=self._device) def _sim_env_ids_view(self, count: int) -> wp.array: """Return a cached prefix view of the simulator-index scratch buffer.""" @@ -4505,43 +4448,19 @@ def _sim_env_ids_view(self, count: int) -> wp.array: def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array: """Resolve a body mask to a ``wp.bool`` array on ``self._device`` (Newton-style).""" - if body_mask is None: - return self._ALL_TRUE_BODY_MASK - if isinstance(body_mask, torch.Tensor): - return wp.from_torch(body_mask.to(torch.bool), dtype=wp.bool) - if isinstance(body_mask, wp.array) and str(body_mask.device) != self._device: - body_mask = wp.clone(body_mask, device=self._device) - return body_mask + return resolve_1d_mask(mask=body_mask, all_mask=self._ALL_TRUE_BODY_MASK, device=self._device) def _resolve_joint_mask(self, joint_mask: wp.array | None) -> wp.array: """Resolve a joint mask to a ``wp.bool`` array on ``self._device``.""" - if joint_mask is None: - return self._ALL_TRUE_JOINT_MASK - if isinstance(joint_mask, torch.Tensor): - return wp.from_torch(joint_mask.to(torch.bool), dtype=wp.bool) - if isinstance(joint_mask, wp.array) and str(joint_mask.device) != self._device: - joint_mask = wp.clone(joint_mask, device=self._device) - return joint_mask + return resolve_1d_mask(mask=joint_mask, all_mask=self._ALL_TRUE_JOINT_MASK, device=self._device) def _resolve_fixed_tendon_mask(self, tendon_mask: wp.array | None) -> wp.array: """Resolve a fixed-tendon mask to a ``wp.bool`` array on ``self._device``.""" - if tendon_mask is None: - return self._ALL_TRUE_FIXED_TENDON_MASK - if isinstance(tendon_mask, torch.Tensor): - return wp.from_torch(tendon_mask.to(torch.bool), dtype=wp.bool) - if isinstance(tendon_mask, wp.array) and str(tendon_mask.device) != self._device: - tendon_mask = wp.clone(tendon_mask, device=self._device) - return tendon_mask + return resolve_1d_mask(mask=tendon_mask, all_mask=self._ALL_TRUE_FIXED_TENDON_MASK, device=self._device) def _resolve_spatial_tendon_mask(self, tendon_mask: wp.array | None) -> wp.array: """Resolve a spatial-tendon mask to a ``wp.bool`` array on ``self._device``.""" - if tendon_mask is None: - return self._ALL_TRUE_SPATIAL_TENDON_MASK - if isinstance(tendon_mask, torch.Tensor): - return wp.from_torch(tendon_mask.to(torch.bool), dtype=wp.bool) - if isinstance(tendon_mask, wp.array) and str(tendon_mask.device) != self._device: - tendon_mask = wp.clone(tendon_mask, device=self._device) - return tendon_mask + return resolve_1d_mask(mask=tendon_mask, all_mask=self._ALL_TRUE_SPATIAL_TENDON_MASK, device=self._device) def _get_cpu_env_mask(self, env_mask: wp.array) -> wp.array: """Return a pinned-host CPU copy of :paramref:`env_mask` for a CPU-only binding write. diff --git a/source/isaaclab_ov/isaaclab_ov/assets/deformable_object/deformable_object.py b/source/isaaclab_ov/isaaclab_ov/assets/deformable_object/deformable_object.py index ccca84b3cd0a..ceef9f0e2049 100644 --- a/source/isaaclab_ov/isaaclab_ov/assets/deformable_object/deformable_object.py +++ b/source/isaaclab_ov/isaaclab_ov/assets/deformable_object/deformable_object.py @@ -22,6 +22,7 @@ from isaaclab.assets.deformable_object.base_deformable_object import BaseDeformableObject from isaaclab.markers import VisualizationMarkers from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_ids from isaaclab_ov import tensor_types as TT from isaaclab_ov.sim.views.ovphysx_view import OvPhysxView @@ -528,22 +529,7 @@ def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array(dtyp dtype=wp.int32 ): """Resolve environment indices to a device-resident int32 Warp array.""" - if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)): - return self._ALL_INDICES - if isinstance(env_ids, slice): - # OVPhysX consumes these IDs directly and requires contiguous DLPack tensors. - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids].contiguous()) - if isinstance(env_ids, torch.Tensor): - values = env_ids.to(device=self.device, dtype=torch.int32).contiguous() - return wp.from_torch(values, dtype=wp.int32) - if isinstance(env_ids, wp.array): - if env_ids.dtype != wp.int32: - values = wp.to_torch(env_ids).to(device=self.device, dtype=torch.int32).contiguous() - return wp.from_torch(values, dtype=wp.int32) - if str(env_ids.device) != self.device: - return wp.clone(env_ids, device=self.device) - return env_ids - return wp.array(list(env_ids), dtype=wp.int32, device=self.device) + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device, dtype=wp.int32) @staticmethod def _arrays_overlap(first: wp.array, second: wp.array) -> bool: diff --git a/source/isaaclab_ov/isaaclab_ov/assets/rigid_object/rigid_object.py b/source/isaaclab_ov/isaaclab_ov/assets/rigid_object/rigid_object.py index 7e19bc75109d..a44665429701 100644 --- a/source/isaaclab_ov/isaaclab_ov/assets/rigid_object/rigid_object.py +++ b/source/isaaclab_ov/isaaclab_ov/assets/rigid_object/rigid_object.py @@ -22,6 +22,7 @@ from isaaclab.sim.utils.queries import path_expr_to_glob, resolve_matching_prims_from_source from isaaclab.utils.string import resolve_matching_names from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_ov import tensor_types as TT @@ -1076,17 +1077,7 @@ def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | No Returns: Environment indices on ``self._device``. """ - if env_ids is None or env_ids == slice(None): - return self._ALL_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self._device) - if isinstance(env_ids, torch.Tensor): - return env_ids.to(device=self._device) - if isinstance(env_ids, wp.array) and str(env_ids.device) != self._device: - env_ids = wp.clone(env_ids, device=self._device) - return env_ids + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self._device) def _sim_env_ids_view(self, count: int) -> wp.array: """Return a cached prefix of the simulator-index scratch buffer.""" @@ -1109,13 +1100,7 @@ def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | Returns: Body indices. """ - if isinstance(body_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if body_ids is None or body_ids == slice(None): - return self._ALL_BODY_INDICES - if isinstance(body_ids, list): - return wp.array(body_ids, dtype=wp.int32, device=self._device) - return body_ids + return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self._device) def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array: """Resolve an environment mask to a ``wp.bool`` array. @@ -1126,13 +1111,7 @@ def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array: Returns: A warp ``wp.bool`` array on ``self._device``. """ - if env_mask is None: - return self._ALL_TRUE_ENV_MASK - if isinstance(env_mask, torch.Tensor): - return wp.from_torch(env_mask.to(torch.bool), dtype=wp.bool) - if isinstance(env_mask, wp.array) and str(env_mask.device) != self._device: - env_mask = wp.clone(env_mask, device=self._device) - return env_mask + return resolve_1d_mask(mask=env_mask, all_mask=self._ALL_TRUE_ENV_MASK, device=self._device) def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array: """Resolve a body mask to a ``wp.bool`` array. @@ -1143,13 +1122,7 @@ def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array: Returns: A warp ``wp.bool`` array on ``self._device``. """ - if body_mask is None: - return self._ALL_TRUE_BODY_MASK - if isinstance(body_mask, torch.Tensor): - return wp.from_torch(body_mask.to(torch.bool), dtype=wp.bool) - if isinstance(body_mask, wp.array) and str(body_mask.device) != self._device: - body_mask = wp.clone(body_mask, device=self._device) - return body_mask + return resolve_1d_mask(mask=body_mask, all_mask=self._ALL_TRUE_BODY_MASK, device=self._device) def _get_cpu_env_mask(self, env_mask: wp.array) -> wp.array: """Return a pinned-host CPU copy of *env_mask* for a CPU-only binding write. diff --git a/source/isaaclab_ov/isaaclab_ov/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_ov/isaaclab_ov/assets/rigid_object_collection/rigid_object_collection.py index c39c8c6d2f2b..1d4066c2709a 100644 --- a/source/isaaclab_ov/isaaclab_ov/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_ov/isaaclab_ov/assets/rigid_object_collection/rigid_object_collection.py @@ -21,6 +21,7 @@ from isaaclab.utils import clone, validate from isaaclab.utils.string import resolve_matching_names from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_ov import tensor_types as TT @@ -439,16 +440,8 @@ def write_body_link_pose_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask - env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0]) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_mask_t = wp.to_torch(body_mask) if isinstance(body_mask, wp.array) else body_mask - body_ids = self._resolve_body_ids(torch.nonzero(body_mask_t)[:, 0]) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device) + body_ids = resolve_ids(mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self._device) self.assert_shape_and_dtype(body_poses, (self._num_instances, self._num_bodies), wp.transformf, "body_poses") wp.launch( shared_kernels.set_body_link_pose_to_sim_kernel(env_ids, body_ids), @@ -535,16 +528,8 @@ def write_body_com_pose_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask - env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0]) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_mask_t = wp.to_torch(body_mask) if isinstance(body_mask, wp.array) else body_mask - body_ids = self._resolve_body_ids(torch.nonzero(body_mask_t)[:, 0]) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device) + body_ids = resolve_ids(mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self._device) self.assert_shape_and_dtype(body_poses, (self._num_instances, self._num_bodies), wp.transformf, "body_poses") wp.launch( shared_kernels.set_body_com_pose_to_sim_kernel(env_ids, body_ids), @@ -643,16 +628,8 @@ def write_body_com_velocity_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask - env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0]) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_mask_t = wp.to_torch(body_mask) if isinstance(body_mask, wp.array) else body_mask - body_ids = self._resolve_body_ids(torch.nonzero(body_mask_t)[:, 0]) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device) + body_ids = resolve_ids(mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self._device) self.assert_shape_and_dtype( body_velocities, (self._num_instances, self._num_bodies), wp.spatial_vectorf, "body_velocities" ) @@ -758,16 +735,8 @@ def write_body_link_velocity_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask - env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0]) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_mask_t = wp.to_torch(body_mask) if isinstance(body_mask, wp.array) else body_mask - body_ids = self._resolve_body_ids(torch.nonzero(body_mask_t)[:, 0]) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device) + body_ids = resolve_ids(mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self._device) self.assert_shape_and_dtype( body_velocities, (self._num_instances, self._num_bodies), wp.spatial_vectorf, "body_velocities" ) @@ -869,11 +838,7 @@ def set_masses_mask( env_mask: Environment mask. If None, all instances are updated. Shape is (num_instances,). """ - if env_mask is not None: - env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask - env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0]) - else: - env_ids = self._ALL_ENV_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device) self.assert_shape_and_dtype(masses, (self._num_instances, self._num_bodies), wp.float32, "masses") wp.launch( shared_kernels.write_2d_data_to_buffer_with_mask, @@ -955,11 +920,7 @@ def set_coms_mask( env_mask: Environment mask. If None, all instances are updated. Shape is (num_instances,). """ - if env_mask is not None: - env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask - env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0]) - else: - env_ids = self._ALL_ENV_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device) self.assert_shape_and_dtype(coms, (self._num_instances, self._num_bodies), wp.transformf, "coms") wp.launch( shared_kernels.write_body_com_pose_to_buffer_mask, @@ -1051,11 +1012,7 @@ def set_inertias_mask( env_mask: Environment mask. If None, all instances are updated. Shape is (num_instances,). """ - if env_mask is not None: - env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask - env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0]) - else: - env_ids = self._ALL_ENV_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device) self.assert_shape_and_dtype(inertias, (self._num_instances, self._num_bodies, 9), wp.float32, "inertias") wp.launch( shared_kernels.write_body_inertia_to_buffer_mask, @@ -1388,27 +1345,11 @@ def _resolve_env_ids(self, env_ids) -> wp.array | torch.Tensor: resolved array onto ``self._device`` so kernel launches don't fail on a device mismatch. """ - if env_ids is None or env_ids == slice(None): - return self._ALL_ENV_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_ENV_INDICES)[env_ids]) - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self._device) - if isinstance(env_ids, torch.Tensor): - return env_ids.to(device=self._device) - if isinstance(env_ids, wp.array) and str(env_ids.device) != self._device: - env_ids = wp.clone(env_ids, device=self._device) - return env_ids + return resolve_ids(ids=env_ids, all_ids=self._ALL_ENV_INDICES, device=self._device) def _resolve_body_ids(self, body_ids) -> wp.array | torch.Tensor: """Resolve body indices on ``self._device``.""" - if isinstance(body_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if body_ids is None or body_ids == slice(None): - return self._ALL_BODY_INDICES - if isinstance(body_ids, list): - return wp.array(body_ids, dtype=wp.int32, device=self._device) - return body_ids + return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self._device) def _env_body_ids_to_view_ids( self, env_ids: torch.Tensor | wp.array, body_ids: torch.Tensor | wp.array, device: str = "cuda:0" @@ -1418,6 +1359,8 @@ def _env_body_ids_to_view_ids( Computes ``view_id = body_id * num_instances + env_id`` for each (env_id, body_id) pair. The output array is laid out column-major over the (env, body) grid, matching the PhysX ``root_view`` ordering. + The mask writers convert their masks to ids for this: a binding write mask must cover all + ``num_bodies * num_instances`` rows, so an env mask does not fit. Args: env_ids: Environment indices. @@ -1466,13 +1409,7 @@ def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array: Returns: A ``wp.bool`` array of shape (num_instances,) on ``self._device``. """ - if env_mask is None: - return self._ALL_TRUE_ENV_MASK - if isinstance(env_mask, torch.Tensor): - return wp.from_torch(env_mask.to(torch.bool), dtype=wp.bool) - if isinstance(env_mask, wp.array) and str(env_mask.device) != self._device: - env_mask = wp.clone(env_mask, device=self._device) - return env_mask + return resolve_1d_mask(mask=env_mask, all_mask=self._ALL_TRUE_ENV_MASK, device=self._device) def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array: """Resolve a body mask to a ``wp.bool`` array on ``self._device``. @@ -1485,13 +1422,7 @@ def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array: Returns: A ``wp.bool`` array of shape (num_bodies,) on ``self._device``. """ - if body_mask is None: - return self._ALL_TRUE_BODY_MASK - if isinstance(body_mask, torch.Tensor): - return wp.from_torch(body_mask.to(torch.bool), dtype=wp.bool) - if isinstance(body_mask, wp.array) and str(body_mask.device) != self._device: - body_mask = wp.clone(body_mask, device=self._device) - return body_mask + return resolve_1d_mask(mask=body_mask, all_mask=self._ALL_TRUE_BODY_MASK, device=self._device) def _sim_env_ids_view(self, count: int) -> wp.array: """Return a cached prefix of the simulator-index scratch buffer.""" diff --git a/source/isaaclab_ov/test/assets/test_deformable_object.py b/source/isaaclab_ov/test/assets/test_deformable_object.py index bc9fc6a2e04b..5ea0b8ea84f6 100644 --- a/source/isaaclab_ov/test/assets/test_deformable_object.py +++ b/source/isaaclab_ov/test/assets/test_deformable_object.py @@ -386,6 +386,28 @@ def test_indexed_partial_write_preserves_retained_aliased_slice( torch.testing.assert_close(retained.torch, expected) +@pytest.mark.usefixtures("host_shell") +def test_mask_writes_update_only_masked_rows(): + """The mask writers take full-size data and write only the environments the mask selects.""" + asset = _make_asset_shell(deformable_type="volume", num_instances=3, num_vertices=2) + zeros = torch.zeros((3, 2, 3), device=asset.device) + asset.root_view.positions = wp.from_torch(zeros.clone(), dtype=wp.float32) + asset.root_view.velocities = wp.from_torch(zeros.clone(), dtype=wp.float32) + asset.update(0.1) + env_mask = wp.array([False, True, False], dtype=wp.bool, device=asset.device) + + state = torch.cat((torch.full((3, 2, 3), 5.0), torch.full((3, 2, 3), -5.0)), dim=-1).to(asset.device) + asset.write_nodal_state_to_sim_mask(state, env_mask=env_mask) + expected = torch.zeros((3, 2, 6), device=asset.device) + expected[1] = state[1] + torch.testing.assert_close(asset.data.nodal_state_w.torch, expected) + + positions = torch.full((3, 2, 3), 9.0, device=asset.device) + asset.write_nodal_pos_to_sim_mask(positions, env_mask=env_mask) + expected[1, :, :3] = 9.0 + torch.testing.assert_close(asset.data.nodal_pos_w.torch, expected[..., :3]) + + @pytest.mark.usefixtures("host_shell") def test_full_overwrite_stale_cache_does_not_read_simulator() -> None: asset = _make_asset_shell(deformable_type="volume", num_instances=3, num_vertices=2) From 3d3661f29616ab829138b9dbdbbbf81f1b79e4e0 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Mon, 26 Oct 2026 02:48:22 -0700 Subject: [PATCH 6/6] Resolve PhysX selections through the shared helpers The PhysX assets resolve env, body, joint and tendon selections through resolve_ids and resolve_1d_mask instead of their own copies. The sensor test doubles get the env masks a real sensor creates. --- .../changelog.d/jichuanh-id-mask.changed.rst | 7 + .../assets/articulation/articulation.py | 282 ++++++------------ .../assets/rigid_object/rigid_object.py | 113 ++----- .../rigid_object_collection.py | 119 ++------ .../assets/surface_gripper/surface_gripper.py | 29 +- .../test/sensors/test_frame_transformer.py | 2 + .../sensors/test_imu_pva_recorded_launch.py | 2 + 7 files changed, 154 insertions(+), 400 deletions(-) create mode 100644 source/isaaclab_physx/changelog.d/jichuanh-id-mask.changed.rst diff --git a/source/isaaclab_physx/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab_physx/changelog.d/jichuanh-id-mask.changed.rst new file mode 100644 index 000000000000..fb8da731a750 --- /dev/null +++ b/source/isaaclab_physx/changelog.d/jichuanh-id-mask.changed.rst @@ -0,0 +1,7 @@ +* Changed the env selection helpers of the PhysX assets to resolve ``env_ids`` through + :func:`~isaaclab.utils.warp.utils.resolve_ids`. Any integer sequence and slices with a negative step are accepted, + ids on another device are moved to the asset's device, and a :class:`~isaaclab.utils.warp.ProxyArray` is rejected + with a :class:`TypeError`. +* Changed the body, joint and tendon selection helpers of the PhysX assets to resolve ids and masks through + :func:`~isaaclab.utils.warp.utils.resolve_ids`, as the env selection does, so slices other than ``slice(None)`` are + accepted. diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 3d2cb4f1d946..b1be80852be8 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -27,6 +27,7 @@ from isaaclab.utils.string import resolve_matching_names, resolve_matching_names_values from isaaclab.utils.version import get_isaac_sim_version, has_kit from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_physx.assets import kernels as shared_kernels @@ -590,7 +591,7 @@ def write_root_link_pose_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) # Set full data to True to ensure the the right code path is taken inside the kernel. self.write_root_link_pose_to_sim_index( root_pose=root_pose, env_ids=env_ids, full_data=True, skip_forward=skip_forward @@ -692,7 +693,7 @@ def write_root_com_pose_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) # Set full data to True to ensure the the right code path is taken inside the kernel. self.write_root_com_pose_to_sim_index( root_pose=root_pose, env_ids=env_ids, full_data=True, skip_forward=skip_forward @@ -866,7 +867,7 @@ def write_root_com_velocity_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) # Set full data to True to ensure the the right code path is taken inside the kernel. self.write_root_com_velocity_to_sim_index( root_velocity=root_velocity, env_ids=env_ids, full_data=True, skip_forward=skip_forward @@ -977,7 +978,7 @@ def write_root_link_velocity_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) # Set full data to True to ensure the the right code path is taken inside the kernel. self.write_root_link_velocity_to_sim_index( root_velocity=root_velocity, env_ids=env_ids, full_data=True, skip_forward=skip_forward @@ -1091,8 +1092,8 @@ def write_joint_state_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. """ # set into simulation - env_ids = self._resolve_env_mask(env_mask) - joint_ids = None if joint_mask is None else self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = None if joint_mask is None else self._resolve_joint_ids(joint_mask=joint_mask) self.write_joint_state_to_sim_index( position=position, velocity=velocity, @@ -1196,8 +1197,8 @@ def write_joint_position_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) - joint_ids = None if joint_mask is None else self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = None if joint_mask is None else self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the the right code path is taken inside the kernel. self.write_joint_position_to_sim_index( position=position, joint_ids=joint_ids, env_ids=env_ids, full_data=True, skip_forward=skip_forward @@ -1301,8 +1302,8 @@ def write_joint_velocity_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) - joint_ids = None if joint_mask is None else self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = None if joint_mask is None else self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the the right code path is taken inside the kernel. self.write_joint_velocity_to_sim_index( velocity=velocity, joint_ids=joint_ids, env_ids=env_ids, full_data=True, skip_forward=skip_forward @@ -1412,8 +1413,8 @@ def write_joint_stiffness_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - joint_ids = self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the the right code path is taken inside the kernel. self.write_joint_stiffness_to_sim_index( stiffness=stiffness, joint_ids=joint_ids, env_ids=env_ids, full_data=True @@ -1560,8 +1561,8 @@ def write_joint_damping_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) - joint_ids = self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.write_joint_damping_to_sim_index(damping=damping, joint_ids=joint_ids, env_ids=env_ids, full_data=True) @@ -1672,8 +1673,8 @@ def write_joint_position_limit_to_sim_mask( exceed the new limits. Defaults to True. """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) - joint_ids = self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.write_joint_position_limit_to_sim_index( limits=limits, @@ -1787,8 +1788,8 @@ def write_joint_velocity_limit_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) - joint_ids = self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.write_joint_velocity_limit_to_sim_index( limits=limits, joint_ids=joint_ids, env_ids=env_ids, full_data=True @@ -1897,8 +1898,8 @@ def write_joint_effort_limit_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) - joint_ids = self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.write_joint_effort_limit_to_sim_index(limits=limits, joint_ids=joint_ids, env_ids=env_ids, full_data=True) @@ -2005,8 +2006,8 @@ def write_joint_armature_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) - joint_ids = self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.write_joint_armature_to_sim_index(armature=armature, joint_ids=joint_ids, env_ids=env_ids, full_data=True) @@ -2173,8 +2174,8 @@ def write_joint_friction_coefficient_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) - joint_ids = self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.write_joint_friction_coefficient_to_sim_index( joint_friction_coeff=joint_friction_coeff, @@ -2283,8 +2284,8 @@ def write_joint_dynamic_friction_coefficient_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) - joint_ids = self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.write_joint_dynamic_friction_coefficient_to_sim_index( joint_dynamic_friction_coeff=joint_dynamic_friction_coeff, @@ -2394,8 +2395,8 @@ def write_joint_viscous_friction_coefficient_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) - joint_ids = self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.write_joint_viscous_friction_coefficient_to_sim_index( joint_viscous_friction_coeff=joint_viscous_friction_coeff, @@ -2496,8 +2497,8 @@ def set_masses_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - body_ids = self._resolve_body_mask(body_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_masses_index(masses=masses, body_ids=body_ids, env_ids=env_ids, full_data=True) @@ -2603,8 +2604,8 @@ def set_coms_mask( all_envs_selected = env_mask is None all_bodies_selected = body_mask is None # Resolve masks. - env_ids = None if all_envs_selected else self._resolve_env_mask(env_mask) - body_ids = None if all_bodies_selected else self._resolve_body_mask(body_mask) + env_ids = None if all_envs_selected else self._resolve_env_ids(env_mask=env_mask) + body_ids = None if all_bodies_selected else self._resolve_body_ids(body_mask=body_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_coms_index(coms=coms, body_ids=body_ids, env_ids=env_ids, full_data=True) @@ -2700,8 +2701,8 @@ def set_inertias_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - body_ids = self._resolve_body_mask(body_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_inertias_index(inertias=inertias, body_ids=body_ids, env_ids=env_ids, full_data=True) @@ -2806,8 +2807,8 @@ def set_fixed_tendon_stiffness_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + fixed_tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_fixed_tendon_stiffness_index( stiffness=stiffness, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True @@ -2906,8 +2907,8 @@ def set_fixed_tendon_damping_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + fixed_tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_fixed_tendon_damping_index( damping=damping, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True @@ -3010,8 +3011,8 @@ def set_fixed_tendon_limit_stiffness_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + fixed_tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_fixed_tendon_limit_stiffness_index( limit_stiffness=limit_stiffness, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True @@ -3088,8 +3089,8 @@ def set_fixed_tendon_position_limit_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + fixed_tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_fixed_tendon_position_limit_index( limit=limit, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True @@ -3192,8 +3193,8 @@ def set_fixed_tendon_rest_length_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + fixed_tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_fixed_tendon_rest_length_index( rest_length=rest_length, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True @@ -3343,8 +3344,8 @@ def set_fixed_tendon_position_target_mask( """ self.set_fixed_tendon_position_target_index( target=target, - fixed_tendon_ids=self._resolve_fixed_tendon_mask(fixed_tendon_mask), - env_ids=self._resolve_env_mask(env_mask), + fixed_tendon_ids=self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask), + env_ids=self._resolve_env_ids(env_mask=env_mask), full_data=True, ) @@ -3374,8 +3375,8 @@ def set_fixed_tendon_offset_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + fixed_tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_fixed_tendon_offset_index( offset=offset, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True @@ -3430,7 +3431,7 @@ def write_fixed_tendon_properties_to_sim_mask( only the selected cached values before this method pushes the complete rows. .. tip:: - For maximum performance we recommend using the mask method. This is because in PhysX, the tensor API + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API is only supporting indexing, hence masks need to be converted to indices. Args: @@ -3439,7 +3440,7 @@ def write_fixed_tendon_properties_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) self.write_fixed_tendon_properties_to_sim_index(env_ids=env_ids) def set_spatial_tendon_stiffness_index( @@ -3539,8 +3540,8 @@ def set_spatial_tendon_stiffness_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - spatial_tendon_ids = self._resolve_spatial_tendon_mask(spatial_tendon_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + spatial_tendon_ids = self._resolve_spatial_tendon_ids(spatial_tendon_mask=spatial_tendon_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_spatial_tendon_stiffness_index( stiffness=stiffness, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids, full_data=True @@ -3639,8 +3640,8 @@ def set_spatial_tendon_damping_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - spatial_tendon_ids = self._resolve_spatial_tendon_mask(spatial_tendon_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + spatial_tendon_ids = self._resolve_spatial_tendon_ids(spatial_tendon_mask=spatial_tendon_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_spatial_tendon_damping_index( damping=damping, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids, full_data=True @@ -3744,8 +3745,8 @@ def set_spatial_tendon_limit_stiffness_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - spatial_tendon_ids = self._resolve_spatial_tendon_mask(spatial_tendon_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + spatial_tendon_ids = self._resolve_spatial_tendon_ids(spatial_tendon_mask=spatial_tendon_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_spatial_tendon_limit_stiffness_index( limit_stiffness=limit_stiffness, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids, full_data=True @@ -3844,8 +3845,8 @@ def set_spatial_tendon_offset_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - spatial_tendon_ids = self._resolve_spatial_tendon_mask(spatial_tendon_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + spatial_tendon_ids = self._resolve_spatial_tendon_ids(spatial_tendon_mask=spatial_tendon_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_spatial_tendon_offset_index( offset=offset, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids, full_data=True @@ -3898,7 +3899,7 @@ def write_spatial_tendon_properties_to_sim_mask( this method pushes the complete rows. .. tip:: - For maximum performance we recommend using the mask method. This is because in PhysX, the tensor API + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API is only supporting indexing, hence masks need to be converted to indices. Args: @@ -3907,7 +3908,7 @@ def write_spatial_tendon_properties_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) # Write spatial tendon properties to the simulation. self.write_spatial_tendon_properties_to_sim_index(env_ids=env_ids) @@ -4463,62 +4464,23 @@ def _cpu_env_ids_view(self, count: int) -> wp.array: ) return self._cpu_env_ids_views[count] - def _resolve_env_mask(self, env_mask: wp.array | None) -> torch.Tensor | wp.array: - """ - Resolve environment mask to a torch tensor. - - Args: - env_mask: Environment mask. If None, then all indices are used. - - Returns: - A torch tensor of environment indices. - """ - # resolve masks - if env_mask is not None: - if isinstance(env_mask, wp.array): - env_mask = wp.to_torch(env_mask) - env_ids = torch.nonzero(env_mask)[:, 0] - else: - env_ids = self._ALL_INDICES - return env_ids - - def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: + def _resolve_env_ids( + self, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_mask: wp.array | None = None + ) -> wp.array | torch.Tensor: """Resolve environment indices. Args: env_ids: Environment indices. If None, then all indices are used. + env_mask: Boolean mask; takes precedence over ``env_ids``. Returns: Environment indices. """ - if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)): - return self._ALL_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - if isinstance(env_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self.device) - return env_ids - - def _resolve_joint_mask(self, joint_mask: wp.array | None) -> torch.Tensor | wp.array: - """Resolve joint mask to a torch tensor. - - Args: - joint_mask: Joint mask. If None, then all indices are used. + return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) - Returns: - A torch tensor of joint indices. - """ - if joint_mask is not None: - if isinstance(joint_mask, wp.array): - joint_mask = wp.to_torch(joint_mask) - joint_ids = torch.nonzero(joint_mask)[:, 0] - else: - joint_ids = self._ALL_JOINT_INDICES - return joint_ids - - def _resolve_joint_ids(self, joint_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: + def _resolve_joint_ids( + self, joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, joint_mask: wp.array | None = None + ) -> wp.array | torch.Tensor: """Resolve joint indices to a warp array or tensor. .. note:: @@ -4526,123 +4488,65 @@ def _resolve_joint_ids(self, joint_ids: Sequence[int] | torch.Tensor | wp.array Args: joint_ids: Joint indices. If None, then all indices are used. + joint_mask: Boolean mask; takes precedence over ``joint_ids``. Returns: A warp array of joint indices or a tensor of joint indices. """ - if isinstance(joint_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(joint_ids, list): - return wp.array(joint_ids, dtype=wp.int32, device=self.device) - if (joint_ids is None) or (joint_ids == slice(None)): - return self._ALL_JOINT_INDICES - return joint_ids - - def _resolve_body_mask(self, body_mask: wp.array | None) -> torch.Tensor | wp.array: - """Resolve body mask to a torch tensor. - - Args: - body_mask: Body mask. If None, then all indices are used. - - Returns: - A torch tensor of body indices. - """ - if body_mask is not None: - if isinstance(body_mask, wp.array): - body_mask = wp.to_torch(body_mask) - body_ids = torch.nonzero(body_mask)[:, 0] - else: - body_ids = self._ALL_BODY_INDICES - return body_ids + return resolve_ids(ids=joint_ids, mask=joint_mask, all_ids=self._ALL_JOINT_INDICES, device=self.device) - def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: + def _resolve_body_ids( + self, body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, body_mask: wp.array | None = None + ) -> wp.array | torch.Tensor: """Resolve body indices to a warp array or tensor. Args: body_ids: Body indices. If None, then all indices are used. + body_mask: Boolean mask; takes precedence over ``body_ids``. Returns: A warp array of body indices or a tensor of body indices. """ - if isinstance(body_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(body_ids, list): - return wp.array(body_ids, dtype=wp.int32, device=self.device) - if (body_ids is None) or (body_ids == slice(None)): - return self._ALL_BODY_INDICES - return body_ids - - def _resolve_fixed_tendon_mask(self, fixed_tendon_mask: wp.array | None) -> torch.Tensor | wp.array: - """Resolve fixed tendon mask to a torch tensor. - - Args: - fixed_tendon_mask: Fixed tendon mask. If None, then all indices are used. - - Returns: - A torch tensor of fixed tendon indices. - """ - if fixed_tendon_mask is not None: - if isinstance(fixed_tendon_mask, wp.array): - fixed_tendon_mask = wp.to_torch(fixed_tendon_mask) - fixed_tendon_ids = torch.nonzero(fixed_tendon_mask)[:, 0] - else: - fixed_tendon_ids = self._ALL_FIXED_TENDON_INDICES - return fixed_tendon_ids + return resolve_ids(ids=body_ids, mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self.device) def _resolve_fixed_tendon_ids( - self, tendon_ids: Sequence[int] | torch.Tensor | wp.array | None + self, + tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + fixed_tendon_mask: wp.array | None = None, ) -> wp.array | torch.Tensor: """Resolve tendon indices to a warp array or tensor. Args: tendon_ids: Tendon indices. If None, then all indices are used. + fixed_tendon_mask: Boolean mask; takes precedence over ``tendon_ids``. Returns: A warp array of tendon indices or a tensor of tendon indices. """ - if isinstance(tendon_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(tendon_ids, list): - return wp.array(tendon_ids, dtype=wp.int32, device=self.device) - if (tendon_ids is None) or (tendon_ids == slice(None)): - return self._ALL_FIXED_TENDON_INDICES - return tendon_ids - - def _resolve_spatial_tendon_mask(self, spatial_tendon_mask: wp.array | None) -> torch.Tensor | wp.array: - """Resolve spatial tendon mask to a torch tensor. - - Args: - spatial_tendon_mask: Spatial tendon mask. If None, then all indices are used. - - Returns: - A torch tensor of spatial tendon indices. - """ - if spatial_tendon_mask is not None: - if isinstance(spatial_tendon_mask, wp.array): - spatial_tendon_mask = wp.to_torch(spatial_tendon_mask) - spatial_tendon_ids = torch.nonzero(spatial_tendon_mask)[:, 0] - else: - spatial_tendon_ids = self._ALL_SPATIAL_TENDON_INDICES - return spatial_tendon_ids + return resolve_ids( + ids=tendon_ids, mask=fixed_tendon_mask, all_ids=self._ALL_FIXED_TENDON_INDICES, device=self.device + ) def _resolve_spatial_tendon_ids( - self, spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None + self, + spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + spatial_tendon_mask: wp.array | None = None, ) -> wp.array | torch.Tensor: """Resolve spatial tendon indices to a warp array or tensor. Args: spatial_tendon_ids: Spatial tendon indices. If None, then all indices are used. + spatial_tendon_mask: Boolean mask; takes precedence over ``spatial_tendon_ids``. Returns: A warp array of spatial tendon indices or a tensor of spatial tendon indices. """ - if isinstance(spatial_tendon_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(spatial_tendon_ids, list): - return wp.array(spatial_tendon_ids, dtype=wp.int32, device=self.device) - if (spatial_tendon_ids is None) or (spatial_tendon_ids == slice(None)): - return self._ALL_SPATIAL_TENDON_INDICES - return spatial_tendon_ids + return resolve_ids( + ids=spatial_tendon_ids, + mask=spatial_tendon_mask, + all_ids=self._ALL_SPATIAL_TENDON_INDICES, + device=self.device, + ) """ Deprecated methods. diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py index 47c7240c510e..48d1345349ea 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py @@ -19,6 +19,7 @@ from isaaclab.assets.rigid_object.base_rigid_object import BaseRigidObject from isaaclab.sim.utils.queries import path_expr_to_glob, resolve_matching_prims_from_source from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_physx.assets import kernels as shared_kernels @@ -397,10 +398,7 @@ def write_root_link_pose_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) self.write_root_link_pose_to_sim_index( root_pose=root_pose, env_ids=env_ids, full_data=True, skip_forward=skip_forward ) @@ -491,10 +489,7 @@ def write_root_com_pose_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) self.write_root_com_pose_to_sim_index( root_pose=root_pose, env_ids=env_ids, full_data=True, skip_forward=skip_forward ) @@ -593,10 +588,7 @@ def write_root_com_velocity_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) self.write_root_com_velocity_to_sim_index( root_velocity=root_velocity, env_ids=env_ids, full_data=True, skip_forward=skip_forward ) @@ -699,10 +691,7 @@ def write_root_link_velocity_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) self.write_root_link_velocity_to_sim_index( root_velocity=root_velocity, env_ids=env_ids, full_data=True, skip_forward=skip_forward ) @@ -793,14 +782,8 @@ def set_masses_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_masses_index(masses=masses, body_ids=body_ids, env_ids=env_ids, full_data=True) @@ -885,14 +868,8 @@ def set_coms_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_coms_index(coms=coms, body_ids=body_ids, env_ids=env_ids, full_data=True) @@ -980,14 +957,8 @@ def set_inertias_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_inertias_index(inertias=inertias, body_ids=body_ids, env_ids=env_ids, full_data=True) @@ -1068,40 +1039,6 @@ def _process_cfg(self) -> None: self._data.default_root_pose = wp.array(default_root_pose, dtype=wp.transformf, device=self.device) self._data.default_root_vel = wp.array(default_root_vel, dtype=wp.spatial_vectorf, device=self.device) - def _resolve_env_mask(self, env_mask: wp.array | None) -> torch.Tensor | wp.array: - """Resolve environment mask to a torch tensor. - - Args: - env_mask: Environment mask. If None, then all indices are used. - - Returns: - A torch tensor of environment indices. - """ - if env_mask is not None: - if isinstance(env_mask, wp.array): - env_mask = wp.to_torch(env_mask) - env_ids = torch.nonzero(env_mask)[:, 0] - else: - env_ids = self._ALL_INDICES - return env_ids - - def _resolve_body_mask(self, body_mask: wp.array | None) -> torch.Tensor | wp.array: - """Resolve body mask to a torch tensor. - - Args: - body_mask: Body mask. If None, then all indices are used. - - Returns: - A torch tensor of body indices. - """ - if body_mask is not None: - if isinstance(body_mask, wp.array): - body_mask = wp.to_torch(body_mask) - body_ids = torch.nonzero(body_mask)[:, 0] - else: - body_ids = self._ALL_BODY_INDICES - return body_ids - def _get_cpu_env_ids(self, env_ids: wp.array | torch.Tensor, sim_env_ids: wp.array | None = None) -> wp.array: """Get CPU environment indices. Uses pre-allocated pinned buffer for full-index case. @@ -1164,39 +1101,33 @@ def _sim_env_ids_view(self, count: int) -> wp.array: ) return self._sim_env_ids_views[count] - def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: + def _resolve_env_ids( + self, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_mask: wp.array | None = None + ) -> wp.array | torch.Tensor: """Resolve environment indices to a warp array or tensor. Args: env_ids: Environment indices. If None, then all indices are used. + env_mask: Boolean mask; takes precedence over ``env_ids``. Returns: A warp array of environment indices or a tensor of environment indices. """ - if (env_ids is None) or (env_ids == slice(None)): - return self._ALL_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self.device) - return env_ids - - def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: + return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) + + def _resolve_body_ids( + self, body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, body_mask: wp.array | None = None + ) -> wp.array | torch.Tensor: """Resolve body indices to a warp array or tensor. Args: body_ids: Body indices. If None, then all indices are used. + body_mask: Boolean mask; takes precedence over ``body_ids``. Returns: A warp array of body indices or a tensor of body indices. """ - if isinstance(body_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(body_ids, list): - return wp.array(body_ids, dtype=wp.int32, device=self.device) - if (body_ids is None) or (body_ids == slice(None)): - return self._ALL_BODY_INDICES - return body_ids + return resolve_ids(ids=body_ids, mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self.device) """ Internal simulation callbacks. diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py index 47d5e1cb460b..e79295d59f04 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py @@ -22,6 +22,7 @@ from isaaclab.assets.rigid_object_collection.base_rigid_object_collection import BaseRigidObjectCollection from isaaclab.utils import clone, validate from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_physx.assets import kernels as shared_kernels @@ -307,14 +308,8 @@ def write_body_pose_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) self.write_body_link_pose_to_sim_index( body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward ) @@ -391,14 +386,8 @@ def write_body_velocity_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) self.write_body_com_velocity_to_sim_index( body_velocities=body_velocities, env_ids=env_ids, @@ -509,10 +498,7 @@ def write_body_link_pose_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. body_ids: Deprecated, use :attr:`body_mask` instead. Body indices. Defaults to None. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) body_ids = self._resolve_mask_writer_body_ids(body_mask, body_ids) self.write_body_link_pose_to_sim_index( body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward @@ -624,10 +610,7 @@ def write_body_com_pose_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. body_ids: Deprecated, use :attr:`body_mask` instead. Body indices. Defaults to None. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) body_ids = self._resolve_mask_writer_body_ids(body_mask, body_ids) self.write_body_com_pose_to_sim_index( body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward @@ -744,10 +727,7 @@ def write_body_com_velocity_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. body_ids: Deprecated, use :attr:`body_mask` instead. Body indices. Defaults to None. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) body_ids = self._resolve_mask_writer_body_ids(body_mask, body_ids) self.write_body_com_velocity_to_sim_index( body_velocities=body_velocities, @@ -871,10 +851,7 @@ def write_body_link_velocity_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. body_ids: Deprecated, use :attr:`body_mask` instead. Body indices. Defaults to None. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) body_ids = self._resolve_mask_writer_body_ids(body_mask, body_ids) self.write_body_link_velocity_to_sim_index( body_velocities=body_velocities, @@ -968,14 +945,8 @@ def set_masses_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_masses_index(masses=masses, body_ids=body_ids, env_ids=env_ids, full_data=True) @@ -1062,14 +1033,8 @@ def set_coms_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_coms_index(coms=coms, body_ids=body_ids, env_ids=env_ids, full_data=True) @@ -1157,14 +1122,8 @@ def set_inertias_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_inertias_index(inertias=inertias, body_ids=body_ids, env_ids=env_ids, full_data=True) @@ -1311,46 +1270,20 @@ def reshape_data_to_view_3d( Internal helper. """ - def _resolve_env_ids(self, env_ids) -> wp.array | torch.Tensor: + def _resolve_env_ids(self, env_ids=None, env_mask: wp.array | None = None) -> wp.array | torch.Tensor: """Resolve environment indices.""" - if (env_ids is None) or (env_ids == slice(None)): - return self._ALL_ENV_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_ENV_INDICES)[env_ids]) - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self.device) - return env_ids - - def _resolve_body_ids(self, body_ids) -> wp.array | torch.Tensor: - """Resolve body indices.""" - if isinstance(body_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(body_ids, list): - return wp.array(body_ids, dtype=wp.int32, device=self.device) - if body_ids is None or (body_ids == slice(None)): - return self._ALL_BODY_INDICES - if isinstance(body_ids, slice): - return wp.from_torch( - torch.arange(self.num_bodies, dtype=torch.int32, device=self.device)[body_ids], dtype=wp.int32 - ) - return body_ids + return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self.device) - def _resolve_env_mask(self, env_mask: wp.array | None) -> torch.Tensor | wp.array: - """Resolve environment mask to indices via torch.nonzero.""" - if env_mask is not None: - if isinstance(env_mask, wp.array): - env_mask = wp.to_torch(env_mask) - env_ids = torch.nonzero(env_mask)[:, 0] - else: - env_ids = self._ALL_ENV_INDICES - return env_ids + def _resolve_body_ids(self, body_ids=None, body_mask: wp.array | None = None) -> wp.array | torch.Tensor: + """Resolve body indices.""" + return resolve_ids(ids=body_ids, mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self.device) def _resolve_mask_writer_body_ids( self, body_mask: wp.array | None, body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None ) -> Sequence[int] | torch.Tensor | wp.array | slice: """Resolve the bodies of a mask writer, accepting the deprecated ``body_ids`` in place of ``body_mask``.""" if body_ids is None: - return self._resolve_body_mask(body_mask) + return self._resolve_body_ids(body_mask=body_mask) if body_mask is not None: raise ValueError("Pass either 'body_mask' or the deprecated 'body_ids', not both.") warnings.warn( @@ -1361,16 +1294,6 @@ def _resolve_mask_writer_body_ids( ) return body_ids - def _resolve_body_mask(self, body_mask: wp.array | None) -> torch.Tensor | wp.array: - """Resolve body mask to indices via torch.nonzero.""" - if body_mask is not None: - if isinstance(body_mask, wp.array): - body_mask = wp.to_torch(body_mask) - body_ids = torch.nonzero(body_mask)[:, 0] - else: - body_ids = self._ALL_BODY_INDICES - return body_ids - def _initialize_impl(self): # clear body names list to prevent double counting on re-initialization self._body_names_list.clear() diff --git a/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper.py b/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper.py index e37a30f7a505..a2f5cb930c8f 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper.py @@ -16,6 +16,7 @@ from isaaclab.assets import AssetBase from isaaclab.utils import clone from isaaclab.utils.version import get_isaac_sim_version, has_kit +from isaaclab.utils.warp.utils import resolve_ids if TYPE_CHECKING: from isaacsim.robot.surface_gripper import GripperView @@ -187,10 +188,7 @@ def set_grippers_command_mask(self, states: torch.Tensor | wp.array, env_mask: t states: A tensor/array of floats representing the gripper command (full size). env_mask: Boolean mask of shape (num_envs,). Defaults to None (all environments). """ - if env_mask is not None: - env_ids = wp.from_torch(torch.argwhere(env_mask).view(-1).to(torch.int32), dtype=wp.int32) - else: - env_ids = self._ALL_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) self.set_grippers_command_index(states, env_ids, full_data=True) def set_grippers_command(self, states: torch.Tensor, indices: torch.Tensor | None = None) -> None: @@ -274,10 +272,7 @@ def update_gripper_properties_mask( retry_interval: The retry interval. Shape (num_envs,). env_mask: Boolean mask of shape (num_envs,). Defaults to None (all environments). """ - if env_mask is not None: - env_ids = wp.from_torch(torch.argwhere(env_mask).view(-1).to(torch.int32), dtype=wp.int32) - else: - env_ids = self._ALL_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) self.update_gripper_properties_index( max_grip_distance=max_grip_distance, coaxial_force_limit=coaxial_force_limit, @@ -407,10 +402,7 @@ def reset_mask(self, env_mask: torch.Tensor | None = None) -> None: Args: env_mask: Boolean mask of shape (num_envs,). Defaults to None (all environments). """ - if env_mask is not None: - env_ids = wp.from_torch(torch.argwhere(env_mask).view(-1).to(torch.int32), dtype=wp.int32) - else: - env_ids = self._ALL_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) self.reset_index(env_ids) def reset(self, indices: torch.Tensor | None = None) -> None: @@ -552,24 +544,17 @@ def _process_cfg(self) -> None: Helper functions. """ - def _resolve_env_ids(self, env_ids) -> wp.array: + def _resolve_env_ids(self, env_ids=None, env_mask: wp.array | torch.Tensor | None = None) -> wp.array: """Resolve environment indices to a warp array. Args: env_ids: Environment indices. Can be None, a slice, a list, or a torch.Tensor. + env_mask: Boolean mask; takes precedence over ``env_ids``. Returns: A warp array of int32 indices. """ - if env_ids is None or env_ids == slice(None): - return self._ALL_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - elif isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self._device) - elif isinstance(env_ids, torch.Tensor): - return wp.from_torch(env_ids.to(torch.int32).contiguous(), dtype=wp.int32) - return env_ids + return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_INDICES, device=self._device, dtype=wp.int32) def parse_gripper_parameter( self, cfg_value: float | int | tuple | None, default_value: float | int | tuple | None, ndim: int = 0 diff --git a/source/isaaclab_physx/test/sensors/test_frame_transformer.py b/source/isaaclab_physx/test/sensors/test_frame_transformer.py index 7ba9a23d3de4..50e52bfbc646 100644 --- a/source/isaaclab_physx/test/sensors/test_frame_transformer.py +++ b/source/isaaclab_physx/test/sensors/test_frame_transformer.py @@ -655,6 +655,8 @@ def _make_frame_transformer(use_recorded_launch: bool = True): sensor.cfg = SimpleNamespace(prim_path="/World/Source") sensor._device = device sensor._num_envs = 1 + sensor._ALL_ENV_MASK = wp.ones(1, dtype=wp.bool, device=device) + sensor._reset_mask = wp.zeros(1, dtype=wp.bool, device=device) sensor._num_target_frames = 1 sensor._frame_physx_view = transform_view sensor._backend_manager = SimpleNamespace(update_kinematics=lambda: None) diff --git a/source/isaaclab_physx/test/sensors/test_imu_pva_recorded_launch.py b/source/isaaclab_physx/test/sensors/test_imu_pva_recorded_launch.py index ea2a1b6df989..80c2956113de 100644 --- a/source/isaaclab_physx/test/sensors/test_imu_pva_recorded_launch.py +++ b/source/isaaclab_physx/test/sensors/test_imu_pva_recorded_launch.py @@ -122,6 +122,8 @@ def _make_sensor(sensor_type: str, use_recorded_launch: bool = True): sensor._device = device sensor._backend_manager = SimpleNamespace(update_kinematics=lambda: None) sensor._num_envs = 2 + sensor._ALL_ENV_MASK = wp.ones(2, dtype=wp.bool, device=device) + sensor._reset_mask = wp.zeros(2, dtype=wp.bool, device=device) sensor._view = rigid_view sensor._offset_pos_b = wp.zeros(2, dtype=wp.vec3f, device=device) sensor._offset_quat_b = wp.array(