diff --git a/scripts/run_ovphysx.sh b/scripts/run_ovphysx.sh index 72d18a1bec82..c3bb76c2d581 100755 --- a/scripts/run_ovphysx.sh +++ b/scripts/run_ovphysx.sh @@ -3,6 +3,14 @@ # Use when ovphysx is installed into Kit's Python. # # Usage: ./scripts/run_ovphysx.sh [your_script.py or -m pytest ...] +# +# CI note: the OVPhysX wheel's device mode is a process-global C++/Carbonite +# static (gap G5 in docs/superpowers/specs/2026-04-28-ovphysx-wheel-gaps-for-marco.md). +# To exercise both CPU and GPU coverage, invoke this script TWICE in separate +# processes -- e.g. +# ./scripts/run_ovphysx.sh -m pytest -k 'cpu' +# ./scripts/run_ovphysx.sh -m pytest -k 'cuda:0' +# A single invocation locks to whichever device is requested first. set -e ISAACLAB_PATH="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" diff --git a/source/isaaclab/changelog.d/antoiner-feat-ovphysx_rigidobject.skip b/source/isaaclab/changelog.d/antoiner-feat-ovphysx_rigidobject.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab/test/assets/test_articulation_iface.py b/source/isaaclab/test/assets/test_articulation_iface.py index 3e471026a9dc..56f69226e139 100644 --- a/source/isaaclab/test/assets/test_articulation_iface.py +++ b/source/isaaclab/test/assets/test_articulation_iface.py @@ -20,9 +20,13 @@ from unittest.mock import MagicMock # When running kitless (e.g., ovphysx backend via run_ovphysx.sh), AppLauncher -# will try to boot Kit and hang. Skip it entirely when LD_PRELOAD is cleared -# (the signature of run_ovphysx.sh) or when EXP_PATH is not set. -_kitless = os.environ.get("LD_PRELOAD", "") == "" and "EXP_PATH" not in os.environ +# will try to boot Kit and hang. Skip it entirely: run_ovphysx.sh sets +# LD_PRELOAD to the ovphysx libcarb.so, which is the signature of a kitless +# ovphysx run. Also guard the case where neither LD_PRELOAD nor EXP_PATH is +# set (bare Python, no Kit at all). +_kitless = "ovphysx" in os.environ.get("LD_PRELOAD", "") or ( + os.environ.get("LD_PRELOAD", "") == "" and "EXP_PATH" not in os.environ +) if not _kitless: from isaaclab.app import AppLauncher @@ -30,6 +34,19 @@ simulation_app = AppLauncher(headless=True).app else: simulation_app = None + # Stub out the Kit/Omniverse modules that are not present under + # run_ovphysx.sh (pxr, carb, omni, omni.kit[.app] are real on PYTHONPATH). + # ``omni`` is a real namespace package, so missing submodules also need + # to be installed as attributes on it -- ``sys.modules`` alone is not + # enough because attribute access on the real ``omni`` won't fall + # through to ``sys.modules``. + import omni as _omni + + for _mod in ("physics", "physics.tensors", "physx", "timeline", "usd"): + _stub = MagicMock() + sys.modules[f"omni.{_mod}"] = _stub + # Bind the leaf attribute so that ``omni.`` resolves. + setattr(_omni, _mod.split(".", 1)[0], _stub) for _mod in ("isaacsim.core", "isaacsim.core.simulation_manager"): sys.modules.setdefault(_mod, MagicMock()) @@ -271,20 +288,27 @@ def create_ovphysx_articulation( object.__setattr__(articulation, "_num_spatial_tendons", num_spatial_tendons) # Create ArticulationData - data = OvPhysxArticulationData(mock_bindings.bindings, device) - data._num_instances = num_instances - data._num_joints = num_joints - data._num_bodies = num_bodies - data._num_fixed_tendons = num_fixed_tendons - data._num_spatial_tendons = num_spatial_tendons + data = OvPhysxArticulationData( + mock_bindings.bindings, + device, + num_instances=num_instances, + num_bodies=num_bodies, + num_joints=num_joints, + num_fixed_tendons=num_fixed_tendons, + num_spatial_tendons=num_spatial_tendons, + body_names=body_names, + joint_names=joint_names, + fixed_tendon_names=fixed_tendon_names, + spatial_tendon_names=spatial_tendon_names, + ) data._is_fixed_base = False - data.body_names = body_names - data.joint_names = joint_names - data.fixed_tendon_names = fixed_tendon_names - data.spatial_tendon_names = spatial_tendon_names - data._create_buffers() object.__setattr__(articulation, "_data", data) + # Allocate the articulation-side index/mask caches and wrench buffer that + # _initialize_impl would normally populate. Wrench composers created here + # are immediately overwritten by the mocks below. + articulation._create_buffers() + # Wrench composers mock_inst_wrench = MockWrenchComposer(articulation) mock_perm_wrench = MockWrenchComposer(articulation) diff --git a/source/isaaclab/test/assets/test_rigid_object_collection_iface.py b/source/isaaclab/test/assets/test_rigid_object_collection_iface.py index 42c2e8d731b5..fb848baf34c8 100644 --- a/source/isaaclab/test/assets/test_rigid_object_collection_iface.py +++ b/source/isaaclab/test/assets/test_rigid_object_collection_iface.py @@ -14,16 +14,42 @@ The setup is a bit convoluted so that we can run these tests without requiring Isaac Sim or GPU simulation. """ -"""Launch Isaac Sim Simulator first.""" +"""Launch Isaac Sim Simulator first (when available).""" -from isaaclab.app import AppLauncher - -HEADLESS = True +import os +import sys +from unittest.mock import MagicMock -# launch omniverse app -simulation_app = AppLauncher(headless=True).app +# When running kitless (e.g., ovphysx backend via run_ovphysx.sh), AppLauncher +# will try to boot Kit and hang. Skip it entirely: run_ovphysx.sh sets +# LD_PRELOAD to the ovphysx libcarb.so, which is the signature of a kitless +# ovphysx run. Also guard the case where neither LD_PRELOAD nor EXP_PATH is +# set (bare Python, no Kit at all). +_kitless = "ovphysx" in os.environ.get("LD_PRELOAD", "") or ( + os.environ.get("LD_PRELOAD", "") == "" and "EXP_PATH" not in os.environ +) -from unittest.mock import MagicMock +if not _kitless: + from isaaclab.app import AppLauncher + + simulation_app = AppLauncher(headless=True).app +else: + simulation_app = None + # Stub out the Kit/Omniverse modules that are not present under + # run_ovphysx.sh (pxr, carb, omni, omni.kit[.app] are real on PYTHONPATH). + # ``omni`` is a real namespace package, so missing submodules also need + # to be installed as attributes on it -- ``sys.modules`` alone is not + # enough because attribute access on the real ``omni`` won't fall + # through to ``sys.modules``. + import omni as _omni + + for _mod in ("physics", "physics.tensors", "physx", "timeline", "usd"): + _stub = MagicMock() + sys.modules[f"omni.{_mod}"] = _stub + # Bind the leaf attribute so that ``omni.`` resolves. + setattr(_omni, _mod.split(".", 1)[0], _stub) + for _mod in ("isaacsim.core", "isaacsim.core.simulation_manager"): + sys.modules.setdefault(_mod, MagicMock()) import numpy as np import pytest @@ -75,6 +101,23 @@ except ImportError: pass +try: + from isaaclab_ovphysx.assets.rigid_object_collection.rigid_object_collection import ( + RigidObjectCollection as OvPhysxRigidObjectCollection, + ) + from isaaclab_ovphysx.assets.rigid_object_collection.rigid_object_collection_data import ( + RigidObjectCollectionData as OvPhysxRigidObjectCollectionData, + ) + from isaaclab_ovphysx.test.mock_interfaces.views import MockOvPhysxBindingSet + + # Guard against stub implementations (not yet functional). + if not hasattr(OvPhysxRigidObjectCollection, "_create_buffers"): + raise AttributeError("OvPhysxRigidObjectCollection is a stub; skipping ovphysx backend") + + BACKENDS.append("ovphysx") +except (ImportError, AttributeError): + pass + def create_physx_rigid_object_collection( num_instances: int = 2, @@ -209,6 +252,61 @@ def create_newton_rigid_object_collection( return collection, mock_view +def create_ovphysx_rigid_object_collection( + num_instances: int = 2, + num_bodies: int = 3, + device: str = "cuda:0", +): + """Create a test OVPhysX RigidObjectCollection instance with mocked tensor bindings.""" + body_names = [f"object_{i}" for i in range(num_bodies)] + + collection = object.__new__(OvPhysxRigidObjectCollection) + + rigid_objects = {f"object_{i}": RigidObjectCfg(prim_path=f"/World/Object_{i}") for i in range(num_bodies)} + collection.cfg = RigidObjectCollectionCfg(rigid_objects=rigid_objects) + + # Use articulation-mode bindings with num_joints=0 to get (N, B, ...) shaped tensors. + mock_bindings = MockOvPhysxBindingSet( + num_instances=num_instances, + num_joints=0, + num_bodies=num_bodies, + body_names=body_names, + asset_kind="articulation", + ) + mock_bindings.set_random_data() + + object.__setattr__(collection, "_device", device) + object.__setattr__(collection, "_ovphysx", MagicMock()) + object.__setattr__(collection, "_bindings", mock_bindings.bindings) + object.__setattr__(collection, "_num_instances", num_instances) + object.__setattr__(collection, "_num_bodies", num_bodies) + object.__setattr__(collection, "_body_names_list", body_names) + + # Create RigidObjectCollectionData + data = OvPhysxRigidObjectCollectionData(mock_bindings.bindings, num_bodies, device) + data.num_instances = num_instances + data.num_bodies = num_bodies + data._is_primed = True + object.__setattr__(collection, "_data", data) + + # Allocate the buffers that RigidObjectCollection normally allocates in _initialize_impl. + collection._create_buffers() + + # Replace the real wrench composers with mocks for iface coverage. + mock_inst_wrench = MockWrenchComposer(collection) + mock_perm_wrench = MockWrenchComposer(collection) + object.__setattr__(collection, "_instantaneous_wrench_composer", mock_inst_wrench) + object.__setattr__(collection, "_permanent_wrench_composer", mock_perm_wrench) + + # Prevent __del__ / _clear_callbacks from raising + object.__setattr__(collection, "_initialize_handle", None) + object.__setattr__(collection, "_invalidate_initialize_handle", None) + object.__setattr__(collection, "_prim_deletion_handle", None) + object.__setattr__(collection, "_debug_vis_handle", None) + + return collection, mock_bindings + + def create_mock_rigid_object_collection( num_instances: int = 2, num_bodies: int = 3, @@ -232,6 +330,8 @@ def get_rigid_object_collection( ): if backend == "physx": return create_physx_rigid_object_collection(num_instances, num_bodies, device) + elif backend == "ovphysx": + return create_ovphysx_rigid_object_collection(num_instances, num_bodies, device) elif backend == "newton": return create_newton_rigid_object_collection(num_instances, num_bodies, device) elif backend.lower() == "mock": diff --git a/source/isaaclab/test/assets/test_rigid_object_iface.py b/source/isaaclab/test/assets/test_rigid_object_iface.py index 178feeddb603..772130149ee3 100644 --- a/source/isaaclab/test/assets/test_rigid_object_iface.py +++ b/source/isaaclab/test/assets/test_rigid_object_iface.py @@ -13,16 +13,42 @@ The setup is a bit convoluted so that we can run these tests without requiring Isaac Sim or GPU simulation. """ -"""Launch Isaac Sim Simulator first.""" +"""Launch Isaac Sim Simulator first (when available).""" -from isaaclab.app import AppLauncher - -HEADLESS = True +import os +import sys +from unittest.mock import MagicMock -# launch omniverse app -simulation_app = AppLauncher(headless=True).app +# When running kitless (e.g., ovphysx backend via run_ovphysx.sh), AppLauncher +# will try to boot Kit and hang. Skip it entirely: run_ovphysx.sh sets +# LD_PRELOAD to the ovphysx libcarb.so, which is the signature of a kitless +# ovphysx run. Also guard the case where neither LD_PRELOAD nor EXP_PATH is +# set (bare Python, no Kit at all). +_kitless = "ovphysx" in os.environ.get("LD_PRELOAD", "") or ( + os.environ.get("LD_PRELOAD", "") == "" and "EXP_PATH" not in os.environ +) -from unittest.mock import MagicMock +if not _kitless: + from isaaclab.app import AppLauncher + + simulation_app = AppLauncher(headless=True).app +else: + simulation_app = None + # Stub out the Kit/Omniverse modules that are not present under + # run_ovphysx.sh (pxr, carb, omni, omni.kit[.app] are real on PYTHONPATH). + # ``omni`` is a real namespace package, so missing submodules also need + # to be installed as attributes on it -- ``sys.modules`` alone is not + # enough because attribute access on the real ``omni`` won't fall + # through to ``sys.modules``. + import omni as _omni + + for _mod in ("physics", "physics.tensors", "physx", "timeline", "usd"): + _stub = MagicMock() + sys.modules[f"omni.{_mod}"] = _stub + # Bind the leaf attribute so that ``omni.`` resolves. + setattr(_omni, _mod.split(".", 1)[0], _stub) + for _mod in ("isaacsim.core", "isaacsim.core.simulation_manager"): + sys.modules.setdefault(_mod, MagicMock()) import numpy as np import pytest @@ -66,6 +92,15 @@ except ImportError: pass +try: + from isaaclab_ovphysx.assets.rigid_object.rigid_object import RigidObject as OvPhysxRigidObject + from isaaclab_ovphysx.assets.rigid_object.rigid_object_data import RigidObjectData as OvPhysxRigidObjectData + from isaaclab_ovphysx.test.mock_interfaces.views import MockOvPhysxBindingSet + + BACKENDS.append("ovphysx") +except (ImportError, AttributeError): + pass + def create_physx_rigid_object( num_instances: int = 2, @@ -206,6 +241,62 @@ def create_newton_rigid_object( return rigid_object, mock_view +def create_ovphysx_rigid_object( + num_instances: int = 2, + device: str = "cuda:0", +): + """Create a test OvPhysX RigidObject instance with mocked tensor bindings.""" + body_names = ["base_link"] + + obj = object.__new__(OvPhysxRigidObject) + + obj.cfg = RigidObjectCfg(prim_path="/World/object") + + # Create mock binding set + mock_bindings = MockOvPhysxBindingSet( + num_instances=num_instances, + num_joints=0, + num_bodies=1, + body_names=body_names, + asset_kind="rigid_object", + ) + mock_bindings.set_random_data() + + object.__setattr__(obj, "_device", device) + object.__setattr__(obj, "_ovphysx", MagicMock()) + object.__setattr__(obj, "_bindings", mock_bindings.bindings) + object.__setattr__(obj, "_num_instances", num_instances) + object.__setattr__(obj, "_num_bodies", 1) + object.__setattr__(obj, "_body_names", body_names) + + # Create RigidObjectData + data = OvPhysxRigidObjectData(mock_bindings.bindings, device) + data.num_instances = num_instances + data.num_bodies = 1 + data._is_primed = True + object.__setattr__(obj, "_data", data) + + # Build the buffers RigidObject normally allocates in _initialize_impl + # (_ALL_INDICES, _ALL_*_MASK, pinned CPU staging buffers, wrench buf). + # _create_buffers also instantiates real WrenchComposers; those get + # replaced with mocks just below. + obj._create_buffers() + + # Replace the real wrench composers with mocks for iface coverage. + mock_inst_wrench = MockWrenchComposer(obj) + mock_perm_wrench = MockWrenchComposer(obj) + object.__setattr__(obj, "_instantaneous_wrench_composer", mock_inst_wrench) + object.__setattr__(obj, "_permanent_wrench_composer", mock_perm_wrench) + + # Prevent __del__ / _clear_callbacks from raising + object.__setattr__(obj, "_initialize_handle", None) + object.__setattr__(obj, "_invalidate_initialize_handle", None) + object.__setattr__(obj, "_prim_deletion_handle", None) + object.__setattr__(obj, "_debug_vis_handle", None) + + return obj, mock_bindings + + def create_mock_rigid_object( num_instances: int = 2, device: str = "cuda:0", @@ -226,6 +317,8 @@ def get_rigid_object( ): if backend == "physx": return create_physx_rigid_object(num_instances, device) + elif backend == "ovphysx": + return create_ovphysx_rigid_object(num_instances, device) elif backend == "newton": return create_newton_rigid_object(num_instances, device) elif backend.lower() == "mock": diff --git a/source/isaaclab_ovphysx/changelog.d/antoiner-feat-ovphysx_articulation.major.rst b/source/isaaclab_ovphysx/changelog.d/antoiner-feat-ovphysx_articulation.major.rst new file mode 100644 index 000000000000..24913760aa9b --- /dev/null +++ b/source/isaaclab_ovphysx/changelog.d/antoiner-feat-ovphysx_articulation.major.rst @@ -0,0 +1,109 @@ +Added +^^^^^ + +* Added :class:`~isaaclab_ovphysx.assets.Articulation` and + :class:`~isaaclab_ovphysx.assets.ArticulationData` for multi-DOF articulated + robots against the OVPhysX backend, satisfying the + :class:`~isaaclab.assets.BaseArticulation` and + :class:`~isaaclab.assets.BaseArticulationData` contracts. Public surface + matches the PhysX/Newton conventions: kwarg-only ``write_*_to_sim_index`` / + ``write_*_to_sim_mask`` writers and ``set_*_index`` / ``set_*_mask`` setters + for root state, joint state, joint properties, body properties, joint + command targets, fixed/spatial tendon properties, and external wrenches via + :attr:`~isaaclab_ovphysx.assets.Articulation.instantaneous_wrench_composer` + / :attr:`~isaaclab_ovphysx.assets.Articulation.permanent_wrench_composer`. + The full IsaacLab actuator pipeline (``compute`` / + ``_apply_actuator_model`` / ``_process_actuators_cfg``) is implemented on + top of the wheel's ``DOF_ACTUATION_FORCE`` / + ``DOF_POSITION_TARGET`` / ``DOF_VELOCITY_TARGET`` bindings. +* Added articulation-specific Warp kernels in + :mod:`isaaclab_ovphysx.assets.articulation.kernels`: soft-limit refresh, + default-joint-pos clamp, friction-component scatter (index + mask + variants). Six articulation kernels were also folded into the shared + :mod:`isaaclab_ovphysx.assets.kernels` module for reuse with + :class:`~isaaclab_ovphysx.assets.RigidObject` and + :class:`~isaaclab_ovphysx.assets.RigidObjectCollection`. +* Added init-time validation in + :meth:`~isaaclab_ovphysx.assets.Articulation._initialize_impl` that raises + ``RuntimeError`` when ``cfg.prim_path`` resolves to no + ``UsdPhysics.ArticulationRootAPI`` prim or to multiple roots, and + ``ValueError`` (via :meth:`_validate_cfg`) when any default joint + position is outside ``[lower, upper]`` or any default joint velocity + exceeds the per-joint maximum. Mirrors the PhysX backend. +* Added support for ``cfg.articulation_root_prim_path`` in + :meth:`~isaaclab_ovphysx.assets.Articulation._initialize_impl`: when the + user supplies an explicit subpath the binding pattern is extended + directly instead of running the auto-discovery walk, and a + ``RuntimeError`` is raised when the resulting expression resolves to no + prim in the USD stage. + +Changed +^^^^^^^ + +* **Breaking:** Renamed ``Articulation`` write/set methods to the dual + ``*_index`` / ``*_mask`` form and dropped the legacy ``full_data`` + flag. Index methods accept partial data shaped + ``(len(env_ids), len(joint_or_body_ids), ...)``; mask methods accept + full-shape data and a ``wp.bool`` mask. All keyword-only arguments live + after ``*,``; no positional fall-through. Migration: replace + ``write_X_to_sim(..., from_mask=True)`` with ``write_X_to_sim_mask(..., mask=...)``. +* **Breaking:** Removed the ``_write_body_state`` plumbing layer. + Deprecated state-writer shims (``write_root_state_to_sim``, + ``write_root_com_state_to_sim``, ``write_root_link_state_to_sim``, + joint-state equivalents) now call the public ``write_*_to_sim_index`` + methods directly. Behaviour is preserved. +* Changed ``Articulation.root_view`` to return the per-tensor-type bindings + dict (``self._bindings``). The OVPhysX wheel does not expose a single + ``ArticulationView`` object; callers that previously walked + ``root_view.shared_metatype`` / ``root_view.max_dofs`` should read from + :attr:`~isaaclab_ovphysx.assets.Articulation.num_joints` / + :attr:`~isaaclab_ovphysx.assets.Articulation.num_bodies` / + :attr:`~isaaclab_ovphysx.assets.Articulation.body_names` / + :attr:`~isaaclab_ovphysx.assets.Articulation.joint_names` instead. +* Changed every ``ArticulationData`` public property to return a + :class:`~isaaclab.utils.ProxyArray` (warp + torch dual view); raw + ``wp.array`` is reserved for one-shot config buffers. Eager + ``TimestampedBufferWarp`` allocation in :meth:`_create_buffers` makes + every buffer a single source of truth — no + ``_invalidate_caches`` / ``_ensure_*_buffers`` machinery. +* Changed ``Articulation`` body and DOF property writers to honor the + wheel's actual binding device. Tensor-type membership in + :data:`isaaclab_ovphysx.tensor_types._CPU_ONLY_TYPES` now reflects what + the wheel exposes: ``BODY_MASS``, ``BODY_COM_POSE``, ``BODY_INERTIA``, + ``DOF_STIFFNESS``, ``DOF_DAMPING``, ``DOF_LIMIT``, ``DOF_MAX_VELOCITY``, + ``DOF_MAX_FORCE``, ``DOF_ARMATURE``, ``DOF_FRICTION_PROPERTIES`` are + CPU-only (write goes through pinned-host staging); fixed and spatial + tendon bindings write directly from sim-device buffers. +* Changed :meth:`~isaaclab_ovphysx.assets.Articulation.write_joint_friction_coefficient_to_sim_index` + / ``_mask`` to accept ``joint_dynamic_friction_coeff`` and + ``joint_viscous_friction_coeff`` keyword arguments (each + ``float | torch.Tensor | wp.array | None``). ``None`` preserves the + existing component on the wheel; matches the PhysX backend. +* Changed :meth:`~isaaclab_ovphysx.assets.Articulation.write_joint_position_limit_to_sim_index` + / ``_mask`` to clamp ``default_joint_pos`` and refresh + ``soft_joint_pos_limits`` when the new hard limits invalidate the + defaults, matching the PhysX backend (with a + ``warn_limit_violation`` log). +* Changed every fixed/spatial tendon ``set_*_index`` / ``set_*_mask`` setter + to accept a scalar :class:`float` for the value argument; broadcast is + materialized via :meth:`_broadcast_scalar_to_2d`. Mirrors PhysX. +* Implemented the previously stubbed + :meth:`~isaaclab_ovphysx.assets.Articulation.write_fixed_tendon_properties_to_sim_index` + / ``_mask`` and + :meth:`~isaaclab_ovphysx.assets.Articulation.write_spatial_tendon_properties_to_sim_index` + / ``_mask``: each iterates the per-tensor bindings since the OVPhysX + wheel has no batch ``set_*_tendon_properties`` setter. + +Removed +^^^^^^^ + +* **Breaking:** Removed the ``full_data`` keyword-argument from every + ``Articulation`` ``*_index`` writer/setter. Index methods now strictly + accept partial data; full-data callers should use the matching + ``*_mask`` overload. +* Removed the stop-gap :mod:`isaaclab_ovphysx.assets.kernels_old` module; + the six articulation kernels it housed + (``_compose_root_com_pose``, ``_compute_heading``, ``_copy_first_body``, + ``_projected_gravity``, ``_world_vel_to_body_ang``, + ``_world_vel_to_body_lin``) are now in + :mod:`isaaclab_ovphysx.assets.kernels`. diff --git a/source/isaaclab_ovphysx/changelog.d/antoiner-feat-ovphysx_rigidobject.major.rst b/source/isaaclab_ovphysx/changelog.d/antoiner-feat-ovphysx_rigidobject.major.rst new file mode 100644 index 000000000000..2a43911295c2 --- /dev/null +++ b/source/isaaclab_ovphysx/changelog.d/antoiner-feat-ovphysx_rigidobject.major.rst @@ -0,0 +1,50 @@ +Added +^^^^^ + +* Added :class:`~isaaclab_ovphysx.assets.RigidObject` and + :class:`~isaaclab_ovphysx.assets.RigidObjectData` for single-actor rigid-body + simulation against the OVPhysX backend, satisfying the + :class:`~isaaclab.assets.BaseRigidObject` and + :class:`~isaaclab.assets.BaseRigidObjectData` contracts. Public surface + matches the PhysX/Newton conventions: ``write_root_*_to_sim_index`` / + ``write_root_*_to_sim_mask`` writers (link- and com-frame variants), + ``set_masses_*``, ``set_coms_*``, ``set_inertias_*`` setters, and the + external-wrench composers exposed via + :meth:`~isaaclab_ovphysx.assets.RigidObject.set_external_force_and_torque`. +* Added the ``RIGID_BODY_*`` :class:`TensorType` aliases in + :mod:`isaaclab_ovphysx.tensor_types` (``POSE``, ``VELOCITY``, ``WRENCH``, + ``MASS``, ``COM_POSE``, ``INERTIA``; plus ``ACCELERATION``, ``INV_MASS``, + ``INV_INERTIA`` declared for forward compatibility once the wheel ships + them). +* Added :class:`~isaaclab_ovphysx.assets.kernels` as a shared Warp-kernel + module (frame conversions, state concatenation, finite-difference + acceleration, index- and mask-style scatter writers) consumed by both the + rigid-object and articulation assets. +* Added USD prim-scan validation in + :meth:`~isaaclab_ovphysx.assets.RigidObject._initialize_impl`: a clear + ``RuntimeError`` is raised when ``cfg.prim_path`` resolves to no + ``UsdPhysics.RigidBodyAPI`` prim, multiple rigid-body prims, or a prim with + an enabled ``UsdPhysics.ArticulationRootAPI``. + +Changed +^^^^^^^ + +* Changed :meth:`~isaaclab_ovphysx.physics.OvPhysxManager._release_physx` to + perform a soft reset (``physx.reset()``) and keep the cached + :class:`ovphysx.PhysX` reference alive across + :class:`~isaaclab.sim.SimulationContext` lifetimes, instead of dropping the + reference and triggering the wheel's dual-Carbonite static-destructor race. + :meth:`~isaaclab_ovphysx.physics.OvPhysxManager._warmup_and_load` now reuses + the cached instance on subsequent calls. +* Changed :meth:`~isaaclab_ovphysx.physics.OvPhysxManager._warmup_and_load` to + raise a clear ``RuntimeError`` when a later + :class:`~isaaclab.sim.SimulationContext` requests a different device than + the one the process is locked to, surfacing the wheel's process-global + device-mode lock as a Python error before + :exc:`ovphysx.types.PhysXDeviceError` would fire. +* Changed :meth:`~isaaclab_ovphysx.physics.OvPhysxManager._configure_physx_scene_prim` + to apply the ``UsdPhysics.PhysxSceneAPI`` schema and + ``enableSceneQuerySupport`` on both CPU and GPU; GPU-only attributes + (``enableGPUDynamics``, ``broadphaseType``, the ``gpu*`` capacity attributes + from :class:`~isaaclab_ovphysx.physics.OvPhysxCfg`) remain gated on + ``device == "gpu"``. diff --git a/source/isaaclab_ovphysx/changelog.d/antoiner-feat-ovphysx_rigidobjectcollection.major.rst b/source/isaaclab_ovphysx/changelog.d/antoiner-feat-ovphysx_rigidobjectcollection.major.rst new file mode 100644 index 000000000000..460fdd789f28 --- /dev/null +++ b/source/isaaclab_ovphysx/changelog.d/antoiner-feat-ovphysx_rigidobjectcollection.major.rst @@ -0,0 +1,13 @@ +Added +^^^^^ + +* Added :class:`~isaaclab_ovphysx.assets.RigidObjectCollection` and + :class:`~isaaclab_ovphysx.assets.RigidObjectCollectionData` for the + OVPhysX backend, completing the rigid-body asset surface alongside + :class:`~isaaclab_ovphysx.assets.RigidObject` and + :class:`~isaaclab_ovphysx.assets.Articulation`. Supports + ``(env, body)`` dual indexing and per-body property setters. Uses the + ovphysx 0.4.3+ native fused multi-prim binding API + (``create_tensor_binding(prim_paths=[...])``) so one binding spans all + ``num_instances * num_bodies`` prims per tensor type, mirroring the + strided-view reshape pattern used by the PhysX collection. diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/__init__.pyi b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/__init__.pyi index 516e15c5ef6a..658ced63d2a0 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/__init__.pyi +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/__init__.pyi @@ -6,6 +6,12 @@ __all__ = [ "Articulation", "ArticulationData", + "RigidObject", + "RigidObjectCollection", + "RigidObjectCollectionData", + "RigidObjectData", ] from .articulation import Articulation, ArticulationData +from .rigid_object import RigidObject, RigidObjectData +from .rigid_object_collection import RigidObjectCollection, RigidObjectCollectionData diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py index bea4345ca5ba..86a398c5f3b1 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py @@ -3,50 +3,90 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Articulation implementation backed by ovphysx TensorBindingsAPI.""" +# Flag for pyright to ignore type errors in this file. +# pyright: reportPrivateUsage=false from __future__ import annotations import logging import re +import warnings from collections.abc import Sequence -from typing import TYPE_CHECKING, Any +from typing import Any import numpy as np import torch import warp as wp +from pxr import UsdPhysics + +import isaaclab.sim as sim_utils +from isaaclab.assets.articulation.articulation_cfg import ArticulationCfg from isaaclab.assets.articulation.base_articulation import BaseArticulation from isaaclab.physics import PhysicsManager from isaaclab.utils.string import resolve_matching_names from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_ovphysx import tensor_types as TT +from isaaclab_ovphysx.assets import kernels as shared_kernels +from isaaclab_ovphysx.assets.kernels import _body_wrench_to_world from isaaclab_ovphysx.physics import OvPhysxManager from .articulation_data import ArticulationData -from .kernels import _body_wrench_to_world, _scatter_rows_partial, update_soft_joint_pos_limits - -if TYPE_CHECKING: - from isaaclab.actuators import ActuatorBase - from isaaclab.assets.articulation.articulation_cfg import ArticulationCfg - +from .kernels import ( + clamp_default_joint_pos_and_update_soft_limits_index, + clamp_default_joint_pos_and_update_soft_limits_mask, + update_soft_joint_pos_limits, + write_joint_friction_data_to_buffer_index, + write_joint_friction_data_to_buffer_mask, +) + +# import logger logger = logging.getLogger(__name__) class Articulation(BaseArticulation): - """Articulation backed by the ovphysx TensorBindingsAPI. + """An articulation asset class. - Reads and writes simulation state through ovphysx.TensorBinding objects created - from the OvPhysxManager's PhysX instance. - """ + An articulation is a collection of rigid bodies connected by joints. The joints can be either + fixed or actuated. The joints can be of different types, such as revolute, prismatic, D-6, etc. + However, the articulation class has currently been tested with revolute and prismatic joints. + The class supports both floating-base and fixed-base articulations. The type of articulation + is determined based on the root joint of the articulation. If the root joint is fixed, then + the articulation is considered a fixed-base system. Otherwise, it is considered a floating-base + system. This can be checked using the :attr:`Articulation.is_fixed_base` attribute. - __backend_name__ = "ovphysx" + For an asset to be considered an articulation, the root prim of the asset must have the + `USD ArticulationRootAPI`_. This API is used to define the sub-tree of the articulation using + the reduced coordinate formulation. On playing the simulation, the physics engine parses the + articulation root prim and creates the corresponding articulation in the physics engine. The + articulation root prim can be specified using the :attr:`AssetBaseCfg.prim_path` attribute. + + OVPhysX exposes per-tensor-type :class:`ovphysx.TensorBinding` objects rather than a single + opaque view; binding handles are created eagerly in :meth:`_initialize_impl` and reused across + reads and writes. CPU-only bindings (mass, CoM, inertia, joint properties, tendon properties) + are routed through pinned-host staging buffers managed by :class:`ArticulationData`. + + .. _`USD ArticulationRootAPI`: https://openusd.org/dev/api/class_usd_physics_articulation_root_a_p_i.html + + """ cfg: ArticulationCfg + """Configuration instance for the articulation.""" + + __backend_name__: str = "ovphysx" + """The name of the backend for the articulation.""" def __init__(self, cfg: ArticulationCfg): + """Initialize the articulation. + + Args: + cfg: A configuration instance. + """ super().__init__(cfg) + # bindings are populated eagerly in ``_initialize_impl``; the dict + # also caches any tensor type the user explicitly queries later + self._bindings: dict[int, Any] = {} """ Properties @@ -54,12 +94,10 @@ def __init__(self, cfg: ArticulationCfg): @property def data(self) -> ArticulationData: - """Data container with simulation state for this articulation.""" return self._data @property def num_instances(self) -> int: - """Number of articulation instances (environments).""" return self._num_instances @property @@ -69,64 +107,88 @@ def is_fixed_base(self) -> bool: @property def num_joints(self) -> int: - """Number of joints in the articulation.""" + """Number of joints in articulation.""" return self._num_joints @property def num_fixed_tendons(self) -> int: - """Number of fixed tendons in the articulation.""" - return getattr(self, "_num_fixed_tendons", 0) + """Number of fixed tendons in articulation.""" + return self._num_fixed_tendons @property def num_spatial_tendons(self) -> int: - """Number of spatial tendons in the articulation.""" - return getattr(self, "_num_spatial_tendons", 0) + """Number of spatial tendons in articulation.""" + return self._num_spatial_tendons @property def num_bodies(self) -> int: - """Number of bodies (links) in the articulation.""" + """Number of bodies in articulation.""" return self._num_bodies @property def joint_names(self) -> list[str]: - """Ordered names of joints in the articulation.""" + """Ordered names of joints in articulation.""" return self._joint_names @property def fixed_tendon_names(self) -> list[str]: - """Ordered names of fixed tendons in the articulation.""" - return getattr(self, "_fixed_tendon_names", []) + """Ordered names of fixed tendons in articulation.""" + return self._fixed_tendon_names @property def spatial_tendon_names(self) -> list[str]: - """Ordered names of spatial tendons in the articulation.""" - return getattr(self, "_spatial_tendon_names", []) + """Ordered names of spatial tendons in articulation.""" + return self._spatial_tendon_names @property def body_names(self) -> list[str]: - """Ordered names of bodies (links) in the articulation.""" + """Ordered names of bodies in articulation.""" return self._body_names @property - def root_view(self) -> Any: - """Root articulation view (not available for ovphysx backend).""" - return None + def root_view(self) -> dict[int, Any]: + """Root view for the asset. + + OVPhysX exposes per-tensor-type bindings rather than a single opaque view object + as used by the PhysX and Newton backends. Callers that need low-level binding + access should call :meth:`_get_binding` rather than iterating this dict directly. + For high-level state access (instance counts, prim paths, transforms), use the + :attr:`num_instances`, :attr:`body_names`, and :attr:`~ArticulationData.root_link_pose_w` + accessors instead. + + .. note:: + Use this view with caution. It requires handling of tensors in a specific way. + """ + return self._bindings @property - def instantaneous_wrench_composer(self) -> WrenchComposer | None: - """Wrench composer for forces applied only during the current step.""" + def instantaneous_wrench_composer(self) -> WrenchComposer: + """Instantaneous wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are only valid for the current simulation step. At the end of the simulation step, the wrenches set + to this object are discarded. This is useful to apply forces that change all the time, things like drag forces + for instance. + """ return self._instantaneous_wrench_composer @property - def permanent_wrench_composer(self) -> WrenchComposer | None: - """Wrench composer for forces applied persistently every step.""" + def permanent_wrench_composer(self) -> WrenchComposer: + """Permanent wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are persistent and are applied to the simulation at every step. This is useful to apply forces that + are constant over a period of time, things like the thrust of a motor for instance. + """ return self._permanent_wrench_composer """ Operations. """ - def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None) -> None: + def reset( + self, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_mask: wp.array | None = None + ) -> None: """Reset the articulation. .. caution:: @@ -136,26 +198,55 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None env_ids: Environment indices. If None, then all indices are used. env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - # use ellipses object to skip initial indices. if (env_ids is None) or (env_ids == slice(None)): env_ids = slice(None) - # reset actuators - for actuator in self.actuators.values(): - actuator.reset(env_ids) # reset external wrenches. self._instantaneous_wrench_composer.reset(env_ids, env_mask) self._permanent_wrench_composer.reset(env_ids, env_mask) def write_data_to_sim(self) -> None: - """Apply external wrenches, actuator model, and write commands into the simulation.""" - # Apply external wrenches (before actuators, same as PhysX backend). - self._apply_external_wrenches() + """Write external wrenches and joint commands to the simulation. + + If any explicit actuators are present, then the actuator models are used to compute the + joint commands. Otherwise, the joint commands are directly set into the simulation. + + .. note:: + We write external wrench to the simulation here since this function is called before the simulation step. + This ensures that the external wrench is applied at every simulation step. + """ + # write external wrench + inst = self._instantaneous_wrench_composer + perm = self._permanent_wrench_composer + if inst.active or perm.active: + if inst.active: + if perm.active: + inst.add_raw_buffers_from(perm) + force_b = inst.out_force_b.warp + torque_b = inst.out_torque_b.warp + else: + force_b = perm.out_force_b.warp + torque_b = perm.out_torque_b.warp + + # rotate body-frame wrenches into the world frame expected by ``LINK_WRENCH`` + poses = self._data.body_link_pose_w.warp + wp.launch( + _body_wrench_to_world, + dim=(self._num_instances, self._num_bodies), + inputs=[force_b, torque_b, poses], + outputs=[self._wrench_buf], + device=self._device, + ) + binding = self._get_binding(TT.LINK_WRENCH) + if binding is not None: + binding.write(self._wrench_buf) + inst.reset() + # apply actuator models self._apply_actuator_model() - # Write effort tensor to simulation. + # write actions into simulation (zeros are safe when no actuators are active) if self._effort_binding is not None: self._effort_binding.write(self._effort_write_view) - # Write position and velocity targets in one shot (not per-actuator). + # position and velocity targets only for implicit actuators if self._has_implicit_actuators: if self._pos_target_binding is not None: self._pos_target_binding.write(self._pos_target_write_view) @@ -163,10 +254,10 @@ def write_data_to_sim(self) -> None: self._vel_target_binding.write(self._vel_target_write_view) def update(self, dt: float) -> None: - """Update internal data buffers after a simulation step. + """Updates the simulation data. Args: - dt: The simulation time step [s] used for finite-difference quantities. + dt: The time step size in seconds. """ self._data.update(dt) @@ -187,7 +278,7 @@ def find_bodies(self, name_keys: str | Sequence[str], preserve_order: bool = Fal Returns: A tuple of lists containing the body indices and names. """ - return resolve_matching_names(name_keys, self._body_names, preserve_order) + return resolve_matching_names(name_keys, self.body_names, preserve_order) def find_joints( self, @@ -197,8 +288,8 @@ def find_joints( ) -> tuple[list[int], list[str]]: """Find joints in the articulation based on the name keys. - Please check the :func:`isaaclab.utils.string.resolve_matching_names` function for more - information on the name matching. + Please see the :func:`isaaclab.utils.string.resolve_matching_names` function for more information + on the name matching. Args: name_keys: A regular expression or a list of regular expressions to match the joint names. @@ -210,7 +301,8 @@ def find_joints( A tuple of lists containing the joint indices and names. """ if joint_subset is None: - joint_subset = self._joint_names + joint_subset = self.joint_names + # find joints return resolve_matching_names(name_keys, joint_subset, preserve_order) def find_fixed_tendons( @@ -221,20 +313,23 @@ def find_fixed_tendons( ) -> tuple[list[int], list[str]]: """Find fixed tendons in the articulation based on the name keys. + Please see the :func:`isaaclab.utils.string.resolve_matching_names` function for more information + on the name matching. + Args: - name_keys: A regular expression or a list of regular expressions - to match the joint names with fixed tendons. - tendon_subsets: A subset of joints with fixed tendons to search - for. Defaults to None, which means all joints in the - articulation are searched. - preserve_order: Whether to preserve the order of the name keys in - the output. Defaults to False. + name_keys: A regular expression or a list of regular expressions to match the + joint names with fixed tendons. + tendon_subsets: A subset of joints with fixed tendons to search for. Defaults to None, which means + all joints in the articulation are searched. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. Returns: A tuple of lists containing the tendon indices and names. """ if tendon_subsets is None: + # tendons follow the joint names they are attached to tendon_subsets = self.fixed_tendon_names + # find tendons return resolve_matching_names(name_keys, tendon_subsets, preserve_order) def find_spatial_tendons( @@ -245,19 +340,21 @@ def find_spatial_tendons( ) -> tuple[list[int], list[str]]: """Find spatial tendons in the articulation based on the name keys. + Please see the :func:`isaaclab.utils.string.resolve_matching_names` function for more information + on the name matching. + Args: - name_keys: A regular expression or a list of regular expressions - to match the tendon names. - tendon_subsets: A subset of tendons to search for. Defaults to - None, which means all tendons in the articulation are searched. - preserve_order: Whether to preserve the order of the name keys in - the output. Defaults to False. + name_keys: A regular expression or a list of regular expressions to match the tendon names. + tendon_subsets: A subset of tendons to search for. Defaults to None, which means all tendons + in the articulation are searched. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. Returns: A tuple of lists containing the tendon indices and names. """ if tendon_subsets is None: tendon_subsets = self.spatial_tendon_names + # find tendons return resolve_matching_names(name_keys, tendon_subsets, preserve_order) """ @@ -268,19 +365,25 @@ def write_root_pose_to_sim_index( self, *, root_pose: torch.Tensor | wp.array, - env_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set the root pose over selected environment indices into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + Args: - root_pose: Root poses in simulation frame. Shape is (len(env_ids),) with dtype wp.transformf. + root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7) + or (len(env_ids),) with dtype wp.transformf. env_ids: Environment indices. If None, then all indices are used. """ - n = self._n_envs_index(env_ids) - self.assert_shape_and_dtype(root_pose, (n,), wp.transformf, "root_pose") - self._write_root_state(TT.ROOT_POSE, root_pose, env_ids) + self.write_root_link_pose_to_sim_index(root_pose=root_pose, env_ids=env_ids) def write_root_pose_to_sim_mask( self, @@ -288,34 +391,61 @@ def write_root_pose_to_sim_mask( root_pose: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root pose over masked environments into the simulation. + """Set the root pose over selected environment mask into the simulation. - The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. Args: - root_pose: Root poses in simulation frame. Shape is (num_instances,) with dtype wp.transformf. - env_mask: Environment mask. If None, then all instances are updated. + root_pose: Root poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - self.assert_shape_and_dtype(root_pose, (self._num_instances,), wp.transformf, "root_pose") - self._write_root_state(TT.ROOT_POSE, root_pose, mask=env_mask) + self.write_root_link_pose_to_sim_mask(root_pose=root_pose, env_mask=env_mask) def write_root_link_pose_to_sim_index( self, *, root_pose: torch.Tensor | wp.array, - env_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set the root link pose over selected environment indices into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + Args: - root_pose: Root link poses in simulation frame. Shape is (len(env_ids),) with dtype wp.transformf. + root_pose: Root link poses in simulation frame. Shape is (len(env_ids), 7) + or (len(env_ids),) with dtype wp.transformf. env_ids: Environment indices. If None, then all indices are used. """ - n = self._n_envs_index(env_ids) - self.assert_shape_and_dtype(root_pose, (n,), wp.transformf, "root_pose") - self._write_root_state(TT.ROOT_POSE, root_pose, env_ids) + env_ids = self._resolve_env_ids(env_ids) + self.assert_shape_and_dtype(root_pose, (env_ids.shape[0],), wp.transformf, "root_pose") + wp.launch( + shared_kernels.set_root_link_pose_to_sim_index, + dim=env_ids.shape[0], + inputs=[root_pose, env_ids], + outputs=[self.data.root_link_pose_w], + device=self._device, + ) + # invalidate dependent timestamps: root link pose changes the body + # kinematics chain, so all body-pose buffers go stale + self.data._root_com_pose_w.timestamp = -1.0 + self.data._body_link_pose_w.timestamp = -1.0 + self.data._body_com_pose_w.timestamp = -1.0 + # push cache to the simulation via an indexed write + binding = self._get_binding(TT.ROOT_POSE) + binding.write(self.data._root_link_pose_w.data.view(wp.float32), indices=env_ids) def write_root_link_pose_to_sim_mask( self, @@ -323,36 +453,76 @@ def write_root_link_pose_to_sim_mask( root_pose: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root link pose over masked environments into the simulation. + """Set the root link pose over selected environment mask into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + Args: - root_pose: Root link poses in simulation frame. Shape is (num_instances,) with dtype wp.transformf. - env_mask: Environment mask. If None, then all instances are updated. + root_pose: Root poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ + env_mask_wp = self._resolve_env_mask(env_mask) self.assert_shape_and_dtype(root_pose, (self._num_instances,), wp.transformf, "root_pose") - self._write_root_state(TT.ROOT_POSE, root_pose, mask=env_mask) + wp.launch( + shared_kernels.set_root_link_pose_to_sim_mask, + dim=self._num_instances, + inputs=[root_pose, env_mask_wp], + outputs=[self.data.root_link_pose_w], + device=self._device, + ) + # invalidate dependent timestamps (see :meth:`write_root_link_pose_to_sim_index`) + self.data._root_com_pose_w.timestamp = -1.0 + self.data._body_link_pose_w.timestamp = -1.0 + self.data._body_com_pose_w.timestamp = -1.0 + binding = self._get_binding(TT.ROOT_POSE) + binding.write(self.data._root_link_pose_w.data.view(wp.float32), mask=env_mask_wp) def write_root_com_pose_to_sim_index( self, *, root_pose: torch.Tensor | wp.array, - env_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set the root center of mass pose over selected environment indices into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). The orientation is the orientation of the principal axes of inertia. + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + Args: - root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids),) - with dtype wp.transformf. + root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7) + or (len(env_ids),) with dtype wp.transformf. env_ids: Environment indices. If None, then all indices are used. """ - n = self._n_envs_index(env_ids) - self.assert_shape_and_dtype(root_pose, (n,), wp.transformf, "root_pose") - self._write_root_state(TT.ROOT_POSE, root_pose, env_ids) + env_ids = self._resolve_env_ids(env_ids) + self.assert_shape_and_dtype(root_pose, (env_ids.shape[0],), wp.transformf, "root_pose") + wp.launch( + shared_kernels.set_root_com_pose_to_sim_index, + dim=env_ids.shape[0], + inputs=[root_pose, self.data.body_com_pose_b, env_ids], + outputs=[self.data.root_com_pose_w, self.data.root_link_pose_w], + device=self._device, + ) + # writing the root CoM pose updates the inferred root link pose, which + # in turn invalidates the body kinematics chain + self.data._body_link_pose_w.timestamp = -1.0 + self.data._body_com_pose_w.timestamp = -1.0 + binding = self._get_binding(TT.ROOT_POSE) + binding.write(self.data._root_link_pose_w.data.view(wp.float32), indices=env_ids) def write_root_com_pose_to_sim_mask( self, @@ -360,37 +530,64 @@ def write_root_com_pose_to_sim_mask( root_pose: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root center of mass pose over masked environments into the simulation. + """Set the root center of mass pose over selected environment mask into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). The orientation is the orientation of the principal axes of inertia. + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + Args: - root_pose: Root center of mass poses in simulation frame. Shape is (num_instances,) - with dtype wp.transformf. - env_mask: Environment mask. If None, then all instances are updated. + root_pose: Root center of mass poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ + env_mask_wp = self._resolve_env_mask(env_mask) self.assert_shape_and_dtype(root_pose, (self._num_instances,), wp.transformf, "root_pose") - self._write_root_state(TT.ROOT_POSE, root_pose, mask=env_mask) + wp.launch( + shared_kernels.set_root_com_pose_to_sim_mask, + dim=self._num_instances, + inputs=[root_pose, self.data.body_com_pose_b, env_mask_wp], + outputs=[self.data.root_com_pose_w, self.data.root_link_pose_w], + device=self._device, + ) + # invalidate dependent timestamps (see :meth:`write_root_com_pose_to_sim_index`) + self.data._body_link_pose_w.timestamp = -1.0 + self.data._body_com_pose_w.timestamp = -1.0 + binding = self._get_binding(TT.ROOT_POSE) + binding.write(self.data._root_link_pose_w.data.view(wp.float32), mask=env_mask_wp) def write_root_velocity_to_sim_index( self, *, root_velocity: torch.Tensor | wp.array, - env_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set the root velocity over selected environment indices into the simulation. + """Set the root center of mass velocity over selected environment indices into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + Args: - root_velocity: Root velocities in simulation world frame. Shape is (len(env_ids),) - with dtype wp.spatial_vectorf. + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6) + or (len(env_ids),) with dtype wp.spatial_vectorf. env_ids: Environment indices. If None, then all indices are used. """ - n = self._n_envs_index(env_ids) - self.assert_shape_and_dtype(root_velocity, (n,), wp.spatial_vectorf, "root_velocity") - self._write_root_state(TT.ROOT_VELOCITY, root_velocity, env_ids) + self.write_root_com_velocity_to_sim_index(root_velocity=root_velocity, env_ids=env_ids) def write_root_velocity_to_sim_mask( self, @@ -398,36 +595,60 @@ def write_root_velocity_to_sim_mask( root_velocity: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root velocity over masked environments into the simulation. + """Set the root center of mass velocity over selected environment mask into the simulation. - The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. Args: - root_velocity: Root velocities in simulation world frame. Shape is (num_instances,) - with dtype wp.spatial_vectorf. - env_mask: Environment mask. If None, then all instances are updated. + root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - self.assert_shape_and_dtype(root_velocity, (self._num_instances,), wp.spatial_vectorf, "root_velocity") - self._write_root_state(TT.ROOT_VELOCITY, root_velocity, mask=env_mask) + self.write_root_com_velocity_to_sim_mask(root_velocity=root_velocity, env_mask=env_mask) def write_root_com_velocity_to_sim_index( self, *, root_velocity: torch.Tensor | wp.array, - env_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set the root center of mass velocity over selected environment indices into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + Args: - root_velocity: Root center of mass velocities in simulation world frame. - Shape is (len(env_ids),) with dtype wp.spatial_vectorf. + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6) + or (len(env_ids),) with dtype wp.spatial_vectorf. env_ids: Environment indices. If None, then all indices are used. """ - n = self._n_envs_index(env_ids) - self.assert_shape_and_dtype(root_velocity, (n,), wp.spatial_vectorf, "root_velocity") - self._write_root_state(TT.ROOT_VELOCITY, root_velocity, env_ids) + env_ids = self._resolve_env_ids(env_ids) + self.assert_shape_and_dtype(root_velocity, (env_ids.shape[0],), wp.spatial_vectorf, "root_velocity") + wp.launch( + shared_kernels.set_root_com_velocity_to_sim_index, + dim=env_ids.shape[0], + inputs=[root_velocity, env_ids, self._num_bodies], + outputs=[self.data.root_com_vel_w, self.data.body_com_acc_w], + device=self._device, + ) + # Invalidate dependent root_link_vel timestamp. + self.data._root_link_vel_w.timestamp = -1.0 + binding = self._get_binding(TT.ROOT_VELOCITY) + binding.write(self.data._root_com_vel_w.data.view(wp.float32), indices=env_ids) def write_root_com_velocity_to_sim_mask( self, @@ -435,23 +656,43 @@ def write_root_com_velocity_to_sim_mask( root_velocity: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root center of mass velocity over masked environments into the simulation. + """Set the root center of mass velocity over selected environment mask into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + Args: - root_velocity: Root center of mass velocities in simulation world frame. - Shape is (num_instances,) with dtype wp.spatial_vectorf. - env_mask: Environment mask. If None, then all instances are updated. + root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ + env_mask_wp = self._resolve_env_mask(env_mask) self.assert_shape_and_dtype(root_velocity, (self._num_instances,), wp.spatial_vectorf, "root_velocity") - self._write_root_state(TT.ROOT_VELOCITY, root_velocity, mask=env_mask) + wp.launch( + shared_kernels.set_root_com_velocity_to_sim_mask, + dim=self._num_instances, + inputs=[root_velocity, env_mask_wp, self._num_bodies], + outputs=[self.data.root_com_vel_w, self.data.body_com_acc_w], + device=self._device, + ) + self.data._root_link_vel_w.timestamp = -1.0 + binding = self._get_binding(TT.ROOT_VELOCITY) + binding.write(self.data._root_com_vel_w.data.view(wp.float32), mask=env_mask_wp) def write_root_link_velocity_to_sim_index( self, *, root_velocity: torch.Tensor | wp.array, - env_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set the root link velocity over selected environment indices into the simulation. @@ -460,14 +701,35 @@ def write_root_link_velocity_to_sim_index( .. note:: This sets the velocity of the root's frame rather than the root's center of mass. + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + Args: - root_velocity: Root frame velocities in simulation world frame. - Shape is (len(env_ids),) with dtype wp.spatial_vectorf. + root_velocity: Root frame velocities in simulation world frame. Shape is (len(env_ids), 6) + or (len(env_ids),) with dtype wp.spatial_vectorf. env_ids: Environment indices. If None, then all indices are used. """ - n = self._n_envs_index(env_ids) - self.assert_shape_and_dtype(root_velocity, (n,), wp.spatial_vectorf, "root_velocity") - self._write_root_state(TT.ROOT_VELOCITY, root_velocity, env_ids) + env_ids = self._resolve_env_ids(env_ids) + self.assert_shape_and_dtype(root_velocity, (env_ids.shape[0],), wp.spatial_vectorf, "root_velocity") + wp.launch( + shared_kernels.set_root_link_velocity_to_sim_index, + dim=env_ids.shape[0], + inputs=[ + root_velocity, + self.data.body_com_pose_b, + self.data.root_link_pose_w, + env_ids, + self._num_bodies, + ], + outputs=[self.data.root_link_vel_w, self.data.root_com_vel_w, self.data.body_com_acc_w], + device=self._device, + ) + binding = self._get_binding(TT.ROOT_VELOCITY) + binding.write(self.data._root_com_vel_w.data.view(wp.float32), indices=env_ids) def write_root_link_velocity_to_sim_mask( self, @@ -475,118 +737,255 @@ def write_root_link_velocity_to_sim_mask( root_velocity: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root link velocity over masked environments into the simulation. + """Set the root link velocity over selected environment mask into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. .. note:: This sets the velocity of the root's frame rather than the root's center of mass. - Args: - root_velocity: Root frame velocities in simulation world frame. - Shape is (num_instances,) with dtype wp.spatial_vectorf. - env_mask: Environment mask. If None, then all instances are updated. - """ - self.assert_shape_and_dtype(root_velocity, (self._num_instances,), wp.spatial_vectorf, "root_velocity") - self._write_root_state(TT.ROOT_VELOCITY, root_velocity, mask=env_mask) - - """ - Operations - Joint State Writers. - """ + .. note:: + This method expects full data. - def write_joint_state_to_sim_mask( - self, - joint_pos: torch.Tensor | wp.array, - joint_vel: torch.Tensor | wp.array, - env_mask: wp.array | None = None, - joint_mask: wp.array | None = None, - ) -> None: - """Write joint positions and velocities over masked environments into the simulation. + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. Args: - joint_pos: Joint positions. Shape is (num_instances, num_joints). - joint_vel: Joint velocities. Shape is (num_instances, num_joints). - env_mask: Environment mask. If None, then all instances are updated. - joint_mask: Joint mask. If None, then all joints are updated. + root_velocity: Root frame velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - self.write_joint_position_to_sim_mask(position=joint_pos, env_mask=env_mask, joint_mask=joint_mask) - self.write_joint_velocity_to_sim_mask(velocity=joint_vel, env_mask=env_mask, joint_mask=joint_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + self.assert_shape_and_dtype(root_velocity, (self._num_instances,), wp.spatial_vectorf, "root_velocity") + wp.launch( + shared_kernels.set_root_link_velocity_to_sim_mask, + dim=self._num_instances, + inputs=[ + root_velocity, + self.data.body_com_pose_b, + self.data.root_link_pose_w, + env_mask_wp, + self._num_bodies, + ], + outputs=[self.data.root_link_vel_w, self.data.root_com_vel_w, self.data.body_com_acc_w], + device=self._device, + ) + binding = self._get_binding(TT.ROOT_VELOCITY) + binding.write(self.data._root_com_vel_w.data.view(wp.float32), mask=env_mask_wp) def write_joint_position_to_sim_index( self, *, position: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write joint positions over selected environment and joint indices into the simulation. + """Set joint positions over selected env / joint indices into the simulation. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - position: Joint positions. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + position: Joint positions [m or rad, depending on joint type]. Shape is + (len(env_ids), len(joint_ids)) with dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - d = len(joint_ids) if joint_ids is not None else self._num_joints - self.assert_shape_and_dtype(position, (n, d), wp.float32, "position") - self._write_flat_tensor(TT.DOF_POSITION, position, env_ids, joint_ids) - self.data._joint_pos_buf.timestamp = -1.0 + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + self.assert_shape_and_dtype(position, (env_ids.shape[0], joint_ids.shape[0]), wp.float32, "position") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[position, env_ids, joint_ids], + outputs=[self._data._joint_pos_buf.data], + device=self._device, + ) + # invalidate body-state buffers so the next read re-fetches FK from the + # wheel using the new joint positions + self._data._body_com_vel_w.timestamp = -1.0 + self._data._body_link_vel_w.timestamp = -1.0 + self._data._body_com_pose_b.timestamp = -1.0 + self._data._body_com_pose_w.timestamp = -1.0 + self._data._body_link_pose_w.timestamp = -1.0 + self._data._joint_acc.timestamp = -1.0 + binding = self._get_binding(TT.DOF_POSITION) + binding.write(self._data._joint_pos_buf.data, indices=env_ids) def write_joint_position_to_sim_mask( self, *, position: torch.Tensor | wp.array, - joint_mask: wp.array | None = None, env_mask: wp.array | None = None, + joint_mask: wp.array | None = None, ) -> None: - """Write joint positions over masked environments into the simulation. + """Set joint positions over selected env / joint masks into the simulation. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - position: Joint positions. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + position: Joint positions [m or rad, depending on joint type]. Shape is + (num_instances, num_joints) with dtype wp.float32. + env_mask: Environment mask. If None, all instances are updated. Shape is + (num_instances,). + joint_mask: Joint mask. If None, all joints are updated. Shape is + (num_joints,). """ + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) self.assert_shape_and_dtype(position, (self._num_instances, self._num_joints), wp.float32, "position") - self._write_flat_tensor_mask(TT.DOF_POSITION, position, env_mask, joint_mask) - self.data._joint_pos_buf.timestamp = -1.0 + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(self._num_instances, self._num_joints), + inputs=[position, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_pos_buf.data], + device=self._device, + ) + # invalidate body-state buffers (see :meth:`write_joint_position_to_sim_index`) + self._data._body_com_vel_w.timestamp = -1.0 + self._data._body_link_vel_w.timestamp = -1.0 + self._data._body_com_pose_b.timestamp = -1.0 + self._data._body_com_pose_w.timestamp = -1.0 + self._data._body_link_pose_w.timestamp = -1.0 + self._data._joint_acc.timestamp = -1.0 + binding = self._get_binding(TT.DOF_POSITION) + binding.write(self._data._joint_pos_buf.data, mask=env_mask_wp) def write_joint_velocity_to_sim_index( self, *, velocity: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write joint velocities over selected environment and joint indices into the simulation. + """Set joint velocities over selected env / joint indices into the simulation. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - velocity: Joint velocities. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + velocity: Joint velocities [m/s or rad/s, depending on joint type]. Shape is + (len(env_ids), len(joint_ids)) with dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - d = len(joint_ids) if joint_ids is not None else self._num_joints - self.assert_shape_and_dtype(velocity, (n, d), wp.float32, "velocity") - self._write_flat_tensor(TT.DOF_VELOCITY, velocity, env_ids, joint_ids) - self.data._joint_vel_buf.timestamp = -1.0 + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + self.assert_shape_and_dtype(velocity, (env_ids.shape[0], joint_ids.shape[0]), wp.float32, "velocity") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[velocity, env_ids, joint_ids], + outputs=[self._data._joint_vel_buf.data], + device=self._device, + ) + # Sync previous_joint_vel to the new values so the next FD step does not + # produce a spurious acceleration spike. + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[velocity, env_ids, joint_ids], + outputs=[self._data._previous_joint_vel], + device=self._device, + ) + self._data._joint_acc.timestamp = -1.0 + binding = self._get_binding(TT.DOF_VELOCITY) + binding.write(self._data._joint_vel_buf.data, indices=env_ids) def write_joint_velocity_to_sim_mask( self, *, velocity: torch.Tensor | wp.array, - joint_mask: wp.array | None = None, env_mask: wp.array | None = None, + joint_mask: wp.array | None = None, ) -> None: - """Write joint velocities over masked environments into the simulation. + """Set joint velocities over selected env / joint masks into the simulation. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - velocity: Joint velocities. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + velocity: Joint velocities [m/s or rad/s, depending on joint type]. Shape is + (num_instances, num_joints) with dtype wp.float32. + env_mask: Environment mask. If None, all instances are updated. Shape is + (num_instances,). + joint_mask: Joint mask. If None, all joints are updated. Shape is + (num_joints,). """ + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) self.assert_shape_and_dtype(velocity, (self._num_instances, self._num_joints), wp.float32, "velocity") - self._write_flat_tensor_mask(TT.DOF_VELOCITY, velocity, env_mask, joint_mask) - self.data._joint_vel_buf.timestamp = -1.0 + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(self._num_instances, self._num_joints), + inputs=[velocity, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_vel_buf.data], + device=self._device, + ) + # Sync previous_joint_vel so the next FD step does not produce a spurious spike. + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(self._num_instances, self._num_joints), + inputs=[velocity, env_mask_wp, joint_mask_wp], + outputs=[self._data._previous_joint_vel], + device=self._device, + ) + self._data._joint_acc.timestamp = -1.0 + binding = self._get_binding(TT.DOF_VELOCITY) + binding.write(self._data._joint_vel_buf.data, mask=env_mask_wp) + + def write_joint_state_to_sim_mask( + self, + *, + position: torch.Tensor | wp.array, + velocity: torch.Tensor | wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write joint positions and velocities over selected environment mask into the simulation. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. + + Args: + position: Joint positions [m or rad, depending on joint type]. Shape is + (num_instances, num_joints) with dtype wp.float32. + velocity: Joint velocities [m/s or rad/s, depending on joint type]. Shape is + (num_instances, num_joints) with dtype wp.float32. + joint_mask: Joint mask. If None, all joints are updated. Shape is + (num_joints,). + env_mask: Environment mask. If None, all instances are updated. Shape is + (num_instances,). + """ + self.write_joint_position_to_sim_mask(position=position, env_mask=env_mask, joint_mask=joint_mask) + self.write_joint_velocity_to_sim_mask(velocity=velocity, env_mask=env_mask, joint_mask=joint_mask) """ Operations - Simulation Parameters Writers. @@ -595,99 +994,278 @@ def write_joint_velocity_to_sim_mask( def write_joint_stiffness_to_sim_index( self, *, - stiffness: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + stiffness: float | torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write joint stiffness over selected indices into the simulation. + """Set joint stiffness over selected env / joint indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``DOF_STIFFNESS`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - stiffness: Joint stiffness. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + stiffness: Joint stiffness [N/m or N·m/rad, depending on joint type]. + May be a scalar :class:`float` (broadcast), or shape + (len(env_ids), len(joint_ids)) with dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - d = len(joint_ids) if joint_ids is not None else self._num_joints - self.assert_shape_and_dtype(stiffness, (n, d), wp.float32, "stiffness") - self._write_flat_tensor(TT.DOF_STIFFNESS, stiffness, env_ids, joint_ids) + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + shape = (env_ids.shape[0], joint_ids.shape[0]) + stiffness = self._broadcast_scalar_to_2d(stiffness, shape) + self.assert_shape_and_dtype(stiffness, shape, wp.float32, "stiffness") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=shape, + inputs=[stiffness, env_ids, joint_ids], + outputs=[self._data._joint_stiffness.data], + device=self._device, + ) + cpu_env_ids = self._get_cpu_env_ids(env_ids) + wp.copy(self.data._cpu_joint_stiffness, self._data._joint_stiffness.data) + binding = self._get_binding(TT.DOF_STIFFNESS) + binding.write(self.data._cpu_joint_stiffness, indices=cpu_env_ids) def write_joint_stiffness_to_sim_mask( self, *, - stiffness: torch.Tensor | wp.array, + stiffness: float | torch.Tensor | wp.array, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Write joint stiffness over masked environments into the simulation. + """Set joint stiffness over selected env / joint masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``DOF_STIFFNESS`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - stiffness: Joint stiffness. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + stiffness: Joint stiffness [N/m or N·m/rad, depending on joint type]. + May be a scalar :class:`float` (broadcast), or shape + (num_instances, num_joints) with dtype wp.float32. + joint_mask: Joint mask. If None, all joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(stiffness, (self._num_instances, self._num_joints), wp.float32, "stiffness") - self._write_flat_tensor_mask(TT.DOF_STIFFNESS, stiffness, env_mask, joint_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + shape = (self._num_instances, self._num_joints) + stiffness = self._broadcast_scalar_to_2d(stiffness, shape) + self.assert_shape_and_dtype(stiffness, shape, wp.float32, "stiffness") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[stiffness, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_stiffness.data], + device=self._device, + ) + wp.copy(self.data._cpu_joint_stiffness, self._data._joint_stiffness.data) + binding = self._get_binding(TT.DOF_STIFFNESS) + binding.write(self.data._cpu_joint_stiffness, mask=self._get_cpu_env_mask(env_mask_wp)) def write_joint_damping_to_sim_index( self, *, - damping: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + damping: float | torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write joint damping over selected indices into the simulation. + """Set joint damping over selected env / joint indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``DOF_DAMPING`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - damping: Joint damping. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + damping: Joint damping [N·s/m or N·m·s/rad, depending on joint type]. + May be a scalar :class:`float` (broadcast), or shape + (len(env_ids), len(joint_ids)) with dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - d = len(joint_ids) if joint_ids is not None else self._num_joints - self.assert_shape_and_dtype(damping, (n, d), wp.float32, "damping") - self._write_flat_tensor(TT.DOF_DAMPING, damping, env_ids, joint_ids) + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + shape = (env_ids.shape[0], joint_ids.shape[0]) + damping = self._broadcast_scalar_to_2d(damping, shape) + self.assert_shape_and_dtype(damping, shape, wp.float32, "damping") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=shape, + inputs=[damping, env_ids, joint_ids], + outputs=[self._data._joint_damping.data], + device=self._device, + ) + cpu_env_ids = self._get_cpu_env_ids(env_ids) + wp.copy(self.data._cpu_joint_damping, self._data._joint_damping.data) + binding = self._get_binding(TT.DOF_DAMPING) + binding.write(self.data._cpu_joint_damping, indices=cpu_env_ids) def write_joint_damping_to_sim_mask( self, *, - damping: torch.Tensor | wp.array, + damping: float | torch.Tensor | wp.array, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Write joint damping over masked environments into the simulation. + """Set joint damping over selected env / joint masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``DOF_DAMPING`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - damping: Joint damping. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + damping: Joint damping [N·s/m or N·m·s/rad, depending on joint type]. + May be a scalar :class:`float` (broadcast), or shape + (num_instances, num_joints) with dtype wp.float32. + joint_mask: Joint mask. If None, all joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(damping, (self._num_instances, self._num_joints), wp.float32, "damping") - self._write_flat_tensor_mask(TT.DOF_DAMPING, damping, env_mask, joint_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + shape = (self._num_instances, self._num_joints) + damping = self._broadcast_scalar_to_2d(damping, shape) + self.assert_shape_and_dtype(damping, shape, wp.float32, "damping") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[damping, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_damping.data], + device=self._device, + ) + wp.copy(self.data._cpu_joint_damping, self._data._joint_damping.data) + binding = self._get_binding(TT.DOF_DAMPING) + binding.write(self.data._cpu_joint_damping, mask=self._get_cpu_env_mask(env_mask_wp)) def write_joint_position_limit_to_sim_index( self, *, limits: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, warn_limit_violation: bool = True, ) -> None: - """Write joint position limits over selected environment indices into the simulation. + """Set joint position limits over selected env / joint indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``DOF_LIMIT`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limits: Joint position limits [rad or m]. Shape is (len(env_ids), len(joint_ids)) - with dtype wp.vec2f. - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. - warn_limit_violation: Whether to use warning or info level logging when default joint - positions exceed the new limits. Defaults to True. + limits: Joint position limits ``[lower, upper]`` + [m or rad, depending on joint type]. Either shape + (len(env_ids), len(joint_ids), 2) with dtype wp.float32, or + shape (len(env_ids), len(joint_ids)) with dtype wp.vec2f. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). + warn_limit_violation: If True, log a warning when the provided limits + are inconsistent (lower > upper). Defaults to True. """ - if isinstance(limits, (int, float)): - raise ValueError("Float scalars are not supported for position limits (vec2f dtype)") - n = self._n_envs_index(env_ids) - d = len(joint_ids) if joint_ids is not None else self._num_joints - self.assert_shape_and_dtype(limits, (n, d), wp.vec2f, "limits") - self._write_flat_tensor(TT.DOF_LIMIT, limits, env_ids, joint_ids) + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + # Position limits cannot be scalar-broadcast (they pair lower/upper); + # match PhysX which explicitly rejects floats here. + if isinstance(limits, float): + raise ValueError("Joint position limits must be a tensor or array, not a float.") + # Accept both wp.vec2f shape (N, J) and the legacy (N, J, 2) wp.float32 + # form (canonical PhysX/Newton layout uses vec2f). + if isinstance(limits, wp.array) and limits.dtype == wp.vec2f: + self.assert_shape_and_dtype(limits, (env_ids.shape[0], joint_ids.shape[0]), wp.vec2f, "limits") + # Reinterpret the vec2f input as a (N, J, 2) float32 view for the kernel. + kernel_limits = wp.array( + ptr=limits.ptr, + shape=(env_ids.shape[0], joint_ids.shape[0], 2), + dtype=wp.float32, + device=str(limits.device), + copy=False, + ) + else: + self.assert_shape_and_dtype(limits, (env_ids.shape[0], joint_ids.shape[0], 2), wp.float32, "limits") + kernel_limits = limits + # Scatter [lower, upper] pairs into the vec2f cache buffer. + wp.launch( + shared_kernels.write_joint_position_limit_to_buffer_index, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[kernel_limits, env_ids, joint_ids], + outputs=[self._data._joint_pos_limits.data], + device=self._device, + ) + # Clamp default_joint_pos to the new limits and refresh soft_joint_pos_limits. + clamped_count = wp.zeros(1, dtype=wp.int32, device=self._device) + wp.launch( + clamp_default_joint_pos_and_update_soft_limits_index, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[ + self._data._joint_pos_limits.data, + env_ids, + joint_ids, + self.cfg.soft_joint_pos_limit_factor, + ], + outputs=[ + self._data._default_joint_pos, + self._data._soft_joint_pos_limits, + clamped_count, + ], + device=self._device, + ) + if clamped_count.numpy()[0] > 0: + violation_message = ( + "Some default joint positions are outside of the range of the new joint limits. Default joint" + " positions will be clamped to be within the new joint limits." + ) + if warn_limit_violation: + logger.warning(violation_message) + else: + logger.info(violation_message) + # Stage to pinned-host CPU: flatten the vec2f buffer to float32 view. + cpu_env_ids = self._get_cpu_env_ids(env_ids) + flat_src = wp.array( + ptr=self._data._joint_pos_limits.data.ptr, + shape=(self._num_instances, self._num_joints, 2), + dtype=wp.float32, + device=self._device, + copy=False, + ) + wp.copy(self.data._cpu_joint_position_limit, flat_src) + binding = self._get_binding(TT.DOF_LIMIT) + binding.write(self.data._cpu_joint_position_limit, indices=cpu_env_ids) def write_joint_position_limit_to_sim_mask( self, @@ -697,180 +1275,700 @@ def write_joint_position_limit_to_sim_mask( env_mask: wp.array | None = None, warn_limit_violation: bool = True, ) -> None: - """Write joint position limits over masked environments into the simulation. + """Set joint position limits over selected env / joint masks into the simulation. - Args: - limits: Joint position limits [rad or m]. Shape is (num_instances, num_joints) - with dtype wp.vec2f. - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. - warn_limit_violation: Whether to use warning or info level logging when default joint - positions exceed the new limits. Defaults to True. - """ - if isinstance(limits, (int, float)): - raise ValueError("Float scalars are not supported for position limits (vec2f dtype)") - self.assert_shape_and_dtype(limits, (self._num_instances, self._num_joints), wp.vec2f, "limits") - self._write_flat_tensor_mask(TT.DOF_LIMIT, limits, env_mask, joint_mask) + This is a CPU-only write routed through pinned-host staging because + ``DOF_LIMIT`` is a CPU-only OVPhysX binding. - def write_joint_velocity_limit_to_sim_index( - self, - *, - limits: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, - ) -> None: - """Write joint max velocity over selected environment indices into the simulation. + .. note:: + This method expects full data. - The velocity limit is used to constrain the joint velocities in the physics engine. + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limits: Joint max velocity [rad/s or m/s]. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + limits: Joint position limits ``[lower, upper]`` + [m or rad, depending on joint type]. Either shape + (num_instances, num_joints, 2) with dtype wp.float32, or shape + (num_instances, num_joints) with dtype wp.vec2f. + joint_mask: Joint mask. If None, all joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). + warn_limit_violation: If True, log a warning when the provided limits + are inconsistent (lower > upper). Defaults to True. """ - n = self._n_envs_index(env_ids) - d = len(joint_ids) if joint_ids is not None else self._num_joints - self.assert_shape_and_dtype(limits, (n, d), wp.float32, "limits") - self._write_flat_tensor(TT.DOF_MAX_VELOCITY, limits, env_ids, joint_ids) - - def write_joint_velocity_limit_to_sim_mask( - self, - *, - limits: torch.Tensor | wp.array, - joint_mask: wp.array | None = None, - env_mask: wp.array | None = None, + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + # Position limits cannot be scalar-broadcast (they pair lower/upper); + # match PhysX which explicitly rejects floats here. + if isinstance(limits, float): + raise ValueError("Joint position limits must be a tensor or array, not a float.") + # Accept both wp.vec2f shape (N, J) and the legacy (N, J, 2) wp.float32 + # form (canonical PhysX/Newton layout uses vec2f). + if isinstance(limits, wp.array) and limits.dtype == wp.vec2f: + self.assert_shape_and_dtype(limits, (self._num_instances, self._num_joints), wp.vec2f, "limits") + kernel_limits = wp.array( + ptr=limits.ptr, + shape=(self._num_instances, self._num_joints, 2), + dtype=wp.float32, + device=str(limits.device), + copy=False, + ) + else: + self.assert_shape_and_dtype(limits, (self._num_instances, self._num_joints, 2), wp.float32, "limits") + kernel_limits = limits + wp.launch( + shared_kernels.write_joint_position_limit_to_buffer_mask, + dim=(self._num_instances, self._num_joints), + inputs=[kernel_limits, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_pos_limits.data], + device=self._device, + ) + # Clamp default_joint_pos to the new limits and refresh soft_joint_pos_limits. + clamped_count = wp.zeros(1, dtype=wp.int32, device=self._device) + wp.launch( + clamp_default_joint_pos_and_update_soft_limits_mask, + dim=(self._num_instances, self._num_joints), + inputs=[ + self._data._joint_pos_limits.data, + env_mask_wp, + joint_mask_wp, + self.cfg.soft_joint_pos_limit_factor, + ], + outputs=[ + self._data._default_joint_pos, + self._data._soft_joint_pos_limits, + clamped_count, + ], + device=self._device, + ) + if clamped_count.numpy()[0] > 0: + violation_message = ( + "Some default joint positions are outside of the range of the new joint limits. Default joint" + " positions will be clamped to be within the new joint limits." + ) + if warn_limit_violation: + logger.warning(violation_message) + else: + logger.info(violation_message) + flat_src = wp.array( + ptr=self._data._joint_pos_limits.data.ptr, + shape=(self._num_instances, self._num_joints, 2), + dtype=wp.float32, + device=self._device, + copy=False, + ) + wp.copy(self.data._cpu_joint_position_limit, flat_src) + binding = self._get_binding(TT.DOF_LIMIT) + binding.write(self.data._cpu_joint_position_limit, mask=self._get_cpu_env_mask(env_mask_wp)) + + def write_joint_velocity_limit_to_sim_index( + self, + *, + limits: float | torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set joint velocity limits over selected env / joint indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``DOF_MAX_VELOCITY`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. + + Args: + limits: Joint velocity limits [m/s or rad/s, depending on joint type]. + May be a scalar :class:`float` (broadcast), or shape + (len(env_ids), len(joint_ids)) with dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). + """ + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + shape = (env_ids.shape[0], joint_ids.shape[0]) + limits = self._broadcast_scalar_to_2d(limits, shape) + self.assert_shape_and_dtype(limits, shape, wp.float32, "limits") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=shape, + inputs=[limits, env_ids, joint_ids], + outputs=[self._data._joint_vel_limits.data], + device=self._device, + ) + cpu_env_ids = self._get_cpu_env_ids(env_ids) + wp.copy(self.data._cpu_joint_velocity_limit, self._data._joint_vel_limits.data) + binding = self._get_binding(TT.DOF_MAX_VELOCITY) + binding.write(self.data._cpu_joint_velocity_limit, indices=cpu_env_ids) + + def write_joint_velocity_limit_to_sim_mask( + self, + *, + limits: float | torch.Tensor | wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: - """Write joint max velocity over masked environments into the simulation. + """Set joint velocity limits over selected env / joint masks into the simulation. - The velocity limit is used to constrain the joint velocities in the physics engine. + This is a CPU-only write routed through pinned-host staging because + ``DOF_MAX_VELOCITY`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limits: Joint max velocity [rad/s or m/s]. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + limits: Joint velocity limits [m/s or rad/s, depending on joint type]. + May be a scalar :class:`float` (broadcast), or shape + (num_instances, num_joints) with dtype wp.float32. + joint_mask: Joint mask. If None, all joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(limits, (self._num_instances, self._num_joints), wp.float32, "limits") - self._write_flat_tensor_mask(TT.DOF_MAX_VELOCITY, limits, env_mask, joint_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + shape = (self._num_instances, self._num_joints) + limits = self._broadcast_scalar_to_2d(limits, shape) + self.assert_shape_and_dtype(limits, shape, wp.float32, "limits") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[limits, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_vel_limits.data], + device=self._device, + ) + wp.copy(self.data._cpu_joint_velocity_limit, self._data._joint_vel_limits.data) + binding = self._get_binding(TT.DOF_MAX_VELOCITY) + binding.write(self.data._cpu_joint_velocity_limit, mask=self._get_cpu_env_mask(env_mask_wp)) def write_joint_effort_limit_to_sim_index( self, *, - limits: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + limits: float | torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write joint effort limits over selected environment indices into the simulation. + """Set joint effort limits over selected env / joint indices into the simulation. - The effort limit is used to constrain the computed joint efforts in the physics engine. + This is a CPU-only write routed through pinned-host staging because + ``DOF_MAX_FORCE`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limits: Joint effort limits [N or N*m]. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + limits: Joint effort limits [N or N·m, depending on joint type]. + May be a scalar :class:`float` (broadcast), or shape + (len(env_ids), len(joint_ids)) with dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - d = len(joint_ids) if joint_ids is not None else self._num_joints - self.assert_shape_and_dtype(limits, (n, d), wp.float32, "limits") - self._write_flat_tensor(TT.DOF_MAX_FORCE, limits, env_ids, joint_ids) + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + shape = (env_ids.shape[0], joint_ids.shape[0]) + limits = self._broadcast_scalar_to_2d(limits, shape) + self.assert_shape_and_dtype(limits, shape, wp.float32, "limits") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=shape, + inputs=[limits, env_ids, joint_ids], + outputs=[self._data._joint_effort_limits.data], + device=self._device, + ) + cpu_env_ids = self._get_cpu_env_ids(env_ids) + wp.copy(self.data._cpu_joint_effort_limit, self._data._joint_effort_limits.data) + binding = self._get_binding(TT.DOF_MAX_FORCE) + binding.write(self.data._cpu_joint_effort_limit, indices=cpu_env_ids) def write_joint_effort_limit_to_sim_mask( self, *, - limits: torch.Tensor | wp.array, + limits: float | torch.Tensor | wp.array, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Write joint effort limits over masked environments into the simulation. + """Set joint effort limits over selected env / joint masks into the simulation. - The effort limit is used to constrain the computed joint efforts in the physics engine. + This is a CPU-only write routed through pinned-host staging because + ``DOF_MAX_FORCE`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limits: Joint effort limits [N or N*m]. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + limits: Joint effort limits [N or N·m, depending on joint type]. + May be a scalar :class:`float` (broadcast), or shape + (num_instances, num_joints) with dtype wp.float32. + joint_mask: Joint mask. If None, all joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(limits, (self._num_instances, self._num_joints), wp.float32, "limits") - self._write_flat_tensor_mask(TT.DOF_MAX_FORCE, limits, env_mask, joint_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + shape = (self._num_instances, self._num_joints) + limits = self._broadcast_scalar_to_2d(limits, shape) + self.assert_shape_and_dtype(limits, shape, wp.float32, "limits") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[limits, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_effort_limits.data], + device=self._device, + ) + wp.copy(self.data._cpu_joint_effort_limit, self._data._joint_effort_limits.data) + binding = self._get_binding(TT.DOF_MAX_FORCE) + binding.write(self.data._cpu_joint_effort_limit, mask=self._get_cpu_env_mask(env_mask_wp)) def write_joint_armature_to_sim_index( self, *, - armature: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + armature: float | torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write joint armature over selected environment indices into the simulation. + """Set joint armature over selected env / joint indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``DOF_ARMATURE`` is a CPU-only OVPhysX binding. - The armature is directly added to the corresponding joint-space inertia. It helps improve the - simulation stability by reducing the joint velocities. + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - armature: Joint armature [kg*m^2 or kg]. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + armature: Joint armature [kg·m²]. May be a scalar :class:`float` + (broadcast), or shape (len(env_ids), len(joint_ids)) with + dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - d = len(joint_ids) if joint_ids is not None else self._num_joints - self.assert_shape_and_dtype(armature, (n, d), wp.float32, "armature") - self._write_flat_tensor(TT.DOF_ARMATURE, armature, env_ids, joint_ids) + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + shape = (env_ids.shape[0], joint_ids.shape[0]) + armature = self._broadcast_scalar_to_2d(armature, shape) + self.assert_shape_and_dtype(armature, shape, wp.float32, "armature") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=shape, + inputs=[armature, env_ids, joint_ids], + outputs=[self._data._joint_armature.data], + device=self._device, + ) + cpu_env_ids = self._get_cpu_env_ids(env_ids) + wp.copy(self.data._cpu_joint_armature, self._data._joint_armature.data) + binding = self._get_binding(TT.DOF_ARMATURE) + binding.write(self.data._cpu_joint_armature, indices=cpu_env_ids) def write_joint_armature_to_sim_mask( self, *, - armature: torch.Tensor | wp.array, + armature: float | torch.Tensor | wp.array, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Write joint armature over masked environments into the simulation. + """Set joint armature over selected env / joint masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``DOF_ARMATURE`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. - The armature is directly added to the corresponding joint-space inertia. It helps improve the - simulation stability by reducing the joint velocities. + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - armature: Joint armature [kg*m^2 or kg]. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + armature: Joint armature [kg·m²]. May be a scalar :class:`float` + (broadcast), or shape (num_instances, num_joints) with dtype + wp.float32. + joint_mask: Joint mask. If None, all joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(armature, (self._num_instances, self._num_joints), wp.float32, "armature") - self._write_flat_tensor_mask(TT.DOF_ARMATURE, armature, env_mask, joint_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + shape = (self._num_instances, self._num_joints) + armature = self._broadcast_scalar_to_2d(armature, shape) + self.assert_shape_and_dtype(armature, shape, wp.float32, "armature") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[armature, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_armature.data], + device=self._device, + ) + wp.copy(self.data._cpu_joint_armature, self._data._joint_armature.data) + binding = self._get_binding(TT.DOF_ARMATURE) + binding.write(self.data._cpu_joint_armature, mask=self._get_cpu_env_mask(env_mask_wp)) def write_joint_friction_coefficient_to_sim_index( self, *, - joint_friction_coeff: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + joint_friction_coeff: float | torch.Tensor | wp.array, + joint_dynamic_friction_coeff: float | torch.Tensor | wp.array | None = None, + joint_viscous_friction_coeff: float | torch.Tensor | wp.array | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write joint friction coefficients over selected indices into the simulation. + r"""Write joint friction coefficients over selected env / joint indices into the simulation. + + Mirrors :meth:`isaaclab_physx.assets.Articulation.write_joint_friction_coefficient_to_sim_index`: + Coulomb (static & dynamic) friction with an optional viscous term. Any of the three + components can be left unset by passing ``None``; the corresponding slot in the + combined ``DOF_FRICTION_PROPERTIES`` ``(N, J, 3)`` binding is preserved. + + ``DOF_FRICTION_PROPERTIES`` is a CPU-only OVPhysX binding, so the + write is routed through pinned-host staging. + + .. note:: + This method expects partial data. Each component, if provided, + may be a scalar :class:`float` (broadcast to + ``(len(env_ids), len(joint_ids))``) or a 2D tensor / warp array. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - joint_friction_coeff: Joint friction coefficients. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + joint_friction_coeff: Static friction coefficient :math:`\mu_s` [dimensionless]. + joint_dynamic_friction_coeff: Dynamic (Coulomb) friction coefficient + :math:`\mu_d`. If ``None``, the dynamic component is preserved. + joint_viscous_friction_coeff: Viscous friction coefficient :math:`c_v`. + If ``None``, the viscous component is preserved. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - d = len(joint_ids) if joint_ids is not None else self._num_joints - self.assert_shape_and_dtype(joint_friction_coeff, (n, d), wp.float32, "joint_friction_coeff") - self._write_friction_column(joint_friction_coeff, env_ids, joint_ids) + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + shape = (env_ids.shape[0], joint_ids.shape[0]) + joint_friction_coeff = self._broadcast_scalar_to_2d(joint_friction_coeff, shape) + if joint_dynamic_friction_coeff is not None: + joint_dynamic_friction_coeff = self._broadcast_scalar_to_2d(joint_dynamic_friction_coeff, shape) + if joint_viscous_friction_coeff is not None: + joint_viscous_friction_coeff = self._broadcast_scalar_to_2d(joint_viscous_friction_coeff, shape) + self.assert_shape_and_dtype(joint_friction_coeff, shape, wp.float32, "joint_friction_coeff") + if joint_dynamic_friction_coeff is not None: + self.assert_shape_and_dtype(joint_dynamic_friction_coeff, shape, wp.float32, "joint_dynamic_friction_coeff") + if joint_viscous_friction_coeff is not None: + self.assert_shape_and_dtype(joint_viscous_friction_coeff, shape, wp.float32, "joint_viscous_friction_coeff") + # refresh the combined (N, J, 3) buffer from the binding so unchanged + # components are preserved on the round-trip + self._data._read_scalar_binding(TT.DOF_FRICTION_PROPERTIES, self._data._joint_friction_props_buf) + wp.launch( + write_joint_friction_data_to_buffer_index, + dim=shape, + inputs=[ + joint_friction_coeff, + joint_dynamic_friction_coeff, + joint_viscous_friction_coeff, + env_ids, + joint_ids, + ], + outputs=[self._data._joint_friction_props_buf.data], + device=self._device, + ) + # Stage the combined (N, J, 3) buffer to pinned-host CPU and write to the binding. + cpu_env_ids = self._get_cpu_env_ids(env_ids) + cpu_friction = self._data._stage_to_pinned_cpu( + TT.DOF_FRICTION_PROPERTIES, "write", self._data._joint_friction_props_buf.data + ) + binding = self._get_binding(TT.DOF_FRICTION_PROPERTIES) + binding.write(cpu_friction, indices=cpu_env_ids) def write_joint_friction_coefficient_to_sim_mask( self, *, - joint_friction_coeff: torch.Tensor | wp.array, + joint_friction_coeff: float | torch.Tensor | wp.array, + joint_dynamic_friction_coeff: float | torch.Tensor | wp.array | None = None, + joint_viscous_friction_coeff: float | torch.Tensor | wp.array | None = None, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Write joint friction coefficients over masked environments into the simulation. + r"""Mask variant of :meth:`write_joint_friction_coefficient_to_sim_index`. Args: - joint_friction_coeff: Joint friction coefficients. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + joint_friction_coeff: Static friction coefficient :math:`\mu_s`. Full data, + shape ``(num_instances, num_joints)``. May be a scalar :class:`float`. + joint_dynamic_friction_coeff: Dynamic friction. ``None`` to preserve. + joint_viscous_friction_coeff: Viscous friction. ``None`` to preserve. + joint_mask: Joint mask. If None, all joints are updated. + env_mask: Environment mask. If None, all instances are updated. """ - self.assert_shape_and_dtype( - joint_friction_coeff, (self._num_instances, self._num_joints), wp.float32, "joint_friction_coeff" + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + shape = (self._num_instances, self._num_joints) + joint_friction_coeff = self._broadcast_scalar_to_2d(joint_friction_coeff, shape) + if joint_dynamic_friction_coeff is not None: + joint_dynamic_friction_coeff = self._broadcast_scalar_to_2d(joint_dynamic_friction_coeff, shape) + if joint_viscous_friction_coeff is not None: + joint_viscous_friction_coeff = self._broadcast_scalar_to_2d(joint_viscous_friction_coeff, shape) + self.assert_shape_and_dtype(joint_friction_coeff, shape, wp.float32, "joint_friction_coeff") + if joint_dynamic_friction_coeff is not None: + self.assert_shape_and_dtype(joint_dynamic_friction_coeff, shape, wp.float32, "joint_dynamic_friction_coeff") + if joint_viscous_friction_coeff is not None: + self.assert_shape_and_dtype(joint_viscous_friction_coeff, shape, wp.float32, "joint_viscous_friction_coeff") + # refresh the (N, J, 3) buffer first (see ``_index`` variant) + self._data._read_scalar_binding(TT.DOF_FRICTION_PROPERTIES, self._data._joint_friction_props_buf) + wp.launch( + write_joint_friction_data_to_buffer_mask, + dim=shape, + inputs=[ + joint_friction_coeff, + joint_dynamic_friction_coeff, + joint_viscous_friction_coeff, + env_mask_wp, + joint_mask_wp, + ], + outputs=[self._data._joint_friction_props_buf.data], + device=self._device, + ) + cpu_friction = self._data._stage_to_pinned_cpu( + TT.DOF_FRICTION_PROPERTIES, "write", self._data._joint_friction_props_buf.data + ) + binding = self._get_binding(TT.DOF_FRICTION_PROPERTIES) + binding.write(cpu_friction, mask=self._get_cpu_env_mask(env_mask_wp)) + + def write_joint_dynamic_friction_coefficient_to_sim_index( + self, + *, + joint_dynamic_friction_coeff: float | torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + r"""Write joint dynamic friction coefficients over selected env / joint indices into the simulation. + + Mirrors :meth:`isaaclab_physx.assets.Articulation.write_joint_dynamic_friction_coefficient_to_sim_index`: + updates only the dynamic (Coulomb) slot of the combined ``DOF_FRICTION_PROPERTIES`` ``(N, J, 3)`` + binding; the static and viscous components are preserved. + + ``DOF_FRICTION_PROPERTIES`` is a CPU-only OVPhysX binding, so the + write is routed through pinned-host staging. + + .. note:: + This method expects partial data. ``joint_dynamic_friction_coeff`` may be a + scalar :class:`float` (broadcast to ``(len(env_ids), len(joint_ids))``) or a + 2D tensor / warp array. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. + + Args: + joint_dynamic_friction_coeff: Dynamic (Coulomb) friction coefficient + :math:`\mu_d` [dimensionless]. Shape is ``(len(env_ids), len(joint_ids))`` + with dtype wp.float32, or a scalar that is broadcast. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). + """ + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + shape = (env_ids.shape[0], joint_ids.shape[0]) + joint_dynamic_friction_coeff = self._broadcast_scalar_to_2d(joint_dynamic_friction_coeff, shape) + self.assert_shape_and_dtype(joint_dynamic_friction_coeff, shape, wp.float32, "joint_dynamic_friction_coeff") + # refresh the combined (N, J, 3) buffer from the binding so unchanged + # components are preserved on the round-trip + self._data._read_scalar_binding(TT.DOF_FRICTION_PROPERTIES, self._data._joint_friction_props_buf) + wp.launch( + write_joint_friction_data_to_buffer_index, + dim=shape, + inputs=[ + None, # in_static — preserved + joint_dynamic_friction_coeff, + None, # in_viscous — preserved + env_ids, + joint_ids, + ], + outputs=[self._data._joint_friction_props_buf.data], + device=self._device, + ) + cpu_env_ids = self._get_cpu_env_ids(env_ids) + cpu_friction = self._data._stage_to_pinned_cpu( + TT.DOF_FRICTION_PROPERTIES, "write", self._data._joint_friction_props_buf.data + ) + binding = self._get_binding(TT.DOF_FRICTION_PROPERTIES) + binding.write(cpu_friction, indices=cpu_env_ids) + + def write_joint_dynamic_friction_coefficient_to_sim_mask( + self, + *, + joint_dynamic_friction_coeff: float | torch.Tensor | wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + r"""Mask variant of :meth:`write_joint_dynamic_friction_coefficient_to_sim_index`. + + Updates only the dynamic (Coulomb) slot of the combined ``DOF_FRICTION_PROPERTIES`` + ``(N, J, 3)`` binding; the static and viscous components are preserved. + + Args: + joint_dynamic_friction_coeff: Dynamic (Coulomb) friction coefficient + :math:`\mu_d` [dimensionless]. Full data, shape + ``(num_instances, num_joints)``. May be a scalar :class:`float`. + joint_mask: Joint mask. If None, all joints are updated. + env_mask: Environment mask. If None, all instances are updated. + """ + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + shape = (self._num_instances, self._num_joints) + joint_dynamic_friction_coeff = self._broadcast_scalar_to_2d(joint_dynamic_friction_coeff, shape) + self.assert_shape_and_dtype(joint_dynamic_friction_coeff, shape, wp.float32, "joint_dynamic_friction_coeff") + # refresh the (N, J, 3) buffer first (see ``_index`` variant) + self._data._read_scalar_binding(TT.DOF_FRICTION_PROPERTIES, self._data._joint_friction_props_buf) + wp.launch( + write_joint_friction_data_to_buffer_mask, + dim=shape, + inputs=[ + None, # in_static — preserved + joint_dynamic_friction_coeff, + None, # in_viscous — preserved + env_mask_wp, + joint_mask_wp, + ], + outputs=[self._data._joint_friction_props_buf.data], + device=self._device, + ) + cpu_friction = self._data._stage_to_pinned_cpu( + TT.DOF_FRICTION_PROPERTIES, "write", self._data._joint_friction_props_buf.data + ) + binding = self._get_binding(TT.DOF_FRICTION_PROPERTIES) + binding.write(cpu_friction, mask=self._get_cpu_env_mask(env_mask_wp)) + + def write_joint_viscous_friction_coefficient_to_sim_index( + self, + *, + joint_viscous_friction_coeff: float | torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + r"""Write joint viscous friction coefficients over selected env / joint indices into the simulation. + + Mirrors :meth:`isaaclab_physx.assets.Articulation.write_joint_viscous_friction_coefficient_to_sim_index`: + updates only the viscous slot of the combined ``DOF_FRICTION_PROPERTIES`` ``(N, J, 3)`` + binding; the static and dynamic components are preserved. + + ``DOF_FRICTION_PROPERTIES`` is a CPU-only OVPhysX binding, so the + write is routed through pinned-host staging. + + .. note:: + This method expects partial data. ``joint_viscous_friction_coeff`` may be a + scalar :class:`float` (broadcast to ``(len(env_ids), len(joint_ids))``) or a + 2D tensor / warp array. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. + + Args: + joint_viscous_friction_coeff: Viscous friction coefficient + :math:`c_v` [N·m·s/rad or N·s/m, depending on joint type]. + Shape is ``(len(env_ids), len(joint_ids))`` with dtype wp.float32, or + a scalar that is broadcast. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). + """ + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + shape = (env_ids.shape[0], joint_ids.shape[0]) + joint_viscous_friction_coeff = self._broadcast_scalar_to_2d(joint_viscous_friction_coeff, shape) + self.assert_shape_and_dtype(joint_viscous_friction_coeff, shape, wp.float32, "joint_viscous_friction_coeff") + # refresh the combined (N, J, 3) buffer from the binding so unchanged + # components are preserved on the round-trip + self._data._read_scalar_binding(TT.DOF_FRICTION_PROPERTIES, self._data._joint_friction_props_buf) + wp.launch( + write_joint_friction_data_to_buffer_index, + dim=shape, + inputs=[ + None, # in_static — preserved + None, # in_dynamic — preserved + joint_viscous_friction_coeff, + env_ids, + joint_ids, + ], + outputs=[self._data._joint_friction_props_buf.data], + device=self._device, + ) + cpu_env_ids = self._get_cpu_env_ids(env_ids) + cpu_friction = self._data._stage_to_pinned_cpu( + TT.DOF_FRICTION_PROPERTIES, "write", self._data._joint_friction_props_buf.data + ) + binding = self._get_binding(TT.DOF_FRICTION_PROPERTIES) + binding.write(cpu_friction, indices=cpu_env_ids) + + def write_joint_viscous_friction_coefficient_to_sim_mask( + self, + *, + joint_viscous_friction_coeff: float | torch.Tensor | wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + r"""Mask variant of :meth:`write_joint_viscous_friction_coefficient_to_sim_index`. + + Updates only the viscous slot of the combined ``DOF_FRICTION_PROPERTIES`` + ``(N, J, 3)`` binding; the static and dynamic components are preserved. + + Args: + joint_viscous_friction_coeff: Viscous friction coefficient + :math:`c_v` [N·m·s/rad or N·s/m, depending on joint type]. + Full data, shape ``(num_instances, num_joints)``. May be a + scalar :class:`float`. + joint_mask: Joint mask. If None, all joints are updated. + env_mask: Environment mask. If None, all instances are updated. + """ + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + shape = (self._num_instances, self._num_joints) + joint_viscous_friction_coeff = self._broadcast_scalar_to_2d(joint_viscous_friction_coeff, shape) + self.assert_shape_and_dtype(joint_viscous_friction_coeff, shape, wp.float32, "joint_viscous_friction_coeff") + # refresh the (N, J, 3) buffer first (see ``_index`` variant) + self._data._read_scalar_binding(TT.DOF_FRICTION_PROPERTIES, self._data._joint_friction_props_buf) + wp.launch( + write_joint_friction_data_to_buffer_mask, + dim=shape, + inputs=[ + None, # in_static — preserved + None, # in_dynamic — preserved + joint_viscous_friction_coeff, + env_mask_wp, + joint_mask_wp, + ], + outputs=[self._data._joint_friction_props_buf.data], + device=self._device, + ) + cpu_friction = self._data._stage_to_pinned_cpu( + TT.DOF_FRICTION_PROPERTIES, "write", self._data._joint_friction_props_buf.data ) - self._write_friction_column_mask(joint_friction_coeff, env_mask, joint_mask) + binding = self._get_binding(TT.DOF_FRICTION_PROPERTIES) + binding.write(cpu_friction, mask=self._get_cpu_env_mask(env_mask_wp)) """ Operations - Setters. @@ -880,20 +1978,42 @@ def set_masses_index( self, *, masses: torch.Tensor | wp.array, - body_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set masses of all bodies using indices. + """Set body masses over selected env / body indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``BODY_MASS`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized + implementations. Performance is similar for both. However, to + allow graphed pipelines, the mask method must be used. Args: - masses: Masses of all bodies [kg]. Shape is (len(env_ids), len(body_ids)). - body_ids: The body indices to set the masses for. Defaults to None (all bodies). - env_ids: The environment indices to set the masses for. Defaults to None (all environments). + masses: Body masses [kg]. Shape is (len(env_ids), len(body_ids)) + with dtype wp.float32. + body_ids: Body indices. Defaults to None (all bodies). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - b = len(body_ids) if body_ids is not None else self._num_bodies - self.assert_shape_and_dtype(masses, (n, b), wp.float32, "masses") - self._write_flat_tensor(TT.BODY_MASS, masses, env_ids, body_ids) + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + self.assert_shape_and_dtype(masses, (env_ids.shape[0], body_ids.shape[0]), wp.float32, "masses") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[masses, env_ids, body_ids], + outputs=[self._data._body_mass.data], + device=self._device, + ) + cpu_env_ids = self._get_cpu_env_ids(env_ids) + wp.copy(self.data._cpu_body_mass, self._data._body_mass.data) + binding = self._get_binding(TT.BODY_MASS) + binding.write(self.data._cpu_body_mass, indices=cpu_env_ids) def set_masses_mask( self, @@ -902,36 +2022,84 @@ def set_masses_mask( body_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set masses of all bodies using masks. + """Set body masses over selected env / body masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``BODY_MASS`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized + implementations. Performance is similar for both. However, to + allow graphed pipelines, the mask method must be used. Args: - masses: Masses of all bodies [kg]. Shape is (num_instances, num_bodies). - body_mask: Body mask. If None, then all bodies are used. - env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + masses: Body masses [kg]. Shape is (num_instances, num_bodies) + with dtype wp.float32. + body_mask: Body mask. If None, all bodies are updated. + Shape is (num_bodies,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ + env_mask_wp = self._resolve_env_mask(env_mask) + body_mask_wp = self._resolve_body_mask(body_mask) self.assert_shape_and_dtype(masses, (self._num_instances, self._num_bodies), wp.float32, "masses") - self._write_flat_tensor_mask(TT.BODY_MASS, masses, env_mask, body_mask) + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(self._num_instances, self._num_bodies), + inputs=[masses, env_mask_wp, body_mask_wp], + outputs=[self._data._body_mass.data], + device=self._device, + ) + wp.copy(self.data._cpu_body_mass, self._data._body_mass.data) + binding = self._get_binding(TT.BODY_MASS) + binding.write(self.data._cpu_body_mass, mask=self._get_cpu_env_mask(env_mask_wp)) def set_coms_index( self, *, coms: torch.Tensor | wp.array, - body_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set center of mass pose of all bodies using indices. + """Set body center-of-mass poses over selected env / body indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``BODY_COM_POSE`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized + implementations. Performance is similar for both. However, to + allow graphed pipelines, the mask method must be used. Args: - coms: Center of mass pose of all bodies. Shape is (len(env_ids), len(body_ids)) - with dtype wp.transformf. - body_ids: The body indices to set the center of mass pose for. Defaults to None (all bodies). - env_ids: The environment indices to set the center of mass pose for. - Defaults to None (all environments). + coms: Body center-of-mass poses [m, quaternion (w, x, y, z)]. + Shape is (len(env_ids), len(body_ids)) with dtype wp.transformf. + body_ids: Body indices. Defaults to None (all bodies). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - b = len(body_ids) if body_ids is not None else self._num_bodies - self.assert_shape_and_dtype(coms, (n, b), wp.transformf, "coms") - self._write_flat_tensor(TT.BODY_COM_POSE, coms, env_ids, body_ids) + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + self.assert_shape_and_dtype(coms, (env_ids.shape[0], body_ids.shape[0]), wp.transformf, "coms") + wp.launch( + shared_kernels.write_body_com_pose_to_buffer_index, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[coms, env_ids, body_ids], + outputs=[self._data._body_com_pose_b.data], + device=self._device, + ) + # Invalidate derived buffers that depend on body_com_pose_b. + self.data._root_com_pose_w.timestamp = -1.0 + self.data._body_com_pose_w.timestamp = -1.0 + cpu_env_ids = self._get_cpu_env_ids(env_ids) + wp.copy(self.data._cpu_body_coms, self._data._body_com_pose_b.data) + binding = self._get_binding(TT.BODY_COM_POSE) + binding.write(self.data._cpu_body_coms, indices=cpu_env_ids) def set_coms_mask( self, @@ -940,35 +2108,84 @@ def set_coms_mask( body_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set center of mass pose of all bodies using masks. + """Set body center-of-mass poses over selected env / body masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``BODY_COM_POSE`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized + implementations. Performance is similar for both. However, to + allow graphed pipelines, the mask method must be used. Args: - coms: Center of mass pose of all bodies. Shape is (num_instances, num_bodies) - with dtype wp.transformf. - body_mask: Body mask. If None, then all bodies are used. - env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + coms: Body center-of-mass poses [m, quaternion (w, x, y, z)]. + Shape is (num_instances, num_bodies) with dtype wp.transformf. + body_mask: Body mask. If None, all bodies are updated. + Shape is (num_bodies,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ + env_mask_wp = self._resolve_env_mask(env_mask) + body_mask_wp = self._resolve_body_mask(body_mask) self.assert_shape_and_dtype(coms, (self._num_instances, self._num_bodies), wp.transformf, "coms") - self._write_flat_tensor_mask(TT.BODY_COM_POSE, coms, env_mask, body_mask) + wp.launch( + shared_kernels.write_body_com_pose_to_buffer_mask, + dim=(self._num_instances, self._num_bodies), + inputs=[coms, env_mask_wp, body_mask_wp], + outputs=[self._data._body_com_pose_b.data], + device=self._device, + ) + # Invalidate derived buffers that depend on body_com_pose_b. + self.data._root_com_pose_w.timestamp = -1.0 + self.data._body_com_pose_w.timestamp = -1.0 + wp.copy(self.data._cpu_body_coms, self._data._body_com_pose_b.data) + binding = self._get_binding(TT.BODY_COM_POSE) + binding.write(self.data._cpu_body_coms, mask=self._get_cpu_env_mask(env_mask_wp)) def set_inertias_index( self, *, inertias: torch.Tensor | wp.array, - body_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set inertias of all bodies using indices. + """Set body inertia tensors over selected env / body indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``BODY_INERTIA`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized + implementations. Performance is similar for both. However, to + allow graphed pipelines, the mask method must be used. Args: - inertias: Inertias of all bodies [kg*m^2]. Shape is (len(env_ids), len(body_ids), 9). - body_ids: The body indices to set the inertias for. Defaults to None (all bodies). - env_ids: The environment indices to set the inertias for. Defaults to None (all environments). + inertias: Body inertia tensors [kg·m²]. Shape is + (len(env_ids), len(body_ids), 9) with dtype wp.float32. + body_ids: Body indices. Defaults to None (all bodies). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - b = len(body_ids) if body_ids is not None else self._num_bodies - self.assert_shape_and_dtype(inertias, (n, b, 9), wp.float32, "inertias") - self._write_flat_tensor(TT.BODY_INERTIA, inertias, env_ids, body_ids) + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + self.assert_shape_and_dtype(inertias, (env_ids.shape[0], body_ids.shape[0], 9), wp.float32, "inertias") + wp.launch( + shared_kernels.write_body_inertia_to_buffer_index, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[inertias, env_ids, body_ids], + outputs=[self._data._body_inertia.data], + device=self._device, + ) + cpu_env_ids = self._get_cpu_env_ids(env_ids) + wp.copy(self.data._cpu_body_inertia, self._data._body_inertia.data) + binding = self._get_binding(TT.BODY_INERTIA) + binding.write(self.data._cpu_body_inertia, indices=cpu_env_ids) def set_inertias_mask( self, @@ -977,38 +2194,80 @@ def set_inertias_mask( body_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set inertias of all bodies using masks. + """Set body inertia tensors over selected env / body masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``BODY_INERTIA`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized + implementations. Performance is similar for both. However, to + allow graphed pipelines, the mask method must be used. Args: - inertias: Inertias of all bodies [kg*m^2]. Shape is (num_instances, num_bodies, 9). - body_mask: Body mask. If None, then all bodies are used. - env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + inertias: Body inertia tensors [kg·m²]. Shape is + (num_instances, num_bodies, 9) with dtype wp.float32. + body_mask: Body mask. If None, all bodies are updated. + Shape is (num_bodies,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ + env_mask_wp = self._resolve_env_mask(env_mask) + body_mask_wp = self._resolve_body_mask(body_mask) self.assert_shape_and_dtype(inertias, (self._num_instances, self._num_bodies, 9), wp.float32, "inertias") - self._write_flat_tensor_mask(TT.BODY_INERTIA, inertias, env_mask, body_mask) - - """ - Operations - Target Setters. - """ + wp.launch( + shared_kernels.write_body_inertia_to_buffer_mask, + dim=(self._num_instances, self._num_bodies), + inputs=[inertias, env_mask_wp, body_mask_wp], + outputs=[self._data._body_inertia.data], + device=self._device, + ) + wp.copy(self.data._cpu_body_inertia, self._data._body_inertia.data) + binding = self._get_binding(TT.BODY_INERTIA) + binding.write(self.data._cpu_body_inertia, mask=self._get_cpu_env_mask(env_mask_wp)) def set_joint_position_target_index( self, *, target: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set joint position targets into internal buffers using indices. - This function does not apply the joint targets to the simulation. It only fills the buffers with - the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. + This function does not apply the joint targets to the simulation. It only fills the + buffers with the desired values. To apply the joint targets, call + :meth:`write_data_to_sim`. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - target: Joint position targets [rad or m]. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + target: Joint position targets [m or rad, depending on joint type]. Shape is + (len(env_ids), len(joint_ids)) with dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - self._set_target_into_buffer(self._data._joint_pos_target, target, env_ids, joint_ids) + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + self.assert_shape_and_dtype(target, (env_ids.shape[0], joint_ids.shape[0]), wp.float32, "target") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[target, env_ids, joint_ids], + outputs=[self._data._joint_pos_target], + device=self._device, + ) + binding = self._get_binding(TT.DOF_POSITION_TARGET) + binding.write(self._data._joint_pos_target, indices=env_ids) def set_joint_position_target_mask( self, @@ -1019,31 +2278,73 @@ def set_joint_position_target_mask( ) -> None: """Set joint position targets into internal buffers using masks. + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. + Args: - target: Joint position targets [rad or m]. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + target: Joint position targets [m or rad, depending on joint type]. Shape is + (num_instances, num_joints) with dtype wp.float32. + joint_mask: Joint mask. If None, all joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, all instances are updated. Shape is + (num_instances,). """ - self._set_target_into_buffer_mask(self._data._joint_pos_target, target, env_mask, joint_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + self.assert_shape_and_dtype(target, (self._num_instances, self._num_joints), wp.float32, "target") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(self._num_instances, self._num_joints), + inputs=[target, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_pos_target], + device=self._device, + ) + binding = self._get_binding(TT.DOF_POSITION_TARGET) + binding.write(self._data._joint_pos_target, mask=env_mask_wp) def set_joint_velocity_target_index( self, *, target: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set joint velocity targets into internal buffers using indices. - This function does not apply the joint targets to the simulation. It only fills the buffers with - the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. + This function does not apply the joint targets to the simulation. It only fills the + buffers with the desired values. To apply the joint targets, call + :meth:`write_data_to_sim`. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - target: Joint velocity targets [rad/s or m/s]. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + target: Joint velocity targets [m/s or rad/s, depending on joint type]. Shape is + (len(env_ids), len(joint_ids)) with dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - self._set_target_into_buffer(self._data._joint_vel_target, target, env_ids, joint_ids) + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + self.assert_shape_and_dtype(target, (env_ids.shape[0], joint_ids.shape[0]), wp.float32, "target") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[target, env_ids, joint_ids], + outputs=[self._data._joint_vel_target], + device=self._device, + ) + binding = self._get_binding(TT.DOF_VELOCITY_TARGET) + binding.write(self._data._joint_vel_target, indices=env_ids) def set_joint_velocity_target_mask( self, @@ -1054,31 +2355,73 @@ def set_joint_velocity_target_mask( ) -> None: """Set joint velocity targets into internal buffers using masks. + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. + Args: - target: Joint velocity targets [rad/s or m/s]. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + target: Joint velocity targets [m/s or rad/s, depending on joint type]. Shape is + (num_instances, num_joints) with dtype wp.float32. + joint_mask: Joint mask. If None, all joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, all instances are updated. Shape is + (num_instances,). """ - self._set_target_into_buffer_mask(self._data._joint_vel_target, target, env_mask, joint_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + self.assert_shape_and_dtype(target, (self._num_instances, self._num_joints), wp.float32, "target") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(self._num_instances, self._num_joints), + inputs=[target, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_vel_target], + device=self._device, + ) + binding = self._get_binding(TT.DOF_VELOCITY_TARGET) + binding.write(self._data._joint_vel_target, mask=env_mask_wp) def set_joint_effort_target_index( self, *, target: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set joint efforts into internal buffers using indices. + """Set joint effort targets into internal buffers using indices. + + This function does not apply the joint targets to the simulation. It only fills the + buffers with the desired values. To apply the joint targets, call + :meth:`write_data_to_sim`. + + .. note:: + This method expects partial data. - This function does not apply the joint targets to the simulation. It only fills the buffers with - the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - target: Joint effort targets [N or N*m]. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + target: Joint effort targets [N or N·m, depending on joint type]. Shape is + (len(env_ids), len(joint_ids)) with dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - self._set_target_into_buffer(self._data._joint_effort_target, target, env_ids, joint_ids) + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + self.assert_shape_and_dtype(target, (env_ids.shape[0], joint_ids.shape[0]), wp.float32, "target") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[target, env_ids, joint_ids], + outputs=[self._data._joint_effort_target], + device=self._device, + ) + binding = self._get_binding(TT.DOF_ACTUATION_FORCE) + binding.write(self._data._joint_effort_target, indices=env_ids) def set_joint_effort_target_mask( self, @@ -1087,14 +2430,35 @@ def set_joint_effort_target_mask( joint_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set joint efforts into internal buffers using masks. + """Set joint effort targets into internal buffers using masks. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - target: Joint effort targets [N or N*m]. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + target: Joint effort targets [N or N·m, depending on joint type]. Shape is + (num_instances, num_joints) with dtype wp.float32. + joint_mask: Joint mask. If None, all joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, all instances are updated. Shape is + (num_instances,). """ - self._set_target_into_buffer_mask(self._data._joint_effort_target, target, env_mask, joint_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + self.assert_shape_and_dtype(target, (self._num_instances, self._num_joints), wp.float32, "target") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(self._num_instances, self._num_joints), + inputs=[target, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_effort_target], + device=self._device, + ) + binding = self._get_binding(TT.DOF_ACTUATION_FORCE) + binding.write(self._data._joint_effort_target, mask=env_mask_wp) """ Operations - Tendons. @@ -1107,23 +2471,41 @@ def set_fixed_tendon_stiffness_index( fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set fixed tendon stiffness into internal buffers using indices. + """Set fixed-tendon stiffness over selected env / tendon indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_STIFFNESS`` is a CPU-only OVPhysX binding. - This function does not apply the tendon stiffness to the simulation. It only fills the buffers with - the desired values. To apply the tendon stiffness, call the - :meth:`write_fixed_tendon_properties_to_sim_index` method. + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - stiffness: Fixed tendon stiffness. Shape is (len(env_ids), len(fixed_tendon_ids)). - fixed_tendon_ids: The tendon indices to set the stiffness for. Defaults to None (all fixed tendons). - env_ids: Environment indices. If None, then all indices are used. + stiffness: Fixed-tendon stiffness [N/m]. May be a scalar + :class:`float` (broadcast), or shape + (len(env_ids), len(fixed_tendon_ids)) with dtype wp.float32. + fixed_tendon_ids: Fixed-tendon indices. Defaults to None (all fixed tendons). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - t = len(fixed_tendon_ids) if fixed_tendon_ids else self._nft() - self.assert_shape_and_dtype(stiffness, (n, t), wp.float32, "stiffness") - if self._data._fixed_tendon_stiffness is not None: - self._set_target_into_buffer(self._data._fixed_tendon_stiffness, stiffness, env_ids, fixed_tendon_ids) - + env_ids = self._resolve_env_ids(env_ids) + tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_ids) + shape = (env_ids.shape[0], tendon_ids.shape[0]) + stiffness = self._broadcast_scalar_to_2d(stiffness, shape) + self.assert_shape_and_dtype(stiffness, shape, wp.float32, "stiffness") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=shape, + inputs=[stiffness, env_ids, tendon_ids], + outputs=[self._data._fixed_tendon_stiffness.data], + device=self._device, + ) + binding = self._get_binding(TT.FIXED_TENDON_STIFFNESS) + binding.write(self._data._fixed_tendon_stiffness.data, indices=env_ids) + def set_fixed_tendon_stiffness_mask( self, *, @@ -1131,18 +2513,42 @@ def set_fixed_tendon_stiffness_mask( fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set fixed tendon stiffness into internal buffers using masks. + """Set fixed-tendon stiffness over selected env / tendon masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_STIFFNESS`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - stiffness: Fixed tendon stiffness. Shape is (num_instances, num_fixed_tendons). - fixed_tendon_mask: Tendon mask. If None, then all fixed tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + stiffness: Fixed-tendon stiffness [N/m]. May be a scalar + :class:`float` (broadcast), or shape + (num_instances, num_fixed_tendons) with dtype wp.float32. + fixed_tendon_mask: Fixed-tendon mask. If None, all fixed tendons are updated. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(stiffness, (self._num_instances, self._nft()), wp.float32, "stiffness") - if self._data._fixed_tendon_stiffness is not None: - self._set_target_into_buffer_mask( - self._data._fixed_tendon_stiffness, stiffness, env_mask, fixed_tendon_mask - ) + env_mask_wp = self._resolve_env_mask(env_mask) + tendon_mask_wp = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + shape = (self._num_instances, self._num_fixed_tendons) + stiffness = self._broadcast_scalar_to_2d(stiffness, shape) + self.assert_shape_and_dtype(stiffness, shape, wp.float32, "stiffness") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[stiffness, env_mask_wp, tendon_mask_wp], + outputs=[self._data._fixed_tendon_stiffness.data], + device=self._device, + ) + binding = self._get_binding(TT.FIXED_TENDON_STIFFNESS) + binding.write(self._data._fixed_tendon_stiffness.data, mask=env_mask_wp) def set_fixed_tendon_damping_index( self, @@ -1151,18 +2557,40 @@ def set_fixed_tendon_damping_index( fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set fixed tendon damping into internal buffers using indices. + """Set fixed-tendon damping over selected env / tendon indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_DAMPING`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - damping: Fixed tendon damping. Shape is (len(env_ids), len(fixed_tendon_ids)). - fixed_tendon_ids: The tendon indices to set the damping for. Defaults to None (all fixed tendons). - env_ids: Environment indices. If None, then all indices are used. + damping: Fixed-tendon damping [N·s/m]. May be a scalar :class:`float` + (broadcast), or shape (len(env_ids), len(fixed_tendon_ids)) with + dtype wp.float32. + fixed_tendon_ids: Fixed-tendon indices. Defaults to None (all fixed tendons). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - t = len(fixed_tendon_ids) if fixed_tendon_ids else self._nft() - self.assert_shape_and_dtype(damping, (n, t), wp.float32, "damping") - if self._data._fixed_tendon_damping is not None: - self._set_target_into_buffer(self._data._fixed_tendon_damping, damping, env_ids, fixed_tendon_ids) + env_ids = self._resolve_env_ids(env_ids) + tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_ids) + shape = (env_ids.shape[0], tendon_ids.shape[0]) + damping = self._broadcast_scalar_to_2d(damping, shape) + self.assert_shape_and_dtype(damping, shape, wp.float32, "damping") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=shape, + inputs=[damping, env_ids, tendon_ids], + outputs=[self._data._fixed_tendon_damping.data], + device=self._device, + ) + binding = self._get_binding(TT.FIXED_TENDON_DAMPING) + binding.write(self._data._fixed_tendon_damping.data, indices=env_ids) def set_fixed_tendon_damping_mask( self, @@ -1171,16 +2599,42 @@ def set_fixed_tendon_damping_mask( fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set fixed tendon damping into internal buffers using masks. + """Set fixed-tendon damping over selected env / tendon masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_DAMPING`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - damping: Fixed tendon damping. Shape is (num_instances, num_fixed_tendons). - fixed_tendon_mask: Tendon mask. If None, then all fixed tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + damping: Fixed-tendon damping [N·s/m]. May be a scalar :class:`float` + (broadcast), or shape (num_instances, num_fixed_tendons) with + dtype wp.float32. + fixed_tendon_mask: Fixed-tendon mask. If None, all fixed tendons are updated. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(damping, (self._num_instances, self._nft()), wp.float32, "damping") - if self._data._fixed_tendon_damping is not None: - self._set_target_into_buffer_mask(self._data._fixed_tendon_damping, damping, env_mask, fixed_tendon_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + tendon_mask_wp = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + shape = (self._num_instances, self._num_fixed_tendons) + damping = self._broadcast_scalar_to_2d(damping, shape) + self.assert_shape_and_dtype(damping, shape, wp.float32, "damping") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[damping, env_mask_wp, tendon_mask_wp], + outputs=[self._data._fixed_tendon_damping.data], + device=self._device, + ) + binding = self._get_binding(TT.FIXED_TENDON_DAMPING) + binding.write(self._data._fixed_tendon_damping.data, mask=env_mask_wp) def set_fixed_tendon_limit_stiffness_index( self, @@ -1189,20 +2643,40 @@ def set_fixed_tendon_limit_stiffness_index( fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set fixed tendon limit stiffness into internal buffers using indices. + """Set fixed-tendon limit stiffness over selected env / tendon indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_LIMIT_STIFFNESS`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limit_stiffness: Fixed tendon limit stiffness. Shape is (len(env_ids), len(fixed_tendon_ids)). - fixed_tendon_ids: The tendon indices. Defaults to None (all fixed tendons). - env_ids: Environment indices. If None, then all indices are used. + limit_stiffness: Fixed-tendon limit stiffness [N/m]. May be a + scalar :class:`float` (broadcast), or shape + (len(env_ids), len(fixed_tendon_ids)) with dtype wp.float32. + fixed_tendon_ids: Fixed-tendon indices. Defaults to None (all fixed tendons). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - t = len(fixed_tendon_ids) if fixed_tendon_ids else self._nft() - self.assert_shape_and_dtype(limit_stiffness, (n, t), wp.float32, "limit_stiffness") - if self._data._fixed_tendon_limit_stiffness is not None: - self._set_target_into_buffer( - self._data._fixed_tendon_limit_stiffness, limit_stiffness, env_ids, fixed_tendon_ids - ) + env_ids = self._resolve_env_ids(env_ids) + tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_ids) + shape = (env_ids.shape[0], tendon_ids.shape[0]) + limit_stiffness = self._broadcast_scalar_to_2d(limit_stiffness, shape) + self.assert_shape_and_dtype(limit_stiffness, shape, wp.float32, "limit_stiffness") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=shape, + inputs=[limit_stiffness, env_ids, tendon_ids], + outputs=[self._data._fixed_tendon_limit_stiffness.data], + device=self._device, + ) + binding = self._get_binding(TT.FIXED_TENDON_LIMIT_STIFFNESS) + binding.write(self._data._fixed_tendon_limit_stiffness.data, indices=env_ids) def set_fixed_tendon_limit_stiffness_mask( self, @@ -1211,18 +2685,42 @@ def set_fixed_tendon_limit_stiffness_mask( fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set fixed tendon limit stiffness into internal buffers using masks. + """Set fixed-tendon limit stiffness over selected env / tendon masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_LIMIT_STIFFNESS`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limit_stiffness: Fixed tendon limit stiffness. Shape is (num_instances, num_fixed_tendons). - fixed_tendon_mask: Tendon mask. If None, then all fixed tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + limit_stiffness: Fixed-tendon limit stiffness [N/m]. May be a + scalar :class:`float` (broadcast), or shape + (num_instances, num_fixed_tendons) with dtype wp.float32. + fixed_tendon_mask: Fixed-tendon mask. If None, all fixed tendons are updated. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(limit_stiffness, (self._num_instances, self._nft()), wp.float32, "limit_stiffness") - if self._data._fixed_tendon_limit_stiffness is not None: - self._set_target_into_buffer_mask( - self._data._fixed_tendon_limit_stiffness, limit_stiffness, env_mask, fixed_tendon_mask - ) + env_mask_wp = self._resolve_env_mask(env_mask) + tendon_mask_wp = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + shape = (self._num_instances, self._num_fixed_tendons) + limit_stiffness = self._broadcast_scalar_to_2d(limit_stiffness, shape) + self.assert_shape_and_dtype(limit_stiffness, shape, wp.float32, "limit_stiffness") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[limit_stiffness, env_mask_wp, tendon_mask_wp], + outputs=[self._data._fixed_tendon_limit_stiffness.data], + device=self._device, + ) + binding = self._get_binding(TT.FIXED_TENDON_LIMIT_STIFFNESS) + binding.write(self._data._fixed_tendon_limit_stiffness.data, mask=env_mask_wp) def set_fixed_tendon_position_limit_index( self, @@ -1231,19 +2729,46 @@ def set_fixed_tendon_position_limit_index( fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set fixed tendon position limits into internal buffers using indices. + """Set fixed-tendon position limits over selected env / tendon indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_LIMIT`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limit: Fixed tendon position limits. Shape is (len(env_ids), len(fixed_tendon_ids)) - with dtype wp.vec2f. - fixed_tendon_ids: The tendon indices. Defaults to None (all fixed tendons). - env_ids: Environment indices. If None, then all indices are used. + limit: Fixed-tendon position limits ``[lower, upper]`` [m]. + Shape is (len(env_ids), len(fixed_tendon_ids), 2) with dtype wp.float32. + fixed_tendon_ids: Fixed-tendon indices. Defaults to None (all fixed tendons). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - t = len(fixed_tendon_ids) if fixed_tendon_ids else self._nft() - self.assert_shape_and_dtype(limit, (n, t), wp.vec2f, "limit") - if self._data._fixed_tendon_pos_limits is not None: - self._set_target_into_buffer(self._data._fixed_tendon_pos_limits, limit, env_ids, fixed_tendon_ids) + env_ids = self._resolve_env_ids(env_ids) + tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_ids) + self.assert_shape_and_dtype(limit, (env_ids.shape[0], tendon_ids.shape[0], 2), wp.float32, "limit") + # Scatter [lower, upper] pairs into the vec2f cache buffer. + wp.launch( + shared_kernels.write_joint_position_limit_to_buffer_index, + dim=(env_ids.shape[0], tendon_ids.shape[0]), + inputs=[limit, env_ids, tendon_ids], + outputs=[self._data._fixed_tendon_pos_limits.data], + device=self._device, + ) + # reinterpret the vec2f buffer as a (N, T, 2) float32 view for the binding + flat_src = wp.array( + ptr=self._data._fixed_tendon_pos_limits.data.ptr, + shape=(self._num_instances, self._num_fixed_tendons, 2), + dtype=wp.float32, + device=self._device, + copy=False, + ) + binding = self._get_binding(TT.FIXED_TENDON_LIMIT) + binding.write(flat_src, indices=env_ids) def set_fixed_tendon_position_limit_mask( self, @@ -1252,17 +2777,46 @@ def set_fixed_tendon_position_limit_mask( fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set fixed tendon position limits into internal buffers using masks. + """Set fixed-tendon position limits over selected env / tendon masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_LIMIT`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limit: Fixed tendon position limits. Shape is (num_instances, num_fixed_tendons) - with dtype wp.vec2f. - fixed_tendon_mask: Tendon mask. If None, then all fixed tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + limit: Fixed-tendon position limits ``[lower, upper]`` [m]. + Shape is (num_instances, num_fixed_tendons, 2) with dtype wp.float32. + fixed_tendon_mask: Fixed-tendon mask. If None, all fixed tendons are updated. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(limit, (self._num_instances, self._nft()), wp.vec2f, "limit") - if self._data._fixed_tendon_pos_limits is not None: - self._set_target_into_buffer_mask(self._data._fixed_tendon_pos_limits, limit, env_mask, fixed_tendon_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + tendon_mask_wp = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + self.assert_shape_and_dtype(limit, (self._num_instances, self._num_fixed_tendons, 2), wp.float32, "limit") + wp.launch( + shared_kernels.write_joint_position_limit_to_buffer_mask, + dim=(self._num_instances, self._num_fixed_tendons), + inputs=[limit, env_mask_wp, tendon_mask_wp], + outputs=[self._data._fixed_tendon_pos_limits.data], + device=self._device, + ) + flat_src = wp.array( + ptr=self._data._fixed_tendon_pos_limits.data.ptr, + shape=(self._num_instances, self._num_fixed_tendons, 2), + dtype=wp.float32, + device=self._device, + copy=False, + ) + binding = self._get_binding(TT.FIXED_TENDON_LIMIT) + binding.write(flat_src, mask=env_mask_wp) def set_fixed_tendon_rest_length_index( self, @@ -1271,18 +2825,40 @@ def set_fixed_tendon_rest_length_index( fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set fixed tendon rest length into internal buffers using indices. + """Set fixed-tendon rest lengths over selected env / tendon indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_REST_LENGTH`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - rest_length: Fixed tendon rest length. Shape is (len(env_ids), len(fixed_tendon_ids)). - fixed_tendon_ids: The tendon indices. Defaults to None (all fixed tendons). - env_ids: Environment indices. If None, then all indices are used. + rest_length: Fixed-tendon rest lengths [m]. May be a scalar + :class:`float` (broadcast), or shape + (len(env_ids), len(fixed_tendon_ids)) with dtype wp.float32. + fixed_tendon_ids: Fixed-tendon indices. Defaults to None (all fixed tendons). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - t = len(fixed_tendon_ids) if fixed_tendon_ids else self._nft() - self.assert_shape_and_dtype(rest_length, (n, t), wp.float32, "rest_length") - if self._data._fixed_tendon_rest_length is not None: - self._set_target_into_buffer(self._data._fixed_tendon_rest_length, rest_length, env_ids, fixed_tendon_ids) + env_ids = self._resolve_env_ids(env_ids) + tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_ids) + shape = (env_ids.shape[0], tendon_ids.shape[0]) + rest_length = self._broadcast_scalar_to_2d(rest_length, shape) + self.assert_shape_and_dtype(rest_length, shape, wp.float32, "rest_length") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=shape, + inputs=[rest_length, env_ids, tendon_ids], + outputs=[self._data._fixed_tendon_rest_length.data], + device=self._device, + ) + binding = self._get_binding(TT.FIXED_TENDON_REST_LENGTH) + binding.write(self._data._fixed_tendon_rest_length.data, indices=env_ids) def set_fixed_tendon_rest_length_mask( self, @@ -1291,18 +2867,42 @@ def set_fixed_tendon_rest_length_mask( fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set fixed tendon rest length into internal buffers using masks. + """Set fixed-tendon rest lengths over selected env / tendon masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_REST_LENGTH`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - rest_length: Fixed tendon rest length. Shape is (num_instances, num_fixed_tendons). - fixed_tendon_mask: Tendon mask. If None, then all fixed tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + rest_length: Fixed-tendon rest lengths [m]. May be a scalar + :class:`float` (broadcast), or shape + (num_instances, num_fixed_tendons) with dtype wp.float32. + fixed_tendon_mask: Fixed-tendon mask. If None, all fixed tendons are updated. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(rest_length, (self._num_instances, self._nft()), wp.float32, "rest_length") - if self._data._fixed_tendon_rest_length is not None: - self._set_target_into_buffer_mask( - self._data._fixed_tendon_rest_length, rest_length, env_mask, fixed_tendon_mask - ) + env_mask_wp = self._resolve_env_mask(env_mask) + tendon_mask_wp = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + shape = (self._num_instances, self._num_fixed_tendons) + rest_length = self._broadcast_scalar_to_2d(rest_length, shape) + self.assert_shape_and_dtype(rest_length, shape, wp.float32, "rest_length") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[rest_length, env_mask_wp, tendon_mask_wp], + outputs=[self._data._fixed_tendon_rest_length.data], + device=self._device, + ) + binding = self._get_binding(TT.FIXED_TENDON_REST_LENGTH) + binding.write(self._data._fixed_tendon_rest_length.data, mask=env_mask_wp) def set_fixed_tendon_offset_index( self, @@ -1311,18 +2911,40 @@ def set_fixed_tendon_offset_index( fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set fixed tendon offset into internal buffers using indices. + """Set fixed-tendon offsets over selected env / tendon indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_OFFSET`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - offset: Fixed tendon offset. Shape is (len(env_ids), len(fixed_tendon_ids)). - fixed_tendon_ids: The tendon indices. Defaults to None (all fixed tendons). - env_ids: Environment indices. If None, then all indices are used. + offset: Fixed-tendon offsets [m]. May be a scalar :class:`float` + (broadcast), or shape (len(env_ids), len(fixed_tendon_ids)) + with dtype wp.float32. + fixed_tendon_ids: Fixed-tendon indices. Defaults to None (all fixed tendons). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - t = len(fixed_tendon_ids) if fixed_tendon_ids else self._nft() - self.assert_shape_and_dtype(offset, (n, t), wp.float32, "offset") - if self._data._fixed_tendon_offset is not None: - self._set_target_into_buffer(self._data._fixed_tendon_offset, offset, env_ids, fixed_tendon_ids) + env_ids = self._resolve_env_ids(env_ids) + tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_ids) + shape = (env_ids.shape[0], tendon_ids.shape[0]) + offset = self._broadcast_scalar_to_2d(offset, shape) + self.assert_shape_and_dtype(offset, shape, wp.float32, "offset") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=shape, + inputs=[offset, env_ids, tendon_ids], + outputs=[self._data._fixed_tendon_offset.data], + device=self._device, + ) + binding = self._get_binding(TT.FIXED_TENDON_OFFSET) + binding.write(self._data._fixed_tendon_offset.data, indices=env_ids) def set_fixed_tendon_offset_mask( self, @@ -1331,16 +2953,42 @@ def set_fixed_tendon_offset_mask( fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set fixed tendon offset into internal buffers using masks. + """Set fixed-tendon offsets over selected env / tendon masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_OFFSET`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - offset: Fixed tendon offset. Shape is (num_instances, num_fixed_tendons). - fixed_tendon_mask: Tendon mask. If None, then all fixed tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + offset: Fixed-tendon offsets [m]. May be a scalar :class:`float` + (broadcast), or shape (num_instances, num_fixed_tendons) with + dtype wp.float32. + fixed_tendon_mask: Fixed-tendon mask. If None, all fixed tendons are updated. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(offset, (self._num_instances, self._nft()), wp.float32, "offset") - if self._data._fixed_tendon_offset is not None: - self._set_target_into_buffer_mask(self._data._fixed_tendon_offset, offset, env_mask, fixed_tendon_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + tendon_mask_wp = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + shape = (self._num_instances, self._num_fixed_tendons) + offset = self._broadcast_scalar_to_2d(offset, shape) + self.assert_shape_and_dtype(offset, shape, wp.float32, "offset") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[offset, env_mask_wp, tendon_mask_wp], + outputs=[self._data._fixed_tendon_offset.data], + device=self._device, + ) + binding = self._get_binding(TT.FIXED_TENDON_OFFSET) + binding.write(self._data._fixed_tendon_offset.data, mask=env_mask_wp) def write_fixed_tendon_properties_to_sim_index( self, @@ -1348,25 +2996,44 @@ def write_fixed_tendon_properties_to_sim_index( fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write fixed tendon properties into the simulation using indices. + """Push the cached fixed-tendon properties to the simulation in a single batch. + + PhysX exposes a single ``root_view.set_fixed_tendon_properties`` that writes all + six tendon property buffers at once. OVPhysX has no such batch setter, so this + method writes each ``FIXED_TENDON_*`` binding individually from the matching + ``self._data._fixed_tendon_*`` buffer. + + .. note:: + Only env indices apply to the simulation write; ``fixed_tendon_ids`` is + accepted for API parity with PhysX but is unused (the simulation + writes all tendons of the selected envs). Args: - fixed_tendon_ids: The fixed tendon indices to write the properties for. Defaults to None - (all fixed tendons). - env_ids: Environment indices. If None, then all indices are used. + fixed_tendon_ids: Accepted for PhysX API parity; ignored. + env_ids: Environment indices. If None, all environments are written. """ - if self._nft() == 0: - return - for tt, buf in [ + env_ids = self._resolve_env_ids(env_ids) + for tt, buf in ( (TT.FIXED_TENDON_STIFFNESS, self._data._fixed_tendon_stiffness), (TT.FIXED_TENDON_DAMPING, self._data._fixed_tendon_damping), (TT.FIXED_TENDON_LIMIT_STIFFNESS, self._data._fixed_tendon_limit_stiffness), - (TT.FIXED_TENDON_LIMIT, self._data._fixed_tendon_pos_limits), (TT.FIXED_TENDON_REST_LENGTH, self._data._fixed_tendon_rest_length), (TT.FIXED_TENDON_OFFSET, self._data._fixed_tendon_offset), - ]: - if buf is not None: - self._write_flat_tensor(tt, buf, env_ids, fixed_tendon_ids) + ): + binding = self._get_binding(tt) + if binding is not None: + binding.write(buf.data, indices=env_ids) + # Position-limit binding consumes a flat (N, T, 2) float32 view. + binding = self._get_binding(TT.FIXED_TENDON_LIMIT) + if binding is not None: + flat_src = wp.array( + ptr=self._data._fixed_tendon_pos_limits.data.ptr, + shape=(self._num_instances, self._num_fixed_tendons, 2), + dtype=wp.float32, + device=self._device, + copy=False, + ) + binding.write(flat_src, indices=env_ids) def write_fixed_tendon_properties_to_sim_mask( self, @@ -1374,24 +3041,33 @@ def write_fixed_tendon_properties_to_sim_mask( fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Write fixed tendon properties into the simulation using masks. + """Mask variant of :meth:`write_fixed_tendon_properties_to_sim_index`. Args: - 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. + fixed_tendon_mask: Accepted for PhysX API parity; ignored. + env_mask: Environment mask. If None, all environments are written. """ - if self._nft() == 0: - return - for tt, buf in [ + env_mask_wp = self._resolve_env_mask(env_mask) + for tt, buf in ( (TT.FIXED_TENDON_STIFFNESS, self._data._fixed_tendon_stiffness), (TT.FIXED_TENDON_DAMPING, self._data._fixed_tendon_damping), (TT.FIXED_TENDON_LIMIT_STIFFNESS, self._data._fixed_tendon_limit_stiffness), - (TT.FIXED_TENDON_LIMIT, self._data._fixed_tendon_pos_limits), (TT.FIXED_TENDON_REST_LENGTH, self._data._fixed_tendon_rest_length), (TT.FIXED_TENDON_OFFSET, self._data._fixed_tendon_offset), - ]: - if buf is not None: - self._write_flat_tensor_mask(tt, buf, env_mask, fixed_tendon_mask) + ): + binding = self._get_binding(tt) + if binding is not None: + binding.write(buf.data, mask=env_mask_wp) + binding = self._get_binding(TT.FIXED_TENDON_LIMIT) + if binding is not None: + flat_src = wp.array( + ptr=self._data._fixed_tendon_pos_limits.data.ptr, + shape=(self._num_instances, self._num_fixed_tendons, 2), + dtype=wp.float32, + device=self._device, + copy=False, + ) + binding.write(flat_src, mask=env_mask_wp) def set_spatial_tendon_stiffness_index( self, @@ -1400,18 +3076,41 @@ def set_spatial_tendon_stiffness_index( spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set spatial tendon stiffness into internal buffers using indices. + """Set spatial-tendon stiffness over selected env / tendon indices into the simulation. + + ``SPATIAL_TENDON_STIFFNESS`` is a sim-device binding on OVPhysX + (tendon properties are applied without a CPU clone), so the write + goes directly from the sim-device buffer to the binding. + + .. note:: + This method expects partial data. A scalar :class:`float` is + broadcast to ``(len(env_ids), len(spatial_tendon_ids))``. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - stiffness: Spatial tendon stiffness. Shape is (len(env_ids), len(spatial_tendon_ids)). - spatial_tendon_ids: The tendon indices. Defaults to None (all spatial tendons). - env_ids: Environment indices. If None, then all indices are used. + stiffness: Spatial-tendon stiffness [N/m]. Scalar :class:`float`, + or shape ``(len(env_ids), len(spatial_tendon_ids))`` with + dtype wp.float32. + spatial_tendon_ids: Spatial-tendon indices. Defaults to None (all spatial tendons). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - t = len(spatial_tendon_ids) if spatial_tendon_ids else self._nst() - self.assert_shape_and_dtype(stiffness, (n, t), wp.float32, "stiffness") - if self._data._spatial_tendon_stiffness is not None: - self._set_target_into_buffer(self._data._spatial_tendon_stiffness, stiffness, env_ids, spatial_tendon_ids) + env_ids = self._resolve_env_ids(env_ids) + tendon_ids = self._resolve_spatial_tendon_ids(spatial_tendon_ids) + stiffness = self._broadcast_scalar_to_2d(stiffness, (env_ids.shape[0], tendon_ids.shape[0])) + self.assert_shape_and_dtype(stiffness, (env_ids.shape[0], tendon_ids.shape[0]), wp.float32, "stiffness") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], tendon_ids.shape[0]), + inputs=[stiffness, env_ids, tendon_ids], + outputs=[self._data._spatial_tendon_stiffness.data], + device=self._device, + ) + binding = self._get_binding(TT.SPATIAL_TENDON_STIFFNESS) + binding.write(self._data._spatial_tendon_stiffness.data, indices=env_ids) def set_spatial_tendon_stiffness_mask( self, @@ -1420,18 +3119,42 @@ def set_spatial_tendon_stiffness_mask( spatial_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set spatial tendon stiffness into internal buffers using masks. + """Set spatial-tendon stiffness over selected env / tendon masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``SPATIAL_TENDON_STIFFNESS`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - stiffness: Spatial tendon stiffness. Shape is (num_instances, num_spatial_tendons). - spatial_tendon_mask: Tendon mask. If None, then all spatial tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + stiffness: Spatial-tendon stiffness [N/m]. May be a scalar + :class:`float` (broadcast), or shape + (num_instances, num_spatial_tendons) with dtype wp.float32. + spatial_tendon_mask: Spatial-tendon mask. If None, all spatial tendons are updated. + Shape is (num_spatial_tendons,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(stiffness, (self._num_instances, self._nst()), wp.float32, "stiffness") - if self._data._spatial_tendon_stiffness is not None: - self._set_target_into_buffer_mask( - self._data._spatial_tendon_stiffness, stiffness, env_mask, spatial_tendon_mask - ) + env_mask_wp = self._resolve_env_mask(env_mask) + tendon_mask_wp = self._resolve_spatial_tendon_mask(spatial_tendon_mask) + shape = (self._num_instances, self._num_spatial_tendons) + stiffness = self._broadcast_scalar_to_2d(stiffness, shape) + self.assert_shape_and_dtype(stiffness, shape, wp.float32, "stiffness") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[stiffness, env_mask_wp, tendon_mask_wp], + outputs=[self._data._spatial_tendon_stiffness.data], + device=self._device, + ) + binding = self._get_binding(TT.SPATIAL_TENDON_STIFFNESS) + binding.write(self._data._spatial_tendon_stiffness.data, mask=env_mask_wp) def set_spatial_tendon_damping_index( self, @@ -1440,18 +3163,38 @@ def set_spatial_tendon_damping_index( spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set spatial tendon damping into internal buffers using indices. + """Set spatial-tendon damping over selected env / tendon indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``SPATIAL_TENDON_DAMPING`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - damping: Spatial tendon damping. Shape is (len(env_ids), len(spatial_tendon_ids)). - spatial_tendon_ids: The tendon indices. Defaults to None (all spatial tendons). - env_ids: Environment indices. If None, then all indices are used. + damping: Spatial-tendon damping [N·s/m]. Shape is + (len(env_ids), len(spatial_tendon_ids)) with dtype wp.float32. + spatial_tendon_ids: Spatial-tendon indices. Defaults to None (all spatial tendons). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - t = len(spatial_tendon_ids) if spatial_tendon_ids else self._nst() - self.assert_shape_and_dtype(damping, (n, t), wp.float32, "damping") - if self._data._spatial_tendon_damping is not None: - self._set_target_into_buffer(self._data._spatial_tendon_damping, damping, env_ids, spatial_tendon_ids) + env_ids = self._resolve_env_ids(env_ids) + tendon_ids = self._resolve_spatial_tendon_ids(spatial_tendon_ids) + damping = self._broadcast_scalar_to_2d(damping, (env_ids.shape[0], tendon_ids.shape[0])) + self.assert_shape_and_dtype(damping, (env_ids.shape[0], tendon_ids.shape[0]), wp.float32, "damping") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], tendon_ids.shape[0]), + inputs=[damping, env_ids, tendon_ids], + outputs=[self._data._spatial_tendon_damping.data], + device=self._device, + ) + binding = self._get_binding(TT.SPATIAL_TENDON_DAMPING) + binding.write(self._data._spatial_tendon_damping.data, indices=env_ids) def set_spatial_tendon_damping_mask( self, @@ -1460,18 +3203,42 @@ def set_spatial_tendon_damping_mask( spatial_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set spatial tendon damping into internal buffers using masks. + """Set spatial-tendon damping over selected env / tendon masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``SPATIAL_TENDON_DAMPING`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - damping: Spatial tendon damping. Shape is (num_instances, num_spatial_tendons). - spatial_tendon_mask: Tendon mask. If None, then all spatial tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + damping: Spatial-tendon damping [N·s/m]. May be a scalar + :class:`float` (broadcast), or shape + (num_instances, num_spatial_tendons) with dtype wp.float32. + spatial_tendon_mask: Spatial-tendon mask. If None, all spatial tendons are updated. + Shape is (num_spatial_tendons,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(damping, (self._num_instances, self._nst()), wp.float32, "damping") - if self._data._spatial_tendon_damping is not None: - self._set_target_into_buffer_mask( - self._data._spatial_tendon_damping, damping, env_mask, spatial_tendon_mask - ) + env_mask_wp = self._resolve_env_mask(env_mask) + tendon_mask_wp = self._resolve_spatial_tendon_mask(spatial_tendon_mask) + shape = (self._num_instances, self._num_spatial_tendons) + damping = self._broadcast_scalar_to_2d(damping, shape) + self.assert_shape_and_dtype(damping, shape, wp.float32, "damping") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[damping, env_mask_wp, tendon_mask_wp], + outputs=[self._data._spatial_tendon_damping.data], + device=self._device, + ) + binding = self._get_binding(TT.SPATIAL_TENDON_DAMPING) + binding.write(self._data._spatial_tendon_damping.data, mask=env_mask_wp) def set_spatial_tendon_limit_stiffness_index( self, @@ -1480,20 +3247,42 @@ def set_spatial_tendon_limit_stiffness_index( spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set spatial tendon limit stiffness into internal buffers using indices. + """Set spatial-tendon limit stiffness over selected env / tendon indices into the simulation. + + ``SPATIAL_TENDON_LIMIT_STIFFNESS`` is a sim-device binding on OVPhysX; + the write goes directly from the sim-device buffer to the binding. + + .. note:: + This method expects partial data. A scalar :class:`float` is + broadcast to ``(len(env_ids), len(spatial_tendon_ids))``. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limit_stiffness: Spatial tendon limit stiffness. Shape is (len(env_ids), len(spatial_tendon_ids)). - spatial_tendon_ids: The tendon indices. Defaults to None (all spatial tendons). - env_ids: Environment indices. If None, then all indices are used. + limit_stiffness: Spatial-tendon limit stiffness [N/m]. Scalar + :class:`float`, or shape ``(len(env_ids), len(spatial_tendon_ids))`` + with dtype wp.float32. + spatial_tendon_ids: Spatial-tendon indices. Defaults to None (all spatial tendons). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - t = len(spatial_tendon_ids) if spatial_tendon_ids else self._nst() - self.assert_shape_and_dtype(limit_stiffness, (n, t), wp.float32, "limit_stiffness") - if self._data._spatial_tendon_limit_stiffness is not None: - self._set_target_into_buffer( - self._data._spatial_tendon_limit_stiffness, limit_stiffness, env_ids, spatial_tendon_ids - ) + env_ids = self._resolve_env_ids(env_ids) + tendon_ids = self._resolve_spatial_tendon_ids(spatial_tendon_ids) + limit_stiffness = self._broadcast_scalar_to_2d(limit_stiffness, (env_ids.shape[0], tendon_ids.shape[0])) + self.assert_shape_and_dtype( + limit_stiffness, (env_ids.shape[0], tendon_ids.shape[0]), wp.float32, "limit_stiffness" + ) + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], tendon_ids.shape[0]), + inputs=[limit_stiffness, env_ids, tendon_ids], + outputs=[self._data._spatial_tendon_limit_stiffness.data], + device=self._device, + ) + binding = self._get_binding(TT.SPATIAL_TENDON_LIMIT_STIFFNESS) + binding.write(self._data._spatial_tendon_limit_stiffness.data, indices=env_ids) def set_spatial_tendon_limit_stiffness_mask( self, @@ -1502,18 +3291,42 @@ def set_spatial_tendon_limit_stiffness_mask( spatial_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set spatial tendon limit stiffness into internal buffers using masks. + """Set spatial-tendon limit stiffness over selected env / tendon masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``SPATIAL_TENDON_LIMIT_STIFFNESS`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limit_stiffness: Spatial tendon limit stiffness. Shape is (num_instances, num_spatial_tendons). - spatial_tendon_mask: Tendon mask. If None, then all spatial tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + limit_stiffness: Spatial-tendon limit stiffness [N/m]. May be a + scalar :class:`float` (broadcast), or shape + (num_instances, num_spatial_tendons) with dtype wp.float32. + spatial_tendon_mask: Spatial-tendon mask. If None, all spatial tendons are updated. + Shape is (num_spatial_tendons,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(limit_stiffness, (self._num_instances, self._nst()), wp.float32, "limit_stiffness") - if self._data._spatial_tendon_limit_stiffness is not None: - self._set_target_into_buffer_mask( - self._data._spatial_tendon_limit_stiffness, limit_stiffness, env_mask, spatial_tendon_mask - ) + env_mask_wp = self._resolve_env_mask(env_mask) + tendon_mask_wp = self._resolve_spatial_tendon_mask(spatial_tendon_mask) + shape = (self._num_instances, self._num_spatial_tendons) + limit_stiffness = self._broadcast_scalar_to_2d(limit_stiffness, shape) + self.assert_shape_and_dtype(limit_stiffness, shape, wp.float32, "limit_stiffness") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[limit_stiffness, env_mask_wp, tendon_mask_wp], + outputs=[self._data._spatial_tendon_limit_stiffness.data], + device=self._device, + ) + binding = self._get_binding(TT.SPATIAL_TENDON_LIMIT_STIFFNESS) + binding.write(self._data._spatial_tendon_limit_stiffness.data, mask=env_mask_wp) def set_spatial_tendon_offset_index( self, @@ -1522,18 +3335,40 @@ def set_spatial_tendon_offset_index( spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set spatial tendon offset into internal buffers using indices. + """Set spatial-tendon offsets over selected env / tendon indices into the simulation. + + ``SPATIAL_TENDON_OFFSET`` is a sim-device binding on OVPhysX; the + write goes directly from the sim-device buffer to the binding. + + .. note:: + This method expects partial data. A scalar :class:`float` is + broadcast to ``(len(env_ids), len(spatial_tendon_ids))``. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - offset: Spatial tendon offset. Shape is (len(env_ids), len(spatial_tendon_ids)). - spatial_tendon_ids: The tendon indices. Defaults to None (all spatial tendons). - env_ids: Environment indices. If None, then all indices are used. + offset: Spatial-tendon offsets [m]. Scalar :class:`float`, or + shape ``(len(env_ids), len(spatial_tendon_ids))`` with + dtype wp.float32. + spatial_tendon_ids: Spatial-tendon indices. Defaults to None (all spatial tendons). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - t = len(spatial_tendon_ids) if spatial_tendon_ids else self._nst() - self.assert_shape_and_dtype(offset, (n, t), wp.float32, "offset") - if self._data._spatial_tendon_offset is not None: - self._set_target_into_buffer(self._data._spatial_tendon_offset, offset, env_ids, spatial_tendon_ids) + env_ids = self._resolve_env_ids(env_ids) + tendon_ids = self._resolve_spatial_tendon_ids(spatial_tendon_ids) + offset = self._broadcast_scalar_to_2d(offset, (env_ids.shape[0], tendon_ids.shape[0])) + self.assert_shape_and_dtype(offset, (env_ids.shape[0], tendon_ids.shape[0]), wp.float32, "offset") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], tendon_ids.shape[0]), + inputs=[offset, env_ids, tendon_ids], + outputs=[self._data._spatial_tendon_offset.data], + device=self._device, + ) + binding = self._get_binding(TT.SPATIAL_TENDON_OFFSET) + binding.write(self._data._spatial_tendon_offset.data, indices=env_ids) def set_spatial_tendon_offset_mask( self, @@ -1542,16 +3377,42 @@ def set_spatial_tendon_offset_mask( spatial_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set spatial tendon offset into internal buffers using masks. + """Set spatial-tendon offsets over selected env / tendon masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``SPATIAL_TENDON_OFFSET`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - offset: Spatial tendon offset. Shape is (num_instances, num_spatial_tendons). - spatial_tendon_mask: Tendon mask. If None, then all spatial tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + offset: Spatial-tendon offsets [m]. May be a scalar :class:`float` + (broadcast), or shape (num_instances, num_spatial_tendons) with + dtype wp.float32. + spatial_tendon_mask: Spatial-tendon mask. If None, all spatial tendons are updated. + Shape is (num_spatial_tendons,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(offset, (self._num_instances, self._nst()), wp.float32, "offset") - if self._data._spatial_tendon_offset is not None: - self._set_target_into_buffer_mask(self._data._spatial_tendon_offset, offset, env_mask, spatial_tendon_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + tendon_mask_wp = self._resolve_spatial_tendon_mask(spatial_tendon_mask) + shape = (self._num_instances, self._num_spatial_tendons) + offset = self._broadcast_scalar_to_2d(offset, shape) + self.assert_shape_and_dtype(offset, shape, wp.float32, "offset") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[offset, env_mask_wp, tendon_mask_wp], + outputs=[self._data._spatial_tendon_offset.data], + device=self._device, + ) + binding = self._get_binding(TT.SPATIAL_TENDON_OFFSET) + binding.write(self._data._spatial_tendon_offset.data, mask=env_mask_wp) def write_spatial_tendon_properties_to_sim_index( self, @@ -1559,23 +3420,28 @@ def write_spatial_tendon_properties_to_sim_index( spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write spatial tendon properties into the simulation using indices. + """Push the cached spatial-tendon properties to the simulation in a single batch. + + Mirrors :meth:`write_fixed_tendon_properties_to_sim_index` for + spatial tendons. Only the four wheel-supported tensor types are + written; ``ARTICULATION_SPATIAL_TENDON_LIMIT`` and + ``ARTICULATION_SPATIAL_TENDON_REST_LENGTH`` are forward-compat + stubs (see ``docs/superpowers/specs/2026-04-28-ovphysx-wheel-gaps-for-marco.md``). Args: - spatial_tendon_ids: The spatial tendon indices to write the properties for. Defaults to None - (all spatial tendons). - env_ids: Environment indices. If None, then all indices are used. + spatial_tendon_ids: Accepted for PhysX API parity; ignored. + env_ids: Environment indices. If None, all environments are written. """ - if self._nst() == 0: - return - for tt, buf in [ + env_ids = self._resolve_env_ids(env_ids) + for tt, buf in ( (TT.SPATIAL_TENDON_STIFFNESS, self._data._spatial_tendon_stiffness), (TT.SPATIAL_TENDON_DAMPING, self._data._spatial_tendon_damping), (TT.SPATIAL_TENDON_LIMIT_STIFFNESS, self._data._spatial_tendon_limit_stiffness), (TT.SPATIAL_TENDON_OFFSET, self._data._spatial_tendon_offset), - ]: - if buf is not None: - self._write_flat_tensor(tt, buf, env_ids, spatial_tendon_ids) + ): + binding = self._get_binding(tt) + if binding is not None: + binding.write(buf.data, indices=env_ids) def write_spatial_tendon_properties_to_sim_mask( self, @@ -1583,138 +3449,94 @@ def write_spatial_tendon_properties_to_sim_mask( spatial_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Write spatial tendon properties into the simulation using masks. - - Args: - spatial_tendon_mask: Spatial tendon mask. If None, then all spatial tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. - """ - if self._nst() == 0: - return - for tt, buf in [ + """Mask variant of :meth:`write_spatial_tendon_properties_to_sim_index`.""" + env_mask_wp = self._resolve_env_mask(env_mask) + for tt, buf in ( (TT.SPATIAL_TENDON_STIFFNESS, self._data._spatial_tendon_stiffness), (TT.SPATIAL_TENDON_DAMPING, self._data._spatial_tendon_damping), (TT.SPATIAL_TENDON_LIMIT_STIFFNESS, self._data._spatial_tendon_limit_stiffness), (TT.SPATIAL_TENDON_OFFSET, self._data._spatial_tendon_offset), - ]: - if buf is not None: - self._write_flat_tensor_mask(tt, buf, env_mask, spatial_tendon_mask) - - """ - Deprecated methods. - """ - - def write_root_state_to_sim( - self, - root_state: torch.Tensor | wp.array, - env_ids: Sequence[int] | wp.array | None = None, - ) -> None: - """Deprecated in base class. Use :meth:`write_root_pose_to_sim_index` and - :meth:`write_root_velocity_to_sim_index` instead.""" - self._write_root_state(TT.ROOT_POSE, root_state[:, :7], env_ids) - self._write_root_state(TT.ROOT_VELOCITY, root_state[:, 7:], env_ids) - - def write_root_com_state_to_sim( - self, - root_state: torch.Tensor | wp.array, - env_ids: Sequence[int] | wp.array | None = None, - ) -> None: - """Deprecated in base class. Use :meth:`write_root_com_pose_to_sim_index` and - :meth:`write_root_com_velocity_to_sim_index` instead.""" - self._write_root_state(TT.ROOT_POSE, root_state[:, :7], env_ids) - self._write_root_state(TT.ROOT_VELOCITY, root_state[:, 7:], env_ids) - - def write_root_link_state_to_sim( - self, - root_state: torch.Tensor | wp.array, - env_ids: Sequence[int] | wp.array | None = None, - ) -> None: - """Deprecated in base class. Use :meth:`write_root_link_pose_to_sim_index` and - :meth:`write_root_link_velocity_to_sim_index` instead.""" - self._write_root_state(TT.ROOT_POSE, root_state[:, :7], env_ids) - self._write_root_state(TT.ROOT_VELOCITY, root_state[:, 7:], env_ids) - - def write_joint_state_to_sim( - self, - position: torch.Tensor | wp.array, - velocity: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, - ) -> None: - """Deprecated in base class. Use :meth:`write_joint_position_to_sim_index` and - :meth:`write_joint_velocity_to_sim_index` instead.""" - self.write_joint_position_to_sim_index(position=position, joint_ids=joint_ids, env_ids=env_ids) - self.write_joint_velocity_to_sim_index(velocity=velocity, joint_ids=joint_ids, env_ids=env_ids) + ): + binding = self._get_binding(tt) + if binding is not None: + binding.write(buf.data, mask=env_mask_wp) """ Internal helper. """ def _initialize_impl(self) -> None: + """Initialize the articulation from the OVPhysX simulation backend.""" + # obtain global simulation view physx_instance = OvPhysxManager.get_physx_instance() if physx_instance is None: raise RuntimeError("OvPhysxManager has not been initialized yet.") + self._ovphysx = physx_instance + self._device = OvPhysxManager.get_device() + # IsaacLab uses two conventions for env-glob prim paths: + # /World/envs/env_.*/Robot -- regex dot-star for "any env index" + # /World/envs/{ENV_REGEX_NS}/... -- explicit placeholder + # ovphysx ``create_tensor_binding`` expects fnmatch-style globs, so both map to '*'. prim_path = self.cfg.prim_path - # Convert IsaacLab prim-path notation to the glob patterns ovphysx expects. - # IsaacLab uses two conventions: - # /World/envs/env_.*/Robot -- regex dot-star for "any env index" - # /World/envs/{ENV_REGEX_NS}/Robot -- explicit placeholder - # ovphysx create_tensor_binding() uses fnmatch-style globs, so both map to '*'. pattern = re.sub(r"\{ENV_REGEX_NS\}", "*", prim_path) - pattern = re.sub(r"\.\*", "*", pattern) # env_.* -> env_* - - # The pattern above points to the ArticulationCfg prim (e.g. /World/envs/env_*/Robot). - # However, PhysicsArticulationRootAPI may be on a CHILD prim (e.g. /Robot/torso) - # rather than on the prim itself. create_tensor_binding() only matches prims that - # *have* PhysicsArticulationRootAPI, so we need to extend the pattern to the actual - # articulation root. Mirror the PhysX backend's discovery logic: find the first - # matching prim in the USD stage, walk its subtree for the articulation root, and - # append the relative suffix to the glob pattern. - from pxr import UsdPhysics - - from isaaclab.sim.utils.queries import find_first_matching_prim, get_all_matching_child_prims + pattern = re.sub(r"\.\*", "*", pattern) + # ``PhysicsArticulationRootAPI`` may live on a CHILD prim rather than on + # the cfg prim itself. ``create_tensor_binding`` only matches prims that + # have the API applied, so the pattern must be extended to the actual + # articulation root. stage = PhysicsManager._sim.stage - first_prim = find_first_matching_prim(prim_path, stage=stage) - if first_prim is None: - raise RuntimeError(f"OvPhysxManager: no prim found for path '{prim_path}'.") - first_prim_path = first_prim.GetPath().pathString - - root_prims = get_all_matching_child_prims( - first_prim_path, - predicate=lambda p: p.HasAPI(UsdPhysics.ArticulationRootAPI), - traverse_instance_prims=False, - ) - if len(root_prims) == 0: - raise RuntimeError( - f"No prim with PhysicsArticulationRootAPI found under '{first_prim_path}'." - " Check that the articulation has 'PhysicsArticulationRootAPI' applied." - ) - if len(root_prims) > 1: - raise RuntimeError( - f"Multiple articulation roots found under '{first_prim_path}': {root_prims}." - " There must be exactly one articulation root per prim path." - ) - self._articulation_root_path = root_prims[0].GetPath().pathString - root_relative = self._articulation_root_path[len(first_prim_path) :] - if root_relative: - # e.g. first_prim_path=/World/envs/env_0/Robot, root_relative=/torso - # pattern becomes /World/envs/env_*/Robot/torso + if self.cfg.articulation_root_prim_path is not None: + # explicit subpath: skip auto-discovery but validate the prim exists + root_relative = self.cfg.articulation_root_prim_path + self._articulation_root_path = prim_path + root_relative + if sim_utils.find_first_matching_prim(self._articulation_root_path, stage=stage) is None: + raise RuntimeError( + f"Failed to find articulation root prim at '{self._articulation_root_path}'." + " Check that ``cfg.articulation_root_prim_path`` points at a prim that exists" + " in the USD stage." + ) pattern = pattern + root_relative - logger.info("OvPhysxManager: articulation root at '%s' (pattern extended to '%s')", root_relative, pattern) + logger.info("OvPhysxManager: explicit articulation root '%s' (pattern '%s')", root_relative, pattern) + else: + first_prim = sim_utils.find_first_matching_prim(prim_path, stage=stage) + if first_prim is None: + raise RuntimeError(f"Failed to find prim for expression: '{prim_path}'.") + first_prim_path = first_prim.GetPath().pathString + + root_prims = sim_utils.get_all_matching_child_prims( + first_prim_path, + predicate=lambda p: p.HasAPI(UsdPhysics.ArticulationRootAPI), + traverse_instance_prims=False, + ) + if len(root_prims) == 0: + raise RuntimeError( + f"Failed to find an articulation root when resolving '{prim_path}'." + " Ensure the prim has 'USD ArticulationRootAPI' applied." + ) + if len(root_prims) > 1: + raise RuntimeError( + f"Failed to find a single articulation root when resolving '{prim_path}'." + f" Found multiple under '{first_prim_path}'." + ) + + self._articulation_root_path = root_prims[0].GetPath().pathString + root_relative = self._articulation_root_path[len(first_prim_path) :] + if root_relative: + pattern = pattern + root_relative + logger.info( + "OvPhysxManager: articulation root at '%s' (pattern extended to '%s')", root_relative, pattern + ) - # Bindings are created lazily (on first access) to avoid allocating - # handles for tensor types the user never queries. Only the root-pose - # binding is created eagerly because we need it to read articulation - # metadata (joint count, body count, names, fixed-base flag). - self._bindings: dict[int, Any] = {} - self._physx_instance = physx_instance self._binding_pattern = pattern + # eagerly create every binding the data container reads at init, so + # failures surface here rather than as KeyError downstream eager_types = [ TT.ROOT_POSE, TT.DOF_POSITION, + TT.DOF_VELOCITY, TT.DOF_STIFFNESS, TT.DOF_DAMPING, TT.DOF_LIMIT, @@ -1732,6 +3554,14 @@ def _initialize_impl(self) -> None: except Exception: logger.debug("Could not create tensor binding for type %s on pattern %s", tt, pattern) + if not self._bindings: + raise RuntimeError( + f"OVPhysX could not create any articulation bindings for pattern {pattern!r}. " + f"Check that prim_path={prim_path!r} matches at least one " + "UsdPhysics.ArticulationRootAPI prim." + ) + + # read metadata from the first available binding sample = next(iter(self._bindings.values())) self._num_instances = sample.count self._num_joints = sample.dof_count @@ -1740,23 +3570,37 @@ def _initialize_impl(self) -> None: self._joint_names = list(sample.dof_names) self._body_names = list(sample.body_names) - # Create data container. - self._data = ArticulationData(self._bindings, self._device, binding_getter=self._get_binding) - - # Discover tendon counts/names before buffer allocation so that - # _create_buffers can size the tendon property arrays. + # tendon counts/names must be resolved before buffer allocation self._process_tendons() + # construct the data container + self._data = ArticulationData( + self._bindings, + self._device, + self._num_instances, + self._num_bodies, + self._num_joints, + self._num_fixed_tendons, + self._num_spatial_tendons, + self._body_names, + self._joint_names, + self._fixed_tendon_names, + self._spatial_tendon_names, + binding_getter=self._get_binding, + ) + + # allocate asset-side buffers self._create_buffers() + # apply initial state from config self._process_cfg() + + # build actuator instances and write drive properties to PhysX self._process_actuators_cfg() - self._validate_cfg() - self._log_articulation_info() - # Cache the effort binding and a stable float32 view of the applied_torque - # buffer for write_data_to_sim(). The binding's internal write cache - # (keyed on object identity) handles the fast path automatically. + # cache effort / target bindings and write-views for write_data_to_sim(). + # The effort view aliases applied_torque so the binding gets the actuator + # output without an extra copy. self._effort_binding = self._get_binding(TT.DOF_ACTUATION_FORCE) if self._effort_binding is not None: torque = self._data._applied_torque @@ -1771,7 +3615,6 @@ def _initialize_impl(self) -> None: else: self._effort_write_view = None - # Cache position/velocity target bindings + views for one-shot writes. def _make_write_view(tt, buf): b = self._get_binding(tt) if b is None or buf is None: @@ -1786,24 +3629,59 @@ def _make_write_view(tt, buf): TT.DOF_VELOCITY_TARGET, self._data._joint_vel_target ) - # Let the articulation data know that it is fully instantiated and ready to use. - self.data.is_primed = True + # validate the resolved configuration AFTER actuator/tendon processing + # so the values reflect any overrides applied by the actuator models + self._validate_cfg() - def _create_buffers(self) -> None: - self._data._create_buffers() + # prime the data by performing the first read + self.update(0.0) - self._ALL_INDICES = wp.array(np.arange(self._num_instances, dtype=np.int32), device=self._device) + # mark data as ready + self._data.is_primed = True + def _create_buffers(self) -> None: + """Allocate asset-side buffers (index/mask constants, wrench buf, pinned CPU staging).""" + N = self._num_instances + B = self._num_bodies + J = self._num_joints + FT = self._num_fixed_tendons + ST = self._num_spatial_tendons + device = self._device + + # Index constants. + self._ALL_INDICES = wp.array(np.arange(N, dtype=np.int32), device=device) + self._ALL_BODY_INDICES = wp.array(np.arange(B, dtype=np.int32), device=device) + self._ALL_JOINT_INDICES = wp.array(np.arange(J, dtype=np.int32), device=device) + self._ALL_FIXED_TENDON_INDICES = wp.array(np.arange(FT, dtype=np.int32), device=device) + self._ALL_SPATIAL_TENDON_INDICES = wp.array(np.arange(ST, dtype=np.int32), device=device) + + # All-true masks. + self._ALL_TRUE_ENV_MASK = wp.array(np.ones(N, dtype=bool), dtype=wp.bool, device=device) + self._ALL_TRUE_BODY_MASK = wp.array(np.ones(B, dtype=bool), dtype=wp.bool, device=device) + self._ALL_TRUE_JOINT_MASK = wp.array(np.ones(J, dtype=bool), dtype=wp.bool, device=device) + self._ALL_TRUE_FIXED_TENDON_MASK = wp.array(np.ones(FT, dtype=bool), dtype=wp.bool, device=device) + self._ALL_TRUE_SPATIAL_TENDON_MASK = wp.array(np.ones(ST, dtype=bool), dtype=wp.bool, device=device) + + # Wrench buffer (force, torque, position) per body, written by the + # ``_body_wrench_to_world`` kernel and consumed by the + # ``LINK_WRENCH`` binding which expects the 3D ``(N, B, 9)`` shape. + self._wrench_buf = wp.zeros((N, B, 9), dtype=wp.float32, device=device) + + # Wrench composers. self._instantaneous_wrench_composer = WrenchComposer(self) self._permanent_wrench_composer = WrenchComposer(self) - self._wrench_buf = wp.zeros((self._num_instances, self._num_bodies, 9), dtype=wp.float32, device=self._device) - # Joint-index arrays for each actuator (filled by _process_actuators_cfg). + # Wrench scratch buffer (used by _apply_external_wrenches, not yet allocated above). + # Joint-index arrays for each actuator (populated by _process_actuators_cfg). self._joint_ids_per_actuator: dict[str, list[int]] = {} - self._write_scratch: dict[int, wp.array] = {} + + # Pinned-host CPU staging for env ids/masks (PR #5329 pattern). + self._cpu_env_ids_all = wp.zeros(N, dtype=wp.int32, device="cpu", pinned=True) + wp.copy(self._cpu_env_ids_all, self._ALL_INDICES) + self._cpu_env_mask = wp.zeros(N, dtype=wp.bool, device="cpu", pinned=True) def _process_cfg(self) -> None: - """Process the articulation configuration (initial state, soft limits, etc.).""" + """Populate default state buffers from the config (mirrors RigidObject and Newton Articulation).""" cfg = self.cfg N = self._num_instances D = self._num_joints @@ -1814,92 +3692,36 @@ def _process_cfg(self) -> None: default_root_vel = tuple(cfg.init_state.lin_vel) + tuple(cfg.init_state.ang_vel) np_pose = np.tile(np.array(default_root_pose, dtype=np.float32), (N, 1)) np_vel = np.tile(np.array(default_root_vel, dtype=np.float32), (N, 1)) - wp.copy( - self._data._default_root_pose, - wp.from_numpy(np_pose, dtype=wp.transformf, device=dev), - ) - wp.copy( - self._data._default_root_vel, - wp.from_numpy(np_vel, dtype=wp.spatial_vectorf, device=dev), - ) + self._data.default_root_pose = wp.array(np_pose, dtype=wp.transformf, device=dev) + self._data.default_root_vel = wp.array(np_vel, dtype=wp.spatial_vectorf, device=dev) # Default joint positions / velocities from config patterns. + # cfg.init_state.joint_pos is a dict[str, float] where keys are regex patterns + # matching joint names. We expand this into a (N, D) buffer. self._resolve_joint_values(cfg.init_state.joint_pos, self._data._default_joint_pos) self._resolve_joint_values(cfg.init_state.joint_vel, self._data._default_joint_vel) - # Keep soft-limit computation on-device, matching the PhysX/Newton path. - wp.launch( - update_soft_joint_pos_limits, - dim=(N, D), - inputs=[self._data.joint_pos_limits, cfg.soft_joint_pos_limit_factor], - outputs=[self._data._soft_joint_pos_limits], - device=dev, - ) - - def _invalidate_initialize_callback(self, event) -> None: - self._is_initialized = False - - def _process_actuators_cfg(self) -> None: - """Build actuator instances from the config and write drive properties to PhysX. + # Compute soft joint position limits from the hard limits read from the binding + # (or zeros if no joints). This matches the PhysX/Newton path. + if D > 0: + wp.launch( + update_soft_joint_pos_limits, + dim=(N, D), + inputs=[self._data.joint_pos_limits, cfg.soft_joint_pos_limit_factor], + outputs=[self._data._soft_joint_pos_limits], + device=dev, + ) - Mirrors what the legacy PhysX backend does in its own _process_actuators_cfg: - - For ImplicitActuator: write the configured stiffness / damping to the PhysX - drive so the solver uses exactly the values from the actuator config. - - For all explicit actuators: zero out PhysX stiffness / damping so the - USD-authored drive gains cannot interfere with the explicit torque path. - - For all actuators: write effort_limit_sim and velocity_limit_sim. + def _process_tendons(self) -> None: + """Discover tendon counts from binding metadata and names from USD. - These writes happen via TensorBinding (GPU-resident) after warmup has - allocated the GPU buffers (MODEL_INIT fires post-warmup). + Tendon counts come from the ovphysx binding metadata. Tendon names are + recovered from the exported USD articulation subtree because ovphysx + exposes joint names/counts, but not the per-joint USD paths that the + PhysX backend can query directly. """ - from isaaclab.actuators import ImplicitActuator - - self.actuators: dict[str, ActuatorBase] = {} - self._has_implicit_actuators = False - for name, act_cfg in self.cfg.actuators.items(): - joint_ids, joint_names = self.find_joints(act_cfg.joint_names_expr) - if not joint_ids: - logger.warning("Actuator '%s': no joints matched '%s'", name, act_cfg.joint_names_expr) - continue - act_cfg_copy = act_cfg.copy() - act = act_cfg_copy.class_type( - act_cfg_copy, - joint_names=joint_names, - joint_ids=joint_ids, - num_envs=self._num_instances, - device=self._device, - ) - self.actuators[name] = act - self._joint_ids_per_actuator[name] = joint_ids - - # Write drive gains and limits to PhysX to match the actuator config. - # Without this, PhysX retains whatever stiffness/damping was authored in the - # USD file, which can produce large restoring forces if the USD gains differ - # from the actuator config (e.g. a position-controlled robot exported with - # non-zero drive stiffness but configured with ImplicitActuator(stiffness=0)). - jids = list(joint_ids) - if isinstance(act, ImplicitActuator): - self._has_implicit_actuators = True - stiffness = act.stiffness # torch (N, J) - damping = act.damping # torch (N, J) - else: - stiffness = wp.zeros((self._num_instances, len(jids)), dtype=wp.float32, device=self._device) - damping = wp.zeros((self._num_instances, len(jids)), dtype=wp.float32, device=self._device) - self.write_joint_stiffness_to_sim_index(stiffness=stiffness, joint_ids=jids) - self.write_joint_damping_to_sim_index(damping=damping, joint_ids=jids) - self.write_joint_effort_limit_to_sim_index(limits=act.effort_limit_sim, joint_ids=jids) - self.write_joint_velocity_limit_to_sim_index(limits=act.velocity_limit_sim, joint_ids=jids) - - def _process_tendons(self) -> None: - """Discover tendon counts from binding metadata and names from USD. - - Tendon counts come from the ovphysx binding metadata. Tendon names are - recovered from the exported USD articulation subtree because ovphysx - exposes joint names/counts, but not the per-joint USD paths that the - PhysX backend can query directly. - """ - self._fixed_tendon_names = [] - self._spatial_tendon_names = [] + self._fixed_tendon_names = [] + self._spatial_tendon_names = [] sample = next(iter(self._bindings.values())) self._num_fixed_tendons = getattr(sample, "fixed_tendon_count", 0) @@ -1909,7 +3731,7 @@ def _process_tendons(self) -> None: stage_path = OvPhysxManager._stage_path if stage_path is not None: try: - from pxr import Usd, UsdPhysics + from pxr import Usd from isaaclab.sim.utils.queries import get_all_matching_child_prims @@ -1946,55 +3768,164 @@ def _process_tendons(self) -> None: except Exception: logger.debug("Could not parse USD stage for tendon names at %s", stage_path) - self._data._num_fixed_tendons = self._num_fixed_tendons - self._data._num_spatial_tendons = self._num_spatial_tendons - self._data.fixed_tendon_names = self._fixed_tendon_names - self._data.spatial_tendon_names = self._spatial_tendon_names + # Push the discovered names into the data container if it already exists. + # (During _initialize_impl the data container is created AFTER _process_tendons, + # so self._data may not be set yet; the ArticulationData constructor receives the + # counts/names directly.) + if hasattr(self, "_data"): + self._data._num_fixed_tendons = self._num_fixed_tendons + self._data._num_spatial_tendons = self._num_spatial_tendons + self._data.fixed_tendon_names = self._fixed_tendon_names + self._data.spatial_tendon_names = self._spatial_tendon_names - def _apply_external_wrenches(self) -> None: - """Compose and write external wrenches to the LINK_WRENCH binding. + def _get_binding(self, tensor_type: int): + """Return a cached TensorBinding, creating it on first access. - WrenchComposer accumulates forces/torques in body (link) frame. - The LINK_WRENCH binding expects world-frame [fx,fy,fz,tx,ty,tz,px,py,pz]. - We rotate the body-frame vectors to world frame using the link quaternion - and pack them into the [N, L, 9] tensor with application position = origin. + Bindings are lightweight handles (a pointer + shape metadata into + PhysX's shared GPU buffer). Creating one does NOT allocate new GPU + memory -- the underlying simulation buffers are allocated once by PhysX + regardless of how many bindings point into them. Still, we defer + creation so that tensor types the user never queries are never looked up. + + Args: + tensor_type: The TensorType constant identifying which simulation + buffer to bind (e.g. :attr:`~isaaclab_ovphysx.tensor_types.ROOT_POSE`). + + Returns: + A TensorBinding object, or ``None`` if the binding could not be created. """ - inst = self._instantaneous_wrench_composer - perm = self._permanent_wrench_composer - if not inst.active and not perm.active: - return - if inst.active: - if perm.active: - inst.add_forces_and_torques_index( - forces=perm.composed_force, - torques=perm.composed_torque, - body_ids=list(range(self._num_bodies)), - env_ids=list(range(self._num_instances)), - ) - force_b = inst.composed_force - torque_b = inst.composed_torque - else: - force_b = perm.composed_force - torque_b = perm.composed_torque + binding = self._bindings.get(tensor_type) + if binding is not None: + return binding + try: + binding = self._ovphysx.create_tensor_binding(pattern=self._binding_pattern, tensor_type=tensor_type) + self._bindings[tensor_type] = binding + return binding + except Exception: + logger.debug("Could not create tensor binding for type %s", tensor_type) + return None - poses = self._data.body_link_pose_w - wp.launch( - _body_wrench_to_world, - dim=(self._num_instances, self._num_bodies), - inputs=[force_b, torque_b, poses], - outputs=[self._wrench_buf], - device=self._device, - ) - wrench_binding = self._get_binding(TT.LINK_WRENCH) - if wrench_binding is not None: - wrench_binding.write(self._wrench_buf) - inst.reset() + def _resolve_joint_values(self, pattern_dict: dict[str, float], buffer: wp.array) -> None: + """Resolve a ``{pattern: value}`` dict into a per-joint buffer. - def _apply_actuator_model(self) -> None: - """Run the actuator model to compute torques from user targets. + Builds values on CPU then copies to buffer's device (GPU arrays' + ``.numpy()`` returns a read-only copy, not a writable view). - IsaacLab actuators are torch-based. We convert warp -> torch via - DLPack (zero-copy on GPU), run the actuator, then write results back. + Args: + pattern_dict: A mapping from regex pattern strings to scalar values. + Matches joint names returned by :attr:`joint_names`. + buffer: Target warp array of shape ``(num_instances, num_joints)`` + to populate. + """ + buf_np = buffer.numpy() + modified = False + for pattern, value in pattern_dict.items(): + for j, name in enumerate(self._joint_names): + if re.fullmatch(pattern, name): + buf_np[:, j] = value + modified = True + if modified: + wp.copy(buffer, wp.from_numpy(buf_np, dtype=buffer.dtype, device=str(buffer.device))) + + def _n_envs_index(self, env_ids) -> int: + """Return the number of environments from an ``env_ids`` argument.""" + if env_ids is None: + return self._num_instances + if isinstance(env_ids, (list, tuple)): + return len(env_ids) + return env_ids.shape[0] if hasattr(env_ids, "shape") else len(env_ids) + + def _nft(self) -> int: + """Return the number of fixed tendons (0 if none).""" + return self._num_fixed_tendons + + def _nst(self) -> int: + """Return the number of spatial tendons (0 if none).""" + return self._num_spatial_tendons + + """ + Internal simulation callbacks. + """ + + def _invalidate_initialize_callback(self, event) -> None: + """Invalidate the asset on simulation reset.""" + super()._invalidate_initialize_callback(event) + + """ + Internal helpers -- Actuators. + """ + + def _process_actuators_cfg(self) -> None: + """Build actuator instances from the config and write drive properties to PhysX. + + Mirrors the PhysX backend's ``_process_actuators_cfg``: + + * For :class:`~isaaclab.actuators.ImplicitActuator`: write the configured + stiffness/damping to the PhysX drive so the solver uses exactly those values. + * For all explicit actuators: zero out PhysX stiffness/damping so USD-authored + drive gains cannot interfere with the explicit torque path. + * For all actuators: write :attr:`~isaaclab.actuators.ActuatorBase.effort_limit_sim` + and :attr:`~isaaclab.actuators.ActuatorBase.velocity_limit_sim`. + """ + from isaaclab.actuators import ImplicitActuator + + self.actuators: dict[str, Any] = {} + self._has_implicit_actuators = False + for name, act_cfg in self.cfg.actuators.items(): + joint_ids, joint_names = self.find_joints(act_cfg.joint_names_expr) + if not joint_ids: + logger.warning("Actuator '%s': no joints matched '%s'", name, act_cfg.joint_names_expr) + continue + act_cfg_copy = act_cfg.copy() + # seed the actuator with the simulation's already-correct DOF defaults + # (USD-authored ``physxJoint:maxJointVelocity`` etc. parsed at scene-load). + # Without these the ActuatorBase constructor falls back to ``inf`` for unset + # cfg fields, and the ``write_joint_*_to_sim_index`` calls below then + # overwrite the correct values with ``inf``. + act = act_cfg_copy.class_type( + act_cfg_copy, + joint_names=joint_names, + joint_ids=joint_ids, + num_envs=self._num_instances, + device=self._device, + stiffness=self._data.joint_stiffness.torch[:, joint_ids], + damping=self._data.joint_damping.torch[:, joint_ids], + armature=self._data.joint_armature.torch[:, joint_ids], + friction=self._data.joint_friction_coeff.torch[:, joint_ids], + dynamic_friction=self._data.joint_dynamic_friction_coeff.torch[:, joint_ids], + viscous_friction=self._data.joint_viscous_friction_coeff.torch[:, joint_ids], + effort_limit=self._data.joint_effort_limits.torch[:, joint_ids].clone(), + velocity_limit=self._data.joint_vel_limits.torch[:, joint_ids], + ) + self.actuators[name] = act + self._joint_ids_per_actuator[name] = joint_ids + + # Write drive gains and limits to PhysX to match the actuator config. + # Without this, PhysX retains whatever stiffness/damping was authored in the + # USD file, which can produce large restoring forces when the USD gains differ + # from the actuator config. + jids = list(joint_ids) + if isinstance(act, ImplicitActuator): + self._has_implicit_actuators = True + stiffness = act.stiffness # torch (N, J) + damping = act.damping # torch (N, J) + else: + stiffness = wp.zeros((self._num_instances, len(jids)), dtype=wp.float32, device=self._device) + damping = wp.zeros((self._num_instances, len(jids)), dtype=wp.float32, device=self._device) + self.write_joint_stiffness_to_sim_index(stiffness=stiffness, joint_ids=jids) + self.write_joint_damping_to_sim_index(damping=damping, joint_ids=jids) + self.write_joint_effort_limit_to_sim_index(limits=act.effort_limit_sim, joint_ids=jids) + self.write_joint_velocity_limit_to_sim_index(limits=act.velocity_limit_sim, joint_ids=jids) + + def _apply_actuator_model(self) -> None: + """Run the actuator model to compute joint torques from user-supplied targets. + + IsaacLab actuators are torch-based. The method converts Warp buffers to + torch via DLPack (zero-copy on GPU), runs each actuator's + :meth:`~isaaclab.actuators.ActuatorBase.compute` method, then writes the + computed effort back to the private ``_computed_torque`` / ``_applied_torque`` + buffers of the data container. :meth:`write_data_to_sim` then pushes + ``_applied_torque`` to the ``DOF_ACTUATION_FORCE`` binding in one shot. """ from isaaclab.utils.types import ArticulationActions @@ -2005,7 +3936,7 @@ def _apply_actuator_model(self) -> None: jids_t = jids if isinstance(jids, list) else list(jids) all_joints = len(jids_t) == self._num_joints - # warp -> torch (zero-copy on same device via DLPack) + # Warp -> torch (zero-copy on same device via DLPack). jp_target_full = self._data.joint_pos_target.torch jv_target_full = self._data.joint_vel_target.torch je_target_full = self._data.joint_effort_target.torch @@ -2036,640 +3967,380 @@ def _apply_actuator_model(self) -> None: ct[:, jids_t] = act.computed_effort at[:, jids_t] = act.applied_effort + """ + Internal helpers -- Debugging. + """ + def _validate_cfg(self) -> None: - pass + """Validate the configuration after processing. - def _log_articulation_info(self) -> None: - """Log information about the articulation. - - .. note:: We purposefully read the values from the simulator to ensure that the values are configured as - expected. - """ - from prettytable import PrettyTable - - def format_large_number(_, v: float) -> str: - if abs(v) >= 1e3: - return f"{v:.1e}" - return f"{v:.3f}" - - def format_limits(_, v: tuple[float, float]) -> str: - if abs(v[0]) >= 1e3 or abs(v[1]) >= 1e3: - return f"[{v[0]:.1e}, {v[1]:.1e}]" - return f"[{v[0]:.3f}, {v[1]:.3f}]" - - stiffnesses = self.data.joint_stiffness.warp.numpy()[0].tolist() - dampings = self.data.joint_damping.warp.numpy()[0].tolist() - armatures = self.data.joint_armature.warp.numpy()[0].tolist() - frictions = self.data.joint_friction_coeff.warp.numpy()[0].tolist() - pos_limits_np = self.data.joint_pos_limits.warp.numpy().reshape(self._num_instances, self._num_joints, 2) - position_limits = [tuple(pos_limits_np[0, j].tolist()) for j in range(self._num_joints)] - velocity_limits = self.data.joint_vel_limits.warp.numpy()[0].tolist() - effort_limits = self.data.joint_effort_limits.warp.numpy()[0].tolist() - - joint_table = PrettyTable() - joint_table.title = f"Simulation Joint Information (Prim path: {self.cfg.prim_path})" - joint_table.field_names = [ - "Index", - "Name", - "Stiffness", - "Damping", - "Armature", - "Friction", - "Position Limits", - "Velocity Limits", - "Effort Limits", - ] - joint_table.custom_format["Stiffness"] = format_large_number - joint_table.custom_format["Damping"] = format_large_number - joint_table.custom_format["Armature"] = format_large_number - joint_table.custom_format["Friction"] = format_large_number - joint_table.custom_format["Position Limits"] = format_limits - joint_table.custom_format["Velocity Limits"] = format_large_number - joint_table.custom_format["Effort Limits"] = format_large_number - joint_table.align["Name"] = "l" - - for index, name in enumerate(self.joint_names): - joint_table.add_row( - [ - index, - name, - stiffnesses[index], - dampings[index], - armatures[index], - frictions[index], - position_limits[index], - velocity_limits[index], - effort_limits[index], - ] - ) - logger.info(f"Simulation parameters for joints in {self.cfg.prim_path}:\n" + joint_table.get_string()) - - if self.num_fixed_tendons > 0: - ft_stiffnesses = self.data.fixed_tendon_stiffness.warp.numpy()[0].tolist() - ft_dampings = self.data.fixed_tendon_damping.warp.numpy()[0].tolist() - ft_limit_stiffnesses = self.data.fixed_tendon_limit_stiffness.warp.numpy()[0].tolist() - ft_limits_np = self.data.fixed_tendon_pos_limits.warp.numpy().reshape( - self._num_instances, self.num_fixed_tendons, 2 - ) - ft_limits = [tuple(ft_limits_np[0, t].tolist()) for t in range(self.num_fixed_tendons)] - ft_rest_lengths = self.data.fixed_tendon_rest_length.warp.numpy()[0].tolist() - ft_offsets = self.data.fixed_tendon_offset.warp.numpy()[0].tolist() - - tendon_table = PrettyTable() - tendon_table.title = f"Simulation Fixed Tendon Information (Prim path: {self.cfg.prim_path})" - tendon_table.field_names = [ - "Index", - "Stiffness", - "Damping", - "Limit Stiffness", - "Limits", - "Rest Length", - "Offset", - ] - tendon_table.custom_format["Stiffness"] = format_large_number - tendon_table.custom_format["Damping"] = format_large_number - tendon_table.custom_format["Limit Stiffness"] = format_large_number - tendon_table.custom_format["Limits"] = format_limits - tendon_table.custom_format["Rest Length"] = format_large_number - tendon_table.custom_format["Offset"] = format_large_number - for index in range(self.num_fixed_tendons): - tendon_table.add_row( - [ - index, - ft_stiffnesses[index], - ft_dampings[index], - ft_limit_stiffnesses[index], - ft_limits[index], - ft_rest_lengths[index], - ft_offsets[index], - ] - ) - logger.info( - f"Simulation parameters for fixed tendons in {self.cfg.prim_path}:\n" + tendon_table.get_string() - ) + Mirrors :meth:`isaaclab_physx.assets.Articulation._validate_cfg` (raises + ``ValueError`` with a per-joint message when any default joint position + is outside ``[lower, upper]`` or any default joint velocity exceeds the + per-joint max velocity). Reads come from :attr:`ArticulationData` + accessors instead of PhysX's ``root_view.get_dof_limits`` / + ``get_dof_max_velocities`` because OVPhysX's ``root_view`` is the + per-tensor-type bindings dict. - if self.num_spatial_tendons > 0: - st_stiffnesses = self.data.spatial_tendon_stiffness.warp.numpy()[0].tolist() - st_dampings = self.data.spatial_tendon_damping.warp.numpy()[0].tolist() - st_limit_stiffnesses = self.data.spatial_tendon_limit_stiffness.warp.numpy()[0].tolist() - st_offsets = self.data.spatial_tendon_offset.warp.numpy()[0].tolist() - - tendon_table = PrettyTable() - tendon_table.title = f"Simulation Spatial Tendon Information (Prim path: {self.cfg.prim_path})" - tendon_table.field_names = [ - "Index", - "Stiffness", - "Damping", - "Limit Stiffness", - "Offset", - ] - tendon_table.float_format = ".3" - for index in range(self.num_spatial_tendons): - tendon_table.add_row( - [ - index, - st_stiffnesses[index], - st_dampings[index], - st_limit_stiffnesses[index], - st_offsets[index], - ] - ) - logger.info( - f"Simulation parameters for spatial tendons in {self.cfg.prim_path}:\n" + tendon_table.get_string() - ) + .. note:: + Must be called only after :meth:`_create_buffers` / + :meth:`_process_cfg` / :meth:`_process_actuators_cfg`, otherwise + limits and defaults may not yet reflect the final values. + """ + # check that the default joint positions are within the limits + joint_pos_limits = self._data.joint_pos_limits.torch[0] # (num_joints, 2) + default_joint_pos = self._data.default_joint_pos.torch[0] # (num_joints,) + out_of_range = default_joint_pos < joint_pos_limits[:, 0] + out_of_range |= default_joint_pos > joint_pos_limits[:, 1] + violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) + if len(violated_indices) > 0: + msg = "The following joints have default positions out of the limits: \n" + for idx in violated_indices: + joint_name = self._data.joint_names[idx] + joint_limit = joint_pos_limits[idx] + joint_pos = default_joint_pos[idx] + msg += f"\t- '{joint_name}': {joint_pos:.3f} not in [{joint_limit[0]:.3f}, {joint_limit[1]:.3f}]\n" + raise ValueError(msg) + + # check that the default joint velocities are within the limits + joint_max_vel = self._data.joint_vel_limits.torch[0] # (num_joints,) + default_joint_vel = self._data.default_joint_vel.torch[0] # (num_joints,) + out_of_range = torch.abs(default_joint_vel) > joint_max_vel + violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) + if len(violated_indices) > 0: + msg = "The following joints have default velocities out of the limits: \n" + for idx in violated_indices: + joint_name = self._data.joint_names[idx] + joint_limit = [-joint_max_vel[idx], joint_max_vel[idx]] + joint_vel = default_joint_vel[idx] + msg += f"\t- '{joint_name}': {joint_vel:.3f} not in [{joint_limit[0]:.3f}, {joint_limit[1]:.3f}]\n" + raise ValueError(msg) - """ - Internal helpers -- Bindings. - """ + def _log_articulation_info(self) -> None: + pass - def _get_binding(self, tensor_type: int): - """Return a cached TensorBinding, creating it on first access. + def _resolve_env_ids(self, env_ids) -> wp.array: + """Resolve environment indices to a warp int32 array on ``self._device`` (mirrors PhysX). - Bindings are lightweight handles (a pointer + shape metadata into - PhysX's shared GPU buffer). Creating one does NOT allocate new GPU - memory -- the underlying simulation buffers are allocated once by PhysX - regardless of how many bindings point into them. Still, we defer - creation so that tensor types the user never queries are never looked up. + Tests sometimes hand us indices on CPU even when the sim runs on GPU; we move the + resolved array onto ``self._device`` so kernel launches don't fail on a device + mismatch. """ - binding = self._bindings.get(tensor_type) - if binding is not None: - return binding - try: - binding = self._physx_instance.create_tensor_binding(pattern=self._binding_pattern, tensor_type=tensor_type) - self._bindings[tensor_type] = binding - return binding - except Exception: - logger.debug("Could not create tensor binding for type %s", tensor_type) - return None + if env_ids is None or env_ids == slice(None): + return self._ALL_INDICES + if isinstance(env_ids, list): + return wp.array(env_ids, dtype=wp.int32, device=self._device) + if isinstance(env_ids, torch.Tensor): + return wp.from_torch(env_ids.to(torch.int32), dtype=wp.int32) + 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 + + def _resolve_body_ids(self, body_ids) -> wp.array: + """Resolve body indices to a warp int32 array on ``self._device`` (mirrors PhysX).""" + 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 wp.from_torch(body_ids.to(torch.int32), dtype=wp.int32) + 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 + + def _resolve_joint_ids(self, joint_ids) -> wp.array: + """Resolve joint indices to a warp int32 array on ``self._device``.""" + 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 wp.from_torch(joint_ids.to(torch.int32), dtype=wp.int32) + 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 + + def _resolve_fixed_tendon_ids(self, tendon_ids) -> wp.array: + """Resolve fixed-tendon indices to a warp int32 array on ``self._device``.""" + 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 wp.from_torch(tendon_ids.to(torch.int32), dtype=wp.int32) + 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 + + def _resolve_spatial_tendon_ids(self, tendon_ids) -> wp.array: + """Resolve spatial-tendon indices to a warp int32 array on ``self._device``.""" + 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 wp.from_torch(tendon_ids.to(torch.int32), dtype=wp.int32) + 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 + + def _broadcast_scalar_to_2d( + self, value: float | torch.Tensor | wp.array, shape: tuple[int, int] + ) -> torch.Tensor | wp.array: + """Broadcast a scalar :class:`float` to a ``(rows, cols)`` torch ``float32`` tensor. + + Tendon and joint setters accept ``float | torch.Tensor | wp.array``; the + underlying ``shared_kernels.write_2d_data_to_buffer_*`` kernels only + accept 2D arrays. This helper expands a Python float into a constant + tensor on :attr:`_device`; tensor / warp inputs are returned as-is. + + Mirrors the PhysX backend's ``isinstance(value, float)`` branching, + which dispatches to ``articulation_kernels.float_data_to_buffer_with_*``. + OVPhysX does not have those scalar kernels, so we materialize the + broadcast on the Python side. - """ - Internal helpers -- Write. - """ + Args: + value: Scalar float or 2D tensor / warp array. + shape: ``(rows, cols)`` target shape used when broadcasting a + scalar. - def _to_flat_f32(self, data, target_shape: tuple[int, ...] | None = None) -> wp.array | np.ndarray: - """Ensure data is a contiguous float32 tensor suitable for binding I/O. + Returns: + A 2D :class:`torch.Tensor` on ``self._device`` if *value* was a + float; otherwise *value* unchanged. + """ + if isinstance(value, float): + return torch.full(shape, value, dtype=torch.float32, device=self._device) + return value - State tensor bindings (positions, velocities, poses) live on the - simulation device (GPU in GPU mode). We always return data on - self._device so the binding device check passes. + def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array: + """Resolve an environment mask to a ``wp.bool`` array on ``self._device``. - For structured warp dtypes (transformf, spatial_vectorf, etc.) a - zero-copy flat float32 view is created instead of roundtripping - through CPU numpy. + OVPhysX (like Newton) uses the binding's native ``binding.write(mask=...)`` path, + so the mask is preserved end-to-end; no ``torch.nonzero`` conversion is needed. + ``None`` returns the pre-allocated all-true mask. """ - dev = self._device - if isinstance(data, wp.array): - if str(data.device) != dev: - data = wp.clone(data, device=dev) - if data.dtype == wp.float32: - return data - # Structured dtype: zero-copy flat float32 view. - # transformf -> [N, 7], spatial_vectorf -> [N, 6], etc. - floats_per_elem = data.strides[0] // 4 - return wp.array( - ptr=data.ptr, - shape=(data.shape[0], floats_per_elem), - dtype=wp.float32, - device=dev, - copy=False, - ) - elif isinstance(data, torch.Tensor): - if data.is_cuda and dev.startswith("cuda"): - return wp.from_torch(data.detach().contiguous().float()) - np_data = data.detach().cpu().numpy().astype(np.float32) - return wp.from_numpy(np_data, dtype=wp.float32, device=dev) - elif isinstance(data, np.ndarray): - return wp.from_numpy(data.astype(np.float32), dtype=wp.float32, device=dev) - elif isinstance(data, (int, float)): - return wp.from_numpy(np.array(data, dtype=np.float32), dtype=wp.float32, device=dev) - return wp.from_numpy(np.asarray(data, dtype=np.float32), dtype=wp.float32, device=dev) - - def _as_gpu_f32_2d(self, data, cols: int) -> wp.array: - """View/convert data as 2D [rows, cols] float32 on self._device. - - For warp arrays with structured dtypes (transformf, spatial_vectorf), - creates a zero-copy flat float32 view. For torch/numpy, converts to - warp on the simulation 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 + + 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 + + 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 + + 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 + + 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 + + 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. + + :paramref:`env_mask` is normally on ``self._device``; ``binding.write(mask=...)`` + requires the mask on the binding's device, which is CPU for mass / CoMs / inertia. + Reuses the pre-allocated ``_cpu_env_mask`` pinned buffer. """ - dev = self._device - if isinstance(data, wp.array): - if str(data.device) != dev: - data = wp.clone(data, device=dev) - if data.dtype == wp.float32 and data.ndim == 2: - return data - n = data.shape[0] - return wp.array( - ptr=data.ptr, - shape=(n, cols), - dtype=wp.float32, - device=dev, - copy=False, - ) - if isinstance(data, torch.Tensor) and data.is_cuda and dev.startswith("cuda"): - return wp.from_torch(data.detach().contiguous().float().reshape(-1, cols)) - np_data = self._to_cpu_numpy(data).reshape(-1, cols) - return wp.from_numpy(np_data, dtype=wp.float32, device=dev) - - def _get_write_scratch(self, tensor_type: int, binding) -> wp.array: - """Return a cached GPU scratch buffer for read-modify-write.""" - if not hasattr(self, "_write_scratch"): - self._write_scratch = {} - buf = self._write_scratch.get(tensor_type) - if buf is None: - buf = wp.zeros(binding.shape, dtype=wp.float32, device=self._device) - self._write_scratch[tensor_type] = buf - return buf - - def _write_root_state(self, tensor_type: int, data, env_ids=None, mask=None, _ids_gpu=None) -> None: - """GPU-native write for root pose [N,7] or velocity [N,6]. - - Three paths, fastest first: - - Full write (no env_ids, no mask): zero-copy DLPack. - - Indexed write with full-size data: zero-copy view + indices. - The binding API only copies the indexed rows from the full buffer, - so no read-modify-write is needed when data is already [N,...]. - - Indexed write with partial data [K,...]: scatter kernel into a GPU - scratch buffer, then write with indices. - - Masked write: data is always full [N,...], pass directly with mask. - - Args: - _ids_gpu: Pre-converted GPU warp int32 array of env indices. - When provided, skips the per-call GPU->CPU->GPU conversion - of env_ids. - """ - binding = self._get_binding(tensor_type) - if binding is None: - return - N, C = binding.shape - - if env_ids is None and _ids_gpu is None and mask is None: - binding.write(self._to_flat_f32(data)) - self._invalidate_root_caches(tensor_type) - return - - src = self._as_gpu_f32_2d(data, C) - - if env_ids is not None or _ids_gpu is not None: - if _ids_gpu is None: - _ids_gpu = self._env_ids_to_gpu_warp(env_ids) - K = _ids_gpu.shape[0] - if src.shape[0] == N: - binding.write(src, indices=_ids_gpu) - else: - scratch = self._get_write_scratch(tensor_type, binding) - binding.read(scratch) - wp.launch( - _scatter_rows_partial, - dim=(K, C), - inputs=[scratch, src, _ids_gpu], - device=self._device, - ) - binding.write(scratch, indices=_ids_gpu) - else: - mask_u8 = wp.from_numpy( - self._to_cpu_numpy(mask).astype(np.uint8), - device=self._device, - ) - binding.write(src, mask=mask_u8) - self._invalidate_root_caches(tensor_type) - - def _invalidate_root_caches(self, tensor_type: int) -> None: - """Force re-read from GPU on next property access after a binding write.""" - if tensor_type == TT.ROOT_POSE: - self.data._root_link_pose_w.timestamp = -1.0 - self.data._root_com_pose_w.timestamp = -1.0 - elif tensor_type == TT.ROOT_VELOCITY: - self.data._root_link_vel_w.timestamp = -1.0 - self.data._root_com_vel_w.timestamp = -1.0 - - def _write_flat_tensor(self, tensor_type: int, data, env_ids=None, joint_ids=None, _ids_gpu=None) -> None: - """Write a 2-D tensor to a binding, with optional env/joint index subsetting.""" - if isinstance(data, (int, float)): - return - binding = self._get_binding(tensor_type) - if binding is None: - return - from isaaclab_ovphysx.tensor_types import _CPU_ONLY_TYPES - - is_cpu_only = tensor_type in _CPU_ONLY_TYPES - - # CPU-only types or column scatter must go through numpy. - if is_cpu_only or joint_ids is not None: - target_device = "cpu" if is_cpu_only else self._device - np_data = self._to_cpu_numpy(data) - if joint_ids is not None: - if is_cpu_only: - full = np.zeros(binding.shape, dtype=np.float32) - binding.read(full) - else: - scratch = self._get_write_scratch(tensor_type, binding) - binding.read(scratch) - full = scratch.numpy() - jids = self._to_cpu_indices(joint_ids, np.intp) - if env_ids is not None: - eids = self._to_cpu_indices(env_ids, np.intp) - full[np.ix_(eids, jids)] = np_data.reshape(len(eids), len(jids), *np_data.shape[2:]) - else: - full[:, jids] = np_data.reshape(full.shape[0], len(jids), *np_data.shape[2:]) - binding.write(wp.from_numpy(full, dtype=wp.float32, device=target_device)) - elif env_ids is not None: - if is_cpu_only: - full = np.zeros(binding.shape, dtype=np.float32) - binding.read(full) - else: - scratch = self._get_write_scratch(tensor_type, binding) - binding.read(scratch) - full = scratch.numpy() - eids = self._to_cpu_indices(env_ids, np.intp) - full[eids] = np_data if np_data.shape[0] == len(eids) else np_data[eids] - flat = wp.from_numpy(full.astype(np.float32), dtype=wp.float32, device=target_device) - idx = _ids_gpu if _ids_gpu is not None else self._env_ids_to_gpu_warp(env_ids) - binding.write(flat, indices=idx) - else: - binding.write(wp.from_numpy(np_data.astype(np.float32), dtype=wp.float32, device=target_device)) - return - - # GPU path: data stays on device. - if env_ids is None and _ids_gpu is None: - binding.write(self._to_flat_f32(data)) - return - - N, C = binding.shape[0], binding.shape[1] - src = self._as_gpu_f32_2d(data, C) - if _ids_gpu is None: - _ids_gpu = self._env_ids_to_gpu_warp(env_ids) - K = _ids_gpu.shape[0] - if src.shape[0] == N: - binding.write(src, indices=_ids_gpu) - else: - scratch = self._get_write_scratch(tensor_type, binding) - binding.read(scratch) - wp.launch( - _scatter_rows_partial, - dim=(K, C), - inputs=[scratch, src, _ids_gpu], - device=self._device, - ) - binding.write(scratch, indices=_ids_gpu) - - def _write_flat_tensor_mask(self, tensor_type: int, data, env_mask=None, joint_mask=None) -> None: - """Write a 2-D tensor to a binding, with optional env/joint mask subsetting.""" - if isinstance(data, (int, float)): - return - binding = self._get_binding(tensor_type) - if binding is None: - return - from isaaclab_ovphysx.tensor_types import _CPU_ONLY_TYPES - - is_cpu_only = tensor_type in _CPU_ONLY_TYPES - - # CPU-only types or column-mask scatter must go through numpy. - if is_cpu_only or joint_mask is not None: - target_device = "cpu" if is_cpu_only else self._device - np_data = self._to_cpu_numpy(data) - if joint_mask is not None: - # GPU bindings cannot read into numpy directly; read into GPU - # scratch first, then pull to CPU for column scatter. - if is_cpu_only: - full = np.zeros(binding.shape, dtype=np.float32) - binding.read(full) - else: - scratch = self._get_write_scratch(tensor_type, binding) - binding.read(scratch) - full = scratch.numpy() - jmask = self._to_cpu_numpy(joint_mask).astype(bool) - cols = np.where(jmask)[0] - if env_mask is not None: - emask = self._to_cpu_numpy(env_mask).astype(bool) - rows = np.where(emask)[0] - full[rows[:, None], cols] = np_data[rows[:, None], cols] - else: - full[:, cols] = np_data[:, cols] - binding.write(wp.from_numpy(full.astype(np.float32), dtype=wp.float32, device=target_device)) - elif env_mask is not None: - flat = wp.from_numpy(np_data.astype(np.float32), dtype=wp.float32, device=target_device) - mask_u8 = wp.from_numpy( - self._to_cpu_numpy(env_mask).astype(np.uint8), - device=target_device, - ) - binding.write(flat, mask=mask_u8) - else: - binding.write(wp.from_numpy(np_data.astype(np.float32), dtype=wp.float32, device=target_device)) - return + wp.copy(self._cpu_env_mask, env_mask) + return self._cpu_env_mask - # GPU path: data stays on device. - if env_mask is None: - binding.write(self._to_flat_f32(data)) - return + def _get_cpu_env_ids(self, env_ids: wp.array | torch.Tensor) -> wp.array: + """Return CPU int32 indices, using the pre-allocated pinned ``_cpu_env_ids_all`` + fast path when *env_ids* matches ``_ALL_INDICES`` (PR #5329 pattern). + """ + if isinstance(env_ids, torch.Tensor): + env_ids = wp.from_torch(env_ids, dtype=wp.int32) + if env_ids.ptr == self._ALL_INDICES.ptr: + return self._cpu_env_ids_all + return wp.clone(env_ids, device="cpu") - # Data is full [N, D], the binding API selects rows via the mask. - mask_u8 = wp.from_numpy( - self._to_cpu_numpy(env_mask).astype(np.uint8), - device=self._device, + """ + Deprecated methods. + """ + + def write_root_state_to_sim( + self, + root_state: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Deprecated; use :meth:`write_root_link_pose_to_sim_index` and + :meth:`write_root_com_velocity_to_sim_index` instead. + + Args: + root_state: Root state [m, m, m, qw, qx, qy, qz, m/s, m/s, m/s, rad/s, rad/s, rad/s]. + Shape is (len(env_ids), 13) with dtype wp.float32. + env_ids: Environment indices. Defaults to None (all environments). + """ + warnings.warn( + "The function 'write_root_state_to_sim' will be deprecated in a future release. Please" + " use 'write_root_link_pose_to_sim_index' and 'write_root_com_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, ) - binding.write(self._to_flat_f32(data), mask=mask_u8) + self.write_root_link_pose_to_sim_index(root_pose=root_state[:, :7], env_ids=env_ids) + self.write_root_com_velocity_to_sim_index(root_velocity=root_state[:, 7:], env_ids=env_ids) - def _write_friction_column(self, data, env_ids=None, joint_ids=None) -> None: - """Write static friction coefficient into column 0 of DOF_FRICTION_PROPERTIES [N,D,3].""" - binding = self._get_binding(TT.DOF_FRICTION_PROPERTIES) - if binding is None: - return - full = np.zeros(binding.shape, dtype=np.float32) - binding.read(full) - if isinstance(data, (int, float)): - if env_ids is not None and joint_ids is not None: - eids = self._to_cpu_numpy(env_ids).astype(np.intp) - jids = self._to_cpu_indices(joint_ids, np.intp) - full[np.ix_(eids, jids, [0])] = data - elif env_ids is not None: - eids = self._to_cpu_numpy(env_ids).astype(np.intp) - full[eids, :, 0] = data - elif joint_ids is not None: - jids = self._to_cpu_indices(joint_ids, np.intp) - full[:, jids, 0] = data - else: - full[..., 0] = data - binding.write(wp.from_numpy(full.astype(np.float32), dtype=wp.float32, device="cpu")) - return - np_data = self._to_cpu_numpy(data) - if env_ids is not None and joint_ids is not None: - eids = self._to_cpu_numpy(env_ids).astype(np.intp) - jids = self._to_cpu_indices(joint_ids, np.intp) - full[np.ix_(eids, jids, [0])] = np_data.reshape(len(eids), len(jids), 1) - elif env_ids is not None: - eids = self._to_cpu_numpy(env_ids).astype(np.intp) - full[eids, :, 0] = np_data.reshape(len(eids), -1) - elif joint_ids is not None: - jids = self._to_cpu_indices(joint_ids, np.intp) - full[:, jids, 0] = np_data.reshape(full.shape[0], len(jids)) - else: - full[..., 0] = np_data.reshape(full.shape[0], full.shape[1]) - binding.write(wp.from_numpy(full.astype(np.float32), dtype=wp.float32, device="cpu")) + def write_root_com_state_to_sim( + self, + root_state: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Deprecated; use :meth:`write_root_com_pose_to_sim_index` and + :meth:`write_root_com_velocity_to_sim_index` instead. - def _write_friction_column_mask(self, data, env_mask=None, joint_mask=None) -> None: - """Write static friction coefficient via mask into column 0 of DOF_FRICTION_PROPERTIES.""" - binding = self._get_binding(TT.DOF_FRICTION_PROPERTIES) - if binding is None: - return - full = np.zeros(binding.shape, dtype=np.float32) - binding.read(full) - if isinstance(data, (int, float)): - new_col = np.full((full.shape[0], full.shape[1]), data, dtype=np.float32) - else: - new_col = self._to_cpu_numpy(data).reshape(full.shape[0], full.shape[1]) - if env_mask is not None: - emask = self._to_cpu_numpy(env_mask).astype(bool) - if joint_mask is not None: - jmask = self._to_cpu_numpy(joint_mask).astype(bool) - rows = np.where(emask)[0] - cols = np.where(jmask)[0] - full[rows[:, None], cols, 0] = new_col[rows[:, None], cols] - else: - full[emask, :, 0] = new_col[emask] - elif joint_mask is not None: - jmask = self._to_cpu_numpy(joint_mask).astype(bool) - full[:, jmask, 0] = new_col[:, jmask] - else: - full[..., 0] = new_col - binding.write(wp.from_numpy(full.astype(np.float32), dtype=wp.float32, device="cpu")) - - def _write_joint_subset(self, tensor_type: int, buffer: wp.array, joint_ids: list[int]) -> None: - """Write a full-width joint buffer into the simulation for an actuator's joints.""" - binding = self._get_binding(tensor_type) - if binding is None: - return - if not hasattr(self, "_write_dltensor_cache"): - self._write_dltensor_cache = {} - cache_key = (tensor_type, buffer.ptr) - cached = self._write_dltensor_cache.get(cache_key) - if cached is None: - flat = self._to_flat_f32(buffer) - from ovphysx._dlpack_utils import acquire_dltensor - - dl, keepalive = acquire_dltensor(flat) - self._write_dltensor_cache[cache_key] = (dl, keepalive, flat) - cached = self._write_dltensor_cache[cache_key] - binding.write(cached[0]) - - @staticmethod - def _to_cpu_numpy(data) -> np.ndarray: - """Convert data (warp, torch, numpy, scalar) to a CPU numpy array.""" - if isinstance(data, wp.array): - return data.numpy().astype(np.float32) - if isinstance(data, torch.Tensor): - return data.detach().cpu().numpy().astype(np.float32) - return np.asarray(data, dtype=np.float32) - - @staticmethod - def _to_cpu_indices(data, dtype=np.int32) -> np.ndarray: - """Convert index array (warp, torch, list, numpy) to CPU numpy int array.""" - if isinstance(data, torch.Tensor): - return data.detach().cpu().numpy().astype(dtype) - if isinstance(data, wp.array): - return data.numpy().astype(dtype) - return np.asarray(data, dtype=dtype) - - def _env_ids_to_gpu_warp(self, env_ids) -> wp.array: - """Convert env_ids to a GPU int32 warp array, with single-entry caching. - - The cache avoids repeated GPU -> CPU -> GPU round-trips when the same - ``env_ids`` object is passed to multiple binding writes in a single step - (e.g. reset writes root_pose, root_vel, joint_pos, joint_vel). A new - object identity (``id()``) or shape change invalidates the cache. - """ - if hasattr(env_ids, "data_ptr"): - key = (env_ids.data_ptr(), env_ids.shape[0]) - elif isinstance(env_ids, wp.array): - key = (env_ids.ptr, env_ids.shape[0]) - else: - key = None - - if key is not None and hasattr(self, "_ids_cache_key") and self._ids_cache_key == key: - return self._ids_cache_val - - result = wp.array(self._to_cpu_indices(env_ids, np.int32), device=self._device) - if key is not None: - self._ids_cache_key = key - self._ids_cache_val = result - return result - - def _set_target_into_buffer(self, buffer: wp.array, data, env_ids=None, joint_ids=None) -> None: - """Set user-provided target data into a warp command buffer. - - For the common case (no index subset), this uses wp.copy to stay on - the simulation device. Subset writes (specific env_ids or joint_ids) - fall back to CPU because warp does not support scatter indexing. - """ - # Fast path: all-joints shortcut. When joint_ids covers every joint - # and env_ids is None, the subset is equivalent to a full copy. - if joint_ids is not None and env_ids is None: - n_joints = buffer.shape[1] if len(buffer.shape) > 1 else 1 - if hasattr(joint_ids, "__len__") and len(joint_ids) == n_joints: - joint_ids = None - if env_ids is None and joint_ids is None: - src = self._to_flat_f32(data) - if isinstance(src, np.ndarray): - src = wp.from_numpy(src, dtype=wp.float32, device=buffer.device) - wp.copy(buffer, src) - else: - np_data = self._to_cpu_numpy(data) - buf_np = buffer.numpy() - env_idx = self._to_cpu_numpy(env_ids).astype(np.intp) if env_ids is not None else None - jnt_idx = self._to_cpu_numpy(joint_ids).astype(np.intp) if joint_ids is not None else None - if env_idx is not None and jnt_idx is not None: - buf_np[np.ix_(env_idx, jnt_idx)] = np_data - elif env_idx is not None: - buf_np[env_idx] = np_data - else: - buf_np[:, jnt_idx] = np_data - wp.copy(buffer, wp.from_numpy(buf_np, dtype=wp.float32, device=buffer.device)) + Args: + root_state: Root CoM state [m, m, m, qw, qx, qy, qz, m/s, m/s, m/s, rad/s, rad/s, rad/s]. + Shape is (len(env_ids), 13) with dtype wp.float32. + env_ids: Environment indices. Defaults to None (all environments). + """ + warnings.warn( + "The function 'write_root_com_state_to_sim' will be deprecated in a future release. Please" + " use 'write_root_com_pose_to_sim_index' and 'write_root_com_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + self.write_root_com_pose_to_sim_index(root_pose=root_state[:, :7], env_ids=env_ids) + self.write_root_com_velocity_to_sim_index(root_velocity=root_state[:, 7:], env_ids=env_ids) - def _set_target_into_buffer_mask(self, buffer: wp.array, data, env_mask=None, joint_mask=None) -> None: - """Set user-provided target data into a warp command buffer using masks.""" - if env_mask is None: - src = self._to_flat_f32(data) - if isinstance(src, np.ndarray): - src = wp.from_numpy(src, dtype=wp.float32, device=buffer.device) - wp.copy(buffer, src) - else: - np_data = self._to_cpu_numpy(data) - buf_np = buffer.numpy() - mask_np = self._to_cpu_numpy(env_mask).astype(bool) - buf_np[mask_np] = np_data[mask_np] - wp.copy(buffer, wp.from_numpy(buf_np, dtype=wp.float32, device=buffer.device)) + def write_root_link_state_to_sim( + self, + root_state: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Deprecated; use :meth:`write_root_link_pose_to_sim_index` and + :meth:`write_root_link_velocity_to_sim_index` instead. - """ - Internal helpers -- Utilities. - """ + Args: + root_state: Root link state [m, m, m, qw, qx, qy, qz, m/s, m/s, m/s, rad/s, rad/s, rad/s]. + Shape is (len(env_ids), 13) with dtype wp.float32. + env_ids: Environment indices. Defaults to None (all environments). + """ + warnings.warn( + "The function 'write_root_link_state_to_sim' will be deprecated in a future release. Please" + " use 'write_root_link_pose_to_sim_index' and 'write_root_link_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + self.write_root_link_pose_to_sim_index(root_pose=root_state[:, :7], env_ids=env_ids) + self.write_root_link_velocity_to_sim_index(root_velocity=root_state[:, 7:], env_ids=env_ids) - def _n_envs_index(self, env_ids): - """Return the number of environments from an env_ids argument.""" - if env_ids is None: - return self._num_instances - if isinstance(env_ids, (list, tuple)): - return len(env_ids) - return env_ids.shape[0] if hasattr(env_ids, "shape") else len(env_ids) + def write_joint_state_to_sim( + self, + position: torch.Tensor | wp.array, + velocity: torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Deprecated combined joint-state write; use :meth:`write_joint_position_to_sim_index` + and :meth:`write_joint_velocity_to_sim_index` instead. - def _nft(self): - """Return the number of fixed tendons (0 if none).""" - return getattr(self, "_num_fixed_tendons", 0) + Args: + position: Joint positions [m or rad, depending on joint type]. Shape is + (len(env_ids), len(joint_ids)) with dtype wp.float32. + velocity: Joint velocities [m/s or rad/s, depending on joint type]. Shape is + (len(env_ids), len(joint_ids)) with dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). + """ + warnings.warn( + "write_joint_state_to_sim is deprecated; use write_joint_position_to_sim_index" + " and write_joint_velocity_to_sim_index instead.", + DeprecationWarning, + stacklevel=2, + ) + self.write_joint_position_to_sim_index(position=position, joint_ids=joint_ids, env_ids=env_ids) + self.write_joint_velocity_to_sim_index(velocity=velocity, joint_ids=joint_ids, env_ids=env_ids) - def _nst(self): - """Return the number of spatial tendons (0 if none).""" - return getattr(self, "_num_spatial_tendons", 0) + def write_joint_friction_coefficient_to_sim( + self, + joint_friction_coeff: torch.Tensor | wp.array | float, + joint_dynamic_friction_coeff: torch.Tensor | wp.array | float | None = None, + joint_viscous_friction_coeff: torch.Tensor | wp.array | float | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Deprecated, same as :meth:`write_joint_friction_coefficient_to_sim_index`.""" + warnings.warn( + "The function 'write_joint_friction_coefficient_to_sim' will be deprecated in a future release. Please" + " use 'write_joint_friction_coefficient_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + self.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff=joint_friction_coeff, + joint_dynamic_friction_coeff=joint_dynamic_friction_coeff, + joint_viscous_friction_coeff=joint_viscous_friction_coeff, + joint_ids=joint_ids, + env_ids=env_ids, + ) - def _resolve_joint_values(self, pattern_dict: dict[str, float], buffer: wp.array) -> None: - """Resolve a {pattern: value} dict into a per-joint buffer. + def write_joint_dynamic_friction_coefficient_to_sim( + self, + joint_dynamic_friction_coeff: torch.Tensor | wp.array | float, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Deprecated, same as :meth:`write_joint_dynamic_friction_coefficient_to_sim_index`.""" + warnings.warn( + "The function 'write_joint_dynamic_friction_coefficient_to_sim' will be deprecated in a future release. " + "Please use 'write_joint_dynamic_friction_coefficient_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + self.write_joint_dynamic_friction_coefficient_to_sim_index( + joint_dynamic_friction_coeff=joint_dynamic_friction_coeff, + joint_ids=joint_ids, + env_ids=env_ids, + ) - Builds values on CPU then copies to buffer's device (GPU arrays' - .numpy() returns a read-only copy, not a writable view). - """ - buf_np = buffer.numpy() - modified = False - for pattern, value in pattern_dict.items(): - for j, name in enumerate(self._joint_names): - if re.fullmatch(pattern, name): - buf_np[:, j] = value - modified = True - if modified: - wp.copy(buffer, wp.from_numpy(buf_np, dtype=buffer.dtype, device=str(buffer.device))) + def write_joint_viscous_friction_coefficient_to_sim( + self, + joint_viscous_friction_coeff: torch.Tensor | wp.array | float, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Deprecated, same as :meth:`write_joint_viscous_friction_coefficient_to_sim_index`.""" + warnings.warn( + "The function 'write_joint_viscous_friction_coefficient_to_sim' will be deprecated in a future release. " + "Please use 'write_joint_viscous_friction_coefficient_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + self.write_joint_viscous_friction_coefficient_to_sim_index( + joint_viscous_friction_coeff=joint_viscous_friction_coeff, + joint_ids=joint_ids, + env_ids=env_ids, + ) diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation_data.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation_data.py index e5be6c05328b..018fe0cb1507 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation_data.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation_data.py @@ -3,11 +3,11 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Articulation data backed by ovphysx TensorBindingsAPI.""" - from __future__ import annotations +import logging import warnings +from collections.abc import Callable from typing import Any import numpy as np @@ -16,23 +16,31 @@ from isaaclab.assets.articulation.base_articulation_data import BaseArticulationData from isaaclab.utils.buffers import TimestampedBufferWarp as TimestampedBuffer from isaaclab.utils.warp import ProxyArray +from isaaclab_ovphysx.physics import OvPhysxManager from isaaclab_ovphysx import tensor_types as TT - -from .kernels import ( - _compose_body_com_poses, +from isaaclab_ovphysx.assets.kernels import ( _compose_root_com_pose, _compute_heading, _copy_first_body, - _fd_joint_acc, _projected_gravity, _world_vel_to_body_ang, _world_vel_to_body_lin, + concat_body_pose_and_vel_to_state, + concat_root_pose_and_vel_to_state, + get_body_com_pose_from_body_link_pose, + get_body_link_vel_from_body_com_vel, + vec13f, ) +from .kernels import _fd_joint_acc + +# import logger +logger = logging.getLogger(__name__) + class ArticulationData(BaseArticulationData): - """Data container for an articulation backed by ovphysx tensor bindings. + """Data container for an articulation. This class contains the data for an articulation in the simulation. The data includes the state of the root rigid body, the state of all the bodies in the articulation, and the joint state. The data is @@ -48,84 +56,106 @@ class ArticulationData(BaseArticulationData): Depending on the settings, the two frames may not coincide with each other. In the robotics sense, the actor frame can be interpreted as the link frame. - Uses ovphysx :class:`TensorBinding` objects to lazily read simulation state into warp - arrays. Writes happen via the :class:`Articulation` class. + .. note:: + **Pull-to-refresh model.** OVPhysX state properties are *not* automatically updated each + simulation step. Each property getter pulls fresh data from the OVPhysX ``TensorBinding`` + on first access per timestamp, then caches the result until the next step. This differs + from the Newton backend, where buffers are refreshed automatically by the simulation. + + .. note:: + **CPU-only bindings.** OVPhysX exposes a subset of bindings (``BODY_MASS``, ``BODY_COM_POSE``, + ``BODY_INERTIA``, and most ``DOF_*`` property bindings) on CPU only. These are routed through + pinned-host staging buffers via :meth:`_binding_read` so that GPU-resident consumers see the + data without per-step host allocations. """ __backend_name__: str = "ovphysx" """The name of the backend for the articulation data.""" - def __init__(self, bindings: dict[int, Any], device: str, binding_getter=None): - """Initialize the articulation data. + def __init__( + self, + bindings: dict[int, Any], + device: str, + num_instances: int, + num_bodies: int, + num_joints: int, + num_fixed_tendons: int, + num_spatial_tendons: int, + body_names: list[str], + joint_names: list[str], + fixed_tendon_names: list[str], + spatial_tendon_names: list[str], + binding_getter: Callable[[int], Any] | None = None, + ) -> None: + """Initialize the articulation data container. Args: - bindings: Mapping from ovphysx tensor type constant to a - live TensorBinding for this articulation. - device: The compute device (``"cpu"`` or ``"cuda:N"``). + bindings: Dictionary of OVPhysX :class:`TensorBinding` objects keyed + by :class:`isaaclab_ovphysx.tensor_types.TensorType`. + device: Simulation device string (e.g., ``"cuda:0"`` or ``"cpu"``). + num_instances: Number of articulation instances. + num_bodies: Number of bodies per articulation. + num_joints: Number of degrees of freedom per articulation. + num_fixed_tendons: Number of fixed tendons per articulation. + num_spatial_tendons: Number of spatial tendons per articulation. + body_names: Ordered list of body names. + joint_names: Ordered list of joint names. + fixed_tendon_names: Ordered list of fixed tendon names. + spatial_tendon_names: Ordered list of spatial tendon names. binding_getter: Optional callable(tensor_type) -> TensorBinding that lazily creates bindings on first access. When provided, - ``_get_binding()`` delegates to this instead of only checking + :meth:`_get_binding` delegates to this instead of only checking the static ``bindings`` dict. """ super().__init__(root_view=None, device=device) self._bindings = bindings self._binding_getter = binding_getter + # counts and names are plain instance attributes (no @property indirection) + self.num_instances = num_instances + self.num_bodies = num_bodies + self.num_joints = num_joints + self.num_fixed_tendons = num_fixed_tendons + self.num_spatial_tendons = num_spatial_tendons + self.body_names = body_names + self.joint_names = joint_names + self.fixed_tendon_names = fixed_tendon_names + self.spatial_tendon_names = spatial_tendon_names + # private aliases used throughout _create_buffers and property bodies + self._num_instances = num_instances + self._num_bodies = num_bodies + self._num_joints = num_joints + self._num_fixed_tendons = num_fixed_tendons + self._num_spatial_tendons = num_spatial_tendons + + # Set initial time stamp self._sim_timestamp: float = 0.0 - self._is_primed = False - - # Metadata from an arbitrary articulation binding. - sample = next(iter(bindings.values())) - self._num_instances = sample.count - self._num_joints = sample.dof_count - self._num_bodies = sample.body_count - self._is_fixed_base = sample.is_fixed_base - - self.body_names = list(sample.body_names) - self.joint_names = list(sample.dof_names) - self.fixed_tendon_names: list[str] = [] - self.spatial_tendon_names: list[str] = [] - - self._num_fixed_tendons = 0 - self._num_spatial_tendons = 0 + self._is_primed: bool = False + # pinned-host staging buffers for CPU-only bindings (keyed by tensor_type) + self._cpu_staging_buffers: dict[int, wp.array] = {} + # scratch buffers for _get_read_view cache (keyed by (tensor_type, ptr)) + self._read_scratch: dict = {} - # Initialize parametric gravity and forward vectors (matching PhysX/Newton pattern). - # Guard against None sim context (e.g. mock/test environments). + # obtain gravity from the simulation configuration (fall back to standard + # gravity when the simulation has not been configured yet, e.g. in unit tests) + gravity = (0.0, 0.0, -9.81) from isaaclab.physics import PhysicsManager - gravity = (0.0, 0.0, -9.81) if PhysicsManager._sim is not None and hasattr(PhysicsManager._sim, "cfg"): gravity = PhysicsManager._sim.cfg.gravity gravity_np = np.array(gravity, dtype=np.float32) - gravity_mag = np.linalg.norm(gravity_np) + gravity_mag = float(np.linalg.norm(gravity_np)) if gravity_mag == 0.0: gravity_dir = np.array([0.0, 0.0, -1.0], dtype=np.float32) else: gravity_dir = gravity_np / gravity_mag - gravity_dir_tiled = np.tile(gravity_dir, (self._num_instances, 1)) - forward_tiled = np.tile(np.array([1.0, 0.0, 0.0], dtype=np.float32), (self._num_instances, 1)) + gravity_dir_tiled = np.tile(gravity_dir, (num_instances, 1)) + forward_tiled = np.tile(np.array([1.0, 0.0, 0.0], dtype=np.float32), (num_instances, 1)) + # Initialize constants self.GRAVITY_VEC_W = ProxyArray(wp.from_numpy(gravity_dir_tiled, dtype=wp.vec3f, device=device)) self.FORWARD_VEC_B = ProxyArray(wp.from_numpy(forward_tiled, dtype=wp.vec3f, device=device)) - def update(self, dt: float) -> None: - """Update the data for the articulation. - - Args: - dt: The time step for the update [s]. This must be a positive value. - """ - self._sim_timestamp += dt - - # Finite-difference joint acceleration from velocity. - if dt > 0.0 and self._previous_joint_vel is not None: - cur_vel = self.joint_vel - wp.launch( - _fd_joint_acc, - dim=(self._num_instances, self._num_joints), - inputs=[cur_vel, self._previous_joint_vel, 1.0 / dt], - outputs=[self._joint_acc.data], - device=self.device, - ) - self._joint_acc.timestamp = self._sim_timestamp + self._create_buffers() @property def is_primed(self) -> bool: @@ -149,21 +179,30 @@ def is_primed(self, value: bool) -> None: raise ValueError("The articulation data is already primed.") self._is_primed = True - """ - Names. - """ - - body_names: list[str] = None - """Body names in the order parsed by the simulation view.""" - - joint_names: list[str] = None - """Joint names in the order parsed by the simulation view.""" - - fixed_tendon_names: list[str] = None - """Fixed tendon names in the order parsed by the simulation view.""" + def update(self, dt: float) -> None: + """Updates the data for the articulation. - spatial_tendon_names: list[str] = None - """Spatial tendon names in the order parsed by the simulation view.""" + Args: + dt: The time step for the update. This must be a positive value. + """ + # update the simulation timestamp + self._sim_timestamp += dt + if not self._is_primed: + return + # trigger an update of the joint acceleration buffer via finite differencing + if dt > 0.0 and self._previous_joint_vel is not None: + cur_vel_buf = self._joint_vel_buf + # ensure joint vel buffer is fresh before differencing + self._read_binding_into_buf(TT.DOF_VELOCITY, cur_vel_buf) + wp.launch( + _fd_joint_acc, + dim=(self._num_instances, self._num_joints), + inputs=[cur_vel_buf.data, self._previous_joint_vel, 1.0 / dt], + outputs=[self._joint_acc.data], + device=self.device, + ) + self._joint_acc.timestamp = self._sim_timestamp + wp.copy(self._previous_joint_vel, cur_vel_buf.data) """ Defaults - Initial state. @@ -171,10 +210,12 @@ def is_primed(self, value: bool) -> None: @property def default_root_pose(self) -> ProxyArray: - """Default root pose ``[pos, quat]`` in the local environment frame. + """Default root pose ``[pos, quat]`` in local environment frame [m, -]. + + Shape is (num_instances,), dtype = wp.transformf. + In torch this resolves to (num_instances, 7). - The position and quaternion are of the articulation root's actor frame. - Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). + Populated from :attr:`ArticulationCfg.init_state` during initialisation. """ if self._default_root_pose_ta is None: self._default_root_pose_ta = ProxyArray(self._default_root_pose) @@ -185,21 +226,23 @@ def default_root_pose(self, value: wp.array) -> None: """Set the default root pose. Args: - value: The default root pose. Shape is (num_instances, 7). + value: The default root pose, shape (num_instances, 7). Raises: ValueError: If the articulation data is already primed. """ - if self.is_primed: + if self._is_primed: raise ValueError("The articulation data is already primed.") self._default_root_pose.assign(value) @property def default_root_vel(self) -> ProxyArray: - """Default root velocity ``[lin_vel, ang_vel]`` in the local environment frame. + """Default root velocity ``[lin_vel, ang_vel]`` in local environment frame [m/s, rad/s]. - The linear and angular velocities are of the articulation root's center of mass frame. - Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). + Shape is (num_instances,), dtype = wp.spatial_vectorf. + In torch this resolves to (num_instances, 6). + + Populated from :attr:`ArticulationCfg.init_state` during initialisation. """ if self._default_root_vel_ta is None: self._default_root_vel_ta = ProxyArray(self._default_root_vel) @@ -210,12 +253,12 @@ def default_root_vel(self, value: wp.array) -> None: """Set the default root velocity. Args: - value: The default root velocity. Shape is (num_instances, 6). + value: The default root velocity, shape (num_instances, 6). Raises: ValueError: If the articulation data is already primed. """ - if self.is_primed: + if self._is_primed: raise ValueError("The articulation data is already primed.") self._default_root_vel.assign(value) @@ -223,10 +266,7 @@ def default_root_vel(self, value: wp.array) -> None: def default_joint_pos(self) -> ProxyArray: """Default joint positions of all joints [m or rad, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). - - This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. + Shape is (num_instances, num_joints), dtype = wp.float32. """ if self._default_joint_pos_ta is None: self._default_joint_pos_ta = ProxyArray(self._default_joint_pos) @@ -237,12 +277,12 @@ def default_joint_pos(self, value: wp.array) -> None: """Set the default joint positions. Args: - value: The default joint positions. Shape is (num_instances, num_joints). + value: The default joint positions, shape (num_instances, num_joints). Raises: ValueError: If the articulation data is already primed. """ - if self.is_primed: + if self._is_primed: raise ValueError("The articulation data is already primed.") self._default_joint_pos.assign(value) @@ -250,10 +290,7 @@ def default_joint_pos(self, value: wp.array) -> None: def default_joint_vel(self) -> ProxyArray: """Default joint velocities of all joints [m/s or rad/s, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). - - This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. + Shape is (num_instances, num_joints), dtype = wp.float32. """ if self._default_joint_vel_ta is None: self._default_joint_vel_ta = ProxyArray(self._default_joint_vel) @@ -264,12 +301,12 @@ def default_joint_vel(self, value: wp.array) -> None: """Set the default joint velocities. Args: - value: The default joint velocities. Shape is (num_instances, num_joints). + value: The default joint velocities, shape (num_instances, num_joints). Raises: ValueError: If the articulation data is already primed. """ - if self.is_primed: + if self._is_primed: raise ValueError("The articulation data is already primed.") self._default_joint_vel.assign(value) @@ -281,12 +318,7 @@ def default_joint_vel(self, value: wp.array) -> None: def joint_pos_target(self) -> ProxyArray: """Joint position targets commanded by the user [m or rad, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). - - For an implicit actuator model, the targets are directly set into the simulation. - For an explicit actuator model, the targets are used to compute the joint torques - (see :attr:`applied_torque`), which are then set into the simulation. + Shape is (num_instances, num_joints), dtype = wp.float32. """ if self._joint_pos_target_ta is None: self._joint_pos_target_ta = ProxyArray(self._joint_pos_target) @@ -296,12 +328,7 @@ def joint_pos_target(self) -> ProxyArray: def joint_vel_target(self) -> ProxyArray: """Joint velocity targets commanded by the user [m/s or rad/s, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). - - For an implicit actuator model, the targets are directly set into the simulation. - For an explicit actuator model, the targets are used to compute the joint torques - (see :attr:`applied_torque`), which are then set into the simulation. + Shape is (num_instances, num_joints), dtype = wp.float32. """ if self._joint_vel_target_ta is None: self._joint_vel_target_ta = ProxyArray(self._joint_vel_target) @@ -311,12 +338,7 @@ def joint_vel_target(self) -> ProxyArray: def joint_effort_target(self) -> ProxyArray: """Joint effort targets commanded by the user [N or N*m, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). - - For an implicit actuator model, the targets are directly set into the simulation. - For an explicit actuator model, the targets are used to compute the joint torques - (see :attr:`applied_torque`), which are then set into the simulation. + Shape is (num_instances, num_joints), dtype = wp.float32. """ if self._joint_effort_target_ta is None: self._joint_effort_target_ta = ProxyArray(self._joint_effort_target) @@ -330,12 +352,7 @@ def joint_effort_target(self) -> ProxyArray: def computed_torque(self) -> ProxyArray: """Joint torques computed from the actuator model (before clipping) [N*m]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). - - This quantity is the raw torque output from the actuator model, before any clipping is applied. - It is exposed for users who want to inspect the computations inside the actuator model. - For instance, to penalize the learning agent for a difference between the computed and applied torques. + Shape is (num_instances, num_joints), dtype = wp.float32. """ if self._computed_torque_ta is None: self._computed_torque_ta = ProxyArray(self._computed_torque) @@ -345,11 +362,7 @@ def computed_torque(self) -> ProxyArray: def applied_torque(self) -> ProxyArray: """Joint torques applied from the actuator model (after clipping) [N*m]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). - - These torques are set into the simulation, after clipping the :attr:`computed_torque` based on the - actuator model. + Shape is (num_instances, num_joints), dtype = wp.float32. """ if self._applied_torque_ta is None: self._applied_torque_ta = ProxyArray(self._applied_torque) @@ -361,85 +374,132 @@ def applied_torque(self) -> ProxyArray: @property def joint_stiffness(self) -> ProxyArray: - """Joint stiffness provided to the simulation. + """Joint stiffness provided to the simulation [N*m/rad or N/m, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + Shape is (num_instances, num_joints), dtype = wp.float32. - In the case of explicit actuators, the value for the corresponding joints is zero. + Routed through pinned-host staging because ``DOF_STIFFNESS`` is a + CPU-only OVPhysX binding. """ + self._read_scalar_binding(TT.DOF_STIFFNESS, self._joint_stiffness) if self._joint_stiffness_ta is None: - self._joint_stiffness_ta = ProxyArray(self._joint_stiffness) + self._joint_stiffness_ta = ProxyArray(self._joint_stiffness.data) return self._joint_stiffness_ta @property def joint_damping(self) -> ProxyArray: - """Joint damping provided to the simulation. + """Joint damping provided to the simulation [N*m*s/rad or N*s/m, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + Shape is (num_instances, num_joints), dtype = wp.float32. - In the case of explicit actuators, the value for the corresponding joints is zero. + Routed through pinned-host staging because ``DOF_DAMPING`` is a + CPU-only OVPhysX binding. """ + self._read_scalar_binding(TT.DOF_DAMPING, self._joint_damping) if self._joint_damping_ta is None: - self._joint_damping_ta = ProxyArray(self._joint_damping) + self._joint_damping_ta = ProxyArray(self._joint_damping.data) return self._joint_damping_ta @property def joint_armature(self) -> ProxyArray: - """Joint armature provided to the simulation. + """Joint armature provided to the simulation [kg*m^2]. + + Shape is (num_instances, num_joints), dtype = wp.float32. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + Routed through pinned-host staging because ``DOF_ARMATURE`` is a + CPU-only OVPhysX binding. """ + self._read_scalar_binding(TT.DOF_ARMATURE, self._joint_armature) if self._joint_armature_ta is None: - self._joint_armature_ta = ProxyArray(self._joint_armature) + self._joint_armature_ta = ProxyArray(self._joint_armature.data) return self._joint_armature_ta @property def joint_friction_coeff(self) -> ProxyArray: - """Joint static friction coefficient provided to the simulation. + """Joint static friction coefficient [dimensionless]. + + Shape is (num_instances, num_joints), dtype = wp.float32. + Component ``[..., 0]`` of the ``DOF_FRICTION_PROPERTIES`` binding. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + Routed through pinned-host staging because ``DOF_FRICTION_PROPERTIES`` + is a CPU-only OVPhysX binding. """ + self._read_scalar_binding(TT.DOF_FRICTION_PROPERTIES, self._joint_friction_props_buf) if self._joint_friction_coeff_ta is None: self._joint_friction_coeff_ta = ProxyArray(self._joint_friction_coeff) return self._joint_friction_coeff_ta + @property + def joint_dynamic_friction_coeff(self) -> ProxyArray: + """Joint dynamic friction coefficient [dimensionless]. + + Shape is (num_instances, num_joints), dtype = wp.float32. + Component ``[..., 1]`` of the ``DOF_FRICTION_PROPERTIES`` binding. + + Routed through pinned-host staging because ``DOF_FRICTION_PROPERTIES`` + is a CPU-only OVPhysX binding. + """ + self._read_scalar_binding(TT.DOF_FRICTION_PROPERTIES, self._joint_friction_props_buf) + if self._joint_dynamic_friction_coeff_ta is None: + self._joint_dynamic_friction_coeff_ta = ProxyArray(self._joint_dynamic_friction_coeff) + return self._joint_dynamic_friction_coeff_ta + + @property + def joint_viscous_friction_coeff(self) -> ProxyArray: + """Joint viscous friction coefficient [N*m*s/rad or N*s/m, depending on joint type]. + + Shape is (num_instances, num_joints), dtype = wp.float32. + Component ``[..., 2]`` of the ``DOF_FRICTION_PROPERTIES`` binding. + + Routed through pinned-host staging because ``DOF_FRICTION_PROPERTIES`` + is a CPU-only OVPhysX binding. + """ + self._read_scalar_binding(TT.DOF_FRICTION_PROPERTIES, self._joint_friction_props_buf) + if self._joint_viscous_friction_coeff_ta is None: + self._joint_viscous_friction_coeff_ta = ProxyArray(self._joint_viscous_friction_coeff) + return self._joint_viscous_friction_coeff_ta + @property def joint_pos_limits(self) -> ProxyArray: - """Joint position limits provided to the simulation. + """Joint position limits provided to the simulation [m or rad, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.vec2f. In torch this resolves to - (num_instances, num_joints, 2). + Shape is (num_instances, num_joints), dtype = wp.vec2f. + In torch this resolves to (num_instances, num_joints, 2). - The limits are in the order :math:`[lower, upper]`. + The limits are in the order :math:`[lower, upper]`. Routed through + pinned-host staging because ``DOF_LIMIT`` is a CPU-only OVPhysX binding. """ + self._read_scalar_binding(TT.DOF_LIMIT, self._joint_pos_limits) if self._joint_pos_limits_ta is None: - self._joint_pos_limits_ta = ProxyArray(self._joint_pos_limits) + self._joint_pos_limits_ta = ProxyArray(self._joint_pos_limits.data) return self._joint_pos_limits_ta @property def joint_vel_limits(self) -> ProxyArray: """Joint maximum velocity provided to the simulation [m/s or rad/s, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + Shape is (num_instances, num_joints), dtype = wp.float32. + + Routed through pinned-host staging because ``DOF_MAX_VELOCITY`` is a + CPU-only OVPhysX binding. """ + self._read_scalar_binding(TT.DOF_MAX_VELOCITY, self._joint_vel_limits) if self._joint_vel_limits_ta is None: - self._joint_vel_limits_ta = ProxyArray(self._joint_vel_limits) + self._joint_vel_limits_ta = ProxyArray(self._joint_vel_limits.data) return self._joint_vel_limits_ta @property def joint_effort_limits(self) -> ProxyArray: """Joint maximum effort provided to the simulation [N or N*m, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + Shape is (num_instances, num_joints), dtype = wp.float32. + + Routed through pinned-host staging because ``DOF_MAX_FORCE`` is a + CPU-only OVPhysX binding. """ + self._read_scalar_binding(TT.DOF_MAX_FORCE, self._joint_effort_limits) if self._joint_effort_limits_ta is None: - self._joint_effort_limits_ta = ProxyArray(self._joint_effort_limits) + self._joint_effort_limits_ta = ProxyArray(self._joint_effort_limits.data) return self._joint_effort_limits_ta """ @@ -448,25 +508,12 @@ def joint_effort_limits(self) -> ProxyArray: @property def soft_joint_pos_limits(self) -> ProxyArray: - r"""Soft joint position limits for all joints. + r"""Soft joint position limits for all joints [m or rad, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.vec2f. In torch this resolves to - (num_instances, num_joints, 2). + Shape is (num_instances, num_joints), dtype = wp.vec2f. + In torch this resolves to (num_instances, num_joints, 2). - The limits are in the order :math:`[lower, upper]`. The soft joint position limits are computed as - a sub-region of the :attr:`joint_pos_limits` based on the - :attr:`~isaaclab.assets.ArticulationCfg.soft_joint_pos_limit_factor` parameter. - - Consider the joint position limits :math:`[lower, upper]` and the soft joint position limits - :math:`[soft\_lower, soft\_upper]`. The soft joint position limits are computed as: - - .. math:: - - soft\_lower = (lower + upper) / 2 - factor * (upper - lower) / 2 - soft\_upper = (lower + upper) / 2 + factor * (upper - lower) / 2 - - The soft joint position limits help specify a safety region around the joint limits. It isn't used by the - simulation, but is useful for learning agents to prevent the joint positions from violating the limits. + The limits are in the order :math:`[lower, upper]`. """ if self._soft_joint_pos_limits_ta is None: self._soft_joint_pos_limits_ta = ProxyArray(self._soft_joint_pos_limits) @@ -474,13 +521,9 @@ def soft_joint_pos_limits(self) -> ProxyArray: @property def soft_joint_vel_limits(self) -> ProxyArray: - """Soft joint velocity limits for all joints. - - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + """Soft joint velocity limits for all joints [m/s or rad/s, depending on joint type]. - These are obtained from the actuator model. It may differ from :attr:`joint_vel_limits` if the actuator model - has a variable velocity limit model. For instance, in a variable gear ratio actuator model. + Shape is (num_instances, num_joints), dtype = wp.float32. """ if self._soft_joint_vel_limits_ta is None: self._soft_joint_vel_limits_ta = ProxyArray(self._soft_joint_vel_limits) @@ -490,8 +533,7 @@ def soft_joint_vel_limits(self) -> ProxyArray: def gear_ratio(self) -> ProxyArray: """Gear ratio for relating motor torques to applied joint torques. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + Shape is (num_instances, num_joints), dtype = wp.float32. """ if self._gear_ratio_ta is None: self._gear_ratio_ta = ProxyArray(self._gear_ratio) @@ -503,68 +545,84 @@ def gear_ratio(self) -> ProxyArray: @property def fixed_tendon_stiffness(self) -> ProxyArray: - """Fixed tendon stiffness provided to the simulation. + """Fixed-tendon stiffness gains [N*m/rad]. - Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to - (num_instances, num_fixed_tendons). + Shape is (num_instances, num_fixed_tendons), dtype = ``wp.float32``. + + Routed through pinned-host staging (CPU-only binding). """ + self._read_scalar_binding(TT.FIXED_TENDON_STIFFNESS, self._fixed_tendon_stiffness) if self._fixed_tendon_stiffness_ta is None: - self._fixed_tendon_stiffness_ta = ProxyArray(self._fixed_tendon_stiffness) + self._fixed_tendon_stiffness_ta = ProxyArray(self._fixed_tendon_stiffness.data) return self._fixed_tendon_stiffness_ta @property def fixed_tendon_damping(self) -> ProxyArray: - """Fixed tendon damping provided to the simulation. + """Fixed-tendon damping coefficients [N*m*s/rad]. + + Shape is (num_instances, num_fixed_tendons), dtype = ``wp.float32``. - Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to - (num_instances, num_fixed_tendons). + Routed through pinned-host staging (CPU-only binding). """ + self._read_scalar_binding(TT.FIXED_TENDON_DAMPING, self._fixed_tendon_damping) if self._fixed_tendon_damping_ta is None: - self._fixed_tendon_damping_ta = ProxyArray(self._fixed_tendon_damping) + self._fixed_tendon_damping_ta = ProxyArray(self._fixed_tendon_damping.data) return self._fixed_tendon_damping_ta @property def fixed_tendon_limit_stiffness(self) -> ProxyArray: - """Fixed tendon limit stiffness provided to the simulation. + """Fixed-tendon limit stiffness [N*m/rad]. - Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to - (num_instances, num_fixed_tendons). + Shape is (num_instances, num_fixed_tendons), dtype = ``wp.float32``. + + Routed through pinned-host staging (CPU-only binding). """ + self._read_scalar_binding(TT.FIXED_TENDON_LIMIT_STIFFNESS, self._fixed_tendon_limit_stiffness) if self._fixed_tendon_limit_stiffness_ta is None: - self._fixed_tendon_limit_stiffness_ta = ProxyArray(self._fixed_tendon_limit_stiffness) + self._fixed_tendon_limit_stiffness_ta = ProxyArray(self._fixed_tendon_limit_stiffness.data) return self._fixed_tendon_limit_stiffness_ta @property def fixed_tendon_rest_length(self) -> ProxyArray: - """Fixed tendon rest length provided to the simulation. + """Fixed-tendon rest lengths [m]. + + Shape is (num_instances, num_fixed_tendons), dtype = ``wp.float32``. - Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to - (num_instances, num_fixed_tendons). + Routed through pinned-host staging (CPU-only binding). """ + self._read_scalar_binding(TT.FIXED_TENDON_REST_LENGTH, self._fixed_tendon_rest_length) if self._fixed_tendon_rest_length_ta is None: - self._fixed_tendon_rest_length_ta = ProxyArray(self._fixed_tendon_rest_length) + self._fixed_tendon_rest_length_ta = ProxyArray(self._fixed_tendon_rest_length.data) return self._fixed_tendon_rest_length_ta @property def fixed_tendon_offset(self) -> ProxyArray: - """Fixed tendon offset provided to the simulation. + """Fixed-tendon offsets [m]. - Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to - (num_instances, num_fixed_tendons). + Shape is (num_instances, num_fixed_tendons), dtype = ``wp.float32``. + + Routed through pinned-host staging (CPU-only binding). """ + self._read_scalar_binding(TT.FIXED_TENDON_OFFSET, self._fixed_tendon_offset) if self._fixed_tendon_offset_ta is None: - self._fixed_tendon_offset_ta = ProxyArray(self._fixed_tendon_offset) + self._fixed_tendon_offset_ta = ProxyArray(self._fixed_tendon_offset.data) return self._fixed_tendon_offset_ta @property def fixed_tendon_pos_limits(self) -> ProxyArray: - """Fixed tendon position limits provided to the simulation. + """Fixed tendon position limits provided to the simulation [m or rad]. + + Shape is (num_instances, num_fixed_tendons), dtype = ``wp.vec2f``. + In torch this resolves to (num_instances, num_fixed_tendons, 2). - Shape is (num_instances, num_fixed_tendons), dtype = wp.vec2f. In torch this resolves to - (num_instances, num_fixed_tendons, 2). + .. deprecated:: + Use :attr:`fixed_tendon_limit` (shape ``(N, T, 2)``, dtype + ``wp.float32``) instead. This alias is kept for backwards + compatibility and reads the same underlying data. """ + self._read_scalar_binding(TT.FIXED_TENDON_LIMIT, self._fixed_tendon_pos_limits) if self._fixed_tendon_pos_limits_ta is None: - self._fixed_tendon_pos_limits_ta = ProxyArray(self._fixed_tendon_pos_limits) + self._fixed_tendon_pos_limits_ta = ProxyArray(self._fixed_tendon_pos_limits.data) return self._fixed_tendon_pos_limits_ta """ @@ -573,46 +631,54 @@ def fixed_tendon_pos_limits(self) -> ProxyArray: @property def spatial_tendon_stiffness(self) -> ProxyArray: - """Spatial tendon stiffness provided to the simulation. + """Spatial-tendon stiffness gains [N/m]. + + Shape is (num_instances, num_spatial_tendons), dtype = ``wp.float32``. - Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to - (num_instances, num_spatial_tendons). + Routed through pinned-host staging (CPU-only binding). """ + self._read_scalar_binding(TT.SPATIAL_TENDON_STIFFNESS, self._spatial_tendon_stiffness) if self._spatial_tendon_stiffness_ta is None: - self._spatial_tendon_stiffness_ta = ProxyArray(self._spatial_tendon_stiffness) + self._spatial_tendon_stiffness_ta = ProxyArray(self._spatial_tendon_stiffness.data) return self._spatial_tendon_stiffness_ta @property def spatial_tendon_damping(self) -> ProxyArray: - """Spatial tendon damping provided to the simulation. + """Spatial-tendon damping coefficients [N*s/m]. - Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to - (num_instances, num_spatial_tendons). + Shape is (num_instances, num_spatial_tendons), dtype = ``wp.float32``. + + Routed through pinned-host staging (CPU-only binding). """ + self._read_scalar_binding(TT.SPATIAL_TENDON_DAMPING, self._spatial_tendon_damping) if self._spatial_tendon_damping_ta is None: - self._spatial_tendon_damping_ta = ProxyArray(self._spatial_tendon_damping) + self._spatial_tendon_damping_ta = ProxyArray(self._spatial_tendon_damping.data) return self._spatial_tendon_damping_ta @property def spatial_tendon_limit_stiffness(self) -> ProxyArray: - """Spatial tendon limit stiffness provided to the simulation. + """Spatial-tendon limit stiffness [N/m]. + + Shape is (num_instances, num_spatial_tendons), dtype = ``wp.float32``. - Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to - (num_instances, num_spatial_tendons). + Routed through pinned-host staging (CPU-only binding). """ + self._read_scalar_binding(TT.SPATIAL_TENDON_LIMIT_STIFFNESS, self._spatial_tendon_limit_stiffness) if self._spatial_tendon_limit_stiffness_ta is None: - self._spatial_tendon_limit_stiffness_ta = ProxyArray(self._spatial_tendon_limit_stiffness) + self._spatial_tendon_limit_stiffness_ta = ProxyArray(self._spatial_tendon_limit_stiffness.data) return self._spatial_tendon_limit_stiffness_ta @property def spatial_tendon_offset(self) -> ProxyArray: - """Spatial tendon offset provided to the simulation. + """Spatial-tendon offsets [m]. + + Shape is (num_instances, num_spatial_tendons), dtype = ``wp.float32``. - Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to - (num_instances, num_spatial_tendons). + Routed through pinned-host staging (CPU-only binding). """ + self._read_scalar_binding(TT.SPATIAL_TENDON_OFFSET, self._spatial_tendon_offset) if self._spatial_tendon_offset_ta is None: - self._spatial_tendon_offset_ta = ProxyArray(self._spatial_tendon_offset) + self._spatial_tendon_offset_ta = ProxyArray(self._spatial_tendon_offset.data) return self._spatial_tendon_offset_ta """ @@ -621,8 +687,10 @@ def spatial_tendon_offset(self) -> ProxyArray: @property def root_link_pose_w(self) -> ProxyArray: - """Root link pose ``[pos, quat]`` in simulation world frame. - Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). + """Root link pose ``[pos, quat]`` in simulation world frame [m, -]. + + Shape is (num_instances,), dtype = wp.transformf. + In torch this resolves to (num_instances, 7). This quantity is the pose of the articulation root's actor frame relative to the world. The orientation is provided in (x, y, z, w) format. @@ -632,10 +700,21 @@ def root_link_pose_w(self) -> ProxyArray: self._root_link_pose_w_ta = ProxyArray(self._root_link_pose_w.data) return self._root_link_pose_w_ta + @property + def root_pose_w(self) -> ProxyArray: + """Alias for :attr:`root_link_pose_w` matching Newton's convention. + + Shape is (num_instances,), dtype = wp.transformf. + In torch this resolves to (num_instances, 7). + """ + return self.root_link_pose_w + @property def root_link_vel_w(self) -> ProxyArray: - """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. - Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). + """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame [m/s, rad/s]. + + Shape is (num_instances,), dtype = wp.spatial_vectorf. + In torch this resolves to (num_instances, 6). This quantity contains the linear and angular velocities of the articulation root's actor frame relative to the world. @@ -646,7 +725,7 @@ def root_link_vel_w(self) -> ProxyArray: if self._root_link_vel_w.timestamp < self._sim_timestamp: wp.launch( _copy_first_body, - dim=self._num_instances, + dim=self.num_instances, inputs=[self._body_link_vel_w.data], outputs=[self._root_link_vel_w.data], device=self.device, @@ -658,8 +737,10 @@ def root_link_vel_w(self) -> ProxyArray: @property def root_com_pose_w(self) -> ProxyArray: - """Root center of mass pose ``[pos, quat]`` in simulation world frame. - Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). + """Root center of mass pose ``[pos, quat]`` in simulation world frame [m, -]. + + Shape is (num_instances,), dtype = wp.transformf. + In torch this resolves to (num_instances, 7). This quantity is the pose of the articulation root's center of mass frame relative to the world. The orientation is provided in (x, y, z, w) format. @@ -667,7 +748,7 @@ def root_com_pose_w(self) -> ProxyArray: if self._root_com_pose_w.timestamp < self._sim_timestamp: wp.launch( _compose_root_com_pose, - dim=self._num_instances, + dim=self.num_instances, inputs=[self.root_link_pose_w, self.body_com_pose_b], outputs=[self._root_com_pose_w.data], device=self.device, @@ -679,8 +760,10 @@ def root_com_pose_w(self) -> ProxyArray: @property def root_com_vel_w(self) -> ProxyArray: - """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. - Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). + """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame [m/s, rad/s]. + + Shape is (num_instances,), dtype = wp.spatial_vectorf. + In torch this resolves to (num_instances, 6). This quantity contains the linear and angular velocities of the articulation root's center of mass frame relative to the world. @@ -696,98 +779,127 @@ def root_com_vel_w(self) -> ProxyArray: @property def body_mass(self) -> ProxyArray: - """Body mass in the world frame [kg]. + """Body masses [kg]. - Shape is (num_instances, num_bodies), dtype = wp.float32. In torch this resolves to - (num_instances, num_bodies). + Shape is (num_instances, num_bodies), dtype = ``wp.float32``. + + Routed through pinned-host staging because the underlying OVPhysX + binding is CPU-only (``ARTICULATION_BODY_MASS``). """ + self._read_scalar_binding(TT.BODY_MASS, self._body_mass) if self._body_mass_ta is None: - self._body_mass_ta = ProxyArray(self._body_mass) + self._body_mass_ta = ProxyArray(self._body_mass.data) return self._body_mass_ta @property def body_inertia(self) -> ProxyArray: - """Flattened body inertia in the world frame [kg*m^2]. + """Body inertia tensors [kg*m^2]. - Shape is (num_instances, num_bodies, 9), dtype = wp.float32. In torch this resolves to - (num_instances, num_bodies, 9). + Shape is (num_instances, num_bodies, 9), dtype = ``wp.float32``; the + trailing 9 is the row-major 3×3 inertia tensor. - Stored as a flattened 3x3 inertia matrix per body. + Routed through pinned-host staging (``ARTICULATION_BODY_INERTIA`` is + a CPU-only binding). """ + self._read_scalar_binding(TT.BODY_INERTIA, self._body_inertia) if self._body_inertia_ta is None: - self._body_inertia_ta = ProxyArray(self._body_inertia) + self._body_inertia_ta = ProxyArray(self._body_inertia.data) return self._body_inertia_ta @property def body_link_pose_w(self) -> ProxyArray: - """Body link pose ``[pos, quat]`` in simulation world frame. - Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to - (num_instances, num_bodies, 7). + """Body link pose ``[pos, quat]`` in simulation world frame [m, -]. + + Shape is (num_instances, num_bodies), dtype = wp.transformf. + In torch this resolves to (num_instances, num_bodies, 7). This quantity is the pose of the articulation links' actor frame relative to the world. The orientation is provided in (x, y, z, w) format. """ + if self._body_link_pose_w.timestamp < self._sim_timestamp: + # perform forward kinematics (shouldn't cause overhead if it happened already) + OvPhysxManager.get_physx_instance().update_articulations_kinematic() self._read_transform_binding(TT.LINK_POSE, self._body_link_pose_w) if self._body_link_pose_w_ta is None: self._body_link_pose_w_ta = ProxyArray(self._body_link_pose_w.data) return self._body_link_pose_w_ta @property - def body_link_vel_w(self) -> ProxyArray: - """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. - Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to - (num_instances, num_bodies, 6). + def body_com_vel_w(self) -> ProxyArray: + """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame [m/s, rad/s]. - This quantity contains the linear and angular velocities of the articulation links' actor frame - relative to the world. + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. + In torch this resolves to (num_instances, num_bodies, 6). """ - self._read_spatial_vector_binding(TT.LINK_VELOCITY, self._body_link_vel_w) + self._read_spatial_vector_binding(TT.LINK_VELOCITY, self._body_com_vel_w) + if self._body_com_vel_w_ta is None: + self._body_com_vel_w_ta = ProxyArray(self._body_com_vel_w.data) + return self._body_com_vel_w_ta + + @property + def body_link_vel_w(self) -> ProxyArray: + """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame [m/s, rad/s]. + + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. + In torch this resolves to (num_instances, num_bodies, 6). + + Derived from :attr:`body_com_vel_w` and :attr:`body_com_pose_b` via + :func:`~isaaclab_ovphysx.assets.kernels.get_body_link_vel_from_body_com_vel`. + """ + if self._body_link_vel_w.timestamp >= self._sim_timestamp: + if self._body_link_vel_w_ta is None: + self._body_link_vel_w_ta = ProxyArray(self._body_link_vel_w.data) + return self._body_link_vel_w_ta + _ = self.body_com_vel_w + _ = self.body_link_pose_w + _ = self.body_com_pose_b + wp.launch( + get_body_link_vel_from_body_com_vel, + dim=(self.num_instances, self.num_bodies), + inputs=[self._body_com_vel_w.data, self._body_link_pose_w.data, self._body_com_pose_b.data], + outputs=[self._body_link_vel_w.data], + device=self.device, + ) + self._body_link_vel_w.timestamp = self._sim_timestamp if self._body_link_vel_w_ta is None: self._body_link_vel_w_ta = ProxyArray(self._body_link_vel_w.data) return self._body_link_vel_w_ta @property def body_com_pose_w(self) -> ProxyArray: - """Body center of mass pose ``[pos, quat]`` in simulation world frame. - Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to - (num_instances, num_bodies, 7). + """Body center of mass pose ``[pos, quat]`` in simulation world frame [m, -]. + + Shape is (num_instances, num_bodies), dtype = wp.transformf. + In torch this resolves to (num_instances, num_bodies, 7). - This quantity is the pose of the center of mass frame of the articulation links relative to the world. + Derived from :attr:`body_link_pose_w` and :attr:`body_com_pose_b` via + :func:`~isaaclab_ovphysx.assets.kernels.get_body_com_pose_from_body_link_pose`. The orientation is provided in (x, y, z, w) format. """ - if self._body_com_pose_w.timestamp < self._sim_timestamp: - wp.launch( - _compose_body_com_poses, - dim=(self._num_instances, self._num_bodies), - inputs=[self.body_link_pose_w, self.body_com_pose_b], - outputs=[self._body_com_pose_w.data], - device=self.device, - ) - self._body_com_pose_w.timestamp = self._sim_timestamp + if self._body_com_pose_w.timestamp >= self._sim_timestamp: + if self._body_com_pose_w_ta is None: + self._body_com_pose_w_ta = ProxyArray(self._body_com_pose_w.data) + return self._body_com_pose_w_ta + _ = self.body_link_pose_w + _ = self.body_com_pose_b + wp.launch( + get_body_com_pose_from_body_link_pose, + dim=(self.num_instances, self.num_bodies), + inputs=[self._body_link_pose_w.data, self._body_com_pose_b.data], + outputs=[self._body_com_pose_w.data], + device=self.device, + ) + self._body_com_pose_w.timestamp = self._sim_timestamp if self._body_com_pose_w_ta is None: self._body_com_pose_w_ta = ProxyArray(self._body_com_pose_w.data) return self._body_com_pose_w_ta - @property - def body_com_vel_w(self) -> ProxyArray: - """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. - Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to - (num_instances, num_bodies, 6). - - This quantity contains the linear and angular velocities of the articulation links' center of mass frame - relative to the world. - - .. note:: - This is currently approximated using the link velocity. A proper COM velocity derivation - accounting for the COM offset is not yet implemented. - """ - return self.body_link_vel_w - @property def body_com_acc_w(self) -> ProxyArray: """Acceleration of all bodies center of mass ``[lin_acc, ang_acc]`` [m/s^2, rad/s^2]. - Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to - (num_instances, num_bodies, 6). + + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. + In torch this resolves to (num_instances, num_bodies, 6). All values are relative to the world. """ @@ -798,9 +910,10 @@ def body_com_acc_w(self) -> ProxyArray: @property def body_com_pose_b(self) -> ProxyArray: - """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. - Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to - (num_instances, num_bodies, 7). + """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames [m, -]. + + Shape is (num_instances, num_bodies), dtype = wp.transformf. + In torch this resolves to (num_instances, num_bodies, 7). This quantity is the pose of the center of mass frame of the rigid body relative to the body's link frame. The orientation is provided in (x, y, z, w) format. @@ -810,6 +923,23 @@ def body_com_pose_b(self) -> ProxyArray: self._body_com_pose_b_ta = ProxyArray(self._body_com_pose_b.data) return self._body_com_pose_b_ta + @property + def body_incoming_joint_wrench_b(self) -> ProxyArray: + """Incoming joint wrenches on each body in the body frame [N, N*m]. + + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. + In torch this resolves to (num_instances, num_bodies, 6). + + All body reaction wrenches are provided including the root body to the world of an articulation. + """ + self._read_spatial_vector_binding( + TT.LINK_INCOMING_JOINT_FORCE, + self._body_incoming_joint_wrench_buf, + ) + if self._body_incoming_joint_wrench_b_ta is None: + self._body_incoming_joint_wrench_b_ta = ProxyArray(self._body_incoming_joint_wrench_buf.data) + return self._body_incoming_joint_wrench_b_ta + """ Joint state properties. """ @@ -818,8 +948,7 @@ def body_com_pose_b(self) -> ProxyArray: def joint_pos(self) -> ProxyArray: """Joint positions of all joints [m or rad, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + Shape is (num_instances, num_joints), dtype = wp.float32. """ self._read_binding_into_buf(TT.DOF_POSITION, self._joint_pos_buf) if self._joint_pos_ta is None: @@ -830,8 +959,7 @@ def joint_pos(self) -> ProxyArray: def joint_vel(self) -> ProxyArray: """Joint velocities of all joints [m/s or rad/s, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + Shape is (num_instances, num_joints), dtype = wp.float32. """ self._read_binding_into_buf(TT.DOF_VELOCITY, self._joint_vel_buf) if self._joint_vel_ta is None: @@ -842,8 +970,7 @@ def joint_vel(self) -> ProxyArray: def joint_acc(self) -> ProxyArray: """Joint acceleration of all joints [m/s^2 or rad/s^2, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + Shape is (num_instances, num_joints), dtype = wp.float32. .. note:: This quantity is computed via finite differencing of joint velocities. @@ -859,12 +986,13 @@ def joint_acc(self) -> ProxyArray: @property def projected_gravity_b(self) -> ProxyArray: """Projection of the gravity direction on base frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ if self._projected_gravity_b.timestamp < self._sim_timestamp: wp.launch( _projected_gravity, - dim=self._num_instances, + dim=self.num_instances, inputs=[self.GRAVITY_VEC_W, self.root_link_pose_w], outputs=[self._projected_gravity_b.data], device=self.device, @@ -876,7 +1004,8 @@ def projected_gravity_b(self) -> ProxyArray: @property def heading_w(self) -> ProxyArray: - """Yaw heading of the base frame (in radians). + """Yaw heading of the base frame (in radians) [rad]. + Shape is (num_instances,), dtype = wp.float32. .. note:: @@ -886,7 +1015,7 @@ def heading_w(self) -> ProxyArray: if self._heading_w.timestamp < self._sim_timestamp: wp.launch( _compute_heading, - dim=self._num_instances, + dim=self.num_instances, inputs=[self.FORWARD_VEC_B, self.root_link_pose_w], outputs=[self._heading_w.data], device=self.device, @@ -899,6 +1028,7 @@ def heading_w(self) -> ProxyArray: @property def root_link_lin_vel_b(self) -> ProxyArray: """Root link linear velocity in base frame [m/s]. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the articulation root's actor frame with respect to its actor frame. @@ -906,7 +1036,7 @@ def root_link_lin_vel_b(self) -> ProxyArray: if self._root_link_lin_vel_b.timestamp < self._sim_timestamp: wp.launch( _world_vel_to_body_lin, - dim=self._num_instances, + dim=self.num_instances, inputs=[self.root_link_pose_w, self.root_link_vel_w], outputs=[self._root_link_lin_vel_b.data], device=self.device, @@ -919,6 +1049,7 @@ def root_link_lin_vel_b(self) -> ProxyArray: @property def root_link_ang_vel_b(self) -> ProxyArray: """Root link angular velocity in base frame [rad/s]. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the articulation root's actor frame with respect to its actor frame. @@ -926,7 +1057,7 @@ def root_link_ang_vel_b(self) -> ProxyArray: if self._root_link_ang_vel_b.timestamp < self._sim_timestamp: wp.launch( _world_vel_to_body_ang, - dim=self._num_instances, + dim=self.num_instances, inputs=[self.root_link_pose_w, self.root_link_vel_w], outputs=[self._root_link_ang_vel_b.data], device=self.device, @@ -939,6 +1070,7 @@ def root_link_ang_vel_b(self) -> ProxyArray: @property def root_com_lin_vel_b(self) -> ProxyArray: """Root center of mass linear velocity in base frame [m/s]. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the articulation root's center of mass frame @@ -947,7 +1079,7 @@ def root_com_lin_vel_b(self) -> ProxyArray: if self._root_com_lin_vel_b.timestamp < self._sim_timestamp: wp.launch( _world_vel_to_body_lin, - dim=self._num_instances, + dim=self.num_instances, inputs=[self.root_link_pose_w, self.root_com_vel_w], outputs=[self._root_com_lin_vel_b.data], device=self.device, @@ -960,6 +1092,7 @@ def root_com_lin_vel_b(self) -> ProxyArray: @property def root_com_ang_vel_b(self) -> ProxyArray: """Root center of mass angular velocity in base frame [rad/s]. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the articulation root's center of mass frame @@ -968,7 +1101,7 @@ def root_com_ang_vel_b(self) -> ProxyArray: if self._root_com_ang_vel_b.timestamp < self._sim_timestamp: wp.launch( _world_vel_to_body_ang, - dim=self._num_instances, + dim=self.num_instances, inputs=[self.root_link_pose_w, self.root_com_vel_w], outputs=[self._root_com_ang_vel_b.data], device=self.device, @@ -984,10 +1117,9 @@ def root_com_ang_vel_b(self) -> ProxyArray: @property def root_link_pos_w(self) -> ProxyArray: - """Root link position in simulation world frame. - Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + """Root link position in simulation world frame [m]. - This quantity is the position of the actor frame of the root rigid body relative to the world. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ parent = self.root_link_pose_w if self._root_link_pos_w_ta is None: @@ -997,9 +1129,8 @@ def root_link_pos_w(self) -> ProxyArray: @property def root_link_quat_w(self) -> ProxyArray: """Root link orientation (x, y, z, w) in simulation world frame. - Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). - This quantity is the orientation of the actor frame of the root rigid body. + Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). """ parent = self.root_link_pose_w if self._root_link_quat_w_ta is None: @@ -1008,10 +1139,9 @@ def root_link_quat_w(self) -> ProxyArray: @property def root_link_lin_vel_w(self) -> ProxyArray: - """Root linear velocity in simulation world frame [m/s]. - Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + """Root link linear velocity in simulation world frame [m/s]. - This quantity is the linear velocity of the root rigid body's actor frame relative to the world. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ parent = self.root_link_vel_w if self._root_link_lin_vel_w_ta is None: @@ -1021,9 +1151,8 @@ def root_link_lin_vel_w(self) -> ProxyArray: @property def root_link_ang_vel_w(self) -> ProxyArray: """Root link angular velocity in simulation world frame [rad/s]. - Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). - This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ parent = self.root_link_vel_w if self._root_link_ang_vel_w_ta is None: @@ -1032,10 +1161,9 @@ def root_link_ang_vel_w(self) -> ProxyArray: @property def root_com_pos_w(self) -> ProxyArray: - """Root center of mass position in simulation world frame. - Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + """Root center of mass position in simulation world frame [m]. - This quantity is the position of the center of mass frame of the root rigid body relative to the world. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ parent = self.root_com_pose_w if self._root_com_pos_w_ta is None: @@ -1045,9 +1173,8 @@ def root_com_pos_w(self) -> ProxyArray: @property def root_com_quat_w(self) -> ProxyArray: """Root center of mass orientation (x, y, z, w) in simulation world frame. - Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). - This quantity is the orientation of the principal axes of inertia of the root rigid body relative to the world. + Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). """ parent = self.root_com_pose_w if self._root_com_quat_w_ta is None: @@ -1057,9 +1184,8 @@ def root_com_quat_w(self) -> ProxyArray: @property def root_com_lin_vel_w(self) -> ProxyArray: """Root center of mass linear velocity in simulation world frame [m/s]. - Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). - This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ parent = self.root_com_vel_w if self._root_com_lin_vel_w_ta is None: @@ -1069,9 +1195,8 @@ def root_com_lin_vel_w(self) -> ProxyArray: @property def root_com_ang_vel_w(self) -> ProxyArray: """Root center of mass angular velocity in simulation world frame [rad/s]. - Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). - This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ parent = self.root_com_vel_w if self._root_com_ang_vel_w_ta is None: @@ -1080,11 +1205,10 @@ def root_com_ang_vel_w(self) -> ProxyArray: @property def body_link_pos_w(self) -> ProxyArray: - """Positions of all bodies in simulation world frame. - Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to - (num_instances, num_bodies, 3). + """Positions of all bodies in simulation world frame [m]. - This quantity is the position of the articulation bodies' actor frame relative to the world. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. + In torch this resolves to (num_instances, num_bodies, 3). """ parent = self.body_link_pose_w if self._body_link_pos_w_ta is None: @@ -1094,10 +1218,9 @@ def body_link_pos_w(self) -> ProxyArray: @property def body_link_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of all bodies in simulation world frame. - Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to - (num_instances, num_bodies, 4). - This quantity is the orientation of the articulation bodies' actor frame relative to the world. + Shape is (num_instances, num_bodies), dtype = wp.quatf. + In torch this resolves to (num_instances, num_bodies, 4). """ parent = self.body_link_pose_w if self._body_link_quat_w_ta is None: @@ -1107,10 +1230,9 @@ def body_link_quat_w(self) -> ProxyArray: @property def body_link_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame [m/s]. - Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to - (num_instances, num_bodies, 3). - This quantity is the linear velocity of the articulation bodies' actor frame relative to the world. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. + In torch this resolves to (num_instances, num_bodies, 3). """ parent = self.body_link_vel_w if self._body_link_lin_vel_w_ta is None: @@ -1120,10 +1242,9 @@ def body_link_lin_vel_w(self) -> ProxyArray: @property def body_link_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame [rad/s]. - Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to - (num_instances, num_bodies, 3). - This quantity is the angular velocity of the articulation bodies' actor frame relative to the world. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. + In torch this resolves to (num_instances, num_bodies, 3). """ parent = self.body_link_vel_w if self._body_link_ang_vel_w_ta is None: @@ -1132,11 +1253,10 @@ def body_link_ang_vel_w(self) -> ProxyArray: @property def body_com_pos_w(self) -> ProxyArray: - """Positions of all bodies' center of mass in simulation world frame. - Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to - (num_instances, num_bodies, 3). + """Positions of all bodies' center of mass in simulation world frame [m]. - This quantity is the position of the articulation bodies' center of mass frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. + In torch this resolves to (num_instances, num_bodies, 3). """ parent = self.body_com_pose_w if self._body_com_pos_w_ta is None: @@ -1146,10 +1266,9 @@ def body_com_pos_w(self) -> ProxyArray: @property def body_com_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. - Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to - (num_instances, num_bodies, 4). - This quantity is the orientation of the articulation bodies' principal axes of inertia. + Shape is (num_instances, num_bodies), dtype = wp.quatf. + In torch this resolves to (num_instances, num_bodies, 4). """ parent = self.body_com_pose_w if self._body_com_quat_w_ta is None: @@ -1159,10 +1278,9 @@ def body_com_quat_w(self) -> ProxyArray: @property def body_com_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame [m/s]. - Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to - (num_instances, num_bodies, 3). - This quantity is the linear velocity of the articulation bodies' center of mass frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. + In torch this resolves to (num_instances, num_bodies, 3). """ parent = self.body_com_vel_w if self._body_com_lin_vel_w_ta is None: @@ -1172,10 +1290,9 @@ def body_com_lin_vel_w(self) -> ProxyArray: @property def body_com_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame [rad/s]. - Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to - (num_instances, num_bodies, 3). - This quantity is the angular velocity of the articulation bodies' center of mass frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. + In torch this resolves to (num_instances, num_bodies, 3). """ parent = self.body_com_vel_w if self._body_com_ang_vel_w_ta is None: @@ -1185,10 +1302,9 @@ def body_com_ang_vel_w(self) -> ProxyArray: @property def body_com_lin_acc_w(self) -> ProxyArray: """Linear acceleration of all bodies in simulation world frame [m/s^2]. - Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to - (num_instances, num_bodies, 3). - This quantity is the linear acceleration of the articulation bodies' center of mass frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. + In torch this resolves to (num_instances, num_bodies, 3). """ parent = self.body_com_acc_w if self._body_com_lin_acc_w_ta is None: @@ -1198,10 +1314,9 @@ def body_com_lin_acc_w(self) -> ProxyArray: @property def body_com_ang_acc_w(self) -> ProxyArray: """Angular acceleration of all bodies in simulation world frame [rad/s^2]. - Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to - (num_instances, num_bodies, 3). - This quantity is the angular acceleration of the articulation bodies' center of mass frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. + In torch this resolves to (num_instances, num_bodies, 3). """ parent = self.body_com_acc_w if self._body_com_ang_acc_w_ta is None: @@ -1210,11 +1325,10 @@ def body_com_ang_acc_w(self) -> ProxyArray: @property def body_com_pos_b(self) -> ProxyArray: - """Center of mass position of all of the bodies in their respective link frames. - Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to - (num_instances, num_bodies, 3). + """Center of mass position of all of the bodies in their respective link frames [m]. - This quantity is the center of mass location relative to its body's link frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. + In torch this resolves to (num_instances, num_bodies, 3). """ parent = self.body_com_pose_b if self._body_com_pos_b_ta is None: @@ -1225,10 +1339,9 @@ def body_com_pos_b(self) -> ProxyArray: def body_com_quat_b(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their respective link frames. - Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to - (num_instances, num_bodies, 4). - This quantity is the orientation of the principal axes of inertia relative to its body's link frame. + Shape is (num_instances, num_bodies), dtype = wp.quatf. + In torch this resolves to (num_instances, num_bodies, 4). """ parent = self.body_com_pose_b if self._body_com_quat_b_ta is None: @@ -1236,95 +1349,12 @@ def body_com_quat_b(self) -> ProxyArray: return self._body_com_quat_b_ta """ - Deprecated in base class (required by ABC for backward compatibility). - """ - - @property - def default_root_state(self) -> ProxyArray: - """Deprecated. Use :attr:`default_root_pose` and :attr:`default_root_vel` instead.""" - warnings.warn( - "default_root_state is deprecated. Use default_root_pose and default_root_vel.", - DeprecationWarning, - stacklevel=2, - ) - if self._root_state_w_buf is None: - self._root_state_w_buf = wp.zeros( - self._num_instances, dtype=wp.types.vector(13, wp.float32), device=self.device - ) - if self._default_root_state_ta is None: - self._default_root_state_ta = ProxyArray(self._root_state_w_buf) - return self._default_root_state_ta - - @property - def root_state_w(self) -> ProxyArray: - """Deprecated. Use :attr:`root_link_pose_w` and :attr:`root_com_vel_w` instead.""" - warnings.warn( - "root_state_w is deprecated. Use root_link_pose_w and root_com_vel_w.", - DeprecationWarning, - stacklevel=2, - ) - return self.root_link_pose_w - - @property - def root_link_state_w(self) -> ProxyArray: - """Deprecated. Use :attr:`root_link_pose_w` and :attr:`root_link_vel_w` instead.""" - warnings.warn( - "root_link_state_w is deprecated. Use root_link_pose_w and root_link_vel_w.", - DeprecationWarning, - stacklevel=2, - ) - return self.root_link_pose_w - - @property - def root_com_state_w(self) -> ProxyArray: - """Deprecated. Use :attr:`root_com_pose_w` and :attr:`root_com_vel_w` instead.""" - warnings.warn( - "root_com_state_w is deprecated. Use root_com_pose_w and root_com_vel_w.", - DeprecationWarning, - stacklevel=2, - ) - return self.root_com_pose_w - - @property - def body_state_w(self) -> ProxyArray: - """Deprecated. Use :attr:`body_link_pose_w` and :attr:`body_com_vel_w` instead.""" - warnings.warn( - "body_state_w is deprecated. Use body_link_pose_w and body_com_vel_w.", - DeprecationWarning, - stacklevel=2, - ) - return self.body_link_pose_w - - @property - def body_link_state_w(self) -> ProxyArray: - """Deprecated. Use :attr:`body_link_pose_w` and :attr:`body_link_vel_w` instead.""" - warnings.warn( - "body_link_state_w is deprecated. Use body_link_pose_w and body_link_vel_w.", - DeprecationWarning, - stacklevel=2, - ) - return self.body_link_pose_w - - @property - def body_com_state_w(self) -> ProxyArray: - """Deprecated. Use :attr:`body_com_pose_w` and :attr:`body_com_vel_w` instead.""" - warnings.warn( - "body_com_state_w is deprecated. Use body_com_pose_w and body_com_vel_w.", - DeprecationWarning, - stacklevel=2, - ) - return self.body_com_pose_w - - """ - Internal helper. + Internal helpers. """ def _create_buffers(self) -> None: # noqa: C901 + """Eagerly allocate every TimestampedBuffer and pinned CPU staging buffer.""" super()._create_buffers() - # Scratch buffers for _read_binding_into_* methods, allocated lazily - # on first use and reused every subsequent step to avoid per-step - # allocation overhead on the hot RL path. - self._read_scratch: dict = {} N = self._num_instances D = self._num_joints @@ -1344,26 +1374,33 @@ def _create_buffers(self) -> None: # noqa: C901 self._body_com_pose_w = TimestampedBuffer((N, L), dev, wp.transformf) self._body_com_vel_w = TimestampedBuffer((N, L), dev, wp.spatial_vectorf) self._body_com_acc_w = TimestampedBuffer((N, L), dev, wp.spatial_vectorf) + self._body_incoming_joint_wrench_buf = TimestampedBuffer((N, L), dev, wp.spatial_vectorf) # -- Joint state buffers self._joint_pos_buf = TimestampedBuffer((N, D), dev, wp.float32) self._joint_vel_buf = TimestampedBuffer((N, D), dev, wp.float32) self._joint_acc = TimestampedBuffer((N, D), dev, wp.float32) self._previous_joint_vel = wp.zeros((N, D), dtype=wp.float32, device=dev) - # -- Joint properties - self._joint_stiffness = wp.zeros((N, D), dtype=wp.float32, device=dev) - self._joint_damping = wp.zeros((N, D), dtype=wp.float32, device=dev) - self._joint_armature = wp.zeros((N, D), dtype=wp.float32, device=dev) - self._joint_friction_coeff = wp.zeros((N, D), dtype=wp.float32, device=dev) - self._joint_pos_limits = wp.zeros((N, D), dtype=wp.vec2f, device=dev) - self._joint_vel_limits = wp.zeros((N, D), dtype=wp.float32, device=dev) - self._joint_effort_limits = wp.zeros((N, D), dtype=wp.float32, device=dev) - - # -- Body properties - self._body_mass = wp.zeros((N, L), dtype=wp.float32, device=dev) - self._body_inertia = wp.zeros((N, L, 9), dtype=wp.float32, device=dev) - - # -- Soft limits / custom properties + # -- Joint properties (CPU-only; timestamped so they can be re-read after writes) + self._joint_stiffness = TimestampedBuffer((N, D), dev, wp.float32) + self._joint_damping = TimestampedBuffer((N, D), dev, wp.float32) + self._joint_armature = TimestampedBuffer((N, D), dev, wp.float32) + self._joint_pos_limits = TimestampedBuffer((N, D), dev, wp.vec2f) + self._joint_vel_limits = TimestampedBuffer((N, D), dev, wp.float32) + self._joint_effort_limits = TimestampedBuffer((N, D), dev, wp.float32) + # Friction: single (N, D, 3) TimestampedBuffer; per-component views are created lazily. + self._joint_friction_props_buf = TimestampedBuffer((N, D, 3), dev, wp.float32) + # These are strided wp.array views into _joint_friction_props_buf.data; created in + # _pin_proxy_arrays after the buffer exists. + self._joint_friction_coeff: wp.array | None = None + self._joint_dynamic_friction_coeff: wp.array | None = None + self._joint_viscous_friction_coeff: wp.array | None = None + + # -- Body properties (CPU-only; read once at init, re-read via _read_scalar_binding) + self._body_mass = TimestampedBuffer((N, L), dev, wp.float32) + self._body_inertia = TimestampedBuffer((N, L, 9), dev, wp.float32) + + # -- Soft limits / custom joint properties self._soft_joint_pos_limits = wp.zeros((N, D), dtype=wp.vec2f, device=dev) self._soft_joint_vel_limits = wp.zeros((N, D), dtype=wp.float32, device=dev) self._gear_ratio = wp.ones((N, D), dtype=wp.float32, device=dev) @@ -1389,48 +1426,180 @@ def _create_buffers(self) -> None: # noqa: C901 self._root_com_lin_vel_b = TimestampedBuffer(N, dev, wp.vec3f) self._root_com_ang_vel_b = TimestampedBuffer(N, dev, wp.vec3f) - # -- Deprecated combined state buffers - self._root_state_w_buf = None - self._root_link_state_w_buf = None - self._root_com_state_w_buf = None - self._body_state_w_buf = None - self._body_link_state_w_buf = None - self._body_com_state_w_buf = None - - # -- Tendon property buffers - T_fix = getattr(self, "_num_fixed_tendons", 0) - T_spa = getattr(self, "_num_spatial_tendons", 0) + # -- Deprecated combined state buffers (TimestampedBuffer; lazily filled on first access) + self._root_state_w_buf = TimestampedBuffer(N, dev, vec13f) + self._root_link_state_w_buf = TimestampedBuffer(N, dev, vec13f) + self._root_com_state_w_buf = TimestampedBuffer(N, dev, vec13f) + self._default_root_state_buf = wp.zeros(N, dtype=vec13f, device=dev) + # -- Deprecated body combined state buffers (TimestampedBuffer; lazily filled on first access) + self._body_state_w_buf = TimestampedBuffer((N, L), dev, vec13f) + self._body_link_state_w_buf = TimestampedBuffer((N, L), dev, vec13f) + self._body_com_state_w_buf = TimestampedBuffer((N, L), dev, vec13f) + + # -- Tendon property buffers (always allocated; empty shape when T==0 so + # properties never return None). Routed through _read_scalar_binding. + T_fix = self._num_fixed_tendons + T_spa = self._num_spatial_tendons + self._fixed_tendon_stiffness = TimestampedBuffer((N, T_fix), dev, wp.float32) + self._fixed_tendon_damping = TimestampedBuffer((N, T_fix), dev, wp.float32) + self._fixed_tendon_limit_stiffness = TimestampedBuffer((N, T_fix), dev, wp.float32) + self._fixed_tendon_rest_length = TimestampedBuffer((N, T_fix), dev, wp.float32) + self._fixed_tendon_offset = TimestampedBuffer((N, T_fix), dev, wp.float32) + # Legacy alias kept for any internal callers that used the old vec2f buffer. + self._fixed_tendon_pos_limits = TimestampedBuffer((N, T_fix), dev, wp.vec2f) + self._spatial_tendon_stiffness = TimestampedBuffer((N, T_spa), dev, wp.float32) + self._spatial_tendon_damping = TimestampedBuffer((N, T_spa), dev, wp.float32) + self._spatial_tendon_limit_stiffness = TimestampedBuffer((N, T_spa), dev, wp.float32) + self._spatial_tendon_offset = TimestampedBuffer((N, T_spa), dev, wp.float32) + + # -- CPU staging buffers for CPU-only bindings. + # Pre-allocate all of them so there is no per-step allocation on the hot path. + # These are keyed by tensor_type in self._cpu_staging_buffers; _binding_read + # selects the right one at read time. The sizes must match the binding shapes + # (flat float32). On a GPU sim the buffers are pinned-host (page-locked) so + # the wheel can dispatch async copies; on a CPU sim the staging copy is + # functionally redundant but the buffer must still exist for the write + # helpers, so we allocate unpinned and pay only the intra-CPU memcpy. + pinned = dev != "cpu" + self._cpu_body_mass = wp.zeros((N, L), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_body_coms = wp.zeros((N, L, 7), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_body_inertia = wp.zeros((N, L, 9), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_joint_stiffness = wp.zeros((N, D), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_joint_damping = wp.zeros((N, D), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_joint_position_limit = wp.zeros((N, D, 2), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_joint_velocity_limit = wp.zeros((N, D), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_joint_effort_limit = wp.zeros((N, D), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_joint_armature = wp.zeros((N, D), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_joint_friction_coeff = wp.zeros((N, D), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_joint_dynamic_friction_coeff = wp.zeros((N, D), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_joint_viscous_friction_coeff = wp.zeros((N, D), dtype=wp.float32, device="cpu", pinned=pinned) if T_fix > 0: - self._fixed_tendon_stiffness = wp.zeros((N, T_fix), dtype=wp.float32, device=dev) - self._fixed_tendon_damping = wp.zeros((N, T_fix), dtype=wp.float32, device=dev) - self._fixed_tendon_limit_stiffness = wp.zeros((N, T_fix), dtype=wp.float32, device=dev) - self._fixed_tendon_rest_length = wp.zeros((N, T_fix), dtype=wp.float32, device=dev) - self._fixed_tendon_offset = wp.zeros((N, T_fix), dtype=wp.float32, device=dev) - self._fixed_tendon_pos_limits = wp.zeros((N, T_fix), dtype=wp.vec2f, device=dev) - else: - self._fixed_tendon_stiffness = None - self._fixed_tendon_damping = None - self._fixed_tendon_limit_stiffness = None - self._fixed_tendon_rest_length = None - self._fixed_tendon_offset = None - self._fixed_tendon_pos_limits = None + self._cpu_fixed_tendon_stiffness = wp.zeros((N, T_fix), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_fixed_tendon_damping = wp.zeros((N, T_fix), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_fixed_tendon_limit_stiffness = wp.zeros((N, T_fix), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_fixed_tendon_rest_length = wp.zeros((N, T_fix), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_fixed_tendon_offset = wp.zeros((N, T_fix), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_fixed_tendon_pos_limits = wp.zeros((N, T_fix, 2), dtype=wp.float32, device="cpu", pinned=pinned) if T_spa > 0: - self._spatial_tendon_stiffness = wp.zeros((N, T_spa), dtype=wp.float32, device=dev) - self._spatial_tendon_damping = wp.zeros((N, T_spa), dtype=wp.float32, device=dev) - self._spatial_tendon_limit_stiffness = wp.zeros((N, T_spa), dtype=wp.float32, device=dev) - self._spatial_tendon_offset = wp.zeros((N, T_spa), dtype=wp.float32, device=dev) - else: - self._spatial_tendon_stiffness = None - self._spatial_tendon_damping = None - self._spatial_tendon_limit_stiffness = None - self._spatial_tendon_offset = None + self._cpu_spatial_tendon_stiffness = wp.zeros((N, T_spa), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_spatial_tendon_damping = wp.zeros((N, T_spa), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_spatial_tendon_limit_stiffness = wp.zeros( + (N, T_spa), dtype=wp.float32, device="cpu", pinned=pinned + ) + self._cpu_spatial_tendon_offset = wp.zeros((N, T_spa), dtype=wp.float32, device="cpu", pinned=pinned) - # Read initial joint properties from bindings + # Read initial joint/body properties from bindings (one-time CPU reads). self._read_initial_properties() - - # Initialize ProxyArray wrappers (lazily created on first access) + # Initialize ProxyArray wrappers (lazily created on first property access). self._pin_proxy_arrays() + def _binding_read(self, tensor_type: int, binding: Any, dst: wp.array) -> None: + """Read *binding* into *dst*, staging through a pinned-host buffer for CPU-only bindings. + + For GPU-resident state bindings (pose, velocity, etc.) the read goes directly + into the destination array. For CPU-only property bindings (mass, COM, limits, + stiffness, …) the wheel writes into a pinned-host staging buffer first, then + :func:`wp.copy` moves the data to the simulation device asynchronously. + + Args: + tensor_type: TensorType key identifying the binding. + binding: OVPhysX TensorBinding whose ``read`` method is called. + dst: Destination :class:`wp.array` on the simulation device. + """ + if tensor_type not in TT._CPU_ONLY_TYPES or self.device == "cpu": + binding.read(dst) + return + # Route through a lazily-allocated pinned-host staging buffer. + staging = self._cpu_staging_buffers.get(tensor_type) + if staging is None: + staging = wp.zeros(binding.shape, dtype=wp.float32, device="cpu", pinned=True) + self._cpu_staging_buffers[tensor_type] = staging + binding.read(staging) + # Build a flat float32 view of dst matching the binding's flat shape. + if dst.dtype == wp.float32: + view = dst + else: + view = wp.array( + ptr=dst.ptr, + shape=binding.shape, + dtype=wp.float32, + device=str(dst.device), + copy=False, + ) + wp.copy(view, staging) + + def _binding_write( + self, + tensor_type: int, + binding: Any, + src: wp.array, + *, + indices: wp.array | None = None, + mask: wp.array | None = None, + ) -> None: + """Write *src* to *binding*, staging through pinned-host buffers for CPU-only bindings. + + Args: + tensor_type: TensorType key identifying the binding. + binding: OVPhysX TensorBinding whose ``write`` method is called. + src: Source :class:`wp.array` on the simulation device. + indices: Optional environment indices for partial writes. + mask: Optional boolean mask for partial writes. + """ + if tensor_type not in TT._CPU_ONLY_TYPES or self.device == "cpu": + binding.write(src, indices=indices, mask=mask) + return + # Stage through a pinned-host buffer. + staging = self._cpu_staging_buffers.get(tensor_type) + if staging is None: + staging = wp.zeros(binding.shape, dtype=wp.float32, device="cpu", pinned=True) + self._cpu_staging_buffers[tensor_type] = staging + if src.dtype == wp.float32: + src_view = src + else: + src_view = wp.array( + ptr=src.ptr, + shape=binding.shape, + dtype=wp.float32, + device=str(src.device), + copy=False, + ) + wp.copy(staging, src_view) + binding.write(staging, indices=indices, mask=mask) + + def _stage_to_pinned_cpu(self, tensor_type: int, role: str, src: wp.array) -> wp.array: + """Copy *src* into a lazily-allocated pinned-host :class:`wp.array`. + + Keyed on *(tensor_type, role)* so the same pair always reuses the same + buffer, avoiding per-call allocation on the hot path. + + Args: + tensor_type: TensorType identifying the binding. + role: Disambiguating string when the same tensor_type may serve + multiple purposes (e.g. ``"read"`` vs ``"write"``). + src: Source array on the simulation device. + + Returns: + Pinned-host wp.array containing a copy of *src*. + """ + key = (tensor_type, role) + staging = self._cpu_staging_buffers.get(key) # type: ignore[call-overload] + if staging is None: + if src.dtype == wp.float32: + shape = src.shape + else: + # Flatten to float32 shape matching the element byte size. + elem_floats = src.dtype.size // 4 + shape = src.shape + (elem_floats,) + staging = wp.zeros(shape, dtype=wp.float32, device="cpu", pinned=True) + self._cpu_staging_buffers[key] = staging # type: ignore[index] + if src.dtype == wp.float32: + wp.copy(staging, src) + else: + flat_src = wp.array(ptr=src.ptr, shape=staging.shape, dtype=wp.float32, device=str(src.device), copy=False) + wp.copy(staging, flat_src) + return staging + def _read_initial_properties(self) -> None: """Read static/initial joint and body properties from ovphysx bindings. @@ -1440,7 +1609,6 @@ def _read_initial_properties(self) -> None: simulation device. """ - # Property reads always use CPU numpy (property tensors are host-side). def _read_cpu(tensor_type): binding = self._get_binding(tensor_type) if binding is None: @@ -1449,75 +1617,97 @@ def _read_cpu(tensor_type): binding.read(np_buf) return np_buf - for tt, dst in [ + # Joint scalar properties — write to .data since buffers are now TimestampedBuffer. + for tt, buf in [ (TT.DOF_STIFFNESS, self._joint_stiffness), (TT.DOF_DAMPING, self._joint_damping), (TT.DOF_ARMATURE, self._joint_armature), (TT.DOF_MAX_VELOCITY, self._joint_vel_limits), (TT.DOF_MAX_FORCE, self._joint_effort_limits), - (TT.BODY_MASS, self._body_mass), ]: np_buf = _read_cpu(tt) if np_buf is not None: - wp.copy(dst, wp.from_numpy(np_buf, dtype=wp.float32, device=self.device)) + wp.copy(buf.data, wp.from_numpy(np_buf, dtype=wp.float32, device=self.device)) + buf.timestamp = self._sim_timestamp + + # Body mass (now a TimestampedBuffer). + np_buf = _read_cpu(TT.BODY_MASS) + if np_buf is not None: + wp.copy(self._body_mass.data, wp.from_numpy(np_buf, dtype=wp.float32, device=self.device)) + self._body_mass.timestamp = self._sim_timestamp - # Joint position limits: [N, D, 2] -> (N, D) wp.vec2f + # Joint position limits: [N, D, 2] -> (N, D) wp.vec2f stored in TimestampedBuffer.data np_lim = _read_cpu(TT.DOF_LIMIT) if np_lim is not None: - self._joint_pos_limits = wp.from_numpy( + src = wp.from_numpy( np_lim.reshape(self._num_instances, self._num_joints, 2), dtype=wp.vec2f, device=self.device ) + wp.copy(self._joint_pos_limits.data, src) + self._joint_pos_limits.timestamp = self._sim_timestamp - # Body inertia: [N, L, 9] + # Body inertia (now a TimestampedBuffer): [N, L, 9] np_iner = _read_cpu(TT.BODY_INERTIA) if np_iner is not None: - self._body_inertia = wp.from_numpy(np_iner, dtype=wp.float32, device=self.device) + wp.copy( + self._body_inertia.data, + wp.from_numpy(np_iner, dtype=wp.float32, device=self.device), + ) + self._body_inertia.timestamp = self._sim_timestamp - # Friction: [N, D, 3] -> extract static friction (column 0) + # Friction: [N, D, 3] -> load directly into the combined TimestampedBuffer. + # The strided per-component views (_joint_friction_coeff/dynamic/viscous) are + # created later in _pin_proxy_arrays, so we write to the combined buffer here. np_fric = _read_cpu(TT.DOF_FRICTION_PROPERTIES) if np_fric is not None: - self._joint_friction_coeff = wp.from_numpy(np_fric[..., 0].copy(), dtype=wp.float32, device=self.device) + fric_contiguous = np.ascontiguousarray(np_fric.reshape(self._num_instances, self._num_joints, 3)) + wp.copy( + self._joint_friction_props_buf.data, + wp.from_numpy(fric_contiguous, dtype=wp.float32, device=self.device), + ) + self._joint_friction_props_buf.timestamp = self._sim_timestamp - # Fixed tendon properties (CPU-side, read once) - T_fix = getattr(self, "_num_fixed_tendons", 0) + # Fixed tendon properties. PhysX exposes tendons on the simulation + # device (no ``device="cpu"`` clone in its ``set_fixed_tendon_properties`` + # call); the OVPhysX wheel mirrors that, so we read directly into the + # sim-device buffer rather than via a numpy round-trip. + T_fix = self._num_fixed_tendons if T_fix > 0: - for tt, dst in [ + for tt, buf in [ (TT.FIXED_TENDON_STIFFNESS, self._fixed_tendon_stiffness), (TT.FIXED_TENDON_DAMPING, self._fixed_tendon_damping), (TT.FIXED_TENDON_LIMIT_STIFFNESS, self._fixed_tendon_limit_stiffness), (TT.FIXED_TENDON_REST_LENGTH, self._fixed_tendon_rest_length), (TT.FIXED_TENDON_OFFSET, self._fixed_tendon_offset), ]: - np_buf = _read_cpu(tt) - if np_buf is not None and dst is not None: - wp.copy(dst, wp.from_numpy(np_buf, dtype=wp.float32, device=self.device)) - # Fixed tendon limits: [N, T, 2] -> (N, T) wp.vec2f - np_tlim = _read_cpu(TT.FIXED_TENDON_LIMIT) - if np_tlim is not None and self._fixed_tendon_pos_limits is not None: - self._fixed_tendon_pos_limits = wp.from_numpy( - np_tlim.reshape(self._num_instances, T_fix, 2), dtype=wp.vec2f, device=self.device - ) - - # Spatial tendon properties (CPU-side, read once) - T_spa = getattr(self, "_num_spatial_tendons", 0) + binding = self._get_binding(tt) + if binding is not None: + self._binding_read(tt, binding, buf.data) + buf.timestamp = self._sim_timestamp + binding = self._get_binding(TT.FIXED_TENDON_LIMIT) + if binding is not None: + self._binding_read(TT.FIXED_TENDON_LIMIT, binding, self._fixed_tendon_pos_limits.data) + self._fixed_tendon_pos_limits.timestamp = self._sim_timestamp + + # Spatial tendon properties (sim-device, see fixed-tendon comment above). + T_spa = self._num_spatial_tendons if T_spa > 0: - for tt, dst in [ + for tt, buf in [ (TT.SPATIAL_TENDON_STIFFNESS, self._spatial_tendon_stiffness), (TT.SPATIAL_TENDON_DAMPING, self._spatial_tendon_damping), (TT.SPATIAL_TENDON_LIMIT_STIFFNESS, self._spatial_tendon_limit_stiffness), (TT.SPATIAL_TENDON_OFFSET, self._spatial_tendon_offset), ]: - np_buf = _read_cpu(tt) - if np_buf is not None and dst is not None: - wp.copy(dst, wp.from_numpy(np_buf, dtype=wp.float32, device=self.device)) + binding = self._get_binding(tt) + if binding is not None: + self._binding_read(tt, binding, buf.data) + buf.timestamp = self._sim_timestamp def _pin_proxy_arrays(self) -> None: """Create pinned ProxyArray wrappers for all data buffers. - This is called once from :meth:`_create_buffers` during initialization. + Called once from :meth:`_create_buffers` during initialization. All ``_ta`` fields are lazily populated on first property access. """ - # -- Pinned ProxyArray cache (one per read property, lazily created on first access) # Defaults self._default_root_pose_ta: ProxyArray | None = None self._default_root_vel_ta: ProxyArray | None = None @@ -1535,6 +1725,8 @@ def _pin_proxy_arrays(self) -> None: self._joint_damping_ta: ProxyArray | None = None self._joint_armature_ta: ProxyArray | None = None self._joint_friction_coeff_ta: ProxyArray | None = None + self._joint_dynamic_friction_coeff_ta: ProxyArray | None = None + self._joint_viscous_friction_coeff_ta: ProxyArray | None = None self._joint_pos_limits_ta: ProxyArray | None = None self._joint_vel_limits_ta: ProxyArray | None = None self._joint_effort_limits_ta: ProxyArray | None = None @@ -1563,8 +1755,10 @@ def _pin_proxy_arrays(self) -> None: self._body_link_pose_w_ta: ProxyArray | None = None self._body_link_vel_w_ta: ProxyArray | None = None self._body_com_pose_w_ta: ProxyArray | None = None + self._body_com_vel_w_ta: ProxyArray | None = None self._body_com_acc_w_ta: ProxyArray | None = None self._body_com_pose_b_ta: ProxyArray | None = None + self._body_incoming_joint_wrench_b_ta: ProxyArray | None = None # Body properties self._body_mass_ta: ProxyArray | None = None self._body_inertia_ta: ProxyArray | None = None @@ -1606,13 +1800,75 @@ def _pin_proxy_arrays(self) -> None: self._body_com_quat_b_ta: ProxyArray | None = None # Deprecated state-concat properties self._default_root_state_ta: ProxyArray | None = None + self._root_state_w_ta: ProxyArray | None = None + self._root_link_state_w_ta: ProxyArray | None = None + self._root_com_state_w_ta: ProxyArray | None = None + # Deprecated body state-concat properties + self._body_state_w_ta: ProxyArray | None = None + self._body_link_state_w_ta: ProxyArray | None = None + self._body_com_state_w_ta: ProxyArray | None = None + + # Create strided wp.array views into _joint_friction_props_buf.data so that + # each friction component is accessible without copying data. The combined + # buffer has shape (N, D, 3) and contiguous float32 storage, so component k + # lives at byte offset k*4 with strides (D*3*4, 3*4). + N = self._num_instances + D = self._num_joints + _fp = self._joint_friction_props_buf.data + _float_bytes = 4 # sizeof(float32) + _stride_row = D * 3 * _float_bytes # bytes between rows + _stride_col = 3 * _float_bytes # bytes between columns (elements) + _dev = str(_fp.device) + self._joint_friction_coeff = wp.array( + ptr=_fp.ptr, + shape=(N, D), + strides=(_stride_row, _stride_col), + dtype=wp.float32, + device=_dev, + copy=False, + ) + self._joint_dynamic_friction_coeff = wp.array( + ptr=_fp.ptr + _float_bytes, + shape=(N, D), + strides=(_stride_row, _stride_col), + dtype=wp.float32, + device=_dev, + copy=False, + ) + self._joint_viscous_friction_coeff = wp.array( + ptr=_fp.ptr + 2 * _float_bytes, + shape=(N, D), + strides=(_stride_row, _stride_col), + dtype=wp.float32, + device=_dev, + copy=False, + ) - """ - Internal helpers -- Bindings. - """ + def _invalidate_initialize_callback(self, event) -> None: + """Invalidate cached buffers when the simulation is reinitialized. + + Args: + event: Simulation event (unused). + """ + self._is_primed = False + self._sim_timestamp = 0.0 + # Reset every TimestampedBuffer timestamp so the next property access + # triggers a fresh pull from the binding. + for attr_name in dir(self): + if attr_name.startswith("_") and not attr_name.startswith("__"): + val = getattr(self, attr_name, None) + if isinstance(val, TimestampedBuffer): + val.timestamp = -1.0 def _get_binding(self, tensor_type: int): - """Return a binding, lazily creating it if a binding_getter was provided.""" + """Return a binding, lazily creating it if a binding_getter was provided. + + Args: + tensor_type: TensorType key. + + Returns: + The TensorBinding, or ``None`` if not available. + """ b = self._bindings.get(tensor_type) if b is not None: return b @@ -1623,26 +1879,6 @@ def _get_binding(self, tensor_type: int): return b return None - def _get_read_scratch(self, tensor_type: int) -> wp.array | None: - """Return a pre-allocated flat float32 scratch buffer for a binding. - - Allocated once on first use, then reused every step. CPU-only - bindings (body properties, DOF properties) get CPU scratch; GPU - bindings get GPU scratch. wp.copy handles cross-device transfer - when the destination buffer lives on a different device. - """ - if tensor_type in self._read_scratch: - return self._read_scratch[tensor_type] - binding = self._get_binding(tensor_type) - if binding is None: - return None - from isaaclab_ovphysx.tensor_types import _CPU_ONLY_TYPES - - dev = "cpu" if tensor_type in _CPU_ONLY_TYPES else self.device - buf = wp.zeros(binding.shape, dtype=wp.float32, device=dev) - self._read_scratch[tensor_type] = buf - return buf - def _get_read_view(self, tensor_type: int, wp_array: wp.array, floats_per_elem: int = 0) -> wp.array | None: """Return a stable float32 view of a warp buffer for reading from a binding. @@ -1652,6 +1888,16 @@ def _get_read_view(self, tensor_type: int, wp_array: wp.array, floats_per_elem: The returned view is cached so that ``binding.read(view)`` sees the same object on every call, enabling the binding's internal read cache. + + Args: + tensor_type: TensorType key. + wp_array: Destination warp array. + floats_per_elem: Number of float32 elements per logical element + (e.g. 7 for transformf, 6 for spatial_vectorf). Pass 0 to + return the array as-is. + + Returns: + Float32 view suitable for ``binding.read()``, or ``None``. """ if not hasattr(self, "_read_view_cache"): self._read_view_cache = {} @@ -1679,65 +1925,89 @@ def _get_read_view(self, tensor_type: int, wp_array: wp.array, floats_per_elem: self._read_view_cache[cache_key] = view return view - def _read_binding_into_flat(self, tensor_type: int, wp_array: wp.array) -> None: - """Read a flat binding (no structured dtype) into an existing warp array. + def _read_binding_into_buf(self, tensor_type: int, buf: TimestampedBuffer) -> None: + """Read from an ovphysx binding into a :class:`TimestampedBuffer`, skipping if fresh. - Reads directly into the target array -- no scratch buffer, no extra copy. + Args: + tensor_type: TensorType key. + buf: Timestamped buffer to refresh. """ - self._read_binding_into_view(tensor_type, wp_array) - - def _read_binding_into_view(self, tensor_type: int, view: wp.array) -> None: - """Read an ovphysx binding into a float32 warp view.""" - binding = self._get_binding(tensor_type) - if binding is None: - return - - from isaaclab_ovphysx.tensor_types import _CPU_ONLY_TYPES - - if tensor_type in _CPU_ONLY_TYPES and str(view.device) != "cpu": - scratch = self._get_read_scratch(tensor_type) - if scratch is None: - return - binding.read(scratch) - wp.copy(view, scratch) - else: - binding.read(view) - - def _read_binding_into_buf(self, tensor_type: int, buf: TimestampedBuffer) -> None: - """Read from an ovphysx binding into a TimestampedBuffer, skipping if fresh.""" if buf.timestamp >= self._sim_timestamp: return view = self._get_read_view(tensor_type, buf.data) if view is None: return - self._read_binding_into_view(tensor_type, view) + self._get_binding(tensor_type).read(view) buf.timestamp = self._sim_timestamp def _read_transform_binding(self, tensor_type: int, buf: TimestampedBuffer) -> None: - """Read a pose binding (float32 view of transformf buffer), skipping if fresh.""" + """Read a pose binding (float32 view of transformf buffer), skipping if fresh. + + CPU-only bindings (e.g. ``BODY_COM_POSE``) are routed through a + pinned-host staging buffer via :meth:`_binding_read` so the wheel's + device-match requirement is satisfied even on a GPU sim. + + Args: + tensor_type: TensorType key. + buf: Timestamped :class:`wp.transformf` buffer to refresh. + """ if buf.timestamp >= self._sim_timestamp: return + binding = self._get_binding(tensor_type) + if binding is None: + return view = self._get_read_view(tensor_type, buf.data, 7) if view is None: return - self._read_binding_into_view(tensor_type, view) + self._binding_read(tensor_type, binding, view) buf.timestamp = self._sim_timestamp def _read_spatial_vector_binding(self, tensor_type: int, buf: TimestampedBuffer) -> None: - """Read a velocity binding (float32 view of spatial_vectorf buffer), skipping if fresh.""" + """Read a velocity binding (float32 view of spatial_vectorf buffer), skipping if fresh. + + Args: + tensor_type: TensorType key. + buf: Timestamped :class:`wp.spatial_vectorf` buffer to refresh. + """ if buf.timestamp >= self._sim_timestamp: return view = self._get_read_view(tensor_type, buf.data, 6) if view is None: return - self._read_binding_into_view(tensor_type, view) + self._get_binding(tensor_type).read(view) buf.timestamp = self._sim_timestamp - """ - Internal helpers -- Extraction. - """ + def _read_scalar_binding(self, tensor_type: int, buf: TimestampedBuffer) -> None: + """Refresh a scalar or flat float32 buffer from the matching binding if stale. + + Identical timestamp-gating contract as :meth:`_read_transform_binding` + but without a structured-dtype reinterpret cast. CPU-only bindings + (e.g. ``DOF_STIFFNESS``, ``DOF_LIMIT``) are routed through a + pre-allocated pinned-host staging buffer via :meth:`_binding_read` so + the wheel's device-match requirement is satisfied even on a GPU sim. + + Args: + tensor_type: TensorType key identifying the binding. + buf: Timestamped buffer whose :attr:`~TimestampedBuffer.data` field + will be refreshed. + """ + if buf.timestamp >= self._sim_timestamp: + return + binding = self._get_binding(tensor_type) + if binding is None: + return + self._binding_read(tensor_type, binding, buf.data) + buf.timestamp = self._sim_timestamp def _get_pos_from_transform(self, transform: wp.array) -> wp.array: + """Return a position view aliased into a transform array. + + Args: + transform: Source transform array. + + Returns: + vec3f view into the position component. + """ return wp.array( ptr=transform.ptr, shape=transform.shape, @@ -1747,6 +2017,14 @@ def _get_pos_from_transform(self, transform: wp.array) -> wp.array: ) def _get_quat_from_transform(self, transform: wp.array) -> wp.array: + """Return a quaternion view aliased into a transform array. + + Args: + transform: Source transform array. + + Returns: + quatf view into the quaternion component (offset 3 floats = 12 bytes). + """ return wp.array( ptr=transform.ptr + 3 * 4, shape=transform.shape, @@ -1756,6 +2034,14 @@ def _get_quat_from_transform(self, transform: wp.array) -> wp.array: ) def _get_lin_vel_from_spatial_vector(self, sv: wp.array) -> wp.array: + """Return a linear velocity view aliased into a spatial vector array. + + Args: + sv: Source spatial vector array. + + Returns: + vec3f view into the linear velocity component. + """ return wp.array( ptr=sv.ptr, shape=sv.shape, @@ -1765,6 +2051,14 @@ def _get_lin_vel_from_spatial_vector(self, sv: wp.array) -> wp.array: ) def _get_ang_vel_from_spatial_vector(self, sv: wp.array) -> wp.array: + """Return an angular velocity view aliased into a spatial vector array. + + Args: + sv: Source spatial vector array. + + Returns: + vec3f view into the angular velocity component (offset 3 floats = 12 bytes). + """ return wp.array( ptr=sv.ptr + 3 * 4, shape=sv.shape, @@ -1772,3 +2066,197 @@ def _get_ang_vel_from_spatial_vector(self, sv: wp.array) -> wp.array: strides=sv.strides, device=self.device, ) + + """ + Deprecated properties. + """ + + @property + def default_root_state(self) -> ProxyArray: + """Deprecated. Use :attr:`default_root_pose` and :attr:`default_root_vel` instead. + + Shape is (num_instances,), dtype = ``vec13f``. In torch this resolves to (num_instances, 13). + """ + warnings.warn( + "default_root_state is deprecated. Use default_root_pose and default_root_vel.", + DeprecationWarning, + stacklevel=2, + ) + wp.launch( + concat_root_pose_and_vel_to_state, + dim=self.num_instances, + inputs=[self._default_root_pose, self._default_root_vel], + outputs=[self._default_root_state_buf], + device=self.device, + ) + if self._default_root_state_ta is None: + self._default_root_state_ta = ProxyArray(self._default_root_state_buf) + return self._default_root_state_ta + + @property + def root_state_w(self) -> ProxyArray: + """Deprecated. Use :attr:`root_link_pose_w` and :attr:`root_com_vel_w` instead. + + Shape is (num_instances,), dtype = ``vec13f``. In torch this resolves to (num_instances, 13). + """ + warnings.warn( + "The `root_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_link_pose_w` and " + "`root_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._root_state_w_buf.timestamp < self._sim_timestamp: + wp.launch( + concat_root_pose_and_vel_to_state, + dim=self.num_instances, + inputs=[self.root_link_pose_w, self.root_com_vel_w], + outputs=[self._root_state_w_buf.data], + device=self.device, + ) + self._root_state_w_buf.timestamp = self._sim_timestamp + if self._root_state_w_ta is None: + self._root_state_w_ta = ProxyArray(self._root_state_w_buf.data) + return self._root_state_w_ta + + @property + def root_link_state_w(self) -> ProxyArray: + """Deprecated. Use :attr:`root_link_pose_w` and :attr:`root_link_vel_w` instead. + + Shape is (num_instances,), dtype = ``vec13f``. In torch this resolves to (num_instances, 13). + """ + warnings.warn( + "The `root_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_link_pose_w` and " + "`root_link_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._root_link_state_w_buf.timestamp < self._sim_timestamp: + wp.launch( + concat_root_pose_and_vel_to_state, + dim=self.num_instances, + inputs=[self.root_link_pose_w, self.root_link_vel_w], + outputs=[self._root_link_state_w_buf.data], + device=self.device, + ) + self._root_link_state_w_buf.timestamp = self._sim_timestamp + if self._root_link_state_w_ta is None: + self._root_link_state_w_ta = ProxyArray(self._root_link_state_w_buf.data) + return self._root_link_state_w_ta + + @property + def root_com_state_w(self) -> ProxyArray: + """Deprecated. Use :attr:`root_com_pose_w` and :attr:`root_com_vel_w` instead. + + Shape is (num_instances,), dtype = ``vec13f``. In torch this resolves to (num_instances, 13). + """ + warnings.warn( + "The `root_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_com_pose_w` and " + "`root_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._root_com_state_w_buf.timestamp < self._sim_timestamp: + wp.launch( + concat_root_pose_and_vel_to_state, + dim=self.num_instances, + inputs=[self.root_com_pose_w, self.root_com_vel_w], + outputs=[self._root_com_state_w_buf.data], + device=self.device, + ) + self._root_com_state_w_buf.timestamp = self._sim_timestamp + if self._root_com_state_w_ta is None: + self._root_com_state_w_ta = ProxyArray(self._root_com_state_w_buf.data) + return self._root_com_state_w_ta + + @property + def body_state_w(self) -> ProxyArray: + """Deprecated. Use :attr:`body_link_pose_w` and :attr:`body_com_vel_w` instead. + + Shape is (num_instances, num_bodies), dtype = ``vec13f``. + In torch this resolves to (num_instances, num_bodies, 13). + """ + warnings.warn( + "The `body_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " + "`body_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_state_w_buf.timestamp >= self._sim_timestamp: + if self._body_state_w_ta is None: + self._body_state_w_ta = ProxyArray(self._body_state_w_buf.data) + return self._body_state_w_ta + _ = self.body_link_pose_w + _ = self.body_com_vel_w + wp.launch( + concat_body_pose_and_vel_to_state, + dim=(self.num_instances, self.num_bodies), + inputs=[self._body_link_pose_w.data, self._body_com_vel_w.data], + outputs=[self._body_state_w_buf.data], + device=self.device, + ) + self._body_state_w_buf.timestamp = self._sim_timestamp + if self._body_state_w_ta is None: + self._body_state_w_ta = ProxyArray(self._body_state_w_buf.data) + return self._body_state_w_ta + + @property + def body_link_state_w(self) -> ProxyArray: + """Deprecated. Use :attr:`body_link_pose_w` and :attr:`body_link_vel_w` instead. + + Shape is (num_instances, num_bodies), dtype = ``vec13f``. + In torch this resolves to (num_instances, num_bodies, 13). + """ + warnings.warn( + "The `body_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " + "`body_link_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_link_state_w_buf.timestamp >= self._sim_timestamp: + if self._body_link_state_w_ta is None: + self._body_link_state_w_ta = ProxyArray(self._body_link_state_w_buf.data) + return self._body_link_state_w_ta + _ = self.body_link_pose_w + _ = self.body_link_vel_w + wp.launch( + concat_body_pose_and_vel_to_state, + dim=(self.num_instances, self.num_bodies), + inputs=[self._body_link_pose_w.data, self._body_link_vel_w.data], + outputs=[self._body_link_state_w_buf.data], + device=self.device, + ) + self._body_link_state_w_buf.timestamp = self._sim_timestamp + if self._body_link_state_w_ta is None: + self._body_link_state_w_ta = ProxyArray(self._body_link_state_w_buf.data) + return self._body_link_state_w_ta + + @property + def body_com_state_w(self) -> ProxyArray: + """Deprecated. Use :attr:`body_com_pose_w` and :attr:`body_com_vel_w` instead. + + Shape is (num_instances, num_bodies), dtype = ``vec13f``. + In torch this resolves to (num_instances, num_bodies, 13). + """ + warnings.warn( + "The `body_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_com_pose_w` and " + "`body_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_com_state_w_buf.timestamp >= self._sim_timestamp: + if self._body_com_state_w_ta is None: + self._body_com_state_w_ta = ProxyArray(self._body_com_state_w_buf.data) + return self._body_com_state_w_ta + _ = self.body_com_pose_w + _ = self.body_com_vel_w + wp.launch( + concat_body_pose_and_vel_to_state, + dim=(self.num_instances, self.num_bodies), + inputs=[self._body_com_pose_w.data, self._body_com_vel_w.data], + outputs=[self._body_com_state_w_buf.data], + device=self.device, + ) + self._body_com_state_w_buf.timestamp = self._sim_timestamp + if self._body_com_state_w_ta is None: + self._body_com_state_w_ta = ProxyArray(self._body_com_state_w_buf.data) + return self._body_com_state_w_ta diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/kernels.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/kernels.py index b93c4e6d4b41..9b9d9ff3d3c5 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/kernels.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/kernels.py @@ -3,44 +3,11 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Warp kernels for the ovphysx articulation.""" - import warp as wp - -@wp.kernel -def _body_wrench_to_world( - force_b: wp.array(dtype=wp.vec3f, ndim=2), - torque_b: wp.array(dtype=wp.vec3f, ndim=2), - poses: wp.array(dtype=wp.transformf, ndim=2), - wrench_out: wp.array(dtype=wp.float32, ndim=3), -): - """Rotate body-frame force/torque to world frame and pack into [N, L, 9].""" - i, j = wp.tid() - q = wp.transform_get_rotation(poses[i, j]) - f_w = wp.quat_rotate(q, force_b[i, j]) - t_w = wp.quat_rotate(q, torque_b[i, j]) - wrench_out[i, j, 0] = f_w[0] - wrench_out[i, j, 1] = f_w[1] - wrench_out[i, j, 2] = f_w[2] - wrench_out[i, j, 3] = t_w[0] - wrench_out[i, j, 4] = t_w[1] - wrench_out[i, j, 5] = t_w[2] - p_w = wp.transform_get_translation(poses[i, j]) - wrench_out[i, j, 6] = p_w[0] - wrench_out[i, j, 7] = p_w[1] - wrench_out[i, j, 8] = p_w[2] - - -@wp.kernel -def _scatter_rows_partial( - dst: wp.array2d(dtype=wp.float32), - src: wp.array2d(dtype=wp.float32), - ids: wp.array(dtype=wp.int32), -): - """dst[ids[i], j] = src[i, j] -- scatter partial [K,C] into full [N,C] on GPU.""" - i, j = wp.tid() - dst[ids[i], j] = src[i, j] +""" +Articulation-specific warp functions. +""" @wp.func @@ -48,7 +15,18 @@ def compute_soft_joint_pos_limits_func( joint_pos_limits: wp.vec2f, soft_limit_factor: wp.float32, ): - """Compute soft joint position limits from hard limits.""" + """Compute the soft joint position limits. + + Args: + joint_pos_limits: Hard joint position limits as ``(lower, upper)`` [m or rad, + depending on joint type]. + soft_limit_factor: Scale factor in [0, 1] shrinking the soft range around + the midpoint of the hard range; ``1.0`` makes the soft limits equal the + hard limits, smaller values create a tighter window. + + Returns: + The soft joint position limits as ``(lower, upper)``. + """ joint_pos_mean = (joint_pos_limits[0] + joint_pos_limits[1]) / 2.0 joint_pos_range = joint_pos_limits[1] - joint_pos_limits[0] return wp.vec2f( @@ -57,19 +35,8 @@ def compute_soft_joint_pos_limits_func( ) -@wp.kernel -def update_soft_joint_pos_limits( - joint_pos_limits: wp.array2d(dtype=wp.vec2f), - soft_limit_factor: wp.float32, - soft_joint_pos_limits: wp.array2d(dtype=wp.vec2f), -): - """Update soft joint position limits from hard limits and a scale factor.""" - i, j = wp.tid() - soft_joint_pos_limits[i, j] = compute_soft_joint_pos_limits_func(joint_pos_limits[i, j], soft_limit_factor) - - """ -Data-layer kernels (used by ArticulationData). +Articulation-specific warp kernels. """ @@ -80,13 +47,20 @@ def _fd_joint_acc( inv_dt: float, out: wp.array2d(dtype=wp.float32), ): - """Compute joint acceleration via finite differencing and update previous velocity. + """Compute the joint acceleration via finite differencing and update the previous velocity. + + Diverges from PhysX's :func:`get_joint_acc_from_joint_vel` in taking the inverse + time step rather than ``dt`` itself; the multiply-by-reciprocal avoids per-element + division inside the kernel. Args: - cur_vel: Current joint velocities. Shape is (num_envs, num_joints). - prev_vel: Previous joint velocities (updated in-place). Shape is (num_envs, num_joints). - inv_dt: Inverse time step (1/dt) [1/s]. - out: Output joint accelerations. Shape is (num_envs, num_joints). + cur_vel: Current joint velocities [m/s or rad/s, depending on joint type]. + Shape is (num_envs, num_joints). + prev_vel: Previous joint velocities (updated in-place). Same shape and units + as :paramref:`cur_vel`. + inv_dt: Inverse time step ``1 / dt`` [1/s]. + out: Output joint accelerations [m/s^2 or rad/s^2, depending on joint type]. + Shape is (num_envs, num_joints). """ i, j = wp.tid() out[i, j] = (cur_vel[i, j] - prev_vel[i, j]) * inv_dt @@ -94,124 +68,201 @@ def _fd_joint_acc( @wp.kernel -def _copy_first_body( - body_vel: wp.array(dtype=wp.spatial_vectorf, ndim=2), - root_vel: wp.array(dtype=wp.spatial_vectorf), +def _compose_body_com_poses( + link_pose: wp.array(dtype=wp.transformf, ndim=2), + com_pose_b: wp.array(dtype=wp.transformf, ndim=2), + com_pose_w: wp.array(dtype=wp.transformf, ndim=2), ): - """Copy the first body's velocity to the root velocity buffer. + """Compose body link poses with body-frame CoM offsets to get world-frame CoM poses. Args: - body_vel: Body velocities. Shape is (num_envs, num_bodies). - root_vel: Output root velocities. Shape is (num_envs,). + link_pose: Body link poses in world frame [m, m, m, qx, qy, qz, qw]. + Shape is (num_envs, num_bodies). + com_pose_b: Body-frame CoM offsets [m, m, m, qx, qy, qz, qw]. + Shape is (num_envs, num_bodies). + com_pose_w: Output world-frame body CoM poses [m, m, m, qx, qy, qz, qw]. + Shape is (num_envs, num_bodies). """ - i = wp.tid() - root_vel[i] = body_vel[i, 0] + i, j = wp.tid() + com_pose_w[i, j] = wp.transform_multiply(link_pose[i, j], com_pose_b[i, j]) @wp.kernel -def _compose_root_com_pose( - link_pose: wp.array(dtype=wp.transformf), - com_pose_b: wp.array(dtype=wp.transformf, ndim=2), - com_pose_w: wp.array(dtype=wp.transformf), +def update_soft_joint_pos_limits( + joint_pos_limits: wp.array2d(dtype=wp.vec2f), + soft_limit_factor: wp.float32, + soft_joint_pos_limits: wp.array2d(dtype=wp.vec2f), ): - """Compose root link pose with body-frame CoM offset to get world-frame root CoM pose. + """Update soft joint position limits from hard limits and a soft limit factor. + + Soft limits provide a safety margin before reaching the hard joint position + limits. See :func:`compute_soft_joint_pos_limits_func` for the per-joint + formula. Args: - link_pose: Root link poses in world frame. Shape is (num_envs,). - com_pose_b: Body-frame CoM offsets. Shape is (num_envs, num_bodies). - com_pose_w: Output world-frame root CoM poses. Shape is (num_envs,). + joint_pos_limits: Hard joint position limits as vec2f ``(lower, upper)`` + [m or rad, depending on joint type]. Shape is (num_envs, num_joints). + soft_limit_factor: Scale factor in [0, 1]. ``1.0`` makes the soft limits + equal the hard limits; smaller values create a tighter window. + soft_joint_pos_limits: Output array. Shape is (num_envs, num_joints). """ - i = wp.tid() - com_pose_w[i] = wp.transform_multiply(link_pose[i], com_pose_b[i, 0]) + i, j = wp.tid() + soft_joint_pos_limits[i, j] = compute_soft_joint_pos_limits_func(joint_pos_limits[i, j], soft_limit_factor) @wp.kernel -def _compose_body_com_poses( - link_pose: wp.array(dtype=wp.transformf, ndim=2), - com_pose_b: wp.array(dtype=wp.transformf, ndim=2), - com_pose_w: wp.array(dtype=wp.transformf, ndim=2), +def clamp_default_joint_pos_and_update_soft_limits_index( + joint_pos_limits: wp.array2d(dtype=wp.vec2f), + env_ids: wp.array(dtype=wp.int32), + joint_ids: wp.array(dtype=wp.int32), + soft_limit_factor: wp.float32, + default_joint_pos: wp.array2d(dtype=wp.float32), + soft_joint_pos_limits: wp.array2d(dtype=wp.vec2f), + clamped_count: wp.array(dtype=wp.int32), ): - """Compose body link poses with body-frame CoM offsets to get world-frame CoM poses. + """Clamp default joint positions to new limits and refresh soft limits over (env_ids x joint_ids). - Args: - link_pose: Body link poses in world frame. Shape is (num_envs, num_bodies). - com_pose_b: Body-frame CoM offsets. Shape is (num_envs, num_bodies). - com_pose_w: Output world-frame body CoM poses. Shape is (num_envs, num_bodies). - """ - i, j = wp.tid() - com_pose_w[i, j] = wp.transform_multiply(link_pose[i, j], com_pose_b[i, j]) + Mirrors PhysX's :func:`isaaclab_physx.assets.articulation.kernels.write_joint_limit_data_to_buffer` + side-effects, minus the limit-write itself (the existing + :func:`shared_kernels.write_joint_position_limit_to_buffer_index` launch handles that). + For each ``(i, j)`` thread the kernel: -@wp.kernel -def _projected_gravity( - gravity_vec_w: wp.array(dtype=wp.vec3f), - root_pose: wp.array(dtype=wp.transformf), - out: wp.array(dtype=wp.vec3f), -): - """Project world-frame gravity direction into the root body frame. + * Clamps :paramref:`default_joint_pos` ``[env_ids[i], joint_ids[j]]`` if it falls outside + the new limits, atomically incrementing :paramref:`clamped_count`. + * Recomputes :paramref:`soft_joint_pos_limits` ``[env_ids[i], joint_ids[j]]`` from the new + hard limits and :paramref:`soft_limit_factor`. Args: - gravity_vec_w: Gravity unit vector per instance in world frame. Shape is (num_envs,). - root_pose: Root link poses in world frame. Shape is (num_envs,). - out: Output projected gravity in body frame. Shape is (num_envs,). + joint_pos_limits: Hard joint position limits as vec2f ``(lower, upper)`` + [m or rad, depending on joint type]. Shape is (num_envs, num_joints). + env_ids: Environment indices to update. Shape is (num_selected_envs,). + joint_ids: Joint indices to update. Shape is (num_selected_joints,). + soft_limit_factor: Scale factor in [0, 1] for the soft limit window. + default_joint_pos: In/out default joint positions [m or rad, depending on joint type]. + Shape is (num_envs, num_joints). + soft_joint_pos_limits: Out soft joint position limits as vec2f ``(lower, upper)`` + [m or rad, depending on joint type]. Shape is (num_envs, num_joints). + clamped_count: One-element output counter incremented atomically each time a + default joint position was clamped. Shape is (1,). """ - i = wp.tid() - q = wp.transform_get_rotation(root_pose[i]) - out[i] = wp.quat_rotate_inv(q, gravity_vec_w[i]) + i, j = wp.tid() + e = env_ids[i] + k = joint_ids[j] + lo = joint_pos_limits[e, k][0] + hi = joint_pos_limits[e, k][1] + if (default_joint_pos[e, k] < lo) or (default_joint_pos[e, k] > hi): + wp.atomic_add(clamped_count, 0, 1) + default_joint_pos[e, k] = wp.clamp(default_joint_pos[e, k], lo, hi) + soft_joint_pos_limits[e, k] = compute_soft_joint_pos_limits_func(joint_pos_limits[e, k], soft_limit_factor) @wp.kernel -def _compute_heading( - forward_vec_b: wp.array(dtype=wp.vec3f), - root_pose: wp.array(dtype=wp.transformf), - out: wp.array(dtype=wp.float32), +def clamp_default_joint_pos_and_update_soft_limits_mask( + joint_pos_limits: wp.array2d(dtype=wp.vec2f), + env_mask: wp.array(dtype=wp.bool), + joint_mask: wp.array(dtype=wp.bool), + soft_limit_factor: wp.float32, + default_joint_pos: wp.array2d(dtype=wp.float32), + soft_joint_pos_limits: wp.array2d(dtype=wp.vec2f), + clamped_count: wp.array(dtype=wp.int32), ): - """Compute yaw heading angle from the forward direction rotated into the world frame. + """Mask variant of :func:`clamp_default_joint_pos_and_update_soft_limits_index`. + + Iterates the full ``(num_envs, num_joints)`` grid and applies the clamp / + soft-limit refresh only where both :paramref:`env_mask` and :paramref:`joint_mask` + are ``True``. Args: - forward_vec_b: Forward direction in body frame per instance. Shape is (num_envs,). - root_pose: Root link poses in world frame. Shape is (num_envs,). - out: Output heading angles [rad]. Shape is (num_envs,). + joint_pos_limits: Hard joint position limits as vec2f ``(lower, upper)`` + [m or rad, depending on joint type]. Shape is (num_envs, num_joints). + env_mask: Boolean mask over environments. Shape is (num_envs,). + joint_mask: Boolean mask over joints. Shape is (num_joints,). + soft_limit_factor: Scale factor in [0, 1] for the soft limit window. + default_joint_pos: In/out default joint positions [m or rad, depending on joint type]. + Shape is (num_envs, num_joints). + soft_joint_pos_limits: Out soft joint position limits as vec2f ``(lower, upper)`` + [m or rad, depending on joint type]. Shape is (num_envs, num_joints). + clamped_count: One-element output counter incremented atomically each time a + default joint position was clamped. Shape is (1,). """ - i = wp.tid() - q = wp.transform_get_rotation(root_pose[i]) - forward = wp.quat_rotate(q, forward_vec_b[i]) - out[i] = wp.atan2(forward[1], forward[0]) + i, j = wp.tid() + if not env_mask[i] or not joint_mask[j]: + return + lo = joint_pos_limits[i, j][0] + hi = joint_pos_limits[i, j][1] + if (default_joint_pos[i, j] < lo) or (default_joint_pos[i, j] > hi): + wp.atomic_add(clamped_count, 0, 1) + default_joint_pos[i, j] = wp.clamp(default_joint_pos[i, j], lo, hi) + soft_joint_pos_limits[i, j] = compute_soft_joint_pos_limits_func(joint_pos_limits[i, j], soft_limit_factor) @wp.kernel -def _world_vel_to_body_lin( - root_pose: wp.array(dtype=wp.transformf), - vel_w: wp.array(dtype=wp.spatial_vectorf), - out: wp.array(dtype=wp.vec3f), +def write_joint_friction_data_to_buffer_index( + in_static: wp.array2d(dtype=wp.float32), + in_dynamic: wp.array2d(dtype=wp.float32), + in_viscous: wp.array2d(dtype=wp.float32), + env_ids: wp.array(dtype=wp.int32), + joint_ids: wp.array(dtype=wp.int32), + out_buffer: wp.array3d(dtype=wp.float32), ): - """Rotate world-frame linear velocity into the root body frame. + """Conditionally update the static / dynamic / viscous slots of the friction buffer. + + Mirrors :func:`isaaclab_physx.assets.articulation.kernels.write_joint_friction_data_to_buffer`: + each of the three input arrays is optional (``None`` translates to a null pointer + which evaluates ``False`` inside the kernel), so callers can update any subset + of the friction components without disturbing the others. Args: - root_pose: Root link poses in world frame. Shape is (num_envs,). - vel_w: Spatial velocities in world frame. Shape is (num_envs,). - out: Output linear velocity in body frame. Shape is (num_envs,). + in_static: Static friction coefficients, or ``None`` to leave that component + unchanged. Shape is (num_selected_envs, num_selected_joints). + in_dynamic: Dynamic friction coefficients, or ``None``. Same shape as + :paramref:`in_static`. + in_viscous: Viscous friction coefficients [N·s/m or N·m·s/rad, depending on + joint type], or ``None``. Same shape as :paramref:`in_static`. + env_ids: Environment indices to write. Shape is (num_selected_envs,). + joint_ids: Joint indices to write. Shape is (num_selected_joints,). + out_buffer: Combined friction buffer. Shape is (num_envs, num_joints, 3) with + slots [0] static, [1] dynamic, [2] viscous. """ - i = wp.tid() - q = wp.transform_get_rotation(root_pose[i]) - lin = wp.spatial_top(vel_w[i]) - out[i] = wp.quat_rotate_inv(q, lin) + i, j = wp.tid() + if in_static: + out_buffer[env_ids[i], joint_ids[j], 0] = in_static[i, j] + if in_dynamic: + out_buffer[env_ids[i], joint_ids[j], 1] = in_dynamic[i, j] + if in_viscous: + out_buffer[env_ids[i], joint_ids[j], 2] = in_viscous[i, j] @wp.kernel -def _world_vel_to_body_ang( - root_pose: wp.array(dtype=wp.transformf), - vel_w: wp.array(dtype=wp.spatial_vectorf), - out: wp.array(dtype=wp.vec3f), +def write_joint_friction_data_to_buffer_mask( + in_static: wp.array2d(dtype=wp.float32), + in_dynamic: wp.array2d(dtype=wp.float32), + in_viscous: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), + joint_mask: wp.array(dtype=wp.bool), + out_buffer: wp.array3d(dtype=wp.float32), ): - """Rotate world-frame angular velocity into the root body frame. + """Mask variant of :func:`write_joint_friction_data_to_buffer_index`. Args: - root_pose: Root link poses in world frame. Shape is (num_envs,). - vel_w: Spatial velocities in world frame. Shape is (num_envs,). - out: Output angular velocity in body frame. Shape is (num_envs,). + in_static: Static friction coefficients, or ``None`` to leave that component + unchanged. Shape is (num_envs, num_joints). + in_dynamic: Dynamic friction coefficients, or ``None``. Same shape as + :paramref:`in_static`. + in_viscous: Viscous friction coefficients [N·s/m or N·m·s/rad, depending on + joint type], or ``None``. Same shape as :paramref:`in_static`. + env_mask: Boolean mask over environments. Shape is (num_envs,). + joint_mask: Boolean mask over joints. Shape is (num_joints,). + out_buffer: Combined friction buffer. Shape is (num_envs, num_joints, 3) with + slots [0] static, [1] dynamic, [2] viscous. """ - i = wp.tid() - q = wp.transform_get_rotation(root_pose[i]) - ang = wp.spatial_bottom(vel_w[i]) - out[i] = wp.quat_rotate_inv(q, ang) + i, j = wp.tid() + if not env_mask[i] or not joint_mask[j]: + return + if in_static: + out_buffer[i, j, 0] = in_static[i, j] + if in_dynamic: + out_buffer[i, j, 1] = in_dynamic[i, j] + if in_viscous: + out_buffer[i, j, 2] = in_viscous[i, j] diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/kernels.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/kernels.py new file mode 100644 index 000000000000..5951ad7be31e --- /dev/null +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/kernels.py @@ -0,0 +1,1401 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import warp as wp + +vec13f = wp.types.vector(length=13, dtype=wp.float32) + +""" +Shared @wp.func helpers. +""" + + +@wp.func +def get_link_vel_from_root_com_vel_func( + com_vel: wp.spatial_vectorf, + link_pose: wp.transformf, + body_com_pose: wp.transformf, +): + """Compute link velocity from center-of-mass velocity. + + Transforms a COM spatial velocity into a link-frame velocity by projecting + the angular velocity contribution from the COM offset relative to the link frame. + + Args: + com_vel: COM spatial velocity (angular, linear). + link_pose: Link pose in world frame. + body_com_pose: COM pose in body (link) frame. + + Returns: + Link spatial velocity (angular, linear). + """ + projected_vel = wp.cross( + wp.spatial_bottom(com_vel), + wp.quat_rotate(wp.transform_get_rotation(link_pose), -wp.transform_get_translation(body_com_pose)), + ) + return wp.spatial_vector(wp.spatial_top(com_vel) + projected_vel, wp.spatial_bottom(com_vel)) + + +@wp.func +def get_com_pose_from_link_pose_func( + link_pose: wp.transformf, + body_com_pose: wp.transformf, +): + """Compute COM pose in world frame from link pose and body-frame COM offset. + + Args: + link_pose: Link pose in world frame. + body_com_pose: COM pose in body (link) frame. + + Returns: + COM pose in world frame. + """ + return link_pose * body_com_pose + + +@wp.func +def concat_pose_and_vel_to_state_func( + pose: wp.transformf, + vel: wp.spatial_vectorf, +) -> vec13f: + """Concatenate a pose and velocity into a 13-element state vector. + + The state vector layout is [pos(3), quat(4), ang_vel(3), lin_vel(3)]. + + Args: + pose: Pose as a transform (position + quaternion). + vel: Spatial velocity (angular, linear). + + Returns: + 13-element state vector. + """ + return vec13f( + pose[0], pose[1], pose[2], pose[3], pose[4], pose[5], pose[6], vel[0], vel[1], vel[2], vel[3], vel[4], vel[5] + ) + + +@wp.func +def compute_heading_w_func( + forward_vec: wp.vec3f, + quat: wp.quatf, +): + """Compute heading angle (yaw) in world frame from a forward vector and orientation. + + Rotates the forward vector by the quaternion and computes atan2(y, x). + + Args: + forward_vec: Forward direction vector in body frame. + quat: Orientation quaternion. + + Returns: + Heading angle in radians. + """ + forward_w = wp.quat_rotate(quat, forward_vec) + return wp.atan2(forward_w[1], forward_w[0]) + + +@wp.func +def set_state_transforms_func( + state: vec13f, + transform: wp.transformf, +) -> vec13f: + """Set the pose portion (first 7 elements) of a 13-element state vector. + + Overwrites elements [0..6] (position + quaternion) with the given transform, + leaving the velocity portion [7..12] unchanged. + + Args: + state: 13-element state vector to modify. + transform: New pose (position + quaternion). + + Returns: + Updated 13-element state vector. + """ + state[0] = transform[0] + state[1] = transform[1] + state[2] = transform[2] + state[3] = transform[3] + state[4] = transform[4] + state[5] = transform[5] + state[6] = transform[6] + return state + + +@wp.func +def set_state_velocities_func( + state: vec13f, + velocity: wp.spatial_vectorf, +) -> vec13f: + """Set the velocity portion (last 6 elements) of a 13-element state vector. + + Overwrites elements [7..12] (angular + linear velocity) with the given spatial velocity, + leaving the pose portion [0..6] unchanged. + + Args: + state: 13-element state vector to modify. + velocity: New spatial velocity (angular, linear). + + Returns: + Updated 13-element state vector. + """ + state[7] = velocity[0] + state[8] = velocity[1] + state[9] = velocity[2] + state[10] = velocity[3] + state[11] = velocity[4] + state[12] = velocity[5] + return state + + +@wp.func +def get_link_velocity_in_com_frame_func( + link_velocity_w: wp.spatial_vectorf, + link_pose_w: wp.transformf, + body_com_pose_b: wp.transformf, +): + """Compute COM velocity from link velocity by accounting for the COM offset. + + Transforms a link-frame spatial velocity into a COM-frame velocity by adding + the cross-product contribution of the COM offset rotated into the world frame. + + Args: + link_velocity_w: Link spatial velocity in world frame (angular, linear). + link_pose_w: Link pose in world frame. + body_com_pose_b: COM pose in body (link) frame. + + Returns: + COM spatial velocity in world frame (angular, linear). + """ + return wp.spatial_vector( + wp.spatial_top(link_velocity_w) + + wp.cross( + wp.spatial_bottom(link_velocity_w), + wp.quat_rotate(wp.transform_get_rotation(link_pose_w), wp.transform_get_translation(body_com_pose_b)), + ), + wp.spatial_bottom(link_velocity_w), + ) + + +@wp.func +def get_com_pose_in_link_frame_func( + com_pose_w: wp.transformf, + com_pose_b: wp.transformf, +): + """Compute link pose in world frame from COM pose by inverting the body-frame COM offset. + + This is the inverse of ``get_com_pose_from_link_pose_func``. Given the COM pose in + world frame and the COM offset in body frame, it recovers the link pose in world frame. + + Args: + com_pose_w: COM pose in world frame. + com_pose_b: COM pose in body (link) frame. + + Returns: + Link pose in world frame. + """ + T2 = wp.transform( + wp.quat_rotate( + wp.quat_inverse(wp.transform_get_rotation(com_pose_b)), -wp.transform_get_translation(com_pose_b) + ), + wp.quat_inverse(wp.transform_get_rotation(com_pose_b)), + ) + link_pose_w = com_pose_w * T2 + return link_pose_w + + +""" +Root-level @wp.kernel (1D — used by RigidObject + Articulation). +""" + + +@wp.kernel +def get_root_link_vel_from_root_com_vel( + com_vel: wp.array(dtype=wp.spatial_vectorf), + link_pose: wp.array(dtype=wp.transformf), + body_com_pose_b: wp.array2d(dtype=wp.transformf), + link_vel: wp.array(dtype=wp.spatial_vectorf), +): + """Compute root link velocity from root center-of-mass velocity. + + This kernel transforms the root COM velocity into link-frame velocity by projecting + the angular velocity contribution from the COM offset. + + Args: + com_vel: Input array of root COM spatial velocities. Shape is (num_envs,). + link_pose: Input array of root link poses in world frame. Shape is (num_envs,). + body_com_pose_b: Input array of body COM poses in body frame. Shape is (num_envs, num_bodies). + Only the first body (index 0) is used for the root. + link_vel: Output array where root link velocities are written. Shape is (num_envs,). + """ + i = wp.tid() + link_vel[i] = get_link_vel_from_root_com_vel_func(com_vel[i], link_pose[i], body_com_pose_b[i, 0]) + + +@wp.kernel +def get_root_com_pose_from_root_link_pose( + link_pose: wp.array(dtype=wp.transformf), + body_com_pose_b: wp.array2d(dtype=wp.transformf), + com_pose_w: wp.array(dtype=wp.transformf), +): + """Compute root COM pose from root link pose. + + This kernel transforms the root link pose to the root COM pose using the body COM offset. + + Args: + link_pose: Input array of root link poses in world frame. Shape is (num_envs,). + body_com_pose_b: Input array of body COM poses in body frame. Shape is (num_envs, num_bodies). + Only the first body (index 0) is used for the root. + com_pose_w: Output array where root COM poses are written. Shape is (num_envs,). + """ + i = wp.tid() + com_pose_w[i] = get_com_pose_from_link_pose_func(link_pose[i], body_com_pose_b[i, 0]) + + +@wp.kernel +def concat_root_pose_and_vel_to_state( + pose: wp.array(dtype=wp.transformf), + vel: wp.array(dtype=wp.spatial_vectorf), + state: wp.array(dtype=vec13f), +): + """Concatenate root pose and velocity into a 13-element state vector. + + This kernel combines a 7-element pose (pos + quat) and a 6-element velocity + (angular + linear) into a single 13-element state vector. + + Args: + pose: Input array of root poses in world frame. Shape is (num_envs,). + vel: Input array of root spatial velocities. Shape is (num_envs,). + state: Output array where concatenated state vectors are written. Shape is (num_envs,). + """ + i = wp.tid() + state[i] = concat_pose_and_vel_to_state_func(pose[i], vel[i]) + + +@wp.kernel +def split_state_to_root_pose_and_vel( + state: wp.array2d(dtype=wp.float32), + pose: wp.array(dtype=wp.transformf), + vel: wp.array(dtype=wp.spatial_vectorf), +): + """Split a 13-element state vector into root pose and velocity. + + This kernel extracts a 7-element pose (pos + quat) and a 6-element velocity + (angular + linear) from a 13-element state vector. + + Args: + state: Input array of root states. Shape is (num_envs, 13). + pose: Output array where root poses are written. Shape is (num_envs,). + vel: Output array where root spatial velocities are written. Shape is (num_envs,). + """ + i = wp.tid() + # Extract pose: [pos(3), quat(4)] = state[0:7] + pose[i] = wp.transform( + wp.vec3f(state[i, 0], state[i, 1], state[i, 2]), wp.quatf(state[i, 3], state[i, 4], state[i, 5], state[i, 6]) + ) + # Extract velocity: [ang_vel(3), lin_vel(3)] = state[7:13] + vel[i] = wp.spatial_vector( + wp.vec3f(state[i, 7], state[i, 8], state[i, 9]), # angular velocity + wp.vec3f(state[i, 10], state[i, 11], state[i, 12]), # linear velocity + ) + + +""" +Body-level @wp.kernel (2D — used by Articulation + RigidObjectCollection). +""" + + +@wp.kernel +def get_body_link_vel_from_body_com_vel( + body_com_vel: wp.array2d(dtype=wp.spatial_vectorf), + body_link_pose: wp.array2d(dtype=wp.transformf), + body_com_pose: wp.array2d(dtype=wp.transformf), + body_link_vel: wp.array2d(dtype=wp.spatial_vectorf), +): + """Compute body link velocities from body COM velocities for all bodies. + + This kernel transforms COM velocities into link-frame velocities by projecting + the angular velocity contribution from the COM offset, for each body in each environment. + + Args: + body_com_vel: Input array of body COM spatial velocities. Shape is (num_envs, num_bodies). + body_link_pose: Input array of body link poses in world frame. Shape is (num_envs, num_bodies). + body_com_pose: Input array of body COM poses in body frame. Shape is (num_envs, num_bodies). + body_link_vel: Output array where body link velocities are written. Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + body_link_vel[i, j] = get_link_vel_from_root_com_vel_func( + body_com_vel[i, j], body_link_pose[i, j], body_com_pose[i, j] + ) + + +@wp.kernel +def get_body_com_pose_from_body_link_pose( + body_link_pose: wp.array2d(dtype=wp.transformf), + body_com_pose_b: wp.array2d(dtype=wp.transformf), + body_com_pose_w: wp.array2d(dtype=wp.transformf), +): + """Compute body COM poses from body link poses for all bodies. + + This kernel transforms link poses to COM poses using the body COM offset in the body frame. + + Args: + body_link_pose: Input array of body link poses in world frame. Shape is (num_envs, num_bodies). + body_com_pose_b: Input array of body COM poses in body frame. Shape is (num_envs, num_bodies). + body_com_pose_w: Output array where body COM poses in world frame are written. + Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + body_com_pose_w[i, j] = get_com_pose_from_link_pose_func(body_link_pose[i, j], body_com_pose_b[i, j]) + + +@wp.kernel +def concat_body_pose_and_vel_to_state( + pose: wp.array2d(dtype=wp.transformf), + vel: wp.array2d(dtype=wp.spatial_vectorf), + state: wp.array2d(dtype=vec13f), +): + """Concatenate body pose and velocity into 13-element state vectors for all bodies. + + This kernel combines a 7-element pose (pos + quat) and a 6-element velocity + (angular + linear) into a single 13-element state vector, for each body in each environment. + + Args: + pose: Input array of body poses in world frame. Shape is (num_envs, num_bodies). + vel: Input array of body spatial velocities. Shape is (num_envs, num_bodies). + state: Output array where concatenated state vectors are written. + Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + state[i, j] = concat_pose_and_vel_to_state_func(pose[i, j], vel[i, j]) + + +""" +Derived property kernels. +""" + + +@wp.kernel +def quat_apply_inverse_1D_kernel( + gravity: wp.array(dtype=wp.vec3f), + quat: wp.array(dtype=wp.quatf), + projected_gravity: wp.array(dtype=wp.vec3f), +): + """Apply inverse quaternion rotation to gravity vectors (1D). + + This kernel rotates gravity vectors into the local frame of each environment + using the inverse of the provided quaternion. + + Args: + gravity: Input array of gravity vectors in world frame. Shape is (num_envs,). + quat: Input array of quaternions representing orientations. Shape is (num_envs,). + projected_gravity: Output array where projected gravity vectors are written. + Shape is (num_envs,). + """ + i = wp.tid() + projected_gravity[i] = wp.quat_rotate_inv(quat[i], gravity[i]) + + +@wp.kernel +def root_heading_w( + forward_vec: wp.array(dtype=wp.vec3f), + quat: wp.array(dtype=wp.quatf), + heading_w: wp.array(dtype=wp.float32), +): + """Compute root heading angle in the world frame. + + This kernel computes the heading angle (yaw) by rotating the forward vector + by the root quaternion and computing atan2 of the resulting x and y components. + + Args: + forward_vec: Input array of forward direction vectors. Shape is (num_envs,). + quat: Input array of root quaternions. Shape is (num_envs,). + heading_w: Output array where heading angles (radians) are written. Shape is (num_envs,). + """ + i = wp.tid() + heading_w[i] = compute_heading_w_func(forward_vec[i], quat[i]) + + +@wp.kernel +def quat_apply_inverse_2D_kernel( + vec: wp.array2d(dtype=wp.vec3f), + quat: wp.array2d(dtype=wp.quatf), + result: wp.array2d(dtype=wp.vec3f), +): + """Apply inverse quaternion rotation to vectors (2D). + + This kernel rotates vectors into the local frame of each body in each environment + using the inverse of the provided quaternion. + + Args: + vec: Input array of vectors in world frame. Shape is (num_envs, num_bodies). + quat: Input array of quaternions representing orientations. Shape is (num_envs, num_bodies). + result: Output array where rotated vectors are written. Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + result[i, j] = wp.quat_rotate_inv(quat[i, j], vec[i, j]) + + +@wp.kernel +def body_heading_w( + forward_vec: wp.array2d(dtype=wp.vec3f), + quat: wp.array2d(dtype=wp.quatf), + heading_w: wp.array2d(dtype=wp.float32), +): + """Compute body heading angles in the world frame for all bodies. + + This kernel computes heading angles (yaw) by rotating forward vectors + by body quaternions and computing atan2 of the resulting x and y components. + + Args: + forward_vec: Input array of forward direction vectors. Shape is (num_envs, num_bodies). + quat: Input array of body quaternions. Shape is (num_envs, num_bodies). + heading_w: Output array where heading angles (radians) are written. + Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + heading_w[i, j] = compute_heading_w_func(forward_vec[i, j], quat[i, j]) + + +""" +Root-level write kernels (1D — used by RigidObject + Articulation). +""" + + +@wp.kernel +def set_root_link_pose_to_sim_index( + data: wp.array(dtype=wp.transformf), + env_ids: wp.array(dtype=wp.int32), + root_link_pose_w: wp.array(dtype=wp.transformf), +): + """Write root link pose data to simulation buffers. + + This kernel scatters root link poses from the partial input array into the cached + world-frame buffer at the specified environment indices. + + Args: + data: Input array of root link poses. Shape is (num_selected_envs,). + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + root_link_pose_w: Output array where root link poses are written. Shape is (num_envs,). + """ + i = wp.tid() + root_link_pose_w[env_ids[i]] = data[i] + + +@wp.kernel +def set_root_com_pose_to_sim_index( + data: wp.array(dtype=wp.transformf), + body_com_pose_b: wp.array2d(dtype=wp.transformf), + env_ids: wp.array(dtype=wp.int32), + root_com_pose_w: wp.array(dtype=wp.transformf), + root_link_pose_w: wp.array(dtype=wp.transformf), +): + """Write root COM pose data to simulation buffers. + + This kernel scatters root COM poses from the partial input array into the cached + world-frame buffer at the specified environment indices and derives the + corresponding link pose via the body-frame COM offset. + + Args: + data: Input array of root COM poses. Shape is (num_selected_envs,). + body_com_pose_b: Input array of body COM poses in body frame. Shape is + (num_envs, num_bodies). Only the first body (index 0) is used for the root. + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + root_com_pose_w: Output array where root COM poses are written. Shape is (num_envs,). + root_link_pose_w: Output array where root link poses (derived from COM) are written. + Shape is (num_envs,). + """ + i = wp.tid() + root_com_pose_w[env_ids[i]] = data[i] + # Get the com pose in the link frame + root_link_pose_w[env_ids[i]] = get_com_pose_in_link_frame_func( + root_com_pose_w[env_ids[i]], body_com_pose_b[env_ids[i], 0] + ) + + +@wp.kernel +def set_root_com_velocity_to_sim_index( + data: wp.array(dtype=wp.spatial_vectorf), + env_ids: wp.array(dtype=wp.int32), + num_bodies: wp.int32, + root_com_velocity_w: wp.array(dtype=wp.spatial_vectorf), + body_acc_w: wp.array2d(dtype=wp.spatial_vectorf), +): + """Write root COM velocity data to simulation buffers. + + This kernel scatters root COM velocities from the partial input array into the cached + world-frame buffer at the specified environment indices and zeros the body acceleration + buffer to prevent reporting stale values. + + Args: + data: Input array of root COM spatial velocities. Shape is (num_selected_envs,). + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + num_bodies: Input scalar number of bodies per environment. + root_com_velocity_w: Output array where root COM velocities are written. Shape is (num_envs,). + body_acc_w: Output array where body accelerations are zeroed. Shape is + (num_envs, num_bodies). + """ + i = wp.tid() + root_com_velocity_w[env_ids[i]] = data[i] + # Make the acceleration zero to prevent reporting old values + for j in range(num_bodies): + body_acc_w[env_ids[i], j] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + + +@wp.kernel +def set_root_link_velocity_to_sim_index( + data: wp.array(dtype=wp.spatial_vectorf), + body_com_pose_b: wp.array2d(dtype=wp.transformf), + link_pose_w: wp.array(dtype=wp.transformf), + env_ids: wp.array(dtype=wp.int32), + num_bodies: wp.int32, + root_link_velocity_w: wp.array(dtype=wp.spatial_vectorf), + root_com_velocity_w: wp.array(dtype=wp.spatial_vectorf), + body_acc_w: wp.array2d(dtype=wp.spatial_vectorf), +): + """Write root link velocity data to simulation buffers. + + This kernel scatters root link velocities from the partial input array into the cached + world-frame buffer at the specified environment indices, derives the corresponding + COM velocity via the lever-arm transform, and zeros the body acceleration buffer. + + Args: + data: Input array of root link spatial velocities. Shape is (num_selected_envs,). + body_com_pose_b: Input array of body COM poses in body frame. Shape is + (num_envs, num_bodies). Only the first body (index 0) is used for the root. + link_pose_w: Input array of root link poses in world frame. Shape is (num_envs,). + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + num_bodies: Input scalar number of bodies per environment. + root_link_velocity_w: Output array where root link velocities are written. + Shape is (num_envs,). + root_com_velocity_w: Output array where root COM velocities (derived from link) + are written. Shape is (num_envs,). + body_acc_w: Output array where body accelerations are zeroed. + Shape is (num_envs, num_bodies). + """ + i = wp.tid() + root_link_velocity_w[env_ids[i]] = data[i] + # Get the link velocity in the com frame + root_com_velocity_w[env_ids[i]] = get_link_velocity_in_com_frame_func( + root_link_velocity_w[env_ids[i]], link_pose_w[env_ids[i]], body_com_pose_b[env_ids[i], 0] + ) + # Make the acceleration zero to prevent reporting old values + for j in range(num_bodies): + body_acc_w[env_ids[i], j] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + + +""" +Body-level write kernels (2D — used by RigidObjectCollection). +""" + + +@wp.kernel +def set_body_link_pose_to_sim( + data: wp.array2d(dtype=wp.transformf), + env_ids: wp.array(dtype=wp.int32), + body_ids: wp.array(dtype=wp.int32), + from_mask: bool, + body_link_pose_w: wp.array2d(dtype=wp.transformf), + body_link_state_w: wp.array2d(dtype=vec13f), + body_state_w: wp.array2d(dtype=vec13f), +): + """Write body link pose data to simulation buffers. + + This kernel writes body link poses from the input array to the output buffers + and optionally updates the corresponding state vectors, for each body in each environment. + + Args: + data: Input array of body link poses. Shape is (num_envs, num_bodies) or + (num_selected_envs, num_selected_bodies) depending on from_mask. + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + body_ids: Input array of body indices to write to. Shape is (num_selected_bodies,). + from_mask: Input flag indicating whether to use masked indexing. + body_link_pose_w: Output array where body link poses are written. + Shape is (num_envs, num_bodies). + body_link_state_w: Output array where body link states are updated (pose portion). + Shape is (num_envs, num_bodies). Can be None if not needed. + body_state_w: Output array where body states are updated (pose portion). + Shape is (num_envs, num_bodies). Can be None if not needed. + """ + i, j = wp.tid() + if from_mask: + body_link_pose_w[env_ids[i], body_ids[j]] = data[env_ids[i], body_ids[j]] + if body_link_state_w: + body_link_state_w[env_ids[i], body_ids[j]] = set_state_transforms_func( + body_link_state_w[env_ids[i], body_ids[j]], data[env_ids[i], body_ids[j]] + ) + if body_state_w: + body_state_w[env_ids[i], body_ids[j]] = set_state_transforms_func( + body_state_w[env_ids[i], body_ids[j]], data[env_ids[i], body_ids[j]] + ) + else: + body_link_pose_w[env_ids[i], body_ids[j]] = data[i, j] + if body_link_state_w: + body_link_state_w[env_ids[i], body_ids[j]] = set_state_transforms_func( + body_link_state_w[env_ids[i], body_ids[j]], data[i, j] + ) + if body_state_w: + body_state_w[env_ids[i], body_ids[j]] = set_state_transforms_func( + body_state_w[env_ids[i], body_ids[j]], data[i, j] + ) + + +@wp.kernel +def set_body_com_pose_to_sim( + data: wp.array2d(dtype=wp.transformf), + body_com_pose_b: wp.array2d(dtype=wp.transformf), + env_ids: wp.array(dtype=wp.int32), + body_ids: wp.array(dtype=wp.int32), + from_mask: bool, + body_com_pose_w: wp.array2d(dtype=wp.transformf), + body_link_pose_w: wp.array2d(dtype=wp.transformf), + body_com_state_w: wp.array2d(dtype=vec13f), + body_link_state_w: wp.array2d(dtype=vec13f), + body_state_w: wp.array2d(dtype=vec13f), +): + """Write body COM pose data to simulation buffers. + + This kernel writes body COM poses from the input array to the output buffers, + computes the corresponding link poses from the COM poses, and optionally updates + the corresponding state vectors, for each body in each environment. + + Args: + data: Input array of body COM poses. Shape is (num_envs, num_bodies) or + (num_selected_envs, num_selected_bodies) depending on from_mask. + body_com_pose_b: Input array of body COM poses in body frame. Shape is + (num_envs, num_bodies). + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + body_ids: Input array of body indices to write to. Shape is (num_selected_bodies,). + from_mask: Input flag indicating whether to use masked indexing. + 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). + body_com_state_w: Output array where body COM states are updated (pose portion). + Shape is (num_envs, num_bodies). Can be None if not needed. + body_link_state_w: Output array where body link states are updated (pose portion). + Shape is (num_envs, num_bodies). Can be None if not needed. + body_state_w: Output array where body states are updated (pose portion). + Shape is (num_envs, num_bodies). Can be None if not needed. + """ + i, j = wp.tid() + if from_mask: + body_com_pose_w[env_ids[i], body_ids[j]] = data[env_ids[i], body_ids[j]] + if body_com_state_w: + body_com_state_w[env_ids[i], body_ids[j]] = set_state_transforms_func( + body_com_state_w[env_ids[i], body_ids[j]], data[env_ids[i], body_ids[j]] + ) + else: + body_com_pose_w[env_ids[i], body_ids[j]] = data[i, j] + if body_com_state_w: + body_com_state_w[env_ids[i], body_ids[j]] = set_state_transforms_func( + body_com_state_w[env_ids[i], body_ids[j]], data[i, j] + ) + # Get the link pose from com pose + body_link_pose_w[env_ids[i], body_ids[j]] = get_com_pose_in_link_frame_func( + body_com_pose_w[env_ids[i], body_ids[j]], body_com_pose_b[env_ids[i], body_ids[j]] + ) + if body_link_state_w: + body_link_state_w[env_ids[i], body_ids[j]] = set_state_transforms_func( + body_link_state_w[env_ids[i], body_ids[j]], body_link_pose_w[env_ids[i], body_ids[j]] + ) + if body_state_w: + body_state_w[env_ids[i], body_ids[j]] = set_state_transforms_func( + body_state_w[env_ids[i], body_ids[j]], body_link_pose_w[env_ids[i], body_ids[j]] + ) + + +@wp.kernel +def set_body_com_velocity_to_sim( + data: wp.array2d(dtype=wp.spatial_vectorf), + env_ids: wp.array(dtype=wp.int32), + body_ids: wp.array(dtype=wp.int32), + from_mask: bool, + body_com_velocity_w: wp.array2d(dtype=wp.spatial_vectorf), + body_acc_w: wp.array2d(dtype=wp.spatial_vectorf), + body_state_w: wp.array2d(dtype=vec13f), + body_com_state_w: wp.array2d(dtype=vec13f), +): + """Write body COM velocity data to simulation buffers. + + This kernel writes body COM velocities from the input array to the output buffers, + optionally updates the corresponding state vectors, and zeros out the body + acceleration buffer, for each body in each environment. + + Args: + data: Input array of body COM spatial velocities. Shape is (num_envs, num_bodies) or + (num_selected_envs, num_selected_bodies) depending on from_mask. + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + body_ids: Input array of body indices to write to. Shape is (num_selected_bodies,). + from_mask: Input flag indicating whether to use masked indexing. + body_com_velocity_w: Output array where body COM velocities are written. + Shape is (num_envs, num_bodies). + body_acc_w: Output array where body accelerations are zeroed. + Shape is (num_envs, num_bodies). + body_state_w: Output array where body states are updated (velocity portion). + Shape is (num_envs, num_bodies). Can be None if not needed. + body_com_state_w: Output array where body COM states are updated (velocity portion). + Shape is (num_envs, num_bodies). Can be None if not needed. + """ + i, j = wp.tid() + if from_mask: + body_com_velocity_w[env_ids[i], body_ids[j]] = data[env_ids[i], body_ids[j]] + if body_state_w: + body_state_w[env_ids[i], body_ids[j]] = set_state_velocities_func( + body_state_w[env_ids[i], body_ids[j]], data[env_ids[i], body_ids[j]] + ) + if body_com_state_w: + body_com_state_w[env_ids[i], body_ids[j]] = set_state_velocities_func( + body_com_state_w[env_ids[i], body_ids[j]], data[env_ids[i], body_ids[j]] + ) + else: + body_com_velocity_w[env_ids[i], body_ids[j]] = data[i, j] + if body_state_w: + body_state_w[env_ids[i], body_ids[j]] = set_state_velocities_func( + body_state_w[env_ids[i], body_ids[j]], data[i, j] + ) + if body_com_state_w: + body_com_state_w[env_ids[i], body_ids[j]] = set_state_velocities_func( + body_com_state_w[env_ids[i], body_ids[j]], data[i, j] + ) + # Make the acceleration zero to prevent reporting old values + body_acc_w[env_ids[i], body_ids[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( + data: wp.array2d(dtype=wp.spatial_vectorf), + body_com_pose_b: wp.array2d(dtype=wp.transformf), + body_link_pose_w: wp.array2d(dtype=wp.transformf), + env_ids: wp.array(dtype=wp.int32), + body_ids: wp.array(dtype=wp.int32), + from_mask: 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), + body_link_state_w: wp.array2d(dtype=vec13f), + body_state_w: wp.array2d(dtype=vec13f), + body_com_state_w: wp.array2d(dtype=vec13f), +): + """Write body link velocity data to simulation buffers. + + This kernel writes body link velocities from the input array to the output buffers, + computes the corresponding COM velocities from the link velocities, optionally updates + the corresponding state vectors, and zeros out the body acceleration buffer. + + Args: + data: Input array of body link spatial velocities. Shape is (num_envs, num_bodies) + or (num_selected_envs, num_selected_bodies) depending on from_mask. + body_com_pose_b: Input array of body COM poses 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_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + body_ids: Input array of body indices to write to. Shape is (num_selected_bodies,). + from_mask: Input flag indicating whether to use masked indexing. + 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). + body_link_state_w: Output array where body link states are updated (velocity portion). + Shape is (num_envs, num_bodies). Can be None if not needed. + body_state_w: Output array where body states are updated (velocity portion). + Shape is (num_envs, num_bodies). Can be None if not needed. + body_com_state_w: Output array where body COM states are updated (velocity portion). + Shape is (num_envs, num_bodies). Can be None if not needed. + """ + i, j = wp.tid() + if from_mask: + body_link_velocity_w[env_ids[i], body_ids[j]] = data[env_ids[i], body_ids[j]] + if body_link_state_w: + body_link_state_w[env_ids[i], body_ids[j]] = set_state_velocities_func( + body_link_state_w[env_ids[i], body_ids[j]], data[env_ids[i], body_ids[j]] + ) + else: + body_link_velocity_w[env_ids[i], body_ids[j]] = data[i, j] + if body_link_state_w: + body_link_state_w[env_ids[i], body_ids[j]] = set_state_velocities_func( + body_link_state_w[env_ids[i], body_ids[j]], data[i, j] + ) + # Get the link velocity in the com frame + body_com_velocity_w[env_ids[i], body_ids[j]] = get_link_velocity_in_com_frame_func( + body_link_velocity_w[env_ids[i], body_ids[j]], + body_link_pose_w[env_ids[i], body_ids[j]], + body_com_pose_b[env_ids[i], body_ids[j]], + ) + if body_com_state_w: + body_com_state_w[env_ids[i], body_ids[j]] = set_state_velocities_func( + body_com_state_w[env_ids[i], body_ids[j]], body_com_velocity_w[env_ids[i], body_ids[j]] + ) + if body_state_w: + body_state_w[env_ids[i], body_ids[j]] = set_state_velocities_func( + body_state_w[env_ids[i], body_ids[j]], body_com_velocity_w[env_ids[i], body_ids[j]] + ) + # Make the acceleration zero to prevent reporting old values + body_acc_w[env_ids[i], body_ids[j]] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + + +""" +Generic buffer-writing kernels (used by Articulation + RigidObject + RigidObjectCollection). +""" + + +@wp.kernel +def write_2d_data_to_buffer_with_indices( + in_data: wp.array2d(dtype=wp.float32), + env_ids: wp.array(dtype=wp.int32), + joint_ids: wp.array(dtype=wp.int32), + out_data: wp.array2d(dtype=wp.float32), +): + """Write 2D float data to a buffer at specified indices. + + This kernel copies float data from a partial input array to an output buffer at the + specified environment and joint/body indices. + + Args: + in_data: Input array containing float data. Shape is (num_selected_envs, num_selected_joints). + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + joint_ids: Input array of joint/body indices to write to. Shape is (num_selected_joints,). + out_data: Output array where data is written. Shape is (num_envs, num_joints). + """ + i, j = wp.tid() + out_data[env_ids[i], joint_ids[j]] = in_data[i, j] + + +@wp.kernel +def write_body_inertia_to_buffer_index( + in_data: wp.array3d(dtype=wp.float32), + env_ids: wp.array(dtype=wp.int32), + body_ids: wp.array(dtype=wp.int32), + out_data: wp.array3d(dtype=wp.float32), +): + """Write body inertia data to a buffer at specified indices. + + This kernel copies 3x3 inertia tensor data (stored as 9 floats) from a partial input + array to an output buffer at the specified environment and body indices. + + Args: + in_data: Input array containing inertia data. Shape is (num_selected_envs, num_selected_bodies, 9). + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + body_ids: Input array of body indices to write to. Shape is (num_selected_bodies,). + out_data: Output array where inertia data is written. Shape is (num_envs, num_bodies, 9). + """ + i, j = wp.tid() + for k in range(9): + out_data[env_ids[i], body_ids[j], k] = in_data[i, j, k] + + +@wp.kernel +def write_body_com_pose_to_buffer_index( + in_data: wp.array2d(dtype=wp.transformf), + env_ids: wp.array(dtype=wp.int32), + body_ids: wp.array(dtype=wp.int32), + out_data: wp.array2d(dtype=wp.transformf), +): + """Write body COM pose data to a buffer at specified indices. + + This kernel copies body COM pose data from a partial input array to an output buffer + at the specified environment and body indices. + + Args: + in_data: Input array containing body COM poses. Shape is (num_selected_envs, num_selected_bodies). + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + body_ids: Input array of body indices to write to. Shape is (num_selected_bodies,). + out_data: Output array where body COM poses are written. Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + out_data[env_ids[i], body_ids[j]] = in_data[i, j] + + +@wp.kernel +def derive_body_acceleration_from_body_com_velocities( + body_com_vel: wp.array2d(dtype=wp.spatial_vectorf), + dt: wp.float32, + prev_body_com_vel: wp.array2d(dtype=wp.spatial_vectorf), + body_acc: wp.array2d(dtype=wp.spatial_vectorf), +): + """Derive body acceleration from body COM velocities. + + This kernel derives body acceleration from body COM velocities using finite differencing. + + Args: + body_com_vel: Input array of body COM velocities. Shape is (num_envs, num_bodies). + dt: Input time step (scalar) used for finite differencing. + prev_body_com_vel: Input/output array of previous body COM velocities. Shape is (num_envs, num_bodies). + body_acc: Output array where body accelerations are written. Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + # Compute the acceleration + body_acc[i, j] = (body_com_vel[i, j] - prev_body_com_vel[i, j]) / dt + # Update the previous body COM velocity + prev_body_com_vel[i, j] = body_com_vel[i, j] + + +@wp.kernel +def _body_wrench_to_world( + force_b: wp.array(dtype=wp.vec3f, ndim=2), + torque_b: wp.array(dtype=wp.vec3f, ndim=2), + poses: wp.array(dtype=wp.transformf, ndim=2), + wrench_out: wp.array(dtype=wp.float32, ndim=3), +): + """Rotate body-frame force/torque to world frame and pack into a flat output array. + + Output layout per ``(i, j)`` slice (9 floats total): + + * ``[0:3]`` -- world-frame force ``[N]`` + * ``[3:6]`` -- world-frame torque ``[N*m]`` + * ``[6:9]`` -- world-frame link position ``[m]`` + + Args: + force_b: Body-frame applied forces ``[N]``. Shape is ``(N, L)``. + torque_b: Body-frame applied torques ``[N*m]``. Shape is ``(N, L)``. + poses: Link poses in world frame. Shape is ``(N, L)``. + wrench_out: Output packed wrench array. Shape is ``(N, L, 9)``. + """ + i, j = wp.tid() + q = wp.transform_get_rotation(poses[i, j]) + f_w = wp.quat_rotate(q, force_b[i, j]) + t_w = wp.quat_rotate(q, torque_b[i, j]) + wrench_out[i, j, 0] = f_w[0] + wrench_out[i, j, 1] = f_w[1] + wrench_out[i, j, 2] = f_w[2] + wrench_out[i, j, 3] = t_w[0] + wrench_out[i, j, 4] = t_w[1] + wrench_out[i, j, 5] = t_w[2] + p_w = wp.transform_get_translation(poses[i, j]) + wrench_out[i, j, 6] = p_w[0] + wrench_out[i, j, 7] = p_w[1] + wrench_out[i, j, 8] = p_w[2] + + +@wp.kernel +def _scatter_rows_partial( + dst: wp.array2d(dtype=wp.float32), + src: wp.array2d(dtype=wp.float32), + ids: wp.array(dtype=wp.int32), +): + """Scatter a partial row-indexed source array into a larger destination array. + + For each thread ``(i, j)`` writes ``dst[ids[i], j] = src[i, j]``. + + Args: + dst: Destination array of shape ``(N, C)`` to scatter values into. + src: Source array of shape ``(K, C)`` containing the values to scatter. + ids: Row indices into ``dst`` for each row of ``src``. Shape is ``(K,)``. + """ + i, j = wp.tid() + dst[ids[i], j] = src[i, j] + + +""" +Native-mask scatter kernels (mirrors Newton; the OVPhysX wheel's ``binding.write`` natively +supports a boolean mask via the ``mask=`` argument, so the ``*_mask`` setters update the cache +in-place and pass the mask straight through to the wheel without a ``torch.nonzero`` round-trip). +""" + + +@wp.kernel +def set_root_link_pose_to_sim_mask( + data: wp.array(dtype=wp.transformf), + env_mask: wp.array(dtype=wp.bool), + root_link_pose_w: wp.array(dtype=wp.transformf), +): + """Mask-scatter root link poses into the cache; rows where ``env_mask[i]`` is False are untouched.""" + i = wp.tid() + if env_mask[i]: + root_link_pose_w[i] = data[i] + + +@wp.kernel +def set_root_com_pose_to_sim_mask( + data: wp.array(dtype=wp.transformf), + body_com_pose_b: wp.array2d(dtype=wp.transformf), + env_mask: wp.array(dtype=wp.bool), + root_com_pose_w: wp.array(dtype=wp.transformf), + root_link_pose_w: wp.array(dtype=wp.transformf), +): + """Mask-scatter root COM poses into the cache and derive the corresponding link poses.""" + i = wp.tid() + if env_mask[i]: + root_com_pose_w[i] = data[i] + # link_pose = com_pose * inverse(com_pose_b) + root_link_pose_w[i] = wp.transform_multiply(root_com_pose_w[i], wp.transform_inverse(body_com_pose_b[i, 0])) + + +@wp.kernel +def set_root_com_velocity_to_sim_mask( + data: wp.array(dtype=wp.spatial_vectorf), + env_mask: wp.array(dtype=wp.bool), + num_bodies: wp.int32, + root_com_velocity_w: wp.array(dtype=wp.spatial_vectorf), + body_acc_w: wp.array2d(dtype=wp.spatial_vectorf), +): + """Mask-scatter root COM velocities into the cache and zero the dependent body acceleration.""" + i = wp.tid() + if env_mask[i]: + root_com_velocity_w[i] = data[i] + for j in range(num_bodies): + body_acc_w[i, j] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + + +@wp.kernel +def set_root_link_velocity_to_sim_mask( + data: wp.array(dtype=wp.spatial_vectorf), + body_com_pose_b: wp.array2d(dtype=wp.transformf), + link_pose_w: wp.array(dtype=wp.transformf), + env_mask: wp.array(dtype=wp.bool), + num_bodies: wp.int32, + root_link_velocity_w: wp.array(dtype=wp.spatial_vectorf), + root_com_velocity_w: wp.array(dtype=wp.spatial_vectorf), + body_acc_w: wp.array2d(dtype=wp.spatial_vectorf), +): + """Mask-scatter root link velocities into the cache and derive the corresponding COM velocities + via the lever-arm transform: ``com_lin = link_lin + omega x rot(link_rot, com_offset)``. + """ + i = wp.tid() + if env_mask[i]: + root_link_velocity_w[i] = data[i] + ang = wp.spatial_bottom(data[i]) + lever = wp.quat_rotate( + wp.transform_get_rotation(link_pose_w[i]), wp.transform_get_translation(body_com_pose_b[i, 0]) + ) + com_lin = wp.spatial_top(data[i]) + wp.cross(ang, lever) + root_com_velocity_w[i] = wp.spatial_vector(com_lin, ang) + for j in range(num_bodies): + body_acc_w[i, j] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + + +@wp.kernel +def write_2d_data_to_buffer_with_mask( + in_data: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), + body_mask: wp.array(dtype=wp.bool), + out_data: wp.array2d(dtype=wp.float32), +): + """Mask-scatter 2D float data into the cache where both ``env_mask[i]`` and ``body_mask[j]`` are True.""" + i, j = wp.tid() + if env_mask[i] and body_mask[j]: + out_data[i, j] = in_data[i, j] + + +@wp.kernel +def write_body_inertia_to_buffer_mask( + in_data: wp.array3d(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), + body_mask: wp.array(dtype=wp.bool), + out_data: wp.array3d(dtype=wp.float32), +): + """Mask-scatter body inertia (3x3 = 9 floats per body) into the cache.""" + i, j = wp.tid() + if env_mask[i] and body_mask[j]: + for k in range(9): + out_data[i, j, k] = in_data[i, j, k] + + +@wp.kernel +def write_body_com_pose_to_buffer_mask( + in_data: wp.array2d(dtype=wp.transformf), + env_mask: wp.array(dtype=wp.bool), + body_mask: wp.array(dtype=wp.bool), + out_data: wp.array2d(dtype=wp.transformf), +): + """Mask-scatter body COM poses (transformf) into the cache.""" + i, j = wp.tid() + if env_mask[i] and body_mask[j]: + out_data[i, j] = in_data[i, j] + + +""" +Articulation-only kernels (used by isaaclab_ovphysx.assets.articulation). +""" + + +@wp.kernel +def _copy_first_body( + body_vel: wp.array(dtype=wp.spatial_vectorf, ndim=2), + root_vel: wp.array(dtype=wp.spatial_vectorf), +): + """Copy the first body's spatial velocity to the root velocity buffer. + + For single rigid-body assets, index 0 is always the root body. This + kernel extracts that slice without allocating an intermediate buffer. + + Args: + body_vel: Body spatial velocities ``[m/s, rad/s]``. Shape is + ``(num_envs, num_bodies)`` with dtype ``wp.spatial_vectorf``. + root_vel: Output root spatial velocities ``[m/s, rad/s]``. Shape is + ``(num_envs,)`` with dtype ``wp.spatial_vectorf``. + """ + i = wp.tid() + root_vel[i] = body_vel[i, 0] + + +@wp.kernel +def _compose_root_com_pose( + link_pose: wp.array(dtype=wp.transformf), + com_pose_b: wp.array(dtype=wp.transformf, ndim=2), + com_pose_w: wp.array(dtype=wp.transformf), +): + """Compose root link pose with the body-frame COM offset to get the world-frame COM pose. + + Implements the forward transform: + + ``com_pose_w = link_pose_w * com_pose_b[0]`` + + where ``*`` denotes ``wp.transform_multiply``. Only the first body + (index ``0``) is used; for articulations this is the base link body. + + Args: + link_pose: Root link poses in world frame ``[m, -]``. Shape is + ``(num_envs,)`` with dtype ``wp.transformf``. + com_pose_b: Body-frame COM offsets ``[m, -]`` from the + ``RIGID_BODY_COM_POSE`` binding. Shape is ``(num_envs, num_bodies)`` + with dtype ``wp.transformf``. + com_pose_w: Output world-frame root COM poses ``[m, -]``. Shape is + ``(num_envs,)`` with dtype ``wp.transformf``. + """ + i = wp.tid() + com_pose_w[i] = wp.transform_multiply(link_pose[i], com_pose_b[i, 0]) + + +@wp.kernel +def _projected_gravity( + gravity_vec_w: wp.array(dtype=wp.vec3f), + root_pose: wp.array(dtype=wp.transformf), + out: wp.array(dtype=wp.vec3f), +): + """Project the world-frame gravity direction into the root body frame. + + Applies the inverse of the root orientation quaternion to the world-frame + gravity vector, yielding the gravity direction expressed in the body frame. + The magnitude is preserved (unit vector in, unit vector out if input is a + unit vector). + + Args: + gravity_vec_w: Gravity direction per instance in world frame ``[-]`` + (typically the normalised ``(0, 0, -1)`` gravitational acceleration + direction). Shape is ``(num_envs,)`` with dtype ``wp.vec3f``. + root_pose: Root link poses in world frame ``[m, -]``. Only the + rotation component is used. Shape is ``(num_envs,)`` with dtype + ``wp.transformf``. + out: Output gravity direction in body frame ``[-]``. Shape is + ``(num_envs,)`` with dtype ``wp.vec3f``. + """ + i = wp.tid() + q = wp.transform_get_rotation(root_pose[i]) + out[i] = wp.quat_rotate_inv(q, gravity_vec_w[i]) + + +@wp.kernel +def _compute_heading( + forward_vec_b: wp.array(dtype=wp.vec3f), + root_pose: wp.array(dtype=wp.transformf), + out: wp.array(dtype=wp.float32), +): + """Compute the yaw heading angle by rotating a body-frame forward vector to world frame. + + Rotates ``forward_vec_b`` by the root orientation quaternion and then computes the + heading as ``atan2(forward_w.y, forward_w.x)`` ``[rad]``, i.e. the signed angle + from the world X-axis to the projected forward direction in the XY plane. + + Args: + forward_vec_b: Forward direction in body frame per instance ``[-]``. + Shape is ``(num_envs,)`` with dtype ``wp.vec3f``. + root_pose: Root link poses in world frame ``[m, -]``. Only the rotation + component is used. Shape is ``(num_envs,)`` with dtype ``wp.transformf``. + out: Output heading angles ``[rad]`` in ``[-π, π]``. Shape is + ``(num_envs,)`` with dtype ``wp.float32``. + """ + i = wp.tid() + q = wp.transform_get_rotation(root_pose[i]) + forward = wp.quat_rotate(q, forward_vec_b[i]) + out[i] = wp.atan2(forward[1], forward[0]) + + +@wp.kernel +def _world_vel_to_body_lin( + root_pose: wp.array(dtype=wp.transformf), + vel_w: wp.array(dtype=wp.spatial_vectorf), + out: wp.array(dtype=wp.vec3f), +): + """Rotate the world-frame linear velocity component into the root body frame. + + Extracts the linear velocity from the top three components of the spatial + velocity vector (``wp.spatial_top``) and rotates it by the inverse of the + root orientation quaternion. + + Args: + root_pose: Root link poses in world frame ``[m, -]``. Only the rotation + component is used. Shape is ``(num_envs,)`` with dtype ``wp.transformf``. + vel_w: Root spatial velocities in world frame ``[m/s, rad/s]``. + Shape is ``(num_envs,)`` with dtype ``wp.spatial_vectorf``. + out: Output linear velocity in body frame ``[m/s]``. Shape is + ``(num_envs,)`` with dtype ``wp.vec3f``. + """ + i = wp.tid() + q = wp.transform_get_rotation(root_pose[i]) + lin = wp.spatial_top(vel_w[i]) + out[i] = wp.quat_rotate_inv(q, lin) + + +@wp.kernel +def _world_vel_to_body_ang( + root_pose: wp.array(dtype=wp.transformf), + vel_w: wp.array(dtype=wp.spatial_vectorf), + out: wp.array(dtype=wp.vec3f), +): + """Rotate the world-frame angular velocity component into the root body frame. + + Extracts the angular velocity from the bottom three components of the spatial + velocity vector (``wp.spatial_bottom``) and rotates it by the inverse of the + root orientation quaternion. + + Args: + root_pose: Root link poses in world frame ``[m, -]``. Only the rotation + component is used. Shape is ``(num_envs,)`` with dtype ``wp.transformf``. + vel_w: Root spatial velocities in world frame ``[m/s, rad/s]``. + Shape is ``(num_envs,)`` with dtype ``wp.spatial_vectorf``. + out: Output angular velocity in body frame ``[rad/s]``. Shape is + ``(num_envs,)`` with dtype ``wp.vec3f``. + """ + i = wp.tid() + q = wp.transform_get_rotation(root_pose[i]) + ang = wp.spatial_bottom(vel_w[i]) + out[i] = wp.quat_rotate_inv(q, ang) + + +@wp.kernel +def write_joint_position_limit_to_buffer_index( + in_data: wp.array3d(dtype=wp.float32), + env_ids: wp.array(dtype=wp.int32), + joint_ids: wp.array(dtype=wp.int32), + out_data: wp.array(dtype=wp.vec2f, ndim=2), +): + """Write joint position-limit data to a vec2f buffer at specified indices. + + This kernel copies ``[lower, upper]`` limit pairs from a partial float32 input + array into the output ``wp.vec2f`` buffer at the specified environment and joint + indices. + + Args: + in_data: Input array containing limit pairs ``[lower, upper]`` [m or rad]. + Shape is (num_selected_envs, num_selected_joints, 2). + env_ids: Input array of environment indices to write to. + Shape is (num_selected_envs,). + joint_ids: Input array of joint indices to write to. + Shape is (num_selected_joints,). + out_data: Output array where limit data is written. Shape is + (num_envs, num_joints) with dtype ``wp.vec2f``. + """ + i, j = wp.tid() + out_data[env_ids[i], joint_ids[j]] = wp.vec2f(in_data[i, j, 0], in_data[i, j, 1]) + + +@wp.kernel +def write_joint_position_limit_to_buffer_mask( + in_data: wp.array3d(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), + joint_mask: wp.array(dtype=wp.bool), + out_data: wp.array(dtype=wp.vec2f, ndim=2), +): + """Mask-scatter joint position-limit data into the vec2f cache buffer. + + Copies ``[lower, upper]`` limit pairs where both ``env_mask[i]`` and + ``joint_mask[j]`` are True. + + Args: + in_data: Input array containing limit pairs ``[lower, upper]`` [m or rad]. + Shape is (num_envs, num_joints, 2). + env_mask: Boolean environment mask. Shape is (num_envs,). + joint_mask: Boolean joint mask. Shape is (num_joints,). + out_data: Output array where limit data is written. Shape is + (num_envs, num_joints) with dtype ``wp.vec2f``. + """ + i, j = wp.tid() + if env_mask[i] and joint_mask[j]: + out_data[i, j] = wp.vec2f(in_data[i, j, 0], in_data[i, j, 1]) + + +@wp.kernel +def write_joint_friction_to_buffer_index( + in_data: wp.array2d(dtype=wp.float32), + env_ids: wp.array(dtype=wp.int32), + joint_ids: wp.array(dtype=wp.int32), + out_data: wp.array3d(dtype=wp.float32), +): + """Write joint friction coefficient to all three components of the friction buffer. + + Broadcasts a single friction value into the static (index 0), dynamic (index 1), + and viscous (index 2) components of the ``(N, D, 3)`` friction properties buffer + at the specified environment and joint indices. + + Args: + in_data: Input friction coefficients [dimensionless]. Shape is + (num_selected_envs, num_selected_joints). + env_ids: Input array of environment indices to write to. + Shape is (num_selected_envs,). + joint_ids: Input array of joint indices to write to. + Shape is (num_selected_joints,). + out_data: Output friction properties buffer. Shape is (num_envs, num_joints, 3). + """ + i, j = wp.tid() + val = in_data[i, j] + out_data[env_ids[i], joint_ids[j], 0] = val + out_data[env_ids[i], joint_ids[j], 1] = val + out_data[env_ids[i], joint_ids[j], 2] = val + + +@wp.kernel +def resolve_view_ids( + env_ids: wp.array(dtype=wp.int32), + body_ids: wp.array(dtype=wp.int32), + num_query_envs: wp.int32, + num_total_envs: wp.int32, + view_ids: wp.array(dtype=wp.int32), +) -> None: + """Resolve flat view indices from environment and body index pairs. + + Computes flat view indices from (env_id, body_id) pairs using body-major ordering: + ``view_id = body_id * num_total_envs + env_id``. The output array is laid out in + column-major order over the (env, body) grid. + + Args: + env_ids: Input environment indices. Shape is (num_query_envs,). + body_ids: Input body indices. Shape is (num_query_bodies,). + num_query_envs: Total number of queried environments. + num_total_envs: Total number of environments in the simulation. + view_ids: Output flat view indices. Shape is (num_query_bodies * num_query_envs,). + """ + i, j = wp.tid() + view_ids[j * num_query_envs + i] = body_ids[j] * num_total_envs + env_ids[i] + + +@wp.kernel +def write_joint_friction_to_buffer_mask( + in_data: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), + joint_mask: wp.array(dtype=wp.bool), + out_data: wp.array3d(dtype=wp.float32), +): + """Mask-scatter joint friction coefficient into all three components of the friction buffer. + + Broadcasts a single friction value into the static (index 0), dynamic (index 1), + and viscous (index 2) components where both ``env_mask[i]`` and ``joint_mask[j]`` + are True. + + Args: + in_data: Input friction coefficients [dimensionless]. Shape is + (num_envs, num_joints). + env_mask: Boolean environment mask. Shape is (num_envs,). + joint_mask: Boolean joint mask. Shape is (num_joints,). + out_data: Output friction properties buffer. Shape is (num_envs, num_joints, 3). + """ + i, j = wp.tid() + if env_mask[i] and joint_mask[j]: + val = in_data[i, j] + out_data[i, j, 0] = val + out_data[i, j, 1] = val + out_data[i, j, 2] = val diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/__init__.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/__init__.py new file mode 100644 index 000000000000..441515cd83a9 --- /dev/null +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/__init__.py @@ -0,0 +1,10 @@ +# 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 + +"""Sub-module for ovphysx-backed rigid object assets.""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/__init__.pyi b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/__init__.pyi new file mode 100644 index 000000000000..1c96a5aa4550 --- /dev/null +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/__init__.pyi @@ -0,0 +1,12 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +__all__ = [ + "RigidObject", + "RigidObjectData", +] + +from .rigid_object import RigidObject +from .rigid_object_data import RigidObjectData diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/rigid_object.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/rigid_object.py new file mode 100644 index 000000000000..015c8f102f44 --- /dev/null +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/rigid_object.py @@ -0,0 +1,1173 @@ +# 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 + +"""OVPhysX-backed RigidObject implementation.""" + +from __future__ import annotations + +import re +import warnings +from collections.abc import Sequence +from typing import Any + +import numpy as np +import torch +import warp as wp + +from pxr import UsdPhysics + +import isaaclab.sim as sim_utils +from isaaclab.assets.rigid_object.base_rigid_object import BaseRigidObject +from isaaclab.assets.rigid_object.rigid_object_cfg import RigidObjectCfg +from isaaclab.utils.string import resolve_matching_names +from isaaclab.utils.wrench_composer import WrenchComposer + +from isaaclab_ovphysx import tensor_types as TT +from isaaclab_ovphysx.assets import kernels as shared_kernels +from isaaclab_ovphysx.assets.kernels import _body_wrench_to_world +from isaaclab_ovphysx.physics import OvPhysxManager + +from .rigid_object_data import RigidObjectData + + +class RigidObject(BaseRigidObject): + """A rigid object asset class. + + Rigid objects are assets comprising of rigid bodies. They can be used to represent dynamic objects + such as boxes, spheres, etc. A rigid body is described by its pose, velocity and mass distribution. + + For an asset to be considered a rigid object, the root prim of the asset must have the `USD RigidBodyAPI`_ + applied to it. This API is used to define the simulation properties of the rigid body. On playing the + simulation, the physics engine will automatically register the rigid body and create a corresponding + rigid body handle. State is read and written through ovphysx ``TensorBinding`` objects acquired from + the :class:`~isaaclab_ovphysx.physics.OvPhysxManager`. Only free (non-articulated) rigid bodies are + supported; prims under an ``ArticulationRootAPI`` should use + :class:`~isaaclab_ovphysx.assets.articulation.Articulation` instead. + + .. _`USD RigidBodyAPI`: https://openusd.org/dev/api/class_usd_physics_rigid_body_a_p_i.html + """ + + cfg: RigidObjectCfg + """Configuration instance for the rigid object.""" + + __backend_name__: str = "ovphysx" + """The name of the backend for the rigid object.""" + + def __init__(self, cfg: RigidObjectCfg): + """Initialize the rigid object. + + Args: + cfg: A configuration instance. + """ + super().__init__(cfg) + # Bindings are created lazily (on first access) to avoid allocating + # handles for tensor types the user never queries. + self._bindings: dict[int, Any] = {} + + """ + Properties + """ + + @property + def data(self) -> RigidObjectData: + return self._data + + @property + def num_instances(self) -> int: + return self._num_instances + + @property + def num_bodies(self) -> int: + """Number of bodies in the asset. + + This is always 1 since each object is a single rigid body. + """ + return self._num_bodies + + @property + def body_names(self) -> list[str]: + """Ordered names of bodies in the rigid object.""" + return self._body_names + + @property + def root_view(self) -> dict[int, Any]: + """Root view for the asset. + + OVPhysX exposes per-tensor-type bindings rather than a single opaque view object + as used by the PhysX and Newton backends. Callers that need low-level binding + access should call :meth:`_get_binding` rather than iterating this dict directly. + For high-level state access (instance counts, prim paths, transforms), use the + :attr:`num_instances`, :attr:`body_names`, and + :attr:`~RigidObjectData.root_link_pose_w` accessors instead. + + .. note:: + Use this view with caution. It requires handling of tensors in a specific way. + """ + return self._bindings + + @property + def instantaneous_wrench_composer(self) -> WrenchComposer | None: + """Instantaneous wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are only valid for the current simulation step. At the end of the simulation step, the wrenches set + to this object are discarded. This is useful to apply forces that change all the time, things like drag forces + for instance. + """ + return self._instantaneous_wrench_composer + + @property + def permanent_wrench_composer(self) -> WrenchComposer | None: + """Permanent wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are persistent and are applied to the simulation at every step. This is useful to apply forces that + are constant over a period of time, things like the thrust of a motor for instance. + """ + return self._permanent_wrench_composer + + """ + Operations. + """ + + def reset( + self, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_mask: wp.array | None = None + ) -> None: + """Reset the rigid object. + + Args: + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + # resolve all indices + if (env_ids is None) or (env_ids == slice(None)): + env_ids = slice(None) + # reset external wrench + self._instantaneous_wrench_composer.reset(env_ids, env_mask) + self._permanent_wrench_composer.reset(env_ids, env_mask) + + def write_data_to_sim(self) -> None: + """Write external wrench to the simulation. + + .. note:: + We write external wrench to the simulation here since this function is called before the simulation step. + This ensures that the external wrench is applied at every simulation step. + """ + inst = self._instantaneous_wrench_composer + perm = self._permanent_wrench_composer + if not inst.active and not perm.active: + return + if inst.active: + if perm.active: + inst.add_raw_buffers_from(perm) + force_b = inst.out_force_b.warp + torque_b = inst.out_torque_b.warp + else: + force_b = perm.out_force_b.warp + torque_b = perm.out_torque_b.warp + + poses = self._data.body_link_pose_w.warp # (N, 1) wp.transformf + wp.launch( + _body_wrench_to_world, + dim=(self._num_instances, 1), + inputs=[force_b, torque_b, poses], + outputs=[self._wrench_buf], + device=self._device, + ) + binding = self._get_binding(TT.RIGID_BODY_WRENCH) + binding.write(self._wrench_buf_flat) + inst.reset() + + def update(self, dt: float) -> None: + """Updates the simulation data. + + Args: + dt: The time step size in seconds. + """ + self._data.update(dt) + + """ + Operations - Finders. + """ + + def find_bodies(self, name_keys: str | Sequence[str], preserve_order: bool = False) -> tuple[list[int], list[str]]: + """Find bodies in the rigid body based on the name keys. + + Please check the :meth:`isaaclab.utils.string.resolve_matching_names` function for more + information on the name matching. + + Args: + name_keys: A regular expression or a list of regular expressions to match the body names. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the body indices and names. + """ + return resolve_matching_names(name_keys, self._body_names, preserve_order) + + """ + Operations - Write to simulation. + """ + + def write_root_pose_to_sim_index( + self, + *, + root_pose: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7) + or (len(env_ids),) with dtype wp.transformf. + env_ids: Environment indices. If None, then all indices are used. + """ + self.write_root_link_pose_to_sim_index(root_pose=root_pose, env_ids=env_ids) + + def write_root_pose_to_sim_mask( + self, + *, + root_pose: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + ) -> None: + """Set the root pose over selected environment mask into the simulation. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_pose: Root poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + self.write_root_link_pose_to_sim_mask(root_pose=root_pose, env_mask=env_mask) + + def write_root_velocity_to_sim_index( + self, + *, + root_velocity: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root center of mass velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6) + or (len(env_ids),) with dtype wp.spatial_vectorf. + env_ids: Environment indices. If None, then all indices are used. + """ + self.write_root_com_velocity_to_sim_index(root_velocity=root_velocity, env_ids=env_ids) + + def write_root_velocity_to_sim_mask( + self, + *, + root_velocity: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + ) -> None: + """Set the root center of mass velocity over selected environment mask into the simulation. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + self.write_root_com_velocity_to_sim_mask(root_velocity=root_velocity, env_mask=env_mask) + + def write_root_link_pose_to_sim_index( + self, + *, + root_pose: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root link pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_pose: Root link poses in simulation frame. Shape is (len(env_ids), 7) + or (len(env_ids),) with dtype wp.transformf. + env_ids: Environment indices. If None, then all indices are used. + """ + env_ids = self._resolve_env_ids(env_ids) + self.assert_shape_and_dtype(root_pose, (env_ids.shape[0],), wp.transformf, "root_pose") + wp.launch( + shared_kernels.set_root_link_pose_to_sim_index, + dim=env_ids.shape[0], + inputs=[root_pose, env_ids], + outputs=[self.data.root_link_pose_w], + device=self._device, + ) + # Invalidate dependent root_com_pose timestamp so the next read recomposes it. + self.data._root_com_pose_w.timestamp = -1.0 + # Push cache to the wheel via an indexed write. + binding = self._get_binding(TT.RIGID_BODY_POSE) + binding.write(self.data._root_link_pose_w.data.view(wp.float32), indices=env_ids) + + def write_root_link_pose_to_sim_mask( + self, + *, + root_pose: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + ) -> None: + """Set the root link pose over selected environment mask into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_pose: Root poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + env_mask_wp = self._resolve_env_mask(env_mask) + self.assert_shape_and_dtype(root_pose, (self._num_instances,), wp.transformf, "root_pose") + wp.launch( + shared_kernels.set_root_link_pose_to_sim_mask, + dim=self._num_instances, + inputs=[root_pose, env_mask_wp], + outputs=[self.data.root_link_pose_w], + device=self._device, + ) + self.data._root_com_pose_w.timestamp = -1.0 + binding = self._get_binding(TT.RIGID_BODY_POSE) + binding.write(self.data._root_link_pose_w.data.view(wp.float32), mask=env_mask_wp) + + def write_root_com_pose_to_sim_index( + self, + *, + root_pose: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root center of mass pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + The orientation is the orientation of the principal axes of inertia. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7) + or (len(env_ids),) with dtype wp.transformf. + env_ids: Environment indices. If None, then all indices are used. + """ + env_ids = self._resolve_env_ids(env_ids) + self.assert_shape_and_dtype(root_pose, (env_ids.shape[0],), wp.transformf, "root_pose") + wp.launch( + shared_kernels.set_root_com_pose_to_sim_index, + dim=env_ids.shape[0], + inputs=[root_pose, self.data.body_com_pose_b, env_ids], + outputs=[self.data.root_com_pose_w, self.data.root_link_pose_w], + device=self._device, + ) + binding = self._get_binding(TT.RIGID_BODY_POSE) + binding.write(self.data._root_link_pose_w.data.view(wp.float32), indices=env_ids) + + def write_root_com_pose_to_sim_mask( + self, + *, + root_pose: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + ) -> None: + """Set the root center of mass pose over selected environment mask into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + The orientation is the orientation of the principal axes of inertia. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_pose: Root center of mass poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + env_mask_wp = self._resolve_env_mask(env_mask) + self.assert_shape_and_dtype(root_pose, (self._num_instances,), wp.transformf, "root_pose") + wp.launch( + shared_kernels.set_root_com_pose_to_sim_mask, + dim=self._num_instances, + inputs=[root_pose, self.data.body_com_pose_b, env_mask_wp], + outputs=[self.data.root_com_pose_w, self.data.root_link_pose_w], + device=self._device, + ) + binding = self._get_binding(TT.RIGID_BODY_POSE) + binding.write(self.data._root_link_pose_w.data.view(wp.float32), mask=env_mask_wp) + + def write_root_com_velocity_to_sim_index( + self, + *, + root_velocity: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root center of mass velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6) + or (len(env_ids),) with dtype wp.spatial_vectorf. + env_ids: Environment indices. If None, then all indices are used. + """ + env_ids = self._resolve_env_ids(env_ids) + self.assert_shape_and_dtype(root_velocity, (env_ids.shape[0],), wp.spatial_vectorf, "root_velocity") + wp.launch( + shared_kernels.set_root_com_velocity_to_sim_index, + dim=env_ids.shape[0], + inputs=[root_velocity, env_ids, 1], + outputs=[self.data.root_com_vel_w, self.data.body_com_acc_w], + device=self._device, + ) + # Invalidate dependent root_link_vel timestamp. + self.data._root_link_vel_w.timestamp = -1.0 + binding = self._get_binding(TT.RIGID_BODY_VELOCITY) + binding.write(self.data._root_com_vel_w.data.view(wp.float32), indices=env_ids) + + def write_root_com_velocity_to_sim_mask( + self, + *, + root_velocity: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + ) -> None: + """Set the root center of mass velocity over selected environment mask into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + env_mask_wp = self._resolve_env_mask(env_mask) + self.assert_shape_and_dtype(root_velocity, (self._num_instances,), wp.spatial_vectorf, "root_velocity") + wp.launch( + shared_kernels.set_root_com_velocity_to_sim_mask, + dim=self._num_instances, + inputs=[root_velocity, env_mask_wp, 1], + outputs=[self.data.root_com_vel_w, self.data.body_com_acc_w], + device=self._device, + ) + self.data._root_link_vel_w.timestamp = -1.0 + binding = self._get_binding(TT.RIGID_BODY_VELOCITY) + binding.write(self.data._root_com_vel_w.data.view(wp.float32), mask=env_mask_wp) + + def write_root_link_velocity_to_sim_index( + self, + *, + root_velocity: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root link velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the root's frame rather than the root's center of mass. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_velocity: Root frame velocities in simulation world frame. Shape is (len(env_ids), 6) + or (len(env_ids),) with dtype wp.spatial_vectorf. + env_ids: Environment indices. If None, then all indices are used. + """ + env_ids = self._resolve_env_ids(env_ids) + self.assert_shape_and_dtype(root_velocity, (env_ids.shape[0],), wp.spatial_vectorf, "root_velocity") + wp.launch( + shared_kernels.set_root_link_velocity_to_sim_index, + dim=env_ids.shape[0], + inputs=[ + root_velocity, + self.data.body_com_pose_b, + self.data.root_link_pose_w, + env_ids, + 1, # num_bodies is always 1 for RigidObject + ], + outputs=[self.data.root_link_vel_w, self.data.root_com_vel_w, self.data.body_com_acc_w], + device=self._device, + ) + binding = self._get_binding(TT.RIGID_BODY_VELOCITY) + binding.write(self.data._root_com_vel_w.data.view(wp.float32), indices=env_ids) + + def write_root_link_velocity_to_sim_mask( + self, + *, + root_velocity: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + ) -> None: + """Set the root link velocity over selected environment mask into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the root's frame rather than the root's center of mass. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_velocity: Root frame velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + env_mask_wp = self._resolve_env_mask(env_mask) + self.assert_shape_and_dtype(root_velocity, (self._num_instances,), wp.spatial_vectorf, "root_velocity") + wp.launch( + shared_kernels.set_root_link_velocity_to_sim_mask, + dim=self._num_instances, + inputs=[root_velocity, self.data.body_com_pose_b, self.data.root_link_pose_w, env_mask_wp, 1], + outputs=[self.data.root_link_vel_w, self.data.root_com_vel_w, self.data.body_com_acc_w], + device=self._device, + ) + binding = self._get_binding(TT.RIGID_BODY_VELOCITY) + binding.write(self.data._root_com_vel_w.data.view(wp.float32), mask=env_mask_wp) + + """ + Operations - Setters. + """ + + def set_masses_index( + self, + *, + masses: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set masses of all bodies using indices. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + masses: Masses of all bodies. Shape is (len(env_ids), len(body_ids)). + body_ids: The body indices to set the masses for. Defaults to None (all bodies). + env_ids: The environment indices to set the masses for. Defaults to None (all environments). + """ + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + self.assert_shape_and_dtype(masses, (env_ids.shape[0], body_ids.shape[0]), wp.float32, "masses") + # Scatter user data into the cached _body_mass at (env_ids, body_ids). + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[masses, env_ids, body_ids], + outputs=[self.data._body_mass], + device=self._device, + ) + # Push cache to the wheel via pinned-CPU staging (RIGID_BODY_MASS is CPU-only). + cpu_env_ids = self._get_cpu_env_ids(env_ids) + wp.copy(self._cpu_body_mass, self.data._body_mass) + binding = self._get_binding(TT.RIGID_BODY_MASS) + binding.write(self._cpu_body_mass.flatten(), indices=cpu_env_ids) + + def set_masses_mask( + self, + *, + masses: torch.Tensor | wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set masses of all bodies using masks. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + masses: Masses of all bodies. Shape is (num_instances, num_bodies). + body_mask: Body mask. If None, then all bodies are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + env_mask_wp = self._resolve_env_mask(env_mask) + body_mask_wp = self._resolve_body_mask(body_mask) + self.assert_shape_and_dtype(masses, (self._num_instances, self._num_bodies), wp.float32, "masses") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(self._num_instances, self._num_bodies), + inputs=[masses, env_mask_wp, body_mask_wp], + outputs=[self.data._body_mass], + device=self._device, + ) + wp.copy(self._cpu_body_mass, self.data._body_mass) + binding = self._get_binding(TT.RIGID_BODY_MASS) + binding.write(self._cpu_body_mass.flatten(), mask=self._get_cpu_env_mask(env_mask_wp)) + + def set_coms_index( + self, + *, + coms: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set center of mass pose of all bodies using indices. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + coms: Center of mass pose of all bodies. Shape is (len(env_ids), len(body_ids), 7). + body_ids: The body indices to set the center of mass pose for. Defaults to None (all bodies). + env_ids: The environment indices to set the center of mass pose for. Defaults to None + (all environments). + """ + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + self.assert_shape_and_dtype(coms, (env_ids.shape[0], body_ids.shape[0]), wp.transformf, "coms") + wp.launch( + shared_kernels.write_body_com_pose_to_buffer_index, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[coms, env_ids, body_ids], + outputs=[self.data._body_com_pose_b.data], + device=self._device, + ) + # Invalidate dependent root_com_pose timestamp -- it's derived from body_com_pose_b. + self.data._root_com_pose_w.timestamp = -1.0 + # Push cache to the wheel via pinned-CPU staging (RIGID_BODY_COM_POSE is CPU-only). + cpu_env_ids = self._get_cpu_env_ids(env_ids) + wp.copy(self._cpu_body_coms, self.data._body_com_pose_b.data) + binding = self._get_binding(TT.RIGID_BODY_COM_POSE) + # Wheel binding shape is (N, 7); squeeze singleton body dim with a flat float32 view. + binding.write(self._cpu_body_coms.reshape((self._num_instances, 7)), indices=cpu_env_ids) + + def set_coms_mask( + self, + *, + coms: torch.Tensor | wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set center of mass pose of all bodies using masks. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + coms: Center of mass pose of all bodies. Shape is (num_instances, num_bodies, 7). + body_mask: Body mask. If None, then all bodies are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + env_mask_wp = self._resolve_env_mask(env_mask) + body_mask_wp = self._resolve_body_mask(body_mask) + self.assert_shape_and_dtype(coms, (self._num_instances, self._num_bodies), wp.transformf, "coms") + wp.launch( + shared_kernels.write_body_com_pose_to_buffer_mask, + dim=(self._num_instances, self._num_bodies), + inputs=[coms, env_mask_wp, body_mask_wp], + outputs=[self.data._body_com_pose_b.data], + device=self._device, + ) + self.data._root_com_pose_w.timestamp = -1.0 + wp.copy(self._cpu_body_coms, self.data._body_com_pose_b.data) + binding = self._get_binding(TT.RIGID_BODY_COM_POSE) + binding.write(self._cpu_body_coms.reshape((self._num_instances, 7)), mask=self._get_cpu_env_mask(env_mask_wp)) + + def set_inertias_index( + self, + *, + inertias: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set inertias of all bodies using indices. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + inertias: Inertias of all bodies. Shape is (len(env_ids), len(body_ids), 9). + body_ids: The body indices to set the inertias for. Defaults to None (all bodies). + env_ids: The environment indices to set the inertias for. Defaults to None (all environments). + """ + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + self.assert_shape_and_dtype(inertias, (env_ids.shape[0], body_ids.shape[0], 9), wp.float32, "inertias") + wp.launch( + shared_kernels.write_body_inertia_to_buffer_index, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[inertias, env_ids, body_ids], + outputs=[self.data._body_inertia], + device=self._device, + ) + # Push cache to the wheel via pinned-CPU staging (RIGID_BODY_INERTIA is CPU-only). + cpu_env_ids = self._get_cpu_env_ids(env_ids) + wp.copy(self._cpu_body_inertia, self.data._body_inertia) + binding = self._get_binding(TT.RIGID_BODY_INERTIA) + # Wheel binding shape is (N, 9); flatten the singleton body dim. + binding.write(self._cpu_body_inertia.reshape((self._num_instances, 9)), indices=cpu_env_ids) + + def set_inertias_mask( + self, + *, + inertias: torch.Tensor | wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set inertias of all bodies using masks. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + inertias: Inertias of all bodies. Shape is (num_instances, num_bodies, 9). + body_mask: Body mask. If None, then all bodies are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + env_mask_wp = self._resolve_env_mask(env_mask) + body_mask_wp = self._resolve_body_mask(body_mask) + self.assert_shape_and_dtype(inertias, (self._num_instances, self._num_bodies, 9), wp.float32, "inertias") + wp.launch( + shared_kernels.write_body_inertia_to_buffer_mask, + dim=(self._num_instances, self._num_bodies), + inputs=[inertias, env_mask_wp, body_mask_wp], + outputs=[self.data._body_inertia], + device=self._device, + ) + wp.copy(self._cpu_body_inertia, self.data._body_inertia) + binding = self._get_binding(TT.RIGID_BODY_INERTIA) + binding.write( + self._cpu_body_inertia.reshape((self._num_instances, 9)), mask=self._get_cpu_env_mask(env_mask_wp) + ) + + """ + Internal helper. + """ + + def _initialize_impl(self) -> None: + # acquire ovphysx instance + physx_instance = OvPhysxManager.get_physx_instance() + if physx_instance is None: + raise RuntimeError("OvPhysxManager has not been initialized yet.") + self._ovphysx = physx_instance + # Derive the device from PhysicsManager (which mirrors SimulationContext.cfg.device). + # The ovphysx PhysX object does not expose a .device property; reading it would + # raise AttributeError (masked by hasattr) and fall back to "cuda:0" even when the + # simulation is running on CPU, causing a device mismatch in binding.read(). + self._device = OvPhysxManager.get_device() + # Convert IsaacLab prim-path notation to the glob patterns ovphysx expects. + # IsaacLab uses two conventions: + # /World/envs/env_.*/object -- regex dot-star for "any env index" + # /World/envs/{ENV_REGEX_NS}/object -- explicit placeholder + # ovphysx ``create_tensor_binding`` uses fnmatch-style globs, so both map to ``*``. + pattern = re.sub(r"\{ENV_REGEX_NS\}", "*", self.cfg.prim_path) + pattern = re.sub(r"\.\*", "*", pattern) + self._binding_pattern = pattern + + # Validate the prim tree before creating tensor bindings -- the wheel silently + # produces a 0-prim binding when the pattern matches nothing, which surfaces as an + # obscure ``TypeError`` deep in property accessors. + # obtain the first prim in the regex expression (all others are assumed to be a copy of this) + template_prim = sim_utils.find_first_matching_prim(self.cfg.prim_path) + if template_prim is None: + raise RuntimeError(f"Failed to find prim for expression: '{self.cfg.prim_path}'.") + template_prim_path = template_prim.GetPath().pathString + + # find rigid root prims + root_prims = sim_utils.get_all_matching_child_prims( + template_prim_path, + predicate=lambda prim: prim.HasAPI(UsdPhysics.RigidBodyAPI), + traverse_instance_prims=False, + ) + if len(root_prims) == 0: + raise RuntimeError( + f"Failed to find a rigid body when resolving '{self.cfg.prim_path}'." + " Please ensure that the prim has 'USD RigidBodyAPI' applied." + ) + if len(root_prims) > 1: + raise RuntimeError( + f"Failed to find a single rigid body when resolving '{self.cfg.prim_path}'." + f" Found multiple '{root_prims}' under '{template_prim_path}'." + " Please ensure that there is only one rigid body in the prim path tree." + ) + articulation_prims = sim_utils.get_all_matching_child_prims( + template_prim_path, + predicate=lambda prim: prim.HasAPI(UsdPhysics.ArticulationRootAPI), + traverse_instance_prims=False, + ) + if len(articulation_prims) != 0: + if articulation_prims[0].GetAttribute("physxArticulation:articulationEnabled").Get(): + raise RuntimeError( + f"Found an articulation root when resolving '{self.cfg.prim_path}' for rigid" + f" objects. These are located at: '{articulation_prims}' under" + f" '{template_prim_path}'. Please disable the articulation root in the USD" + " or from code by setting the parameter" + " 'ArticulationRootPropertiesCfg.articulation_enabled' to False in the spawn" + " configuration." + ) + + # Eagerly create every binding the data container reads at init, so failures + # surface here with a helpful message rather than as a raw wheel exception + # (or a KeyError) at first writer call. + for tt in ( + TT.RIGID_BODY_POSE, + TT.RIGID_BODY_VELOCITY, + TT.RIGID_BODY_WRENCH, + TT.RIGID_BODY_MASS, + TT.RIGID_BODY_COM_POSE, + TT.RIGID_BODY_INERTIA, + ): + try: + self._get_binding(tt) + except Exception as e: + raise RuntimeError( + f"OVPhysX could not create rigid-body binding {tt!r}. " + f"Check that prim_path={self._binding_pattern!r} matches " + f"at least one UsdPhysics.RigidBodyAPI prim and that the " + f"ovphysx wheel exposes the RIGID_BODY_* TensorType. " + f"Note: pattern resolution may currently include articulation " + f"links; an explicit selection policy is on the wheel-side roadmap." + ) from e + + # read counts and body names from the root-pose binding + root_pose = self._bindings[TT.RIGID_BODY_POSE] + self._num_instances = root_pose.count + self._num_bodies = 1 + try: + body_names_value = root_pose.body_names + # body_names may be an empty list for non-articulation bindings; fall back to + # the documented single-body default in that case. + self._body_names = list(body_names_value) if body_names_value else ["base_link"] + except (AttributeError, TypeError): + # ovphysx TensorBinding raises TypeError (not AttributeError) when body_names + # is queried on a non-articulation tensor type such as RIGID_BODY_POSE: + # "Articulation metadata … is not available for tensor type 'RIGID_BODY_POSE'." + # For a single-body rigid object the default ["base_link"] is always correct. + self._body_names = ["base_link"] + + # container for data access + self._data = RigidObjectData(self._bindings, self._device, check_shapes=self._check_shapes) + + # create buffers + self._create_buffers() + # process configuration + self._process_cfg() + # update the rigid body data + self.update(0.0) + # Let the rigid object data know that it is fully instantiated and ready to use. + self._data.is_primed = True + + def _create_buffers(self) -> None: + """Create buffers for storing data.""" + N = self._num_instances + B = 1 # rigid object always has a single body + device = self._device + + # constants + self._ALL_INDICES = wp.array(np.arange(N, dtype=np.int32), device=device) + self._ALL_BODY_INDICES = wp.array(np.arange(B, dtype=np.int32), device=device) + # All-true masks for default mask paths. These let ``binding.write(..., mask=...)`` + # cover all instances when no env_mask is supplied, without converting back to indices. + self._ALL_TRUE_ENV_MASK = wp.array(np.ones(N, dtype=bool), dtype=wp.bool, device=device) + self._ALL_TRUE_BODY_MASK = wp.array(np.ones(B, dtype=bool), dtype=wp.bool, device=device) + + # external wrench composer + # The kernel writes into the (N, 1, 9) view; the binding consumes the (N, 9) view -- + # both alias the same allocation, so we cache the flat reshape once. + self._wrench_buf = wp.zeros((N, 1, 9), dtype=wp.float32, device=device) + self._wrench_buf_flat = wp.array( + ptr=self._wrench_buf.ptr, + shape=(N, 9), + dtype=wp.float32, + device=device, + copy=False, + ) + self._instantaneous_wrench_composer = WrenchComposer(self) + self._permanent_wrench_composer = WrenchComposer(self) + + # set information about rigid body into data + self._data.body_names = self._body_names + + # Pre-allocated pinned CPU buffers for OVPhysX TensorBinding writes. The wheel + # requires CPU arrays for "model" property updates (mass / coms / inertia); pinned + # host memory enables DMA fast path and avoids per-call ``wp.clone`` allocation. + self._cpu_env_ids_all = wp.zeros(N, dtype=wp.int32, device="cpu", pinned=True) + wp.copy(self._cpu_env_ids_all, self._ALL_INDICES) + self._cpu_body_mass = wp.zeros((N, B), dtype=wp.float32, device="cpu", pinned=True) + self._cpu_body_coms = wp.zeros((N, B, 7), dtype=wp.float32, device="cpu", pinned=True) + self._cpu_body_inertia = wp.zeros((N, B, 9), dtype=wp.float32, device="cpu", pinned=True) + # Pinned-host mask staging for CPU-only binding writes (mass/coms/inertia). + self._cpu_env_mask = wp.zeros(N, dtype=wp.bool, device="cpu", pinned=True) + + def _process_cfg(self) -> None: + """Post processing of configuration parameters.""" + # default state + # -- root state + # note: we cast to tuple to avoid torch/numpy type mismatch. + default_root_pose = tuple(self.cfg.init_state.pos) + tuple(self.cfg.init_state.rot) + default_root_vel = tuple(self.cfg.init_state.lin_vel) + tuple(self.cfg.init_state.ang_vel) + default_root_pose = np.tile(np.array(default_root_pose, dtype=np.float32), (self._num_instances, 1)) + default_root_vel = np.tile(np.array(default_root_vel, dtype=np.float32), (self._num_instances, 1)) + 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_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array: + """Resolve environment indices to a warp array. + + Args: + env_ids: Environment indices. If None, then all indices are used. + + Returns: + A warp array of environment indices on ``self._device``. + """ + if env_ids is None or env_ids == slice(None): + return self._ALL_INDICES + if isinstance(env_ids, list): + return wp.array(env_ids, dtype=wp.int32, device=self._device) + if isinstance(env_ids, torch.Tensor): + return wp.from_torch(env_ids.to(torch.int32), dtype=wp.int32) + 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 + + def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array: + """Resolve body indices to a warp array. + + Args: + body_ids: Body indices. If None, then all indices are used. + + Returns: + A warp array of body indices on ``self._device``. + """ + 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 + + def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array: + """Resolve an environment mask to a ``wp.bool`` array. + + Args: + env_mask: Environment mask. If None, then the pre-allocated all-true mask is used. + + 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 + + def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array: + """Resolve a body mask to a ``wp.bool`` array. + + Args: + body_mask: Body mask. If None, then the pre-allocated all-true mask is used. + + 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 + + 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. + + The wheel's ``binding.write(mask=...)`` requires the mask on the binding's + device, which is CPU for mass / coms / inertia. Reuses the pre-allocated + ``_cpu_env_mask`` pinned buffer. + """ + wp.copy(self._cpu_env_mask, env_mask) + return self._cpu_env_mask + + def _get_cpu_env_ids(self, env_ids: wp.array | torch.Tensor) -> wp.array: + """Return CPU int32 indices, using the pre-allocated pinned ``_cpu_env_ids_all`` + fast path when *env_ids* matches ``_ALL_INDICES``. + """ + if isinstance(env_ids, torch.Tensor): + env_ids = wp.from_torch(env_ids, dtype=wp.int32) + if env_ids.ptr == self._ALL_INDICES.ptr: + return self._cpu_env_ids_all + return wp.clone(env_ids, device="cpu") + + def _get_binding(self, tensor_type: int): + """Return a cached TensorBinding, creating it on first access. + + Bindings are lightweight handles (a pointer + shape metadata into PhysX's + shared GPU buffer). Creating one does NOT allocate new GPU memory -- the + underlying simulation buffers are allocated once by PhysX regardless of how + many bindings point into them. Still, we defer creation so that tensor types + the user never queries are never looked up. + + Args: + tensor_type: The TensorType constant identifying which simulation buffer + to bind (e.g. :attr:`~isaaclab_ovphysx.tensor_types.RIGID_BODY_POSE`). + + Returns: + The cached TensorBinding for ``tensor_type``. + """ + binding = self._bindings.get(tensor_type) + if binding is not None: + return binding + binding = self._ovphysx.create_tensor_binding(pattern=self._binding_pattern, tensor_type=tensor_type) + self._bindings[tensor_type] = binding + return binding + + """ + Internal simulation callbacks. + """ + + def _invalidate_initialize_callback(self, event) -> None: + """Invalidates the scene elements.""" + super()._invalidate_initialize_callback(event) + + def write_root_state_to_sim( + self, + root_state: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Deprecated, same as :meth:`write_root_link_pose_to_sim_index` and + :meth:`write_root_com_velocity_to_sim_index`.""" + warnings.warn( + "The function 'write_root_state_to_sim' will be deprecated in a future release. Please" + " use 'write_root_link_pose_to_sim_index' and 'write_root_com_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + self.write_root_link_pose_to_sim_index(root_pose=root_state[:, :7], env_ids=env_ids) + self.write_root_com_velocity_to_sim_index(root_velocity=root_state[:, 7:], env_ids=env_ids) + + def write_root_com_state_to_sim( + self, + root_state: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Deprecated, same as :meth:`write_root_com_pose_to_sim_index` and + :meth:`write_root_com_velocity_to_sim_index`.""" + warnings.warn( + "The function 'write_root_com_state_to_sim' will be deprecated in a future release. Please" + " use 'write_root_com_pose_to_sim_index' and 'write_root_com_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + self.write_root_com_pose_to_sim_index(root_pose=root_state[:, :7], env_ids=env_ids) + self.write_root_com_velocity_to_sim_index(root_velocity=root_state[:, 7:], env_ids=env_ids) + + def write_root_link_state_to_sim( + self, + root_state: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Deprecated, same as :meth:`write_root_link_pose_to_sim_index` and + :meth:`write_root_link_velocity_to_sim_index`.""" + warnings.warn( + "The function 'write_root_link_state_to_sim' will be deprecated in a future release. Please" + " use 'write_root_link_pose_to_sim_index' and 'write_root_link_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + self.write_root_link_pose_to_sim_index(root_pose=root_state[:, :7], env_ids=env_ids) + self.write_root_link_velocity_to_sim_index(root_velocity=root_state[:, 7:], env_ids=env_ids) diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/rigid_object_data.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/rigid_object_data.py new file mode 100644 index 000000000000..05e5c45a0ebb --- /dev/null +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/rigid_object_data.py @@ -0,0 +1,1198 @@ +# 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 + +"""OVPhysX-backed RigidObjectData implementation.""" + +from __future__ import annotations + +import math +import warnings + +import torch +import warp as wp + +from isaaclab.assets.rigid_object.base_rigid_object_data import BaseRigidObjectData +from isaaclab.utils.buffers import TimestampedBufferWarp as TimestampedBuffer +from isaaclab.utils.math import normalize +from isaaclab.utils.warp import ProxyArray + +from isaaclab_ovphysx import tensor_types as TT +from isaaclab_ovphysx.assets import kernels as shared_kernels +from isaaclab_ovphysx.physics import OvPhysxManager as SimulationManager + + +class RigidObjectData(BaseRigidObjectData): + """Data container for a rigid object. + + This class contains the data for a rigid object in the simulation. The data includes the state of + the root rigid body and the state of all the bodies in the object. The data is stored in the simulation + world frame unless otherwise specified. + + For a rigid body, there are two frames of reference that are used: + + - Actor frame: The frame of reference of the rigid body prim. This typically corresponds to the Xform prim + with the rigid body schema. + - Center of mass frame: The frame of reference of the center of mass of the rigid body. + + Depending on the settings of the simulation, the actor frame and the center of mass frame may be the same. + This needs to be taken into account when interpreting the data. + + The data is lazily updated, meaning that the data is only updated when it is accessed. This is useful + when the data is expensive to compute or retrieve. The data is updated when the timestamp of the buffer + is older than the current simulation timestamp. The timestamp is updated whenever the data is updated. + + .. note:: + **Pull-to-refresh model.** Properties pull fresh data from the PhysX tensor API on first + access per timestamp and cache the result. This differs from Newton, where buffers are + refreshed automatically by the simulation. + + .. note:: + **ProxyArray pointer stability.** Each :class:`ProxyArray` wrapper is created once and + reused because the PhysX tensor API returns views into stable, pre-allocated GPU buffers + whose device pointer does not change across simulation steps. + """ + + __backend_name__: str = "ovphysx" + """The name of the backend for the rigid object data.""" + + def __init__( + self, + bindings: dict, + device: str, + check_shapes: bool = True, + ): + """Initializes the rigid object data. + + Args: + bindings: The OVPhysX tensor bindings dict keyed by tensor-type constant. + ``num_instances`` is read from ``bindings[RIGID_BODY_POSE].count`` and + ``num_bodies`` is fixed at 1; ``body_names`` is set by + :meth:`~isaaclab_ovphysx.assets.RigidObject._initialize_impl`. + device: The device used for processing. + check_shapes: Whether to enforce internal shape/dtype invariants on + lazy reads. Defaults to ``True``; production callers thread this + from :attr:`~isaaclab.assets.AssetBaseCfg.disable_shape_checks`. + """ + super().__init__(bindings, device) + # Set the tensor bindings (OVPhysX exposes per-tensor-type bindings rather than a single view). + self._bindings = bindings + self._check_shapes = check_shapes + # Set initial time stamp + self._sim_timestamp = 0.0 + self._is_primed = False + root_pose = self._bindings[TT.RIGID_BODY_POSE] + self._num_instances = root_pose.count + self._num_bodies = 1 + + if SimulationManager._sim is not None and hasattr(SimulationManager._sim, "cfg"): + gravity = SimulationManager._sim.cfg.gravity + else: + gravity = (0.0, 0.0, -9.81) + + gravity_dir = torch.tensor((gravity[0], gravity[1], gravity[2]), device=self.device) + # When gravity is disabled (cfg.gravity == (0, 0, 0)), normalize() would NaN. + if torch.linalg.norm(gravity_dir) > 0.0: + gravity_dir = normalize(gravity_dir.unsqueeze(0)).squeeze(0) + gravity_dir = gravity_dir.repeat(self._num_instances, 1) + forward_vec = torch.tensor((1.0, 0.0, 0.0), device=self.device).repeat(self._num_instances, 1) + + # Initialize constants + self.GRAVITY_VEC_W = ProxyArray(wp.from_torch(gravity_dir, dtype=wp.vec3f)) + self.FORWARD_VEC_B = ProxyArray(wp.from_torch(forward_vec, dtype=wp.vec3f)) + + self._create_buffers() + + @property + def is_primed(self) -> bool: + """Whether the rigid object data is fully instantiated and ready to use.""" + return self._is_primed + + @is_primed.setter + def is_primed(self, value: bool) -> None: + """Set whether the rigid object data is fully instantiated and ready to use. + + .. note:: + Once this quantity is set to True, it cannot be changed. + + Args: + value: The primed state. + + Raises: + ValueError: If the rigid object data is already primed. + """ + if self._is_primed: + raise ValueError("The rigid object data is already primed.") + self._is_primed = value + + def update(self, dt: float) -> None: + """Updates the data for the rigid object. + + Args: + dt: The time step for the update [s]. This must be a positive value. + """ + # update the simulation timestamp + self._sim_timestamp += dt + # Trigger an update of the body com acceleration buffer at a higher frequency + # since we do finite differencing. + self.body_com_acc_w + + """ + Names. + """ + + body_names: list[str] = None + """Body names in the order parsed by the simulation view.""" + + """ + Defaults. + """ + + @property + def default_root_pose(self) -> ProxyArray: + """Default root pose ``[pos, quat]`` in simulation world frame [m, -]. + Shape is (num_instances,), dtype = wp.transformf. + In torch this resolves to (num_instances, 7). + + Populated from :attr:`RigidObjectCfg.init_state` during initialisation. + """ + if self._default_root_pose_ta is None: + self._default_root_pose_ta = ProxyArray(self._default_root_pose) + return self._default_root_pose_ta + + @default_root_pose.setter + def default_root_pose(self, value: wp.array) -> None: + """Set the default root pose. + + Args: + value: The default root pose. Shape is (num_instances, 7). + + Raises: + ValueError: If the rigid object data is already primed. + """ + if self._is_primed: + raise ValueError("The rigid object data is already primed.") + self._default_root_pose.assign(value) + + @property + def default_root_vel(self) -> ProxyArray: + """Default root velocity ``[lin_vel, ang_vel]`` in simulation world frame [m/s, rad/s]. + Shape is (num_instances,), dtype = wp.spatial_vectorf. + In torch this resolves to (num_instances, 6). + + Populated from :attr:`RigidObjectCfg.init_state` during initialisation. + """ + if self._default_root_vel_ta is None: + self._default_root_vel_ta = ProxyArray(self._default_root_vel) + return self._default_root_vel_ta + + @default_root_vel.setter + def default_root_vel(self, value: wp.array) -> None: + """Set the default root velocity. + + Args: + value: The default root velocity. Shape is (num_instances, 6). + + Raises: + ValueError: If the rigid object data is already primed. + """ + if self._is_primed: + raise ValueError("The rigid object data is already primed.") + self._default_root_vel.assign(value) + + """ + Root state properties. + """ + + @property + def root_link_pose_w(self) -> ProxyArray: + """Root link pose ``[pos, quat]`` in simulation world frame [m, -]. + + Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). + This quantity is the pose of the actor frame of the root rigid body relative to the world. + The orientation is provided in (x, y, z, w) format. + """ + if self._root_link_pose_w.timestamp < self._sim_timestamp: + # read data from simulation + self._read_binding_into(TT.RIGID_BODY_POSE, self._root_link_pose_w.data) + self._root_link_pose_w.timestamp = self._sim_timestamp + if self._root_link_pose_w_ta is None: + self._root_link_pose_w_ta = ProxyArray(self._root_link_pose_w.data) + return self._root_link_pose_w_ta + + @property + def root_link_vel_w(self) -> ProxyArray: + """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame [m/s, rad/s]. + + Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). + This quantity contains the linear and angular velocities of the actor frame of the root + rigid body relative to the world. + """ + if self._root_link_vel_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.get_root_link_vel_from_root_com_vel, + dim=self._num_instances, + inputs=[ + self.root_com_vel_w, + self.root_link_pose_w, + self.body_com_pose_b, + ], + outputs=[self._root_link_vel_w.data], + device=self.device, + ) + self._root_link_vel_w.timestamp = self._sim_timestamp + if self._root_link_vel_w_ta is None: + self._root_link_vel_w_ta = ProxyArray(self._root_link_vel_w.data) + return self._root_link_vel_w_ta + + @property + def root_com_pose_w(self) -> ProxyArray: + """Root center of mass pose ``[pos, quat]`` in simulation world frame [m, -]. + + Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). + This quantity is the pose of the center of mass frame of the root rigid body relative to the world. + The orientation is provided in (x, y, z, w) format. + """ + if self._root_com_pose_w.timestamp < self._sim_timestamp: + # apply local transform to center of mass frame + wp.launch( + shared_kernels.get_root_com_pose_from_root_link_pose, + dim=self._num_instances, + inputs=[ + self.root_link_pose_w, + self.body_com_pose_b, + ], + outputs=[ + self._root_com_pose_w.data, + ], + device=self.device, + ) + self._root_com_pose_w.timestamp = self._sim_timestamp + + if self._root_com_pose_w_ta is None: + self._root_com_pose_w_ta = ProxyArray(self._root_com_pose_w.data) + return self._root_com_pose_w_ta + + @property + def root_com_vel_w(self) -> ProxyArray: + """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame [m/s, rad/s]. + + Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). + This quantity contains the linear and angular velocities of the root rigid body's center of mass frame + relative to the world. + """ + if self._root_com_vel_w.timestamp < self._sim_timestamp: + self._read_binding_into(TT.RIGID_BODY_VELOCITY, self._root_com_vel_w.data) + self._root_com_vel_w.timestamp = self._sim_timestamp + if self._root_com_vel_w_ta is None: + self._root_com_vel_w_ta = ProxyArray(self._root_com_vel_w.data) + return self._root_com_vel_w_ta + + """ + Body state properties. + """ + + @property + def body_mass(self) -> ProxyArray: + """Mass of all bodies [kg]. + + Shape is (num_instances, 1), dtype = wp.float32. + In torch this resolves to (num_instances, 1). + """ + if self._body_mass_ta is None: + self._body_mass_ta = ProxyArray(self._body_mass) + return self._body_mass_ta + + @property + def body_inertia(self) -> ProxyArray: + """Inertia tensor of all bodies, expressed at the center of mass [kg·m²]. + + Shape is (num_instances, 1, 9), dtype = wp.float32. The 9 components are the row-major + flatten of the 3×3 inertia matrix ``(Ixx, Ixy, Ixz, Iyx, Iyy, Iyz, Izx, Izy, Izz)``. + In torch this resolves to (num_instances, 1, 9). + """ + if self._body_inertia_ta is None: + self._body_inertia_ta = ProxyArray(self._body_inertia) + return self._body_inertia_ta + + @property + def body_link_pose_w(self) -> ProxyArray: + """Body link pose ``[pos, quat]`` in simulation world frame [m, -]. + + Shape is (num_instances, 1), dtype = wp.transformf. In torch this resolves to (num_instances, 1, 7). + This quantity is the pose of the actor frame of the rigid body relative to the world. + The orientation is provided in (x, y, z, w) format. + """ + parent = self.root_link_pose_w + if self._body_link_pose_w_ta is None: + self._body_link_pose_w_ta = ProxyArray(parent.warp.reshape((self._num_instances, 1))) + return self._body_link_pose_w_ta + + @property + def body_link_vel_w(self) -> ProxyArray: + """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame [m/s, rad/s]. + + Shape is (num_instances, 1), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 1, 6). + This quantity contains the linear and angular velocities of the body's link (actor) frame + relative to the world. + """ + parent = self.root_link_vel_w + if self._body_link_vel_w_ta is None: + self._body_link_vel_w_ta = ProxyArray(parent.warp.reshape((self._num_instances, 1))) + return self._body_link_vel_w_ta + + @property + def body_com_pose_w(self) -> ProxyArray: + """Body center of mass pose ``[pos, quat]`` in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.transformf. In torch this resolves to (num_instances, 1, 7). + This quantity is the pose of the center of mass frame of the rigid body relative to the world. + The orientation is provided in (x, y, z, w) format. + """ + parent = self.root_com_pose_w + if self._body_com_pose_w_ta is None: + self._body_com_pose_w_ta = ProxyArray(parent.warp.reshape((self._num_instances, 1))) + return self._body_com_pose_w_ta + + @property + def body_com_vel_w(self) -> ProxyArray: + """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame [m/s, rad/s]. + + Shape is (num_instances, 1), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 1, 6). + This quantity contains the linear and angular velocities of the body's center of mass frame + relative to the world. + """ + parent = self.root_com_vel_w + if self._body_com_vel_w_ta is None: + self._body_com_vel_w_ta = ProxyArray(parent.warp.reshape((self._num_instances, 1))) + return self._body_com_vel_w_ta + + @property + def body_com_acc_w(self) -> ProxyArray: + """Acceleration of all bodies ``[lin_acc, ang_acc]`` in the simulation world frame [m/s², rad/s²]. + + Shape is (num_instances, 1), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 1, 6). + This quantity is the acceleration of the rigid bodies' center of mass frame relative to the world. + """ + if self._body_com_acc_w.timestamp < self._sim_timestamp: + if self._previous_body_com_vel is None: + self._previous_body_com_vel = wp.clone(self.body_com_vel_w.warp) + wp.launch( + shared_kernels.derive_body_acceleration_from_body_com_velocities, + dim=(self._num_instances, 1), + device=self.device, + inputs=[ + self.body_com_vel_w.warp, + SimulationManager.get_physics_dt(), + self._previous_body_com_vel, + ], + outputs=[ + self._body_com_acc_w.data, + ], + ) + self._body_com_acc_w.timestamp = self._sim_timestamp + if self._body_com_acc_w_ta is None: + self._body_com_acc_w_ta = ProxyArray(self._body_com_acc_w.data) + return self._body_com_acc_w_ta + + @property + def body_com_pose_b(self) -> ProxyArray: + """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. + + Shape is (num_instances, 1), dtype = wp.transformf. In torch this resolves to (num_instances, 1, 7). + This quantity is the pose of the center of mass frame of the rigid body relative to the body's link frame. + The orientation is provided in (x, y, z, w) format. + """ + if self._body_com_pose_b.timestamp < self._sim_timestamp: + # read data from simulation + self._read_binding_into(TT.RIGID_BODY_COM_POSE, self._body_com_pose_b.data) + self._body_com_pose_b.timestamp = self._sim_timestamp + + if self._body_com_pose_b_ta is None: + self._body_com_pose_b_ta = ProxyArray(self._body_com_pose_b.data) + return self._body_com_pose_b_ta + + """ + Derived Properties. + """ + + @property + def projected_gravity_b(self) -> ProxyArray: + """Projection of the gravity direction on base frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + """ + if self._projected_gravity_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_1D_kernel, + dim=self._num_instances, + inputs=[self.GRAVITY_VEC_W, self.root_link_quat_w], + outputs=[self._projected_gravity_b.data], + device=self.device, + ) + self._projected_gravity_b.timestamp = self._sim_timestamp + if self._projected_gravity_b_ta is None: + self._projected_gravity_b_ta = ProxyArray(self._projected_gravity_b.data) + return self._projected_gravity_b_ta + + @property + def heading_w(self) -> ProxyArray: + """Yaw heading of the base frame (in radians). + + Shape is (num_instances,), dtype = wp.float32. In torch this resolves to (num_instances,). + + .. note:: + This quantity is computed by assuming that the forward-direction of the base + frame is along x-direction, i.e. :math:`(1, 0, 0)`. + """ + if self._heading_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.root_heading_w, + dim=self._num_instances, + inputs=[self.FORWARD_VEC_B, self.root_link_quat_w], + outputs=[self._heading_w.data], + device=self.device, + ) + self._heading_w.timestamp = self._sim_timestamp + if self._heading_w_ta is None: + self._heading_w_ta = ProxyArray(self._heading_w.data) + return self._heading_w_ta + + @property + def root_link_lin_vel_b(self) -> ProxyArray: + """Root link linear velocity in base frame [m/s]. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the linear velocity of the root link frame relative to the world, + expressed in the root link's actor frame. + """ + if self._root_link_lin_vel_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_1D_kernel, + dim=self._num_instances, + inputs=[self.root_link_lin_vel_w, self.root_link_quat_w], + outputs=[self._root_link_lin_vel_b.data], + device=self.device, + ) + self._root_link_lin_vel_b.timestamp = self._sim_timestamp + if self._root_link_lin_vel_b_ta is None: + self._root_link_lin_vel_b_ta = ProxyArray(self._root_link_lin_vel_b.data) + return self._root_link_lin_vel_b_ta + + @property + def root_link_ang_vel_b(self) -> ProxyArray: + """Root link angular velocity in base frame [rad/s]. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the angular velocity of the root link frame relative to the world, + expressed in the root link's actor frame. + """ + if self._root_link_ang_vel_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_1D_kernel, + dim=self._num_instances, + inputs=[self.root_link_ang_vel_w, self.root_link_quat_w], + outputs=[self._root_link_ang_vel_b.data], + device=self.device, + ) + self._root_link_ang_vel_b.timestamp = self._sim_timestamp + if self._root_link_ang_vel_b_ta is None: + self._root_link_ang_vel_b_ta = ProxyArray(self._root_link_ang_vel_b.data) + return self._root_link_ang_vel_b_ta + + @property + def root_com_lin_vel_b(self) -> ProxyArray: + """Root center of mass linear velocity in base frame [m/s]. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the linear velocity of the root center of mass frame relative to the world, + expressed in the root link's actor frame. + """ + if self._root_com_lin_vel_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_1D_kernel, + dim=self._num_instances, + inputs=[self.root_com_lin_vel_w, self.root_link_quat_w], + outputs=[self._root_com_lin_vel_b.data], + device=self.device, + ) + self._root_com_lin_vel_b.timestamp = self._sim_timestamp + if self._root_com_lin_vel_b_ta is None: + self._root_com_lin_vel_b_ta = ProxyArray(self._root_com_lin_vel_b.data) + return self._root_com_lin_vel_b_ta + + @property + def root_com_ang_vel_b(self) -> ProxyArray: + """Root center of mass angular velocity in base frame [rad/s]. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the angular velocity of the root center of mass frame relative to the world, + expressed in the root link's actor frame. + """ + if self._root_com_ang_vel_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_1D_kernel, + dim=self._num_instances, + inputs=[self.root_com_ang_vel_w, self.root_link_quat_w], + outputs=[self._root_com_ang_vel_b.data], + device=self.device, + ) + self._root_com_ang_vel_b.timestamp = self._sim_timestamp + if self._root_com_ang_vel_b_ta is None: + self._root_com_ang_vel_b_ta = ProxyArray(self._root_com_ang_vel_b.data) + return self._root_com_ang_vel_b_ta + + """ + Sliced properties. + """ + + @property + def root_link_pos_w(self) -> ProxyArray: + """Root link position in simulation world frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the position of the actor frame of the root rigid body relative to the world. + """ + parent = self.root_link_pose_w + if self._root_link_pos_w_ta is None: + self._root_link_pos_w_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._root_link_pos_w_ta + + @property + def root_link_quat_w(self) -> ProxyArray: + """Root link orientation (x, y, z, w) in simulation world frame. + + Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). + This quantity is the orientation of the actor frame of the root rigid body. + """ + parent = self.root_link_pose_w + if self._root_link_quat_w_ta is None: + self._root_link_quat_w_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._root_link_quat_w_ta + + @property + def root_link_lin_vel_w(self) -> ProxyArray: + """Root linear velocity in simulation world frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the linear velocity of the root rigid body's actor frame relative to the world. + """ + parent = self.root_link_vel_w + if self._root_link_lin_vel_w_ta is None: + self._root_link_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._root_link_lin_vel_w_ta + + @property + def root_link_ang_vel_w(self) -> ProxyArray: + """Root link angular velocity in simulation world frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. + """ + parent = self.root_link_vel_w + if self._root_link_ang_vel_w_ta is None: + self._root_link_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._root_link_ang_vel_w_ta + + @property + def root_com_pos_w(self) -> ProxyArray: + """Root center of mass position in simulation world frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the position of the center of mass frame of the root rigid body relative to the world. + """ + parent = self.root_com_pose_w + if self._root_com_pos_w_ta is None: + self._root_com_pos_w_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._root_com_pos_w_ta + + @property + def root_com_quat_w(self) -> ProxyArray: + """Root center of mass orientation (x, y, z, w) in simulation world frame. + + Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). + This quantity is the orientation of the principal axes of inertia of the root rigid body relative to the world. + """ + parent = self.root_com_pose_w + if self._root_com_quat_w_ta is None: + self._root_com_quat_w_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._root_com_quat_w_ta + + @property + def root_com_lin_vel_w(self) -> ProxyArray: + """Root center of mass linear velocity in simulation world frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. + """ + parent = self.root_com_vel_w + if self._root_com_lin_vel_w_ta is None: + self._root_com_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._root_com_lin_vel_w_ta + + @property + def root_com_ang_vel_w(self) -> ProxyArray: + """Root center of mass angular velocity in simulation world frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. + """ + parent = self.root_com_vel_w + if self._root_com_ang_vel_w_ta is None: + self._root_com_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._root_com_ang_vel_w_ta + + @property + def body_link_pos_w(self) -> ProxyArray: + """Positions of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the position of the rigid bodies' actor frame relative to the world. + """ + parent = self.body_link_pose_w + if self._body_link_pos_w_ta is None: + self._body_link_pos_w_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._body_link_pos_w_ta + + @property + def body_link_quat_w(self) -> ProxyArray: + """Orientation (x, y, z, w) of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.quatf. In torch this resolves to (num_instances, 1, 4). + This quantity is the orientation of the rigid bodies' actor frame relative to the world. + """ + parent = self.body_link_pose_w + if self._body_link_quat_w_ta is None: + self._body_link_quat_w_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._body_link_quat_w_ta + + @property + def body_link_lin_vel_w(self) -> ProxyArray: + """Linear velocity of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the linear velocity of the rigid bodies' actor frame relative to the world. + """ + parent = self.body_link_vel_w + if self._body_link_lin_vel_w_ta is None: + self._body_link_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._body_link_lin_vel_w_ta + + @property + def body_link_ang_vel_w(self) -> ProxyArray: + """Angular velocity of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the angular velocity of the rigid bodies' actor frame relative to the world. + """ + parent = self.body_link_vel_w + if self._body_link_ang_vel_w_ta is None: + self._body_link_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._body_link_ang_vel_w_ta + + @property + def body_com_pos_w(self) -> ProxyArray: + """Positions of all bodies' center of mass in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the position of the rigid bodies' center of mass frame. + """ + parent = self.body_com_pose_w + if self._body_com_pos_w_ta is None: + self._body_com_pos_w_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._body_com_pos_w_ta + + @property + def body_com_quat_w(self) -> ProxyArray: + """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.quatf. In torch this resolves to (num_instances, 1, 4). + This quantity is the orientation of the principal axes of inertia of the rigid bodies. + """ + parent = self.body_com_pose_w + if self._body_com_quat_w_ta is None: + self._body_com_quat_w_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._body_com_quat_w_ta + + @property + def body_com_lin_vel_w(self) -> ProxyArray: + """Linear velocity of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the linear velocity of the rigid bodies' center of mass frame. + """ + parent = self.body_com_vel_w + if self._body_com_lin_vel_w_ta is None: + self._body_com_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._body_com_lin_vel_w_ta + + @property + def body_com_ang_vel_w(self) -> ProxyArray: + """Angular velocity of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the angular velocity of the rigid bodies' center of mass frame. + """ + parent = self.body_com_vel_w + if self._body_com_ang_vel_w_ta is None: + self._body_com_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._body_com_ang_vel_w_ta + + @property + def body_com_lin_acc_w(self) -> ProxyArray: + """Linear acceleration of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the linear acceleration of the rigid bodies' center of mass frame. + """ + parent = self.body_com_acc_w + if self._body_com_lin_acc_w_ta is None: + self._body_com_lin_acc_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._body_com_lin_acc_w_ta + + @property + def body_com_ang_acc_w(self) -> ProxyArray: + """Angular acceleration of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the angular acceleration of the rigid bodies' center of mass frame. + """ + parent = self.body_com_acc_w + if self._body_com_ang_acc_w_ta is None: + self._body_com_ang_acc_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._body_com_ang_acc_w_ta + + @property + def body_com_pos_b(self) -> ProxyArray: + """Center of mass position of all of the bodies in their respective link frames. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the center of mass location relative to its body's link frame. + """ + parent = self.body_com_pose_b + if self._body_com_pos_b_ta is None: + self._body_com_pos_b_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._body_com_pos_b_ta + + @property + def body_com_quat_b(self) -> ProxyArray: + """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their + respective link frames. + + Shape is (num_instances, 1), dtype = wp.quatf. In torch this resolves to (num_instances, 1, 4). + This quantity is the orientation of the principal axes of inertia relative to its body's link frame. + """ + parent = self.body_com_pose_b + if self._body_com_quat_b_ta is None: + self._body_com_quat_b_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._body_com_quat_b_ta + + def _create_buffers(self) -> None: + super()._create_buffers() + # Initialize the lazy buffers. + # -- link frame w.r.t. world frame + self._root_link_pose_w = TimestampedBuffer((self._num_instances), self.device, wp.transformf) + self._root_link_vel_w = TimestampedBuffer((self._num_instances), self.device, wp.spatial_vectorf) + # -- com frame w.r.t. link frame + self._body_com_pose_b = TimestampedBuffer((self._num_instances, 1), self.device, wp.transformf) + # -- com frame w.r.t. world frame + self._root_com_pose_w = TimestampedBuffer((self._num_instances), self.device, wp.transformf) + self._root_com_vel_w = TimestampedBuffer((self._num_instances), self.device, wp.spatial_vectorf) + self._body_com_acc_w = TimestampedBuffer((self._num_instances, 1), self.device, wp.spatial_vectorf) + # -- combined state (these are cached as they concatenate) + self._root_state_w = TimestampedBuffer((self._num_instances), self.device, shared_kernels.vec13f) + self._root_link_state_w = TimestampedBuffer((self._num_instances), self.device, shared_kernels.vec13f) + self._root_com_state_w = TimestampedBuffer((self._num_instances), self.device, shared_kernels.vec13f) + # -- derived properties (these are cached to avoid repeated memory allocations) + self._projected_gravity_b = TimestampedBuffer((self._num_instances), self.device, wp.vec3f) + self._heading_w = TimestampedBuffer((self._num_instances), self.device, wp.float32) + self._root_link_lin_vel_b = TimestampedBuffer((self._num_instances), self.device, wp.vec3f) + self._root_link_ang_vel_b = TimestampedBuffer((self._num_instances), self.device, wp.vec3f) + self._root_com_lin_vel_b = TimestampedBuffer((self._num_instances), self.device, wp.vec3f) + self._root_com_ang_vel_b = TimestampedBuffer((self._num_instances), self.device, wp.vec3f) + + # -- Default state + self._default_root_pose = wp.zeros((self._num_instances), dtype=wp.transformf, device=self.device) + self._default_root_vel = wp.zeros((self._num_instances), dtype=wp.spatial_vectorf, device=self.device) + self._default_root_state = None + + # -- Previous body com velocity + self._previous_body_com_vel = None + + # -- Pinned-host staging buffers for CPU-only bindings on a non-CPU sim + # (lazily allocated, keyed by tensor type). + self._cpu_staging_buffers: dict[int, wp.array] = {} + + # -- Body properties (semi-static; read once from CPU-only bindings). + # The wheel exposes ``RIGID_BODY_MASS`` as ``(N,)`` and ``RIGID_BODY_INERTIA`` as ``(N, 9)``; + # the ``BaseRigidObjectData`` contract is ``(N, 1)`` and ``(N, 1, 9)`` respectively, so we + # read into a flat buffer and reshape (zero-copy) after the read. + mass_binding = self._bindings[TT.RIGID_BODY_MASS] + inertia_binding = self._bindings[TT.RIGID_BODY_INERTIA] + self._body_mass = wp.zeros(mass_binding.shape, dtype=wp.float32, device=self.device) + self._body_inertia = wp.zeros(inertia_binding.shape, dtype=wp.float32, device=self.device) + self._read_binding_into(TT.RIGID_BODY_MASS, self._body_mass) + self._read_binding_into(TT.RIGID_BODY_INERTIA, self._body_inertia) + self._body_mass = self._body_mass.reshape((self._num_instances, 1)) + self._body_inertia = self._body_inertia.reshape((self._num_instances, 1, 9)) + + # Initialize ProxyArray wrappers + self._pin_proxy_arrays() + + def _pin_proxy_arrays(self) -> None: + """Create pinned ProxyArray wrappers for all data buffers. + + This is called once from :meth:`_create_buffers` during initialization. + PhysX tensor API buffers have stable GPU pointers across simulation steps, + so no rebinding is needed (unlike Newton). + """ + # -- Pinned ProxyArray cache (one per read property, lazily created on first access) + # Defaults + self._default_root_pose_ta: ProxyArray | None = None + self._default_root_vel_ta: ProxyArray | None = None + # Root state (timestamped) + self._root_link_pose_w_ta: ProxyArray | None = None + self._root_link_vel_w_ta: ProxyArray | None = None + self._root_com_pose_w_ta: ProxyArray | None = None + self._root_com_vel_w_ta: ProxyArray | None = None + # Body properties + self._body_mass_ta: ProxyArray | None = None + self._body_inertia_ta: ProxyArray | None = None + # Body state (reshaped from root) + self._body_link_pose_w_ta: ProxyArray | None = None + self._body_link_vel_w_ta: ProxyArray | None = None + self._body_com_pose_w_ta: ProxyArray | None = None + self._body_com_vel_w_ta: ProxyArray | None = None + self._body_com_acc_w_ta: ProxyArray | None = None + self._body_com_pose_b_ta: ProxyArray | None = None + # Derived properties (timestamped) + self._projected_gravity_b_ta: ProxyArray | None = None + self._heading_w_ta: ProxyArray | None = None + self._root_link_lin_vel_b_ta: ProxyArray | None = None + self._root_link_ang_vel_b_ta: ProxyArray | None = None + self._root_com_lin_vel_b_ta: ProxyArray | None = None + self._root_com_ang_vel_b_ta: ProxyArray | None = None + # Sliced properties (root link) + self._root_link_pos_w_ta: ProxyArray | None = None + self._root_link_quat_w_ta: ProxyArray | None = None + self._root_link_lin_vel_w_ta: ProxyArray | None = None + self._root_link_ang_vel_w_ta: ProxyArray | None = None + # Sliced properties (root com) + self._root_com_pos_w_ta: ProxyArray | None = None + self._root_com_quat_w_ta: ProxyArray | None = None + self._root_com_lin_vel_w_ta: ProxyArray | None = None + self._root_com_ang_vel_w_ta: ProxyArray | None = None + # Sliced properties (body link) + self._body_link_pos_w_ta: ProxyArray | None = None + self._body_link_quat_w_ta: ProxyArray | None = None + self._body_link_lin_vel_w_ta: ProxyArray | None = None + self._body_link_ang_vel_w_ta: ProxyArray | None = None + # Sliced properties (body com) + self._body_com_pos_w_ta: ProxyArray | None = None + self._body_com_quat_w_ta: ProxyArray | None = None + self._body_com_lin_vel_w_ta: ProxyArray | None = None + self._body_com_ang_vel_w_ta: ProxyArray | None = None + self._body_com_lin_acc_w_ta: ProxyArray | None = None + self._body_com_ang_acc_w_ta: ProxyArray | None = None + # Sliced properties (body com in body frame) + self._body_com_pos_b_ta: ProxyArray | None = None + self._body_com_quat_b_ta: ProxyArray | None = None + # Deprecated state-concat properties + self._default_root_state_ta: ProxyArray | None = None + self._root_state_w_ta: ProxyArray | None = None + self._root_link_state_w_ta: ProxyArray | None = None + self._root_com_state_w_ta: ProxyArray | None = None + self._body_state_w_ta: ProxyArray | None = None + self._body_link_state_w_ta: ProxyArray | None = None + self._body_com_state_w_ta: ProxyArray | None = None + + """ + Internal helpers. + """ + + def _get_binding(self, tensor_type: int): + """Return the binding for the given tensor type, or None.""" + return self._bindings.get(tensor_type) + + def _read_binding_into(self, tensor_type: int, dst: wp.array) -> None: + """Read the OVPhysX TensorBinding for *tensor_type* into *dst*. + + Adapter that replaces PhysX's view-getter pattern: the wheel exposes + ``binding.read(target)`` rather than a getter returning a wp.array, so + we read into a flat float32 view of *dst*. CPU-only bindings on a + non-CPU sim go through a lazily-allocated pinned-host wp.array to + satisfy the wheel's device match. + """ + binding = self._bindings[tensor_type] + if self._check_shapes: + dst_bytes = dst.size * wp.types.type_size_in_bytes(dst.dtype) + binding_bytes = 4 * math.prod(binding.shape) + assert dst_bytes >= binding_bytes, ( + f"_read_binding_into: dst buffer too small for binding {tensor_type!r} " + f"({dst_bytes} B < {binding_bytes} B). Caller allocated dst with " + f"shape={tuple(dst.shape)}, dtype={dst.dtype}; binding shape={tuple(binding.shape)}." + ) + # Build a flat float32 view of dst matching the binding's shape. + if dst.dtype == wp.float32: + view = dst + else: + view = wp.array( + ptr=dst.ptr, + shape=binding.shape, + dtype=wp.float32, + device=str(dst.device), + copy=False, + ) + if tensor_type in TT._CPU_ONLY_TYPES and str(view.device) != "cpu": + staging = self._cpu_staging_buffers.get(tensor_type) + if staging is None: + staging = wp.zeros(binding.shape, dtype=wp.float32, device="cpu", pinned=True) + self._cpu_staging_buffers[tensor_type] = staging + binding.read(staging) + wp.copy(view, staging) + else: + binding.read(view) + + def _get_pos_from_transform(self, transform: wp.array) -> wp.array: + """Generates a position array from a transform array.""" + return wp.array( + ptr=transform.ptr, + shape=transform.shape, + dtype=wp.vec3f, + strides=transform.strides, + device=self.device, + ) + + def _get_quat_from_transform(self, transform: wp.array) -> wp.array: + """Generates a quaternion array from a transform array.""" + return wp.array( + ptr=transform.ptr + 3 * 4, + shape=transform.shape, + dtype=wp.quatf, + strides=transform.strides, + device=self.device, + ) + + def _get_lin_vel_from_spatial_vector(self, sv: wp.array) -> wp.array: + """Generates a linear velocity array from a spatial vector array.""" + return wp.array( + ptr=sv.ptr, + shape=sv.shape, + dtype=wp.vec3f, + strides=sv.strides, + device=self.device, + ) + + def _get_ang_vel_from_spatial_vector(self, sv: wp.array) -> wp.array: + """Generates an angular velocity array from a spatial vector array.""" + return wp.array( + ptr=sv.ptr + 3 * 4, + shape=sv.shape, + dtype=wp.vec3f, + strides=sv.strides, + device=self.device, + ) + + """ + Deprecated properties. + """ + + @property + def default_root_state(self) -> ProxyArray: + """Default root state ``[pos, quat, lin_vel, ang_vel]`` in local environment frame. + + The position and quaternion are of the rigid body's actor frame. Meanwhile, the linear and angular velocities + are of the center of mass frame. Shape is (num_instances, 13). + """ + warnings.warn( + "Reading the root state directly is deprecated since IsaacLab 3.0 and will be removed in a future version. " + "Please use the default_root_pose and default_root_vel properties instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._default_root_state is None: + self._default_root_state = wp.zeros((self._num_instances), dtype=shared_kernels.vec13f, device=self.device) + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self._default_root_pose, + self._default_root_vel, + ], + outputs=[ + self._default_root_state, + ], + device=self.device, + ) + if self._default_root_state_ta is None: + self._default_root_state_ta = ProxyArray(self._default_root_state) + return self._default_root_state_ta + + @property + def root_state_w(self) -> ProxyArray: + """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_com_vel_w`.""" + warnings.warn( + "The `root_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_link_pose_w` and " + "`root_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._root_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self.root_link_pose_w, + self.root_com_vel_w, + ], + outputs=[ + self._root_state_w.data, + ], + device=self.device, + ) + self._root_state_w.timestamp = self._sim_timestamp + + if self._root_state_w_ta is None: + self._root_state_w_ta = ProxyArray(self._root_state_w.data) + return self._root_state_w_ta + + @property + def root_link_state_w(self) -> ProxyArray: + """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_link_vel_w`.""" + warnings.warn( + "The `root_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_link_pose_w` and " + "`root_link_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._root_link_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self.root_link_pose_w, + self.root_link_vel_w, + ], + outputs=[ + self._root_link_state_w.data, + ], + device=self.device, + ) + self._root_link_state_w.timestamp = self._sim_timestamp + + if self._root_link_state_w_ta is None: + self._root_link_state_w_ta = ProxyArray(self._root_link_state_w.data) + return self._root_link_state_w_ta + + @property + def root_com_state_w(self) -> ProxyArray: + """Deprecated, same as :attr:`root_com_pose_w` and :attr:`root_com_vel_w`.""" + warnings.warn( + "The `root_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_com_pose_w` and " + "`root_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._root_com_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self.root_com_pose_w, + self.root_com_vel_w, + ], + outputs=[ + self._root_com_state_w.data, + ], + device=self.device, + ) + self._root_com_state_w.timestamp = self._sim_timestamp + + if self._root_com_state_w_ta is None: + self._root_com_state_w_ta = ProxyArray(self._root_com_state_w.data) + return self._root_com_state_w_ta + + @property + def body_state_w(self) -> ProxyArray: + """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_com_vel_w`.""" + warnings.warn( + "The `body_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " + "`body_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + # Access internal buffer directly to avoid cascading deprecation warnings from root_state_w + if self._root_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self.root_link_pose_w, + self.root_com_vel_w, + ], + outputs=[ + self._root_state_w.data, + ], + device=self.device, + ) + self._root_state_w.timestamp = self._sim_timestamp + if self._body_state_w_ta is None: + self._body_state_w_ta = ProxyArray(self._root_state_w.data.reshape((self._num_instances, 1))) + return self._body_state_w_ta + + @property + def body_link_state_w(self) -> ProxyArray: + """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_link_vel_w`.""" + warnings.warn( + "The `body_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " + "`body_link_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + # Access internal buffer directly to avoid cascading deprecation warnings from root_link_state_w + if self._root_link_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self.root_link_pose_w, + self.root_link_vel_w, + ], + outputs=[ + self._root_link_state_w.data, + ], + device=self.device, + ) + self._root_link_state_w.timestamp = self._sim_timestamp + if self._body_link_state_w_ta is None: + self._body_link_state_w_ta = ProxyArray(self._root_link_state_w.data.reshape((self._num_instances, 1))) + return self._body_link_state_w_ta + + @property + def body_com_state_w(self) -> ProxyArray: + """Deprecated, same as :attr:`body_com_pose_w` and :attr:`body_com_vel_w`.""" + warnings.warn( + "The `body_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_com_pose_w` and " + "`body_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + # Access internal buffer directly to avoid cascading deprecation warnings from root_com_state_w + if self._root_com_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self.root_com_pose_w, + self.root_com_vel_w, + ], + outputs=[ + self._root_com_state_w.data, + ], + device=self.device, + ) + self._root_com_state_w.timestamp = self._sim_timestamp + if self._body_com_state_w_ta is None: + self._body_com_state_w_ta = ProxyArray(self._root_com_state_w.data.reshape((self._num_instances, 1))) + return self._body_com_state_w_ta diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/__init__.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/__init__.py new file mode 100644 index 000000000000..3fce63a6d0e0 --- /dev/null +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/__init__.py @@ -0,0 +1,10 @@ +# 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 + +"""Sub-module for ovphysx-backed rigid object collection assets.""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/__init__.pyi b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/__init__.pyi new file mode 100644 index 000000000000..8b12ec95e7a2 --- /dev/null +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/__init__.pyi @@ -0,0 +1,12 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +__all__ = [ + "RigidObjectCollection", + "RigidObjectCollectionData", +] + +from .rigid_object_collection import RigidObjectCollection +from .rigid_object_collection_data import RigidObjectCollectionData diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/rigid_object_collection.py new file mode 100644 index 000000000000..7a1d4e21e535 --- /dev/null +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/rigid_object_collection.py @@ -0,0 +1,1525 @@ +# 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 + +from __future__ import annotations + +import re +import warnings +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch +import warp as wp + +from pxr import UsdPhysics + +import isaaclab.sim as sim_utils +from isaaclab.assets.rigid_object_collection.base_rigid_object_collection import BaseRigidObjectCollection +from isaaclab.utils.string import resolve_matching_names +from isaaclab.utils.wrench_composer import WrenchComposer + +from isaaclab_ovphysx import tensor_types as TT +from isaaclab_ovphysx.assets import kernels as shared_kernels +from isaaclab_ovphysx.assets.kernels import _body_wrench_to_world, resolve_view_ids +from isaaclab_ovphysx.physics import OvPhysxManager + +from .rigid_object_collection_data import RigidObjectCollectionData + +if TYPE_CHECKING: + from isaaclab.assets.rigid_object_collection.rigid_object_collection_cfg import RigidObjectCollectionCfg + + +class RigidObjectCollection(BaseRigidObjectCollection): + """A rigid object collection class. + + This class represents a collection of rigid objects in the simulation, where the state of the + rigid objects can be accessed and modified using a batched ``(env_ids, object_ids)`` API. + + For each rigid body in the collection, the root prim of the asset must have the `USD RigidBodyAPI`_ + applied to it. This API is used to define the simulation properties of the rigid bodies. On playing the + simulation, the physics engine will automatically register the rigid bodies and create a corresponding + rigid body handle. This handle can be accessed using the :attr:`root_view` attribute. + + Rigid objects in the collection are uniquely identified via the key of the dictionary + :attr:`~isaaclab.assets.RigidObjectCollectionCfg.rigid_objects` in the + :class:`~isaaclab.assets.RigidObjectCollectionCfg` configuration class. + This differs from the :class:`~isaaclab.assets.RigidObject` class, where a rigid object is identified by + the name of the Xform where the `USD RigidBodyAPI`_ is applied. This would not be possible for the rigid + object collection since the :attr:`~isaaclab.assets.RigidObjectCollectionCfg.rigid_objects` dictionary + could contain the same rigid object multiple times, leading to ambiguity. + + .. _`USD RigidBodyAPI`: https://openusd.org/dev/api/class_usd_physics_rigid_body_a_p_i.html + """ + + cfg: RigidObjectCollectionCfg + """Configuration instance for the rigid object.""" + + __backend_name__: str = "ovphysx" + """The name of the backend for the rigid object.""" + + def __init__(self, cfg: RigidObjectCollectionCfg): + """Initialize the rigid object. + + Args: + cfg: A configuration instance. + """ + # Note: We never call the parent constructor as it tries to call its own spawning which we don't want. + # check that the config is valid + cfg.validate() + # store inputs + self.cfg = cfg.copy() + # flag for whether the asset is initialized + self._is_initialized = False + # spawn the rigid objects + for rigid_body_cfg in self.cfg.rigid_objects.values(): + # spawn the asset + if rigid_body_cfg.spawn is not None: + rigid_body_cfg.spawn.func( + rigid_body_cfg.prim_path, + rigid_body_cfg.spawn, + translation=rigid_body_cfg.init_state.pos, + orientation=rigid_body_cfg.init_state.rot, + ) + # check that spawn was successful + matching_prims = sim_utils.find_matching_prims(rigid_body_cfg.prim_path) + if len(matching_prims) == 0: + raise RuntimeError(f"Could not find prim with path {rigid_body_cfg.prim_path}.") + # stores object names + self._body_names_list: list[str] = [] + # one fused TensorBinding per tensor type, populated in _initialize_impl + self._bindings: dict[int, Any] = {} + + # register various callback functions + self._register_callbacks() + self._debug_vis_handle = None + + """ + Properties + """ + + @property + def data(self) -> RigidObjectCollectionData: + return self._data + + @property + def num_instances(self) -> int: + return self._num_instances + + @property + def num_bodies(self) -> int: + """Number of bodies in the rigid object collection.""" + return self._num_bodies + + @property + def body_names(self) -> list[str]: + """Ordered names of bodies in the rigid object collection.""" + return list(self._body_names_list) + + @property + def root_view(self): + """Root view for the rigid object collection. + + Dictionary keyed by TensorType constant, each value a single fused + :class:`~isaaclab_ovphysx.TensorBinding` spanning all bodies in the collection. + + .. note:: + Use this view with caution. It requires handling of tensors in a specific way. + """ + return self._bindings + + @property + def instantaneous_wrench_composer(self) -> WrenchComposer: + """Instantaneous wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are only valid for the current simulation step. At the end of the simulation step, the wrenches set + to this object are discarded. This is useful to apply forces that change all the time, things like drag forces + for instance. + """ + return self._instantaneous_wrench_composer + + @property + def permanent_wrench_composer(self) -> WrenchComposer: + """Permanent wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are persistent and are applied to the simulation at every step. This is useful to apply forces that + are constant over a period of time, things like the thrust of a motor for instance. + """ + return self._permanent_wrench_composer + + """ + Operations. + """ + + def reset( + self, + env_ids: Sequence[int] | wp.array | None = None, + object_ids: slice | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Resets all internal buffers of selected environments and objects. + + Args: + env_ids: Environment indices. If None, then all indices are used. + object_ids: Object indices. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + # resolve all indices + if (env_ids is None) or (env_ids == slice(None)): + env_ids = slice(None) + # reset external wrench + self._instantaneous_wrench_composer.reset(env_ids, env_mask) + self._permanent_wrench_composer.reset(env_ids, env_mask) + + def write_data_to_sim(self) -> None: + """Write external wrench to the simulation. + + .. note:: + We write external wrench to the simulation here since this function is called before the simulation step. + This ensures that the external wrench is applied at every simulation step. + """ + inst = self._instantaneous_wrench_composer + perm = self._permanent_wrench_composer + if not inst.active and not perm.active: + return + if inst.active: + if perm.active: + inst.add_raw_buffers_from(perm) + force_b = inst.out_force_b.warp + torque_b = inst.out_torque_b.warp + else: + force_b = perm.out_force_b.warp + torque_b = perm.out_torque_b.warp + + poses = self._data.body_link_pose_w.warp # (N, B) wp.transformf + wp.launch( + _body_wrench_to_world, + dim=(self._num_instances, self._num_bodies), + inputs=[force_b, torque_b, poses], + outputs=[self._wrench_buf], + device=self._device, + ) + binding = self._get_binding(TT.LINK_WRENCH) + if binding is not None: + # The articulation-mode mock used by iface tests exposes an instance-major + # ``(N, B, 9)`` view directly; the native fused binding lays elements body- + # major flat as ``(N * B, 9)``. Dispatch via the binding's exposed shape. + if len(binding.shape) >= 2 and binding.shape[1] == self._num_bodies: + binding.write(self._wrench_buf) + else: + view = self.reshape_data_to_view_3d(self._wrench_buf, 9, device=self._device) + binding.write(view) + inst.reset() + + def update(self, dt: float) -> None: + """Updates the simulation data. + + Args: + dt: The time step size in seconds. + """ + self._data.update(dt) + + """ + Operations - Finders. + """ + + def find_bodies( + self, name_keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[torch.Tensor, list[str]]: + """Find bodies in the rigid body collection based on the name keys. + + Please check the :meth:`isaaclab.utils.string_utils.resolve_matching_names` function for more + information on the name matching. + + Args: + name_keys: A regular expression or a list of regular expressions to match the body names. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the body indices and names. + """ + obj_ids, obj_names = resolve_matching_names(name_keys, self.body_names, preserve_order) + return torch.tensor(obj_ids, device=self._device, dtype=torch.int32), obj_names + + """ + Operations - Write to simulation. + """ + + def write_body_pose_to_sim_index( + self, + *, + body_poses: wp.array, + body_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | wp.array | None = None, + ) -> None: + """Set the body pose over selected environment and body indices into the simulation. + + The body pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + .. note:: + This method expects partial data. + + Args: + body_poses: Body poses in simulation frame [m, rad]. Shape is (len(env_ids), len(body_ids), 7) + or (len(env_ids), len(body_ids)) with dtype wp.transformf. + body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. + """ + self.write_body_link_pose_to_sim_index(body_poses=body_poses, body_ids=body_ids, env_ids=env_ids) + + def write_body_pose_to_sim_mask( + self, + *, + body_poses: wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set the body pose over selected environment and body masks into the simulation. + + The body pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + .. note:: + This method expects full data. + + Args: + body_poses: Body poses in simulation frame [m, rad]. Shape is (num_instances, num_bodies, 7) + or (num_instances, num_bodies) with dtype wp.transformf. + body_mask: Body mask. If None, then all bodies are updated. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + self.write_body_link_pose_to_sim_mask(body_poses=body_poses, body_mask=body_mask, env_mask=env_mask) + + def write_body_velocity_to_sim_index( + self, + *, + body_velocities: wp.array, + body_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | wp.array | None = None, + ) -> None: + """Set the body velocity over selected environment and body indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the body's center of mass rather than the body's frame. + + .. note:: + This method expects partial data. + + Args: + body_velocities: Body velocities in simulation world frame [m/s, rad/s]. + Shape is (len(env_ids), len(body_ids)) with dtype wp.spatial_vectorf. + body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. + """ + self.write_body_com_velocity_to_sim_index(body_velocities=body_velocities, body_ids=body_ids, env_ids=env_ids) + + def write_body_velocity_to_sim_mask( + self, + *, + body_velocities: wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set the body velocity over selected environment and body masks into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the body's center of mass rather than the body's frame. + + .. note:: + This method expects full data. + + Args: + body_velocities: Body velocities in simulation world frame [m/s, rad/s]. + Shape is (num_instances, num_bodies) with dtype wp.spatial_vectorf. + body_mask: Body mask. If None, then all bodies are updated. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + self.write_body_com_velocity_to_sim_mask( + body_velocities=body_velocities, body_mask=body_mask, env_mask=env_mask + ) + + def write_body_link_pose_to_sim_index( + self, + *, + body_poses: wp.array, + body_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | wp.array | None = None, + ) -> None: + """Set the body link pose over selected environment and body indices into the simulation. + + The body link pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + .. note:: + This method expects partial data. + + Args: + body_poses: Body link poses in simulation frame [m, rad]. Shape is (len(env_ids), len(body_ids), 7) + or (len(env_ids), len(body_ids)) with dtype wp.transformf. + body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. + """ + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + self.assert_shape_and_dtype(body_poses, (env_ids.shape[0], body_ids.shape[0]), wp.transformf, "body_poses") + wp.launch( + shared_kernels.set_body_link_pose_to_sim, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[body_poses, env_ids, body_ids, False], + outputs=[ + self.data._body_link_pose_w.data, + self.data._body_link_state_w.data, + self.data._body_state_w.data, + ], + device=self._device, + ) + # Mark the link pose fresh so reads within the same step return the + # kernel-written value rather than re-fetching the pre-step OVPhysX state. + self.data._body_link_pose_w.timestamp = self.data._sim_timestamp + self.data._body_com_pose_w.timestamp = -1.0 + self.data._body_com_state_w.timestamp = -1.0 + self.data._body_link_state_w.timestamp = -1.0 + self.data._body_state_w.timestamp = -1.0 + # set into simulation + self._binding_write(TT.LINK_POSE, self.data._body_link_pose_w.data, env_ids=env_ids) + + def write_body_link_pose_to_sim_mask( + self, + *, + body_poses: wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set the body link pose over selected environment and body masks into the simulation. + + The body link pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + .. note:: + This method expects full data. + + Args: + body_poses: Body link poses in simulation frame [m, rad]. Shape is (num_instances, num_bodies, 7) + or (num_instances, num_bodies) with dtype wp.transformf. + body_mask: Body mask. If None, then all bodies are updated. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the 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].to(torch.int32)) + 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].to(torch.int32)) + else: + body_ids = self._ALL_BODY_INDICES + 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, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[body_poses, env_ids, body_ids, True], + outputs=[ + self.data._body_link_pose_w.data, + self.data._body_link_state_w.data, + self.data._body_state_w.data, + ], + device=self._device, + ) + # Invalidate dependent timestamps + self.data._body_com_pose_w.timestamp = -1.0 + self.data._body_com_state_w.timestamp = -1.0 + self.data._body_link_state_w.timestamp = -1.0 + self.data._body_state_w.timestamp = -1.0 + # set into simulation + self._binding_write(TT.LINK_POSE, self.data._body_link_pose_w.data, env_ids=env_ids) + + def write_body_com_pose_to_sim_index( + self, + *, + body_poses: wp.array, + body_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | wp.array | None = None, + ) -> None: + """Set the body center of mass pose over selected environment and body indices into the simulation. + + The body center of mass pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + The orientation is the orientation of the principal axes of inertia. + + .. note:: + This method expects partial data. + + Args: + body_poses: Body center of mass poses in simulation frame [m, rad]. + Shape is (len(env_ids), len(body_ids), 7) or (len(env_ids), len(body_ids)) with dtype wp.transformf. + body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. + """ + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + self.assert_shape_and_dtype(body_poses, (env_ids.shape[0], body_ids.shape[0]), wp.transformf, "body_poses") + wp.launch( + shared_kernels.set_body_com_pose_to_sim, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[body_poses, self.data.body_com_pose_b, env_ids, body_ids, False], + outputs=[ + self.data._body_com_pose_w.data, + self.data._body_link_pose_w.data, + self.data._body_com_state_w.data, + self.data._body_link_state_w.data, + self.data._body_state_w.data, + ], + device=self._device, + ) + # Invalidate dependent timestamps + self.data._body_link_state_w.timestamp = -1.0 + self.data._body_state_w.timestamp = -1.0 + self.data._body_com_state_w.timestamp = -1.0 + # set into simulation (OVPhysX only exposes the link frame) + self._binding_write(TT.LINK_POSE, self.data._body_link_pose_w.data, env_ids=env_ids) + + def write_body_com_pose_to_sim_mask( + self, + *, + body_poses: wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set the body center of mass pose over selected environment and body masks into the simulation. + + The body center of mass pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + The orientation is the orientation of the principal axes of inertia. + + .. note:: + This method expects full data. + + Args: + body_poses: Body center of mass poses in simulation frame [m, rad]. + Shape is (num_instances, num_bodies, 7) or (num_instances, num_bodies) with dtype wp.transformf. + body_mask: Body mask. If None, then all bodies are updated. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the 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].to(torch.int32)) + 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].to(torch.int32)) + else: + body_ids = self._ALL_BODY_INDICES + 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, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[body_poses, self.data.body_com_pose_b, env_ids, body_ids, True], + outputs=[ + self.data._body_com_pose_w.data, + self.data._body_link_pose_w.data, + self.data._body_com_state_w.data, + self.data._body_link_state_w.data, + self.data._body_state_w.data, + ], + device=self._device, + ) + # Invalidate dependent timestamps + self.data._body_link_state_w.timestamp = -1.0 + self.data._body_state_w.timestamp = -1.0 + self.data._body_com_state_w.timestamp = -1.0 + # set into simulation (OVPhysX only exposes the link frame) + self._binding_write(TT.LINK_POSE, self.data._body_link_pose_w.data, env_ids=env_ids) + + def write_body_com_velocity_to_sim_index( + self, + *, + body_velocities: wp.array, + body_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | wp.array | None = None, + ) -> None: + """Set the body center of mass velocity over selected environment and body indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the body's center of mass rather than the body's frame. + + .. note:: + This method expects partial data. + + Args: + body_velocities: Body center of mass velocities in simulation world frame [m/s, rad/s]. + Shape is (len(env_ids), len(body_ids)) with dtype wp.spatial_vectorf. + body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. + """ + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + self.assert_shape_and_dtype( + body_velocities, (env_ids.shape[0], body_ids.shape[0]), wp.spatial_vectorf, "body_velocities" + ) + wp.launch( + shared_kernels.set_body_com_velocity_to_sim, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[body_velocities, env_ids, body_ids, False], + outputs=[ + self.data._body_com_vel_w.data, + self.data._body_com_acc_w.data, + self.data._body_state_w.data, + self.data._body_com_state_w.data, + ], + device=self._device, + ) + # Mark the COM velocity fresh so reads within the same step return the + # kernel-written value rather than re-fetching the pre-step OVPhysX state. + self.data._body_com_vel_w.timestamp = self.data._sim_timestamp + self.data._body_link_vel_w.timestamp = -1.0 + self.data._body_state_w.timestamp = -1.0 + self.data._body_com_state_w.timestamp = -1.0 + self.data._body_link_state_w.timestamp = -1.0 + # set into simulation + self._binding_write(TT.LINK_VELOCITY, self.data._body_com_vel_w.data, env_ids=env_ids) + + def write_body_com_velocity_to_sim_mask( + self, + *, + body_velocities: wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set the body center of mass velocity over selected environment and body masks into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the body's center of mass rather than the body's frame. + + .. note:: + This method expects full data. + + Args: + body_velocities: Body center of mass velocities in simulation world frame [m/s, rad/s]. + Shape is (num_instances, num_bodies) with dtype wp.spatial_vectorf. + body_mask: Body mask. If None, then all bodies are updated. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the 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].to(torch.int32)) + 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].to(torch.int32)) + else: + body_ids = self._ALL_BODY_INDICES + self.assert_shape_and_dtype( + body_velocities, (self._num_instances, self._num_bodies), wp.spatial_vectorf, "body_velocities" + ) + wp.launch( + shared_kernels.set_body_com_velocity_to_sim, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[body_velocities, env_ids, body_ids, True], + outputs=[ + self.data._body_com_vel_w.data, + self.data._body_com_acc_w.data, + self.data._body_state_w.data, + self.data._body_com_state_w.data, + ], + device=self._device, + ) + # Invalidate dependent timestamps + self.data._body_link_vel_w.timestamp = -1.0 + self.data._body_state_w.timestamp = -1.0 + self.data._body_com_state_w.timestamp = -1.0 + self.data._body_link_state_w.timestamp = -1.0 + # set into simulation + self._binding_write(TT.LINK_VELOCITY, self.data._body_com_vel_w.data, env_ids=env_ids) + + def write_body_link_velocity_to_sim_index( + self, + *, + body_velocities: wp.array, + body_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | wp.array | None = None, + ) -> None: + """Set the body link velocity over selected environment and body indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the body's frame rather than the body's center of mass. + + .. note:: + This method expects partial data. + + Args: + body_velocities: Body link velocities in simulation world frame [m/s, rad/s]. + Shape is (len(env_ids), len(body_ids)) with dtype wp.spatial_vectorf. + body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. + """ + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + self.assert_shape_and_dtype( + body_velocities, (env_ids.shape[0], body_ids.shape[0]), wp.spatial_vectorf, "body_velocities" + ) + wp.launch( + shared_kernels.set_body_link_velocity_to_sim, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[ + body_velocities, + self.data.body_com_pose_b, + self.data.body_link_pose_w, + env_ids, + body_ids, + False, + ], + outputs=[ + self.data._body_link_vel_w.data, + self.data._body_com_vel_w.data, + self.data._body_com_acc_w.data, + self.data._body_link_state_w.data, + self.data._body_state_w.data, + self.data._body_com_state_w.data, + ], + device=self._device, + ) + # Invalidate dependent timestamps + self.data._body_link_state_w.timestamp = -1.0 + self.data._body_state_w.timestamp = -1.0 + self.data._body_com_state_w.timestamp = -1.0 + # set into simulation + self._binding_write(TT.LINK_VELOCITY, self.data._body_com_vel_w.data, env_ids=env_ids) + + def write_body_link_velocity_to_sim_mask( + self, + *, + body_velocities: wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set the body link velocity over selected environment and body masks into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the body's frame rather than the body's center of mass. + + .. note:: + This method expects full data. + + Args: + body_velocities: Body link velocities in simulation world frame [m/s, rad/s]. + Shape is (num_instances, num_bodies) with dtype wp.spatial_vectorf. + body_mask: Body mask. If None, then all bodies are updated. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the 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].to(torch.int32)) + 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].to(torch.int32)) + else: + body_ids = self._ALL_BODY_INDICES + self.assert_shape_and_dtype( + body_velocities, (self._num_instances, self._num_bodies), wp.spatial_vectorf, "body_velocities" + ) + wp.launch( + shared_kernels.set_body_link_velocity_to_sim, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[ + body_velocities, + self.data.body_com_pose_b, + self.data.body_link_pose_w, + env_ids, + body_ids, + True, + ], + outputs=[ + self.data._body_link_vel_w.data, + self.data._body_com_vel_w.data, + self.data._body_com_acc_w.data, + self.data._body_link_state_w.data, + self.data._body_state_w.data, + self.data._body_com_state_w.data, + ], + device=self._device, + ) + # Invalidate dependent timestamps + self.data._body_link_state_w.timestamp = -1.0 + self.data._body_state_w.timestamp = -1.0 + self.data._body_com_state_w.timestamp = -1.0 + # set into simulation + self._binding_write(TT.LINK_VELOCITY, self.data._body_com_vel_w.data, env_ids=env_ids) + + """ + Operations - Setters. + """ + + def set_masses_index( + self, + *, + masses: wp.array, + body_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | wp.array | None = None, + ) -> None: + """Set body masses over selected env / body indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``BODY_MASS`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + Args: + masses: Body masses [kg]. Shape is (len(env_ids), len(body_ids)) + with dtype wp.float32. + body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. + """ + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + self.assert_shape_and_dtype(masses, (env_ids.shape[0], body_ids.shape[0]), wp.float32, "masses") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[masses, env_ids, body_ids], + outputs=[self.data._body_mass.data], + device=self._device, + ) + wp.copy(self.data._cpu_body_mass, self.data._body_mass.data) + self._binding_write( + TT.BODY_MASS, self.data._cpu_body_mass, env_ids=self._get_cpu_env_ids(env_ids), device="cpu" + ) + + def set_masses_mask( + self, + *, + masses: wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set body masses over selected env / body masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``BODY_MASS`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + Args: + masses: Body masses [kg]. Shape is (num_instances, num_bodies) + with dtype wp.float32. + body_mask: Body mask. If None, all bodies are updated. + Shape is (num_bodies,). + 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].to(torch.int32)) + else: + env_ids = self._ALL_ENV_INDICES + 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, + dim=(self._num_instances, self._num_bodies), + inputs=[masses, self._resolve_env_mask(env_mask), self._resolve_body_mask(body_mask)], + outputs=[self.data._body_mass.data], + device=self._device, + ) + wp.copy(self.data._cpu_body_mass, self.data._body_mass.data) + self._binding_write( + TT.BODY_MASS, self.data._cpu_body_mass, env_ids=self._get_cpu_env_ids(env_ids), device="cpu" + ) + + def set_coms_index( + self, + *, + coms: wp.array, + body_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | wp.array | None = None, + ) -> None: + """Set body center-of-mass poses over selected env / body indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``BODY_COM_POSE`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + Args: + coms: Body center-of-mass poses [m, quaternion (w, x, y, z)]. + Shape is (len(env_ids), len(body_ids)) with dtype wp.transformf. + body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. + """ + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + self.assert_shape_and_dtype(coms, (env_ids.shape[0], body_ids.shape[0]), wp.transformf, "coms") + wp.launch( + shared_kernels.write_body_com_pose_to_buffer_index, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[coms, env_ids, body_ids], + outputs=[self.data._body_com_pose_b.data], + device=self._device, + ) + # Invalidate derived buffers that depend on body_com_pose_b. + self.data._body_com_pose_w.timestamp = -1.0 + wp.copy(self.data._cpu_body_coms, self.data._body_com_pose_b.data) + self._binding_write( + TT.BODY_COM_POSE, + self.data._cpu_body_coms, + env_ids=self._get_cpu_env_ids(env_ids), + device="cpu", + data_dim=7, + ) + + def set_coms_mask( + self, + *, + coms: wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set body center-of-mass poses over selected env / body masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``BODY_COM_POSE`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + Args: + coms: Body center-of-mass poses [m, quaternion (w, x, y, z)]. + Shape is (num_instances, num_bodies) with dtype wp.transformf. + body_mask: Body mask. If None, all bodies are updated. + Shape is (num_bodies,). + 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].to(torch.int32)) + else: + env_ids = self._ALL_ENV_INDICES + 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, + dim=(self._num_instances, self._num_bodies), + inputs=[coms, self._resolve_env_mask(env_mask), self._resolve_body_mask(body_mask)], + outputs=[self.data._body_com_pose_b.data], + device=self._device, + ) + # Invalidate derived buffers that depend on body_com_pose_b. + self.data._body_com_pose_w.timestamp = -1.0 + wp.copy(self.data._cpu_body_coms, self.data._body_com_pose_b.data) + self._binding_write( + TT.BODY_COM_POSE, + self.data._cpu_body_coms, + env_ids=self._get_cpu_env_ids(env_ids), + device="cpu", + data_dim=7, + ) + + def set_inertias_index( + self, + *, + inertias: wp.array, + body_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | wp.array | None = None, + ) -> None: + """Set body inertia tensors over selected env / body indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``BODY_INERTIA`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + Args: + inertias: Body inertia tensors [kg·m²]. Shape is + (len(env_ids), len(body_ids), 9) with dtype wp.float32. + The 9 components are the row-major flatten of the 3×3 inertia + matrix (Ixx, Ixy, Ixz, Iyx, Iyy, Iyz, Izx, Izy, Izz). + body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. + """ + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + self.assert_shape_and_dtype(inertias, (env_ids.shape[0], body_ids.shape[0], 9), wp.float32, "inertias") + wp.launch( + shared_kernels.write_body_inertia_to_buffer_index, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[inertias, env_ids, body_ids], + outputs=[self.data._body_inertia.data], + device=self._device, + ) + wp.copy(self.data._cpu_body_inertia, self.data._body_inertia.data) + self._binding_write( + TT.BODY_INERTIA, + self.data._cpu_body_inertia, + env_ids=self._get_cpu_env_ids(env_ids), + device="cpu", + data_dim=9, + ) + + def set_inertias_mask( + self, + *, + inertias: wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set body inertia tensors over selected env / body masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``BODY_INERTIA`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + Args: + inertias: Body inertia tensors [kg·m²]. Shape is + (num_instances, num_bodies, 9) with dtype wp.float32. + The 9 components are the row-major flatten of the 3×3 inertia + matrix (Ixx, Ixy, Ixz, Iyx, Iyy, Iyz, Izx, Izy, Izz). + body_mask: Body mask. If None, all bodies are updated. + Shape is (num_bodies,). + 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].to(torch.int32)) + else: + env_ids = self._ALL_ENV_INDICES + 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, + dim=(self._num_instances, self._num_bodies), + inputs=[inertias, self._resolve_env_mask(env_mask), self._resolve_body_mask(body_mask)], + outputs=[self.data._body_inertia.data], + device=self._device, + ) + wp.copy(self.data._cpu_body_inertia, self.data._body_inertia.data) + self._binding_write( + TT.BODY_INERTIA, + self.data._cpu_body_inertia, + env_ids=self._get_cpu_env_ids(env_ids), + device="cpu", + data_dim=9, + ) + + def _initialize_impl(self) -> None: + """Initialize the rigid object collection from the OVPhysX simulation backend. + + For each body in :attr:`cfg.rigid_objects`, validates the prim tree, + converts the IsaacLab prim path to an fnmatch glob, and eagerly creates + a single fused :class:`TensorBinding` per tensor type using the new + ``prim_paths=[...]`` API introduced in ovphysx 0.4.3. + + Then creates the :class:`RigidObjectCollectionData` container and primes + the asset-side buffers. + """ + physx_instance = OvPhysxManager.get_physx_instance() + if physx_instance is None: + raise RuntimeError("OvPhysxManager has not been initialized yet.") + self._ovphysx = physx_instance + self._device = OvPhysxManager.get_device() + + self._prim_paths: list[str] = [] + self._body_names_list: list[str] = [] + + for name, obj_cfg in self.cfg.rigid_objects.items(): + # Convert IsaacLab prim-path notation to the fnmatch-style glob that + # OVPhysX create_tensor_binding expects. Two conventions are in use: + # /World/envs/env_.*/object -- regex dot-star for any env index + # /World/envs/{ENV_REGEX_NS}/object -- explicit placeholder + pattern = re.sub(r"\{ENV_REGEX_NS\}", "*", obj_cfg.prim_path) + pattern = re.sub(r"\.\*", "*", pattern) + + # Validate the prim tree before creating tensor bindings. + # OVPhysX silently returns a zero-count binding when the pattern + # matches nothing; fail fast here with a clear message instead. + template_prim = sim_utils.find_first_matching_prim(obj_cfg.prim_path) + if template_prim is None: + raise RuntimeError(f"Failed to find prim for expression: '{obj_cfg.prim_path}' (body '{name}').") + template_prim_path = template_prim.GetPath().pathString + + root_prims = sim_utils.get_all_matching_child_prims( + template_prim_path, + predicate=lambda prim: prim.HasAPI(UsdPhysics.RigidBodyAPI), + traverse_instance_prims=False, + ) + if len(root_prims) == 0: + raise RuntimeError( + f"Failed to find a rigid body when resolving '{obj_cfg.prim_path}' (body '{name}')." + " Please ensure that the prim has 'USD RigidBodyAPI' applied." + ) + if len(root_prims) > 1: + raise RuntimeError( + f"Failed to find a single rigid body when resolving '{obj_cfg.prim_path}' (body '{name}')." + f" Found multiple '{root_prims}' under '{template_prim_path}'." + " Please ensure that there is only one rigid body in the prim path tree." + ) + + articulation_prims = sim_utils.get_all_matching_child_prims( + template_prim_path, + predicate=lambda prim: prim.HasAPI(UsdPhysics.ArticulationRootAPI), + traverse_instance_prims=False, + ) + if len(articulation_prims) != 0: + if articulation_prims[0].GetAttribute("physxArticulation:articulationEnabled").Get(): + raise RuntimeError( + f"Found an articulation root when resolving '{obj_cfg.prim_path}' (body '{name}') in the" + f" rigid object collection. These are located at: '{articulation_prims}' under" + f" '{template_prim_path}'. Please disable the articulation root in the USD or from code by" + " setting the parameter 'ArticulationRootPropertiesCfg.articulation_enabled' to False in the" + " spawn configuration." + ) + + # resolve root prim back into the regex expression + root_prim_path = root_prims[0].GetPath().pathString + suffix = root_prim_path[len(template_prim_path) :] + if suffix: + pattern = pattern + suffix + + self._prim_paths.append(pattern) + self._body_names_list.append(name) + + self._num_bodies = len(self._prim_paths) + + # ovphysx 0.4.3+ accepts ``prim_paths=[g0, ..., g_{B-1}]`` and returns a single + # binding spanning N*B prims with shape ``(N*B, D)`` in body-major order + # ``(body0_env0, body0_env1, ..., body1_env0, ...)``. Bindings are stored under + # the ``LINK_*``/``BODY_*`` data-class keys so the same key works with the + # articulation-mode mock used by iface tests. + _TT_MAP = ( + (TT.LINK_POSE, TT.RIGID_BODY_POSE), + (TT.LINK_VELOCITY, TT.RIGID_BODY_VELOCITY), + (TT.LINK_WRENCH, TT.RIGID_BODY_WRENCH), + (TT.BODY_MASS, TT.RIGID_BODY_MASS), + (TT.BODY_COM_POSE, TT.RIGID_BODY_COM_POSE), + (TT.BODY_INERTIA, TT.RIGID_BODY_INERTIA), + ) + for store_key, rb_tt in _TT_MAP: + try: + self._bindings[store_key] = self._ovphysx.create_tensor_binding( + prim_paths=self._prim_paths, tensor_type=rb_tt + ) + except Exception as e: + raise RuntimeError( + f"OVPhysX could not create fused RIGID_BODY binding {rb_tt!r} for" + f" prim_paths={self._prim_paths!r}." + f" Check that each prim path matches at least one" + f" UsdPhysics.RigidBodyAPI prim." + ) from e + + # Native fused binding has ``count == N * num_bodies`` (body-major flat). + pose_count = self._bindings[TT.LINK_POSE].count + if pose_count % self._num_bodies != 0: + raise RuntimeError( + f"Fused LINK_POSE binding count {pose_count} is not divisible by" + f" num_bodies {self._num_bodies}. prim_paths={self._prim_paths!r}." + ) + self._num_instances = pose_count // self._num_bodies + + self._data = RigidObjectCollectionData( + root_view=self._bindings, + num_bodies=self._num_bodies, + device=self._device, + ) + + self._create_buffers() + self._process_cfg() + self.update(0.0) + self._data.is_primed = True + + def _create_buffers(self) -> None: + """Pre-allocate asset-side index arrays and CPU staging buffers.""" + N = self._num_instances + B = self._num_bodies + + self._ALL_ENV_INDICES = wp.array(np.arange(N), dtype=wp.int32, device=self._device) + self._ALL_BODY_INDICES = wp.array(np.arange(B), dtype=wp.int32, device=self._device) + + # CPU copy of all-env indices used when calling CPU-only binding.write(). + self._cpu_all_env_ids = wp.zeros(N, dtype=wp.int32, device="cpu", pinned=True) + wp.copy(self._cpu_all_env_ids, self._ALL_ENV_INDICES) + + # All-true boolean masks used as defaults in mask-based kernel calls. + self._ALL_TRUE_ENV_MASK = wp.array(np.ones(N, dtype=bool), dtype=wp.bool, device=self._device) + self._ALL_TRUE_BODY_MASK = wp.array(np.ones(B, dtype=bool), dtype=wp.bool, device=self._device) + + # External wrench buffer: direct (N, B, 9) contiguous allocation. + # The fused LINK_WRENCH binding writes from a single (N, B, 9) buffer. + self._wrench_buf = wp.zeros((N, B, 9), dtype=wp.float32, device=self._device) + + self._instantaneous_wrench_composer = WrenchComposer(self) + self._permanent_wrench_composer = WrenchComposer(self) + + # set information about rigid body into data + self._data.body_names = self._body_names_list + + def _process_cfg(self) -> None: + """Post-processing of configuration parameters. + + Reads the per-body initial state from :attr:`cfg.rigid_objects` and + broadcasts it across all environment instances to produce + ``(num_instances, num_bodies, data_size)`` default-state arrays. + """ + default_body_poses = [] + default_body_vels = [] + + for obj_cfg in self.cfg.rigid_objects.values(): + default_body_pose = tuple(obj_cfg.init_state.pos) + tuple(obj_cfg.init_state.rot) + default_body_vel = tuple(obj_cfg.init_state.lin_vel) + tuple(obj_cfg.init_state.ang_vel) + # Broadcast across num_instances: (data_size,) -> (num_instances, data_size) + default_body_pose = np.tile(np.array(default_body_pose, dtype=np.float32), (self._num_instances, 1)) + default_body_vel = np.tile(np.array(default_body_vel, dtype=np.float32), (self._num_instances, 1)) + default_body_poses.append(default_body_pose) + default_body_vels.append(default_body_vel) + + # Stack per-body arrays: each (num_instances, data_size) -> (num_instances, num_bodies, data_size) + default_body_poses = np.stack(default_body_poses, axis=1) + default_body_vels = np.stack(default_body_vels, axis=1) + self._data.default_body_pose = wp.array(default_body_poses, dtype=wp.transformf, device=self._device) + self._data.default_body_vel = wp.array(default_body_vels, dtype=wp.spatial_vectorf, device=self._device) + + """ + Internal simulation callbacks. + """ + + def _invalidate_initialize_callback(self, event) -> None: + """Invalidates the scene elements.""" + # call parent + super()._invalidate_initialize_callback(event) + + """ + Helper functions. + """ + + def _get_binding(self, tensor_type: int): + """Return the cached fused :class:`TensorBinding` for *tensor_type*. + + All bindings are eagerly created in :meth:`_initialize_impl` and stored + under the ``TT.LINK_*`` / ``TT.BODY_*`` keys that + :class:`RigidObjectCollectionData` uses. + + Args: + tensor_type: The TensorType constant identifying which simulation + buffer to bind (e.g. :attr:`~isaaclab_ovphysx.tensor_types.LINK_POSE`). + + Returns: + The cached :class:`TensorBinding`, or ``None`` if not found. + """ + return self._bindings.get(tensor_type) + + def reshape_data_to_view_2d(self, data: wp.array, device: str | None = None) -> wp.array: + """Reshape instance-major ``(num_instances, num_bodies)`` data to body-major view order. + + The native fused multi-prim binding lays data out as + ``(body0_env0, body0_env1, ..., body1_env0, body1_env1, ...)`` with shape + ``(num_bodies * num_instances,)``. This helper builds a strided view of the + instance-major buffer with the transposed layout and clones it into a + contiguous body-major flat array. + + Args: + data: Source buffer with shape ``(num_instances, num_bodies)`` (any single-element dtype). + device: Optional target device for the cloned output. Defaults to ``data.device``. + + Returns: + Contiguous body-major flat buffer with shape ``(num_bodies * num_instances,)``. + """ + if device is None: + device = str(data.device) + element_size = wp.types.type_size_in_bytes(data.dtype) + strided_view = wp.array( + ptr=data.ptr, + shape=(self.num_bodies, self.num_instances), + dtype=data.dtype, + strides=(element_size, self.num_bodies * element_size), + device=str(data.device), + ) + return wp.clone(strided_view, device=device).reshape((self.num_bodies * self.num_instances,)) + + def reshape_data_to_view_3d(self, data: wp.array, data_dim: int, device: str | None = None) -> wp.array: + """Reshape instance-major ``(num_instances, num_bodies, data_dim)`` data to body-major view order. + + Companion of :meth:`reshape_data_to_view_2d` for 3D buffers (e.g. inertia + tensors). Output shape is ``(num_bodies * num_instances, data_dim)``. + + Args: + data: Source buffer with shape ``(num_instances, num_bodies, data_dim)``. + data_dim: Trailing per-element dimension size. + device: Optional target device for the cloned output. Defaults to ``data.device``. + + Returns: + Contiguous body-major buffer with shape ``(num_bodies * num_instances, data_dim)``. + """ + if device is None: + device = str(data.device) + element_size = wp.types.type_size_in_bytes(data.dtype) + row_size = element_size * data_dim + strided_view = wp.array( + ptr=data.ptr, + shape=(self.num_bodies, self.num_instances, data_dim), + dtype=data.dtype, + strides=(row_size, self.num_bodies * row_size, element_size), + device=str(data.device), + ) + return wp.clone(strided_view, device=device).reshape((self.num_bodies * self.num_instances, data_dim)) + + def _binding_write( + self, + tensor_type: int, + instance_major_data: wp.array, + env_ids: wp.array, + device: str | None = None, + data_dim: int | None = None, + ) -> None: + """Write an instance-major buffer through a fused binding. + + Dispatches to one of two paths depending on the binding's layout: + + * **Native fused binding** (``count == num_instances * num_bodies``, + body-major flat layout): the instance-major buffer is reshaped via + :meth:`reshape_data_to_view_2d` / :meth:`reshape_data_to_view_3d` to a + contiguous body-major view, then written with body-major view indices + ``view_id = body_id * num_instances + env_id``. + * **Articulation-mode mock** (``count == num_instances``, instance-major + ``(N, B[, D])`` shape): the buffer is written directly with the + environment indices, matching the existing mock contract. + + Args: + tensor_type: TensorType key identifying the cached binding. + instance_major_data: Instance-major buffer of shape ``(N, B)`` or + ``(N, B, data_dim)``. May use ``wp.float32`` or a structured dtype. + env_ids: Environment indices (1D ``wp.int32`` on ``self._device`` or + ``"cpu"`` for CPU-only bindings). + device: Destination device for the body-major clone (only used on the + fused-binding path). Defaults to ``self._device``. + data_dim: When provided, treat the buffer as 3D and use + :meth:`reshape_data_to_view_3d`. When ``None`` (default), use + :meth:`reshape_data_to_view_2d`. + """ + binding = self._get_binding(tensor_type) + if binding is None: + return + if device is None: + device = self._device + # Disambiguate via the binding's exposed shape: the articulation-mode + # mock returns a directly instance-major view ``(N, B[, D])`` while the + # native fused multi-prim binding lays elements body-major-flat with + # ``shape == (N * B[, D])``. + is_mock_layout = len(binding.shape) >= 2 and binding.shape[1] == self._num_bodies + if is_mock_layout: + float32_data = ( + instance_major_data if instance_major_data.dtype == wp.float32 else instance_major_data.view(wp.float32) + ) + binding.write(float32_data, indices=env_ids) + return + # Native fused path: body-major flat (N*B[, D]); reshape and use view_ids. + if data_dim is None: + view = self.reshape_data_to_view_2d(instance_major_data, device=device).view(wp.float32) + else: + view = self.reshape_data_to_view_3d(instance_major_data, data_dim, device=device) + view_ids = self._env_body_ids_to_view_ids(env_ids, self._ALL_BODY_INDICES, device=device) + binding.write(view, indices=view_ids) + + """ + Internal helper. + """ + + def _resolve_env_ids(self, env_ids) -> wp.array: + """Resolve environment indices to a warp int32 array on ``self._device``. + + Tests sometimes hand us indices on CPU even when the sim runs on GPU; we move the + resolved array onto ``self._device`` so kernel launches don't fail on a device + mismatch. + """ + if env_ids is None or env_ids == slice(None): + return self._ALL_ENV_INDICES + if isinstance(env_ids, list): + return wp.array(env_ids, dtype=wp.int32, device=self._device) + if isinstance(env_ids, torch.Tensor): + return wp.from_torch(env_ids.to(torch.int32), dtype=wp.int32) + 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 + + def _resolve_body_ids(self, body_ids) -> wp.array: + """Resolve body indices to a warp int32 array on ``self._device``.""" + 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 + + def _env_body_ids_to_view_ids( + self, env_ids: torch.Tensor | wp.array, body_ids: torch.Tensor | wp.array, device: str = "cuda:0" + ) -> wp.array: + """Convert environment and body indices to flat view indices (body-major ordering). + + 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. + + Args: + env_ids: Environment indices. + body_ids: Body indices. + device: Target device for the returned array. + + Returns: + A :class:`wp.array` of shape ``(len(env_ids) * len(body_ids),)`` with + flat view indices on *device*. + """ + if isinstance(env_ids, torch.Tensor): + env_ids = wp.from_torch(env_ids.to(torch.int32), dtype=wp.int32) + if isinstance(body_ids, torch.Tensor): + body_ids = wp.from_torch(body_ids.to(torch.int32), dtype=wp.int32) + if str(env_ids.device) != device: + env_ids = wp.clone(env_ids, device=device) + if str(body_ids.device) != device: + body_ids = wp.clone(body_ids, device=device) + num_query_envs = env_ids.shape[0] + view_ids = wp.zeros(num_query_envs * body_ids.shape[0], dtype=wp.int32, device=device) + wp.launch( + resolve_view_ids, + dim=(num_query_envs, body_ids.shape[0]), + inputs=[env_ids, body_ids, num_query_envs, self.num_instances], + outputs=[view_ids], + device=device, + ) + return view_ids + + def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array: + """Resolve an environment mask to a ``wp.bool`` array on ``self._device``. + + ``None`` returns the pre-allocated all-true mask. + + Args: + env_mask: Boolean environment mask or None. Shape is (num_instances,). + + 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 + + def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array: + """Resolve a body mask to a ``wp.bool`` array on ``self._device``. + + ``None`` returns the pre-allocated all-true mask. + + Args: + body_mask: Boolean body mask or None. Shape is (num_bodies,). + + 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 + + def _get_cpu_env_ids(self, env_ids: wp.array) -> wp.array: + """Return CPU int32 env indices for CPU-only binding writes. + + Uses the pre-allocated pinned ``_cpu_all_env_ids`` fast path when + *env_ids* covers all instances, otherwise clones to CPU. + + Args: + env_ids: A warp int32 array of environment indices on any device. + + Returns: + A warp int32 array guaranteed to live on ``"cpu"``. + """ + if env_ids.ptr == self._ALL_ENV_INDICES.ptr: + return self._cpu_all_env_ids + return wp.clone(env_ids, device="cpu") + + """ + Deprecated properties and methods. + """ + + def write_body_state_to_sim( + self, + body_states: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + body_ids: slice | torch.Tensor | None = None, + ) -> None: + """Deprecated, same as :meth:`write_body_link_pose_to_sim_index` and + :meth:`write_body_com_velocity_to_sim_index`.""" + warnings.warn( + "The function 'write_body_state_to_sim' will be deprecated in a future release. Please" + " use 'write_body_link_pose_to_sim_index' and 'write_body_com_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + if isinstance(body_states, wp.array): + body_states = wp.to_torch(body_states) + self.write_body_link_pose_to_sim_index(body_poses=body_states[:, :, :7], env_ids=env_ids, body_ids=body_ids) + self.write_body_com_velocity_to_sim_index( + body_velocities=body_states[:, :, 7:], env_ids=env_ids, body_ids=body_ids + ) + + def write_body_com_state_to_sim( + self, + body_states: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + body_ids: slice | torch.Tensor | None = None, + ) -> None: + """Deprecated, same as :meth:`write_body_com_pose_to_sim_index` and + :meth:`write_body_com_velocity_to_sim_index`.""" + warnings.warn( + "The function 'write_body_com_state_to_sim' will be deprecated in a future release. Please" + " use 'write_body_com_pose_to_sim_index' and 'write_body_com_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + if isinstance(body_states, wp.array): + body_states = wp.to_torch(body_states) + self.write_body_com_pose_to_sim_index(body_poses=body_states[:, :, :7], env_ids=env_ids, body_ids=body_ids) + self.write_body_com_velocity_to_sim_index( + body_velocities=body_states[:, :, 7:], env_ids=env_ids, body_ids=body_ids + ) + + def write_body_link_state_to_sim( + self, + body_states: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + body_ids: slice | torch.Tensor | None = None, + ) -> None: + """Deprecated, same as :meth:`write_body_link_pose_to_sim_index` and + :meth:`write_body_link_velocity_to_sim_index`.""" + warnings.warn( + "The function 'write_body_link_state_to_sim' will be deprecated in a future release. Please" + " use 'write_body_link_pose_to_sim_index' and 'write_body_link_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + if isinstance(body_states, wp.array): + body_states = wp.to_torch(body_states) + self.write_body_link_pose_to_sim_index(body_poses=body_states[:, :, :7], env_ids=env_ids, body_ids=body_ids) + self.write_body_link_velocity_to_sim_index( + body_velocities=body_states[:, :, 7:], env_ids=env_ids, body_ids=body_ids + ) diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/rigid_object_collection_data.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/rigid_object_collection_data.py new file mode 100644 index 000000000000..8bce18404f89 --- /dev/null +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/rigid_object_collection_data.py @@ -0,0 +1,1132 @@ +# 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 + +from __future__ import annotations + +import warnings +from typing import Any + +import numpy as np +import torch +import warp as wp + +from isaaclab.assets.rigid_object_collection.base_rigid_object_collection_data import BaseRigidObjectCollectionData +from isaaclab.utils.buffers import TimestampedBufferWarp as TimestampedBuffer +from isaaclab.utils.math import normalize +from isaaclab.utils.warp import ProxyArray + +from isaaclab_ovphysx import tensor_types as TT +from isaaclab_ovphysx.assets import kernels as shared_kernels +from isaaclab_ovphysx.physics import OvPhysxManager as SimulationManager + + +class RigidObjectCollectionData(BaseRigidObjectCollectionData): + """Data container for a rigid object collection. + + This class contains the data for a rigid object collection in the simulation. The data includes the state of + all the bodies in the collection. The data is stored in the simulation world frame unless otherwise specified. + The data is in the order ``(num_instances, num_objects, data_size)``, where data_size is the size of the data. + + For a rigid body, there are two frames of reference that are used: + + - Actor frame: The frame of reference of the rigid body prim. This typically corresponds to the Xform prim + with the rigid body schema. + - Center of mass frame: The frame of reference of the center of mass of the rigid body. + + Depending on the settings of the simulation, the actor frame and the center of mass frame may be the same. + This needs to be taken into account when interpreting the data. + + The data is lazily updated, meaning that the data is only updated when it is accessed. This is useful + when the data is expensive to compute or retrieve. The data is updated when the timestamp of the buffer + is older than the current simulation timestamp. The timestamp is updated whenever the data is updated. + + .. note:: + **Pull-to-refresh model.** Properties pull fresh data from the OVPhysX tensor API on first access + per timestamp and cache the result. This differs from Newton, where buffers are refreshed + automatically by the simulation. + + .. note:: + **ProxyArray pointer stability.** Each :class:`ProxyArray` wrapper is created once and reused + because the OVPhysX tensor API returns views into stable, pre-allocated GPU buffers whose device + pointer does not change across simulation steps. + """ + + __backend_name__: str = "ovphysx" + """The name of the backend for the rigid object collection data.""" + + def __init__( + self, + root_view: dict[int, Any], + num_bodies: int, + device: str, + ): + """Initializes the rigid object data. + + Args: + root_view: Fused TensorBinding dict, keyed by TensorType constant. Each value is a single + :class:`TensorBinding` spanning all bodies in the collection. + num_bodies: The number of bodies in the collection. + device: The device used for processing. + """ + super().__init__(root_view, num_bodies, device) + # Store the bindings dict (the equivalent of the root view in PhysX). + self._bindings = root_view + self._binding_getter = None # may be set externally after construction + self.num_bodies = num_bodies + self._num_bodies = num_bodies + # Set initial time stamp + self._sim_timestamp = 0.0 + self._is_primed = False + # Body-major read scratch buffers (keyed by tensor_type). Allocated on the binding's own + # device — pinned host for CPU-only bindings, GPU for the rest — so ``binding.read(scratch)`` + # never crosses devices. + self._cpu_staging_buffers: dict[int, wp.array] = {} + + # Read num_instances from the LINK_POSE binding. The native fused multi-prim binding lays + # elements out body-major-flat with ``shape == (N * B, 7)`` and ``count == N * B``. The + # articulation-mode mock used by iface tests exposes an instance-major view directly with + # ``shape == (N, B, 7)`` and ``count == N``. Dispatch via the binding's exposed shape. + pose_binding = self._bindings[TT.LINK_POSE] + if len(pose_binding.shape) >= 2 and pose_binding.shape[1] == num_bodies: + self.num_instances = pose_binding.count + else: + self.num_instances = pose_binding.count // num_bodies + self._num_instances = self.num_instances + + if SimulationManager._sim is not None and hasattr(SimulationManager._sim, "cfg"): + gravity = SimulationManager._sim.cfg.gravity + else: + gravity = (0.0, 0.0, -9.81) + + gravity_dir = torch.tensor((gravity[0], gravity[1], gravity[2]), device=self.device) + if torch.linalg.norm(gravity_dir) > 0.0: + gravity_dir = normalize(gravity_dir.unsqueeze(0)).squeeze(0) + gravity_dir = gravity_dir.repeat(self.num_instances, self.num_bodies, 1) + forward_vec = torch.tensor((1.0, 0.0, 0.0), device=self.device).repeat(self.num_instances, self.num_bodies, 1) + + # Initialize constants + self.GRAVITY_VEC_W = ProxyArray(wp.from_torch(gravity_dir, dtype=wp.vec3f)) + self.FORWARD_VEC_B = ProxyArray(wp.from_torch(forward_vec, dtype=wp.vec3f)) + + self._create_buffers() + + @property + def is_primed(self) -> bool: + """Whether the rigid object collection data is fully instantiated and ready to use.""" + return self._is_primed + + @is_primed.setter + def is_primed(self, value: bool) -> None: + """Set whether the rigid object collection data is fully instantiated and ready to use. + + .. note:: + Once this quantity is set to True, it cannot be changed. + + Args: + value: The primed state. + + Raises: + ValueError: If the rigid object collection data is already primed. + """ + if self._is_primed: + raise ValueError("The rigid object collection data is already primed.") + self._is_primed = value + + def update(self, dt: float) -> None: + """Updates the data for the rigid object collection. + + Args: + dt: The time step for the update [s]. This must be a positive value. + """ + # update the simulation timestamp + self._sim_timestamp += dt + # Prime the FD-dependent COM acceleration so the first read returns a sensible (zero) value. + _ = self.body_com_acc_w + + """ + Names. + """ + + body_names: list[str] = None + """Body names in the order parsed by the simulation view.""" + + """ + Defaults. + """ + + @property + def default_body_pose(self) -> ProxyArray: + """Default body pose ``[pos, quat]`` in local environment frame. + + Shape is (num_instances, num_bodies), dtype = ``wp.transformf``. + In torch this resolves to (num_instances, num_bodies, 7). + Set by :meth:`~RigidObjectCollection._process_cfg` during initialization. + """ + if self._default_body_pose_ta is None: + self._default_body_pose_ta = ProxyArray(self._default_body_pose) + return self._default_body_pose_ta + + @default_body_pose.setter + def default_body_pose(self, value: wp.array) -> None: + self._default_body_pose.assign(value) + + @property + def default_body_vel(self) -> ProxyArray: + """Default body velocity ``[lin_vel, ang_vel]`` in local environment frame. + + Shape is (num_instances, num_bodies), dtype = ``wp.spatial_vectorf``. + In torch this resolves to (num_instances, num_bodies, 6). + Set by :meth:`~RigidObjectCollection._process_cfg` during initialization. + """ + if self._default_body_vel_ta is None: + self._default_body_vel_ta = ProxyArray(self._default_body_vel) + return self._default_body_vel_ta + + @default_body_vel.setter + def default_body_vel(self, value: wp.array) -> None: + self._default_body_vel.assign(value) + + @property + def default_body_state(self) -> ProxyArray: + """Default root state ``[pos, quat, lin_vel, ang_vel]`` in local environment frame. + + Deprecated. Use :attr:`default_body_pose` and :attr:`default_body_vel` instead. + + Shape is (num_instances, num_bodies), dtype = ``vec13f``. + In torch this resolves to (num_instances, num_bodies, 13). + """ + warnings.warn( + "Reading the body state directly is deprecated since IsaacLab 3.0 and will be removed in a future version. " + "Please use the default_body_pose and default_body_vel properties instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._default_body_state is None: + self._default_body_state = wp.zeros( + (self.num_instances, self.num_bodies), dtype=shared_kernels.vec13f, device=self.device + ) + self._default_body_state_ta = ProxyArray(self._default_body_state) + wp.launch( + shared_kernels.concat_body_pose_and_vel_to_state, + dim=(self.num_instances, self.num_bodies), + inputs=[ + self._default_body_pose, + self._default_body_vel, + ], + outputs=[ + self._default_body_state, + ], + device=self.device, + ) + return self._default_body_state_ta + + """ + Body state properties. + """ + + @property + def body_link_pose_w(self) -> ProxyArray: + """Body link pose ``[pos, quat]`` in simulation world frame [m, -]. + + Shape is (num_instances, num_bodies), dtype = ``wp.transformf``. + In torch this resolves to (num_instances, num_bodies, 7). + This quantity is the pose of the actor frame of the rigid body relative to + the world. The orientation is provided in (x, y, z, w) format. + """ + if self._body_link_pose_w.timestamp < self._sim_timestamp: + self._read_transform_binding(TT.LINK_POSE, self._body_link_pose_w) + # Invalidate sliced sub-component proxies so they are rebuilt from the + # updated buffer on next access. + self._body_link_pos_w_ta = None + self._body_link_quat_w_ta = None + if self._body_link_pose_w_ta is None: + self._body_link_pose_w_ta = ProxyArray(self._body_link_pose_w.data) + return self._body_link_pose_w_ta + + @property + def body_link_vel_w(self) -> ProxyArray: + """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame [m/s, rad/s]. + + Shape is (num_instances, num_bodies), dtype = ``wp.spatial_vectorf``. + In torch this resolves to (num_instances, num_bodies, 6). + This quantity contains the linear and angular velocities of the actor frame + of the rigid body relative to the world. + """ + if self._body_link_vel_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.get_body_link_vel_from_body_com_vel, + dim=(self.num_instances, self.num_bodies), + inputs=[ + self.body_com_vel_w, + self.body_link_pose_w, + self.body_com_pose_b, + ], + outputs=[self._body_link_vel_w.data], + device=self.device, + ) + self._body_link_vel_w.timestamp = self._sim_timestamp + self._body_link_lin_vel_w_ta = None + self._body_link_ang_vel_w_ta = None + if self._body_link_vel_w_ta is None: + self._body_link_vel_w_ta = ProxyArray(self._body_link_vel_w.data) + return self._body_link_vel_w_ta + + @property + def body_com_pose_w(self) -> ProxyArray: + """Body center of mass pose ``[pos, quat]`` in simulation world frame [m, -]. + + Shape is (num_instances, num_bodies), dtype = ``wp.transformf``. + In torch this resolves to (num_instances, num_bodies, 7). + This quantity is the pose of the center of mass frame of the rigid body + relative to the world. The orientation is provided in (x, y, z, w) format. + """ + if self._body_com_pose_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.get_body_com_pose_from_body_link_pose, + dim=(self.num_instances, self.num_bodies), + inputs=[ + self.body_link_pose_w, + self.body_com_pose_b, + ], + outputs=[self._body_com_pose_w.data], + device=self.device, + ) + self._body_com_pose_w.timestamp = self._sim_timestamp + self._body_com_pos_w_ta = None + self._body_com_quat_w_ta = None + if self._body_com_pose_w_ta is None: + self._body_com_pose_w_ta = ProxyArray(self._body_com_pose_w.data) + return self._body_com_pose_w_ta + + @property + def body_com_vel_w(self) -> ProxyArray: + """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame [m/s, rad/s]. + + Shape is (num_instances, num_bodies), dtype = ``wp.spatial_vectorf``. + In torch this resolves to (num_instances, num_bodies, 6). + This quantity contains the linear and angular velocities of the rigid body's + center of mass frame relative to the world. + """ + if self._body_com_vel_w.timestamp < self._sim_timestamp: + self._read_spatial_vector_binding(TT.LINK_VELOCITY, self._body_com_vel_w) + self._body_com_lin_vel_w_ta = None + self._body_com_ang_vel_w_ta = None + if self._body_com_vel_w_ta is None: + self._body_com_vel_w_ta = ProxyArray(self._body_com_vel_w.data) + return self._body_com_vel_w_ta + + @property + def body_com_acc_w(self) -> ProxyArray: + """Acceleration of all bodies ``[lin_acc, ang_acc]`` in the simulation world frame [m/s², rad/s²]. + + Shape is (num_instances, num_bodies), dtype = ``wp.spatial_vectorf``. + In torch this resolves to (num_instances, num_bodies, 6). + This quantity is the acceleration of the rigid bodies' center of mass frame relative + to the world, derived by finite differencing consecutive COM velocities. + """ + if self._body_com_acc_w.timestamp < self._sim_timestamp: + if self._previous_body_com_vel is None: + self._previous_body_com_vel = wp.clone(self.body_com_vel_w.warp) + wp.launch( + shared_kernels.derive_body_acceleration_from_body_com_velocities, + dim=(self.num_instances, self.num_bodies), + device=self.device, + inputs=[ + self.body_com_vel_w.warp, + SimulationManager.get_physics_dt(), + self._previous_body_com_vel, + ], + outputs=[ + self._body_com_acc_w.data, + ], + ) + self._body_com_acc_w.timestamp = self._sim_timestamp + self._body_com_lin_acc_w_ta = None + self._body_com_ang_acc_w_ta = None + if self._body_com_acc_w_ta is None: + self._body_com_acc_w_ta = ProxyArray(self._body_com_acc_w.data) + return self._body_com_acc_w_ta + + @property + def body_com_pose_b(self) -> ProxyArray: + """Center of mass pose ``[pos, quat]`` of all bodies in their respective body link frames [m, -]. + + Shape is (num_instances, num_bodies), dtype = ``wp.transformf``. + In torch this resolves to (num_instances, num_bodies, 7). + This quantity is the pose of the center of mass frame of the rigid body + relative to the body's link frame. The orientation is provided in + (x, y, z, w) format. + """ + if self._body_com_pose_b.timestamp < self._sim_timestamp: + self._read_transform_binding(TT.BODY_COM_POSE, self._body_com_pose_b) + self._body_com_pos_b_ta = None + self._body_com_quat_b_ta = None + if self._body_com_pose_b_ta is None: + self._body_com_pose_b_ta = ProxyArray(self._body_com_pose_b.data) + return self._body_com_pose_b_ta + + @property + def body_mass(self) -> ProxyArray: + """Mass of all bodies [kg]. + + Shape is (num_instances, num_bodies), dtype = ``wp.float32``. + In torch this resolves to (num_instances, num_bodies). + """ + if self._body_mass_ta is None: + self._body_mass_ta = ProxyArray(self._body_mass.data) + return self._body_mass_ta + + @property + def body_inertia(self) -> ProxyArray: + """Inertia tensor of all bodies, expressed at the center of mass [kg·m²]. + + Shape is (num_instances, num_bodies, 9), dtype = ``wp.float32``. + The 9 components are the row-major flatten of the 3×3 inertia matrix + ``(Ixx, Ixy, Ixz, Iyx, Iyy, Iyz, Izx, Izy, Izz)``. + In torch this resolves to (num_instances, num_bodies, 9). + """ + if self._body_inertia_ta is None: + self._body_inertia_ta = ProxyArray(self._body_inertia.data) + return self._body_inertia_ta + + """ + Deprecated state-concat properties. + """ + + @property + def body_state_w(self) -> ProxyArray: + """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_com_vel_w`.""" + warnings.warn( + "The `body_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " + "`body_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_body_pose_and_vel_to_state, + dim=(self.num_instances, self.num_bodies), + inputs=[self.body_link_pose_w, self.body_com_vel_w], + outputs=[self._body_state_w.data], + device=self.device, + ) + self._body_state_w.timestamp = self._sim_timestamp + if self._body_state_w_ta is None: + self._body_state_w_ta = ProxyArray(self._body_state_w.data) + return self._body_state_w_ta + + @property + def body_link_state_w(self) -> ProxyArray: + """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_link_vel_w`.""" + warnings.warn( + "The `body_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " + "`body_link_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_link_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_body_pose_and_vel_to_state, + dim=(self.num_instances, self.num_bodies), + inputs=[self.body_link_pose_w, self.body_link_vel_w], + outputs=[self._body_link_state_w.data], + device=self.device, + ) + self._body_link_state_w.timestamp = self._sim_timestamp + if self._body_link_state_w_ta is None: + self._body_link_state_w_ta = ProxyArray(self._body_link_state_w.data) + return self._body_link_state_w_ta + + @property + def body_com_state_w(self) -> ProxyArray: + """Deprecated, same as :attr:`body_com_pose_w` and :attr:`body_com_vel_w`.""" + warnings.warn( + "The `body_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_com_pose_w` and " + "`body_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_com_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_body_pose_and_vel_to_state, + dim=(self.num_instances, self.num_bodies), + inputs=[self.body_com_pose_w, self.body_com_vel_w], + outputs=[self._body_com_state_w.data], + device=self.device, + ) + self._body_com_state_w.timestamp = self._sim_timestamp + if self._body_com_state_w_ta is None: + self._body_com_state_w_ta = ProxyArray(self._body_com_state_w.data) + return self._body_com_state_w_ta + + """ + Sliced properties. + """ + + @property + def body_link_pos_w(self) -> ProxyArray: + """Positions of all bodies in simulation world frame [m]. + + Shape is (num_instances, num_bodies), dtype = ``wp.vec3f``. + In torch this resolves to (num_instances, num_bodies, 3). + This quantity is the position of the rigid bodies' actor frame relative to + the world. + """ + parent = self.body_link_pose_w + if self._body_link_pos_w_ta is None: + self._body_link_pos_w_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._body_link_pos_w_ta + + @property + def body_link_quat_w(self) -> ProxyArray: + """Orientation (x, y, z, w) of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = ``wp.quatf``. + In torch this resolves to (num_instances, num_bodies, 4). + This quantity is the orientation of the rigid bodies' actor frame relative + to the world. + """ + parent = self.body_link_pose_w + if self._body_link_quat_w_ta is None: + self._body_link_quat_w_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._body_link_quat_w_ta + + @property + def body_link_lin_vel_w(self) -> ProxyArray: + """Linear velocity of all bodies in simulation world frame [m/s]. + + Shape is (num_instances, num_bodies), dtype = ``wp.vec3f``. + In torch this resolves to (num_instances, num_bodies, 3). + This quantity is the linear velocity of the rigid bodies' actor frame + relative to the world. + """ + parent = self.body_link_vel_w + if self._body_link_lin_vel_w_ta is None: + self._body_link_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._body_link_lin_vel_w_ta + + @property + def body_link_ang_vel_w(self) -> ProxyArray: + """Angular velocity of all bodies in simulation world frame [rad/s]. + + Shape is (num_instances, num_bodies), dtype = ``wp.vec3f``. + In torch this resolves to (num_instances, num_bodies, 3). + This quantity is the angular velocity of the rigid bodies' actor frame + relative to the world. + """ + parent = self.body_link_vel_w + if self._body_link_ang_vel_w_ta is None: + self._body_link_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._body_link_ang_vel_w_ta + + @property + def body_link_lin_vel_b(self) -> ProxyArray: + """Linear velocity of all bodies in their respective body (actor) frames [m/s]. + + Shape is (num_instances, num_bodies), dtype = ``wp.vec3f``. + In torch this resolves to (num_instances, num_bodies, 3). + """ + if self._body_link_lin_vel_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_2D_kernel, + dim=(self.num_instances, self.num_bodies), + inputs=[self.body_link_lin_vel_w, self.body_link_quat_w], + outputs=[self._body_link_lin_vel_b.data], + device=self.device, + ) + self._body_link_lin_vel_b.timestamp = self._sim_timestamp + if self._body_link_lin_vel_b_ta is None: + self._body_link_lin_vel_b_ta = ProxyArray(self._body_link_lin_vel_b.data) + return self._body_link_lin_vel_b_ta + + @property + def body_link_ang_vel_b(self) -> ProxyArray: + """Angular velocity of all bodies in their respective body (actor) frames [rad/s]. + + Shape is (num_instances, num_bodies), dtype = ``wp.vec3f``. + In torch this resolves to (num_instances, num_bodies, 3). + """ + if self._body_link_ang_vel_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_2D_kernel, + dim=(self.num_instances, self.num_bodies), + inputs=[self.body_link_ang_vel_w, self.body_link_quat_w], + outputs=[self._body_link_ang_vel_b.data], + device=self.device, + ) + self._body_link_ang_vel_b.timestamp = self._sim_timestamp + if self._body_link_ang_vel_b_ta is None: + self._body_link_ang_vel_b_ta = ProxyArray(self._body_link_ang_vel_b.data) + return self._body_link_ang_vel_b_ta + + @property + def body_com_pos_w(self) -> ProxyArray: + """Positions of all bodies' center of mass in simulation world frame [m]. + + Shape is (num_instances, num_bodies), dtype = ``wp.vec3f``. + In torch this resolves to (num_instances, num_bodies, 3). + """ + parent = self.body_com_pose_w + if self._body_com_pos_w_ta is None: + self._body_com_pos_w_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._body_com_pos_w_ta + + @property + def body_com_quat_w(self) -> ProxyArray: + """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = ``wp.quatf``. + In torch this resolves to (num_instances, num_bodies, 4). + """ + parent = self.body_com_pose_w + if self._body_com_quat_w_ta is None: + self._body_com_quat_w_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._body_com_quat_w_ta + + @property + def body_com_lin_vel_w(self) -> ProxyArray: + """Linear velocity of all bodies' center of mass in simulation world frame [m/s]. + + Shape is (num_instances, num_bodies), dtype = ``wp.vec3f``. + In torch this resolves to (num_instances, num_bodies, 3). + """ + parent = self.body_com_vel_w + if self._body_com_lin_vel_w_ta is None: + self._body_com_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._body_com_lin_vel_w_ta + + @property + def body_com_ang_vel_w(self) -> ProxyArray: + """Angular velocity of all bodies' center of mass in simulation world frame [rad/s]. + + Shape is (num_instances, num_bodies), dtype = ``wp.vec3f``. + In torch this resolves to (num_instances, num_bodies, 3). + """ + parent = self.body_com_vel_w + if self._body_com_ang_vel_w_ta is None: + self._body_com_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._body_com_ang_vel_w_ta + + @property + def body_com_lin_vel_b(self) -> ProxyArray: + """Linear velocity of all bodies' center of mass in their respective body (actor) frames [m/s]. + + Shape is (num_instances, num_bodies), dtype = ``wp.vec3f``. + In torch this resolves to (num_instances, num_bodies, 3). + """ + if self._body_com_lin_vel_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_2D_kernel, + dim=(self.num_instances, self.num_bodies), + inputs=[self.body_com_lin_vel_w, self.body_link_quat_w], + outputs=[self._body_com_lin_vel_b.data], + device=self.device, + ) + self._body_com_lin_vel_b.timestamp = self._sim_timestamp + if self._body_com_lin_vel_b_ta is None: + self._body_com_lin_vel_b_ta = ProxyArray(self._body_com_lin_vel_b.data) + return self._body_com_lin_vel_b_ta + + @property + def body_com_ang_vel_b(self) -> ProxyArray: + """Angular velocity of all bodies' center of mass in their respective body (actor) frames [rad/s]. + + Shape is (num_instances, num_bodies), dtype = ``wp.vec3f``. + In torch this resolves to (num_instances, num_bodies, 3). + """ + if self._body_com_ang_vel_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_2D_kernel, + dim=(self.num_instances, self.num_bodies), + inputs=[self.body_com_ang_vel_w, self.body_link_quat_w], + outputs=[self._body_com_ang_vel_b.data], + device=self.device, + ) + self._body_com_ang_vel_b.timestamp = self._sim_timestamp + if self._body_com_ang_vel_b_ta is None: + self._body_com_ang_vel_b_ta = ProxyArray(self._body_com_ang_vel_b.data) + return self._body_com_ang_vel_b_ta + + @property + def body_com_lin_acc_w(self) -> ProxyArray: + """Linear acceleration of all bodies' center of mass in simulation world frame [m/s²]. + + Shape is (num_instances, num_bodies), dtype = ``wp.vec3f``. + In torch this resolves to (num_instances, num_bodies, 3). + """ + parent = self.body_com_acc_w + if self._body_com_lin_acc_w_ta is None: + self._body_com_lin_acc_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._body_com_lin_acc_w_ta + + @property + def body_com_ang_acc_w(self) -> ProxyArray: + """Angular acceleration of all bodies' center of mass in simulation world frame [rad/s²]. + + Shape is (num_instances, num_bodies), dtype = ``wp.vec3f``. + In torch this resolves to (num_instances, num_bodies, 3). + """ + parent = self.body_com_acc_w + if self._body_com_ang_acc_w_ta is None: + self._body_com_ang_acc_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._body_com_ang_acc_w_ta + + @property + def body_com_pos_b(self) -> ProxyArray: + """Center of mass position of all of the bodies in their respective link frames [m]. + + Shape is (num_instances, num_bodies), dtype = ``wp.vec3f``. + In torch this resolves to (num_instances, num_bodies, 3). + """ + parent = self.body_com_pose_b + if self._body_com_pos_b_ta is None: + self._body_com_pos_b_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._body_com_pos_b_ta + + @property + def body_com_quat_b(self) -> ProxyArray: + """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies + in their respective link frames. + + Shape is (num_instances, num_bodies), dtype = ``wp.quatf``. + In torch this resolves to (num_instances, num_bodies, 4). + """ + parent = self.body_com_pose_b + if self._body_com_quat_b_ta is None: + self._body_com_quat_b_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._body_com_quat_b_ta + + """ + Derived Properties. + """ + + @property + def projected_gravity_b(self) -> ProxyArray: + """Projection of the gravity direction onto each body frame [-]. + + Shape is (num_instances, num_bodies), dtype = ``wp.vec3f``. + In torch this resolves to (num_instances, num_bodies, 3). + """ + if self._projected_gravity_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_2D_kernel, + dim=(self.num_instances, self.num_bodies), + inputs=[self.GRAVITY_VEC_W, self.body_link_quat_w], + outputs=[self._projected_gravity_b.data], + device=self.device, + ) + self._projected_gravity_b.timestamp = self._sim_timestamp + if self._projected_gravity_b_ta is None: + self._projected_gravity_b_ta = ProxyArray(self._projected_gravity_b.data) + return self._projected_gravity_b_ta + + @property + def heading_w(self) -> ProxyArray: + """Yaw heading of each body frame [rad]. + + Shape is (num_instances, num_bodies), dtype = ``wp.float32``. + In torch this resolves to (num_instances, num_bodies). + + .. note:: + This quantity is computed by assuming that the forward-direction of each + body frame is along the x-direction, i.e. :math:`(1, 0, 0)`. + """ + if self._heading_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.body_heading_w, + dim=(self.num_instances, self.num_bodies), + inputs=[self.FORWARD_VEC_B, self.body_link_quat_w], + outputs=[self._heading_w.data], + device=self.device, + ) + self._heading_w.timestamp = self._sim_timestamp + if self._heading_w_ta is None: + self._heading_w_ta = ProxyArray(self._heading_w.data) + return self._heading_w_ta + + def _create_buffers(self) -> None: + """Eagerly allocate every per-body TimestampedBuffer and the slots for + cached :class:`ProxyArray` wrappers. + + Buffers use direct ``(num_instances, num_bodies, D)`` shapes, matching + the fused binding output. No flat+strided tricks are needed because the + fused binding returns a contiguous ``(N, B, D)`` array directly. + """ + super()._create_buffers() + + N = self.num_instances + B = self.num_bodies + + # -- link frame w.r.t. world frame + self._body_link_pose_w = TimestampedBuffer((N, B), self.device, wp.transformf) + self._body_link_vel_w = TimestampedBuffer((N, B), self.device, wp.spatial_vectorf) + # -- com frame w.r.t. link frame + self._body_com_pose_b = TimestampedBuffer((N, B), self.device, wp.transformf) + # -- com frame w.r.t. world frame + self._body_com_pose_w = TimestampedBuffer((N, B), self.device, wp.transformf) + self._body_com_vel_w = TimestampedBuffer((N, B), self.device, wp.spatial_vectorf) + # -- combined state (cached, used by deprecated concat properties) + self._body_state_w = TimestampedBuffer((N, B), self.device, shared_kernels.vec13f) + self._body_link_state_w = TimestampedBuffer((N, B), self.device, shared_kernels.vec13f) + self._body_com_state_w = TimestampedBuffer((N, B), self.device, shared_kernels.vec13f) + # -- derived properties (in-body-frame velocities) + self._body_link_lin_vel_b = TimestampedBuffer((N, B), self.device, wp.vec3f) + self._body_link_ang_vel_b = TimestampedBuffer((N, B), self.device, wp.vec3f) + self._body_com_lin_vel_b = TimestampedBuffer((N, B), self.device, wp.vec3f) + self._body_com_ang_vel_b = TimestampedBuffer((N, B), self.device, wp.vec3f) + # -- derived properties (acceleration via finite differencing) + self._body_com_acc_w = TimestampedBuffer((N, B), self.device, wp.spatial_vectorf) + # Holds the previous-step COM velocity for FD; initialised lazily on first access. + self._previous_body_com_vel: wp.array | None = None + # -- derived properties (projected gravity and heading) + self._projected_gravity_b = TimestampedBuffer((N, B), self.device, wp.vec3f) + self._heading_w = TimestampedBuffer((N, B), self.device, wp.float32) + + # -- Body properties: mass (N, B) and inertia (N, B, 9). + # Initialised eagerly from the CPU-only bindings. + self._body_mass = TimestampedBuffer((N, B), self.device, wp.float32) + self._body_inertia = TimestampedBuffer((N, B, 9), self.device, wp.float32) + + # Pinned CPU staging buffers used by mass/com/inertia setters. + pinned = self.device != "cpu" + self._cpu_body_mass = wp.zeros((N, B), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_body_coms = wp.zeros((N, B, 7), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_body_inertia = wp.zeros((N, B, 9), dtype=wp.float32, device="cpu", pinned=pinned) + + # Eagerly read mass and inertia (CPU-only bindings) at construction time. + # The native fused binding returns body-major flat data ``(N*B[, D])``; + # the articulation-mode mock returns instance-major ``(N, B[, D])``. + # In either case we reshape on the CPU into instance-major numpy arrays. + def _read_cpu(tensor_type, trailing_dim=None): + binding = self._get_binding(tensor_type) + if binding is None: + return None + np_buf = np.zeros(binding.shape, dtype=np.float32) + binding.read(np_buf) + if binding.count == N: + # Mock fast-path: already (N, B[, D]). + return np_buf + # Native fused path: body-major flat -> instance-major via reshape+transpose. + if trailing_dim is None: + return np_buf.reshape(B, N).T.copy() + return np_buf.reshape(B, N, trailing_dim).transpose(1, 0, 2).copy() + + np_mass = _read_cpu(TT.BODY_MASS) + if np_mass is not None: + wp.copy(self._body_mass.data, wp.from_numpy(np_mass, dtype=wp.float32, device=self.device)) + self._body_mass.timestamp = self._sim_timestamp + + np_inertia = _read_cpu(TT.BODY_INERTIA, trailing_dim=9) + if np_inertia is not None: + wp.copy( + self._body_inertia.data, + wp.from_numpy(np_inertia, dtype=wp.float32, device=self.device), + ) + self._body_inertia.timestamp = self._sim_timestamp + + # -- Defaults (allocated here, filled by _process_cfg after __init__). + # Zero-initialized buffers; populated by RigidObjectCollection._process_cfg. + self._default_body_pose = wp.zeros((N, B), dtype=wp.transformf, device=self.device) + self._default_body_vel = wp.zeros((N, B), dtype=wp.spatial_vectorf, device=self.device) + + # Initialize ProxyArray wrappers. + self._pin_proxy_arrays() + + def _pin_proxy_arrays(self) -> None: + """Create pinned :class:`ProxyArray` wrappers for all data buffers. + + This is called once from :meth:`_create_buffers` during initialization. + OVPhysX tensor API buffers have stable GPU pointers across simulation steps, + so no rebinding is needed (unlike Newton). + """ + # Defaults + self._default_body_pose_ta: ProxyArray | None = None + self._default_body_vel_ta: ProxyArray | None = None + # Body state (timestamped) + self._body_link_pose_w_ta: ProxyArray | None = None + self._body_link_vel_w_ta: ProxyArray | None = None + self._body_com_pose_w_ta: ProxyArray | None = None + self._body_com_vel_w_ta: ProxyArray | None = None + self._body_com_pose_b_ta: ProxyArray | None = None + # Body properties + self._body_mass_ta: ProxyArray | None = None + self._body_inertia_ta: ProxyArray | None = None + # Derived properties (in-body-frame velocities) + self._body_link_lin_vel_b_ta: ProxyArray | None = None + self._body_link_ang_vel_b_ta: ProxyArray | None = None + self._body_com_lin_vel_b_ta: ProxyArray | None = None + self._body_com_ang_vel_b_ta: ProxyArray | None = None + # Derived properties (FD acceleration) + self._body_com_acc_w_ta: ProxyArray | None = None + self._body_com_lin_acc_w_ta: ProxyArray | None = None + self._body_com_ang_acc_w_ta: ProxyArray | None = None + # Derived properties (projected gravity and heading) + self._projected_gravity_b_ta: ProxyArray | None = None + self._heading_w_ta: ProxyArray | None = None + # Sliced properties (body link) + self._body_link_pos_w_ta: ProxyArray | None = None + self._body_link_quat_w_ta: ProxyArray | None = None + self._body_link_lin_vel_w_ta: ProxyArray | None = None + self._body_link_ang_vel_w_ta: ProxyArray | None = None + # Sliced properties (body com) + self._body_com_pos_w_ta: ProxyArray | None = None + self._body_com_quat_w_ta: ProxyArray | None = None + self._body_com_lin_vel_w_ta: ProxyArray | None = None + self._body_com_ang_vel_w_ta: ProxyArray | None = None + # Sliced properties (body com in body frame) + self._body_com_pos_b_ta: ProxyArray | None = None + self._body_com_quat_b_ta: ProxyArray | None = None + # Deprecated state-concat properties + self._default_body_state: wp.array | None = None + self._default_body_state_ta: ProxyArray | None = None + self._body_state_w_ta: ProxyArray | None = None + self._body_link_state_w_ta: ProxyArray | None = None + self._body_com_state_w_ta: ProxyArray | None = None + + """ + Helpers. + """ + + def _get_binding(self, tensor_type: int): + """Return the binding for the given tensor type, or None. + + Args: + tensor_type: The TensorType constant identifying which simulation buffer. + + Returns: + The cached :class:`TensorBinding`, or ``None`` if not available. + """ + b = self._bindings.get(tensor_type) + if b is not None: + return b + if self._binding_getter is not None: + b = self._binding_getter(tensor_type) + if b is not None: + self._bindings[tensor_type] = b + return b + return None + + def _read_view_scratch(self, tensor_type: int, binding) -> wp.array: + """Return a cached body-major scratch float32 buffer matching ``binding.shape``. + + Allocated on the binding's own device (GPU bindings → GPU, CPU-only + bindings → pinned host) so that ``binding.read(scratch)`` never crosses + devices. The scratch buffer is body-major flat + ``(num_bodies * num_instances[, D])`` and is reshaped into instance-major + ``(N, B[, D])`` by :meth:`_reshape_view_to_data_2d` / + :meth:`_reshape_view_to_data_3d` before being copied into the user-facing + timestamped buffer. + + Args: + tensor_type: TensorType key (used as cache key). + binding: The fused :class:`TensorBinding` whose shape the scratch + must match. + + Returns: + Cached body-major scratch buffer for reads. + """ + scratch = self._cpu_staging_buffers.get(tensor_type) + if scratch is not None: + return scratch + binding_device = "cpu" if tensor_type in TT._CPU_ONLY_TYPES else self.device + pinned = binding_device == "cpu" and self.device != "cpu" + if pinned: + scratch = wp.zeros(binding.shape, dtype=wp.float32, device="cpu", pinned=True) + else: + scratch = wp.zeros(binding.shape, dtype=wp.float32, device=binding_device) + self._cpu_staging_buffers[tensor_type] = scratch + return scratch + + def _read_binding_into_instance_major(self, tensor_type: int, buf: TimestampedBuffer, floats_per_elem: int) -> None: + """Read a fused binding into the instance-major ``buf.data``. + + The native fused multi-prim binding returns data in body-major flat order + ``(body0_env0, body0_env1, ..., body1_env0, ...)`` with + ``binding.count == num_instances * num_bodies``. The articulation-mode + mock used by iface tests instead exposes a directly instance-major view + with ``binding.count == num_instances`` and shape ``(N, B[, D])``. + + This method dispatches to the right path: + + * **Mock fast-path** (``binding.count == num_instances``): a float32 view + of the destination buffer is filled directly via ``binding.read()``. + * **Native fused path** (``binding.count == num_instances * num_bodies``): + ``binding.read()`` fills a body-major scratch, which is then reshaped + into instance-major order via :meth:`_reshape_view_to_data_2d` (for + single-element-per-body fields like mass) or + :meth:`_reshape_view_to_data_3d` (for fields with a trailing + ``floats_per_elem`` dimension), and the result is copied into the + destination buffer. + + Args: + tensor_type: TensorType key identifying the binding. + buf: Timestamped buffer to refresh. ``buf.data`` is + ``(num_instances, num_bodies)`` for single-element-per-body + fields (e.g. mass), or ``(num_instances, num_bodies, ...)`` for + multi-component fields. + floats_per_elem: Number of trailing ``float32`` elements per body + (e.g. 7 for transformf, 6 for spatial_vectorf, 9 for inertia). + Pass 1 for plain scalar fields like mass. + """ + if buf.timestamp >= self._sim_timestamp: + return + binding = self._get_binding(tensor_type) + if binding is None: + return + + B = self.num_bodies + + # Disambiguate via the binding's exposed shape: the articulation-mode + # mock returns a directly instance-major view ``(N, B[, D])`` while the + # native fused multi-prim binding lays elements body-major-flat with + # ``shape == (N * B[, D])``. + is_mock_layout = len(binding.shape) >= 2 and binding.shape[1] == B + + if is_mock_layout: + if buf.data.dtype == wp.float32: + view = buf.data + else: + view = wp.array( + ptr=buf.data.ptr, + shape=binding.shape, + dtype=wp.float32, + device=str(buf.data.device), + copy=False, + ) + binding.read(view) + buf.timestamp = self._sim_timestamp + return + + # Native fused path: read body-major scratch then strided-view reshape. + scratch = self._read_view_scratch(tensor_type, binding) + binding.read(scratch) + if floats_per_elem <= 1: + reshaped = self._reshape_view_to_data_2d(scratch) + else: + reshaped = self._reshape_view_to_data_3d(scratch, floats_per_elem) + # Copy into buf.data, reinterpreting structured-dtype buffers as float32. + if buf.data.dtype == wp.float32: + dst_view = buf.data + else: + dst_view = wp.array( + ptr=buf.data.ptr, + shape=reshaped.shape, + dtype=wp.float32, + device=str(buf.data.device), + copy=False, + ) + wp.copy(dst_view, reshaped) + buf.timestamp = self._sim_timestamp + + def _read_transform_binding(self, tensor_type: int, buf: TimestampedBuffer) -> None: + """Read a pose binding (``wp.transformf`` buffer), skipping if fresh. + + Args: + tensor_type: TensorType key. + buf: Timestamped :class:`wp.transformf` buffer to refresh. + """ + self._read_binding_into_instance_major(tensor_type, buf, floats_per_elem=7) + + def _read_spatial_vector_binding(self, tensor_type: int, buf: TimestampedBuffer) -> None: + """Read a velocity binding (``wp.spatial_vectorf`` buffer), skipping if fresh. + + Args: + tensor_type: TensorType key. + buf: Timestamped :class:`wp.spatial_vectorf` buffer to refresh. + """ + self._read_binding_into_instance_major(tensor_type, buf, floats_per_elem=6) + + def _reshape_view_to_data_2d(self, data: wp.array) -> wp.array: + """Reshape body-major flat data into instance-major ``(num_instances, num_bodies)``. + + The native fused binding lays elements out body-major: + ``(body0_env0, body0_env1, ..., body1_env0, body1_env1, ...)``. This helper + constructs a strided view that traverses the data in instance-major order + ``(env0_body0, env0_body1, ..., env1_body0, ...)``, then clones it onto + :attr:`device` for contiguity and (when needed) cross-device transfer. + + Args: + data: Body-major flat buffer. Shape is ``(num_bodies * num_instances,)`` + with any single-element dtype. + + Returns: + Contiguous instance-major buffer with shape ``(num_instances, num_bodies)`` + on :attr:`device`. + """ + element_size = wp.types.type_size_in_bytes(data.dtype) + strided_view = wp.array( + ptr=data.ptr, + shape=(self.num_instances, self.num_bodies), + dtype=data.dtype, + strides=(element_size, self.num_instances * element_size), + device=str(data.device), + ) + return wp.clone(strided_view, self.device) + + def _reshape_view_to_data_3d(self, data: wp.array, data_dim: int) -> wp.array: + """Reshape body-major flat data into instance-major ``(num_instances, num_bodies, data_dim)``. + + Companion of :meth:`_reshape_view_to_data_2d` for fields with a trailing + per-body dimension (e.g. pose has 7, spatial velocity has 6, inertia has 9). + + Args: + data: Body-major flat buffer with shape ``(num_bodies * num_instances, data_dim)`` + or ``(num_bodies * num_instances,)`` reinterpreted as ``data_dim``-wide rows. + data_dim: Trailing per-body dimension size. + + Returns: + Contiguous instance-major buffer with shape + ``(num_instances, num_bodies, data_dim)`` on :attr:`device`. + """ + element_size = wp.types.type_size_in_bytes(wp.float32) + row_size = element_size * data_dim + strided_view = wp.array( + ptr=data.ptr, + shape=(self.num_instances, self.num_bodies, data_dim), + dtype=wp.float32, + strides=(row_size, self.num_instances * row_size, element_size), + device=str(data.device), + ) + return wp.clone(strided_view, self.device) + + def _get_pos_from_transform(self, transform: wp.array) -> wp.array: + """Generates a position array from a transform array.""" + return wp.array( + ptr=transform.ptr, + shape=transform.shape, + dtype=wp.vec3f, + strides=transform.strides, + device=self.device, + ) + + def _get_quat_from_transform(self, transform: wp.array) -> wp.array: + """Generates a quaternion array from a transform array.""" + return wp.array( + ptr=transform.ptr + 3 * 4, + shape=transform.shape, + dtype=wp.quatf, + strides=transform.strides, + device=self.device, + ) + + def _get_lin_vel_from_spatial_vector(self, sv: wp.array) -> wp.array: + """Generates a linear velocity array from a spatial vector array.""" + return wp.array( + ptr=sv.ptr, + shape=sv.shape, + dtype=wp.vec3f, + strides=sv.strides, + device=self.device, + ) + + def _get_ang_vel_from_spatial_vector(self, sv: wp.array) -> wp.array: + """Generates an angular velocity array from a spatial vector array.""" + return wp.array( + ptr=sv.ptr + 3 * 4, + shape=sv.shape, + dtype=wp.vec3f, + strides=sv.strides, + device=self.device, + ) diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/physics/ovphysx_manager.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/physics/ovphysx_manager.py index 6caad37ab7bf..e96154bacf59 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/physics/ovphysx_manager.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/physics/ovphysx_manager.py @@ -47,6 +47,12 @@ class OvPhysxManager(PhysicsManager): _stage_path: ClassVar[str | None] = None _warmup_done: ClassVar[bool] = False _tmp_dir: ClassVar[tempfile.TemporaryDirectory | None] = None + # Device the process is locked to once :meth:`_warmup_and_load` constructs the + # ``ovphysx.PhysX`` instance for the first time. ``ovphysx<=0.3.7`` enforces + # a process-global device-mode lock at the C++ layer (see HACK note on + # :meth:`_release_physx`); we mirror it here so a clear Python error is raised + # if a later :class:`~isaaclab.sim.SimulationContext` requests a different device. + _locked_device: ClassVar[str | None] = None # Pending (source, targets, parent_positions) triples registered by # ovphysx_replicate() before the PhysX instance exists. Replayed via # physx.clone() in _warmup_and_load(). @@ -54,6 +60,11 @@ class OvPhysxManager(PhysicsManager): _pending_clones: ClassVar[list[tuple[str, list[str], list[tuple[float, float, float]]]]] = [] _atexit_registered: ClassVar[bool] = False + @classmethod + def get_dt(cls) -> float: + """Get the physics timestep. Alias for get_physics_dt().""" + return cls.get_physics_dt() + @classmethod def register_clone( cls, source: str, targets: list[str], parent_positions: list[tuple[float, float, float]] | None = None @@ -80,13 +91,20 @@ def register_clone( def initialize(cls, sim_context: SimulationContext) -> None: """Initialize the physics manager with simulation context. - This stores the config and device but does not create the ovphysx - instance yet -- the USD stage may not be fully populated at this point. - The actual creation happens lazily in :meth:`reset`. + This stores the config and device but does not load the USD stage yet -- + the stage may not be fully populated at this point. The actual load + happens lazily in :meth:`reset`. + + ``cls._physx`` is intentionally not cleared here: the ovphysx C++ instance + is process-global (see HACK on :meth:`_release_physx`). When a previous + :class:`SimulationContext` has already constructed it, we reuse it rather + than dropping the only Python reference (which would trigger the + destructor race) or re-constructing (which would hit the wheel's + device-mode lock). ``cls._locked_device`` carries the device the cached + instance is bound to. """ super().initialize(sim_context) cls._warmup_done = False - cls._physx = None cls._usd_handle = None cls._stage_path = None cls._pending_clones = [] @@ -120,6 +138,7 @@ def step(cls) -> None: dt = cls.get_physics_dt() sim_time = PhysicsManager._sim_time cls._physx.step_sync(dt=dt, sim_time=sim_time) + cls._physx.update_articulations_kinematic() PhysicsManager._sim_time += dt @classmethod @@ -139,15 +158,27 @@ def close(cls) -> None: @classmethod def _release_physx(cls) -> None: - """Release the ovphysx instance if it exists. Safe to call multiple times. - - With ovphysx<=0.3.7 and Kit's pxr in the same process, physx.release() - deadlocks due to dual-Carbonite static destructor races. Skip the - native release and let os._exit() (registered via atexit) terminate the - process; GPU resources are reclaimed by the driver. + """Soft-reset the ovphysx runtime stage; keep the C++ instance alive. + + Calls ``physx.reset()`` to clear the loaded scene, but does **not** drop + the Python reference. The cached :class:`ovphysx.PhysX` is reused by the + next :class:`~isaaclab.sim.SimulationContext` via the reuse path in + :meth:`_warmup_and_load`. Safe to call multiple times. + + HACK(ovphysx<=0.3.7): the wheel's bundled libcarb.so and Kit's libcarb.so + coexist in the same process whenever ``import pxr`` runs (Kit USD plugins + on ``LD_LIBRARY_PATH`` pull in Kit's Carbonite). Both register C++ static + destructors that race at process exit -- and crucially, also race when + ``ovphysx.PhysX``'s Python destructor fires mid-process via refcount drop. + So we must never let the only Python reference go to zero while the + process is alive. ``os._exit(0)`` (registered via ``atexit`` in + :meth:`_warmup_and_load`) sidesteps the static-destructor phase entirely + at process exit. Remove this workaround once the wheel ships a + namespace-isolated Carbonite (different soname / hidden visibility). """ if cls._physx is not None: - cls._physx = None + op = cls._physx.reset() + cls._physx.wait_op(op) @classmethod def get_physx_instance(cls) -> Any: @@ -160,7 +191,22 @@ def get_physx_instance(cls) -> Any: @classmethod def _warmup_and_load(cls) -> None: - """Export the USD stage, create the ovphysx instance, and load the scene.""" + """Export the USD stage and load it into the ovphysx runtime. + + On the first call per process, constructs the :class:`ovphysx.PhysX` + instance, registers the ``atexit`` handler, and locks the process to + the resolved device. On subsequent calls, reuses the cached instance + (see HACK on :meth:`_release_physx`) -- exporting the new USD, + re-attaching it via ``add_usd``, replaying pending clones, and (on GPU) + re-running ``warmup_gpu`` so the new stage's bodies are resident. + + Raises: + RuntimeError: if ``SimulationContext`` is not set, or if a device + different from the process-locked one is requested. The wheel + enforces a process-global device-mode lock at the C++ layer; + we surface it here as a clear Python error before the wheel + would raise :exc:`ovphysx.types.PhysXDeviceError`. + """ sim = PhysicsManager._sim if sim is None: raise RuntimeError("OvPhysxManager: SimulationContext is not set.") @@ -174,9 +220,16 @@ def _warmup_and_load(cls) -> None: gpu_index = 0 ovphysx_device = "cpu" + if cls._locked_device is not None and ovphysx_device != cls._locked_device: + raise RuntimeError( + f"OvPhysxManager is locked to device {cls._locked_device!r} for the lifetime of this process; " + f"cannot switch to {ovphysx_device!r}. ovphysx<=0.3.7 binds device mode at the C++ layer on the " + "first ovphysx.PhysX(...) construction and it cannot be changed without restarting the process." + ) + scene_prim = sim.stage.GetPrimAtPath(sim.cfg.physics_prim_path) - if scene_prim.IsValid() and ovphysx_device == "gpu": - cls._configure_physx_scene_prim(scene_prim, PhysicsManager._cfg) + if scene_prim.IsValid(): + cls._configure_physx_scene_prim(scene_prim, PhysicsManager._cfg, ovphysx_device) # Export the current USD stage to a temporary file so ovphysx can load it. cls._tmp_dir = tempfile.TemporaryDirectory(prefix="isaaclab_ovphysx_") @@ -185,6 +238,66 @@ def _warmup_and_load(cls) -> None: cls._stage_path = stage_file logger.info("OvPhysxManager: exported USD stage to %s", stage_file) + if cls._physx is None: + cls._construct_physx(ovphysx_device, gpu_index) + cls._locked_device = ovphysx_device + else: + # Reuse path: the cached PhysX may still hold the prior stage (the + # wheel allows only one loaded USD at a time). ``physx.reset()`` is + # idempotent on an already-cleared stage and required when this is + # a second :meth:`_warmup_and_load` within the same SimulationContext + # (e.g. when a caller manually clears ``_warmup_done`` to force a + # re-warmup). + op = cls._physx.reset() + cls._physx.wait_op(op) + + usd_handle, op_idx = cls._physx.add_usd(stage_file) + cls._physx.wait_op(op_idx) + cls._usd_handle = usd_handle + logger.info("OvPhysxManager: loaded USD into ovphysx (device=%s)", ovphysx_device) + + # Replay pending physics clones registered by ovphysx_replicate(). + # The USD stage contains only env_0's physics; env_1..N are empty + # Xform containers. physx.clone() creates the remaining environments + # in the physics runtime without modifying the USD file. + if cls._pending_clones: + # ovphysx_replicate() only registers pending clones when clone_usd=False, + # meaning the USD contains only env_0 physics and physx.clone() is required + # to populate env_1..N in the physics runtime. Execute unconditionally — + # no USD content heuristic is needed. + for source, targets, parent_positions in cls._pending_clones: + logger.info( + "OvPhysxManager: cloning %s -> %d targets (%s ... %s)", + source, + len(targets), + targets[0], + targets[-1], + ) + if parent_positions: + transforms = [(x, y, z, 0.0, 0.0, 0.0, 1.0) for x, y, z in parent_positions] + else: + transforms = None + op_idx = cls._physx.clone(source, targets, transforms) + cls._physx.wait_op(op_idx) + cls._pending_clones = [] + + # GPU bodies must be re-warmed after every add_usd: the cached PhysX + # instance carries its old buffer layout from the previous stage. + if ovphysx_device == "gpu": + cls._physx.warmup_gpu() + + cls.dispatch_event(PhysicsEvent.MODEL_INIT, payload={}) + cls._warmup_done = True + + @classmethod + def _construct_physx(cls, ovphysx_device: str, gpu_index: int) -> None: + """Bootstrap the ``ovphysx`` wheel and create the :class:`ovphysx.PhysX` instance. + + Runs once per process. Configures worker threads, registers the + process-exit ``os._exit(0)`` handler, and stores the result on + ``cls._physx``. See HACK on :meth:`_release_physx` for why the + instance must outlive every individual :class:`SimulationContext`. + """ # HACK (temporary): hide pxr from sys.modules during ovphysx bootstrap. # IsaacSim's pxr reports version 0.25.5 (pip convention) while ovphysx # expects 25.11 (OpenUSD release convention). Hiding pxr causes @@ -270,52 +383,25 @@ def _atexit_release_and_exit(): atexit.register(_atexit_release_and_exit) cls._atexit_registered = True - usd_handle, op_idx = cls._physx.add_usd(stage_file) - cls._physx.wait_op(op_idx) - cls._usd_handle = usd_handle - logger.info("OvPhysxManager: loaded USD into ovphysx (device=%s)", ovphysx_device) - - # Replay pending physics clones registered by ovphysx_replicate(). - # The USD stage contains only env_0's physics; env_1..N are empty - # Xform containers. physx.clone() creates the remaining environments - # in the physics runtime without modifying the USD file. - if cls._pending_clones: - # ovphysx_replicate() only registers pending clones when clone_usd=False, - # meaning the USD contains only env_0 physics and physx.clone() is required - # to populate env_1..N in the physics runtime. Execute unconditionally — - # no USD content heuristic is needed. - for source, targets, parent_positions in cls._pending_clones: - logger.info( - "OvPhysxManager: cloning %s -> %d targets (%s ... %s)", - source, - len(targets), - targets[0], - targets[-1], - ) - if parent_positions: - transforms = [(x, y, z, 0.0, 0.0, 0.0, 1.0) for x, y, z in parent_positions] - else: - transforms = None - op_idx = cls._physx.clone(source, targets, transforms) - cls._physx.wait_op(op_idx) - cls._pending_clones = [] - - if ovphysx_device == "gpu": - cls._physx.warmup_gpu() - - cls.dispatch_event(PhysicsEvent.MODEL_INIT, payload={}) - cls._warmup_done = True - @staticmethod - def _configure_physx_scene_prim(scene_prim, cfg) -> None: - """Apply PhysxSceneAPI schema and GPU dynamics attributes to a scene prim. + def _configure_physx_scene_prim(scene_prim, cfg, device: str) -> None: + """Apply PhysxSceneAPI schema and device-specific scene attributes to the + scene prim. The PhysxSchema USD plugin may not be loaded in standalone ovphysx mode, so we write the apiSchemas list entry and scene attributes directly via raw Sdf metadata manipulation instead of using the high-level USD API. - Without these attributes PhysX defaults to CPU broadphase even when - ovphysx is created with device="gpu". + The schema and scene-query-support attribute are applied regardless of + device. The GPU-specific dynamics/broadphase/capacity attributes are + applied only when ``device == "gpu"`` — without them PhysX defaults to + CPU broadphase even when ovphysx is created with ``device="gpu"``. + + Args: + scene_prim: The /World/PhysicsScene prim to configure. + cfg: The :class:`OvPhysxCfg` carrying GPU buffer-capacity values. + Only consulted when ``device == "gpu"``. + device: Resolved physics device — one of ``"cpu"`` or ``"gpu"``. """ from pxr import Sdf @@ -326,22 +412,24 @@ def _configure_physx_scene_prim(scene_prim, cfg) -> None: items.append("PhysxSceneAPI") schemas.prependedItems = items scene_prim.SetMetadata("apiSchemas", schemas) - scene_prim.CreateAttribute("physxScene:enableGPUDynamics", Sdf.ValueTypeNames.Bool).Set(True) - scene_prim.CreateAttribute("physxScene:broadphaseType", Sdf.ValueTypeNames.String).Set("GPU") - - if cfg is not None: - for attr, val in [ - ("gpuMaxRigidContactCount", cfg.gpu_max_rigid_contact_count), - ("gpuMaxRigidPatchCount", cfg.gpu_max_rigid_patch_count), - ("gpuFoundLostPairsCapacity", cfg.gpu_found_lost_pairs_capacity), - ("gpuFoundLostAggregatePairsCapacity", cfg.gpu_found_lost_aggregate_pairs_capacity), - ("gpuTotalAggregatePairsCapacity", cfg.gpu_total_aggregate_pairs_capacity), - ("gpuCollisionStackSize", cfg.gpu_collision_stack_size), - ]: - scene_prim.CreateAttribute(f"physxScene:{attr}", Sdf.ValueTypeNames.UInt).Set(val) # Propagate scene query support from SimulationCfg so omni.physx creates # the scene with the correct query mode. OvPhysxCfg does not carry this field. sim_cfg = PhysicsManager._sim.cfg if PhysicsManager._sim is not None else None enable_sq = getattr(sim_cfg, "enable_scene_query_support", False) scene_prim.CreateAttribute("physxScene:enableSceneQuerySupport", Sdf.ValueTypeNames.Bool).Set(enable_sq) + + if device == "gpu": + scene_prim.CreateAttribute("physxScene:enableGPUDynamics", Sdf.ValueTypeNames.Bool).Set(True) + scene_prim.CreateAttribute("physxScene:broadphaseType", Sdf.ValueTypeNames.String).Set("GPU") + + if cfg is not None: + for attr, val in [ + ("gpuMaxRigidContactCount", cfg.gpu_max_rigid_contact_count), + ("gpuMaxRigidPatchCount", cfg.gpu_max_rigid_patch_count), + ("gpuFoundLostPairsCapacity", cfg.gpu_found_lost_pairs_capacity), + ("gpuFoundLostAggregatePairsCapacity", cfg.gpu_found_lost_aggregate_pairs_capacity), + ("gpuTotalAggregatePairsCapacity", cfg.gpu_total_aggregate_pairs_capacity), + ("gpuCollisionStackSize", cfg.gpu_collision_stack_size), + ]: + scene_prim.CreateAttribute(f"physxScene:{attr}", Sdf.ValueTypeNames.UInt).Set(val) diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/tensor_types.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/tensor_types.py index 44a5cadeeb0a..9e7df44aef9a 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/tensor_types.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/tensor_types.py @@ -191,6 +191,65 @@ Shape is ``[N, L, 9]``, dtype ``float32``. """ +""" +Rigid-body TensorTypes + +Shapes assume N = number of rigid actor instances matched by the binding +pattern. Components and units are stated per alias below. +""" + +RIGID_BODY_POSE = _TT.RIGID_BODY_POSE +"""Rigid actor root transform — read/write, GPU. Shape ``(N, 7)``, +components ``(px, py, pz, qx, qy, qz, qw)`` [m, dimensionless].""" + +RIGID_BODY_VELOCITY = _TT.RIGID_BODY_VELOCITY +"""Rigid actor root spatial velocity — read/write, GPU. Shape ``(N, 6)``, +components ``(vx, vy, vz, wx, wy, wz)`` [m/s, rad/s].""" + +RIGID_BODY_WRENCH = _TT.RIGID_BODY_WRENCH +"""External wrench applied at a world-frame point — write-only, GPU. +Shape ``(N, 9)``, components ``(fx, fy, fz, tx, ty, tz, px, py, pz)`` +[N, N·m, m]. Cleared after each sim step (instantaneous semantics).""" + +RIGID_BODY_MASS = _TT.RIGID_BODY_MASS +"""Rigid actor mass — read/write, CPU. Shape ``(N,)`` [kg].""" + +RIGID_BODY_COM_POSE = _TT.RIGID_BODY_COM_POSE +"""Center-of-mass pose in actor-link frame — read/write, CPU. Shape +``(N, 7)``, components ``(px, py, pz, qx, qy, qz, qw)`` [m, dimensionless].""" + +RIGID_BODY_INERTIA = _TT.RIGID_BODY_INERTIA +"""Rigid actor inertia tensor in COM frame — read/write, CPU. Shape +``(N, 9)``, row-major flatten of the 3×3 inertia matrix +``(Ixx, Ixy, Ixz, Iyx, Iyy, Iyz, Izx, Izy, Izz)`` [kg·m²].""" + +# These three aliases are pending an upcoming ovphysx wheel update. +# When the wheel ships them, the corresponding ``hasattr`` checks below +# in IsaacLab consumers will start returning True and the bindings will +# become usable; until then, ``isaaclab_ovphysx.tensor_types`` simply +# does not expose the alias. +try: + RIGID_BODY_ACCELERATION = _TT.RIGID_BODY_ACCELERATION + """Rigid actor spatial acceleration — read-only, GPU. Shape ``(N, 6)``, + components ``(ax, ay, az, αx, αy, αz)`` [m/s², rad/s²].""" +except AttributeError: + pass + +try: + RIGID_BODY_INV_MASS = _TT.RIGID_BODY_INV_MASS + """Rigid actor inverse mass — read-only, CPU. Shape ``(N,)`` [1/kg]. + Zero indicates an immovable actor.""" +except AttributeError: + pass + +try: + RIGID_BODY_INV_INERTIA = _TT.RIGID_BODY_INV_INERTIA + """Rigid actor inverse inertia tensor in COM frame — read-only, CPU. + Shape ``(N, 9)``, row-major flatten of the 3×3 matrix [1/(kg·m²)]. + Zero rows indicate locked rotational DOFs.""" +except AttributeError: + pass + """ Dynamics tensors (GPU) """ @@ -306,29 +365,31 @@ # fmt: on # DOF/body property tensor types are CPU-resident even in GPU simulations. # Write helpers check this set to route data through CPU, not self._device. -_CPU_ONLY_TYPES: frozenset[TensorType] = frozenset( - { - DOF_STIFFNESS, - DOF_DAMPING, - DOF_LIMIT, - DOF_MAX_VELOCITY, - DOF_MAX_FORCE, - DOF_ARMATURE, - DOF_FRICTION_PROPERTIES, - BODY_MASS, - BODY_COM_POSE, - BODY_INERTIA, - BODY_INV_MASS, - BODY_INV_INERTIA, - FIXED_TENDON_STIFFNESS, - FIXED_TENDON_DAMPING, - FIXED_TENDON_LIMIT_STIFFNESS, - FIXED_TENDON_LIMIT, - FIXED_TENDON_REST_LENGTH, - FIXED_TENDON_OFFSET, - SPATIAL_TENDON_STIFFNESS, - SPATIAL_TENDON_DAMPING, - SPATIAL_TENDON_LIMIT_STIFFNESS, - SPATIAL_TENDON_OFFSET, - } +# +# Tendon tensor types are NOT in this set: PhysX exposes tendons on the +# simulation device (its ``set_fixed_tendon_properties`` takes ``data.warp`` +# without a ``device="cpu"`` clone, unlike ``set_dof_stiffnesses``), and the +# OVPhysX wheel mirrors that — tendon bindings are GPU-resident on a GPU sim. +_CPU_ONLY_TYPES_CANDIDATES: tuple = ( + DOF_STIFFNESS, + DOF_DAMPING, + DOF_LIMIT, + DOF_MAX_VELOCITY, + DOF_MAX_FORCE, + DOF_ARMATURE, + DOF_FRICTION_PROPERTIES, + BODY_MASS, + BODY_COM_POSE, + BODY_INERTIA, + BODY_INV_MASS, + BODY_INV_INERTIA, + # Rigid-body CPU-only entries (always available) + RIGID_BODY_MASS, + RIGID_BODY_COM_POSE, + RIGID_BODY_INERTIA, +) +# Optional rigid-body CPU entries: only included when the wheel exposes them. +_RIGID_BODY_OPTIONAL_CPU: tuple = tuple( + globals()[name] for name in ("RIGID_BODY_INV_MASS", "RIGID_BODY_INV_INERTIA") if name in globals() ) +_CPU_ONLY_TYPES: frozenset[TensorType] = frozenset(_CPU_ONLY_TYPES_CANDIDATES + _RIGID_BODY_OPTIONAL_CPU) diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/test/mock_interfaces/views/mock_ovphysx_bindings.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/test/mock_interfaces/views/mock_ovphysx_bindings.py index 29472ce74fd0..51e96d9bb427 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/test/mock_interfaces/views/mock_ovphysx_bindings.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/test/mock_interfaces/views/mock_ovphysx_bindings.py @@ -7,6 +7,8 @@ from __future__ import annotations +from typing import Literal + import numpy as np from isaaclab_ovphysx import tensor_types as TT @@ -161,6 +163,10 @@ class MockOvPhysxBindingSet: for a given articulation configuration. Mirrors the tensor types that ``Articulation._initialize_impl`` creates. + + With ``asset_kind='rigid_object'`` it produces the smaller set + consumed by ``RigidObject._initialize_impl``: ``RIGID_BODY_*`` only, + ``num_joints`` must be 0, ``num_bodies`` must be 1, no tendons. """ def __init__( @@ -173,7 +179,54 @@ def __init__( body_names: list[str] | None = None, num_fixed_tendons: int = 0, num_spatial_tendons: int = 0, + *, + asset_kind: Literal["articulation", "rigid_object"] = "articulation", ): + if asset_kind == "rigid_object": + if num_joints != 0 or num_bodies != 1 or num_fixed_tendons != 0 or num_spatial_tendons != 0: + raise ValueError( + "asset_kind='rigid_object' requires num_joints=0, num_bodies=1, " + "num_fixed_tendons=0, num_spatial_tendons=0; got " + f"num_joints={num_joints}, num_bodies={num_bodies}, " + f"num_fixed_tendons={num_fixed_tendons}, " + f"num_spatial_tendons={num_spatial_tendons}" + ) + N = num_instances + if body_names is None: + body_names = ["base_link"] + common = dict( + count=N, + dof_count=0, + body_count=1, + joint_count=0, + is_fixed_base=is_fixed_base, + dof_names=[], + body_names=body_names, + joint_names=[], + fixed_tendon_count=0, + spatial_tendon_count=0, + ) + self.bindings: dict[int, MockTensorBinding] = { + TT.RIGID_BODY_POSE: MockTensorBinding(TT.RIGID_BODY_POSE, (N, 7), **common), + TT.RIGID_BODY_VELOCITY: MockTensorBinding(TT.RIGID_BODY_VELOCITY, (N, 6), **common), + TT.RIGID_BODY_WRENCH: MockTensorBinding(TT.RIGID_BODY_WRENCH, (N, 9), write_only=True, **common), + TT.RIGID_BODY_MASS: MockTensorBinding(TT.RIGID_BODY_MASS, (N,), **common), + TT.RIGID_BODY_COM_POSE: MockTensorBinding(TT.RIGID_BODY_COM_POSE, (N, 7), **common), + TT.RIGID_BODY_INERTIA: MockTensorBinding(TT.RIGID_BODY_INERTIA, (N, 9), **common), + } + # Optional bindings: only present when the wheel exposes the alias. + if hasattr(TT, "RIGID_BODY_ACCELERATION"): + self.bindings[TT.RIGID_BODY_ACCELERATION] = MockTensorBinding( + TT.RIGID_BODY_ACCELERATION, (N, 6), **common + ) + if hasattr(TT, "RIGID_BODY_INV_MASS"): + self.bindings[TT.RIGID_BODY_INV_MASS] = MockTensorBinding(TT.RIGID_BODY_INV_MASS, (N,), **common) + if hasattr(TT, "RIGID_BODY_INV_INERTIA"): + self.bindings[TT.RIGID_BODY_INV_INERTIA] = MockTensorBinding( + TT.RIGID_BODY_INV_INERTIA, (N, 9), **common + ) + return + N = num_instances D = num_joints L = num_bodies @@ -259,17 +312,25 @@ def set_random_data(self) -> None: for b in self.bindings.values(): if not b._write_only: b.set_random_data() - lim = self.bindings[TT.DOF_LIMIT] - lim._data[..., 0] = -3.14 - lim._data[..., 1] = 3.14 - for tt in (TT.ROOT_POSE, TT.LINK_POSE, TT.BODY_COM_POSE): + if TT.DOF_LIMIT in self.bindings: + lim = self.bindings[TT.DOF_LIMIT] + lim._data[..., 0] = -3.14 + lim._data[..., 1] = 3.14 + pose_keys = [ + k + for k in (TT.ROOT_POSE, TT.LINK_POSE, TT.BODY_COM_POSE, TT.RIGID_BODY_POSE, TT.RIGID_BODY_COM_POSE) + if k in self.bindings + ] + for tt in pose_keys: b = self.bindings[tt] b._data[..., 3:6] = 0.0 b._data[..., 6] = 1.0 - self.bindings[TT.BODY_MASS]._data = np.abs(self.bindings[TT.BODY_MASS]._data) + 0.1 - self.bindings[TT.DOF_MAX_VELOCITY]._data = np.abs(self.bindings[TT.DOF_MAX_VELOCITY]._data) + 1.0 - self.bindings[TT.DOF_MAX_FORCE]._data = np.abs(self.bindings[TT.DOF_MAX_FORCE]._data) + 1.0 - # Set sensible defaults for fixed tendon limits + for mass_key in (TT.BODY_MASS, TT.RIGID_BODY_MASS): + if mass_key in self.bindings: + self.bindings[mass_key]._data = np.abs(self.bindings[mass_key]._data) + 0.1 + for max_key in (TT.DOF_MAX_VELOCITY, TT.DOF_MAX_FORCE): + if max_key in self.bindings: + self.bindings[max_key]._data = np.abs(self.bindings[max_key]._data) + 1.0 if TT.FIXED_TENDON_LIMIT in self.bindings: tlim = self.bindings[TT.FIXED_TENDON_LIMIT] tlim._data[..., 0] = -1.0 diff --git a/source/isaaclab_ovphysx/setup.py b/source/isaaclab_ovphysx/setup.py index 4898806285dd..f09519155eb8 100644 --- a/source/isaaclab_ovphysx/setup.py +++ b/source/isaaclab_ovphysx/setup.py @@ -37,6 +37,7 @@ "isaaclab_ovphysx", "isaaclab_ovphysx.assets", "isaaclab_ovphysx.assets.articulation", + "isaaclab_ovphysx.assets.rigid_object_collection", "isaaclab_ovphysx.cloner", "isaaclab_ovphysx.physics", "isaaclab_ovphysx.test", diff --git a/source/isaaclab_ovphysx/test/assets/test_articulation.py b/source/isaaclab_ovphysx/test/assets/test_articulation.py index 52998a2ec5f6..057d720b640c 100644 --- a/source/isaaclab_ovphysx/test/assets/test_articulation.py +++ b/source/isaaclab_ovphysx/test/assets/test_articulation.py @@ -3,114 +3,2428 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Unit tests for ovphysx articulation helpers.""" +# ignore private usage of variables warning +# pyright: reportPrivateUsage=none + + +"""Real-backend tests for the OVPhysX Articulation. + +Mirrors :mod:`isaaclab_physx.test.assets.test_articulation` 1-to-1: same set +of test functions, names, parametrizations, and assertions. + +OVPhysX runs kitless under ``./scripts/run_ovphysx.sh`` so there is no +``AppLauncher`` boot — :class:`~isaaclab.sim.SimulationContext` is driven +directly via ``build_simulation_context(sim_cfg=SimulationCfg(physics=OvPhysxCfg(), ...))`` +which works because :func:`isaaclab.app.has_kit` returns False in this +environment. + +PhysX-specific ``cube_object.root_view.set_X(...)`` / ``get_X(...)`` calls are +adapted to OVPhysX by going through the backend's per-tensor-type binding +dictionary (``cube_object._bindings`` / :meth:`~isaaclab_ovphysx.assets.Articulation._get_binding`) +and the public setters (:meth:`set_masses_index`, :meth:`set_coms_index`, +:meth:`set_inertias_index`). Reads use the data-class properties +(``cube_object.data.body_mass``, ``body_inertia``, ``body_com_pose_b``). + +Process-global device lock +-------------------------- + +``ovphysx<=0.3.7`` binds device mode (CPU vs GPU) at the C++ layer on the +first ``ovphysx.PhysX(device=...)`` call and cannot release/swap it without a +process restart. :class:`~isaaclab_ovphysx.physics.OvPhysxManager` tracks +this on ``_locked_device`` and raises :exc:`RuntimeError` if a later +:class:`SimulationContext` requests a different device. The +``_ovphysx_skip_other_device`` autouse fixture below preempts that error in +parametrized tests by ``pytest.skip``-ing on the unlocked device, so the +session finishes cleanly when only one device is exercised. + +CI note +------- +Because the lock is process-global, full coverage requires **two separate +``./scripts/run_ovphysx.sh -m pytest`` invocations** -- once with ``-k 'cpu'`` +and once with ``-k 'cuda:0'``. Tracked as gap G5 in +``docs/superpowers/specs/2026-04-28-ovphysx-wheel-gaps-for-marco.md``; until +the wheel exposes a way to reset Carbonite device state, this is the supported +pattern. +""" from __future__ import annotations -from types import SimpleNamespace +import sys +import numpy as np import pytest +import torch import warp as wp +from isaaclab_ovphysx import tensor_types as TT +from isaaclab_ovphysx.assets import Articulation +from isaaclab_ovphysx.physics import OvPhysxCfg -from pxr import Sdf, Usd, UsdPhysics - -# The CI isaaclab_ov* pattern unintentionally collects isaaclab_ovphysx tests, -# but the ovphysx wheel is not installed in that environment. Skip gracefully -# so the isaaclab_ov CI pipeline is not blocked by an unrelated dependency. -pytest.importorskip("ovphysx.types", reason="ovphysx wheel not installed") +import isaaclab.sim as sim_utils +import isaaclab.utils.math as math_utils +import isaaclab.utils.string as string_utils +from isaaclab.actuators import ActuatorBase, IdealPDActuatorCfg, ImplicitActuatorCfg +from isaaclab.assets import ArticulationCfg +from isaaclab.envs.mdp.terminations import joint_effort_out_of_limit +from isaaclab.managers import SceneEntityCfg +from isaaclab.sim import SimulationCfg, build_simulation_context +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.version import get_isaac_sim_version, has_kit -from isaaclab_ovphysx.assets.articulation.articulation import Articulation # noqa: E402 -from isaaclab_ovphysx.physics import OvPhysxManager # noqa: E402 -from isaaclab_ovphysx.test.mock_interfaces.views import MockOvPhysxBindingSet # noqa: E402 +## +# Pre-defined configs +## +from isaaclab_assets import ANYMAL_C_CFG, FRANKA_PANDA_CFG, SHADOW_HAND_CFG # isort:skip wp.init() -def _define_tendon_joint(stage: Usd.Stage, path: str, schema_name: str) -> None: - """Define a revolute joint prim with a tendon schema marker.""" - joint = UsdPhysics.RevoluteJoint.Define(stage, path) - schemas = Sdf.TokenListOp() - schemas.explicitItems = [schema_name] - joint.GetPrim().SetMetadata("apiSchemas", schemas) - - -def _make_articulation_root_stage(tmp_path) -> str: - """Create a stage with one relevant articulation subtree and unrelated joints elsewhere.""" - stage = Usd.Stage.CreateInMemory() - stage.DefinePrim("/World", "Xform") - stage.DefinePrim("/World/envs", "Xform") - stage.DefinePrim("/World/envs/env_0", "Xform") - stage.DefinePrim("/World/envs/env_0/Robot", "Xform") - stage.DefinePrim("/World/envs/env_0/Robot/root", "Xform") - stage.DefinePrim("/World/unrelated", "Xform") - - _define_tendon_joint( - stage, - "/World/envs/env_0/Robot/root/fixed_joint", - "PhysxTendonAxisRootAPI:inst0", - ) - _define_tendon_joint( - stage, - "/World/envs/env_0/Robot/root/spatial_joint", - "PhysxTendonAttachmentRootAPI:inst0", - ) - _define_tendon_joint( - stage, - "/World/unrelated/unrelated_fixed_joint", - "PhysxTendonAxisRootAPI:inst0", - ) - _define_tendon_joint( - stage, - "/World/unrelated/unrelated_spatial_joint", - "PhysxTendonAttachmentLeafAPI:inst0", - ) - - stage_path = tmp_path / "scene.usda" - stage.Export(str(stage_path)) - return str(stage_path) - - -def _make_articulation_shell() -> Articulation: - """Create a minimal ovphysx articulation shell for tendon processing tests.""" - articulation = object.__new__(Articulation) - bindings = MockOvPhysxBindingSet( - num_instances=1, - num_joints=2, - num_bodies=2, - num_fixed_tendons=1, - num_spatial_tendons=1, - ) - object.__setattr__(articulation, "_bindings", bindings.bindings) - object.__setattr__(articulation, "_articulation_root_path", "/World/envs/env_0/Robot/root") - object.__setattr__(articulation, "_initialize_handle", None) - object.__setattr__(articulation, "_invalidate_initialize_handle", None) - object.__setattr__(articulation, "_prim_deletion_handle", None) - object.__setattr__(articulation, "_debug_vis_handle", None) - object.__setattr__( - articulation, - "_data", - SimpleNamespace( - _num_fixed_tendons=0, - _num_spatial_tendons=0, - fixed_tendon_names=[], - spatial_tendon_names=[], - ), - ) - return articulation - - -def test_process_tendons_scopes_to_articulation_root(tmp_path): - """Tendon discovery should ignore joints that live outside the current articulation subtree.""" - articulation = _make_articulation_shell() - stage_path = _make_articulation_root_stage(tmp_path) - old_stage_path = OvPhysxManager._stage_path - OvPhysxManager._stage_path = stage_path - try: - articulation._process_tendons() - finally: - OvPhysxManager._stage_path = old_stage_path - - assert articulation.fixed_tendon_names == ["fixed_joint"] - assert articulation.spatial_tendon_names == ["spatial_joint"] - assert articulation._data.fixed_tendon_names == ["fixed_joint"] - assert articulation._data.spatial_tendon_names == ["spatial_joint"] +_OMNI_PHYSX_SCHEMAS_GAP_REASON = ( + "Schema-level fixed-joint creation in :mod:`isaaclab.sim.schemas` imports " + "``omni.physx.scripts.utils``, which is a Kit-only module not shipped by " + "the ovphysx wheel. See " + "docs/superpowers/specs/2026-04-28-ovphysx-wheel-gaps-for-marco.md." +) + +_MATERIAL_GAP_REASON = ( + "Requires a ``RIGID_BODY_MATERIAL`` TensorType (or a view-helper) on the " + "ovphysx wheel side. ``Articulation.root_view`` is a per-tensor-type " + "bindings dict on OVPhysX, so ``root_view.get_material_properties()`` / " + "``set_material_properties()`` / ``max_shapes`` are not available. See " + "docs/superpowers/specs/2026-04-28-ovphysx-wheel-gaps-for-marco.md." +) + +def _read_binding_to_torch(articulation: Articulation, tensor_type: int, device: str | torch.device) -> torch.Tensor: + """Read an OVPhysX TensorBinding into a torch tensor on *device*. + + Test-side adapter for the verbatim PhysX mirror. PhysX cross-checks the + data class against the simulation via ``articulation.root_view.get_X()`` + accessors; on OVPhysX, ``root_view`` is a per-tensor-type bindings dict + (no view-level getters), so we read the binding directly into a CPU + numpy buffer (CPU-only types) and move the result to *device*. + """ + binding = articulation.root_view[tensor_type] + np_buf = np.zeros(binding.shape, dtype=np.float32) + binding.read(np_buf) + return torch.from_numpy(np_buf).to(device) + + +# Session-locked device. Set on the first parametrized test that runs and +# never reassigned -- ovphysx's process-global device lock means subsequent +# tests on the other device must skip. +_LOCKED_DEVICE: list[str | None] = [None] + + +@pytest.fixture(autouse=True) +def _ovphysx_skip_other_device(request): + """Skip tests whose ``device`` parameter mismatches the session-locked device. + + ``ovphysx<=0.3.7`` locks the process-global device mode on the first + ``ovphysx.PhysX(device=...)`` call, so any test parametrized to a different + device after the first ``sim.reset()`` would hit + :exc:`ovphysx.types.PhysXDeviceError`. We detect the locked device on the + first encounter and skip subsequent tests on the other device with a clear + message so the run finishes cleanly rather than producing spurious failures. + """ + callspec = getattr(request.node, "callspec", None) + device = callspec.params.get("device") if callspec is not None else None + if device is None: + # Test does not parametrize on device (e.g. test_warmup_attach_stage_not_called_for_cpu). + return + locked = _LOCKED_DEVICE[0] + if locked is None: + _LOCKED_DEVICE[0] = device + return + if device != locked: + pytest.skip( + f"ovphysx process-global device lock is held by '{locked}'; cannot run '{device}' " + "tests in the same session. Run pytest twice (once per device) for full coverage." + ) + + +def _ovphysx_sim_context(device: str, **kwargs): + """Wrapper around :func:`build_simulation_context` that injects OVPhysX cfg. + + PhysX tests pass ``device=device`` directly and let + :func:`build_simulation_context` build a default :class:`SimulationCfg`. + OVPhysX needs ``physics=OvPhysxCfg()`` set on the cfg so the manager + dispatches to OVPhysX rather than PhysX, so we build the cfg here and + pass it through. ``gravity_enabled`` is consumed locally (it is ignored + by ``build_simulation_context`` once a ``sim_cfg`` is provided). + ``add_ground_plane``, ``auto_add_lighting``, and other kwargs continue + to flow through ``build_simulation_context`` as before. + """ + dt = kwargs.pop("dt", 1.0 / 60.0) + gravity_enabled = kwargs.pop("gravity_enabled", True) + gravity = (0.0, 0.0, -9.81) if gravity_enabled else (0.0, 0.0, 0.0) + sim_cfg = SimulationCfg(physics=OvPhysxCfg(), device=device, dt=dt, gravity=gravity) + return build_simulation_context(device=device, sim_cfg=sim_cfg, **kwargs) + + +def generate_articulation_cfg( + articulation_type: str, + stiffness: float | None = 10.0, + damping: float | None = 2.0, + velocity_limit: float | None = None, + effort_limit: float | None = None, + velocity_limit_sim: float | None = None, + effort_limit_sim: float | None = None, +) -> ArticulationCfg: + """Generate an articulation configuration. + + Args: + articulation_type: Type of articulation to generate. + It should be one of: "humanoid", "panda", "anymal", "shadow_hand", "single_joint_implicit", + "single_joint_explicit". + stiffness: Stiffness value for the articulation's actuators. Only currently used for "humanoid". + Defaults to 10.0. + damping: Damping value for the articulation's actuators. Only currently used for "humanoid". + Defaults to 2.0. + velocity_limit: Velocity limit for the actuators. Only currently used for "single_joint_implicit" + and "single_joint_explicit". + effort_limit: Effort limit for the actuators. Only currently used for "single_joint_implicit" + and "single_joint_explicit". + velocity_limit_sim: Velocity limit for the actuators (set into the simulation). + Only currently used for "single_joint_implicit" and "single_joint_explicit". + effort_limit_sim: Effort limit for the actuators (set into the simulation). + Only currently used for "single_joint_implicit" and "single_joint_explicit". + + Returns: + The articulation configuration for the requested articulation type. + + """ + if articulation_type == "humanoid": + articulation_cfg = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/IsaacSim/Humanoid/humanoid_instanceable.usd" + ), + init_state=ArticulationCfg.InitialStateCfg(pos=(0.0, 0.0, 1.34)), + actuators={"body": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=stiffness, damping=damping)}, + ) + elif articulation_type == "panda": + articulation_cfg = FRANKA_PANDA_CFG + elif articulation_type == "anymal": + articulation_cfg = ANYMAL_C_CFG + elif articulation_type == "shadow_hand": + articulation_cfg = SHADOW_HAND_CFG + elif articulation_type == "single_joint_implicit": + articulation_cfg = ArticulationCfg( + # we set 80.0 default for max force because default in USD is 10e10 which makes testing annoying. + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/IsaacSim/SimpleArticulation/revolute_articulation.usd", + joint_drive_props=sim_utils.JointDrivePropertiesCfg(max_effort=80.0, max_velocity=5.0), + ), + actuators={ + "joint": ImplicitActuatorCfg( + joint_names_expr=[".*"], + effort_limit_sim=effort_limit_sim, + velocity_limit_sim=velocity_limit_sim, + effort_limit=effort_limit, + velocity_limit=velocity_limit, + stiffness=2000.0, + damping=100.0, + ), + }, + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.0, 0.0, 0.0), + joint_pos=({"RevoluteJoint": 1.5708}), + rot=(0.7071081, 0, 0, 0.7071055), + ), + ) + elif articulation_type == "single_joint_explicit": + # we set 80.0 default for max force because default in USD is 10e10 which makes testing annoying. + articulation_cfg = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/IsaacSim/SimpleArticulation/revolute_articulation.usd", + joint_drive_props=sim_utils.JointDrivePropertiesCfg(max_effort=80.0, max_velocity=5.0), + ), + actuators={ + "joint": IdealPDActuatorCfg( + joint_names_expr=[".*"], + effort_limit_sim=effort_limit_sim, + velocity_limit_sim=velocity_limit_sim, + effort_limit=effort_limit, + velocity_limit=velocity_limit, + stiffness=0.0, + damping=10.0, + ), + }, + ) + elif articulation_type == "spatial_tendon_test_asset": + # we set 80.0 default for max force because default in USD is 10e10 which makes testing annoying. + articulation_cfg = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/IsaacLab/Tests/spatial_tendons.usd", + ), + actuators={ + "joint": ImplicitActuatorCfg( + joint_names_expr=[".*"], + stiffness=2000.0, + damping=100.0, + ), + }, + ) + else: + raise ValueError( + f"Invalid articulation type: {articulation_type}, valid options are 'humanoid', 'panda', 'anymal'," + " 'shadow_hand', 'single_joint_implicit', 'single_joint_explicit' or 'spatial_tendon_test_asset'." + ) + + return articulation_cfg + + +def generate_articulation( + articulation_cfg: ArticulationCfg, num_articulations: int, device: str +) -> tuple[Articulation, torch.tensor]: + """Generate an articulation from a configuration. + + Handles the creation of the articulation, the environment prims and the articulation's environment + translations + + Args: + articulation_cfg: Articulation configuration. + num_articulations: Number of articulations to generate. + device: Device to use for the tensors. + + Returns: + The articulation and environment translations. + + """ + # Generate translations of 2.5 m in x for each articulation + translations = torch.zeros(num_articulations, 3, device=device) + translations[:, 0] = torch.arange(num_articulations) * 2.5 + + # Create Top-level Xforms, one for each articulation + for i in range(num_articulations): + sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=translations[i][:3]) + articulation = Articulation(articulation_cfg.replace(prim_path="/World/Env_.*/Robot")) + + return articulation, translations + + +@pytest.fixture +def sim(request): + """Create simulation context with the specified device.""" + device = request.getfixturevalue("device") + if "gravity_enabled" in request.fixturenames: + gravity_enabled = request.getfixturevalue("gravity_enabled") + else: + gravity_enabled = True # default to gravity enabled + if "add_ground_plane" in request.fixturenames: + add_ground_plane = request.getfixturevalue("add_ground_plane") + else: + add_ground_plane = False # default to no ground plane + with _ovphysx_sim_context( + device=device, auto_add_lighting=True, gravity_enabled=gravity_enabled, add_ground_plane=add_ground_plane + ) as sim: + sim._app_control_on_stop_handle = None + yield sim + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.isaacsim_ci +def test_initialization_floating_base_non_root(sim, num_articulations, device, add_ground_plane): + """Test initialization for a floating-base with articulation root on a rigid body. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is not fixed base + 3. All buffers have correct shapes + 4. The articulation can be simulated + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + articulation_cfg = generate_articulation_cfg(articulation_type="humanoid", stiffness=0.0, damping=0.0) + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + + # Check if articulation is initialized + assert articulation.is_initialized + # Check that is fixed base + assert not articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 21) + + # Cross-check binding shapes against cached counts. PhysX does this via + # ``root_view.max_dofs == shared_metatype.dof_count``; on OVPhysX + # ``root_view`` is the per-tensor-type bindings dict, so the equivalent + # invariant is that each per-DOF / per-link binding's shape agrees with + # the count cached on the asset. + for tt in (TT.DOF_POSITION, TT.DOF_VELOCITY, TT.DOF_STIFFNESS): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_joints + for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_bodies + # Body-name ordering check is degenerate on OVPhysX: ``body_names`` is + # sourced from binding metadata (``sample.body_names``), so the PhysX + # ``link_paths[0]`` round-trip is a no-op here and is omitted. + # -- actuator type + for actuator_name, actuator in articulation.actuators.items(): + is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) + assert actuator.is_implicit_model == is_implicit_model_cfg + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.isaacsim_ci +def test_initialization_floating_base(sim, num_articulations, device, add_ground_plane): + """Test initialization for a floating-base with articulation root on provided prim path. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is not fixed base + 3. All buffers have correct shapes + 4. The articulation can be simulated + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + articulation_cfg = generate_articulation_cfg(articulation_type="anymal", stiffness=0.0, damping=0.0) + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that floating base + assert not articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 12) + assert articulation.data.body_mass.torch.shape == (num_articulations, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (num_articulations, articulation.num_bodies, 9) + + # Cross-check binding shapes against cached counts. PhysX does this via + # ``root_view.max_dofs == shared_metatype.dof_count``; on OVPhysX + # ``root_view`` is the per-tensor-type bindings dict, so the equivalent + # invariant is that each per-DOF / per-link binding's shape agrees with + # the count cached on the asset. + for tt in (TT.DOF_POSITION, TT.DOF_VELOCITY, TT.DOF_STIFFNESS): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_joints + for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_bodies + # Body-name ordering check is degenerate on OVPhysX: ``body_names`` is + # sourced from binding metadata (``sample.body_names``), so the PhysX + # ``link_paths[0]`` round-trip is a no-op here and is omitted. + # -- actuator type + for actuator_name, actuator in articulation.actuators.items(): + is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) + assert actuator.is_implicit_model == is_implicit_model_cfg + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_initialization_fixed_base(sim, num_articulations, device): + """Test initialization for fixed base. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is fixed base + 3. All buffers have correct shapes + 4. The articulation maintains its default state + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + articulation_cfg = generate_articulation_cfg(articulation_type="panda") + articulation, translations = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that fixed base + assert articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 9) + assert articulation.data.body_mass.torch.shape == (num_articulations, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (num_articulations, articulation.num_bodies, 9) + + # Cross-check binding shapes against cached counts. PhysX does this via + # ``root_view.max_dofs == shared_metatype.dof_count``; on OVPhysX + # ``root_view`` is the per-tensor-type bindings dict, so the equivalent + # invariant is that each per-DOF / per-link binding's shape agrees with + # the count cached on the asset. + for tt in (TT.DOF_POSITION, TT.DOF_VELOCITY, TT.DOF_STIFFNESS): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_joints + for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_bodies + # Body-name ordering check is degenerate on OVPhysX: ``body_names`` is + # sourced from binding metadata (``sample.body_names``), so the PhysX + # ``link_paths[0]`` round-trip is a no-op here and is omitted. + # -- actuator type + for actuator_name, actuator in articulation.actuators.items(): + is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) + assert actuator.is_implicit_model == is_implicit_model_cfg + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + # check that the root is at the correct state - its default state as it is fixed base + default_root_pose = articulation.data.default_root_pose.torch.clone() + default_root_vel = articulation.data.default_root_vel.torch.clone() + default_root_pose[:, :3] = default_root_pose[:, :3] + translations + + torch.testing.assert_close(articulation.data.root_link_pose_w.torch, default_root_pose) + torch.testing.assert_close(articulation.data.root_com_vel_w.torch, default_root_vel) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.isaacsim_ci +def test_initialization_fixed_base_single_joint(sim, num_articulations, device, add_ground_plane): + """Test initialization for fixed base articulation with a single joint. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is fixed base + 3. All buffers have correct shapes + 4. The articulation maintains its default state + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + articulation_cfg = generate_articulation_cfg(articulation_type="single_joint_implicit") + articulation, translations = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that fixed base + assert articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 1) + assert articulation.data.body_mass.torch.shape == (num_articulations, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (num_articulations, articulation.num_bodies, 9) + + # Cross-check binding shapes against cached counts. PhysX does this via + # ``root_view.max_dofs == shared_metatype.dof_count``; on OVPhysX + # ``root_view`` is the per-tensor-type bindings dict, so the equivalent + # invariant is that each per-DOF / per-link binding's shape agrees with + # the count cached on the asset. + for tt in (TT.DOF_POSITION, TT.DOF_VELOCITY, TT.DOF_STIFFNESS): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_joints + for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_bodies + # Body-name ordering check is degenerate on OVPhysX: ``body_names`` is + # sourced from binding metadata (``sample.body_names``), so the PhysX + # ``link_paths[0]`` round-trip is a no-op here and is omitted. + # -- actuator type + for actuator_name, actuator in articulation.actuators.items(): + is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) + assert actuator.is_implicit_model == is_implicit_model_cfg + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + # check that the root is at the correct state - its default state as it is fixed base + default_root_pose = articulation.data.default_root_pose.torch.clone() + default_root_vel = articulation.data.default_root_vel.torch.clone() + default_root_pose[:, :3] = default_root_pose[:, :3] + translations + + torch.testing.assert_close(articulation.data.root_link_pose_w.torch, default_root_pose) + torch.testing.assert_close(articulation.data.root_com_vel_w.torch, default_root_vel) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_initialization_hand_with_tendons(sim, num_articulations, device): + """Test initialization for fixed base articulated hand with tendons. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is fixed base + 3. All buffers have correct shapes + 4. The articulation can be simulated + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + articulation_cfg = generate_articulation_cfg(articulation_type="shadow_hand") + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that fixed base + assert articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 24) + assert articulation.data.body_mass.torch.shape == (num_articulations, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (num_articulations, articulation.num_bodies, 9) + + # Cross-check binding shapes against cached counts. See the equivalent + # block in test_initialization_fixed_base_single_joint for why the verbatim + # PhysX ``root_view.max_dofs == shared_metatype.dof_count`` identity is + # replaced with binding-shape checks on OVPhysX. + for tt in (TT.DOF_POSITION, TT.DOF_VELOCITY, TT.DOF_STIFFNESS): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_joints + for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_bodies + # -- actuator type + for actuator_name, actuator in articulation.actuators.items(): + is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) + assert actuator.is_implicit_model == is_implicit_model_cfg + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.isaacsim_ci +@pytest.mark.xfail(reason=_OMNI_PHYSX_SCHEMAS_GAP_REASON, strict=False) +def test_initialization_floating_base_made_fixed_base(sim, num_articulations, device, add_ground_plane): + """Test initialization for a floating-base articulation made fixed-base using schema properties. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is fixed base after modification + 3. All buffers have correct shapes + 4. The articulation maintains its default state + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type="anymal").copy() + # Fix root link by making it kinematic + articulation_cfg.spawn.articulation_props.fix_root_link = True + articulation, translations = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that is fixed base + assert articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 12) + + # Cross-check binding shapes against cached counts. PhysX does this via + # ``root_view.max_dofs == shared_metatype.dof_count``; on OVPhysX + # ``root_view`` is the per-tensor-type bindings dict, so the equivalent + # invariant is that each per-DOF / per-link binding's shape agrees with + # the count cached on the asset. + for tt in (TT.DOF_POSITION, TT.DOF_VELOCITY, TT.DOF_STIFFNESS): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_joints + for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_bodies + # Body-name ordering check is degenerate on OVPhysX: ``body_names`` is + # sourced from binding metadata (``sample.body_names``), so the PhysX + # ``link_paths[0]`` round-trip is a no-op here and is omitted. + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + # check that the root is at the correct state - its default state as it is fixed base + default_root_pose = articulation.data.default_root_pose.torch.clone() + default_root_vel = articulation.data.default_root_vel.torch.clone() + default_root_pose[:, :3] = default_root_pose[:, :3] + translations + + torch.testing.assert_close(articulation.data.root_link_pose_w.torch, default_root_pose) + torch.testing.assert_close(articulation.data.root_com_vel_w.torch, default_root_vel) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.isaacsim_ci +def test_initialization_fixed_base_made_floating_base(sim, num_articulations, device, add_ground_plane): + """Test initialization for fixed base made floating-base using schema properties. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is floating base after modification + 3. All buffers have correct shapes + 4. The articulation can be simulated + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type="panda") + # Unfix root link by making it non-kinematic + articulation_cfg.spawn.articulation_props.fix_root_link = False + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that is floating base + assert not articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 9) + + # Cross-check binding shapes against cached counts. PhysX does this via + # ``root_view.max_dofs == shared_metatype.dof_count``; on OVPhysX + # ``root_view`` is the per-tensor-type bindings dict, so the equivalent + # invariant is that each per-DOF / per-link binding's shape agrees with + # the count cached on the asset. + for tt in (TT.DOF_POSITION, TT.DOF_VELOCITY, TT.DOF_STIFFNESS): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_joints + for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_bodies + # Body-name ordering check is degenerate on OVPhysX: ``body_names`` is + # sourced from binding metadata (``sample.body_names``), so the PhysX + # ``link_paths[0]`` round-trip is a no-op here and is omitted. + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.isaacsim_ci +def test_out_of_range_default_joint_pos(sim, num_articulations, device, add_ground_plane): + """Test that the default joint position from configuration is out of range. + + This test verifies that: + 1. The articulation fails to initialize when joint positions are out of range + 2. The error is properly handled + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + # Create articulation + articulation_cfg = generate_articulation_cfg(articulation_type="panda").copy() + articulation_cfg.init_state.joint_pos = { + "panda_joint1": 10.0, + "panda_joint[2, 4]": -20.0, + } + + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + with pytest.raises(ValueError): + sim.reset() + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_out_of_range_default_joint_vel(sim, device): + """Test that the default joint velocity from configuration is out of range. + + This test verifies that: + 1. The articulation fails to initialize when joint velocities are out of range + 2. The error is properly handled + """ + articulation_cfg = FRANKA_PANDA_CFG.replace(prim_path="/World/Robot") + articulation_cfg.init_state.joint_vel = { + "panda_joint1": 100.0, + "panda_joint[2, 4]": -60.0, + } + articulation = Articulation(articulation_cfg) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + with pytest.raises(ValueError): + sim.reset() + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.isaacsim_ci +def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane): + """Test write_joint_limits_to_sim API and when default pos falls outside of the new limits. + + This test verifies that: + 1. Joint limits can be set correctly + 2. Default positions are preserved when setting new limits + 3. Joint limits can be set with indexing + 4. Invalid joint positions are properly handled + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + # Create articulation + articulation_cfg = generate_articulation_cfg(articulation_type="panda") + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device) + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + + # Get current default joint pos + default_joint_pos = articulation._data.default_joint_pos.torch.clone() + + # Set new joint limits + limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) + limits[..., 0] = (torch.rand(num_articulations, articulation.num_joints, device=device) + 5.0) * -1.0 + limits[..., 1] = torch.rand(num_articulations, articulation.num_joints, device=device) + 5.0 + articulation.write_joint_position_limit_to_sim_index(limits=limits) + + # Check new limits are in place + torch.testing.assert_close(articulation._data.joint_pos_limits.torch, limits) + torch.testing.assert_close(articulation._data.default_joint_pos.torch, default_joint_pos) + + # Set new joint limits with indexing + env_ids = torch.arange(1, device=device, dtype=torch.int32) + joint_ids = torch.arange(2, device=device, dtype=torch.int32) + limits = torch.zeros(env_ids.shape[0], joint_ids.shape[0], 2, device=device) + limits[..., 0] = (torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) + 5.0) * -1.0 + limits[..., 1] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) + 5.0 + articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) + + # Check new limits are in place + torch.testing.assert_close(articulation._data.joint_pos_limits.torch[env_ids][:, joint_ids], limits) + torch.testing.assert_close(articulation._data.default_joint_pos.torch, default_joint_pos) + + # Set new joint limits that invalidate default joint pos + limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) + limits[..., 0] = torch.rand(num_articulations, articulation.num_joints, device=device) * -0.1 + limits[..., 1] = torch.rand(num_articulations, articulation.num_joints, device=device) * 0.1 + articulation.write_joint_position_limit_to_sim_index(limits=limits) + + # Check if all values are within the bounds + default_joint_pos_torch = articulation._data.default_joint_pos.torch + within_bounds = (default_joint_pos_torch >= limits[..., 0]) & (default_joint_pos_torch <= limits[..., 1]) + assert torch.all(within_bounds) + + # Set new joint limits that invalidate default joint pos with indexing + limits = torch.zeros(env_ids.shape[0], joint_ids.shape[0], 2, device=device) + limits[..., 0] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) * -0.1 + limits[..., 1] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) * 0.1 + articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) + + # Check if all values are within the bounds + default_joint_pos_torch = articulation._data.default_joint_pos.torch + within_bounds = (default_joint_pos_torch[env_ids][:, joint_ids] >= limits[..., 0]) & ( + default_joint_pos_torch[env_ids][:, joint_ids] <= limits[..., 1] + ) + assert torch.all(within_bounds) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +def test_joint_effort_limits(sim, num_articulations, device, add_ground_plane): + """Validate joint effort limits via joint_effort_out_of_limit().""" + # Create articulation + articulation_cfg = generate_articulation_cfg(articulation_type="panda") + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device) + + # Minimal env wrapper exposing scene["robot"] + class _Env: + def __init__(self, art): + self.scene = {"robot": art} + + env = _Env(articulation) + robot_all = SceneEntityCfg(name="robot") + + sim.reset() + assert articulation.is_initialized + + # Case A: no clipping → should NOT terminate + articulation._data.computed_torque.torch.zero_() + articulation._data.applied_torque.torch.zero_() + out = joint_effort_out_of_limit(env, robot_all) # [N] + assert torch.all(~out) + + # Case B: simulate clipping → should terminate + articulation._data.computed_torque.torch.fill_(100.0) # pretend controller commanded 100 + articulation._data.applied_torque.torch.fill_(50.0) # pretend actuator clipped to 50 + out = joint_effort_out_of_limit(env, robot_all) # [N] + assert torch.all(out) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_external_force_buffer(sim, num_articulations, device): + """Test if external force buffer correctly updates in the force value is zero case. + + This test verifies that: + 1. External forces can be applied correctly + 2. Force buffers are updated properly + 3. Zero forces are handled correctly + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type="anymal") + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + + # play the simulator + sim.reset() + + # find bodies to apply the force + body_ids, _ = articulation.find_bodies("base") + + # reset root state + articulation.write_root_pose_to_sim_index(root_pose=articulation.data.default_root_pose.torch.clone()) + articulation.write_root_velocity_to_sim_index(root_velocity=articulation.data.default_root_vel.torch.clone()) + + # reset dof state + joint_pos, joint_vel = ( + articulation.data.default_joint_pos.torch, + articulation.data.default_joint_vel.torch, + ) + articulation.write_joint_position_to_sim_index(position=joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) + + # reset articulation + articulation.reset() + + # perform simulation + for step in range(5): + # initiate force tensor + external_wrench_b = torch.zeros(articulation.num_instances, len(body_ids), 6, device=sim.device) + + if step == 0 or step == 3: + # set a non-zero force + force = 1 + else: + # set a zero force + force = 0 + + # set force value + external_wrench_b[:, :, 0] = force + external_wrench_b[:, :, 3] = force + + # apply force + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + body_ids=body_ids, + ) + + # check if the articulation's force and torque buffers are correctly updated + for i in range(num_articulations): + assert articulation.permanent_wrench_composer.composed_force.torch[i, 0, 0].item() == force + assert articulation.permanent_wrench_composer.composed_torque.torch[i, 0, 0].item() == force + + # Check if the instantaneous wrench is correctly added to the permanent wrench + articulation.instantaneous_wrench_composer.add_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + body_ids=body_ids, + ) + + # apply action to the articulation + articulation.set_joint_position_target_index(target=articulation.data.default_joint_pos.torch.clone()) + articulation.write_data_to_sim() + + # perform step + sim.step() + + # update buffers + articulation.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_external_force_on_single_body(sim, num_articulations, device): + """Test application of external force on the base of the articulation. + + This test verifies that: + 1. External forces can be applied to specific bodies + 2. The forces affect the articulation's motion correctly + 3. The articulation responds to the forces as expected + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type="anymal") + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + # Play the simulator + sim.reset() + + # Find bodies to apply the force + body_ids, _ = articulation.find_bodies("base") + # Sample a large force + external_wrench_b = torch.zeros(articulation.num_instances, len(body_ids), 6, device=sim.device) + external_wrench_b[..., 1] = 1000.0 + + # Now we are ready! + for _ in range(5): + # reset root state + articulation.write_root_pose_to_sim_index(root_pose=articulation.data.default_root_pose.torch.clone()) + articulation.write_root_velocity_to_sim_index(root_velocity=articulation.data.default_root_vel.torch.clone()) + # reset dof state + joint_pos, joint_vel = ( + articulation.data.default_joint_pos.torch, + articulation.data.default_joint_vel.torch, + ) + articulation.write_joint_position_to_sim_index(position=joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) + # reset articulation + articulation.reset() + # apply force + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], body_ids=body_ids + ) + # perform simulation + for _ in range(100): + # apply action to the articulation + articulation.set_joint_position_target_index(target=articulation.data.default_joint_pos.torch.clone()) + articulation.write_data_to_sim() + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + # check condition that the articulations have fallen down + for i in range(num_articulations): + assert articulation.data.root_pos_w.torch[i, 2].item() < 0.2 + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_external_force_on_single_body_at_position(sim, num_articulations, device): + """Test application of external force on the base of the articulation at a given position. + + This test verifies that: + 1. External forces can be applied to specific bodies at a given position + 2. External forces can be applied to specific bodies in the global frame + 3. External forces are calculated and composed correctly + 4. The forces affect the articulation's motion correctly + 5. The articulation responds to the forces as expected + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type="anymal") + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + # Play the simulator + sim.reset() + + # Find bodies to apply the force + body_ids, _ = articulation.find_bodies("base") + # Sample a large force + external_wrench_b = torch.zeros(articulation.num_instances, len(body_ids), 6, device=sim.device) + external_wrench_b[..., 2] = 500.0 + external_wrench_positions_b = torch.zeros(articulation.num_instances, len(body_ids), 3, device=sim.device) + external_wrench_positions_b[..., 1] = 1.0 + + desired_force = torch.zeros(articulation.num_instances, len(body_ids), 3, device=sim.device) + desired_force[..., 2] = 1000.0 + desired_torque = torch.zeros(articulation.num_instances, len(body_ids), 3, device=sim.device) + desired_torque[..., 0] = 1000.0 + + # Now we are ready! + for i in range(5): + # reset root state + root_pose = articulation.data.default_root_pose.torch.clone() + root_pose[0, 0] = 2.5 # space them apart by 2.5m + + articulation.write_root_pose_to_sim_index(root_pose=root_pose) + articulation.write_root_velocity_to_sim_index(root_velocity=articulation.data.default_root_vel.torch.clone()) + # reset dof state + joint_pos, joint_vel = ( + articulation.data.default_joint_pos.torch, + articulation.data.default_joint_vel.torch, + ) + articulation.write_joint_position_to_sim_index(position=joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) + # reset articulation + articulation.reset() + # apply force + is_global = False + + if i % 2 == 0: + body_com_pos_w = articulation.data.body_com_pos_w.torch[:, body_ids, :3] + # is_global = True + external_wrench_positions_b[..., 0] = 0.0 + external_wrench_positions_b[..., 1] = 1.0 + external_wrench_positions_b[..., 2] = 0.0 + external_wrench_positions_b += body_com_pos_w + else: + external_wrench_positions_b[..., 0] = 0.0 + external_wrench_positions_b[..., 1] = 1.0 + external_wrench_positions_b[..., 2] = 0.0 + + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=external_wrench_positions_b, + body_ids=body_ids, + is_global=is_global, + ) + articulation.permanent_wrench_composer.add_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=external_wrench_positions_b, + body_ids=body_ids, + is_global=is_global, + ) + # perform simulation + for _ in range(100): + # apply action to the articulation + articulation.set_joint_position_target_index(target=articulation.data.default_joint_pos.torch.clone()) + articulation.write_data_to_sim() + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + # check condition that the articulations have fallen down + for i in range(num_articulations): + assert articulation.data.root_pos_w.torch[i, 2].item() < 0.2 + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_external_force_on_multiple_bodies(sim, num_articulations, device): + """Test application of external force on the legs of the articulation. + + This test verifies that: + 1. External forces can be applied to multiple bodies + 2. The forces affect the articulation's motion correctly + 3. The articulation responds to the forces as expected + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type="anymal") + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + + # Play the simulator + sim.reset() + + # Find bodies to apply the force + body_ids, _ = articulation.find_bodies(".*_SHANK") + # Sample a large force + external_wrench_b = torch.zeros(articulation.num_instances, len(body_ids), 6, device=sim.device) + external_wrench_b[..., 1] = 100.0 + + # Now we are ready! + for _ in range(5): + # reset root state + articulation.write_root_pose_to_sim_index(root_pose=articulation.data.default_root_pose.torch.clone()) + articulation.write_root_velocity_to_sim_index(root_velocity=articulation.data.default_root_vel.torch.clone()) + # reset dof state + joint_pos, joint_vel = ( + articulation.data.default_joint_pos.torch, + articulation.data.default_joint_vel.torch, + ) + articulation.write_joint_position_to_sim_index(position=joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) + # reset articulation + articulation.reset() + # apply force + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], body_ids=body_ids + ) + # perform simulation + for _ in range(100): + # apply action to the articulation + articulation.set_joint_position_target_index(target=articulation.data.default_joint_pos.torch.clone()) + articulation.write_data_to_sim() + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + # check condition + for i in range(num_articulations): + # since there is a moment applied on the articulation, the articulation should rotate + assert articulation.data.root_ang_vel_w.torch[i, 2].item() > 0.1 + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_external_force_on_multiple_bodies_at_position(sim, num_articulations, device): + """Test application of external force on the legs of the articulation at a given position. + + This test verifies that: + 1. External forces can be applied to multiple bodies at a given position + 2. External forces can be applied to multiple bodies in the global frame + 3. External forces are calculated and composed correctly + 4. The forces affect the articulation's motion correctly + 5. The articulation responds to the forces as expected + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type="anymal") + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + + # Play the simulator + sim.reset() + + # Find bodies to apply the force + body_ids, _ = articulation.find_bodies(".*_SHANK") + # Sample a large force + external_wrench_b = torch.zeros(articulation.num_instances, len(body_ids), 6, device=sim.device) + external_wrench_b[..., 2] = 500.0 + external_wrench_positions_b = torch.zeros(articulation.num_instances, len(body_ids), 3, device=sim.device) + external_wrench_positions_b[..., 1] = 1.0 + + desired_force = torch.zeros(articulation.num_instances, len(body_ids), 3, device=sim.device) + desired_force[..., 2] = 1000.0 + desired_torque = torch.zeros(articulation.num_instances, len(body_ids), 3, device=sim.device) + desired_torque[..., 0] = 1000.0 + + # Now we are ready! + for i in range(5): + # reset root state + articulation.write_root_pose_to_sim_index(root_pose=articulation.data.default_root_pose.torch.clone()) + articulation.write_root_velocity_to_sim_index(root_velocity=articulation.data.default_root_vel.torch.clone()) + # reset dof state + joint_pos, joint_vel = ( + articulation.data.default_joint_pos.torch, + articulation.data.default_joint_vel.torch, + ) + articulation.write_joint_position_to_sim_index(position=joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) + # reset articulation + articulation.reset() + + is_global = False + if i % 2 == 0: + body_com_pos_w = articulation.data.body_com_pos_w.torch[:, body_ids, :3] + is_global = True + external_wrench_positions_b[..., 0] = 0.0 + external_wrench_positions_b[..., 1] = 1.0 + external_wrench_positions_b[..., 2] = 0.0 + external_wrench_positions_b += body_com_pos_w + else: + external_wrench_positions_b[..., 0] = 0.0 + external_wrench_positions_b[..., 1] = 1.0 + external_wrench_positions_b[..., 2] = 0.0 + + # apply force + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=external_wrench_positions_b, + body_ids=body_ids, + is_global=is_global, + ) + articulation.permanent_wrench_composer.add_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=external_wrench_positions_b, + body_ids=body_ids, + is_global=is_global, + ) + # perform simulation + for _ in range(100): + # apply action to the articulation + articulation.set_joint_position_target_index(target=articulation.data.default_joint_pos.torch.clone()) + articulation.write_data_to_sim() + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + # check condition + for i in range(num_articulations): + # since there is a moment applied on the articulation, the articulation should rotate + assert torch.abs(articulation.data.root_ang_vel_w.torch[i, 2]).item() > 0.1 + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_loading_gains_from_usd(sim, num_articulations, device): + """Test that gains are loaded from USD file if actuator model has them as None. + + This test verifies that: + 1. Gains are loaded correctly from USD file + 2. Default gains are applied when not specified + 3. The gains match the expected values + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type="humanoid", stiffness=None, damping=None) + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + + # Play sim + sim.reset() + + # Expected gains + # -- Stiffness values + expected_stiffness = { + ".*_waist.*": 20.0, + ".*_upper_arm.*": 10.0, + "pelvis": 10.0, + ".*_lower_arm": 2.0, + ".*_thigh:0": 10.0, + ".*_thigh:1": 20.0, + ".*_thigh:2": 10.0, + ".*_shin": 5.0, + ".*_foot.*": 2.0, + } + indices_list, _, values_list = string_utils.resolve_matching_names_values( + expected_stiffness, articulation.joint_names + ) + expected_stiffness = torch.zeros(articulation.num_instances, articulation.num_joints, device=articulation.device) + expected_stiffness[:, indices_list] = torch.tensor(values_list, device=articulation.device) + # -- Damping values + expected_damping = { + ".*_waist.*": 5.0, + ".*_upper_arm.*": 5.0, + "pelvis": 5.0, + ".*_lower_arm": 1.0, + ".*_thigh:0": 5.0, + ".*_thigh:1": 5.0, + ".*_thigh:2": 5.0, + ".*_shin": 0.1, + ".*_foot.*": 1.0, + } + indices_list, _, values_list = string_utils.resolve_matching_names_values( + expected_damping, articulation.joint_names + ) + expected_damping = torch.zeros_like(expected_stiffness) + expected_damping[:, indices_list] = torch.tensor(values_list, device=articulation.device) + + # Check that gains are loaded from USD file + torch.testing.assert_close(articulation.actuators["body"].stiffness, expected_stiffness) + torch.testing.assert_close(articulation.actuators["body"].damping, expected_damping) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.isaacsim_ci +def test_setting_gains_from_cfg(sim, num_articulations, device, add_ground_plane): + """Test that gains are loaded from the configuration correctly. + + This test verifies that: + 1. Gains are loaded correctly from configuration + 2. The gains match the expected values + 3. The gains are applied correctly to the actuators + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type="humanoid") + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=sim.device + ) + + # Play sim + sim.reset() + + # Expected gains + expected_stiffness = torch.full( + (articulation.num_instances, articulation.num_joints), 10.0, device=articulation.device + ) + expected_damping = torch.full_like(expected_stiffness, 2.0) + + # Check that gains are loaded from USD file + torch.testing.assert_close(articulation.actuators["body"].stiffness, expected_stiffness) + torch.testing.assert_close(articulation.actuators["body"].damping, expected_damping) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_setting_gains_from_cfg_dict(sim, num_articulations, device): + """Test that gains are loaded from the configuration dictionary correctly. + + This test verifies that: + 1. Gains are loaded correctly from configuration dictionary + 2. The gains match the expected values + 3. The gains are applied correctly to the actuators + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type="humanoid") + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=sim.device + ) + # Play sim + sim.reset() + + # Expected gains + expected_stiffness = torch.full( + (articulation.num_instances, articulation.num_joints), 10.0, device=articulation.device + ) + expected_damping = torch.full_like(expected_stiffness, 2.0) + + # Check that gains are loaded from USD file + torch.testing.assert_close(articulation.actuators["body"].stiffness, expected_stiffness) + torch.testing.assert_close(articulation.actuators["body"].damping, expected_damping) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("vel_limit_sim", [1e5, None]) +@pytest.mark.parametrize("vel_limit", [1e2, None]) +@pytest.mark.parametrize("add_ground_plane", [False]) +@pytest.mark.isaacsim_ci +def test_setting_velocity_limit_implicit(sim, num_articulations, device, vel_limit_sim, vel_limit, add_ground_plane): + """Test setting of velocity limit for implicit actuators. + + This test verifies that: + 1. Velocity limits can be set correctly for implicit actuators + 2. The limits are applied correctly to the simulation + 3. The limits are handled correctly when both sim and non-sim limits are set + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + vel_limit_sim: The velocity limit to set in simulation + vel_limit: The velocity limit to set in actuator + """ + # create simulation + articulation_cfg = generate_articulation_cfg( + articulation_type="single_joint_implicit", + velocity_limit_sim=vel_limit_sim, + velocity_limit=vel_limit, + ) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, + num_articulations=num_articulations, + device=device, + ) + # Play sim + if vel_limit_sim is not None and vel_limit is not None: + with pytest.raises(ValueError): + sim.reset() + return + sim.reset() + + # read the values set into the simulation + physx_vel_limit = _read_binding_to_torch(articulation, TT.DOF_MAX_VELOCITY, device) + # check data buffer + torch.testing.assert_close(articulation.data.joint_velocity_limits.torch, physx_vel_limit) + # check actuator has simulation velocity limit + torch.testing.assert_close(articulation.actuators["joint"].velocity_limit_sim, physx_vel_limit) + # check that both values match for velocity limit + torch.testing.assert_close( + articulation.actuators["joint"].velocity_limit_sim, + articulation.actuators["joint"].velocity_limit, + ) + + if vel_limit_sim is None: + # Case 2: both velocity limit and velocity limit sim are not set + # This is the case where the velocity limit keeps its USD default value + # Case 3: velocity limit sim is not set but velocity limit is set + # For backwards compatibility, we do not set velocity limit to simulation + # Thus, both default to USD default value. + limit = articulation_cfg.spawn.joint_drive_props.max_joint_velocity + else: + # Case 4: only velocity limit sim is set + # In this case, the velocity limit is set to the USD value + limit = vel_limit_sim + + # check max velocity is what we set + expected_velocity_limit = torch.full_like(physx_vel_limit, limit) + torch.testing.assert_close(physx_vel_limit, expected_velocity_limit) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("vel_limit_sim", [1e5, None]) +@pytest.mark.parametrize("vel_limit", [1e2, None]) +@pytest.mark.isaacsim_ci +def test_setting_velocity_limit_explicit(sim, num_articulations, device, vel_limit_sim, vel_limit): + """Test setting of velocity limit for explicit actuators.""" + articulation_cfg = generate_articulation_cfg( + articulation_type="single_joint_explicit", + velocity_limit_sim=vel_limit_sim, + velocity_limit=vel_limit, + ) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, + num_articulations=num_articulations, + device=device, + ) + # Play sim + sim.reset() + + # collect limit init values + physx_vel_limit = _read_binding_to_torch(articulation, TT.DOF_MAX_VELOCITY, device) + actuator_vel_limit = articulation.actuators["joint"].velocity_limit + actuator_vel_limit_sim = articulation.actuators["joint"].velocity_limit_sim + + # check data buffer for joint_velocity_limits_sim + torch.testing.assert_close(articulation.data.joint_velocity_limits.torch, physx_vel_limit) + # check actuator velocity_limit_sim is set to physx + torch.testing.assert_close(actuator_vel_limit_sim, physx_vel_limit) + + if vel_limit is not None: + expected_actuator_vel_limit = torch.full( + (articulation.num_instances, articulation.num_joints), + vel_limit, + device=articulation.device, + ) + # check actuator is set + torch.testing.assert_close(actuator_vel_limit, expected_actuator_vel_limit) + # check physx is not velocity_limit + assert not torch.allclose(actuator_vel_limit, physx_vel_limit) + else: + # check actuator velocity_limit is the same as the PhysX default + torch.testing.assert_close(actuator_vel_limit, physx_vel_limit) + + # simulation velocity limit is set to USD value unless user overrides + if vel_limit_sim is not None: + limit = vel_limit_sim + else: + limit = articulation_cfg.spawn.joint_drive_props.max_joint_velocity + # check physx is set to expected value + expected_vel_limit = torch.full_like(physx_vel_limit, limit) + torch.testing.assert_close(physx_vel_limit, expected_vel_limit) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("effort_limit_sim", [1e5, None]) +@pytest.mark.parametrize("effort_limit", [1e2, 80.0, None]) +@pytest.mark.isaacsim_ci +def test_setting_effort_limit_implicit(sim, num_articulations, device, effort_limit_sim, effort_limit): + """Test setting of effort limit for implicit actuators. + + This test verifies the effort limit resolution logic for actuator models implemented in :class:`ActuatorBase`: + - Case 1: If USD value == actuator config value: values match correctly + - Case 2: If USD value != actuator config value: actuator config value is used + - Case 3: If actuator config value is None: USD value is used as default + """ + articulation_cfg = generate_articulation_cfg( + articulation_type="single_joint_implicit", + effort_limit_sim=effort_limit_sim, + effort_limit=effort_limit, + ) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, + num_articulations=num_articulations, + device=device, + ) + # Play sim + if effort_limit_sim is not None and effort_limit is not None: + with pytest.raises(ValueError): + sim.reset() + return + sim.reset() + + # obtain the physx effort limits + physx_effort_limit = _read_binding_to_torch(articulation, TT.DOF_MAX_FORCE, device) + + # check that the two are equivalent + torch.testing.assert_close( + articulation.actuators["joint"].effort_limit_sim, + articulation.actuators["joint"].effort_limit, + ) + torch.testing.assert_close(articulation.actuators["joint"].effort_limit_sim, physx_effort_limit) + + # decide the limit based on what is set + if effort_limit_sim is None and effort_limit is None: + limit = articulation_cfg.spawn.joint_drive_props.max_force + elif effort_limit_sim is not None and effort_limit is None: + limit = effort_limit_sim + elif effort_limit_sim is None and effort_limit is not None: + limit = effort_limit + + # check that the max force is what we set + expected_effort_limit = torch.full_like(physx_effort_limit, limit) + torch.testing.assert_close(physx_effort_limit, expected_effort_limit) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("effort_limit_sim", [1e5, None]) +@pytest.mark.parametrize("effort_limit", [80.0, 1e2, None]) +@pytest.mark.isaacsim_ci +def test_setting_effort_limit_explicit(sim, num_articulations, device, effort_limit_sim, effort_limit): + """Test setting of effort limit for explicit actuators. + + This test verifies the effort limit resolution logic for actuator models implemented in :class:`ActuatorBase`: + - Case 1: If USD value == actuator config value: values match correctly + - Case 2: If USD value != actuator config value: actuator config value is used + - Case 3: If actuator config value is None: USD value is used as default + + """ + + articulation_cfg = generate_articulation_cfg( + articulation_type="single_joint_explicit", + effort_limit_sim=effort_limit_sim, + effort_limit=effort_limit, + ) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, + num_articulations=num_articulations, + device=device, + ) + # Play sim + sim.reset() + + # usd default effort limit is set to 80 + usd_default_effort_limit = 80.0 + + # collect limit init values + physx_effort_limit = _read_binding_to_torch(articulation, TT.DOF_MAX_FORCE, device) + actuator_effort_limit = articulation.actuators["joint"].effort_limit + actuator_effort_limit_sim = articulation.actuators["joint"].effort_limit_sim + + # check actuator effort_limit_sim is set to physx + torch.testing.assert_close(actuator_effort_limit_sim, physx_effort_limit) + + if effort_limit is not None: + expected_actuator_effort_limit = torch.full_like(actuator_effort_limit, effort_limit) + # check actuator is set + torch.testing.assert_close(actuator_effort_limit, expected_actuator_effort_limit) + + # check physx effort limit does not match the one explicit actuator has + assert not (torch.allclose(actuator_effort_limit, physx_effort_limit)) + else: + # When effort_limit is None, actuator should use USD default values + expected_actuator_effort_limit = torch.full_like(physx_effort_limit, usd_default_effort_limit) + torch.testing.assert_close(actuator_effort_limit, expected_actuator_effort_limit) + + # when using explicit actuators, the limits are set to high unless user overrides + if effort_limit_sim is not None: + limit = effort_limit_sim + else: + limit = ActuatorBase._DEFAULT_MAX_EFFORT_SIM # type: ignore + # check physx internal value matches the expected sim value + expected_effort_limit = torch.full_like(physx_effort_limit, limit) + torch.testing.assert_close(actuator_effort_limit_sim, expected_effort_limit) + torch.testing.assert_close(physx_effort_limit, expected_effort_limit) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_reset(sim, num_articulations, device): + """Test that reset method works properly.""" + articulation_cfg = generate_articulation_cfg(articulation_type="humanoid") + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=device + ) + + # Play the simulator + sim.reset() + + # Now we are ready! + # reset articulation + articulation.reset() + + # Reset should zero external forces and torques + assert not articulation._instantaneous_wrench_composer.active + assert not articulation._permanent_wrench_composer.active + assert torch.count_nonzero(articulation._instantaneous_wrench_composer.composed_force.torch) == 0 + assert torch.count_nonzero(articulation._instantaneous_wrench_composer.composed_torque.torch) == 0 + assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_force.torch) == 0 + assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_torque.torch) == 0 + + if num_articulations > 1: + num_bodies = articulation.num_bodies + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=torch.ones((num_articulations, num_bodies, 3), device=device), + torques=torch.ones((num_articulations, num_bodies, 3), device=device), + ) + articulation.instantaneous_wrench_composer.add_forces_and_torques_index( + forces=torch.ones((num_articulations, num_bodies, 3), device=device), + torques=torch.ones((num_articulations, num_bodies, 3), device=device), + ) + articulation.reset(env_ids=torch.tensor([0], device=device)) + assert articulation._instantaneous_wrench_composer.active + assert articulation._permanent_wrench_composer.active + assert torch.count_nonzero(articulation._instantaneous_wrench_composer.composed_force.torch) == num_bodies * 3 + assert torch.count_nonzero(articulation._instantaneous_wrench_composer.composed_torque.torch) == num_bodies * 3 + assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_force.torch) == num_bodies * 3 + assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_torque.torch) == num_bodies * 3 + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.isaacsim_ci +def test_apply_joint_command(sim, num_articulations, device, add_ground_plane): + """Test applying of joint position target functions correctly for a robotic arm.""" + articulation_cfg = generate_articulation_cfg(articulation_type="panda") + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=device + ) + + # Play the simulator + sim.reset() + + for _ in range(100): + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + # reset dof state + joint_pos = articulation.data.default_joint_pos.torch.clone() + joint_pos[:, 3] = 0.0 + + # apply action to the articulation + articulation.set_joint_position_target_index(target=joint_pos) + articulation.write_data_to_sim() + + for _ in range(100): + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + # Check that current joint position is not the same as default joint position, meaning + # the articulation moved. We can't check that it reached its desired joint position as the gains + # are not properly tuned + assert not torch.allclose(articulation.data.joint_pos.torch, joint_pos) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("with_offset", [True, False]) +@pytest.mark.isaacsim_ci +def test_body_root_state(sim, num_articulations, device, with_offset): + """Test for reading the `body_state_w` property. + + This test verifies that: + 1. Body states can be read correctly + 2. States are correct with and without offsets + 3. States are consistent across different devices + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + with_offset: Whether to test with offset + """ + sim._app_control_on_stop_handle = None + articulation_cfg = generate_articulation_cfg(articulation_type="single_joint_implicit") + articulation, env_pos = generate_articulation(articulation_cfg, num_articulations, device) + env_idx = torch.tensor([x for x in range(num_articulations)], device=device, dtype=torch.int32) + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10, "Possible reference leak for articulation" + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized, "Articulation is not initialized" + # Check that fixed base + assert articulation.is_fixed_base, "Articulation is not a fixed base" + + # Resolve body indices by name (ordering may differ across physics backends) + root_idx = articulation.body_names.index("CenterPivot") + arm_idx = articulation.body_names.index("Arm") + + # change center of mass offset from link frame + if with_offset: + offset = [0.5, 0.0, 0.0] + else: + offset = [0.0, 0.0, 0.0] + + # create com offsets — apply offset to the Arm body + num_bodies = articulation.num_bodies + com = _read_binding_to_torch(articulation, TT.BODY_COM_POSE, device) + link_offset = [1.0, 0.0, 0.0] # the offset from CenterPivot to Arm frames + new_com = torch.tensor(offset, device=device).repeat(num_articulations, 1, 1) + com[:, arm_idx, :3] = new_com.squeeze(-2) + # PhysX uses ``root_view.set_coms``; OVPhysX wraps the wheel + # ``BODY_COM_POSE`` write in :meth:`set_coms_index` (wp.transformf contract). + articulation.set_coms_index( + coms=wp.from_torch(com.contiguous(), dtype=wp.transformf), + env_ids=wp.from_torch(env_idx, dtype=wp.int32), + ) + + # check they are set + torch.testing.assert_close(_read_binding_to_torch(articulation, TT.BODY_COM_POSE, device), com) + + for i in range(50): + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + # get state properties + root_link_pose_w = articulation.data.root_link_pose_w.torch + root_link_vel_w = articulation.data.root_link_vel_w.torch + root_com_pose_w = articulation.data.root_com_pose_w.torch + root_com_vel_w = articulation.data.root_com_vel_w.torch + body_link_pose_w = articulation.data.body_link_pose_w.torch + body_link_vel_w = articulation.data.body_link_vel_w.torch + body_com_pose_w = articulation.data.body_com_pose_w.torch + body_com_vel_w = articulation.data.body_com_vel_w.torch + + if with_offset: + # get joint state + joint_pos = articulation.data.joint_pos.torch.unsqueeze(-1) + joint_vel = articulation.data.joint_vel.torch.unsqueeze(-1) + + # LINK state + # angular velocity should be the same for both COM and link frames + torch.testing.assert_close(root_com_vel_w[..., 3:], root_link_vel_w[..., 3:]) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) + + # lin_vel arm + lin_vel_gt = torch.zeros(num_articulations, num_bodies, 3, device=device) + vx = -(link_offset[0]) * joint_vel * torch.sin(joint_pos) + vy = torch.zeros(num_articulations, 1, 1, device=device) + vz = (link_offset[0]) * joint_vel * torch.cos(joint_pos) + lin_vel_gt[:, arm_idx, :] = torch.cat([vx, vy, vz], dim=-1).squeeze(-2) + + # linear velocity of root link should be zero + torch.testing.assert_close(lin_vel_gt[:, root_idx, :], root_link_vel_w[..., :3], atol=1e-3, rtol=1e-1) + # linear velocity of pendulum link should be + torch.testing.assert_close(lin_vel_gt, body_link_vel_w[..., :3], atol=1e-3, rtol=1e-1) + + # ang_vel + torch.testing.assert_close(root_com_vel_w[..., 3:], root_link_vel_w[..., 3:]) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) + + # COM state + # position and orientation shouldn't match for the _state_com_w but everything else will + pos_gt = torch.zeros(num_articulations, num_bodies, 3, device=device) + px = (link_offset[0] + offset[0]) * torch.cos(joint_pos) + py = torch.zeros(num_articulations, 1, 1, device=device) + pz = (link_offset[0] + offset[0]) * torch.sin(joint_pos) + pos_gt[:, arm_idx, :] = torch.cat([px, py, pz], dim=-1).squeeze(-2) + pos_gt += env_pos.unsqueeze(-2).repeat(1, num_bodies, 1) + torch.testing.assert_close(pos_gt[:, root_idx, :], root_com_pose_w[..., :3], atol=1e-3, rtol=1e-1) + torch.testing.assert_close(pos_gt, body_com_pose_w[..., :3], atol=1e-3, rtol=1e-1) + + # orientation + com_quat_b = articulation.data.body_com_quat_b.torch + com_quat_w = math_utils.quat_mul(body_link_pose_w[..., 3:], com_quat_b) + torch.testing.assert_close(com_quat_w, body_com_pose_w[..., 3:]) + torch.testing.assert_close(com_quat_w[:, root_idx, :], root_com_pose_w[..., 3:]) + + # angular velocity should be the same for both COM and link frames + torch.testing.assert_close(root_com_vel_w[..., 3:], root_link_vel_w[..., 3:]) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) + else: + # single joint center of masses are at link frames so they will be the same + torch.testing.assert_close(root_link_pose_w, root_com_pose_w) + torch.testing.assert_close(root_com_vel_w, root_link_vel_w) + torch.testing.assert_close(body_link_pose_w, body_com_pose_w) + torch.testing.assert_close(body_com_vel_w, body_link_vel_w) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("with_offset", [True, False]) +@pytest.mark.parametrize("state_location", ["com", "link"]) +@pytest.mark.parametrize("gravity_enabled", [False]) +@pytest.mark.isaacsim_ci +def test_write_root_state(sim, num_articulations, device, with_offset, state_location, gravity_enabled): + """Test the setters for root_state using both the link frame and center of mass as reference frame. + + This test verifies that: + 1. Root states can be written correctly + 2. States are correct with and without offsets + 3. States can be written for both COM and link frames + 4. States are consistent across different devices + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + with_offset: Whether to test with offset + state_location: Whether to test COM or link frame + """ + sim._app_control_on_stop_handle = None + articulation_cfg = generate_articulation_cfg(articulation_type="anymal") + articulation, env_pos = generate_articulation(articulation_cfg, num_articulations, device) + env_idx = torch.tensor([x for x in range(num_articulations)], device=device, dtype=torch.int32) + + # Play sim + sim.reset() + + # change center of mass offset from link frame + if with_offset: + offset = torch.tensor([1.0, 0.0, 0.0]).repeat(num_articulations, 1, 1) + else: + offset = torch.tensor([0.0, 0.0, 0.0]).repeat(num_articulations, 1, 1) + + # create com offsets + com = _read_binding_to_torch(articulation, TT.BODY_COM_POSE, device) + new_com = offset.to(device) + com[:, 0, :3] = new_com.squeeze(-2) + # See test_body_root_state for the PhysX → OVPhysX setter substitution. + articulation.set_coms_index( + coms=wp.from_torch(com.contiguous(), dtype=wp.transformf), + env_ids=wp.from_torch(env_idx, dtype=wp.int32), + ) + + # check they are set + torch.testing.assert_close(_read_binding_to_torch(articulation, TT.BODY_COM_POSE, device), com) + + rand_state = torch.zeros(num_articulations, 13, device=device) + rand_state[..., :7] = articulation.data.default_root_pose.torch + rand_state[..., :3] += env_pos + # make quaternion a unit vector + rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) + + env_idx = env_idx.to(device) + for i in range(10): + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + if state_location == "com": + if i % 2 == 0: + articulation.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7]) + articulation.write_root_com_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) + else: + articulation.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_idx) + articulation.write_root_com_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_idx) + elif state_location == "link": + if i % 2 == 0: + articulation.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7]) + articulation.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) + else: + articulation.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_idx) + articulation.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_idx) + + if state_location == "com": + torch.testing.assert_close(rand_state[..., :7], articulation.data.root_com_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], articulation.data.root_com_vel_w.torch) + elif state_location == "link": + torch.testing.assert_close(rand_state[..., :7], articulation.data.root_link_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], articulation.data.root_link_vel_w.torch) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_body_incoming_joint_wrench_b_single_joint(sim, num_articulations, device): + """Test the data.body_incoming_joint_wrench_b buffer is populated correctly and statically correct for single joint. + + This test verifies that: + 1. The body incoming joint wrench buffer has correct shape + 2. The wrench values are statically correct for a single joint + 3. The wrench values match expected values from gravity and external forces + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + articulation_cfg = generate_articulation_cfg(articulation_type="single_joint_implicit") + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=device + ) + + # Play the simulator + sim.reset() + + # Resolve body indices by name (ordering may differ across physics backends) + arm_idx = articulation.body_names.index("Arm") + root_idx = articulation.body_names.index("CenterPivot") + # apply external force + external_force_vector_b = torch.zeros((num_articulations, articulation.num_bodies, 3), device=device) + external_force_vector_b[:, arm_idx, 1] = 10.0 # 10 N in Y direction + external_torque_vector_b = torch.zeros((num_articulations, articulation.num_bodies, 3), device=device) + external_torque_vector_b[:, arm_idx, 2] = 10.0 # 10 Nm in z direction + + # apply action to the articulation + joint_pos = torch.ones_like(articulation.data.joint_pos.torch) * 1.5708 / 2.0 + articulation.write_joint_position_to_sim_index( + position=torch.ones_like(articulation.data.joint_pos.torch), + ) + articulation.write_joint_velocity_to_sim_index( + velocity=torch.zeros_like(articulation.data.joint_vel.torch), + ) + articulation.set_joint_position_target_index(target=joint_pos) + articulation.write_data_to_sim() + for _ in range(50): + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_force_vector_b, torques=external_torque_vector_b + ) + articulation.write_data_to_sim() + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + # check shape + assert articulation.data.body_incoming_joint_wrench_b.torch.shape == ( + num_articulations, + articulation.num_bodies, + 6, + ) + + # calculate expected static + mass = articulation.data.body_mass.torch.to("cpu") + pos_w = articulation.data.body_pos_w.torch + quat_w = articulation.data.body_quat_w.torch + + mass_link2 = mass[:, arm_idx].view(num_articulations, -1) + gravity = torch.tensor(sim.cfg.gravity, device="cpu").repeat(num_articulations, 1).view((num_articulations, 3)) + + # NOTE: the com and link pose for single joint are colocated + weight_vector_w = mass_link2 * gravity + # expected wrench from link mass and external wrench + # PhysX reports the incoming joint wrench as the force FROM body0 ONTO body1 (body1's frame). + # The USD asset defines body0=CenterPivot, body1=Arm, so the wrench is the constraint/support + # force from CenterPivot onto Arm, expressed in Arm's frame. + # In static equilibrium this equals -(gravity + external forces on Arm). + total_force_w = weight_vector_w.to(device) + math_utils.quat_apply( + quat_w[:, arm_idx, :], external_force_vector_b[:, arm_idx, :] + ) + total_torque_w = torch.cross( + pos_w[:, arm_idx, :].to(device) - pos_w[:, root_idx, :].to(device), + total_force_w, + dim=-1, + ) + math_utils.quat_apply(quat_w[:, arm_idx, :], external_torque_vector_b[:, arm_idx, :]) + expected_wrench = torch.zeros((num_articulations, 6), device=device) + expected_wrench[:, :3] = math_utils.quat_apply( + math_utils.quat_conjugate(quat_w[:, arm_idx, :]), + -total_force_w, + ) + expected_wrench[:, 3:] = math_utils.quat_apply( + math_utils.quat_conjugate(quat_w[:, arm_idx, :]), + -total_torque_w, + ) + + # check value of last joint wrench + torch.testing.assert_close( + expected_wrench, + articulation.data.body_incoming_joint_wrench_b.torch[:, arm_idx, :].squeeze(1), + atol=1e-2, + rtol=1e-3, + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_setting_articulation_root_prim_path(sim, device): + """Test that the articulation root prim path can be set explicitly.""" + sim._app_control_on_stop_handle = None + # Create articulation + articulation_cfg = generate_articulation_cfg(articulation_type="humanoid") + articulation_cfg.articulation_root_prim_path = "/torso" + articulation, _ = generate_articulation(articulation_cfg, 1, device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation._is_initialized + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_setting_invalid_articulation_root_prim_path(sim, device): + """Test that the articulation root prim path can be set explicitly.""" + sim._app_control_on_stop_handle = None + # Create articulation + articulation_cfg = generate_articulation_cfg(articulation_type="humanoid") + articulation_cfg.articulation_root_prim_path = "/non_existing_prim_path" + articulation, _ = generate_articulation(articulation_cfg, 1, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + with pytest.raises(RuntimeError): + sim.reset() + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("gravity_enabled", [False]) +@pytest.mark.isaacsim_ci +def test_write_joint_state_data_consistency(sim, num_articulations, device, gravity_enabled): + """Test the setters for root_state using both the link frame and center of mass as reference frame. + + This test verifies that after write_joint_state_to_sim operations: + 1. state, com_state, link_state value consistency + 2. body_pose, link + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + sim._app_control_on_stop_handle = None + articulation_cfg = generate_articulation_cfg(articulation_type="anymal") + articulation, env_pos = generate_articulation(articulation_cfg, num_articulations, device) + env_idx = torch.tensor([x for x in range(num_articulations)]) + + # Play sim + sim.reset() + + limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) + limits[..., 0] = (torch.rand(num_articulations, articulation.num_joints, device=device) + 5.0) * -1.0 + limits[..., 1] = torch.rand(num_articulations, articulation.num_joints, device=device) + 5.0 + articulation.write_joint_position_limit_to_sim_index(limits=limits) + + from torch.distributions import Uniform + + joint_pos_limits = articulation.data.joint_pos_limits.torch + joint_vel_limits = articulation.data.joint_vel_limits.torch + pos_dist = Uniform(joint_pos_limits[..., 0], joint_pos_limits[..., 1]) + vel_dist = Uniform(-joint_vel_limits, joint_vel_limits) + + original_body_link_pose_w = articulation.data.body_link_pose_w.torch.clone() + original_body_com_vel_w = articulation.data.body_com_vel_w.torch.clone() + + rand_joint_pos = pos_dist.sample() + rand_joint_vel = vel_dist.sample() + + articulation.write_joint_position_to_sim_index(position=rand_joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=rand_joint_vel) + # make sure valued updated + body_link_pose_w = articulation.data.body_link_pose_w.torch + body_com_vel_w = articulation.data.body_com_vel_w.torch + original_body_states = torch.cat([original_body_link_pose_w, original_body_com_vel_w], dim=-1) + body_state_w = torch.cat([body_link_pose_w, body_com_vel_w], dim=-1) + assert torch.count_nonzero(original_body_states[:, 1:] != body_state_w[:, 1:]) > ( + len(original_body_states[:, 1:]) / 2 + ) + # validate body - link consistency + body_link_vel_w = articulation.data.body_link_vel_w.torch + torch.testing.assert_close(body_link_pose_w, articulation.data.body_link_pose_w.torch) + # skip lin_vel because it differs from link frame, this should be fine because we are only checking + # if velocity update is triggered, which can be determined by comparing angular velocity + torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) + + # validate link - com conistency + body_com_pos_b = articulation.data.body_com_pos_b.torch + body_com_quat_b = articulation.data.body_com_quat_b.torch + expected_com_pos, expected_com_quat = math_utils.combine_frame_transforms( + body_link_pose_w[..., :3].view(-1, 3), + body_link_pose_w[..., 3:].view(-1, 4), + body_com_pos_b.view(-1, 3), + body_com_quat_b.view(-1, 4), + ) + body_com_pos_w = articulation.data.body_com_pos_w.torch + body_com_quat_w = articulation.data.body_com_quat_w.torch + torch.testing.assert_close(expected_com_pos.view(len(env_idx), -1, 3), body_com_pos_w) + torch.testing.assert_close(expected_com_quat.view(len(env_idx), -1, 4), body_com_quat_w) + + # validate body - com consistency + body_com_lin_vel_w = articulation.data.body_com_lin_vel_w.torch + body_com_ang_vel_w = articulation.data.body_com_ang_vel_w.torch + torch.testing.assert_close(body_com_vel_w[..., :3], body_com_lin_vel_w) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_com_ang_vel_w) + + # validate pos_w, quat_w, pos_b, quat_b is consistent with pose_w and pose_b + expected_com_pose_w = torch.cat((body_com_pos_w, body_com_quat_w), dim=2) + expected_com_pose_b = torch.cat((body_com_pos_b, body_com_quat_b), dim=2) + body_pos_w = articulation.data.body_pos_w.torch + body_quat_w = articulation.data.body_quat_w.torch + expected_body_pose_w = torch.cat((body_pos_w, body_quat_w), dim=2) + body_link_pos_w = articulation.data.body_link_pos_w.torch + body_link_quat_w = articulation.data.body_link_quat_w.torch + expected_body_link_pose_w = torch.cat((body_link_pos_w, body_link_quat_w), dim=2) + body_com_pose_w = articulation.data.body_com_pose_w.torch + body_com_pose_b = articulation.data.body_com_pose_b.torch + body_pose_w = articulation.data.body_pose_w.torch + body_link_pose_w_fresh = articulation.data.body_link_pose_w.torch + torch.testing.assert_close(body_com_pose_w, expected_com_pose_w) + torch.testing.assert_close(body_com_pose_b, expected_com_pose_b) + torch.testing.assert_close(body_pose_w, expected_body_pose_w) + torch.testing.assert_close(body_link_pose_w_fresh, expected_body_link_pose_w) + + # validate pose_w is consistent with individual properties + body_vel_w = articulation.data.body_vel_w.torch + body_com_vel_w_fresh = articulation.data.body_com_vel_w.torch + torch.testing.assert_close(body_pose_w, body_link_pose_w) + torch.testing.assert_close(body_vel_w, body_com_vel_w) + torch.testing.assert_close(body_link_pose_w_fresh, body_link_pose_w) + torch.testing.assert_close(body_com_pose_w, articulation.data.body_com_pose_w.torch) + torch.testing.assert_close(body_vel_w, body_com_vel_w_fresh) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_spatial_tendons(sim, num_articulations, device): + """Test spatial tendons apis. + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation has spatial tendons + 3. All buffers have correct shapes + 4. The articulation can be simulated + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + # skip test if Isaac Sim version is less than 5.0 + if has_kit() and get_isaac_sim_version().major < 5: + pytest.skip("Spatial tendons are not supported in Isaac Sim < 5.0. Please update to Isaac Sim 5.0 or later.") + return + articulation_cfg = generate_articulation_cfg(articulation_type="spatial_tendon_test_asset") + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that fixed base + assert articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 3) + assert articulation.data.body_mass.torch.shape == (num_articulations, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (num_articulations, articulation.num_bodies, 9) + assert articulation.num_spatial_tendons == 1 + + articulation.set_spatial_tendon_stiffness_index(stiffness=10.0) + articulation.set_spatial_tendon_limit_stiffness_index(limit_stiffness=10.0) + articulation.set_spatial_tendon_damping_index(damping=10.0) + articulation.set_spatial_tendon_offset_index(offset=10.0) + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_write_joint_frictions_to_sim(sim, num_articulations, device, add_ground_plane): + """Test applying of joint position target functions correctly for a robotic arm.""" + articulation_cfg = generate_articulation_cfg(articulation_type="panda") + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=device + ) + + # Play the simulator + sim.reset() + + for _ in range(100): + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + # apply action to the articulation + dynamic_friction = torch.rand(num_articulations, articulation.num_joints, device=device) + viscous_friction = torch.rand(num_articulations, articulation.num_joints, device=device) + friction = torch.rand(num_articulations, articulation.num_joints, device=device) + + # Guarantee that the dynamic friction is not greater than the static friction + dynamic_friction = torch.min(dynamic_friction, friction) + + # The static friction must be set first to be sure the dynamic friction is not greater than static + # when both are set. + articulation.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff=friction, + joint_dynamic_friction_coeff=dynamic_friction, + joint_viscous_friction_coeff=viscous_friction, + ) + articulation.write_data_to_sim() + + for _ in range(100): + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + friction_props_from_sim = _read_binding_to_torch(articulation, TT.DOF_FRICTION_PROPERTIES, "cpu") + joint_friction_coeff_sim = friction_props_from_sim[:, :, 0] + joint_dynamic_friction_coeff_sim = friction_props_from_sim[:, :, 1] + joint_viscous_friction_coeff_sim = friction_props_from_sim[:, :, 2] + assert torch.allclose(joint_dynamic_friction_coeff_sim, dynamic_friction.cpu()) + assert torch.allclose(joint_viscous_friction_coeff_sim, viscous_friction.cpu()) + assert torch.allclose(joint_friction_coeff_sim, friction.cpu()) + + # For Isaac Sim >= 5.0: also test the combined API that can set dynamic and viscous via + # write_joint_friction_coefficient_to_sim; reset the sim to isolate this path. + if has_kit() and get_isaac_sim_version().major >= 5: + # Reset simulator to ensure a clean state for the alternative API path + sim.reset() + + # Warm up a few steps to populate buffers + for _ in range(100): + sim.step() + articulation.update(sim.cfg.dt) + + # New random coefficients + dynamic_friction_2 = torch.rand(num_articulations, articulation.num_joints, device=device) + viscous_friction_2 = torch.rand(num_articulations, articulation.num_joints, device=device) + friction_2 = torch.rand(num_articulations, articulation.num_joints, device=device) + + # Guarantee that the dynamic friction is not greater than the static friction + dynamic_friction_2 = torch.min(dynamic_friction_2, friction_2) + + # Use the combined setter to write all three at once + articulation.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff=friction_2, + joint_dynamic_friction_coeff=dynamic_friction_2, + joint_viscous_friction_coeff=viscous_friction_2, + ) + articulation.write_data_to_sim() + + # Step to let sim ingest new params and refresh data buffers + for _ in range(100): + sim.step() + articulation.update(sim.cfg.dt) + + friction_props_from_sim_2 = _read_binding_to_torch(articulation, TT.DOF_FRICTION_PROPERTIES, "cpu") + joint_friction_coeff_sim_2 = friction_props_from_sim_2[:, :, 0] + friction_dynamic_coef_sim_2 = friction_props_from_sim_2[:, :, 1] + friction_viscous_coeff_sim_2 = friction_props_from_sim_2[:, :, 2] + + # Validate values propagated + assert torch.allclose(friction_viscous_coeff_sim_2, viscous_friction_2.cpu()) + assert torch.allclose(friction_dynamic_coef_sim_2, dynamic_friction_2.cpu()) + assert torch.allclose(joint_friction_coeff_sim_2, friction_2.cpu()) + + +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("articulation_type", ["panda"]) +@pytest.mark.isaacsim_ci +@pytest.mark.xfail(reason=_MATERIAL_GAP_REASON, strict=False) +def test_set_material_properties(sim, num_articulations, device, add_ground_plane, articulation_type): + """Test getting and setting material properties (friction/restitution) of articulation shapes.""" + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=device + ) + + # Play the simulator + sim.reset() + + # Get number of shapes from the articulation + max_shapes = articulation.root_view.max_shapes + + # Generate random material properties: (static_friction, dynamic_friction, restitution) + materials = torch.empty(num_articulations, max_shapes, 3, device="cpu").uniform_(0.0, 1.0) + # Ensure dynamic friction <= static friction + materials[..., 1] = torch.min(materials[..., 0], materials[..., 1]) + + # Set material properties via the PhysX view-level API + env_ids = torch.arange(num_articulations, dtype=torch.int32) + articulation.root_view.set_material_properties( + wp.from_torch(materials, dtype=wp.float32), wp.from_torch(env_ids, dtype=wp.int32) + ) + + # Simulate physics + sim.step() + articulation.update(sim.cfg.dt) + + # Get material properties from simulation + materials_check = wp.to_torch(articulation.root_view.get_material_properties()) + + # Check if material properties are set correctly + torch.testing.assert_close(materials_check, materials) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "--maxfail=1"]) diff --git a/source/isaaclab_ovphysx/test/assets/test_articulation_helpers.py b/source/isaaclab_ovphysx/test/assets/test_articulation_helpers.py new file mode 100644 index 000000000000..9be6185022da --- /dev/null +++ b/source/isaaclab_ovphysx/test/assets/test_articulation_helpers.py @@ -0,0 +1,142 @@ +# 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 + +"""OVPhysX-only unit tests for articulation helpers. + +These tests cover OVPhysX-specific scaffolding (USD tendon-scope resolution, +mock binding-set shape contracts) that has no PhysX equivalent and therefore +does not appear in the PhysX-mirrored ``test_articulation.py``. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import warp as wp + +from pxr import Sdf, Usd, UsdPhysics + +# The CI isaaclab_ov* pattern unintentionally collects isaaclab_ovphysx tests, +# but the ovphysx wheel is not installed in that environment. Skip gracefully +# so the isaaclab_ov CI pipeline is not blocked by an unrelated dependency. +pytest.importorskip("ovphysx.types", reason="ovphysx wheel not installed") + +from isaaclab_ovphysx.assets.articulation.articulation import Articulation # noqa: E402 +from isaaclab_ovphysx.physics import OvPhysxManager # noqa: E402 +from isaaclab_ovphysx.test.mock_interfaces.views import MockOvPhysxBindingSet # noqa: E402 + +wp.init() + + +def _define_tendon_joint(stage: Usd.Stage, path: str, schema_name: str) -> None: + """Define a revolute joint prim with a tendon schema marker.""" + joint = UsdPhysics.RevoluteJoint.Define(stage, path) + schemas = Sdf.TokenListOp() + schemas.explicitItems = [schema_name] + joint.GetPrim().SetMetadata("apiSchemas", schemas) + + +def _make_articulation_root_stage(tmp_path) -> str: + """Create a stage with one relevant articulation subtree and unrelated joints elsewhere.""" + stage = Usd.Stage.CreateInMemory() + stage.DefinePrim("/World", "Xform") + stage.DefinePrim("/World/envs", "Xform") + stage.DefinePrim("/World/envs/env_0", "Xform") + stage.DefinePrim("/World/envs/env_0/Robot", "Xform") + stage.DefinePrim("/World/envs/env_0/Robot/root", "Xform") + stage.DefinePrim("/World/unrelated", "Xform") + + _define_tendon_joint( + stage, + "/World/envs/env_0/Robot/root/fixed_joint", + "PhysxTendonAxisRootAPI:inst0", + ) + _define_tendon_joint( + stage, + "/World/envs/env_0/Robot/root/spatial_joint", + "PhysxTendonAttachmentRootAPI:inst0", + ) + _define_tendon_joint( + stage, + "/World/unrelated/unrelated_fixed_joint", + "PhysxTendonAxisRootAPI:inst0", + ) + _define_tendon_joint( + stage, + "/World/unrelated/unrelated_spatial_joint", + "PhysxTendonAttachmentLeafAPI:inst0", + ) + + stage_path = tmp_path / "scene.usda" + stage.Export(str(stage_path)) + return str(stage_path) + + +def _make_articulation_shell() -> Articulation: + """Create a minimal ovphysx articulation shell for tendon processing tests.""" + articulation = object.__new__(Articulation) + bindings = MockOvPhysxBindingSet( + num_instances=1, + num_joints=2, + num_bodies=2, + num_fixed_tendons=1, + num_spatial_tendons=1, + ) + object.__setattr__(articulation, "_bindings", bindings.bindings) + object.__setattr__(articulation, "_articulation_root_path", "/World/envs/env_0/Robot/root") + object.__setattr__(articulation, "_initialize_handle", None) + object.__setattr__(articulation, "_invalidate_initialize_handle", None) + object.__setattr__(articulation, "_prim_deletion_handle", None) + object.__setattr__(articulation, "_debug_vis_handle", None) + object.__setattr__( + articulation, + "_data", + SimpleNamespace( + _num_fixed_tendons=0, + _num_spatial_tendons=0, + fixed_tendon_names=[], + spatial_tendon_names=[], + ), + ) + return articulation + + +def test_process_tendons_scopes_to_articulation_root(tmp_path): + """Tendon discovery should ignore joints that live outside the current articulation subtree.""" + articulation = _make_articulation_shell() + stage_path = _make_articulation_root_stage(tmp_path) + old_stage_path = OvPhysxManager._stage_path + OvPhysxManager._stage_path = stage_path + try: + articulation._process_tendons() + finally: + OvPhysxManager._stage_path = old_stage_path + + assert articulation.fixed_tendon_names == ["fixed_joint"] + assert articulation.spatial_tendon_names == ["spatial_joint"] + assert articulation._data.fixed_tendon_names == ["fixed_joint"] + assert articulation._data.spatial_tendon_names == ["spatial_joint"] + + +def test_mock_binding_set_rigid_object_shapes(): + pytest.importorskip("isaaclab_ovphysx.tensor_types").RIGID_BODY_POSE # gates on wheel + from isaaclab_ovphysx import tensor_types as TT + from isaaclab_ovphysx.test.mock_interfaces.views import MockOvPhysxBindingSet + + bindings = MockOvPhysxBindingSet( + num_instances=4, + num_joints=0, + num_bodies=1, + asset_kind="rigid_object", + ) + assert bindings.bindings[TT.RIGID_BODY_POSE].shape == (4, 7) + assert bindings.bindings[TT.RIGID_BODY_VELOCITY].shape == (4, 6) + assert bindings.bindings[TT.RIGID_BODY_WRENCH].shape == (4, 9) + assert bindings.bindings[TT.RIGID_BODY_MASS].shape == (4,) + assert bindings.bindings[TT.RIGID_BODY_INERTIA].shape == (4, 9) + # Articulation-only bindings must be absent + assert TT.DOF_POSITION not in bindings.bindings + assert TT.LINK_WRENCH not in bindings.bindings diff --git a/source/isaaclab_ovphysx/test/assets/test_rigid_object.py b/source/isaaclab_ovphysx/test/assets/test_rigid_object.py new file mode 100644 index 000000000000..407cf4b41e22 --- /dev/null +++ b/source/isaaclab_ovphysx/test/assets/test_rigid_object.py @@ -0,0 +1,1134 @@ +# 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 + +# ignore private usage of variables warning +# pyright: reportPrivateUsage=none + + +"""Real-backend tests for the OVPhysX RigidObject. + +Run via ``./scripts/run_ovphysx.sh -m pytest`` (kitless, no ``AppLauncher``). + +``ovphysx<=0.3.7`` binds device mode (CPU vs GPU) at the C++ layer on the +first ``ovphysx.PhysX(device=...)`` construction and cannot swap it without a +process restart. Full coverage therefore requires two separate pytest +invocations -- once with ``-k 'cpu'`` and once with ``-k 'cuda:0'``. The +``_ovphysx_skip_other_device`` autouse fixture below preempts the manager's +:exc:`RuntimeError` by ``pytest.skip``-ing on the unlocked device so +single-device runs finish cleanly. +""" + +from __future__ import annotations + +import logging +import sys +from typing import Literal +from unittest.mock import MagicMock + +import pytest +import torch +import warp as wp +from flaky import flaky + +# The CI isaaclab_ov* pattern unintentionally collects isaaclab_ovphysx tests, +# but the ovphysx wheel is not installed in that environment. Skip gracefully +# so the isaaclab_ov CI pipeline is not blocked by an unrelated dependency. +pytest.importorskip("ovphysx.types", reason="ovphysx wheel not installed") + +from isaaclab_ovphysx.assets import RigidObject # noqa: E402 +from isaaclab_ovphysx.physics import OvPhysxCfg, OvPhysxManager # noqa: E402 + +import isaaclab.sim as sim_utils # noqa: E402 +from isaaclab.assets import RigidObjectCfg # noqa: E402 +from isaaclab.sim import SimulationCfg, build_simulation_context # noqa: E402 +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR # noqa: E402 +from isaaclab.utils.math import ( # noqa: E402 + combine_frame_transforms, + default_orientation, + quat_apply_inverse, + quat_inv, + quat_mul, + quat_rotate, + random_orientation, +) + +wp.init() + +_logger = logging.getLogger(__name__) + + +_LOCKED_DEVICE: list[str | None] = [None] +"""Device the session pins to on the first parametrized test that runs.""" + + +@pytest.fixture(autouse=True) +def _ovphysx_skip_other_device(request): + """Skip parametrized tests on the device the session is not pinned to. + + See the module docstring for the wheel's process-global device-mode lock. + """ + callspec = getattr(request.node, "callspec", None) + device = callspec.params.get("device") if callspec is not None else None + if device is None: + # Test does not parametrize on device (e.g. test_warmup_attach_stage_not_called_for_cpu). + return + locked = _LOCKED_DEVICE[0] + if locked is None: + _LOCKED_DEVICE[0] = device + return + if device != locked: + pytest.skip( + f"ovphysx process-global device lock is held by '{locked}'; cannot run '{device}' " + "tests in the same session. Run pytest twice (once per device) for full coverage." + ) + + +def _ovphysx_sim_context(device: str, **kwargs): + """Wrapper around :func:`build_simulation_context` that injects OVPhysX cfg. + + PhysX tests pass ``device=device`` directly and let + :func:`build_simulation_context` build a default :class:`SimulationCfg`. + OVPhysX needs ``physics=OvPhysxCfg()`` set on the cfg so the manager + dispatches to OVPhysX rather than PhysX, so we build the cfg here and + pass it through. ``gravity_enabled`` is consumed locally (it is ignored + by ``build_simulation_context`` once a ``sim_cfg`` is provided). + ``add_ground_plane``, ``auto_add_lighting``, and other kwargs continue + to flow through ``build_simulation_context`` as before. + """ + dt = kwargs.pop("dt", 1.0 / 60.0) + gravity_enabled = kwargs.pop("gravity_enabled", True) + gravity = (0.0, 0.0, -9.81) if gravity_enabled else (0.0, 0.0, 0.0) + sim_cfg = SimulationCfg(physics=OvPhysxCfg(), device=device, dt=dt, gravity=gravity) + return build_simulation_context(device=device, sim_cfg=sim_cfg, **kwargs) + + +def generate_cubes_scene( + num_cubes: int = 1, + height=1.0, + api: Literal["none", "rigid_body", "articulation_root"] = "rigid_body", + kinematic_enabled: bool = False, + device: str = "cuda:0", +) -> tuple[RigidObject, torch.Tensor]: + """Generate a scene with the provided number of cubes. + + Args: + num_cubes: Number of cubes to generate. + height: Height of the cubes. + api: The type of API that the cubes should have. + kinematic_enabled: Whether the cubes are kinematic. + device: Device to use for the simulation. + + Returns: + A tuple containing the rigid object representing the cubes and the origins of the cubes. + + """ + origins = torch.tensor([(i * 1.0, 0, height) for i in range(num_cubes)]).to(device) + # Create Top-level Xforms, one for each cube + for i, origin in enumerate(origins): + sim_utils.create_prim(f"/World/Table_{i}", "Xform", translation=origin) + + # Resolve spawn configuration + if api == "none": + # since no rigid body properties defined, this is just a static collider + spawn_cfg = sim_utils.CuboidCfg( + size=(0.1, 0.1, 0.1), + collision_props=sim_utils.CollisionPropertiesCfg(), + ) + elif api == "rigid_body": + spawn_cfg = sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", + rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=kinematic_enabled), + ) + elif api == "articulation_root": + spawn_cfg = sim_utils.UsdFileCfg( + usd_path=f"{ISAACLAB_NUCLEUS_DIR}/Tests/RigidObject/Cube/dex_cube_instanceable_with_articulation_root.usd", + rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=kinematic_enabled), + ) + else: + raise ValueError(f"Unknown api: {api}") + + # Create rigid object. OVPhysX matches prim paths via fnmatch globs (not regex), + # so use ``Table_*`` rather than the PhysX ``Table_.*`` form. + cube_object_cfg = RigidObjectCfg( + prim_path="/World/Table_*/Object", + spawn=spawn_cfg, + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, height)), + ) + cube_object = RigidObject(cfg=cube_object_cfg) + + return cube_object, origins + + +# --------------------------------------------------------------------------- +# Material-property gap (xfail reason shared by 5 tests below) +# --------------------------------------------------------------------------- + +_MATERIAL_GAP_REASON = ( + "Requires RIGID_BODY_MATERIAL TensorType (or a view-helper) on the ovphysx " + "wheel side. RigidObject.root_view is a per-tensor-type bindings dict on " + "OVPhysX, so root_view.get_material_properties() / set_material_properties() " + "are not available. See " + "docs/superpowers/specs/2026-04-28-ovphysx-wheel-gaps-for-marco.md." +) + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_initialization(num_cubes, device): + """Test initialization for prim with rigid body API at the provided prim path.""" + with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: + # Generate cubes scene + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(cube_object) < 10 + + # Play sim + sim.reset() + + # Check if object is initialized + assert cube_object.is_initialized + assert len(cube_object.body_names) == 1 + + # Check buffers that exists and have correct shapes + assert cube_object.data.root_pos_w.torch.shape == (num_cubes, 3) + assert cube_object.data.root_quat_w.torch.shape == (num_cubes, 4) + assert cube_object.data.body_mass.torch.shape == (num_cubes, 1) + assert cube_object.data.body_inertia.torch.shape == (num_cubes, 1, 9) + + # Simulate physics + for _ in range(2): + # perform rendering + sim.step() + # update object + cube_object.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_initialization_with_kinematic_enabled(num_cubes, device): + """Test that initialization for prim with kinematic flag enabled.""" + with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: + # Generate cubes scene + cube_object, origins = generate_cubes_scene(num_cubes=num_cubes, kinematic_enabled=True, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(cube_object) < 10 + + # Play sim + sim.reset() + + # Check if object is initialized + assert cube_object.is_initialized + assert len(cube_object.body_names) == 1 + + # Check buffers that exists and have correct shapes + assert cube_object.data.root_pos_w.torch.shape == (num_cubes, 3) + assert cube_object.data.root_quat_w.torch.shape == (num_cubes, 4) + + # Simulate physics + for _ in range(2): + # perform rendering + sim.step() + # update object + cube_object.update(sim.cfg.dt) + # check that the object is kinematic + default_root_pose = cube_object.data.default_root_pose.torch.clone() + default_root_vel = cube_object.data.default_root_vel.torch.clone() + default_root_pose[:, :3] += origins + torch.testing.assert_close(cube_object.data.root_link_pose_w.torch, default_root_pose) + torch.testing.assert_close(cube_object.data.root_com_vel_w.torch, default_root_vel) + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_initialization_with_no_rigid_body(num_cubes, device): + """Test that initialization fails when no rigid body is found at the provided prim path.""" + with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: + # Generate cubes scene + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, api="none", device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(cube_object) < 10 + + # Play sim + with pytest.raises(RuntimeError): + sim.reset() + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_initialization_with_articulation_root(num_cubes, device): + """Test that initialization fails when an articulation root is found at the provided prim path.""" + with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: + # Generate cubes scene + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, api="articulation_root", device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(cube_object) < 10 + + # Play sim + with pytest.raises(RuntimeError): + sim.reset() + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_external_force_buffer(device): + """Test if external force buffer correctly updates in the force value is zero case. + + In this test, we apply a non-zero force, then a zero force, then finally a non-zero force + to an object. We check if the force buffer is properly updated at each step. + """ + + # Generate cubes scene + with _ovphysx_sim_context(device=device, add_ground_plane=True, auto_add_lighting=True) as sim: + cube_object, origins = generate_cubes_scene(num_cubes=1, device=device) + + # play the simulator + sim.reset() + + # find bodies to apply the force + body_ids, body_names = cube_object.find_bodies(".*") + + # reset object + cube_object.reset() + + # perform simulation + for step in range(5): + # initiate force tensor + external_wrench_b = torch.zeros(cube_object.num_instances, len(body_ids), 6, device=sim.device) + + if step == 0 or step == 3: + # set a non-zero force + force = 1 + else: + # set a zero force + force = 0 + + # set force value + external_wrench_b[:, :, 0] = force + external_wrench_b[:, :, 3] = force + + # apply force + cube_object.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + body_ids=body_ids, + ) + + # check if the cube's force and torque buffers are correctly updated + for i in range(cube_object.num_instances): + assert cube_object._permanent_wrench_composer.composed_force.torch[i, 0, 0].item() == force + assert cube_object._permanent_wrench_composer.composed_torque.torch[i, 0, 0].item() == force + + # Check if the instantaneous wrench is correctly added to the permanent wrench + cube_object.permanent_wrench_composer.add_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + body_ids=body_ids, + ) + + # apply action to the object + cube_object.write_data_to_sim() + + # perform step + sim.step() + + # update buffers + cube_object.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_cubes", [2, 4]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_external_force_on_single_body(num_cubes, device): + """Test application of external force on the base of the object. + + In this test, we apply a force equal to the weight of an object on the base of + one of the objects. We check that the object does not move. For the other object, + we do not apply any force and check that it falls down. + + We validate that this works when we apply the force in the global frame and in the local frame. + """ + # Generate cubes scene + with _ovphysx_sim_context(device=device, add_ground_plane=True, auto_add_lighting=True) as sim: + cube_object, origins = generate_cubes_scene(num_cubes=num_cubes, device=device) + + # Play the simulator + sim.reset() + + # Find bodies to apply the force + body_ids, body_names = cube_object.find_bodies(".*") + + # Sample a force equal to the weight of the object. PhysX reads the mass + # from ``root_view.get_masses()``; OVPhysX exposes the same value via + # ``cube_object.data.body_mass`` (shape ``(N, 1)``). + external_wrench_b = torch.zeros(cube_object.num_instances, len(body_ids), 6, device=sim.device) + # Every 2nd cube should have a force applied to it + external_wrench_b[0::2, :, 2] = 9.81 * cube_object.data.body_mass.torch[0] + + # Now we are ready! + for i in range(5): + # reset root state + root_pose = cube_object.data.default_root_pose.torch.clone() + root_vel = cube_object.data.default_root_vel.torch.clone() + + # need to shift the position of the cubes otherwise they will be on top of each other + root_pose[:, :3] = origins + cube_object.write_root_pose_to_sim_index(root_pose=root_pose) + cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) + + # reset object + cube_object.reset() + + is_global = False + if i % 2 == 0: + is_global = True + positions = cube_object.data.body_com_pos_w.torch[:, body_ids, :3] + else: + positions = None + + # apply force + cube_object.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=positions, + body_ids=body_ids, + is_global=is_global, + ) + # perform simulation + for _ in range(5): + # apply action to the object + cube_object.write_data_to_sim() + + # perform step + sim.step() + + # update buffers + cube_object.update(sim.cfg.dt) + + # First object should still be at the same Z position (1.0) + torch.testing.assert_close( + cube_object.data.root_pos_w.torch[0::2, 2], torch.ones(num_cubes // 2, device=sim.device) + ) + # Second object should have fallen, so it's Z height should be less than initial height of 1.0 + assert torch.all(cube_object.data.root_pos_w.torch[1::2, 2] < 1.0) + + +@pytest.mark.parametrize("num_cubes", [2, 4]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_external_force_on_single_body_at_position(num_cubes, device): + """Test application of external force on the base of the object at a specific position. + + In this test, we apply a force equal to the weight of an object on the base of + one of the objects at 1m in the Y direction, we check that the object rotates around it's X axis. + For the other object, we do not apply any force and check that it falls down. + + We validate that this works when we apply the force in the global frame and in the local frame. + """ + # Generate cubes scene + with _ovphysx_sim_context(device=device, add_ground_plane=True, auto_add_lighting=True) as sim: + cube_object, origins = generate_cubes_scene(num_cubes=num_cubes, device=device) + + # Play the simulator + sim.reset() + + # Find bodies to apply the force + body_ids, body_names = cube_object.find_bodies(".*") + + # Sample a force equal to the weight of the object + external_wrench_b = torch.zeros(cube_object.num_instances, len(body_ids), 6, device=sim.device) + external_wrench_positions_b = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=sim.device) + # Every 2nd cube should have a force applied to it + external_wrench_b[0::2, :, 2] = 500.0 + external_wrench_positions_b[0::2, :, 1] = 1.0 + + # Desired force and torque + desired_force = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=sim.device) + desired_force[0::2, :, 2] = 1000.0 + desired_torque = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=sim.device) + desired_torque[0::2, :, 0] = 1000.0 + # Now we are ready! + for i in range(5): + # reset root state + root_pose = cube_object.data.default_root_pose.torch.clone() + root_vel = cube_object.data.default_root_vel.torch.clone() + + # need to shift the position of the cubes otherwise they will be on top of each other + root_pose[:, :3] = origins + cube_object.write_root_pose_to_sim_index(root_pose=root_pose) + cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) + + # reset object + cube_object.reset() + + is_global = False + if i % 2 == 0: + is_global = True + body_com_pos_w = cube_object.data.body_com_pos_w.torch[:, body_ids, :3] + external_wrench_positions_b[..., 0] = 0.0 + external_wrench_positions_b[..., 1] = 1.0 + external_wrench_positions_b[..., 2] = 0.0 + external_wrench_positions_b += body_com_pos_w + else: + external_wrench_positions_b[..., 0] = 0.0 + external_wrench_positions_b[..., 1] = 1.0 + external_wrench_positions_b[..., 2] = 0.0 + + # apply force + cube_object.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=external_wrench_positions_b, + body_ids=body_ids, + is_global=is_global, + ) + cube_object.permanent_wrench_composer.add_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=external_wrench_positions_b, + body_ids=body_ids, + is_global=is_global, + ) + torch.testing.assert_close( + cube_object._permanent_wrench_composer.composed_force.torch[:, 0, :], + desired_force[:, 0, :], + rtol=1e-6, + atol=1e-7, + ) + torch.testing.assert_close( + cube_object._permanent_wrench_composer.composed_torque.torch[:, 0, :], + desired_torque[:, 0, :], + rtol=1e-6, + atol=1e-7, + ) + # perform simulation + for _ in range(5): + # apply action to the object + cube_object.write_data_to_sim() + + # perform step + sim.step() + + # update buffers + cube_object.update(sim.cfg.dt) + + # The first object should be rotating around it's X axis + assert torch.all(torch.abs(cube_object.data.root_ang_vel_b.torch[0::2, 0]) > 0.1) + # Second object should have fallen, so it's Z height should be less than initial height of 1.0 + assert torch.all(cube_object.data.root_pos_w.torch[1::2, 2] < 1.0) + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_set_rigid_object_state(num_cubes, device): + """Test setting the state of the rigid object. + + In this test, we set the state of the rigid object to a random state and check + that the object is in that state after simulation. We set gravity to zero as + we don't want any external forces acting on the object to ensure state remains static. + """ + # Turn off gravity for this test as we don't want any external forces acting on the object + # to ensure state remains static + with _ovphysx_sim_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + # Generate cubes scene + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, device=device) + + # Play the simulator + sim.reset() + + state_types = ["root_pos_w", "root_quat_w", "root_lin_vel_w", "root_ang_vel_w"] + + # Set each state type individually as they are dependent on each other + for state_type_to_randomize in state_types: + state_dict = { + "root_pos_w": torch.zeros_like(cube_object.data.root_pos_w.torch, device=sim.device), + "root_quat_w": default_orientation(num=num_cubes, device=sim.device), + "root_lin_vel_w": torch.zeros_like(cube_object.data.root_lin_vel_w.torch, device=sim.device), + "root_ang_vel_w": torch.zeros_like(cube_object.data.root_ang_vel_w.torch, device=sim.device), + } + + # Now we are ready! + for _ in range(5): + # reset object + cube_object.reset() + + # Set random state + if state_type_to_randomize == "root_quat_w": + state_dict[state_type_to_randomize] = random_orientation(num=num_cubes, device=sim.device) + else: + state_dict[state_type_to_randomize] = torch.randn(num_cubes, 3, device=sim.device) + + # perform simulation + for _ in range(5): + root_pose = torch.cat( + [state_dict["root_pos_w"], state_dict["root_quat_w"]], + dim=-1, + ) + root_vel = torch.cat( + [state_dict["root_lin_vel_w"], state_dict["root_ang_vel_w"]], + dim=-1, + ) + # reset root state + cube_object.write_root_pose_to_sim_index(root_pose=root_pose) + cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) + + sim.step() + + # assert that set root quantities are equal to the ones set in the state_dict + for key, expected_value in state_dict.items(): + value = getattr(cube_object.data, key).torch + torch.testing.assert_close(value, expected_value, rtol=1e-3, atol=1e-3) + + cube_object.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_reset_rigid_object(num_cubes, device): + """Test resetting the state of the rigid object.""" + with _ovphysx_sim_context(device=device, gravity_enabled=True, auto_add_lighting=True) as sim: + # Generate cubes scene + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, device=device) + + # Play the simulator + sim.reset() + + for i in range(5): + # perform rendering + sim.step() + + # update object + cube_object.update(sim.cfg.dt) + + # Move the object to a random position + root_pose = cube_object.data.default_root_pose.torch.clone() + root_pose[:, :3] = torch.randn(num_cubes, 3, device=sim.device) + + # Random orientation + root_pose[:, 3:7] = random_orientation(num=num_cubes, device=sim.device) + cube_object.write_root_pose_to_sim_index(root_pose=root_pose) + root_vel = cube_object.data.default_root_vel.torch.clone() + cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) + + if i % 2 == 0: + # reset object + cube_object.reset() + + # Reset should zero external forces and torques + assert not cube_object._instantaneous_wrench_composer.active + assert not cube_object._permanent_wrench_composer.active + assert torch.count_nonzero(cube_object._instantaneous_wrench_composer.composed_force.torch) == 0 + assert torch.count_nonzero(cube_object._instantaneous_wrench_composer.composed_torque.torch) == 0 + assert torch.count_nonzero(cube_object._permanent_wrench_composer.composed_force.torch) == 0 + assert torch.count_nonzero(cube_object._permanent_wrench_composer.composed_torque.torch) == 0 + + +@pytest.mark.xfail(reason=_MATERIAL_GAP_REASON, strict=False) +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_rigid_body_set_material_properties(num_cubes, device): + """Test getting and setting material properties of rigid object.""" + raise NotImplementedError(_MATERIAL_GAP_REASON) + + +@pytest.mark.xfail(reason=_MATERIAL_GAP_REASON, strict=False) +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_set_material_properties_via_view(num_cubes, device): + """Test setting material properties via the PhysX view-level API.""" + raise NotImplementedError(_MATERIAL_GAP_REASON) + + +@pytest.mark.xfail(reason=_MATERIAL_GAP_REASON, strict=False) +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_rigid_body_no_friction(num_cubes, device): + """Test that a rigid object with no friction will maintain it's velocity when sliding across a plane.""" + raise NotImplementedError(_MATERIAL_GAP_REASON) + + +@pytest.mark.xfail(reason=_MATERIAL_GAP_REASON, strict=False) +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda", "cpu"]) +@pytest.mark.isaacsim_ci +def test_rigid_body_with_static_friction(num_cubes, device): + """Test that static friction applied to rigid object works as expected. + + This test works by applying a force to the object and checking if the object moves or not based on the + mu (coefficient of static friction) value set for the object. We set the static friction to be non-zero and + apply a force to the object. When the force applied is below mu, the object should not move. When the force + applied is above mu, the object should move. + """ + raise NotImplementedError(_MATERIAL_GAP_REASON) + + +@pytest.mark.xfail(reason=_MATERIAL_GAP_REASON, strict=False) +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_rigid_body_with_restitution(num_cubes, device): + """Test that restitution when applied to rigid object works as expected. + + This test works by dropping a block from a height and checking if the block bounces or not based on the + restitution value set for the object. We set the restitution to be non-zero and drop the block from a height. + When the restitution is 0, the block should not bounce. When the restitution is between 0 and 1, the block + should bounce with less energy. + """ + raise NotImplementedError(_MATERIAL_GAP_REASON) + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_rigid_body_set_mass(num_cubes, device): + """Test getting and setting mass of rigid object.""" + with _ovphysx_sim_context( + device=device, gravity_enabled=False, add_ground_plane=True, auto_add_lighting=True + ) as sim: + # Create a scene with random cubes + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=1.0, device=device) + + # Play sim + sim.reset() + + # Get masses before increasing + original_masses = cube_object.data.body_mass.torch.clone() + + assert original_masses.shape == (num_cubes, 1) + + # Randomize mass of the object + masses = original_masses + torch.FloatTensor(num_cubes, 1).uniform_(4, 8).to(sim.device) + + indices = torch.tensor(range(num_cubes), dtype=torch.int32) + + # Set the new masses via the OVPhysX writer (matches PhysX/Newton). + cube_object.set_masses_index( + masses=wp.from_torch(masses.contiguous(), dtype=wp.float32), + env_ids=wp.from_torch(indices, dtype=wp.int32), + ) + + torch.testing.assert_close(cube_object.data.body_mass.torch, masses) + + # Simulate physics + # perform rendering + sim.step() + # update object + cube_object.update(sim.cfg.dt) + + masses_to_check = cube_object.data.body_mass.torch + + # Check if mass is set correctly + torch.testing.assert_close(masses, masses_to_check) + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("gravity_enabled", [True, False]) +@pytest.mark.isaacsim_ci +def test_gravity_vec_w(num_cubes, device, gravity_enabled): + """Test that gravity vector direction is set correctly for the rigid object.""" + with _ovphysx_sim_context(device=device, gravity_enabled=gravity_enabled) as sim: + # Create a scene with random cubes + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, device=device) + + # Obtain gravity direction + if gravity_enabled: + gravity_dir = (0.0, 0.0, -1.0) + else: + gravity_dir = (0.0, 0.0, 0.0) + + # Play sim + sim.reset() + + # Check that gravity is set correctly + assert cube_object.data.GRAVITY_VEC_W.torch[0, 0] == gravity_dir[0] + assert cube_object.data.GRAVITY_VEC_W.torch[0, 1] == gravity_dir[1] + assert cube_object.data.GRAVITY_VEC_W.torch[0, 2] == gravity_dir[2] + + # Simulate physics + for _ in range(2): + # perform rendering + sim.step() + # update object + cube_object.update(sim.cfg.dt) + + # Expected gravity value is the acceleration of the body + gravity = torch.zeros(num_cubes, 1, 6, device=device) + if gravity_enabled: + gravity[:, :, 2] = -9.81 + # Check the body accelerations are correct + torch.testing.assert_close(cube_object.data.body_acc_w.torch, gravity) + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("with_offset", [True, False]) +@pytest.mark.isaacsim_ci +@flaky(max_runs=3, min_passes=1) +def test_body_root_state_properties(num_cubes, device, with_offset): + """Test the root_com_state_w, root_link_state_w, body_com_state_w, and body_link_state_w properties.""" + with _ovphysx_sim_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + # Create a scene with random cubes + cube_object, env_pos = generate_cubes_scene(num_cubes=num_cubes, height=0.0, device=device) + env_idx = torch.tensor([x for x in range(num_cubes)], dtype=torch.int32) + + # Play sim + sim.reset() + + # Check if cube_object is initialized + assert cube_object.is_initialized + + # change center of mass offset from link frame + if with_offset: + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_cubes, 1) + else: + offset = torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_cubes, 1) + + # Read current COMs, mutate the translation, write back via the OVPhysX + # ``set_coms_index`` setter (PhysX uses ``root_view.set_coms`` for the same + # operation; OVPhysX wraps the wheel ``RIGID_BODY_COM_POSE`` write in + # :meth:`set_coms_index`, which follows the PhysX ``wp.transformf`` contract). + com = cube_object.data.body_com_pose_b.torch.clone() # shape (N, 1, 7) + com[..., :3] = offset.to(com.device).unsqueeze(1) + cube_object.set_coms_index( + coms=wp.from_torch(com.contiguous(), dtype=wp.transformf), + env_ids=wp.from_torch(env_idx, dtype=wp.int32), + ) + + # check ceter of mass has been set + torch.testing.assert_close(cube_object.data.body_com_pose_b.torch, com) + + # random z spin velocity + spin_twist = torch.zeros(6, device=device) + spin_twist[5] = torch.randn(1, device=device) + + # Simulate physics + for _ in range(100): + # spin the object around Z axis (com) + cube_object.write_root_velocity_to_sim_index(root_velocity=spin_twist.repeat(num_cubes, 1)) + # perform rendering + sim.step() + # update object + cube_object.update(sim.cfg.dt) + + # get state properties + root_link_pose_w = cube_object.data.root_link_pose_w.torch + root_link_vel_w = cube_object.data.root_link_vel_w.torch + root_com_pose_w = cube_object.data.root_com_pose_w.torch + root_com_vel_w = cube_object.data.root_com_vel_w.torch + body_link_pose_w = cube_object.data.body_link_pose_w.torch + body_link_vel_w = cube_object.data.body_link_vel_w.torch + body_com_pose_w = cube_object.data.body_com_pose_w.torch + body_com_vel_w = cube_object.data.body_com_vel_w.torch + + # if offset is [0,0,0] all root_state_%_w will match and all body_%_w will match + if not with_offset: + torch.testing.assert_close(root_link_pose_w, root_com_pose_w) + torch.testing.assert_close(root_com_vel_w, root_link_vel_w) + torch.testing.assert_close(root_link_pose_w, root_link_pose_w) + torch.testing.assert_close(root_com_vel_w, root_link_vel_w) + torch.testing.assert_close(body_link_pose_w, body_com_pose_w) + torch.testing.assert_close(body_com_vel_w, body_link_vel_w) + torch.testing.assert_close(body_link_pose_w, body_link_pose_w) + torch.testing.assert_close(body_com_vel_w, body_link_vel_w) + else: + # cubes are spinning around center of mass + # position will not match + # center of mass position will be constant (i.e. spinning around com) + torch.testing.assert_close(env_pos + offset, root_com_pose_w[..., :3]) + torch.testing.assert_close(env_pos + offset, body_com_pose_w[..., :3].squeeze(-2)) + # link position will be moving but should stay constant away from center of mass + root_link_state_pos_rel_com = quat_apply_inverse( + root_link_pose_w[..., 3:], + root_link_pose_w[..., :3] - root_com_pose_w[..., :3], + ) + torch.testing.assert_close(-offset, root_link_state_pos_rel_com) + body_link_state_pos_rel_com = quat_apply_inverse( + body_link_pose_w[..., 3:], + body_link_pose_w[..., :3] - body_com_pose_w[..., :3], + ) + torch.testing.assert_close(-offset, body_link_state_pos_rel_com.squeeze(-2)) + + # orientation of com will be a constant rotation from link orientation + com_quat_b = cube_object.data.body_com_quat_b.torch + com_quat_w = quat_mul(body_link_pose_w[..., 3:], com_quat_b) + torch.testing.assert_close(com_quat_w, body_com_pose_w[..., 3:]) + torch.testing.assert_close(com_quat_w.squeeze(-2), root_com_pose_w[..., 3:]) + + # orientation of link will match root state will always match + torch.testing.assert_close(root_link_pose_w[..., 3:], root_link_pose_w[..., 3:]) + torch.testing.assert_close(body_link_pose_w[..., 3:], body_link_pose_w[..., 3:]) + + # lin_vel will not match + # center of mass vel will be constant (i.e. spinning around com) + torch.testing.assert_close(torch.zeros_like(root_com_vel_w[..., :3]), root_com_vel_w[..., :3]) + torch.testing.assert_close(torch.zeros_like(body_com_vel_w[..., :3]), body_com_vel_w[..., :3]) + # link frame will be moving, and should be equal to input angular velocity cross offset + lin_vel_rel_root_gt = quat_apply_inverse(root_link_pose_w[..., 3:], root_link_vel_w[..., :3]) + lin_vel_rel_body_gt = quat_apply_inverse(body_link_pose_w[..., 3:], body_link_vel_w[..., :3]) + lin_vel_rel_gt = torch.linalg.cross(spin_twist.repeat(num_cubes, 1)[..., 3:], -offset) + torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_root_gt, atol=1e-4, rtol=1e-4) + torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_body_gt.squeeze(-2), atol=1e-4, rtol=1e-4) + + # ang_vel will always match + torch.testing.assert_close(root_com_vel_w[..., 3:], root_com_vel_w[..., 3:]) + torch.testing.assert_close(root_com_vel_w[..., 3:], root_link_vel_w[..., 3:]) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_com_vel_w[..., 3:]) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("with_offset", [True, False]) +@pytest.mark.parametrize("state_location", ["com", "link"]) +@pytest.mark.isaacsim_ci +def test_write_root_state(num_cubes, device, with_offset, state_location): + """Test the setters for root_state using both the link frame and center of mass as reference frame.""" + with _ovphysx_sim_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + # Create a scene with random cubes + cube_object, env_pos = generate_cubes_scene(num_cubes=num_cubes, height=0.0, device=device) + env_idx = torch.tensor([x for x in range(num_cubes)], dtype=torch.int32) + + # Play sim + sim.reset() + + # Check if cube_object is initialized + assert cube_object.is_initialized + + # change center of mass offset from link frame + if with_offset: + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_cubes, 1) + else: + offset = torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_cubes, 1) + + com = cube_object.data.body_com_pose_b.torch.clone() # shape (N, 1, 7) + com[..., :3] = offset.to(com.device).unsqueeze(1) + cube_object.set_coms_index( + coms=wp.from_torch(com.contiguous(), dtype=wp.transformf), + env_ids=wp.from_torch(env_idx, dtype=wp.int32), + ) + + # check center of mass has been set + torch.testing.assert_close(cube_object.data.body_com_pose_b.torch, com) + + rand_state = torch.zeros(num_cubes, 13, device=device) + rand_state[..., :7] = cube_object.data.default_root_pose.torch + rand_state[..., :3] += env_pos + # make quaternion a unit vector + rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) + + env_idx = env_idx.to(device) + for i in range(10): + # perform step + sim.step() + # update buffers + cube_object.update(sim.cfg.dt) + + if state_location == "com": + if i % 2 == 0: + cube_object.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7]) + cube_object.write_root_com_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) + else: + cube_object.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_idx) + cube_object.write_root_com_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_idx) + elif state_location == "link": + if i % 2 == 0: + cube_object.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7]) + cube_object.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) + else: + cube_object.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_idx) + cube_object.write_root_link_velocity_to_sim_index( + root_velocity=rand_state[..., 7:], env_ids=env_idx + ) + + if state_location == "com": + torch.testing.assert_close(rand_state[..., :7], cube_object.data.root_com_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], cube_object.data.root_com_vel_w.torch) + elif state_location == "link": + torch.testing.assert_close(rand_state[..., :7], cube_object.data.root_link_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], cube_object.data.root_link_vel_w.torch) + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("with_offset", [True]) +@pytest.mark.parametrize("state_location", ["com", "link", "root"]) +@pytest.mark.isaacsim_ci +def test_write_state_functions_data_consistency(num_cubes, device, with_offset, state_location): + """Test the setters for root_state using both the link frame and center of mass as reference frame.""" + with _ovphysx_sim_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + # Create a scene with random cubes + cube_object, env_pos = generate_cubes_scene(num_cubes=num_cubes, height=0.0, device=device) + env_idx = torch.tensor([x for x in range(num_cubes)], dtype=torch.int32) + + # Play sim + sim.reset() + + # Check if cube_object is initialized + assert cube_object.is_initialized + + # change center of mass offset from link frame + if with_offset: + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_cubes, 1) + else: + offset = torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_cubes, 1) + + com = cube_object.data.body_com_pose_b.torch.clone() # shape (N, 1, 7) + com[..., :3] = offset.to(com.device).unsqueeze(1) + cube_object.set_coms_index( + coms=wp.from_torch(com.contiguous(), dtype=wp.transformf), + env_ids=wp.from_torch(env_idx, dtype=wp.int32), + ) + + # check ceter of mass has been set + torch.testing.assert_close(cube_object.data.body_com_pose_b.torch, com) + + rand_state = torch.rand(num_cubes, 13, device=device) + rand_state[..., :3] += env_pos + # make quaternion a unit vector + rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) + + env_idx = env_idx.to(device) + + # perform step + sim.step() + # update buffers + cube_object.update(sim.cfg.dt) + + if state_location == "com": + cube_object.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7]) + cube_object.write_root_com_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) + elif state_location == "link": + cube_object.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7]) + cube_object.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) + elif state_location == "root": + cube_object.write_root_pose_to_sim_index(root_pose=rand_state[..., :7]) + cube_object.write_root_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) + + if state_location == "com": + root_com_pose_w = cube_object.data.root_com_pose_w.torch + root_com_vel_w = cube_object.data.root_com_vel_w.torch + body_com_pose_b = cube_object.data.body_com_pose_b.torch + expected_root_link_pos, expected_root_link_quat = combine_frame_transforms( + root_com_pose_w[:, :3], + root_com_pose_w[:, 3:], + quat_rotate(quat_inv(body_com_pose_b[:, 0, 3:7]), -body_com_pose_b[:, 0, :3]), + quat_inv(body_com_pose_b[:, 0, 3:7]), + ) + expected_root_link_pose = torch.cat((expected_root_link_pos, expected_root_link_quat), dim=1) + root_link_pose_w = cube_object.data.root_link_pose_w.torch + root_link_vel_w = cube_object.data.root_link_vel_w.torch + # test both root_pose and root_link successfully updated when root_com updates + torch.testing.assert_close(expected_root_link_pose, root_link_pose_w) + # skip lin_vel because it differs from link frame, this should be fine because we are only checking + # if velocity update is triggered, which can be determined by comparing angular velocity + torch.testing.assert_close(root_com_vel_w[:, 3:], root_link_vel_w[:, 3:]) + torch.testing.assert_close(expected_root_link_pose, root_link_pose_w) + torch.testing.assert_close(root_com_vel_w[:, 3:], cube_object.data.root_com_vel_w.torch[:, 3:]) + elif state_location == "link": + root_link_pose_w = cube_object.data.root_link_pose_w.torch + root_link_vel_w = cube_object.data.root_link_vel_w.torch + body_com_pose_b = cube_object.data.body_com_pose_b.torch + expected_com_pos, expected_com_quat = combine_frame_transforms( + root_link_pose_w[:, :3], + root_link_pose_w[:, 3:], + body_com_pose_b[:, 0, :3], + body_com_pose_b[:, 0, 3:7], + ) + expected_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1) + root_com_pose_w = cube_object.data.root_com_pose_w.torch + root_com_vel_w = cube_object.data.root_com_vel_w.torch + # test both root_pose and root_com successfully updated when root_link updates + torch.testing.assert_close(expected_com_pose, root_com_pose_w) + # skip lin_vel because it differs from link frame, this should be fine because we are only checking + # if velocity update is triggered, which can be determined by comparing angular velocity + torch.testing.assert_close(root_link_vel_w[:, 3:], root_com_vel_w[:, 3:]) + torch.testing.assert_close(root_link_pose_w, cube_object.data.root_link_pose_w.torch) + torch.testing.assert_close(root_link_vel_w[:, 3:], cube_object.data.root_com_vel_w.torch[:, 3:]) + elif state_location == "root": + root_link_pose_w = cube_object.data.root_link_pose_w.torch + root_com_vel_w = cube_object.data.root_com_vel_w.torch + body_com_pose_b = cube_object.data.body_com_pose_b.torch + expected_com_pos, expected_com_quat = combine_frame_transforms( + root_link_pose_w[:, :3], + root_link_pose_w[:, 3:], + body_com_pose_b[:, 0, :3], + body_com_pose_b[:, 0, 3:7], + ) + expected_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1) + root_com_pose_w = cube_object.data.root_com_pose_w.torch + root_link_vel_w = cube_object.data.root_link_vel_w.torch + # test both root_com and root_link successfully updated when root_pose updates + torch.testing.assert_close(expected_com_pose, root_com_pose_w) + torch.testing.assert_close(root_com_vel_w, cube_object.data.root_com_vel_w.torch) + torch.testing.assert_close(root_link_pose_w, cube_object.data.root_link_pose_w.torch) + torch.testing.assert_close(root_com_vel_w[:, 3:], root_link_vel_w[:, 3:]) + + +@pytest.mark.isaacsim_ci +def test_warmup_attach_stage_not_called_for_cpu(): + """Regression test: ``physx.warmup_gpu()`` must not be called for CPU. + + OVPhysX-equivalent of PhysX's ``test_warmup_attach_stage_not_called_for_cpu``: + PhysX guards :meth:`attach_stage` with ``if is_gpu:`` so the CPU MBP + broadphase is not double-initialised. The OVPhysX manager has the same + structural guard around :meth:`OvPhysxManager._physx.warmup_gpu`: it is + only invoked when ``ovphysx_device == "gpu"``. + + We monkey-patch ``OvPhysxManager._physx`` with a :class:`MagicMock` + wrapping the live PhysX object so that ``warmup_gpu`` becomes a spy while + other calls continue to forward, then assert ``warmup_gpu.call_count == 0`` + after a CPU-mode :meth:`sim.reset`. + + The test always runs CPU regardless of session parametrization, so it is + skipped when the session-locked device is anything other than CPU. The + skip is enforced inline (rather than in the autouse fixture) so the rest + of the suite can still pin to GPU when invoked together. + """ + if _LOCKED_DEVICE[0] not in (None, "cpu"): + pytest.skip( + f"ovphysx process-global device lock is held by '{_LOCKED_DEVICE[0]}'; cannot run " + "CPU-only regression test in the same session." + ) + _LOCKED_DEVICE[0] = "cpu" + + with _ovphysx_sim_context(device="cpu", add_ground_plane=True, dt=0.01, auto_add_lighting=True) as sim: + # Allocate a single rigid body so the manager has something to load. + generate_cubes_scene(num_cubes=1, height=1.0, device="cpu") + + # First reset constructs (or reuses) the real ovphysx.PhysX so we have + # a live instance to wrap. The PhysX object is a C++ binding, so we + # cannot patch attributes directly — replace the class-level reference + # with a MagicMock(wraps=...) that forwards every call. + sim.reset() + original_physx = OvPhysxManager._physx + assert original_physx is not None, "PhysX should be constructed after sim.reset()" + spy = MagicMock(wraps=original_physx) + OvPhysxManager._physx = spy + # Force _warmup_and_load to run again on the next reset so the spy + # observes the warmup_gpu (or non-call) decision; close() resets + # _warmup_done back to False but we just called sim.reset() above. + OvPhysxManager._warmup_done = False + try: + sim.reset() + finally: + OvPhysxManager._physx = original_physx + + assert spy.warmup_gpu.call_count == 0, ( + f"warmup_gpu() was called {spy.warmup_gpu.call_count} time(s) during CPU warmup. " + "OvPhysxManager._warmup_and_load() must guard warmup_gpu() with " + "ovphysx_device == 'gpu' so the CPU pipeline is not mis-initialised." + ) diff --git a/source/isaaclab_ovphysx/test/assets/test_rigid_object_collection.py b/source/isaaclab_ovphysx/test/assets/test_rigid_object_collection.py new file mode 100644 index 000000000000..d2163989d32b --- /dev/null +++ b/source/isaaclab_ovphysx/test/assets/test_rigid_object_collection.py @@ -0,0 +1,957 @@ +# 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 + +# ignore private usage of variables warning +# pyright: reportPrivateUsage=none + + +"""Real-backend tests for the OVPhysX RigidObjectCollection. + +Run via ``./scripts/run_ovphysx.sh -m pytest`` (kitless, no ``AppLauncher``). + +``ovphysx<=0.3.7`` binds device mode (CPU vs GPU) at the C++ layer on the +first ``ovphysx.PhysX(device=...)`` construction and cannot swap it without a +process restart. Full coverage therefore requires two separate pytest +invocations -- once with ``-k 'cpu'`` and once with ``-k 'cuda:0'``. The +``_ovphysx_skip_other_device`` autouse fixture below preempts the manager's +:exc:`RuntimeError` by ``pytest.skip``-ing on the unlocked device so +single-device runs finish cleanly. +""" + +from __future__ import annotations + +import sys + +import pytest +import torch +import warp as wp +from isaaclab_ovphysx.assets import RigidObjectCollection +from isaaclab_ovphysx.physics import OvPhysxCfg + +import isaaclab.sim as sim_utils +from isaaclab.assets import RigidObjectCfg, RigidObjectCollectionCfg +from isaaclab.sim import SimulationCfg, build_simulation_context +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.math import ( + combine_frame_transforms, + default_orientation, + quat_apply_inverse, + quat_inv, + quat_mul, + quat_rotate, + random_orientation, + subtract_frame_transforms, +) + +wp.init() + + +_LOCKED_DEVICE: list[str | None] = [None] +"""Device the session pins to on the first parametrized test that runs.""" + + +@pytest.fixture(autouse=True) +def _ovphysx_skip_other_device(request): + """Skip parametrized tests on the device the session is not pinned to. + + See the module docstring for the wheel's process-global device-mode lock. + """ + callspec = getattr(request.node, "callspec", None) + device = callspec.params.get("device") if callspec is not None else None + if device is None: + # Test does not parametrize on device. + return + locked = _LOCKED_DEVICE[0] + if locked is None: + _LOCKED_DEVICE[0] = device + return + if device != locked: + pytest.skip( + f"ovphysx process-global device lock is held by '{locked}'; cannot run '{device}' " + "tests in the same session. Run pytest twice (once per device) for full coverage." + ) + + +def _ovphysx_sim_context(device: str, **kwargs): + """Wrapper around :func:`build_simulation_context` that injects OVPhysX cfg. + + PhysX tests pass ``device=device`` directly and let + :func:`build_simulation_context` build a default :class:`SimulationCfg`. + OVPhysX needs ``physics=OvPhysxCfg()`` set on the cfg so the manager + dispatches to OVPhysX rather than PhysX, so we build the cfg here and + pass it through. ``gravity_enabled`` is consumed locally (it is ignored + by ``build_simulation_context`` once a ``sim_cfg`` is provided). + ``add_ground_plane``, ``auto_add_lighting``, and other kwargs continue + to flow through ``build_simulation_context`` as before. + """ + dt = kwargs.pop("dt", 1.0 / 60.0) + gravity_enabled = kwargs.pop("gravity_enabled", True) + gravity = (0.0, 0.0, -9.81) if gravity_enabled else (0.0, 0.0, 0.0) + sim_cfg = SimulationCfg(physics=OvPhysxCfg(), device=device, dt=dt, gravity=gravity) + return build_simulation_context(device=device, sim_cfg=sim_cfg, **kwargs) + + +# --------------------------------------------------------------------------- +# Material-property gap (xfail reason shared by the test below) +# --------------------------------------------------------------------------- + +_MATERIAL_GAP_REASON = ( + "Requires RIGID_BODY_MATERIAL TensorType (or a view-helper) on the ovphysx " + "wheel side. RigidObjectCollection.root_view is a per-tensor-type bindings dict on " + "OVPhysX, so root_view.get_material_properties() / set_material_properties() " + "are not available. See " + "docs/superpowers/specs/2026-04-28-ovphysx-wheel-gaps-for-marco.md." +) + + +def generate_cubes_scene( + num_envs: int = 1, + num_cubes: int = 1, + height=1.0, + has_api: bool = True, + kinematic_enabled: bool = False, + device: str = "cuda:0", +) -> tuple[RigidObjectCollection, torch.Tensor]: + """Generate a scene with the provided number of cubes. + + Args: + num_envs: Number of envs to generate. + num_cubes: Number of cubes to generate. + height: Height of the cubes. + has_api: Whether the cubes have a rigid body API on them. + kinematic_enabled: Whether the cubes are kinematic. + device: Device to use for the simulation. + + Returns: + A tuple containing the rigid object representing the cubes and the origins of the cubes. + + """ + origins = torch.tensor([(i * 3.0, 0, height) for i in range(num_envs)]).to(device) + # Create Top-level Xforms, one for each cube + for i, origin in enumerate(origins): + sim_utils.create_prim(f"/World/Table_{i}", "Xform", translation=origin) + + # Resolve spawn configuration + if has_api: + spawn_cfg = sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", + rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=kinematic_enabled), + ) + else: + # since no rigid body properties defined, this is just a static collider + spawn_cfg = sim_utils.CuboidCfg( + size=(0.1, 0.1, 0.1), + collision_props=sim_utils.CollisionPropertiesCfg(), + ) + + # create the rigid object configs. OVPhysX matches prim paths via fnmatch globs (not regex), + # so use ``Table_*`` rather than the PhysX ``Table_.*`` form. + cube_config_dict = {} + for i in range(num_cubes): + cube_object_cfg = RigidObjectCfg( + prim_path=f"/World/Table_*/Object_{i}", + spawn=spawn_cfg, + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 3 * i, height)), + ) + cube_config_dict[f"cube_{i}"] = cube_object_cfg + # create the rigid object collection + cube_object_collection_cfg = RigidObjectCollectionCfg(rigid_objects=cube_config_dict) + cube_object_colection = RigidObjectCollection(cfg=cube_object_collection_cfg) + + return cube_object_colection, origins + + +@pytest.mark.parametrize("num_envs", [1, 2]) +@pytest.mark.parametrize("num_cubes", [1, 3]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_initialization(num_envs, num_cubes, device): + """Test initialization for prim with rigid body API at the provided prim path.""" + with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: + object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(object_collection) < 10 + + # Play sim + sim.reset() + + # Check if object is initialized + assert object_collection.is_initialized + assert len(object_collection.body_names) == num_cubes + + # Check buffers that exist and have correct shapes + assert object_collection.data.body_link_pos_w.torch.shape == (num_envs, num_cubes, 3) + assert object_collection.data.body_link_quat_w.torch.shape == (num_envs, num_cubes, 4) + assert object_collection.data.body_mass.torch.shape == (num_envs, num_cubes) + assert object_collection.data.body_inertia.torch.shape == (num_envs, num_cubes, 9) + + # Simulate physics + for _ in range(2): + sim.step() + object_collection.update(sim.cfg.dt) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_id_conversion(device): + """Test environment and object index conversion to physics view indices.""" + with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: + object_collection, _ = generate_cubes_scene(num_envs=2, num_cubes=3, device=device) + + # Play sim + sim.reset() + + expected = [ + torch.tensor([4, 5], device=device, dtype=torch.int32), + torch.tensor([4], device=device, dtype=torch.int32), + torch.tensor([0, 2, 4], device=device, dtype=torch.int32), + torch.tensor([1, 3, 5], device=device, dtype=torch.int32), + ] + + torch_all_env_indices = wp.to_torch(object_collection._ALL_ENV_INDICES) + torch_all_body_indices = wp.to_torch(object_collection._ALL_BODY_INDICES) + + view_ids = object_collection._env_body_ids_to_view_ids( + torch_all_env_indices, torch_all_body_indices[None, 2], device=device + ) + assert (wp.to_torch(view_ids) == expected[0]).all() + view_ids = object_collection._env_body_ids_to_view_ids( + torch_all_env_indices[None, 0], torch_all_body_indices[None, 2], device=device + ) + assert (wp.to_torch(view_ids) == expected[1]).all() + view_ids = object_collection._env_body_ids_to_view_ids( + torch_all_env_indices[None, 0], torch_all_body_indices, device=device + ) + assert (wp.to_torch(view_ids) == expected[2]).all() + view_ids = object_collection._env_body_ids_to_view_ids( + torch_all_env_indices[None, 1], torch_all_body_indices, device=device + ) + assert (wp.to_torch(view_ids) == expected[3]).all() + + +@pytest.mark.parametrize("num_envs", [1, 2]) +@pytest.mark.parametrize("num_cubes", [1, 3]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_initialization_with_kinematic_enabled(num_envs, num_cubes, device): + """Test that initialization for prim with kinematic flag enabled.""" + with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: + object_collection, origins = generate_cubes_scene( + num_envs=num_envs, num_cubes=num_cubes, kinematic_enabled=True, device=device + ) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(object_collection) < 10 + + # Play sim + sim.reset() + + # Check if object is initialized + assert object_collection.is_initialized + assert len(object_collection.body_names) == num_cubes + + # Check buffers that exist and have correct shapes + assert object_collection.data.body_link_pos_w.torch.shape == (num_envs, num_cubes, 3) + assert object_collection.data.body_link_quat_w.torch.shape == (num_envs, num_cubes, 4) + + # Simulate physics + for _ in range(2): + sim.step() + object_collection.update(sim.cfg.dt) + # check that the object is kinematic + default_body_pose = object_collection.data.default_body_pose.torch.clone() + default_body_vel = object_collection.data.default_body_vel.torch.clone() + default_body_pose[..., :3] += origins.unsqueeze(1) + torch.testing.assert_close(object_collection.data.body_link_pose_w.torch, default_body_pose) + torch.testing.assert_close(object_collection.data.body_link_vel_w.torch, default_body_vel) + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_initialization_with_no_rigid_body(num_cubes, device): + """Test that initialization fails when no rigid body is found at the provided prim path.""" + with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: + object_collection, _ = generate_cubes_scene(num_cubes=num_cubes, has_api=False, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(object_collection) < 10 + + # Play sim + with pytest.raises(RuntimeError): + sim.reset() + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_external_force_buffer(device): + """Test if external force buffer correctly updates in the force value is zero case.""" + num_envs = 2 + num_cubes = 1 + with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: + object_collection, origins = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) + sim.reset() + + # find objects to apply the force + object_ids, object_names = object_collection.find_bodies(".*") + # reset object + object_collection.reset() + + # perform simulation + for step in range(5): + # initiate force tensor + external_wrench_b = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=sim.device) + + # decide if zero or non-zero force + if step == 0 or step == 3: + force = 1.0 + else: + force = 0.0 + + # apply force to the object + external_wrench_b[:, :, 0] = force + external_wrench_b[:, :, 3] = force + + object_collection.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + body_ids=object_ids, + env_ids=None, + ) + + # check if the object collection's force and torque buffers are correctly updated + for i in range(num_envs): + assert object_collection._permanent_wrench_composer.composed_force.torch[i, 0, 0].item() == force + assert object_collection._permanent_wrench_composer.composed_torque.torch[i, 0, 0].item() == force + + object_collection.instantaneous_wrench_composer.add_forces_and_torques_index( + body_ids=object_ids, + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + ) + + # apply action to the object collection + object_collection.write_data_to_sim() + sim.step() + object_collection.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_envs", [1, 2]) +@pytest.mark.parametrize("num_cubes", [1, 4]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_external_force_on_single_body(num_envs, num_cubes, device): + """Test application of external force on the base of the object.""" + with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: + object_collection, origins = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) + sim.reset() + + # find objects to apply the force + object_ids, object_names = object_collection.find_bodies(".*") + + # Sample a force equal to the weight of the object + external_wrench_b = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=sim.device) + # Every 2nd cube should have a force applied to it + external_wrench_b[:, 0::2, 2] = 9.81 * object_collection.data.body_mass.torch[:, 0::2] + + for i in range(5): + # reset object state + body_pose = object_collection.data.default_body_pose.torch.clone() + body_vel = object_collection.data.default_body_vel.torch.clone() + # need to shift the position of the cubes otherwise they will be on top of each other + body_pose[..., :2] += origins.unsqueeze(1)[..., :2] + object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) + object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) + # reset object + object_collection.reset() + + is_global = False + if i % 2 == 0: + positions = object_collection.data.body_link_pos_w.torch[:, object_ids, :3] + is_global = True + else: + positions = None + + # apply force + object_collection.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=positions, + body_ids=object_ids, + env_ids=None, + is_global=is_global, + ) + for _ in range(10): + # write data to sim + object_collection.write_data_to_sim() + # step sim + sim.step() + # update object collection + object_collection.update(sim.cfg.dt) + + # First object should still be at the same Z position (1.0) + torch.testing.assert_close( + object_collection.data.body_link_pos_w.torch[:, 0::2, 2], + torch.ones_like(object_collection.data.body_link_pos_w.torch[:, 0::2, 2]), + ) + # Second object should have fallen, so it's Z height should be less than initial height of 1.0 + assert torch.all(object_collection.data.body_link_pos_w.torch[:, 1::2, 2] < 1.0) + + +@pytest.mark.parametrize("num_envs", [1, 2]) +@pytest.mark.parametrize("num_cubes", [1, 4]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_external_force_on_single_body_at_position(num_envs, num_cubes, device): + """Test application of external force on the base of the object at a specific position. + + In this test, we apply a force equal to the weight of an object on the base of + one of the objects at 1m in the Y direction, we check that the object rotates around it's X axis. + For the other object, we do not apply any force and check that it falls down. + """ + with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: + object_collection, origins = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) + sim.reset() + + # find objects to apply the force + object_ids, object_names = object_collection.find_bodies(".*") + + # Sample a force equal to the weight of the object + external_wrench_b = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=sim.device) + external_wrench_positions_b = torch.zeros( + object_collection.num_instances, len(object_ids), 3, device=sim.device + ) + # Every 2nd cube should have a force applied to it + external_wrench_b[:, 0::2, 2] = 500.0 + external_wrench_positions_b[:, 0::2, 1] = 1.0 + + # Desired force and torque + for i in range(5): + # reset object state + body_pose = object_collection.data.default_body_pose.torch.clone() + body_vel = object_collection.data.default_body_vel.torch.clone() + # need to shift the position of the cubes otherwise they will be on top of each other + body_pose[..., :2] += origins.unsqueeze(1)[..., :2] + object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) + object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) + # reset object + object_collection.reset() + + is_global = False + if i % 2 == 0: + body_com_pos_w = object_collection.data.body_link_pos_w.torch[:, object_ids, :3] + external_wrench_positions_b[..., 0] = 0.0 + external_wrench_positions_b[..., 1] = 1.0 + external_wrench_positions_b[..., 2] = 0.0 + external_wrench_positions_b += body_com_pos_w + is_global = True + else: + external_wrench_positions_b[..., 0] = 0.0 + external_wrench_positions_b[..., 1] = 1.0 + external_wrench_positions_b[..., 2] = 0.0 + + # apply force + object_collection.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=external_wrench_positions_b, + body_ids=object_ids, + env_ids=None, + is_global=is_global, + ) + object_collection.permanent_wrench_composer.add_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=external_wrench_positions_b, + body_ids=object_ids, + is_global=is_global, + ) + + for _ in range(10): + # write data to sim + object_collection.write_data_to_sim() + # step sim + sim.step() + # update object collection + object_collection.update(sim.cfg.dt) + + # First object should be rotating around it's X axis + assert torch.all(object_collection.data.body_com_ang_vel_b.torch[:, 0::2, 0] > 0.1) + # Second object should have fallen, so it's Z height should be less than initial height of 1.0 + assert torch.all(object_collection.data.body_link_pos_w.torch[:, 1::2, 2] < 1.0) + + +@pytest.mark.parametrize("num_envs", [1, 3]) +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("gravity_enabled", [False]) +@pytest.mark.isaacsim_ci +def test_set_object_state(num_envs, num_cubes, device, gravity_enabled): + """Test setting the state of the object. + + .. note:: + Turn off gravity for this test as we don't want any external forces acting on the object + to ensure state remains static + """ + with _ovphysx_sim_context(device=device, gravity_enabled=gravity_enabled, auto_add_lighting=True) as sim: + object_collection, origins = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) + sim.reset() + + state_types = ["body_link_pos_w", "body_link_quat_w", "body_com_lin_vel_w", "body_com_ang_vel_w"] + + # Set each state type individually as they are dependent on each other + for state_type_to_randomize in state_types: + state_dict = { + "body_link_pos_w": torch.zeros_like(object_collection.data.body_link_pos_w.torch, device=sim.device), + "body_link_quat_w": default_orientation(num=num_cubes * num_envs, device=sim.device).view( + num_envs, num_cubes, 4 + ), + "body_com_lin_vel_w": torch.zeros_like( + object_collection.data.body_com_lin_vel_w.torch, device=sim.device + ), + "body_com_ang_vel_w": torch.zeros_like( + object_collection.data.body_com_ang_vel_w.torch, device=sim.device + ), + } + + for _ in range(5): + # reset object + object_collection.reset() + + # Set random state + if state_type_to_randomize == "body_link_quat_w": + state_dict[state_type_to_randomize] = random_orientation( + num=num_cubes * num_envs, device=sim.device + ).view(num_envs, num_cubes, 4) + else: + state_dict[state_type_to_randomize] = torch.randn(num_envs, num_cubes, 3, device=sim.device) + # make sure objects do not overlap + if state_type_to_randomize == "body_link_pos_w": + state_dict[state_type_to_randomize][..., :2] += origins.unsqueeze(1)[..., :2] + + # perform simulation + for _ in range(5): + body_pose = torch.cat( + [state_dict["body_link_pos_w"], state_dict["body_link_quat_w"]], + dim=-1, + ) + body_vel = torch.cat( + [state_dict["body_com_lin_vel_w"], state_dict["body_com_ang_vel_w"]], + dim=-1, + ) + # reset object state + object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) + object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) + sim.step() + + # assert that set object quantities are equal to the ones set in the state_dict + for key, expected_value in state_dict.items(): + value = getattr(object_collection.data, key).torch + torch.testing.assert_close(value, expected_value, rtol=1e-5, atol=1e-5) + + object_collection.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_envs", [1, 4]) +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("with_offset", [True, False]) +@pytest.mark.parametrize("gravity_enabled", [False]) +@pytest.mark.isaacsim_ci +def test_object_state_properties(num_envs, num_cubes, device, with_offset, gravity_enabled): + """Test the object_com_state_w and object_link_state_w properties.""" + with _ovphysx_sim_context(device=device, gravity_enabled=gravity_enabled, auto_add_lighting=True) as sim: + cube_object, env_pos = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, height=0.0, device=device) + env_ids = torch.tensor([x for x in range(num_envs)], dtype=torch.int32) + + sim.reset() + + # check if cube_object is initialized + assert cube_object.is_initialized + + # change center of mass offset from link frame + offset = ( + torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) + if with_offset + else torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) + ) + + # Read current COMs, mutate the translation, write back via the OVPhysX + # ``set_coms_index`` setter (PhysX uses ``root_view.set_coms`` + reshape helpers + # for the same operation; OVPhysX wraps the wheel RIGID_BODY_COM_POSE write in + # :meth:`set_coms_index`). + com = cube_object.data.body_com_pose_b.torch.clone() # shape (num_envs, num_cubes, 7) + com[..., :3] = offset.to(com.device) + cube_object.set_coms_index( + coms=wp.from_torch(com.contiguous(), dtype=wp.transformf), + env_ids=wp.from_torch(env_ids, dtype=wp.int32), + ) + + # check center of mass has been set + torch.testing.assert_close(cube_object.data.body_com_pose_b.torch, com) + + # random z spin velocity + spin_twist = torch.zeros(6, device=device) + spin_twist[5] = torch.randn(1, device=device) + + # initial spawn point + init_com = cube_object.data.body_com_pose_w.torch[..., :3] + + for i in range(10): + # spin the object around Z axis (com) + cube_object.write_body_com_velocity_to_sim_index(body_velocities=spin_twist.repeat(num_envs, num_cubes, 1)) + sim.step() + cube_object.update(sim.cfg.dt) + + # get state properties + object_link_pose_w = cube_object.data.body_link_pose_w.torch + object_link_vel_w = cube_object.data.body_link_vel_w.torch + object_com_pose_w = cube_object.data.body_com_pose_w.torch + object_com_vel_w = cube_object.data.body_com_vel_w.torch + + # if offset is [0,0,0] all object_state_%_w will match and all body_%_w will match + if not with_offset: + torch.testing.assert_close(object_link_pose_w, object_com_pose_w) + torch.testing.assert_close(object_com_vel_w, object_link_vel_w) + else: + # cubes are spinning around center of mass + # position will not match + # center of mass position will be constant (i.e. spinning around com) + torch.testing.assert_close(init_com, object_com_pose_w[..., :3]) + + # link position will be moving but should stay constant away from center of mass + object_link_state_pos_rel_com = quat_apply_inverse( + object_link_pose_w[..., 3:], + object_link_pose_w[..., :3] - object_com_pose_w[..., :3], + ) + + torch.testing.assert_close(-offset, object_link_state_pos_rel_com) + + # orientation of com will be a constant rotation from link orientation + com_quat_b = cube_object.data.body_com_quat_b.torch + com_quat_w = quat_mul(object_link_pose_w[..., 3:], com_quat_b) + torch.testing.assert_close(com_quat_w, object_com_pose_w[..., 3:]) + + # orientation of link will match object state will always match + torch.testing.assert_close(object_link_pose_w[..., 3:], object_link_pose_w[..., 3:]) + + # lin_vel will not match + # center of mass vel will be constant (i.e. spinning around com) + torch.testing.assert_close( + torch.zeros_like(object_com_vel_w[..., :3]), + object_com_vel_w[..., :3], + ) + + # link frame will be moving, and should be equal to input angular velocity cross offset + lin_vel_rel_object_gt = quat_apply_inverse(object_link_pose_w[..., 3:], object_link_vel_w[..., :3]) + lin_vel_rel_gt = torch.linalg.cross(spin_twist.repeat(num_envs, num_cubes, 1)[..., 3:], -offset) + torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_object_gt, atol=1e-4, rtol=1e-3) + + # ang_vel will always match + torch.testing.assert_close(object_com_vel_w[..., 3:], object_com_vel_w[..., 3:]) + torch.testing.assert_close(object_com_vel_w[..., 3:], object_link_vel_w[..., 3:]) + + +@pytest.mark.parametrize("num_envs", [1, 3]) +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("with_offset", [True, False]) +@pytest.mark.parametrize("state_location", ["com", "link"]) +@pytest.mark.parametrize("gravity_enabled", [False]) +@pytest.mark.isaacsim_ci +def test_write_object_state(num_envs, num_cubes, device, with_offset, state_location, gravity_enabled): + """Test the setters for object_state using both the link frame and center of mass as reference frame.""" + with _ovphysx_sim_context(device=device, gravity_enabled=gravity_enabled, auto_add_lighting=True) as sim: + # Create a scene with random cubes + cube_object, env_pos = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, height=0.0, device=device) + env_ids = torch.tensor([x for x in range(num_envs)], dtype=torch.int32) + object_ids = torch.tensor([x for x in range(num_cubes)], dtype=torch.int32) + + sim.reset() + + # Check if cube_object is initialized + assert cube_object.is_initialized + + # change center of mass offset from link frame + offset = ( + torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) + if with_offset + else torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) + ) + + com = cube_object.data.body_com_pose_b.torch.clone() # shape (num_envs, num_cubes, 7) + com[..., :3] = offset.to(com.device) + cube_object.set_coms_index( + coms=wp.from_torch(com.contiguous(), dtype=wp.transformf), + env_ids=wp.from_torch(env_ids, dtype=wp.int32), + ) + # check center of mass has been set + torch.testing.assert_close(cube_object.data.body_com_pose_b.torch, com) + + rand_state = torch.zeros(num_envs, num_cubes, 13, device=device) + rand_state[..., :7] = cube_object.data.default_body_pose.torch + rand_state[..., :3] += cube_object.data.body_link_pos_w.torch + # make quaternion a unit vector + rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) + + env_ids = env_ids.to(device) + object_ids = object_ids.to(device) + for i in range(10): + sim.step() + cube_object.update(sim.cfg.dt) + + if state_location == "com": + if i % 2 == 0: + cube_object.write_body_com_pose_to_sim_index(body_poses=rand_state[..., :7]) + cube_object.write_body_com_velocity_to_sim_index(body_velocities=rand_state[..., 7:]) + else: + cube_object.write_body_com_pose_to_sim_index( + body_poses=rand_state[..., :7], env_ids=env_ids, body_ids=object_ids + ) + cube_object.write_body_com_velocity_to_sim_index( + body_velocities=rand_state[..., 7:], env_ids=env_ids, body_ids=object_ids + ) + elif state_location == "link": + if i % 2 == 0: + cube_object.write_body_link_pose_to_sim_index(body_poses=rand_state[..., :7]) + cube_object.write_body_link_velocity_to_sim_index(body_velocities=rand_state[..., 7:]) + else: + cube_object.write_body_link_pose_to_sim_index( + body_poses=rand_state[..., :7], env_ids=env_ids, body_ids=object_ids + ) + cube_object.write_body_link_velocity_to_sim_index( + body_velocities=rand_state[..., 7:], env_ids=env_ids, body_ids=object_ids + ) + + if state_location == "com": + torch.testing.assert_close(rand_state[..., :7], cube_object.data.body_com_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], cube_object.data.body_com_vel_w.torch) + elif state_location == "link": + torch.testing.assert_close(rand_state[..., :7], cube_object.data.body_link_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], cube_object.data.body_link_vel_w.torch) + + +@pytest.mark.parametrize("num_envs", [1, 3]) +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_reset_object_collection(num_envs, num_cubes, device): + """Test resetting the state of the rigid object.""" + with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: + object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) + sim.reset() + + for i in range(5): + sim.step() + object_collection.update(sim.cfg.dt) + + # Move the object to a random position + body_pose = object_collection.data.default_body_pose.torch.clone() + body_pose[..., :3] = torch.randn(num_envs, num_cubes, 3, device=sim.device) + # Random orientation + body_pose[..., 3:7] = random_orientation(num=num_cubes, device=sim.device) + object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) + body_vel = object_collection.data.default_body_vel.torch.clone() + object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) + + if i % 2 == 0: + object_collection.reset() + + # Reset should zero external forces and torques + assert not object_collection._instantaneous_wrench_composer.active + assert not object_collection._permanent_wrench_composer.active + assert torch.count_nonzero(object_collection._instantaneous_wrench_composer.composed_force.torch) == 0 + assert torch.count_nonzero(object_collection._instantaneous_wrench_composer.composed_torque.torch) == 0 + assert torch.count_nonzero(object_collection._permanent_wrench_composer.composed_force.torch) == 0 + assert torch.count_nonzero(object_collection._permanent_wrench_composer.composed_torque.torch) == 0 + + +@pytest.mark.xfail(reason=_MATERIAL_GAP_REASON, strict=False) +@pytest.mark.parametrize("num_envs", [1, 3]) +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_set_material_properties(num_envs, num_cubes, device): + """Test getting and setting material properties of rigid object.""" + raise NotImplementedError(_MATERIAL_GAP_REASON) + + +@pytest.mark.parametrize("num_envs", [1, 3]) +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("gravity_enabled", [True, False]) +@pytest.mark.isaacsim_ci +def test_gravity_vec_w(num_envs, num_cubes, device, gravity_enabled): + """Test that gravity vector direction is set correctly for the rigid object.""" + with _ovphysx_sim_context(device=device, gravity_enabled=gravity_enabled, auto_add_lighting=True) as sim: + object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) + + # Obtain gravity direction + gravity_dir = (0.0, 0.0, -1.0) if gravity_enabled else (0.0, 0.0, 0.0) + + sim.reset() + + # Check if gravity vector is set correctly + gravity_vec = object_collection.data.GRAVITY_VEC_W.torch + assert gravity_vec[0, 0, 0] == gravity_dir[0] + assert gravity_vec[0, 0, 1] == gravity_dir[1] + assert gravity_vec[0, 0, 2] == gravity_dir[2] + + # Perform simulation + for _ in range(2): + sim.step() + object_collection.update(sim.cfg.dt) + + # Expected gravity value is the acceleration of the body + gravity = torch.zeros(num_envs, num_cubes, 6, device=device) + if gravity_enabled: + gravity[..., 2] = -9.81 + + # Check the body accelerations are correct + torch.testing.assert_close(object_collection.data.body_com_acc_w.torch, gravity) + + +@pytest.mark.parametrize("num_envs", [1, 3]) +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("with_offset", [True]) +@pytest.mark.parametrize("state_location", ["com", "link", "root"]) +@pytest.mark.parametrize("gravity_enabled", [False]) +@pytest.mark.isaacsim_ci +def test_write_object_state_functions_data_consistency( + num_envs, num_cubes, device, with_offset, state_location, gravity_enabled +): + """Test the setters for object_state using both the link frame and center of mass as reference frame.""" + with _ovphysx_sim_context(device=device, gravity_enabled=gravity_enabled, auto_add_lighting=True) as sim: + # Create a scene with random cubes + cube_object, env_pos = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, height=0.0, device=device) + env_ids = torch.tensor([x for x in range(num_envs)], dtype=torch.int32) + object_ids = torch.tensor([x for x in range(num_cubes)], dtype=torch.int32) + + sim.reset() + + # Check if cube_object is initialized + assert cube_object.is_initialized + + # change center of mass offset from link frame + offset = ( + torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) + if with_offset + else torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) + ) + + com = cube_object.data.body_com_pose_b.torch.clone() # shape (num_envs, num_cubes, 7) + com[..., :3] = offset.to(com.device) + cube_object.set_coms_index( + coms=wp.from_torch(com.contiguous(), dtype=wp.transformf), + env_ids=wp.from_torch(env_ids, dtype=wp.int32), + ) + + # check center of mass has been set + torch.testing.assert_close(cube_object.data.body_com_pose_b.torch, com) + + rand_state = torch.rand(num_envs, num_cubes, 13, device=device) + rand_state[..., :3] += cube_object.data.body_link_pos_w.torch + # make quaternion a unit vector + rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) + + env_ids = env_ids.to(device) + object_ids = object_ids.to(device) + sim.step() + cube_object.update(sim.cfg.dt) + + body_link_pose_w = cube_object.data.body_link_pose_w.torch + body_com_pose_w = cube_object.data.body_com_pose_w.torch + object_link_to_com_pos, object_link_to_com_quat = subtract_frame_transforms( + body_link_pose_w[..., :3].view(-1, 3), + body_link_pose_w[..., 3:7].view(-1, 4), + body_com_pose_w[..., :3].view(-1, 3), + body_com_pose_w[..., 3:7].view(-1, 4), + ) + + if state_location == "com": + cube_object.write_body_com_pose_to_sim_index( + body_poses=rand_state[..., :7], env_ids=env_ids, body_ids=object_ids + ) + cube_object.write_body_com_velocity_to_sim_index( + body_velocities=rand_state[..., 7:], env_ids=env_ids, body_ids=object_ids + ) + elif state_location == "link": + cube_object.write_body_link_pose_to_sim_index( + body_poses=rand_state[..., :7], env_ids=env_ids, body_ids=object_ids + ) + cube_object.write_body_link_velocity_to_sim_index( + body_velocities=rand_state[..., 7:], env_ids=env_ids, body_ids=object_ids + ) + elif state_location == "root": + cube_object.write_body_link_pose_to_sim_index( + body_poses=rand_state[..., :7], env_ids=env_ids, body_ids=object_ids + ) + cube_object.write_body_com_velocity_to_sim_index( + body_velocities=rand_state[..., 7:], env_ids=env_ids, body_ids=object_ids + ) + + if state_location == "com": + com_pose_w = cube_object.data.body_com_pose_w.torch + com_vel_w = cube_object.data.body_com_vel_w.torch + expected_root_link_pos, expected_root_link_quat = combine_frame_transforms( + com_pose_w[..., :3].view(-1, 3), + com_pose_w[..., 3:].view(-1, 4), + quat_rotate(quat_inv(object_link_to_com_quat), -object_link_to_com_pos), + quat_inv(object_link_to_com_quat), + ) + expected_object_link_pose = torch.cat((expected_root_link_pos, expected_root_link_quat), dim=1).view( + num_envs, -1, 7 + ) + link_pose_w = cube_object.data.body_link_pose_w.torch + link_vel_w = cube_object.data.body_link_vel_w.torch + # test both root_pose and root_link successfully updated when root_com updates + torch.testing.assert_close(expected_object_link_pose, link_pose_w) + # skip lin_vel because it differs from link frame, this should be fine because we are only checking + # if velocity update is triggered, which can be determined by comparing angular velocity + torch.testing.assert_close(com_vel_w[..., 3:], link_vel_w[..., 3:]) + torch.testing.assert_close(expected_object_link_pose, link_pose_w) + torch.testing.assert_close(com_vel_w[..., 3:], cube_object.data.body_com_vel_w.torch[..., 3:]) + elif state_location == "link": + link_pose_w = cube_object.data.body_link_pose_w.torch + link_vel_w = cube_object.data.body_link_vel_w.torch + expected_com_pos, expected_com_quat = combine_frame_transforms( + link_pose_w[..., :3].view(-1, 3), + link_pose_w[..., 3:].view(-1, 4), + object_link_to_com_pos, + object_link_to_com_quat, + ) + expected_object_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1).view(num_envs, -1, 7) + com_pose_w = cube_object.data.body_com_pose_w.torch + com_vel_w = cube_object.data.body_com_vel_w.torch + # test both root_pose and root_com successfully updated when root_link updates + torch.testing.assert_close(expected_object_com_pose, com_pose_w) + # skip lin_vel because it differs from link frame, this should be fine because we are only checking + # if velocity update is triggered, which can be determined by comparing angular velocity + torch.testing.assert_close(link_vel_w[..., 3:], com_vel_w[..., 3:]) + torch.testing.assert_close(link_pose_w, cube_object.data.body_link_pose_w.torch) + torch.testing.assert_close(link_vel_w[..., 3:], cube_object.data.body_com_vel_w.torch[..., 3:]) + elif state_location == "root": + body_link_pose_w = cube_object.data.body_link_pose_w.torch + body_com_vel_w = cube_object.data.body_com_vel_w.torch + expected_object_com_pos, expected_object_com_quat = combine_frame_transforms( + body_link_pose_w[..., :3].view(-1, 3), + body_link_pose_w[..., 3:].view(-1, 4), + object_link_to_com_pos, + object_link_to_com_quat, + ) + expected_object_com_pose = torch.cat((expected_object_com_pos, expected_object_com_quat), dim=1).view( + num_envs, -1, 7 + ) + com_pose_w = cube_object.data.body_com_pose_w.torch + com_vel_w = cube_object.data.body_com_vel_w.torch + link_pose_w = cube_object.data.body_link_pose_w.torch + link_vel_w = cube_object.data.body_link_vel_w.torch + # test both root_com and root_link successfully updated when root_pose updates + torch.testing.assert_close(expected_object_com_pose, com_pose_w) + torch.testing.assert_close(body_com_vel_w, com_vel_w) + torch.testing.assert_close(body_link_pose_w, link_pose_w) + torch.testing.assert_close(body_com_vel_w[..., 3:], link_vel_w[..., 3:]) diff --git a/source/isaaclab_ovphysx/test/assets/test_rigid_object_helpers.py b/source/isaaclab_ovphysx/test/assets/test_rigid_object_helpers.py new file mode 100644 index 000000000000..e57d8d2651d7 --- /dev/null +++ b/source/isaaclab_ovphysx/test/assets/test_rigid_object_helpers.py @@ -0,0 +1,45 @@ +# 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 + +"""OVPhysX-only unit tests for rigid-object helpers. + +These tests cover OVPhysX-specific scaffolding (mock binding-set shape +contracts for ``asset_kind="rigid_object"``) that has no PhysX equivalent +and therefore does not appear in the PhysX-mirrored ``test_rigid_object.py``. +""" + +from __future__ import annotations + +import pytest +import warp as wp + +# The CI isaaclab_ov* pattern unintentionally collects isaaclab_ovphysx tests, +# but the ovphysx wheel is not installed in that environment. Skip gracefully +# so the isaaclab_ov CI pipeline is not blocked by an unrelated dependency. +pytest.importorskip("ovphysx.types", reason="ovphysx wheel not installed") + +from isaaclab_ovphysx import tensor_types as TT # noqa: E402 +from isaaclab_ovphysx.test.mock_interfaces.views import MockOvPhysxBindingSet # noqa: E402 + +wp.init() + + +def test_mock_binding_set_rigid_object_shapes(): + pytest.importorskip("isaaclab_ovphysx.tensor_types").RIGID_BODY_POSE # gates on wheel + + bindings = MockOvPhysxBindingSet( + num_instances=4, + num_joints=0, + num_bodies=1, + asset_kind="rigid_object", + ) + assert bindings.bindings[TT.RIGID_BODY_POSE].shape == (4, 7) + assert bindings.bindings[TT.RIGID_BODY_VELOCITY].shape == (4, 6) + assert bindings.bindings[TT.RIGID_BODY_WRENCH].shape == (4, 9) + assert bindings.bindings[TT.RIGID_BODY_MASS].shape == (4,) + assert bindings.bindings[TT.RIGID_BODY_INERTIA].shape == (4, 9) + # Articulation-only bindings must be absent + assert TT.DOF_POSITION not in bindings.bindings + assert TT.LINK_WRENCH not in bindings.bindings diff --git a/source/isaaclab_tasks/changelog.d/antoiner-feat-ovphysx_articulation.rst b/source/isaaclab_tasks/changelog.d/antoiner-feat-ovphysx_articulation.rst new file mode 100644 index 000000000000..26c7d5b8f9ca --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/antoiner-feat-ovphysx_articulation.rst @@ -0,0 +1,10 @@ +Added +^^^^^ + +* Added the ``ovphysx`` preset to ``Isaac-Repose-Cube-Allegro-Direct-v0`` + (``ObjectCfg`` and ``PhysicsCfg`` in + :mod:`isaaclab_tasks.direct.allegro_hand.allegro_hand_env_cfg`), so the + task can be selected with ``presets=ovphysx`` against the OVPhysX + backend. Exercises the OVPhysX :class:`~isaaclab_ovphysx.assets.Articulation` + (Allegro hand) and :class:`~isaaclab_ovphysx.assets.RigidObject` (cube) + in the same scene. diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/allegro_hand/allegro_hand_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/direct/allegro_hand/allegro_hand_env_cfg.py index bd982c5d7105..6065f81a0868 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/allegro_hand/allegro_hand_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/allegro_hand/allegro_hand_env_cfg.py @@ -5,6 +5,7 @@ from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +from isaaclab_ovphysx.physics import OvPhysxCfg from isaaclab_physx.physics import PhysxCfg import isaaclab.sim as sim_utils @@ -56,6 +57,25 @@ class ObjectCfg(PresetCfg): actuators={}, articulation_root_prim_path="", ) + ovphysx = RigidObjectCfg( + prim_path="/World/envs/env_.*/object", + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", + rigid_props=sim_utils.RigidBodyPropertiesCfg( + kinematic_enabled=False, + disable_gravity=False, + enable_gyroscopic_forces=True, + solver_position_iteration_count=8, + solver_velocity_iteration_count=0, + sleep_threshold=0.005, + stabilization_threshold=0.0025, + max_depenetration_velocity=1000.0, + ), + mass_props=sim_utils.MassPropertiesCfg(density=400.0), + scale=(1.2, 1.2, 1.2), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, -0.17, 0.56), rot=(0.0, 0.0, 0.0, 1.0)), + ) default = physx @@ -79,6 +99,7 @@ class PhysicsCfg(PresetCfg): num_substeps=2, debug_mode=False, ) + ovphysx = OvPhysxCfg() default = physx