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_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.minor.rst b/source/isaaclab_ovphysx/changelog.d/antoiner-feat-ovphysx_articulation.minor.rst new file mode 100644 index 000000000000..24913760aa9b --- /dev/null +++ b/source/isaaclab_ovphysx/changelog.d/antoiner-feat-ovphysx_articulation.minor.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.minor.rst b/source/isaaclab_ovphysx/changelog.d/antoiner-feat-ovphysx_rigidobject.minor.rst new file mode 100644 index 000000000000..2a43911295c2 --- /dev/null +++ b/source/isaaclab_ovphysx/changelog.d/antoiner-feat-ovphysx_rigidobject.minor.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/isaaclab_ovphysx/assets/__init__.pyi b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/__init__.pyi index 516e15c5ef6a..52d1a435596a 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/__init__.pyi +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/__init__.pyi @@ -6,6 +6,9 @@ __all__ = [ "Articulation", "ArticulationData", + "RigidObject", + "RigidObjectData", ] from .articulation import Articulation, ArticulationData +from .rigid_object import RigidObject, RigidObjectData diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py index f3919d56da73..89c47039686e 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py @@ -3,50 +3,83 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Articulation implementation backed by ovphysx TensorBindingsAPI.""" +"""OVPhysX-backed Articulation implementation. + +Mirrors the post-refactor :class:`~isaaclab_ovphysx.assets.RigidObject` +shape with the API surface coming from Newton. Eager binding creation +in :meth:`_initialize_impl`; dual mask+index API; pinned-host CPU +staging via the data class's ``_binding_read`` / ``_binding_write`` +helpers (the PR #5329 pattern). +""" 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, +) logger = logging.getLogger(__name__) class Articulation(BaseArticulation): - """Articulation backed by the ovphysx TensorBindingsAPI. + """OVPhysX-backed Articulation asset class. - Reads and writes simulation state through ovphysx.TensorBinding objects created - from the OvPhysxManager's PhysX instance. - """ + Mirrors the OVPhysX RigidObject shape with the public API surface + coming from Newton. Joint, body, tendon, and root state are read + via :class:`ovphysx.TensorBinding` objects acquired from the + :class:`~isaaclab_ovphysx.physics.OvPhysxManager`. Writes go through + the dual mask+index API; CPU-only bindings (mass, COM, inertia, + joint properties, tendon properties) route through pinned-host + staging buffers (PR #5329 pattern). - __backend_name__ = "ovphysx" + .. _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 created lazily (on first access) to avoid allocating + # handles for tensor types the user never queries. + self._bindings: dict[int, Any] = {} """ Properties @@ -63,29 +96,29 @@ def num_instances(self) -> int: return self._num_instances @property - def is_fixed_base(self) -> bool: - """Whether the articulation is a fixed-base or floating-base system.""" - return self._is_fixed_base + def num_bodies(self) -> int: + """Number of bodies (links) per articulation.""" + return self._num_bodies @property def num_joints(self) -> int: - """Number of joints in the articulation.""" + """Number of degrees of freedom per 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 per 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 per articulation.""" + return self._num_spatial_tendons @property - def num_bodies(self) -> int: - """Number of bodies (links) in the articulation.""" - return self._num_bodies + def body_names(self) -> list[str]: + """Ordered names of bodies in the articulation.""" + return self._body_names @property def joint_names(self) -> list[str]: @@ -95,67 +128,121 @@ def joint_names(self) -> list[str]: @property def fixed_tendon_names(self) -> list[str]: """Ordered names of fixed tendons in the articulation.""" - return getattr(self, "_fixed_tendon_names", []) + 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", []) + return self._spatial_tendon_names @property - def body_names(self) -> list[str]: - """Ordered names of bodies (links) in the articulation.""" - return self._body_names + def is_fixed_base(self) -> bool: + """Whether the articulation is a fixed-base or floating-base system.""" + return self._is_fixed_base @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:: - If both `env_ids` and `env_mask` are provided, then `env_mask` takes precedence over `env_ids`. - 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,). + 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() + """Apply external wrenches, run the actuator model, and write commands to the simulation. + + Execution order mirrors the PhysX backend: + + 1. Compose and write external wrenches (body-frame -> world-frame via + :func:`_body_wrench_to_world`, then write to ``LINK_WRENCH``). + 2. Run the actuator model (:meth:`_apply_actuator_model`) to compute + joint efforts from user-supplied targets. + 3. Write the effort tensor to the ``DOF_ACTUATION_FORCE`` binding. + 4. For implicit actuators, write position and velocity targets in a + single shot to ``DOF_POSITION_TARGET`` / ``DOF_VELOCITY_TARGET``. + """ + # 1. External wrenches. + 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 + + poses = self._data.body_link_pose_w.warp # (N, num_bodies) 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: + binding.write(self._wrench_buf) + inst.reset() + # 2. Actuator model. self._apply_actuator_model() - # Write effort tensor to simulation. + + # 3. Write effort (applies even when no actuators: zeros are safe). 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). + + # 4. Write position / velocity targets 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 +250,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. + """Update the simulation data. Args: - dt: The simulation time step [s] used for finite-difference quantities. + dt: Time-step in seconds [s]. """ self._data.update(dt) @@ -197,8 +284,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. @@ -221,20 +308,21 @@ 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: - tendon_subsets = self.fixed_tendon_names + tendon_subsets = self._fixed_tendon_names return resolve_matching_names(name_keys, tendon_subsets, preserve_order) def find_spatial_tendons( @@ -245,19 +333,20 @@ 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 + tendon_subsets = self._spatial_tendon_names return resolve_matching_names(name_keys, tendon_subsets, preserve_order) """ @@ -268,19 +357,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 +383,62 @@ 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 so the next reads re-fetch. Mirrors + # PhysX (articulation.py:457-464): 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 wheel 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 +446,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, + ) + # See write_root_link_pose_to_sim_index for why these are invalidated. + 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, + ) + # Mirrors PhysX (articulation.py:548-556): writing root COM pose updates + # the inferred root link pose, which in turn invalidates the body 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 +523,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, + ) + # See write_root_com_pose_to_sim_index for why these are invalidated. + 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 +588,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 +649,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 +694,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 +730,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. Mirrors PhysX (articulation.py:894-903). + 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, + ) + # See write_joint_position_to_sim_index for why these are invalidated. + 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 +987,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 +1268,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. + + This is a CPU-only write routed through pinned-host staging because + ``DOF_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: - 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. + 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. """ - 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) + 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: 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 max velocity over selected environment indices into the simulation. + """Set joint velocity limits over selected env / joint indices 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 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 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 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). """ - 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) + 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: 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 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 wheel 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") + # See _index variant for why we refresh the (N, J, 3) buffer from the wheel. + 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 wheel 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") + # See _index variant for why we refresh the (N, J, 3) buffer from the wheel. + 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 wheel 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") + # See _index variant for why we refresh the (N, J, 3) buffer from the wheel. + 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 +1971,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 +2015,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 +2101,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 +2187,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 +2271,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 +2348,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 +2423,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,22 +2464,40 @@ 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,19 +2506,43 @@ 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 +2550,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 +2592,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 +2636,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 +2678,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 +2722,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 wheel. + 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 +2770,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 +2818,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 +2860,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 +2904,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 +2946,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 +2989,45 @@ 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's equivalent calls ``root_view.set_fixed_tendon_properties`` with + all six property buffers at once (articulation.py:3078-3086). OVPhysX + does not expose a single wheel-side batch setter; instead, we write + each ``FIXED_TENDON_*`` binding individually from the matching + ``self._data._fixed_tendon_*`` buffer. + + .. note:: + Only env indices apply to the wheel write; ``fixed_tendon_ids`` is + accepted for API parity with PhysX but is unused (the wheel 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 +3035,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 +3070,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 (PhysX + applies tendon properties without a CPU clone), so the write goes + directly from the sim-device buffer to the wheel 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 +3113,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 +3157,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 +3197,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 +3241,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 wheel. + + .. 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 +3285,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 +3329,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 wheel. + + .. 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 +3371,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 +3414,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,80 +3443,48 @@ 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. + + 1. Acquire the OvPhysxManager.PhysX instance and the simulation device. + 2. Validate the USD prim tree (exactly one ArticulationRootAPI under + ``cfg.prim_path``; surface a clear RuntimeError otherwise). + 3. Convert IsaacLab prim-path notation to the glob pattern ovphysx expects, + extending it to the articulation root prim if needed. + 4. Eagerly create the bindings the data container needs and read counts + (num_instances, num_bodies, num_joints) and names from binding metadata. + 5. Discover tendon counts and names via :meth:`_process_tendons`. + 6. Construct the :class:`ArticulationData` container. + 7. Allocate asset-side buffers via :meth:`_create_buffers`. + 8. Apply initial state via :meth:`_process_cfg`. + 9. Prime the data via ``update(0.0)``. + 10. Mark the data as ready. + """ + # Step 1: Acquire PhysX instance and device. 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() prim_path = self.cfg.prim_path - # Convert IsaacLab prim-path notation to the glob patterns ovphysx expects. + # Step 2: 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 @@ -1664,6 +3492,7 @@ def _initialize_impl(self) -> None: pattern = re.sub(r"\{ENV_REGEX_NS\}", "*", prim_path) pattern = re.sub(r"\.\*", "*", pattern) # env_.* -> env_* + # Step 3: Validate the USD prim tree and extend pattern to the articulation root if needed. # 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 @@ -1671,50 +3500,65 @@ def _initialize_impl(self) -> None: # 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 - 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: + # Mirror the PhysX backend (articulation.py:3572): when the user has + # explicitly specified the articulation-root subpath, append it to + # the prim path directly and skip auto-discovery. Validate that the + # resulting expression resolves to at least one prim in the USD + # stage; otherwise surface a RuntimeError before downstream binding + # creation fails opaquely. + 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: + # e.g. first_prim_path=/World/envs/env_0/Robot, root_relative=/torso + # pattern becomes /World/envs/env_*/Robot/torso + 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 + # Step 4: 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 +3576,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,26 +3592,40 @@ 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. + # Step 5: Discover tendon counts and names before buffer allocation. self._process_tendons() + # Step 6: 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, + ) + + # Step 7: Allocate asset-side buffers. self._create_buffers() + # Step 8: Apply initial state from config. self._process_cfg() + + # Step 8b: 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. + # Step 8c: 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 + torque = self._data._applied_torque shape = self._effort_binding.shape self._effort_write_view = wp.array( ptr=torque.ptr, @@ -1771,7 +3637,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: @@ -1780,30 +3645,67 @@ def _make_write_view(tt, buf): return b, v self._pos_target_binding, self._pos_target_write_view = _make_write_view( - TT.DOF_POSITION_TARGET, self._data.joint_pos_target + TT.DOF_POSITION_TARGET, self._data._joint_pos_target ) self._vel_target_binding, self._vel_target_write_view = _make_write_view( - TT.DOF_VELOCITY_TARGET, self._data.joint_vel_target + 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 + # Step 9: Validate the resolved configuration (raises ValueError when + # default joint pos/vel falls outside the wheel-reported limits). + # Mirrors PhysX's ``_validate_cfg`` ordering — after actuator/tendon + # processing, before priming. + self._validate_cfg() - def _create_buffers(self) -> None: - self._data._create_buffers() + # Step 10: 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) + # Step 11: 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,102 +3716,46 @@ 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, - ) + # 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, + ) - def _invalidate_initialize_callback(self, event) -> None: - self._is_initialized = False + def _process_tendons(self) -> None: + """Discover tendon counts from binding metadata and names from USD. - def _process_actuators_cfg(self) -> None: - """Build actuator instances from the config and write drive properties to PhysX. + 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 = [] - 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. + sample = next(iter(self._bindings.values())) + self._num_fixed_tendons = getattr(sample, "fixed_tendon_count", 0) + self._num_spatial_tendons = getattr(sample, "spatial_tendon_count", 0) - These writes happen via TensorBinding (GPU-resident) after warmup has - allocated the GPU buffers (MODEL_INIT fires post-warmup). - """ - 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 = [] - - sample = next(iter(self._bindings.values())) - self._num_fixed_tendons = getattr(sample, "fixed_tendon_count", 0) - self._num_spatial_tendons = getattr(sample, "spatial_tendon_count", 0) - - if self._num_fixed_tendons > 0 or self._num_spatial_tendons > 0: - stage_path = OvPhysxManager._stage_path - if stage_path is not None: - try: - from pxr import Usd, UsdPhysics + if self._num_fixed_tendons > 0 or self._num_spatial_tendons > 0: + stage_path = OvPhysxManager._stage_path + if stage_path is not None: + try: + from pxr import Usd from isaaclab.sim.utils.queries import get_all_matching_child_prims @@ -1946,55 +3792,165 @@ 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). + + 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``: - IsaacLab actuators are torch-based. We convert warp -> torch via - DLPack (zero-copy on GPU), run the actuator, then write results back. + * 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 wheel's already-correct DOF defaults + # (USD-authored ``physxJoint:maxJointVelocity`` etc. that the wheel + # 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 + # wheel's correct values with ``inf``. Mirrors the PhysX backend. + 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 +3961,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 +3992,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 wheel's native ``binding.write(mask=...)`` path, + so the mask is preserved end-to-end -- no ``torch.nonzero`` conversion 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 *env_mask* for a CPU-only binding write. + + ``env_mask`` is normally on ``self._device``; 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. """ - 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 10c7e4b7ecd6..3ada3b36e63b 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,22 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Articulation data backed by ovphysx TensorBindingsAPI.""" +"""OVPhysX-backed ArticulationData implementation. + +Mirrors the post-refactor :class:`~isaaclab_ovphysx.assets.RigidObjectData` +shape with the API surface coming from +:class:`isaaclab_newton.assets.ArticulationData`: eager buffer allocation +in ``_create_buffers``; pinned-host CPU staging for CPU-only bindings via +``_binding_read`` (the PR #5329 pattern); ``ProxyArray`` returns from every +public property; timestamp-cached pull semantics through +``_read_*_binding``. +""" from __future__ import annotations +import logging import warnings +from collections.abc import Callable from typing import Any import numpy as np @@ -18,114 +29,125 @@ from isaaclab.utils.warp import ProxyArray 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 -class ArticulationData(BaseArticulationData): - """Data container for an articulation backed by ovphysx tensor bindings. +logger = logging.getLogger(__name__) - 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 - stored in the simulation world frame unless otherwise specified. - An articulation is comprised of multiple rigid bodies or links. For a rigid body, there are two frames - of reference that are used: +class ArticulationData(BaseArticulationData): + """Data container for an articulation backed by OVPhysX tensor bindings. - - 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. + Reads simulation state on demand through ``ovphysx`` ``TensorBinding`` + objects keyed by :data:`~isaaclab_ovphysx.tensor_types.ARTICULATION_*`. + Buffers are allocated eagerly in :meth:`_create_buffers` and timestamped + so each binding is read at most once per simulation step. - 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. + CPU-only bindings (``BODY_MASS``, ``BODY_COM_POSE``, + ``BODY_INERTIA``, every ``DOF_*`` property except state) route through + pinned-host staging buffers via :meth:`_binding_read`, matching the + PR #5329 PhysX pattern. - Uses ovphysx :class:`TensorBinding` objects to lazily read simulation state into warp - arrays. Writes happen via the :class:`Articulation` class. + Properties return :class:`~isaaclab.utils.warp.ProxyArray` for time-varying + state and raw :class:`wp.array` for one-shot config buffers (defaults). """ __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 (not @property), + # mirroring the post-audit RigidObjectData demotion. + 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 + # Internal 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 + # Simulation timestamp drives the cache validity contract. 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 - - # Initialize parametric gravity and forward vectors (matching PhysX/Newton pattern). - # Guard against None sim context (e.g. mock/test environments). + 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 gravity and forward constants (matching PhysX/Newton pattern). + 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)) 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 + # Allocate every TimestampedBufferWarp and pinned CPU staging buffer. + self._create_buffers() @property def is_primed(self) -> bool: @@ -149,21 +171,31 @@ 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.""" + def update(self, dt: float) -> None: + """Advance the simulation timestamp; trigger FD derivations. - fixed_tendon_names: list[str] = None - """Fixed tendon names in the order parsed by the simulation view.""" + Called by :class:`~isaaclab_ovphysx.assets.Articulation` each step. - spatial_tendon_names: list[str] = None - """Spatial tendon names in the order parsed by the simulation view.""" + Args: + dt: Simulation time-step [s]. + """ + self._sim_timestamp += dt + if not self._is_primed: + return + # FD joint acceleration from velocity. + 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 +203,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, -]. - 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). + 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 +219,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]. + + Shape is (num_instances,), dtype = wp.spatial_vectorf. + In torch this resolves to (num_instances, 6). - 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). + 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 +246,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 +259,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 +270,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 +283,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 +294,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 +311,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 +321,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 +331,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 +345,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 +355,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 +367,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. 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_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 +501,12 @@ def joint_effort_limits(self) -> ProxyArray: @property def soft_joint_pos_limits(self) -> ProxyArray: - r"""Soft joint position limits for all joints. - - 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:: + r"""Soft joint position limits for all joints [m or rad, depending on joint type]. - soft\_lower = (lower + upper) / 2 - factor * (upper - lower) / 2 - soft\_upper = (lower + upper) / 2 + factor * (upper - lower) / 2 + Shape is (num_instances, num_joints), dtype = wp.vec2f. + In torch this resolves to (num_instances, num_joints, 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 +514,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 +526,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 +538,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``. - 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_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 +624,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``. - 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_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. 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_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 +680,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 +693,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 +718,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 +730,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 +741,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 +753,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,33 +772,39 @@ 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``. - Shape is (num_instances, num_bodies), dtype = wp.float32. In torch this resolves to - (num_instances, num_bodies). + 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. @@ -733,61 +815,81 @@ def body_link_pose_w(self) -> ProxyArray: 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, -]. - This quantity is the pose of the center of mass frame of the articulation links relative to the world. + Shape is (num_instances, num_bodies), dtype = wp.transformf. + In torch this resolves to (num_instances, num_bodies, 7). + + 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 +900,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 +913,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 +938,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 +949,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 +960,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 +976,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 +994,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 +1005,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 +1018,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 +1026,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 +1039,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 +1047,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 +1060,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 +1069,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 +1082,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 +1091,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 +1107,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 +1119,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 +1129,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 +1141,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 +1151,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 +1163,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 +1174,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 +1185,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 +1195,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 +1208,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 +1220,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 +1232,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 +1243,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 +1256,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 +1268,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 +1280,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 +1292,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 +1304,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 +1315,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 +1329,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 +1339,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 @@ -1350,20 +1370,26 @@ def _create_buffers(self) -> None: # noqa: C901 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 +1415,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 +1598,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 +1606,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 +1714,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,6 +1744,7 @@ 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 # Body properties @@ -1606,13 +1788,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 +1867,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 +1876,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,18 +1913,13 @@ 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. """ - binding = self._get_binding(tensor_type) - if binding is None: - return - binding.read(wp_array) - - 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) @@ -1700,17 +1929,34 @@ def _read_binding_into_buf(self, tensor_type: int, buf: TimestampedBuffer) -> No 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._get_binding(tensor_type).read(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) @@ -1719,11 +1965,37 @@ def _read_spatial_vector_binding(self, tensor_type: int, buf: TimestampedBuffer) 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, @@ -1733,6 +2005,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, @@ -1742,6 +2022,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, @@ -1751,6 +2039,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, @@ -1758,3 +2054,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..1f11a6d74789 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/kernels.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/kernels.py @@ -3,52 +3,46 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Warp kernels for the ovphysx articulation.""" +"""Warp kernels for the OVPhysX Articulation backend. + +Mirrors the structure of :mod:`isaaclab_physx.assets.articulation.kernels` for +the kernels OVPhysX exposes today. The following PhysX kernels are intentionally +absent because the corresponding write paths are not yet plumbed through the +OVPhysX Python API: + +* ``write_joint_vel_data`` +* ``write_joint_limit_data_to_buffer`` +* ``float_data_to_buffer_with_indices`` +* ``update_default_joint_values`` +* ``update_targets`` +* ``update_actuator_state_model`` +* ``extract_friction_properties`` + +OVPhysX-only data-layer kernels (``_fd_joint_acc``, ``_compose_body_com_poses``) +support the timestamped-buffer / pull-on-demand model used by +:class:`isaaclab_ovphysx.assets.ArticulationData`; they have no PhysX equivalent. +""" 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] - - @wp.func 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( @@ -63,155 +57,183 @@ def update_soft_joint_pos_limits( 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) + """Update soft joint position limits from hard limits and a soft limit factor. - -""" -Data-layer kernels (used by ArticulationData). -""" - - -@wp.kernel -def _fd_joint_acc( - cur_vel: wp.array2d(dtype=wp.float32), - prev_vel: wp.array2d(dtype=wp.float32), - inv_dt: float, - out: wp.array2d(dtype=wp.float32), -): - """Compute joint acceleration via finite differencing and update previous velocity. + 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: - 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). + 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, j = wp.tid() - out[i, j] = (cur_vel[i, j] - prev_vel[i, j]) * inv_dt - prev_vel[i, j] = cur_vel[i, j] + soft_joint_pos_limits[i, j] = compute_soft_joint_pos_limits_func(joint_pos_limits[i, j], soft_limit_factor) @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 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), ): - """Copy the first body's velocity to the root velocity buffer. + """Clamp default joint positions to new limits and refresh soft limits over (env_ids × joint_ids). - Args: - body_vel: Body velocities. Shape is (num_envs, num_bodies). - root_vel: Output root velocities. Shape is (num_envs,). - """ - i = wp.tid() - root_vel[i] = body_vel[i, 0] + 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 _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 body-frame CoM offset to get world-frame root CoM pose. - - 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,). + * Clamps :attr:`default_joint_pos` ``[env_ids[i], joint_ids[j]]`` if it falls outside + the new limits, atomically incrementing :attr:`clamped_count`. + * Recomputes :attr:`soft_joint_pos_limits` ``[env_ids[i], joint_ids[j]]`` from the new + hard limits and :attr:`soft_limit_factor`. """ - i = wp.tid() - com_pose_w[i] = wp.transform_multiply(link_pose[i], com_pose_b[i, 0]) + 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 _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 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), ): - """Compose body link poses with body-frame CoM offsets to get world-frame CoM poses. + """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: - 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). + 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``. + in_viscous: Viscous friction coefficients, or ``None``. + env_ids: Environment indices to write. + joint_ids: Joint indices to write. + out_buffer: The combined ``(num_envs, num_joints, 3)`` friction buffer. """ i, j = wp.tid() - com_pose_w[i, j] = wp.transform_multiply(link_pose[i, j], com_pose_b[i, j]) + 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 _projected_gravity( - gravity_vec_w: wp.array(dtype=wp.vec3f), - root_pose: wp.array(dtype=wp.transformf), - 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), ): - """Project world-frame gravity direction into the root body frame. - - 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,). - """ - i = wp.tid() - q = wp.transform_get_rotation(root_pose[i]) - out[i] = wp.quat_rotate_inv(q, gravity_vec_w[i]) + """Mask variant of :func:`write_joint_friction_data_to_buffer_index`.""" + 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] @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`. - 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,). + 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``. """ - 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) + + +""" +Data-layer kernels (used by ArticulationData). +""" @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 _fd_joint_acc( + cur_vel: wp.array2d(dtype=wp.float32), + prev_vel: wp.array2d(dtype=wp.float32), + inv_dt: float, + out: wp.array2d(dtype=wp.float32), ): - """Rotate world-frame linear velocity into the root body frame. + """Compute joint acceleration via finite differencing and update 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: - 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,). + 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 = 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() + out[i, j] = (cur_vel[i, j] - prev_vel[i, j]) * inv_dt + prev_vel[i, j] = cur_vel[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 _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), ): - """Rotate world-frame angular velocity into the root body frame. + """Compose body link poses with body-frame CoM offsets to get world-frame CoM poses. 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,). + 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 = 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() + com_pose_w[i, j] = wp.transform_multiply(link_pose[i, j], com_pose_b[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..9d1b6ca1b8aa --- /dev/null +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/kernels.py @@ -0,0 +1,1376 @@ +# 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 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..7d8f18c9ef8c --- /dev/null +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/rigid_object.py @@ -0,0 +1,1164 @@ +# 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 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) + ) + + # ------------------------------------------------------------------ + # Deprecated writers + # ------------------------------------------------------------------ + + 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) + + # ------------------------------------------------------------------ + # Internal helpers -- Resolve + # ------------------------------------------------------------------ + + def _resolve_env_ids(self, env_ids) -> wp.array: + """Resolve environment indices to a warp int32 array on ``self._device`` (mirrors PhysX). + + 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_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) + return body_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``. + + OVPhysX (like Newton) uses the wheel's native ``binding.write(mask=...)`` path, + so the mask is preserved end-to-end -- no ``torch.nonzero`` conversion needed. + ``None`` returns the pre-allocated all-true mask. + """ + if env_mask is None: + return self._ALL_TRUE_ENV_MASK + if isinstance(env_mask, torch.Tensor): + return wp.from_torch(env_mask.to(torch.bool), dtype=wp.bool) + if isinstance(env_mask, wp.array) and str(env_mask.device) != self._device: + env_mask = wp.clone(env_mask, device=self._device) + return env_mask + + 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 _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. + + ``env_mask`` is normally on ``self._device``; 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`` (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") + + # ------------------------------------------------------------------ + # Internal helpers -- Lifecycle + # ------------------------------------------------------------------ + + def _initialize_impl(self) -> None: + """Initialize the rigid object from the OVPhysX simulation backend. + + Creates tensor bindings for the rigid-body state tensors, reads counts + and body names, creates the :class:`RigidObjectData` container, and + primes the data buffers. + """ + # Step 1-3: Acquire PhysX instance and build binding pattern. + 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() + self._binding_pattern = self.cfg.prim_path + + # 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. + # Mirror PhysX's prim-scan validation so failures surface here with a + # clear message. + 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 + 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." + ) + + # Step 4: 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 + + # Step 5: 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"] + + # Step 6: Create the data container (mirrors PhysX: takes bindings + device). + self._data = RigidObjectData(self._bindings, self._device, check_shapes=self._check_shapes) + + # Allocate asset-side buffers and apply the initial state from the configuration. + self._create_buffers() + self._process_cfg() + + # Step 9: Prime the data by performing the first read. + self.update(0.0) + + # Step 10: Mark data as ready. + self._data.is_primed = True + + def _create_buffers(self) -> None: + """Create buffers for storing data (mirrors PhysX).""" + 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 (mirrors Newton). 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 (mirrors PhysX). + self._data.body_names = self._body_names + + # Pre-allocated pinned CPU buffers for OVPhysX TensorBinding writes (PR #5329 pattern). + # 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 (mirrors PhysX).""" + # 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 _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``. + + Raises: + Whatever the wheel raises if ``create_tensor_binding`` fails. + :meth:`_initialize_impl` eagerly creates every binding the writers + consume, so post-init calls hit the cache and cannot fail. + """ + 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 + + def _invalidate_initialize_callback(self, event) -> None: + """Invalidates the scene elements.""" + super()._invalidate_initialize_callback(event) 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..cf986689d2c1 --- /dev/null +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/rigid_object_data.py @@ -0,0 +1,1197 @@ +# 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 bindings (equivalent to the view in PhysX) + 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 + # Mirrors Newton's update() pattern (rigid_object_data.py line 126). + self.body_com_acc_w + + """ + Names. + """ + + body_names: list[str] = None + """Body names in the order parsed by the simulation view.""" + + @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: + """Eagerly allocate every per-instance TimestampedBuffer and the slots for + cached :class:`ProxyArray` wrappers. Mirrors + :meth:`isaaclab_physx.assets.RigidObjectData._create_buffers`. + """ + 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/physics/ovphysx_manager.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/physics/ovphysx_manager.py index 9063078e45b3..1b4e8cd6e25a 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/physics/ovphysx_manager.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/physics/ovphysx_manager.py @@ -46,6 +46,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(). @@ -53,6 +59,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 @@ -79,13 +90,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 = [] @@ -138,15 +156,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: @@ -159,7 +189,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.") @@ -173,9 +218,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_") @@ -184,6 +236,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 @@ -253,52 +365,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 @@ -309,22 +394,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/test/assets/test_articulation.py b/source/isaaclab_ovphysx/test/assets/test_articulation.py index 52998a2ec5f6..ad17be4eb7ab 100644 --- a/source/isaaclab_ovphysx/test/assets/test_articulation.py +++ b/source/isaaclab_ovphysx/test/assets/test_articulation.py @@ -3,114 +3,2444 @@ # # 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." +) + +_FK_ON_DEMAND_GAP_REASON = ( + "GPU-only: PhysX's data-class getters call " + "``SimulationView.update_articulations_kinematic`` before reading link " + "transforms (see ``isaaclab_physx.assets.articulation_data:735``), so body " + "poses reflect new joint positions immediately after " + "``write_joint_position_to_sim_*`` without a sim step. The OVPhysX wheel's " + "``ovphysx.PhysX`` class does not expose an equivalent FK-on-demand API " + "(omni.physics.tensors has it; the OVPhysX wrapper does not surface it), " + "so body-pose bindings remain stale on a GPU sim until the next ``step``. " + "On a CPU sim the same write path happens to update the bindings synchronously " + "(no async stream involved), so the test xpasses there — hence ``strict=False``. " + "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_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_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_effort + 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 +@pytest.mark.xfail(reason=_FK_ON_DEMAND_GAP_REASON, strict=False) +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..65943b53fed8 --- /dev/null +++ b/source/isaaclab_ovphysx/test/assets/test_rigid_object.py @@ -0,0 +1,1128 @@ +# 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 +from isaaclab_ovphysx.assets import RigidObject +from isaaclab_ovphysx.physics import OvPhysxCfg, OvPhysxManager + +import isaaclab.sim as sim_utils +from isaaclab.assets import RigidObjectCfg +from isaaclab.sim import SimulationCfg, build_simulation_context +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR +from isaaclab.utils.math import ( + 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_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