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..498091f51058 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()) 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_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..cb459221c681 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py @@ -22,10 +22,11 @@ from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_ovphysx import tensor_types as TT +from isaaclab_ovphysx.assets.kernels import _body_wrench_to_world, _scatter_rows_partial 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 +from .kernels import update_soft_joint_pos_limits if TYPE_CHECKING: from isaaclab.actuators import ActuatorBase 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..5de3057e56cd 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation_data.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation_data.py @@ -18,18 +18,17 @@ 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, ) +from .kernels import _compose_body_com_poses, _fd_joint_acc + class ArticulationData(BaseArticulationData): """Data container for an articulation backed by ovphysx tensor bindings. diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/kernels.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/kernels.py index b93c4e6d4b41..cc9faf15753a 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/kernels.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/kernels.py @@ -8,41 +8,6 @@ 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, @@ -93,38 +58,6 @@ def _fd_joint_acc( prev_vel[i, j] = cur_vel[i, j] -@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 velocity to the root velocity buffer. - - 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] - - -@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,). - """ - i = wp.tid() - com_pose_w[i] = wp.transform_multiply(link_pose[i], com_pose_b[i, 0]) - - @wp.kernel def _compose_body_com_poses( link_pose: wp.array(dtype=wp.transformf, ndim=2), @@ -140,78 +73,3 @@ def _compose_body_com_poses( """ i, j = wp.tid() com_pose_w[i, j] = wp.transform_multiply(link_pose[i, j], com_pose_b[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), -): - """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]) - - -@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 yaw heading angle from the forward direction rotated into the world frame. - - 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,). - """ - 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 world-frame linear velocity into the root body frame. - - 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,). - """ - 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 world-frame angular velocity into the root body frame. - - 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,). - """ - 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) 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..cf49c8362636 --- /dev/null +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/kernels.py @@ -0,0 +1,1108 @@ +# 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] 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..41afe07cf09c 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,36 @@ # 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, - } +_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, + 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, + # 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_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