From 7a12614298936f7a35a6261186640c54a4edb329 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Tue, 5 May 2026 16:50:06 +0200 Subject: [PATCH] [OVPHYSX] Articulation rewrite (data class + asset class + kernels) Add Articulation and ArticulationData for the OVPhysX backend, mirroring the PhysX/Newton public API. Resolves PR #5459. Articulation ^^^^^^^^^^^^ * index/mask split for every state writer, simulation-parameter writer, setter, and tendon setter; OVPhysX exposes both _index and _mask as first-class paths and intentionally drops the PhysX-specific ``full_data`` kwarg. * Dedicated dynamic + viscous friction setters (write_joint_{dynamic,viscous}_friction_coefficient_to_sim_{index,mask}) that touch only their slot of the combined (N, J, 3) DOF_FRICTION_PROPERTIES buffer. The combined write_joint_friction_coefficient_to_sim_index/_mask still accepts all three components as kwargs (Coulomb static + dynamic + optional viscous) for source-compatible PhysX call sites. * Deprecated non-indexed shorthand shims for friction (x3) and root / joint state (x4), forwarding to the index variants with a DeprecationWarning, matching PhysX's deprecated section. * Wrench-composer return types tightened to non-None (instantaneous_wrench_composer / permanent_wrench_composer); composers are always set in _create_buffers, mirroring PhysX/Newton. * Section organisation matches PhysX exactly: Properties -> Operations -> Operations - Finders -> Operations - State Writers -> Operations - Simulation Parameters Writers -> Operations - Setters -> Operations - Tendons -> Internal helper -> Internal simulation callbacks -> Internal helpers -- Actuators -> Internal helpers -- Debugging -> Deprecated methods. Section delimiters use bare """Section.""" docstring blocks (Newton/PhysX convention). ArticulationData ^^^^^^^^^^^^^^^^ * Pull-on-demand timestamped buffers; CPU-only bindings route through pinned-host staging (PR #5329 pattern). * Property names match PhysX exactly: joint_friction_coeff, joint_dynamic_friction_coeff, joint_viscous_friction_coeff (no *_static / *_dynamic / *_viscous renames). * SI units annotated on every public property docstring ([m or rad, depending on joint type], [m/s or rad/s, ...], [N*m], [kg], etc.) per AGENTS.md. * binding_getter parameter on __init__ typed as Callable[[int], Any] | None. * Section organisation matches PhysX (Defaults -> Joint commands -> Joint properties -> Fixed tendon -> Spatial tendon -> Root state -> Body state -> Joint state -> Derived -> Sliced -> Internal helpers -> Deprecated properties). Kernels ^^^^^^^ * Articulation-specific kernels in isaaclab_ovphysx/assets/articulation/kernels.py (soft-limit clamp, friction-data writer, finite-difference joint-acc helper, body-CoM pose composer); shared kernels promoted to isaaclab_ovphysx/assets/kernels.py. * Per-kernel docstrings document purpose, shape/dtype/SI units, and divergence notes where the OVPhysX implementation differs (e.g. _fd_joint_acc takes inv_dt rather than dt to avoid per-element division). Tests ^^^^^ * Real-backend test_articulation.py mirrors isaaclab_physx 1-to-1 under run_ovphysx.sh; 99 tests pass on each of CPU + CUDA. * test_articulation_helpers.py covers the kitless-only helpers (tendon scoping, mock binding shapes). * Cross-backend test_articulation_iface.py runs the OVPhysX path: 544 tests pass, 16 skipped, 0 failed on each of CPU + CUDA. Brings the iface helper up to the actual ArticulationData constructor signature and broadcasts scalar inputs across joint / fixed-tendon / spatial-tendon mask setters that previously rejected them. OVPhysX-only surface dropped ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ * Articulation: set_external_force_and_torque_{index,mask} (use instantaneous_wrench_composer.add_forces_and_torques_*), set_spatial_tendon_limit_{index,mask}, and set_spatial_tendon_rest_length_{index,mask} (NotImplementedError stubs PhysX never had). * ArticulationData: body_pose_w / body_lin_vel_w / body_ang_vel_w / body_acc_w / body_link_acc_w (base class provides matching defaults), body_inv_mass / body_inv_inertia, fixed_tendon_limit (PhysX exposes only fixed_tendon_pos_limits), spatial_tendon_limit / spatial_tendon_rest_length (no PhysX equivalent). --- .../test/assets/test_articulation_iface.py | 29 +- ...toiner-feat-ovphysx_articulation.minor.rst | 109 + .../assets/articulation/articulation.py | 4697 +++++++++++------ .../assets/articulation/articulation_data.py | 1617 ++++-- .../assets/articulation/kernels.py | 178 +- .../isaaclab_ovphysx/assets/kernels.py | 268 + .../isaaclab_ovphysx/tensor_types.py | 15 +- .../test/assets/test_articulation.py | 2528 ++++++++- .../test/assets/test_articulation_helpers.py | 142 + 9 files changed, 7392 insertions(+), 2191 deletions(-) create mode 100644 source/isaaclab_ovphysx/changelog.d/antoiner-feat-ovphysx_articulation.minor.rst create mode 100644 source/isaaclab_ovphysx/test/assets/test_articulation_helpers.py diff --git a/source/isaaclab/test/assets/test_articulation_iface.py b/source/isaaclab/test/assets/test_articulation_iface.py index 498091f51058..56f69226e139 100644 --- a/source/isaaclab/test/assets/test_articulation_iface.py +++ b/source/isaaclab/test/assets/test_articulation_iface.py @@ -288,20 +288,27 @@ def create_ovphysx_articulation( object.__setattr__(articulation, "_num_spatial_tendons", num_spatial_tendons) # Create ArticulationData - data = OvPhysxArticulationData(mock_bindings.bindings, device) - data._num_instances = num_instances - data._num_joints = num_joints - data._num_bodies = num_bodies - data._num_fixed_tendons = num_fixed_tendons - data._num_spatial_tendons = num_spatial_tendons + data = OvPhysxArticulationData( + mock_bindings.bindings, + device, + num_instances=num_instances, + num_bodies=num_bodies, + num_joints=num_joints, + num_fixed_tendons=num_fixed_tendons, + num_spatial_tendons=num_spatial_tendons, + body_names=body_names, + joint_names=joint_names, + fixed_tendon_names=fixed_tendon_names, + spatial_tendon_names=spatial_tendon_names, + ) data._is_fixed_base = False - data.body_names = body_names - data.joint_names = joint_names - data.fixed_tendon_names = fixed_tendon_names - data.spatial_tendon_names = spatial_tendon_names - data._create_buffers() object.__setattr__(articulation, "_data", data) + # Allocate the articulation-side index/mask caches and wrench buffer that + # _initialize_impl would normally populate. Wrench composers created here + # are immediately overwritten by the mocks below. + articulation._create_buffers() + # Wrench composers mock_inst_wrench = MockWrenchComposer(articulation) mock_perm_wrench = MockWrenchComposer(articulation) diff --git a/source/isaaclab_ovphysx/changelog.d/antoiner-feat-ovphysx_articulation.minor.rst b/source/isaaclab_ovphysx/changelog.d/antoiner-feat-ovphysx_articulation.minor.rst new file mode 100644 index 000000000000..24913760aa9b --- /dev/null +++ b/source/isaaclab_ovphysx/changelog.d/antoiner-feat-ovphysx_articulation.minor.rst @@ -0,0 +1,109 @@ +Added +^^^^^ + +* Added :class:`~isaaclab_ovphysx.assets.Articulation` and + :class:`~isaaclab_ovphysx.assets.ArticulationData` for multi-DOF articulated + robots against the OVPhysX backend, satisfying the + :class:`~isaaclab.assets.BaseArticulation` and + :class:`~isaaclab.assets.BaseArticulationData` contracts. Public surface + matches the PhysX/Newton conventions: kwarg-only ``write_*_to_sim_index`` / + ``write_*_to_sim_mask`` writers and ``set_*_index`` / ``set_*_mask`` setters + for root state, joint state, joint properties, body properties, joint + command targets, fixed/spatial tendon properties, and external wrenches via + :attr:`~isaaclab_ovphysx.assets.Articulation.instantaneous_wrench_composer` + / :attr:`~isaaclab_ovphysx.assets.Articulation.permanent_wrench_composer`. + The full IsaacLab actuator pipeline (``compute`` / + ``_apply_actuator_model`` / ``_process_actuators_cfg``) is implemented on + top of the wheel's ``DOF_ACTUATION_FORCE`` / + ``DOF_POSITION_TARGET`` / ``DOF_VELOCITY_TARGET`` bindings. +* Added articulation-specific Warp kernels in + :mod:`isaaclab_ovphysx.assets.articulation.kernels`: soft-limit refresh, + default-joint-pos clamp, friction-component scatter (index + mask + variants). Six articulation kernels were also folded into the shared + :mod:`isaaclab_ovphysx.assets.kernels` module for reuse with + :class:`~isaaclab_ovphysx.assets.RigidObject` and + :class:`~isaaclab_ovphysx.assets.RigidObjectCollection`. +* Added init-time validation in + :meth:`~isaaclab_ovphysx.assets.Articulation._initialize_impl` that raises + ``RuntimeError`` when ``cfg.prim_path`` resolves to no + ``UsdPhysics.ArticulationRootAPI`` prim or to multiple roots, and + ``ValueError`` (via :meth:`_validate_cfg`) when any default joint + position is outside ``[lower, upper]`` or any default joint velocity + exceeds the per-joint maximum. Mirrors the PhysX backend. +* Added support for ``cfg.articulation_root_prim_path`` in + :meth:`~isaaclab_ovphysx.assets.Articulation._initialize_impl`: when the + user supplies an explicit subpath the binding pattern is extended + directly instead of running the auto-discovery walk, and a + ``RuntimeError`` is raised when the resulting expression resolves to no + prim in the USD stage. + +Changed +^^^^^^^ + +* **Breaking:** Renamed ``Articulation`` write/set methods to the dual + ``*_index`` / ``*_mask`` form and dropped the legacy ``full_data`` + flag. Index methods accept partial data shaped + ``(len(env_ids), len(joint_or_body_ids), ...)``; mask methods accept + full-shape data and a ``wp.bool`` mask. All keyword-only arguments live + after ``*,``; no positional fall-through. Migration: replace + ``write_X_to_sim(..., from_mask=True)`` with ``write_X_to_sim_mask(..., mask=...)``. +* **Breaking:** Removed the ``_write_body_state`` plumbing layer. + Deprecated state-writer shims (``write_root_state_to_sim``, + ``write_root_com_state_to_sim``, ``write_root_link_state_to_sim``, + joint-state equivalents) now call the public ``write_*_to_sim_index`` + methods directly. Behaviour is preserved. +* Changed ``Articulation.root_view`` to return the per-tensor-type bindings + dict (``self._bindings``). The OVPhysX wheel does not expose a single + ``ArticulationView`` object; callers that previously walked + ``root_view.shared_metatype`` / ``root_view.max_dofs`` should read from + :attr:`~isaaclab_ovphysx.assets.Articulation.num_joints` / + :attr:`~isaaclab_ovphysx.assets.Articulation.num_bodies` / + :attr:`~isaaclab_ovphysx.assets.Articulation.body_names` / + :attr:`~isaaclab_ovphysx.assets.Articulation.joint_names` instead. +* Changed every ``ArticulationData`` public property to return a + :class:`~isaaclab.utils.ProxyArray` (warp + torch dual view); raw + ``wp.array`` is reserved for one-shot config buffers. Eager + ``TimestampedBufferWarp`` allocation in :meth:`_create_buffers` makes + every buffer a single source of truth — no + ``_invalidate_caches`` / ``_ensure_*_buffers`` machinery. +* Changed ``Articulation`` body and DOF property writers to honor the + wheel's actual binding device. Tensor-type membership in + :data:`isaaclab_ovphysx.tensor_types._CPU_ONLY_TYPES` now reflects what + the wheel exposes: ``BODY_MASS``, ``BODY_COM_POSE``, ``BODY_INERTIA``, + ``DOF_STIFFNESS``, ``DOF_DAMPING``, ``DOF_LIMIT``, ``DOF_MAX_VELOCITY``, + ``DOF_MAX_FORCE``, ``DOF_ARMATURE``, ``DOF_FRICTION_PROPERTIES`` are + CPU-only (write goes through pinned-host staging); fixed and spatial + tendon bindings write directly from sim-device buffers. +* Changed :meth:`~isaaclab_ovphysx.assets.Articulation.write_joint_friction_coefficient_to_sim_index` + / ``_mask`` to accept ``joint_dynamic_friction_coeff`` and + ``joint_viscous_friction_coeff`` keyword arguments (each + ``float | torch.Tensor | wp.array | None``). ``None`` preserves the + existing component on the wheel; matches the PhysX backend. +* Changed :meth:`~isaaclab_ovphysx.assets.Articulation.write_joint_position_limit_to_sim_index` + / ``_mask`` to clamp ``default_joint_pos`` and refresh + ``soft_joint_pos_limits`` when the new hard limits invalidate the + defaults, matching the PhysX backend (with a + ``warn_limit_violation`` log). +* Changed every fixed/spatial tendon ``set_*_index`` / ``set_*_mask`` setter + to accept a scalar :class:`float` for the value argument; broadcast is + materialized via :meth:`_broadcast_scalar_to_2d`. Mirrors PhysX. +* Implemented the previously stubbed + :meth:`~isaaclab_ovphysx.assets.Articulation.write_fixed_tendon_properties_to_sim_index` + / ``_mask`` and + :meth:`~isaaclab_ovphysx.assets.Articulation.write_spatial_tendon_properties_to_sim_index` + / ``_mask``: each iterates the per-tensor bindings since the OVPhysX + wheel has no batch ``set_*_tendon_properties`` setter. + +Removed +^^^^^^^ + +* **Breaking:** Removed the ``full_data`` keyword-argument from every + ``Articulation`` ``*_index`` writer/setter. Index methods now strictly + accept partial data; full-data callers should use the matching + ``*_mask`` overload. +* Removed the stop-gap :mod:`isaaclab_ovphysx.assets.kernels_old` module; + the six articulation kernels it housed + (``_compose_root_com_pose``, ``_compute_heading``, ``_copy_first_body``, + ``_projected_gravity``, ``_world_vel_to_body_ang``, + ``_world_vel_to_body_lin``) are now in + :mod:`isaaclab_ovphysx.assets.kernels`. diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py index cb459221c681..89c47039686e 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py @@ -3,51 +3,83 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Articulation implementation backed by ovphysx TensorBindingsAPI.""" +"""OVPhysX-backed Articulation implementation. + +Mirrors the post-refactor :class:`~isaaclab_ovphysx.assets.RigidObject` +shape with the API surface coming from Newton. Eager binding creation +in :meth:`_initialize_impl`; dual mask+index API; pinned-host CPU +staging via the data class's ``_binding_read`` / ``_binding_write`` +helpers (the PR #5329 pattern). +""" from __future__ import annotations import logging import re +import warnings from collections.abc import Sequence -from typing import TYPE_CHECKING, Any +from typing import Any import numpy as np import torch import warp as wp +from pxr import UsdPhysics + +import isaaclab.sim as sim_utils +from isaaclab.assets.articulation.articulation_cfg import ArticulationCfg from isaaclab.assets.articulation.base_articulation import BaseArticulation from isaaclab.physics import PhysicsManager from isaaclab.utils.string import resolve_matching_names from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_ovphysx import tensor_types as TT -from isaaclab_ovphysx.assets.kernels import _body_wrench_to_world, _scatter_rows_partial +from isaaclab_ovphysx.assets import kernels as shared_kernels +from isaaclab_ovphysx.assets.kernels import _body_wrench_to_world from isaaclab_ovphysx.physics import OvPhysxManager from .articulation_data import ArticulationData -from .kernels import update_soft_joint_pos_limits - -if TYPE_CHECKING: - from isaaclab.actuators import ActuatorBase - from isaaclab.assets.articulation.articulation_cfg import ArticulationCfg +from .kernels import ( + clamp_default_joint_pos_and_update_soft_limits_index, + clamp_default_joint_pos_and_update_soft_limits_mask, + update_soft_joint_pos_limits, + write_joint_friction_data_to_buffer_index, + write_joint_friction_data_to_buffer_mask, +) logger = logging.getLogger(__name__) class Articulation(BaseArticulation): - """Articulation backed by the ovphysx TensorBindingsAPI. + """OVPhysX-backed Articulation asset class. - Reads and writes simulation state through ovphysx.TensorBinding objects created - from the OvPhysxManager's PhysX instance. - """ + Mirrors the OVPhysX RigidObject shape with the public API surface + coming from Newton. Joint, body, tendon, and root state are read + via :class:`ovphysx.TensorBinding` objects acquired from the + :class:`~isaaclab_ovphysx.physics.OvPhysxManager`. Writes go through + the dual mask+index API; CPU-only bindings (mass, COM, inertia, + joint properties, tendon properties) route through pinned-host + staging buffers (PR #5329 pattern). - __backend_name__ = "ovphysx" + .. _USD ArticulationRootAPI: https://openusd.org/dev/api/class_usd_physics_articulation_root_a_p_i.html + """ cfg: ArticulationCfg + """Configuration instance for the articulation.""" + + __backend_name__: str = "ovphysx" + """The name of the backend for the articulation.""" def __init__(self, cfg: ArticulationCfg): + """Initialize the articulation. + + Args: + cfg: A configuration instance. + """ super().__init__(cfg) + # Bindings are created lazily (on first access) to avoid allocating + # handles for tensor types the user never queries. + self._bindings: dict[int, Any] = {} """ Properties @@ -64,29 +96,29 @@ def num_instances(self) -> int: return self._num_instances @property - def is_fixed_base(self) -> bool: - """Whether the articulation is a fixed-base or floating-base system.""" - return self._is_fixed_base + def num_bodies(self) -> int: + """Number of bodies (links) per articulation.""" + return self._num_bodies @property def num_joints(self) -> int: - """Number of joints in the articulation.""" + """Number of degrees of freedom per articulation.""" return self._num_joints @property def num_fixed_tendons(self) -> int: - """Number of fixed tendons in the articulation.""" - return getattr(self, "_num_fixed_tendons", 0) + """Number of fixed tendons per articulation.""" + return self._num_fixed_tendons @property def num_spatial_tendons(self) -> int: - """Number of spatial tendons in the articulation.""" - return getattr(self, "_num_spatial_tendons", 0) + """Number of spatial tendons per articulation.""" + return self._num_spatial_tendons @property - def num_bodies(self) -> int: - """Number of bodies (links) in the articulation.""" - return self._num_bodies + def body_names(self) -> list[str]: + """Ordered names of bodies in the articulation.""" + return self._body_names @property def joint_names(self) -> list[str]: @@ -96,67 +128,121 @@ def joint_names(self) -> list[str]: @property def fixed_tendon_names(self) -> list[str]: """Ordered names of fixed tendons in the articulation.""" - return getattr(self, "_fixed_tendon_names", []) + return self._fixed_tendon_names @property def spatial_tendon_names(self) -> list[str]: """Ordered names of spatial tendons in the articulation.""" - return getattr(self, "_spatial_tendon_names", []) + return self._spatial_tendon_names @property - def body_names(self) -> list[str]: - """Ordered names of bodies (links) in the articulation.""" - return self._body_names + def is_fixed_base(self) -> bool: + """Whether the articulation is a fixed-base or floating-base system.""" + return self._is_fixed_base @property - def root_view(self) -> Any: - """Root articulation view (not available for ovphysx backend).""" - return None + def root_view(self) -> dict[int, Any]: + """Root view for the asset. + + OVPhysX exposes per-tensor-type bindings rather than a single opaque view object + as used by the PhysX and Newton backends. Callers that need low-level binding + access should call :meth:`_get_binding` rather than iterating this dict directly. + For high-level state access (instance counts, prim paths, transforms), use the + :attr:`num_instances`, :attr:`body_names`, and :attr:`~ArticulationData.root_link_pose_w` + accessors instead. + + .. note:: + Use this view with caution. It requires handling of tensors in a specific way. + """ + return self._bindings @property - def instantaneous_wrench_composer(self) -> WrenchComposer | None: - """Wrench composer for forces applied only during the current step.""" + def instantaneous_wrench_composer(self) -> WrenchComposer: + """Instantaneous wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are only valid for the current simulation step. At the end of the simulation step, the wrenches set + to this object are discarded. This is useful to apply forces that change all the time, things like drag forces + for instance. + """ return self._instantaneous_wrench_composer @property - def permanent_wrench_composer(self) -> WrenchComposer | None: - """Wrench composer for forces applied persistently every step.""" + def permanent_wrench_composer(self) -> WrenchComposer: + """Permanent wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are persistent and are applied to the simulation at every step. This is useful to apply forces that + are constant over a period of time, things like the thrust of a motor for instance. + """ return self._permanent_wrench_composer """ Operations. """ - def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None) -> None: + def reset( + self, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_mask: wp.array | None = None + ) -> None: """Reset the articulation. - .. caution:: - If both `env_ids` and `env_mask` are provided, then `env_mask` takes precedence over `env_ids`. - Args: - env_ids: Environment indices. If None, then all indices are used. - env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are + updated. Shape is (num_instances,). """ - # use ellipses object to skip initial indices. if (env_ids is None) or (env_ids == slice(None)): env_ids = slice(None) - # reset actuators - for actuator in self.actuators.values(): - actuator.reset(env_ids) - # reset external wrenches. self._instantaneous_wrench_composer.reset(env_ids, env_mask) self._permanent_wrench_composer.reset(env_ids, env_mask) def write_data_to_sim(self) -> None: - """Apply external wrenches, actuator model, and write commands into the simulation.""" - # Apply external wrenches (before actuators, same as PhysX backend). - self._apply_external_wrenches() + """Apply external wrenches, run the actuator model, and write commands to the simulation. + + Execution order mirrors the PhysX backend: + + 1. Compose and write external wrenches (body-frame -> world-frame via + :func:`_body_wrench_to_world`, then write to ``LINK_WRENCH``). + 2. Run the actuator model (:meth:`_apply_actuator_model`) to compute + joint efforts from user-supplied targets. + 3. Write the effort tensor to the ``DOF_ACTUATION_FORCE`` binding. + 4. For implicit actuators, write position and velocity targets in a + single shot to ``DOF_POSITION_TARGET`` / ``DOF_VELOCITY_TARGET``. + """ + # 1. External wrenches. + inst = self._instantaneous_wrench_composer + perm = self._permanent_wrench_composer + if inst.active or perm.active: + if inst.active: + if perm.active: + inst.add_raw_buffers_from(perm) + force_b = inst.out_force_b.warp + torque_b = inst.out_torque_b.warp + else: + force_b = perm.out_force_b.warp + torque_b = perm.out_torque_b.warp + + poses = self._data.body_link_pose_w.warp # (N, num_bodies) wp.transformf + wp.launch( + _body_wrench_to_world, + dim=(self._num_instances, self._num_bodies), + inputs=[force_b, torque_b, poses], + outputs=[self._wrench_buf], + device=self._device, + ) + binding = self._get_binding(TT.LINK_WRENCH) + if binding is not None: + binding.write(self._wrench_buf) + inst.reset() + # 2. Actuator model. self._apply_actuator_model() - # Write effort tensor to simulation. + + # 3. Write effort (applies even when no actuators: zeros are safe). if self._effort_binding is not None: self._effort_binding.write(self._effort_write_view) - # Write position and velocity targets in one shot (not per-actuator). + + # 4. Write position / velocity targets for implicit actuators. if self._has_implicit_actuators: if self._pos_target_binding is not None: self._pos_target_binding.write(self._pos_target_write_view) @@ -164,10 +250,10 @@ def write_data_to_sim(self) -> None: self._vel_target_binding.write(self._vel_target_write_view) def update(self, dt: float) -> None: - """Update internal data buffers after a simulation step. + """Update the simulation data. Args: - dt: The simulation time step [s] used for finite-difference quantities. + dt: Time-step in seconds [s]. """ self._data.update(dt) @@ -198,8 +284,8 @@ def find_joints( ) -> tuple[list[int], list[str]]: """Find joints in the articulation based on the name keys. - Please check the :func:`isaaclab.utils.string.resolve_matching_names` function for more - information on the name matching. + Please see the :func:`isaaclab.utils.string.resolve_matching_names` function for more information + on the name matching. Args: name_keys: A regular expression or a list of regular expressions to match the joint names. @@ -222,20 +308,21 @@ def find_fixed_tendons( ) -> tuple[list[int], list[str]]: """Find fixed tendons in the articulation based on the name keys. + Please see the :func:`isaaclab.utils.string.resolve_matching_names` function for more information + on the name matching. + Args: - name_keys: A regular expression or a list of regular expressions - to match the joint names with fixed tendons. - tendon_subsets: A subset of joints with fixed tendons to search - for. Defaults to None, which means all joints in the - articulation are searched. - preserve_order: Whether to preserve the order of the name keys in - the output. Defaults to False. + name_keys: A regular expression or a list of regular expressions to match the + joint names with fixed tendons. + tendon_subsets: A subset of joints with fixed tendons to search for. Defaults to None, which means + all joints in the articulation are searched. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. Returns: A tuple of lists containing the tendon indices and names. """ if tendon_subsets is None: - tendon_subsets = self.fixed_tendon_names + tendon_subsets = self._fixed_tendon_names return resolve_matching_names(name_keys, tendon_subsets, preserve_order) def find_spatial_tendons( @@ -246,19 +333,20 @@ def find_spatial_tendons( ) -> tuple[list[int], list[str]]: """Find spatial tendons in the articulation based on the name keys. + Please see the :func:`isaaclab.utils.string.resolve_matching_names` function for more information + on the name matching. + Args: - name_keys: A regular expression or a list of regular expressions - to match the tendon names. - tendon_subsets: A subset of tendons to search for. Defaults to - None, which means all tendons in the articulation are searched. - preserve_order: Whether to preserve the order of the name keys in - the output. Defaults to False. + name_keys: A regular expression or a list of regular expressions to match the tendon names. + tendon_subsets: A subset of tendons to search for. Defaults to None, which means all tendons + in the articulation are searched. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. Returns: A tuple of lists containing the tendon indices and names. """ if tendon_subsets is None: - tendon_subsets = self.spatial_tendon_names + tendon_subsets = self._spatial_tendon_names return resolve_matching_names(name_keys, tendon_subsets, preserve_order) """ @@ -269,19 +357,25 @@ def write_root_pose_to_sim_index( self, *, root_pose: torch.Tensor | wp.array, - env_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set the root pose over selected environment indices into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + Args: - root_pose: Root poses in simulation frame. Shape is (len(env_ids),) with dtype wp.transformf. + root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7) + or (len(env_ids),) with dtype wp.transformf. env_ids: Environment indices. If None, then all indices are used. """ - n = self._n_envs_index(env_ids) - self.assert_shape_and_dtype(root_pose, (n,), wp.transformf, "root_pose") - self._write_root_state(TT.ROOT_POSE, root_pose, env_ids) + self.write_root_link_pose_to_sim_index(root_pose=root_pose, env_ids=env_ids) def write_root_pose_to_sim_mask( self, @@ -289,34 +383,62 @@ def write_root_pose_to_sim_mask( root_pose: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root pose over masked environments into the simulation. + """Set the root pose over selected environment mask into the simulation. - The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. Args: - root_pose: Root poses in simulation frame. Shape is (num_instances,) with dtype wp.transformf. - env_mask: Environment mask. If None, then all instances are updated. + root_pose: Root poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - self.assert_shape_and_dtype(root_pose, (self._num_instances,), wp.transformf, "root_pose") - self._write_root_state(TT.ROOT_POSE, root_pose, mask=env_mask) + self.write_root_link_pose_to_sim_mask(root_pose=root_pose, env_mask=env_mask) def write_root_link_pose_to_sim_index( self, *, root_pose: torch.Tensor | wp.array, - env_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set the root link pose over selected environment indices into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + Args: - root_pose: Root link poses in simulation frame. Shape is (len(env_ids),) with dtype wp.transformf. + root_pose: Root link poses in simulation frame. Shape is (len(env_ids), 7) + or (len(env_ids),) with dtype wp.transformf. env_ids: Environment indices. If None, then all indices are used. """ - n = self._n_envs_index(env_ids) - self.assert_shape_and_dtype(root_pose, (n,), wp.transformf, "root_pose") - self._write_root_state(TT.ROOT_POSE, root_pose, env_ids) + env_ids = self._resolve_env_ids(env_ids) + self.assert_shape_and_dtype(root_pose, (env_ids.shape[0],), wp.transformf, "root_pose") + wp.launch( + shared_kernels.set_root_link_pose_to_sim_index, + dim=env_ids.shape[0], + inputs=[root_pose, env_ids], + outputs=[self.data.root_link_pose_w], + device=self._device, + ) + # Invalidate dependent timestamps so the next reads re-fetch. Mirrors + # PhysX (articulation.py:457-464): root link pose changes the body + # kinematics chain, so all body-pose buffers go stale. + self.data._root_com_pose_w.timestamp = -1.0 + self.data._body_link_pose_w.timestamp = -1.0 + self.data._body_com_pose_w.timestamp = -1.0 + # Push cache to the wheel via an indexed write. + binding = self._get_binding(TT.ROOT_POSE) + binding.write(self.data._root_link_pose_w.data.view(wp.float32), indices=env_ids) def write_root_link_pose_to_sim_mask( self, @@ -324,36 +446,76 @@ def write_root_link_pose_to_sim_mask( root_pose: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root link pose over masked environments into the simulation. + """Set the root link pose over selected environment mask into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + Args: - root_pose: Root link poses in simulation frame. Shape is (num_instances,) with dtype wp.transformf. - env_mask: Environment mask. If None, then all instances are updated. + root_pose: Root poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ + env_mask_wp = self._resolve_env_mask(env_mask) self.assert_shape_and_dtype(root_pose, (self._num_instances,), wp.transformf, "root_pose") - self._write_root_state(TT.ROOT_POSE, root_pose, mask=env_mask) + wp.launch( + shared_kernels.set_root_link_pose_to_sim_mask, + dim=self._num_instances, + inputs=[root_pose, env_mask_wp], + outputs=[self.data.root_link_pose_w], + device=self._device, + ) + # See write_root_link_pose_to_sim_index for why these are invalidated. + self.data._root_com_pose_w.timestamp = -1.0 + self.data._body_link_pose_w.timestamp = -1.0 + self.data._body_com_pose_w.timestamp = -1.0 + binding = self._get_binding(TT.ROOT_POSE) + binding.write(self.data._root_link_pose_w.data.view(wp.float32), mask=env_mask_wp) def write_root_com_pose_to_sim_index( self, *, root_pose: torch.Tensor | wp.array, - env_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set the root center of mass pose over selected environment indices into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). The orientation is the orientation of the principal axes of inertia. + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + Args: - root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids),) - with dtype wp.transformf. + root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7) + or (len(env_ids),) with dtype wp.transformf. env_ids: Environment indices. If None, then all indices are used. """ - n = self._n_envs_index(env_ids) - self.assert_shape_and_dtype(root_pose, (n,), wp.transformf, "root_pose") - self._write_root_state(TT.ROOT_POSE, root_pose, env_ids) + env_ids = self._resolve_env_ids(env_ids) + self.assert_shape_and_dtype(root_pose, (env_ids.shape[0],), wp.transformf, "root_pose") + wp.launch( + shared_kernels.set_root_com_pose_to_sim_index, + dim=env_ids.shape[0], + inputs=[root_pose, self.data.body_com_pose_b, env_ids], + outputs=[self.data.root_com_pose_w, self.data.root_link_pose_w], + device=self._device, + ) + # Mirrors PhysX (articulation.py:548-556): writing root COM pose updates + # the inferred root link pose, which in turn invalidates the body chain. + self.data._body_link_pose_w.timestamp = -1.0 + self.data._body_com_pose_w.timestamp = -1.0 + binding = self._get_binding(TT.ROOT_POSE) + binding.write(self.data._root_link_pose_w.data.view(wp.float32), indices=env_ids) def write_root_com_pose_to_sim_mask( self, @@ -361,37 +523,64 @@ def write_root_com_pose_to_sim_mask( root_pose: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root center of mass pose over masked environments into the simulation. + """Set the root center of mass pose over selected environment mask into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). The orientation is the orientation of the principal axes of inertia. + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + Args: - root_pose: Root center of mass poses in simulation frame. Shape is (num_instances,) - with dtype wp.transformf. - env_mask: Environment mask. If None, then all instances are updated. + root_pose: Root center of mass poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ + env_mask_wp = self._resolve_env_mask(env_mask) self.assert_shape_and_dtype(root_pose, (self._num_instances,), wp.transformf, "root_pose") - self._write_root_state(TT.ROOT_POSE, root_pose, mask=env_mask) + wp.launch( + shared_kernels.set_root_com_pose_to_sim_mask, + dim=self._num_instances, + inputs=[root_pose, self.data.body_com_pose_b, env_mask_wp], + outputs=[self.data.root_com_pose_w, self.data.root_link_pose_w], + device=self._device, + ) + # See write_root_com_pose_to_sim_index for why these are invalidated. + self.data._body_link_pose_w.timestamp = -1.0 + self.data._body_com_pose_w.timestamp = -1.0 + binding = self._get_binding(TT.ROOT_POSE) + binding.write(self.data._root_link_pose_w.data.view(wp.float32), mask=env_mask_wp) def write_root_velocity_to_sim_index( self, *, root_velocity: torch.Tensor | wp.array, - env_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set the root velocity over selected environment indices into the simulation. + """Set the root center of mass velocity over selected environment indices into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + Args: - root_velocity: Root velocities in simulation world frame. Shape is (len(env_ids),) - with dtype wp.spatial_vectorf. + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6) + or (len(env_ids),) with dtype wp.spatial_vectorf. env_ids: Environment indices. If None, then all indices are used. """ - n = self._n_envs_index(env_ids) - self.assert_shape_and_dtype(root_velocity, (n,), wp.spatial_vectorf, "root_velocity") - self._write_root_state(TT.ROOT_VELOCITY, root_velocity, env_ids) + self.write_root_com_velocity_to_sim_index(root_velocity=root_velocity, env_ids=env_ids) def write_root_velocity_to_sim_mask( self, @@ -399,36 +588,60 @@ def write_root_velocity_to_sim_mask( root_velocity: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root velocity over masked environments into the simulation. + """Set the root center of mass velocity over selected environment mask into the simulation. - The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. Args: - root_velocity: Root velocities in simulation world frame. Shape is (num_instances,) - with dtype wp.spatial_vectorf. - env_mask: Environment mask. If None, then all instances are updated. + root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - self.assert_shape_and_dtype(root_velocity, (self._num_instances,), wp.spatial_vectorf, "root_velocity") - self._write_root_state(TT.ROOT_VELOCITY, root_velocity, mask=env_mask) + self.write_root_com_velocity_to_sim_mask(root_velocity=root_velocity, env_mask=env_mask) def write_root_com_velocity_to_sim_index( self, *, root_velocity: torch.Tensor | wp.array, - env_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set the root center of mass velocity over selected environment indices into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + Args: - root_velocity: Root center of mass velocities in simulation world frame. - Shape is (len(env_ids),) with dtype wp.spatial_vectorf. + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6) + or (len(env_ids),) with dtype wp.spatial_vectorf. env_ids: Environment indices. If None, then all indices are used. """ - n = self._n_envs_index(env_ids) - self.assert_shape_and_dtype(root_velocity, (n,), wp.spatial_vectorf, "root_velocity") - self._write_root_state(TT.ROOT_VELOCITY, root_velocity, env_ids) + env_ids = self._resolve_env_ids(env_ids) + self.assert_shape_and_dtype(root_velocity, (env_ids.shape[0],), wp.spatial_vectorf, "root_velocity") + wp.launch( + shared_kernels.set_root_com_velocity_to_sim_index, + dim=env_ids.shape[0], + inputs=[root_velocity, env_ids, self._num_bodies], + outputs=[self.data.root_com_vel_w, self.data.body_com_acc_w], + device=self._device, + ) + # Invalidate dependent root_link_vel timestamp. + self.data._root_link_vel_w.timestamp = -1.0 + binding = self._get_binding(TT.ROOT_VELOCITY) + binding.write(self.data._root_com_vel_w.data.view(wp.float32), indices=env_ids) def write_root_com_velocity_to_sim_mask( self, @@ -436,23 +649,43 @@ def write_root_com_velocity_to_sim_mask( root_velocity: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root center of mass velocity over masked environments into the simulation. + """Set the root center of mass velocity over selected environment mask into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + Args: - root_velocity: Root center of mass velocities in simulation world frame. - Shape is (num_instances,) with dtype wp.spatial_vectorf. - env_mask: Environment mask. If None, then all instances are updated. + root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ + env_mask_wp = self._resolve_env_mask(env_mask) self.assert_shape_and_dtype(root_velocity, (self._num_instances,), wp.spatial_vectorf, "root_velocity") - self._write_root_state(TT.ROOT_VELOCITY, root_velocity, mask=env_mask) + wp.launch( + shared_kernels.set_root_com_velocity_to_sim_mask, + dim=self._num_instances, + inputs=[root_velocity, env_mask_wp, self._num_bodies], + outputs=[self.data.root_com_vel_w, self.data.body_com_acc_w], + device=self._device, + ) + self.data._root_link_vel_w.timestamp = -1.0 + binding = self._get_binding(TT.ROOT_VELOCITY) + binding.write(self.data._root_com_vel_w.data.view(wp.float32), mask=env_mask_wp) def write_root_link_velocity_to_sim_index( self, *, root_velocity: torch.Tensor | wp.array, - env_ids: Sequence[int] | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set the root link velocity over selected environment indices into the simulation. @@ -461,14 +694,35 @@ def write_root_link_velocity_to_sim_index( .. note:: This sets the velocity of the root's frame rather than the root's center of mass. + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + Args: - root_velocity: Root frame velocities in simulation world frame. - Shape is (len(env_ids),) with dtype wp.spatial_vectorf. + root_velocity: Root frame velocities in simulation world frame. Shape is (len(env_ids), 6) + or (len(env_ids),) with dtype wp.spatial_vectorf. env_ids: Environment indices. If None, then all indices are used. """ - n = self._n_envs_index(env_ids) - self.assert_shape_and_dtype(root_velocity, (n,), wp.spatial_vectorf, "root_velocity") - self._write_root_state(TT.ROOT_VELOCITY, root_velocity, env_ids) + env_ids = self._resolve_env_ids(env_ids) + self.assert_shape_and_dtype(root_velocity, (env_ids.shape[0],), wp.spatial_vectorf, "root_velocity") + wp.launch( + shared_kernels.set_root_link_velocity_to_sim_index, + dim=env_ids.shape[0], + inputs=[ + root_velocity, + self.data.body_com_pose_b, + self.data.root_link_pose_w, + env_ids, + self._num_bodies, + ], + outputs=[self.data.root_link_vel_w, self.data.root_com_vel_w, self.data.body_com_acc_w], + device=self._device, + ) + binding = self._get_binding(TT.ROOT_VELOCITY) + binding.write(self.data._root_com_vel_w.data.view(wp.float32), indices=env_ids) def write_root_link_velocity_to_sim_mask( self, @@ -476,118 +730,255 @@ def write_root_link_velocity_to_sim_mask( root_velocity: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root link velocity over masked environments into the simulation. + """Set the root link velocity over selected environment mask into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. .. note:: This sets the velocity of the root's frame rather than the root's center of mass. - Args: - root_velocity: Root frame velocities in simulation world frame. - Shape is (num_instances,) with dtype wp.spatial_vectorf. - env_mask: Environment mask. If None, then all instances are updated. - """ - self.assert_shape_and_dtype(root_velocity, (self._num_instances,), wp.spatial_vectorf, "root_velocity") - self._write_root_state(TT.ROOT_VELOCITY, root_velocity, mask=env_mask) - - """ - Operations - Joint State Writers. - """ + .. note:: + This method expects full data. - def write_joint_state_to_sim_mask( - self, - joint_pos: torch.Tensor | wp.array, - joint_vel: torch.Tensor | wp.array, - env_mask: wp.array | None = None, - joint_mask: wp.array | None = None, - ) -> None: - """Write joint positions and velocities over masked environments into the simulation. + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. Args: - joint_pos: Joint positions. Shape is (num_instances, num_joints). - joint_vel: Joint velocities. Shape is (num_instances, num_joints). - env_mask: Environment mask. If None, then all instances are updated. - joint_mask: Joint mask. If None, then all joints are updated. + root_velocity: Root frame velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - self.write_joint_position_to_sim_mask(position=joint_pos, env_mask=env_mask, joint_mask=joint_mask) - self.write_joint_velocity_to_sim_mask(velocity=joint_vel, env_mask=env_mask, joint_mask=joint_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + self.assert_shape_and_dtype(root_velocity, (self._num_instances,), wp.spatial_vectorf, "root_velocity") + wp.launch( + shared_kernels.set_root_link_velocity_to_sim_mask, + dim=self._num_instances, + inputs=[ + root_velocity, + self.data.body_com_pose_b, + self.data.root_link_pose_w, + env_mask_wp, + self._num_bodies, + ], + outputs=[self.data.root_link_vel_w, self.data.root_com_vel_w, self.data.body_com_acc_w], + device=self._device, + ) + binding = self._get_binding(TT.ROOT_VELOCITY) + binding.write(self.data._root_com_vel_w.data.view(wp.float32), mask=env_mask_wp) def write_joint_position_to_sim_index( self, *, position: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write joint positions over selected environment and joint indices into the simulation. + """Set joint positions over selected env / joint indices into the simulation. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - position: Joint positions. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + position: Joint positions [m or rad, depending on joint type]. Shape is + (len(env_ids), len(joint_ids)) with dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - d = len(joint_ids) if joint_ids is not None else self._num_joints - self.assert_shape_and_dtype(position, (n, d), wp.float32, "position") - self._write_flat_tensor(TT.DOF_POSITION, position, env_ids, joint_ids) - self.data._joint_pos_buf.timestamp = -1.0 + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + self.assert_shape_and_dtype(position, (env_ids.shape[0], joint_ids.shape[0]), wp.float32, "position") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[position, env_ids, joint_ids], + outputs=[self._data._joint_pos_buf.data], + device=self._device, + ) + # Invalidate body-state buffers so the next read re-fetches FK from the + # wheel using the new joint positions. Mirrors PhysX (articulation.py:894-903). + self._data._body_com_vel_w.timestamp = -1.0 + self._data._body_link_vel_w.timestamp = -1.0 + self._data._body_com_pose_b.timestamp = -1.0 + self._data._body_com_pose_w.timestamp = -1.0 + self._data._body_link_pose_w.timestamp = -1.0 + self._data._joint_acc.timestamp = -1.0 + binding = self._get_binding(TT.DOF_POSITION) + binding.write(self._data._joint_pos_buf.data, indices=env_ids) def write_joint_position_to_sim_mask( self, *, position: torch.Tensor | wp.array, - joint_mask: wp.array | None = None, env_mask: wp.array | None = None, + joint_mask: wp.array | None = None, ) -> None: - """Write joint positions over masked environments into the simulation. + """Set joint positions over selected env / joint masks into the simulation. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - position: Joint positions. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + position: Joint positions [m or rad, depending on joint type]. Shape is + (num_instances, num_joints) with dtype wp.float32. + env_mask: Environment mask. If None, all instances are updated. Shape is + (num_instances,). + joint_mask: Joint mask. If None, all joints are updated. Shape is + (num_joints,). """ + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) self.assert_shape_and_dtype(position, (self._num_instances, self._num_joints), wp.float32, "position") - self._write_flat_tensor_mask(TT.DOF_POSITION, position, env_mask, joint_mask) - self.data._joint_pos_buf.timestamp = -1.0 + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(self._num_instances, self._num_joints), + inputs=[position, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_pos_buf.data], + device=self._device, + ) + # See write_joint_position_to_sim_index for why these are invalidated. + self._data._body_com_vel_w.timestamp = -1.0 + self._data._body_link_vel_w.timestamp = -1.0 + self._data._body_com_pose_b.timestamp = -1.0 + self._data._body_com_pose_w.timestamp = -1.0 + self._data._body_link_pose_w.timestamp = -1.0 + self._data._joint_acc.timestamp = -1.0 + binding = self._get_binding(TT.DOF_POSITION) + binding.write(self._data._joint_pos_buf.data, mask=env_mask_wp) def write_joint_velocity_to_sim_index( self, *, velocity: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write joint velocities over selected environment and joint indices into the simulation. + """Set joint velocities over selected env / joint indices into the simulation. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - velocity: Joint velocities. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + velocity: Joint velocities [m/s or rad/s, depending on joint type]. Shape is + (len(env_ids), len(joint_ids)) with dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - d = len(joint_ids) if joint_ids is not None else self._num_joints - self.assert_shape_and_dtype(velocity, (n, d), wp.float32, "velocity") - self._write_flat_tensor(TT.DOF_VELOCITY, velocity, env_ids, joint_ids) - self.data._joint_vel_buf.timestamp = -1.0 + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + self.assert_shape_and_dtype(velocity, (env_ids.shape[0], joint_ids.shape[0]), wp.float32, "velocity") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[velocity, env_ids, joint_ids], + outputs=[self._data._joint_vel_buf.data], + device=self._device, + ) + # Sync previous_joint_vel to the new values so the next FD step does not + # produce a spurious acceleration spike. + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[velocity, env_ids, joint_ids], + outputs=[self._data._previous_joint_vel], + device=self._device, + ) + self._data._joint_acc.timestamp = -1.0 + binding = self._get_binding(TT.DOF_VELOCITY) + binding.write(self._data._joint_vel_buf.data, indices=env_ids) def write_joint_velocity_to_sim_mask( self, *, velocity: torch.Tensor | wp.array, - joint_mask: wp.array | None = None, env_mask: wp.array | None = None, + joint_mask: wp.array | None = None, ) -> None: - """Write joint velocities over masked environments into the simulation. + """Set joint velocities over selected env / joint masks into the simulation. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - velocity: Joint velocities. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + velocity: Joint velocities [m/s or rad/s, depending on joint type]. Shape is + (num_instances, num_joints) with dtype wp.float32. + env_mask: Environment mask. If None, all instances are updated. Shape is + (num_instances,). + joint_mask: Joint mask. If None, all joints are updated. Shape is + (num_joints,). """ + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) self.assert_shape_and_dtype(velocity, (self._num_instances, self._num_joints), wp.float32, "velocity") - self._write_flat_tensor_mask(TT.DOF_VELOCITY, velocity, env_mask, joint_mask) - self.data._joint_vel_buf.timestamp = -1.0 + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(self._num_instances, self._num_joints), + inputs=[velocity, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_vel_buf.data], + device=self._device, + ) + # Sync previous_joint_vel so the next FD step does not produce a spurious spike. + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(self._num_instances, self._num_joints), + inputs=[velocity, env_mask_wp, joint_mask_wp], + outputs=[self._data._previous_joint_vel], + device=self._device, + ) + self._data._joint_acc.timestamp = -1.0 + binding = self._get_binding(TT.DOF_VELOCITY) + binding.write(self._data._joint_vel_buf.data, mask=env_mask_wp) + + def write_joint_state_to_sim_mask( + self, + *, + position: torch.Tensor | wp.array, + velocity: torch.Tensor | wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write joint positions and velocities over selected environment mask into the simulation. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. + + Args: + position: Joint positions [m or rad, depending on joint type]. Shape is + (num_instances, num_joints) with dtype wp.float32. + velocity: Joint velocities [m/s or rad/s, depending on joint type]. Shape is + (num_instances, num_joints) with dtype wp.float32. + joint_mask: Joint mask. If None, all joints are updated. Shape is + (num_joints,). + env_mask: Environment mask. If None, all instances are updated. Shape is + (num_instances,). + """ + self.write_joint_position_to_sim_mask(position=position, env_mask=env_mask, joint_mask=joint_mask) + self.write_joint_velocity_to_sim_mask(velocity=velocity, env_mask=env_mask, joint_mask=joint_mask) """ Operations - Simulation Parameters Writers. @@ -596,99 +987,278 @@ def write_joint_velocity_to_sim_mask( def write_joint_stiffness_to_sim_index( self, *, - stiffness: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + stiffness: float | torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write joint stiffness over selected indices into the simulation. + """Set joint stiffness over selected env / joint indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``DOF_STIFFNESS`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - stiffness: Joint stiffness. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + stiffness: Joint stiffness [N/m or N·m/rad, depending on joint type]. + May be a scalar :class:`float` (broadcast), or shape + (len(env_ids), len(joint_ids)) with dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - d = len(joint_ids) if joint_ids is not None else self._num_joints - self.assert_shape_and_dtype(stiffness, (n, d), wp.float32, "stiffness") - self._write_flat_tensor(TT.DOF_STIFFNESS, stiffness, env_ids, joint_ids) + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + shape = (env_ids.shape[0], joint_ids.shape[0]) + stiffness = self._broadcast_scalar_to_2d(stiffness, shape) + self.assert_shape_and_dtype(stiffness, shape, wp.float32, "stiffness") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=shape, + inputs=[stiffness, env_ids, joint_ids], + outputs=[self._data._joint_stiffness.data], + device=self._device, + ) + cpu_env_ids = self._get_cpu_env_ids(env_ids) + wp.copy(self.data._cpu_joint_stiffness, self._data._joint_stiffness.data) + binding = self._get_binding(TT.DOF_STIFFNESS) + binding.write(self.data._cpu_joint_stiffness, indices=cpu_env_ids) def write_joint_stiffness_to_sim_mask( self, *, - stiffness: torch.Tensor | wp.array, + stiffness: float | torch.Tensor | wp.array, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Write joint stiffness over masked environments into the simulation. + """Set joint stiffness over selected env / joint masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``DOF_STIFFNESS`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - stiffness: Joint stiffness. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + stiffness: Joint stiffness [N/m or N·m/rad, depending on joint type]. + May be a scalar :class:`float` (broadcast), or shape + (num_instances, num_joints) with dtype wp.float32. + joint_mask: Joint mask. If None, all joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(stiffness, (self._num_instances, self._num_joints), wp.float32, "stiffness") - self._write_flat_tensor_mask(TT.DOF_STIFFNESS, stiffness, env_mask, joint_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + shape = (self._num_instances, self._num_joints) + stiffness = self._broadcast_scalar_to_2d(stiffness, shape) + self.assert_shape_and_dtype(stiffness, shape, wp.float32, "stiffness") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[stiffness, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_stiffness.data], + device=self._device, + ) + wp.copy(self.data._cpu_joint_stiffness, self._data._joint_stiffness.data) + binding = self._get_binding(TT.DOF_STIFFNESS) + binding.write(self.data._cpu_joint_stiffness, mask=self._get_cpu_env_mask(env_mask_wp)) def write_joint_damping_to_sim_index( self, *, - damping: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + damping: float | torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write joint damping over selected indices into the simulation. + """Set joint damping over selected env / joint indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``DOF_DAMPING`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - damping: Joint damping. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + damping: Joint damping [N·s/m or N·m·s/rad, depending on joint type]. + May be a scalar :class:`float` (broadcast), or shape + (len(env_ids), len(joint_ids)) with dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - d = len(joint_ids) if joint_ids is not None else self._num_joints - self.assert_shape_and_dtype(damping, (n, d), wp.float32, "damping") - self._write_flat_tensor(TT.DOF_DAMPING, damping, env_ids, joint_ids) + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + shape = (env_ids.shape[0], joint_ids.shape[0]) + damping = self._broadcast_scalar_to_2d(damping, shape) + self.assert_shape_and_dtype(damping, shape, wp.float32, "damping") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=shape, + inputs=[damping, env_ids, joint_ids], + outputs=[self._data._joint_damping.data], + device=self._device, + ) + cpu_env_ids = self._get_cpu_env_ids(env_ids) + wp.copy(self.data._cpu_joint_damping, self._data._joint_damping.data) + binding = self._get_binding(TT.DOF_DAMPING) + binding.write(self.data._cpu_joint_damping, indices=cpu_env_ids) def write_joint_damping_to_sim_mask( self, *, - damping: torch.Tensor | wp.array, + damping: float | torch.Tensor | wp.array, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Write joint damping over masked environments into the simulation. + """Set joint damping over selected env / joint masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``DOF_DAMPING`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - damping: Joint damping. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + damping: Joint damping [N·s/m or N·m·s/rad, depending on joint type]. + May be a scalar :class:`float` (broadcast), or shape + (num_instances, num_joints) with dtype wp.float32. + joint_mask: Joint mask. If None, all joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(damping, (self._num_instances, self._num_joints), wp.float32, "damping") - self._write_flat_tensor_mask(TT.DOF_DAMPING, damping, env_mask, joint_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + shape = (self._num_instances, self._num_joints) + damping = self._broadcast_scalar_to_2d(damping, shape) + self.assert_shape_and_dtype(damping, shape, wp.float32, "damping") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[damping, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_damping.data], + device=self._device, + ) + wp.copy(self.data._cpu_joint_damping, self._data._joint_damping.data) + binding = self._get_binding(TT.DOF_DAMPING) + binding.write(self.data._cpu_joint_damping, mask=self._get_cpu_env_mask(env_mask_wp)) def write_joint_position_limit_to_sim_index( self, *, limits: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, warn_limit_violation: bool = True, ) -> None: - """Write joint position limits over selected environment indices into the simulation. + """Set joint position limits over selected env / joint indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``DOF_LIMIT`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limits: Joint position limits [rad or m]. Shape is (len(env_ids), len(joint_ids)) - with dtype wp.vec2f. - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. - warn_limit_violation: Whether to use warning or info level logging when default joint - positions exceed the new limits. Defaults to True. + limits: Joint position limits ``[lower, upper]`` + [m or rad, depending on joint type]. Either shape + (len(env_ids), len(joint_ids), 2) with dtype wp.float32, or + shape (len(env_ids), len(joint_ids)) with dtype wp.vec2f. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). + warn_limit_violation: If True, log a warning when the provided limits + are inconsistent (lower > upper). Defaults to True. """ - if isinstance(limits, (int, float)): - raise ValueError("Float scalars are not supported for position limits (vec2f dtype)") - n = self._n_envs_index(env_ids) - d = len(joint_ids) if joint_ids is not None else self._num_joints - self.assert_shape_and_dtype(limits, (n, d), wp.vec2f, "limits") - self._write_flat_tensor(TT.DOF_LIMIT, limits, env_ids, joint_ids) + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + # Position limits cannot be scalar-broadcast (they pair lower/upper); + # match PhysX which explicitly rejects floats here. + if isinstance(limits, float): + raise ValueError("Joint position limits must be a tensor or array, not a float.") + # Accept both wp.vec2f shape (N, J) and the legacy (N, J, 2) wp.float32 + # form (canonical PhysX/Newton layout uses vec2f). + if isinstance(limits, wp.array) and limits.dtype == wp.vec2f: + self.assert_shape_and_dtype(limits, (env_ids.shape[0], joint_ids.shape[0]), wp.vec2f, "limits") + # Reinterpret the vec2f input as a (N, J, 2) float32 view for the kernel. + kernel_limits = wp.array( + ptr=limits.ptr, + shape=(env_ids.shape[0], joint_ids.shape[0], 2), + dtype=wp.float32, + device=str(limits.device), + copy=False, + ) + else: + self.assert_shape_and_dtype(limits, (env_ids.shape[0], joint_ids.shape[0], 2), wp.float32, "limits") + kernel_limits = limits + # Scatter [lower, upper] pairs into the vec2f cache buffer. + wp.launch( + shared_kernels.write_joint_position_limit_to_buffer_index, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[kernel_limits, env_ids, joint_ids], + outputs=[self._data._joint_pos_limits.data], + device=self._device, + ) + # Clamp default_joint_pos to the new limits and refresh soft_joint_pos_limits. + clamped_count = wp.zeros(1, dtype=wp.int32, device=self._device) + wp.launch( + clamp_default_joint_pos_and_update_soft_limits_index, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[ + self._data._joint_pos_limits.data, + env_ids, + joint_ids, + self.cfg.soft_joint_pos_limit_factor, + ], + outputs=[ + self._data._default_joint_pos, + self._data._soft_joint_pos_limits, + clamped_count, + ], + device=self._device, + ) + if clamped_count.numpy()[0] > 0: + violation_message = ( + "Some default joint positions are outside of the range of the new joint limits. Default joint" + " positions will be clamped to be within the new joint limits." + ) + if warn_limit_violation: + logger.warning(violation_message) + else: + logger.info(violation_message) + # Stage to pinned-host CPU: flatten the vec2f buffer to float32 view. + cpu_env_ids = self._get_cpu_env_ids(env_ids) + flat_src = wp.array( + ptr=self._data._joint_pos_limits.data.ptr, + shape=(self._num_instances, self._num_joints, 2), + dtype=wp.float32, + device=self._device, + copy=False, + ) + wp.copy(self.data._cpu_joint_position_limit, flat_src) + binding = self._get_binding(TT.DOF_LIMIT) + binding.write(self.data._cpu_joint_position_limit, indices=cpu_env_ids) def write_joint_position_limit_to_sim_mask( self, @@ -698,180 +1268,700 @@ def write_joint_position_limit_to_sim_mask( env_mask: wp.array | None = None, warn_limit_violation: bool = True, ) -> None: - """Write joint position limits over masked environments into the simulation. + """Set joint position limits over selected env / joint masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``DOF_LIMIT`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limits: Joint position limits [rad or m]. Shape is (num_instances, num_joints) - with dtype wp.vec2f. - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. - warn_limit_violation: Whether to use warning or info level logging when default joint - positions exceed the new limits. Defaults to True. + limits: Joint position limits ``[lower, upper]`` + [m or rad, depending on joint type]. Either shape + (num_instances, num_joints, 2) with dtype wp.float32, or shape + (num_instances, num_joints) with dtype wp.vec2f. + joint_mask: Joint mask. If None, all joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). + warn_limit_violation: If True, log a warning when the provided limits + are inconsistent (lower > upper). Defaults to True. """ - if isinstance(limits, (int, float)): - raise ValueError("Float scalars are not supported for position limits (vec2f dtype)") - self.assert_shape_and_dtype(limits, (self._num_instances, self._num_joints), wp.vec2f, "limits") - self._write_flat_tensor_mask(TT.DOF_LIMIT, limits, env_mask, joint_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + # Position limits cannot be scalar-broadcast (they pair lower/upper); + # match PhysX which explicitly rejects floats here. + if isinstance(limits, float): + raise ValueError("Joint position limits must be a tensor or array, not a float.") + # Accept both wp.vec2f shape (N, J) and the legacy (N, J, 2) wp.float32 + # form (canonical PhysX/Newton layout uses vec2f). + if isinstance(limits, wp.array) and limits.dtype == wp.vec2f: + self.assert_shape_and_dtype(limits, (self._num_instances, self._num_joints), wp.vec2f, "limits") + kernel_limits = wp.array( + ptr=limits.ptr, + shape=(self._num_instances, self._num_joints, 2), + dtype=wp.float32, + device=str(limits.device), + copy=False, + ) + else: + self.assert_shape_and_dtype(limits, (self._num_instances, self._num_joints, 2), wp.float32, "limits") + kernel_limits = limits + wp.launch( + shared_kernels.write_joint_position_limit_to_buffer_mask, + dim=(self._num_instances, self._num_joints), + inputs=[kernel_limits, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_pos_limits.data], + device=self._device, + ) + # Clamp default_joint_pos to the new limits and refresh soft_joint_pos_limits. + clamped_count = wp.zeros(1, dtype=wp.int32, device=self._device) + wp.launch( + clamp_default_joint_pos_and_update_soft_limits_mask, + dim=(self._num_instances, self._num_joints), + inputs=[ + self._data._joint_pos_limits.data, + env_mask_wp, + joint_mask_wp, + self.cfg.soft_joint_pos_limit_factor, + ], + outputs=[ + self._data._default_joint_pos, + self._data._soft_joint_pos_limits, + clamped_count, + ], + device=self._device, + ) + if clamped_count.numpy()[0] > 0: + violation_message = ( + "Some default joint positions are outside of the range of the new joint limits. Default joint" + " positions will be clamped to be within the new joint limits." + ) + if warn_limit_violation: + logger.warning(violation_message) + else: + logger.info(violation_message) + flat_src = wp.array( + ptr=self._data._joint_pos_limits.data.ptr, + shape=(self._num_instances, self._num_joints, 2), + dtype=wp.float32, + device=self._device, + copy=False, + ) + wp.copy(self.data._cpu_joint_position_limit, flat_src) + binding = self._get_binding(TT.DOF_LIMIT) + binding.write(self.data._cpu_joint_position_limit, mask=self._get_cpu_env_mask(env_mask_wp)) def write_joint_velocity_limit_to_sim_index( self, *, - limits: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + limits: float | torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write joint max velocity over selected environment indices into the simulation. + """Set joint velocity limits over selected env / joint indices into the simulation. - The velocity limit is used to constrain the joint velocities in the physics engine. + This is a CPU-only write routed through pinned-host staging because + ``DOF_MAX_VELOCITY`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limits: Joint max velocity [rad/s or m/s]. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + limits: Joint velocity limits [m/s or rad/s, depending on joint type]. + May be a scalar :class:`float` (broadcast), or shape + (len(env_ids), len(joint_ids)) with dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - d = len(joint_ids) if joint_ids is not None else self._num_joints - self.assert_shape_and_dtype(limits, (n, d), wp.float32, "limits") - self._write_flat_tensor(TT.DOF_MAX_VELOCITY, limits, env_ids, joint_ids) + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + shape = (env_ids.shape[0], joint_ids.shape[0]) + limits = self._broadcast_scalar_to_2d(limits, shape) + self.assert_shape_and_dtype(limits, shape, wp.float32, "limits") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=shape, + inputs=[limits, env_ids, joint_ids], + outputs=[self._data._joint_vel_limits.data], + device=self._device, + ) + cpu_env_ids = self._get_cpu_env_ids(env_ids) + wp.copy(self.data._cpu_joint_velocity_limit, self._data._joint_vel_limits.data) + binding = self._get_binding(TT.DOF_MAX_VELOCITY) + binding.write(self.data._cpu_joint_velocity_limit, indices=cpu_env_ids) def write_joint_velocity_limit_to_sim_mask( self, *, - limits: torch.Tensor | wp.array, + limits: float | torch.Tensor | wp.array, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Write joint max velocity over masked environments into the simulation. + """Set joint velocity limits over selected env / joint masks into the simulation. - The velocity limit is used to constrain the joint velocities in the physics engine. + This is a CPU-only write routed through pinned-host staging because + ``DOF_MAX_VELOCITY`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limits: Joint max velocity [rad/s or m/s]. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + limits: Joint velocity limits [m/s or rad/s, depending on joint type]. + May be a scalar :class:`float` (broadcast), or shape + (num_instances, num_joints) with dtype wp.float32. + joint_mask: Joint mask. If None, all joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(limits, (self._num_instances, self._num_joints), wp.float32, "limits") - self._write_flat_tensor_mask(TT.DOF_MAX_VELOCITY, limits, env_mask, joint_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + shape = (self._num_instances, self._num_joints) + limits = self._broadcast_scalar_to_2d(limits, shape) + self.assert_shape_and_dtype(limits, shape, wp.float32, "limits") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[limits, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_vel_limits.data], + device=self._device, + ) + wp.copy(self.data._cpu_joint_velocity_limit, self._data._joint_vel_limits.data) + binding = self._get_binding(TT.DOF_MAX_VELOCITY) + binding.write(self.data._cpu_joint_velocity_limit, mask=self._get_cpu_env_mask(env_mask_wp)) def write_joint_effort_limit_to_sim_index( self, *, - limits: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + limits: float | torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write joint effort limits over selected environment indices into the simulation. + """Set joint effort limits over selected env / joint indices into the simulation. - The effort limit is used to constrain the computed joint efforts in the physics engine. + This is a CPU-only write routed through pinned-host staging because + ``DOF_MAX_FORCE`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limits: Joint effort limits [N or N*m]. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + limits: Joint effort limits [N or N·m, depending on joint type]. + May be a scalar :class:`float` (broadcast), or shape + (len(env_ids), len(joint_ids)) with dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - d = len(joint_ids) if joint_ids is not None else self._num_joints - self.assert_shape_and_dtype(limits, (n, d), wp.float32, "limits") - self._write_flat_tensor(TT.DOF_MAX_FORCE, limits, env_ids, joint_ids) + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + shape = (env_ids.shape[0], joint_ids.shape[0]) + limits = self._broadcast_scalar_to_2d(limits, shape) + self.assert_shape_and_dtype(limits, shape, wp.float32, "limits") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=shape, + inputs=[limits, env_ids, joint_ids], + outputs=[self._data._joint_effort_limits.data], + device=self._device, + ) + cpu_env_ids = self._get_cpu_env_ids(env_ids) + wp.copy(self.data._cpu_joint_effort_limit, self._data._joint_effort_limits.data) + binding = self._get_binding(TT.DOF_MAX_FORCE) + binding.write(self.data._cpu_joint_effort_limit, indices=cpu_env_ids) def write_joint_effort_limit_to_sim_mask( self, *, - limits: torch.Tensor | wp.array, + limits: float | torch.Tensor | wp.array, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Write joint effort limits over masked environments into the simulation. + """Set joint effort limits over selected env / joint masks into the simulation. - The effort limit is used to constrain the computed joint efforts in the physics engine. + This is a CPU-only write routed through pinned-host staging because + ``DOF_MAX_FORCE`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limits: Joint effort limits [N or N*m]. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + limits: Joint effort limits [N or N·m, depending on joint type]. + May be a scalar :class:`float` (broadcast), or shape + (num_instances, num_joints) with dtype wp.float32. + joint_mask: Joint mask. If None, all joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(limits, (self._num_instances, self._num_joints), wp.float32, "limits") - self._write_flat_tensor_mask(TT.DOF_MAX_FORCE, limits, env_mask, joint_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + shape = (self._num_instances, self._num_joints) + limits = self._broadcast_scalar_to_2d(limits, shape) + self.assert_shape_and_dtype(limits, shape, wp.float32, "limits") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[limits, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_effort_limits.data], + device=self._device, + ) + wp.copy(self.data._cpu_joint_effort_limit, self._data._joint_effort_limits.data) + binding = self._get_binding(TT.DOF_MAX_FORCE) + binding.write(self.data._cpu_joint_effort_limit, mask=self._get_cpu_env_mask(env_mask_wp)) def write_joint_armature_to_sim_index( self, *, - armature: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + armature: float | torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write joint armature over selected environment indices into the simulation. + """Set joint armature over selected env / joint indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``DOF_ARMATURE`` is a CPU-only OVPhysX binding. - The armature is directly added to the corresponding joint-space inertia. It helps improve the - simulation stability by reducing the joint velocities. + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - armature: Joint armature [kg*m^2 or kg]. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + armature: Joint armature [kg·m²]. May be a scalar :class:`float` + (broadcast), or shape (len(env_ids), len(joint_ids)) with + dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - d = len(joint_ids) if joint_ids is not None else self._num_joints - self.assert_shape_and_dtype(armature, (n, d), wp.float32, "armature") - self._write_flat_tensor(TT.DOF_ARMATURE, armature, env_ids, joint_ids) + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + shape = (env_ids.shape[0], joint_ids.shape[0]) + armature = self._broadcast_scalar_to_2d(armature, shape) + self.assert_shape_and_dtype(armature, shape, wp.float32, "armature") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=shape, + inputs=[armature, env_ids, joint_ids], + outputs=[self._data._joint_armature.data], + device=self._device, + ) + cpu_env_ids = self._get_cpu_env_ids(env_ids) + wp.copy(self.data._cpu_joint_armature, self._data._joint_armature.data) + binding = self._get_binding(TT.DOF_ARMATURE) + binding.write(self.data._cpu_joint_armature, indices=cpu_env_ids) def write_joint_armature_to_sim_mask( self, *, - armature: torch.Tensor | wp.array, + armature: float | torch.Tensor | wp.array, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Write joint armature over masked environments into the simulation. + """Set joint armature over selected env / joint masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``DOF_ARMATURE`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. - The armature is directly added to the corresponding joint-space inertia. It helps improve the - simulation stability by reducing the joint velocities. + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - armature: Joint armature [kg*m^2 or kg]. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + armature: Joint armature [kg·m²]. May be a scalar :class:`float` + (broadcast), or shape (num_instances, num_joints) with dtype + wp.float32. + joint_mask: Joint mask. If None, all joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(armature, (self._num_instances, self._num_joints), wp.float32, "armature") - self._write_flat_tensor_mask(TT.DOF_ARMATURE, armature, env_mask, joint_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + shape = (self._num_instances, self._num_joints) + armature = self._broadcast_scalar_to_2d(armature, shape) + self.assert_shape_and_dtype(armature, shape, wp.float32, "armature") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[armature, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_armature.data], + device=self._device, + ) + wp.copy(self.data._cpu_joint_armature, self._data._joint_armature.data) + binding = self._get_binding(TT.DOF_ARMATURE) + binding.write(self.data._cpu_joint_armature, mask=self._get_cpu_env_mask(env_mask_wp)) def write_joint_friction_coefficient_to_sim_index( self, *, - joint_friction_coeff: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + joint_friction_coeff: float | torch.Tensor | wp.array, + joint_dynamic_friction_coeff: float | torch.Tensor | wp.array | None = None, + joint_viscous_friction_coeff: float | torch.Tensor | wp.array | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write joint friction coefficients over selected indices into the simulation. + r"""Write joint friction coefficients over selected env / joint indices into the simulation. + + Mirrors :meth:`isaaclab_physx.assets.Articulation.write_joint_friction_coefficient_to_sim_index`: + Coulomb (static & dynamic) friction with an optional viscous term. Any of the three + components can be left unset by passing ``None``; the corresponding slot in the + combined ``DOF_FRICTION_PROPERTIES`` ``(N, J, 3)`` binding is preserved. + + ``DOF_FRICTION_PROPERTIES`` is a CPU-only OVPhysX binding, so the + write is routed through pinned-host staging. + + .. note:: + This method expects partial data. Each component, if provided, + may be a scalar :class:`float` (broadcast to + ``(len(env_ids), len(joint_ids))``) or a 2D tensor / warp array. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - joint_friction_coeff: Joint friction coefficients. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + joint_friction_coeff: Static friction coefficient :math:`\mu_s` [dimensionless]. + joint_dynamic_friction_coeff: Dynamic (Coulomb) friction coefficient + :math:`\mu_d`. If ``None``, the dynamic component is preserved. + joint_viscous_friction_coeff: Viscous friction coefficient :math:`c_v`. + If ``None``, the viscous component is preserved. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - d = len(joint_ids) if joint_ids is not None else self._num_joints - self.assert_shape_and_dtype(joint_friction_coeff, (n, d), wp.float32, "joint_friction_coeff") - self._write_friction_column(joint_friction_coeff, env_ids, joint_ids) + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + shape = (env_ids.shape[0], joint_ids.shape[0]) + joint_friction_coeff = self._broadcast_scalar_to_2d(joint_friction_coeff, shape) + if joint_dynamic_friction_coeff is not None: + joint_dynamic_friction_coeff = self._broadcast_scalar_to_2d(joint_dynamic_friction_coeff, shape) + if joint_viscous_friction_coeff is not None: + joint_viscous_friction_coeff = self._broadcast_scalar_to_2d(joint_viscous_friction_coeff, shape) + self.assert_shape_and_dtype(joint_friction_coeff, shape, wp.float32, "joint_friction_coeff") + if joint_dynamic_friction_coeff is not None: + self.assert_shape_and_dtype(joint_dynamic_friction_coeff, shape, wp.float32, "joint_dynamic_friction_coeff") + if joint_viscous_friction_coeff is not None: + self.assert_shape_and_dtype(joint_viscous_friction_coeff, shape, wp.float32, "joint_viscous_friction_coeff") + # Refresh the combined (N, J, 3) buffer from the wheel so unchanged components + # are preserved on the round-trip. + self._data._read_scalar_binding(TT.DOF_FRICTION_PROPERTIES, self._data._joint_friction_props_buf) + wp.launch( + write_joint_friction_data_to_buffer_index, + dim=shape, + inputs=[ + joint_friction_coeff, + joint_dynamic_friction_coeff, + joint_viscous_friction_coeff, + env_ids, + joint_ids, + ], + outputs=[self._data._joint_friction_props_buf.data], + device=self._device, + ) + # Stage the combined (N, J, 3) buffer to pinned-host CPU and write to the binding. + cpu_env_ids = self._get_cpu_env_ids(env_ids) + cpu_friction = self._data._stage_to_pinned_cpu( + TT.DOF_FRICTION_PROPERTIES, "write", self._data._joint_friction_props_buf.data + ) + binding = self._get_binding(TT.DOF_FRICTION_PROPERTIES) + binding.write(cpu_friction, indices=cpu_env_ids) def write_joint_friction_coefficient_to_sim_mask( self, *, - joint_friction_coeff: torch.Tensor | wp.array, + joint_friction_coeff: float | torch.Tensor | wp.array, + joint_dynamic_friction_coeff: float | torch.Tensor | wp.array | None = None, + joint_viscous_friction_coeff: float | torch.Tensor | wp.array | None = None, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Write joint friction coefficients over masked environments into the simulation. + r"""Mask variant of :meth:`write_joint_friction_coefficient_to_sim_index`. Args: - joint_friction_coeff: Joint friction coefficients. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + joint_friction_coeff: Static friction coefficient :math:`\mu_s`. Full data, + shape ``(num_instances, num_joints)``. May be a scalar :class:`float`. + joint_dynamic_friction_coeff: Dynamic friction. ``None`` to preserve. + joint_viscous_friction_coeff: Viscous friction. ``None`` to preserve. + joint_mask: Joint mask. If None, all joints are updated. + env_mask: Environment mask. If None, all instances are updated. """ - self.assert_shape_and_dtype( - joint_friction_coeff, (self._num_instances, self._num_joints), wp.float32, "joint_friction_coeff" + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + shape = (self._num_instances, self._num_joints) + joint_friction_coeff = self._broadcast_scalar_to_2d(joint_friction_coeff, shape) + if joint_dynamic_friction_coeff is not None: + joint_dynamic_friction_coeff = self._broadcast_scalar_to_2d(joint_dynamic_friction_coeff, shape) + if joint_viscous_friction_coeff is not None: + joint_viscous_friction_coeff = self._broadcast_scalar_to_2d(joint_viscous_friction_coeff, shape) + self.assert_shape_and_dtype(joint_friction_coeff, shape, wp.float32, "joint_friction_coeff") + if joint_dynamic_friction_coeff is not None: + self.assert_shape_and_dtype(joint_dynamic_friction_coeff, shape, wp.float32, "joint_dynamic_friction_coeff") + if joint_viscous_friction_coeff is not None: + self.assert_shape_and_dtype(joint_viscous_friction_coeff, shape, wp.float32, "joint_viscous_friction_coeff") + # See _index variant for why we refresh the (N, J, 3) buffer from the wheel. + self._data._read_scalar_binding(TT.DOF_FRICTION_PROPERTIES, self._data._joint_friction_props_buf) + wp.launch( + write_joint_friction_data_to_buffer_mask, + dim=shape, + inputs=[ + joint_friction_coeff, + joint_dynamic_friction_coeff, + joint_viscous_friction_coeff, + env_mask_wp, + joint_mask_wp, + ], + outputs=[self._data._joint_friction_props_buf.data], + device=self._device, + ) + cpu_friction = self._data._stage_to_pinned_cpu( + TT.DOF_FRICTION_PROPERTIES, "write", self._data._joint_friction_props_buf.data + ) + binding = self._get_binding(TT.DOF_FRICTION_PROPERTIES) + binding.write(cpu_friction, mask=self._get_cpu_env_mask(env_mask_wp)) + + def write_joint_dynamic_friction_coefficient_to_sim_index( + self, + *, + joint_dynamic_friction_coeff: float | torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + r"""Write joint dynamic friction coefficients over selected env / joint indices into the simulation. + + Mirrors :meth:`isaaclab_physx.assets.Articulation.write_joint_dynamic_friction_coefficient_to_sim_index`: + updates only the dynamic (Coulomb) slot of the combined ``DOF_FRICTION_PROPERTIES`` ``(N, J, 3)`` + binding; the static and viscous components are preserved. + + ``DOF_FRICTION_PROPERTIES`` is a CPU-only OVPhysX binding, so the + write is routed through pinned-host staging. + + .. note:: + This method expects partial data. ``joint_dynamic_friction_coeff`` may be a + scalar :class:`float` (broadcast to ``(len(env_ids), len(joint_ids))``) or a + 2D tensor / warp array. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. + + Args: + joint_dynamic_friction_coeff: Dynamic (Coulomb) friction coefficient + :math:`\mu_d` [dimensionless]. Shape is ``(len(env_ids), len(joint_ids))`` + with dtype wp.float32, or a scalar that is broadcast. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). + """ + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + shape = (env_ids.shape[0], joint_ids.shape[0]) + joint_dynamic_friction_coeff = self._broadcast_scalar_to_2d(joint_dynamic_friction_coeff, shape) + self.assert_shape_and_dtype(joint_dynamic_friction_coeff, shape, wp.float32, "joint_dynamic_friction_coeff") + # Refresh the combined (N, J, 3) buffer from the wheel so unchanged components + # are preserved on the round-trip. + self._data._read_scalar_binding(TT.DOF_FRICTION_PROPERTIES, self._data._joint_friction_props_buf) + wp.launch( + write_joint_friction_data_to_buffer_index, + dim=shape, + inputs=[ + None, # in_static — preserved + joint_dynamic_friction_coeff, + None, # in_viscous — preserved + env_ids, + joint_ids, + ], + outputs=[self._data._joint_friction_props_buf.data], + device=self._device, + ) + cpu_env_ids = self._get_cpu_env_ids(env_ids) + cpu_friction = self._data._stage_to_pinned_cpu( + TT.DOF_FRICTION_PROPERTIES, "write", self._data._joint_friction_props_buf.data + ) + binding = self._get_binding(TT.DOF_FRICTION_PROPERTIES) + binding.write(cpu_friction, indices=cpu_env_ids) + + def write_joint_dynamic_friction_coefficient_to_sim_mask( + self, + *, + joint_dynamic_friction_coeff: float | torch.Tensor | wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + r"""Mask variant of :meth:`write_joint_dynamic_friction_coefficient_to_sim_index`. + + Updates only the dynamic (Coulomb) slot of the combined ``DOF_FRICTION_PROPERTIES`` + ``(N, J, 3)`` binding; the static and viscous components are preserved. + + Args: + joint_dynamic_friction_coeff: Dynamic (Coulomb) friction coefficient + :math:`\mu_d` [dimensionless]. Full data, shape + ``(num_instances, num_joints)``. May be a scalar :class:`float`. + joint_mask: Joint mask. If None, all joints are updated. + env_mask: Environment mask. If None, all instances are updated. + """ + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + shape = (self._num_instances, self._num_joints) + joint_dynamic_friction_coeff = self._broadcast_scalar_to_2d(joint_dynamic_friction_coeff, shape) + self.assert_shape_and_dtype(joint_dynamic_friction_coeff, shape, wp.float32, "joint_dynamic_friction_coeff") + # See _index variant for why we refresh the (N, J, 3) buffer from the wheel. + self._data._read_scalar_binding(TT.DOF_FRICTION_PROPERTIES, self._data._joint_friction_props_buf) + wp.launch( + write_joint_friction_data_to_buffer_mask, + dim=shape, + inputs=[ + None, # in_static — preserved + joint_dynamic_friction_coeff, + None, # in_viscous — preserved + env_mask_wp, + joint_mask_wp, + ], + outputs=[self._data._joint_friction_props_buf.data], + device=self._device, + ) + cpu_friction = self._data._stage_to_pinned_cpu( + TT.DOF_FRICTION_PROPERTIES, "write", self._data._joint_friction_props_buf.data + ) + binding = self._get_binding(TT.DOF_FRICTION_PROPERTIES) + binding.write(cpu_friction, mask=self._get_cpu_env_mask(env_mask_wp)) + + def write_joint_viscous_friction_coefficient_to_sim_index( + self, + *, + joint_viscous_friction_coeff: float | torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + r"""Write joint viscous friction coefficients over selected env / joint indices into the simulation. + + Mirrors :meth:`isaaclab_physx.assets.Articulation.write_joint_viscous_friction_coefficient_to_sim_index`: + updates only the viscous slot of the combined ``DOF_FRICTION_PROPERTIES`` ``(N, J, 3)`` + binding; the static and dynamic components are preserved. + + ``DOF_FRICTION_PROPERTIES`` is a CPU-only OVPhysX binding, so the + write is routed through pinned-host staging. + + .. note:: + This method expects partial data. ``joint_viscous_friction_coeff`` may be a + scalar :class:`float` (broadcast to ``(len(env_ids), len(joint_ids))``) or a + 2D tensor / warp array. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. + + Args: + joint_viscous_friction_coeff: Viscous friction coefficient + :math:`c_v` [N·m·s/rad or N·s/m, depending on joint type]. + Shape is ``(len(env_ids), len(joint_ids))`` with dtype wp.float32, or + a scalar that is broadcast. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). + """ + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + shape = (env_ids.shape[0], joint_ids.shape[0]) + joint_viscous_friction_coeff = self._broadcast_scalar_to_2d(joint_viscous_friction_coeff, shape) + self.assert_shape_and_dtype(joint_viscous_friction_coeff, shape, wp.float32, "joint_viscous_friction_coeff") + # Refresh the combined (N, J, 3) buffer from the wheel so unchanged components + # are preserved on the round-trip. + self._data._read_scalar_binding(TT.DOF_FRICTION_PROPERTIES, self._data._joint_friction_props_buf) + wp.launch( + write_joint_friction_data_to_buffer_index, + dim=shape, + inputs=[ + None, # in_static — preserved + None, # in_dynamic — preserved + joint_viscous_friction_coeff, + env_ids, + joint_ids, + ], + outputs=[self._data._joint_friction_props_buf.data], + device=self._device, + ) + cpu_env_ids = self._get_cpu_env_ids(env_ids) + cpu_friction = self._data._stage_to_pinned_cpu( + TT.DOF_FRICTION_PROPERTIES, "write", self._data._joint_friction_props_buf.data + ) + binding = self._get_binding(TT.DOF_FRICTION_PROPERTIES) + binding.write(cpu_friction, indices=cpu_env_ids) + + def write_joint_viscous_friction_coefficient_to_sim_mask( + self, + *, + joint_viscous_friction_coeff: float | torch.Tensor | wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + r"""Mask variant of :meth:`write_joint_viscous_friction_coefficient_to_sim_index`. + + Updates only the viscous slot of the combined ``DOF_FRICTION_PROPERTIES`` + ``(N, J, 3)`` binding; the static and dynamic components are preserved. + + Args: + joint_viscous_friction_coeff: Viscous friction coefficient + :math:`c_v` [N·m·s/rad or N·s/m, depending on joint type]. + Full data, shape ``(num_instances, num_joints)``. May be a + scalar :class:`float`. + joint_mask: Joint mask. If None, all joints are updated. + env_mask: Environment mask. If None, all instances are updated. + """ + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + shape = (self._num_instances, self._num_joints) + joint_viscous_friction_coeff = self._broadcast_scalar_to_2d(joint_viscous_friction_coeff, shape) + self.assert_shape_and_dtype(joint_viscous_friction_coeff, shape, wp.float32, "joint_viscous_friction_coeff") + # See _index variant for why we refresh the (N, J, 3) buffer from the wheel. + self._data._read_scalar_binding(TT.DOF_FRICTION_PROPERTIES, self._data._joint_friction_props_buf) + wp.launch( + write_joint_friction_data_to_buffer_mask, + dim=shape, + inputs=[ + None, # in_static — preserved + None, # in_dynamic — preserved + joint_viscous_friction_coeff, + env_mask_wp, + joint_mask_wp, + ], + outputs=[self._data._joint_friction_props_buf.data], + device=self._device, + ) + cpu_friction = self._data._stage_to_pinned_cpu( + TT.DOF_FRICTION_PROPERTIES, "write", self._data._joint_friction_props_buf.data ) - self._write_friction_column_mask(joint_friction_coeff, env_mask, joint_mask) + binding = self._get_binding(TT.DOF_FRICTION_PROPERTIES) + binding.write(cpu_friction, mask=self._get_cpu_env_mask(env_mask_wp)) """ Operations - Setters. @@ -881,20 +1971,42 @@ def set_masses_index( self, *, masses: torch.Tensor | wp.array, - body_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set masses of all bodies using indices. + """Set body masses over selected env / body indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``BODY_MASS`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized + implementations. Performance is similar for both. However, to + allow graphed pipelines, the mask method must be used. Args: - masses: Masses of all bodies [kg]. Shape is (len(env_ids), len(body_ids)). - body_ids: The body indices to set the masses for. Defaults to None (all bodies). - env_ids: The environment indices to set the masses for. Defaults to None (all environments). + masses: Body masses [kg]. Shape is (len(env_ids), len(body_ids)) + with dtype wp.float32. + body_ids: Body indices. Defaults to None (all bodies). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - b = len(body_ids) if body_ids is not None else self._num_bodies - self.assert_shape_and_dtype(masses, (n, b), wp.float32, "masses") - self._write_flat_tensor(TT.BODY_MASS, masses, env_ids, body_ids) + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + self.assert_shape_and_dtype(masses, (env_ids.shape[0], body_ids.shape[0]), wp.float32, "masses") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[masses, env_ids, body_ids], + outputs=[self._data._body_mass.data], + device=self._device, + ) + cpu_env_ids = self._get_cpu_env_ids(env_ids) + wp.copy(self.data._cpu_body_mass, self._data._body_mass.data) + binding = self._get_binding(TT.BODY_MASS) + binding.write(self.data._cpu_body_mass, indices=cpu_env_ids) def set_masses_mask( self, @@ -903,36 +2015,84 @@ def set_masses_mask( body_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set masses of all bodies using masks. + """Set body masses over selected env / body masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``BODY_MASS`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized + implementations. Performance is similar for both. However, to + allow graphed pipelines, the mask method must be used. Args: - masses: Masses of all bodies [kg]. Shape is (num_instances, num_bodies). - body_mask: Body mask. If None, then all bodies are used. - env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + masses: Body masses [kg]. Shape is (num_instances, num_bodies) + with dtype wp.float32. + body_mask: Body mask. If None, all bodies are updated. + Shape is (num_bodies,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ + env_mask_wp = self._resolve_env_mask(env_mask) + body_mask_wp = self._resolve_body_mask(body_mask) self.assert_shape_and_dtype(masses, (self._num_instances, self._num_bodies), wp.float32, "masses") - self._write_flat_tensor_mask(TT.BODY_MASS, masses, env_mask, body_mask) + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(self._num_instances, self._num_bodies), + inputs=[masses, env_mask_wp, body_mask_wp], + outputs=[self._data._body_mass.data], + device=self._device, + ) + wp.copy(self.data._cpu_body_mass, self._data._body_mass.data) + binding = self._get_binding(TT.BODY_MASS) + binding.write(self.data._cpu_body_mass, mask=self._get_cpu_env_mask(env_mask_wp)) def set_coms_index( self, *, coms: torch.Tensor | wp.array, - body_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set center of mass pose of all bodies using indices. + """Set body center-of-mass poses over selected env / body indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``BODY_COM_POSE`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized + implementations. Performance is similar for both. However, to + allow graphed pipelines, the mask method must be used. Args: - coms: Center of mass pose of all bodies. Shape is (len(env_ids), len(body_ids)) - with dtype wp.transformf. - body_ids: The body indices to set the center of mass pose for. Defaults to None (all bodies). - env_ids: The environment indices to set the center of mass pose for. - Defaults to None (all environments). + coms: Body center-of-mass poses [m, quaternion (w, x, y, z)]. + Shape is (len(env_ids), len(body_ids)) with dtype wp.transformf. + body_ids: Body indices. Defaults to None (all bodies). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - b = len(body_ids) if body_ids is not None else self._num_bodies - self.assert_shape_and_dtype(coms, (n, b), wp.transformf, "coms") - self._write_flat_tensor(TT.BODY_COM_POSE, coms, env_ids, body_ids) + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + self.assert_shape_and_dtype(coms, (env_ids.shape[0], body_ids.shape[0]), wp.transformf, "coms") + wp.launch( + shared_kernels.write_body_com_pose_to_buffer_index, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[coms, env_ids, body_ids], + outputs=[self._data._body_com_pose_b.data], + device=self._device, + ) + # Invalidate derived buffers that depend on body_com_pose_b. + self.data._root_com_pose_w.timestamp = -1.0 + self.data._body_com_pose_w.timestamp = -1.0 + cpu_env_ids = self._get_cpu_env_ids(env_ids) + wp.copy(self.data._cpu_body_coms, self._data._body_com_pose_b.data) + binding = self._get_binding(TT.BODY_COM_POSE) + binding.write(self.data._cpu_body_coms, indices=cpu_env_ids) def set_coms_mask( self, @@ -941,35 +2101,84 @@ def set_coms_mask( body_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set center of mass pose of all bodies using masks. + """Set body center-of-mass poses over selected env / body masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``BODY_COM_POSE`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized + implementations. Performance is similar for both. However, to + allow graphed pipelines, the mask method must be used. Args: - coms: Center of mass pose of all bodies. Shape is (num_instances, num_bodies) - with dtype wp.transformf. - body_mask: Body mask. If None, then all bodies are used. - env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + coms: Body center-of-mass poses [m, quaternion (w, x, y, z)]. + Shape is (num_instances, num_bodies) with dtype wp.transformf. + body_mask: Body mask. If None, all bodies are updated. + Shape is (num_bodies,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ + env_mask_wp = self._resolve_env_mask(env_mask) + body_mask_wp = self._resolve_body_mask(body_mask) self.assert_shape_and_dtype(coms, (self._num_instances, self._num_bodies), wp.transformf, "coms") - self._write_flat_tensor_mask(TT.BODY_COM_POSE, coms, env_mask, body_mask) + wp.launch( + shared_kernels.write_body_com_pose_to_buffer_mask, + dim=(self._num_instances, self._num_bodies), + inputs=[coms, env_mask_wp, body_mask_wp], + outputs=[self._data._body_com_pose_b.data], + device=self._device, + ) + # Invalidate derived buffers that depend on body_com_pose_b. + self.data._root_com_pose_w.timestamp = -1.0 + self.data._body_com_pose_w.timestamp = -1.0 + wp.copy(self.data._cpu_body_coms, self._data._body_com_pose_b.data) + binding = self._get_binding(TT.BODY_COM_POSE) + binding.write(self.data._cpu_body_coms, mask=self._get_cpu_env_mask(env_mask_wp)) def set_inertias_index( self, *, inertias: torch.Tensor | wp.array, - body_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set inertias of all bodies using indices. + """Set body inertia tensors over selected env / body indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``BODY_INERTIA`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized + implementations. Performance is similar for both. However, to + allow graphed pipelines, the mask method must be used. Args: - inertias: Inertias of all bodies [kg*m^2]. Shape is (len(env_ids), len(body_ids), 9). - body_ids: The body indices to set the inertias for. Defaults to None (all bodies). - env_ids: The environment indices to set the inertias for. Defaults to None (all environments). + inertias: Body inertia tensors [kg·m²]. Shape is + (len(env_ids), len(body_ids), 9) with dtype wp.float32. + body_ids: Body indices. Defaults to None (all bodies). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - b = len(body_ids) if body_ids is not None else self._num_bodies - self.assert_shape_and_dtype(inertias, (n, b, 9), wp.float32, "inertias") - self._write_flat_tensor(TT.BODY_INERTIA, inertias, env_ids, body_ids) + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + self.assert_shape_and_dtype(inertias, (env_ids.shape[0], body_ids.shape[0], 9), wp.float32, "inertias") + wp.launch( + shared_kernels.write_body_inertia_to_buffer_index, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[inertias, env_ids, body_ids], + outputs=[self._data._body_inertia.data], + device=self._device, + ) + cpu_env_ids = self._get_cpu_env_ids(env_ids) + wp.copy(self.data._cpu_body_inertia, self._data._body_inertia.data) + binding = self._get_binding(TT.BODY_INERTIA) + binding.write(self.data._cpu_body_inertia, indices=cpu_env_ids) def set_inertias_mask( self, @@ -978,38 +2187,80 @@ def set_inertias_mask( body_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set inertias of all bodies using masks. + """Set body inertia tensors over selected env / body masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``BODY_INERTIA`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized + implementations. Performance is similar for both. However, to + allow graphed pipelines, the mask method must be used. Args: - inertias: Inertias of all bodies [kg*m^2]. Shape is (num_instances, num_bodies, 9). - body_mask: Body mask. If None, then all bodies are used. - env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + inertias: Body inertia tensors [kg·m²]. Shape is + (num_instances, num_bodies, 9) with dtype wp.float32. + body_mask: Body mask. If None, all bodies are updated. + Shape is (num_bodies,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ + env_mask_wp = self._resolve_env_mask(env_mask) + body_mask_wp = self._resolve_body_mask(body_mask) self.assert_shape_and_dtype(inertias, (self._num_instances, self._num_bodies, 9), wp.float32, "inertias") - self._write_flat_tensor_mask(TT.BODY_INERTIA, inertias, env_mask, body_mask) - - """ - Operations - Target Setters. - """ + wp.launch( + shared_kernels.write_body_inertia_to_buffer_mask, + dim=(self._num_instances, self._num_bodies), + inputs=[inertias, env_mask_wp, body_mask_wp], + outputs=[self._data._body_inertia.data], + device=self._device, + ) + wp.copy(self.data._cpu_body_inertia, self._data._body_inertia.data) + binding = self._get_binding(TT.BODY_INERTIA) + binding.write(self.data._cpu_body_inertia, mask=self._get_cpu_env_mask(env_mask_wp)) def set_joint_position_target_index( self, *, target: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set joint position targets into internal buffers using indices. - This function does not apply the joint targets to the simulation. It only fills the buffers with - the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. + This function does not apply the joint targets to the simulation. It only fills the + buffers with the desired values. To apply the joint targets, call + :meth:`write_data_to_sim`. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - target: Joint position targets [rad or m]. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + target: Joint position targets [m or rad, depending on joint type]. Shape is + (len(env_ids), len(joint_ids)) with dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - self._set_target_into_buffer(self._data._joint_pos_target, target, env_ids, joint_ids) + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + self.assert_shape_and_dtype(target, (env_ids.shape[0], joint_ids.shape[0]), wp.float32, "target") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[target, env_ids, joint_ids], + outputs=[self._data._joint_pos_target], + device=self._device, + ) + binding = self._get_binding(TT.DOF_POSITION_TARGET) + binding.write(self._data._joint_pos_target, indices=env_ids) def set_joint_position_target_mask( self, @@ -1020,31 +2271,73 @@ def set_joint_position_target_mask( ) -> None: """Set joint position targets into internal buffers using masks. + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. + Args: - target: Joint position targets [rad or m]. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + target: Joint position targets [m or rad, depending on joint type]. Shape is + (num_instances, num_joints) with dtype wp.float32. + joint_mask: Joint mask. If None, all joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, all instances are updated. Shape is + (num_instances,). """ - self._set_target_into_buffer_mask(self._data._joint_pos_target, target, env_mask, joint_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + self.assert_shape_and_dtype(target, (self._num_instances, self._num_joints), wp.float32, "target") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(self._num_instances, self._num_joints), + inputs=[target, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_pos_target], + device=self._device, + ) + binding = self._get_binding(TT.DOF_POSITION_TARGET) + binding.write(self._data._joint_pos_target, mask=env_mask_wp) def set_joint_velocity_target_index( self, *, target: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set joint velocity targets into internal buffers using indices. - This function does not apply the joint targets to the simulation. It only fills the buffers with - the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. + This function does not apply the joint targets to the simulation. It only fills the + buffers with the desired values. To apply the joint targets, call + :meth:`write_data_to_sim`. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - target: Joint velocity targets [rad/s or m/s]. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + target: Joint velocity targets [m/s or rad/s, depending on joint type]. Shape is + (len(env_ids), len(joint_ids)) with dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - self._set_target_into_buffer(self._data._joint_vel_target, target, env_ids, joint_ids) + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + self.assert_shape_and_dtype(target, (env_ids.shape[0], joint_ids.shape[0]), wp.float32, "target") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[target, env_ids, joint_ids], + outputs=[self._data._joint_vel_target], + device=self._device, + ) + binding = self._get_binding(TT.DOF_VELOCITY_TARGET) + binding.write(self._data._joint_vel_target, indices=env_ids) def set_joint_velocity_target_mask( self, @@ -1055,31 +2348,73 @@ def set_joint_velocity_target_mask( ) -> None: """Set joint velocity targets into internal buffers using masks. + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. + Args: - target: Joint velocity targets [rad/s or m/s]. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + target: Joint velocity targets [m/s or rad/s, depending on joint type]. Shape is + (num_instances, num_joints) with dtype wp.float32. + joint_mask: Joint mask. If None, all joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, all instances are updated. Shape is + (num_instances,). """ - self._set_target_into_buffer_mask(self._data._joint_vel_target, target, env_mask, joint_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + self.assert_shape_and_dtype(target, (self._num_instances, self._num_joints), wp.float32, "target") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(self._num_instances, self._num_joints), + inputs=[target, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_vel_target], + device=self._device, + ) + binding = self._get_binding(TT.DOF_VELOCITY_TARGET) + binding.write(self._data._joint_vel_target, mask=env_mask_wp) def set_joint_effort_target_index( self, *, target: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set joint efforts into internal buffers using indices. + """Set joint effort targets into internal buffers using indices. + + This function does not apply the joint targets to the simulation. It only fills the + buffers with the desired values. To apply the joint targets, call + :meth:`write_data_to_sim`. + + .. note:: + This method expects partial data. - This function does not apply the joint targets to the simulation. It only fills the buffers with - the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - target: Joint effort targets [N or N*m]. Shape is (len(env_ids), len(joint_ids)). - joint_ids: Joint indices. If None, then all joints are used. - env_ids: Environment indices. If None, then all indices are used. + target: Joint effort targets [N or N·m, depending on joint type]. Shape is + (len(env_ids), len(joint_ids)) with dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). """ - self._set_target_into_buffer(self._data._joint_effort_target, target, env_ids, joint_ids) + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + self.assert_shape_and_dtype(target, (env_ids.shape[0], joint_ids.shape[0]), wp.float32, "target") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[target, env_ids, joint_ids], + outputs=[self._data._joint_effort_target], + device=self._device, + ) + binding = self._get_binding(TT.DOF_ACTUATION_FORCE) + binding.write(self._data._joint_effort_target, indices=env_ids) def set_joint_effort_target_mask( self, @@ -1088,14 +2423,35 @@ def set_joint_effort_target_mask( joint_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set joint efforts into internal buffers using masks. + """Set joint effort targets into internal buffers using masks. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - target: Joint effort targets [N or N*m]. Shape is (num_instances, num_joints). - joint_mask: Joint mask. If None, then all joints are updated. - env_mask: Environment mask. If None, then all instances are updated. + target: Joint effort targets [N or N·m, depending on joint type]. Shape is + (num_instances, num_joints) with dtype wp.float32. + joint_mask: Joint mask. If None, all joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, all instances are updated. Shape is + (num_instances,). """ - self._set_target_into_buffer_mask(self._data._joint_effort_target, target, env_mask, joint_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + joint_mask_wp = self._resolve_joint_mask(joint_mask) + self.assert_shape_and_dtype(target, (self._num_instances, self._num_joints), wp.float32, "target") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(self._num_instances, self._num_joints), + inputs=[target, env_mask_wp, joint_mask_wp], + outputs=[self._data._joint_effort_target], + device=self._device, + ) + binding = self._get_binding(TT.DOF_ACTUATION_FORCE) + binding.write(self._data._joint_effort_target, mask=env_mask_wp) """ Operations - Tendons. @@ -1108,22 +2464,40 @@ def set_fixed_tendon_stiffness_index( fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set fixed tendon stiffness into internal buffers using indices. + """Set fixed-tendon stiffness over selected env / tendon indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_STIFFNESS`` is a CPU-only OVPhysX binding. - This function does not apply the tendon stiffness to the simulation. It only fills the buffers with - the desired values. To apply the tendon stiffness, call the - :meth:`write_fixed_tendon_properties_to_sim_index` method. + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - stiffness: Fixed tendon stiffness. Shape is (len(env_ids), len(fixed_tendon_ids)). - fixed_tendon_ids: The tendon indices to set the stiffness for. Defaults to None (all fixed tendons). - env_ids: Environment indices. If None, then all indices are used. + stiffness: Fixed-tendon stiffness [N/m]. May be a scalar + :class:`float` (broadcast), or shape + (len(env_ids), len(fixed_tendon_ids)) with dtype wp.float32. + fixed_tendon_ids: Fixed-tendon indices. Defaults to None (all fixed tendons). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - t = len(fixed_tendon_ids) if fixed_tendon_ids else self._nft() - self.assert_shape_and_dtype(stiffness, (n, t), wp.float32, "stiffness") - if self._data._fixed_tendon_stiffness is not None: - self._set_target_into_buffer(self._data._fixed_tendon_stiffness, stiffness, env_ids, fixed_tendon_ids) + env_ids = self._resolve_env_ids(env_ids) + tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_ids) + shape = (env_ids.shape[0], tendon_ids.shape[0]) + stiffness = self._broadcast_scalar_to_2d(stiffness, shape) + self.assert_shape_and_dtype(stiffness, shape, wp.float32, "stiffness") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=shape, + inputs=[stiffness, env_ids, tendon_ids], + outputs=[self._data._fixed_tendon_stiffness.data], + device=self._device, + ) + binding = self._get_binding(TT.FIXED_TENDON_STIFFNESS) + binding.write(self._data._fixed_tendon_stiffness.data, indices=env_ids) def set_fixed_tendon_stiffness_mask( self, @@ -1132,19 +2506,43 @@ def set_fixed_tendon_stiffness_mask( fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set fixed tendon stiffness into internal buffers using masks. + """Set fixed-tendon stiffness over selected env / tendon masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_STIFFNESS`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - stiffness: Fixed tendon stiffness. Shape is (num_instances, num_fixed_tendons). - fixed_tendon_mask: Tendon mask. If None, then all fixed tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + stiffness: Fixed-tendon stiffness [N/m]. May be a scalar + :class:`float` (broadcast), or shape + (num_instances, num_fixed_tendons) with dtype wp.float32. + fixed_tendon_mask: Fixed-tendon mask. If None, all fixed tendons are updated. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(stiffness, (self._num_instances, self._nft()), wp.float32, "stiffness") - if self._data._fixed_tendon_stiffness is not None: - self._set_target_into_buffer_mask( - self._data._fixed_tendon_stiffness, stiffness, env_mask, fixed_tendon_mask - ) - + env_mask_wp = self._resolve_env_mask(env_mask) + tendon_mask_wp = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + shape = (self._num_instances, self._num_fixed_tendons) + stiffness = self._broadcast_scalar_to_2d(stiffness, shape) + self.assert_shape_and_dtype(stiffness, shape, wp.float32, "stiffness") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[stiffness, env_mask_wp, tendon_mask_wp], + outputs=[self._data._fixed_tendon_stiffness.data], + device=self._device, + ) + binding = self._get_binding(TT.FIXED_TENDON_STIFFNESS) + binding.write(self._data._fixed_tendon_stiffness.data, mask=env_mask_wp) + def set_fixed_tendon_damping_index( self, *, @@ -1152,18 +2550,40 @@ def set_fixed_tendon_damping_index( fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set fixed tendon damping into internal buffers using indices. + """Set fixed-tendon damping over selected env / tendon indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_DAMPING`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - damping: Fixed tendon damping. Shape is (len(env_ids), len(fixed_tendon_ids)). - fixed_tendon_ids: The tendon indices to set the damping for. Defaults to None (all fixed tendons). - env_ids: Environment indices. If None, then all indices are used. + damping: Fixed-tendon damping [N·s/m]. May be a scalar :class:`float` + (broadcast), or shape (len(env_ids), len(fixed_tendon_ids)) with + dtype wp.float32. + fixed_tendon_ids: Fixed-tendon indices. Defaults to None (all fixed tendons). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - t = len(fixed_tendon_ids) if fixed_tendon_ids else self._nft() - self.assert_shape_and_dtype(damping, (n, t), wp.float32, "damping") - if self._data._fixed_tendon_damping is not None: - self._set_target_into_buffer(self._data._fixed_tendon_damping, damping, env_ids, fixed_tendon_ids) + env_ids = self._resolve_env_ids(env_ids) + tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_ids) + shape = (env_ids.shape[0], tendon_ids.shape[0]) + damping = self._broadcast_scalar_to_2d(damping, shape) + self.assert_shape_and_dtype(damping, shape, wp.float32, "damping") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=shape, + inputs=[damping, env_ids, tendon_ids], + outputs=[self._data._fixed_tendon_damping.data], + device=self._device, + ) + binding = self._get_binding(TT.FIXED_TENDON_DAMPING) + binding.write(self._data._fixed_tendon_damping.data, indices=env_ids) def set_fixed_tendon_damping_mask( self, @@ -1172,16 +2592,42 @@ def set_fixed_tendon_damping_mask( fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set fixed tendon damping into internal buffers using masks. + """Set fixed-tendon damping over selected env / tendon masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_DAMPING`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - damping: Fixed tendon damping. Shape is (num_instances, num_fixed_tendons). - fixed_tendon_mask: Tendon mask. If None, then all fixed tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + damping: Fixed-tendon damping [N·s/m]. May be a scalar :class:`float` + (broadcast), or shape (num_instances, num_fixed_tendons) with + dtype wp.float32. + fixed_tendon_mask: Fixed-tendon mask. If None, all fixed tendons are updated. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(damping, (self._num_instances, self._nft()), wp.float32, "damping") - if self._data._fixed_tendon_damping is not None: - self._set_target_into_buffer_mask(self._data._fixed_tendon_damping, damping, env_mask, fixed_tendon_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + tendon_mask_wp = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + shape = (self._num_instances, self._num_fixed_tendons) + damping = self._broadcast_scalar_to_2d(damping, shape) + self.assert_shape_and_dtype(damping, shape, wp.float32, "damping") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[damping, env_mask_wp, tendon_mask_wp], + outputs=[self._data._fixed_tendon_damping.data], + device=self._device, + ) + binding = self._get_binding(TT.FIXED_TENDON_DAMPING) + binding.write(self._data._fixed_tendon_damping.data, mask=env_mask_wp) def set_fixed_tendon_limit_stiffness_index( self, @@ -1190,20 +2636,40 @@ def set_fixed_tendon_limit_stiffness_index( fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set fixed tendon limit stiffness into internal buffers using indices. + """Set fixed-tendon limit stiffness over selected env / tendon indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_LIMIT_STIFFNESS`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limit_stiffness: Fixed tendon limit stiffness. Shape is (len(env_ids), len(fixed_tendon_ids)). - fixed_tendon_ids: The tendon indices. Defaults to None (all fixed tendons). - env_ids: Environment indices. If None, then all indices are used. + limit_stiffness: Fixed-tendon limit stiffness [N/m]. May be a + scalar :class:`float` (broadcast), or shape + (len(env_ids), len(fixed_tendon_ids)) with dtype wp.float32. + fixed_tendon_ids: Fixed-tendon indices. Defaults to None (all fixed tendons). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - t = len(fixed_tendon_ids) if fixed_tendon_ids else self._nft() - self.assert_shape_and_dtype(limit_stiffness, (n, t), wp.float32, "limit_stiffness") - if self._data._fixed_tendon_limit_stiffness is not None: - self._set_target_into_buffer( - self._data._fixed_tendon_limit_stiffness, limit_stiffness, env_ids, fixed_tendon_ids - ) + env_ids = self._resolve_env_ids(env_ids) + tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_ids) + shape = (env_ids.shape[0], tendon_ids.shape[0]) + limit_stiffness = self._broadcast_scalar_to_2d(limit_stiffness, shape) + self.assert_shape_and_dtype(limit_stiffness, shape, wp.float32, "limit_stiffness") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=shape, + inputs=[limit_stiffness, env_ids, tendon_ids], + outputs=[self._data._fixed_tendon_limit_stiffness.data], + device=self._device, + ) + binding = self._get_binding(TT.FIXED_TENDON_LIMIT_STIFFNESS) + binding.write(self._data._fixed_tendon_limit_stiffness.data, indices=env_ids) def set_fixed_tendon_limit_stiffness_mask( self, @@ -1212,18 +2678,42 @@ def set_fixed_tendon_limit_stiffness_mask( fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set fixed tendon limit stiffness into internal buffers using masks. + """Set fixed-tendon limit stiffness over selected env / tendon masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_LIMIT_STIFFNESS`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limit_stiffness: Fixed tendon limit stiffness. Shape is (num_instances, num_fixed_tendons). - fixed_tendon_mask: Tendon mask. If None, then all fixed tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + limit_stiffness: Fixed-tendon limit stiffness [N/m]. May be a + scalar :class:`float` (broadcast), or shape + (num_instances, num_fixed_tendons) with dtype wp.float32. + fixed_tendon_mask: Fixed-tendon mask. If None, all fixed tendons are updated. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(limit_stiffness, (self._num_instances, self._nft()), wp.float32, "limit_stiffness") - if self._data._fixed_tendon_limit_stiffness is not None: - self._set_target_into_buffer_mask( - self._data._fixed_tendon_limit_stiffness, limit_stiffness, env_mask, fixed_tendon_mask - ) + env_mask_wp = self._resolve_env_mask(env_mask) + tendon_mask_wp = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + shape = (self._num_instances, self._num_fixed_tendons) + limit_stiffness = self._broadcast_scalar_to_2d(limit_stiffness, shape) + self.assert_shape_and_dtype(limit_stiffness, shape, wp.float32, "limit_stiffness") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[limit_stiffness, env_mask_wp, tendon_mask_wp], + outputs=[self._data._fixed_tendon_limit_stiffness.data], + device=self._device, + ) + binding = self._get_binding(TT.FIXED_TENDON_LIMIT_STIFFNESS) + binding.write(self._data._fixed_tendon_limit_stiffness.data, mask=env_mask_wp) def set_fixed_tendon_position_limit_index( self, @@ -1232,19 +2722,46 @@ def set_fixed_tendon_position_limit_index( fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set fixed tendon position limits into internal buffers using indices. + """Set fixed-tendon position limits over selected env / tendon indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_LIMIT`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limit: Fixed tendon position limits. Shape is (len(env_ids), len(fixed_tendon_ids)) - with dtype wp.vec2f. - fixed_tendon_ids: The tendon indices. Defaults to None (all fixed tendons). - env_ids: Environment indices. If None, then all indices are used. + limit: Fixed-tendon position limits ``[lower, upper]`` [m]. + Shape is (len(env_ids), len(fixed_tendon_ids), 2) with dtype wp.float32. + fixed_tendon_ids: Fixed-tendon indices. Defaults to None (all fixed tendons). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - t = len(fixed_tendon_ids) if fixed_tendon_ids else self._nft() - self.assert_shape_and_dtype(limit, (n, t), wp.vec2f, "limit") - if self._data._fixed_tendon_pos_limits is not None: - self._set_target_into_buffer(self._data._fixed_tendon_pos_limits, limit, env_ids, fixed_tendon_ids) + env_ids = self._resolve_env_ids(env_ids) + tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_ids) + self.assert_shape_and_dtype(limit, (env_ids.shape[0], tendon_ids.shape[0], 2), wp.float32, "limit") + # Scatter [lower, upper] pairs into the vec2f cache buffer. + wp.launch( + shared_kernels.write_joint_position_limit_to_buffer_index, + dim=(env_ids.shape[0], tendon_ids.shape[0]), + inputs=[limit, env_ids, tendon_ids], + outputs=[self._data._fixed_tendon_pos_limits.data], + device=self._device, + ) + # Reinterpret the vec2f buffer as a (N, T, 2) float32 view for the wheel. + flat_src = wp.array( + ptr=self._data._fixed_tendon_pos_limits.data.ptr, + shape=(self._num_instances, self._num_fixed_tendons, 2), + dtype=wp.float32, + device=self._device, + copy=False, + ) + binding = self._get_binding(TT.FIXED_TENDON_LIMIT) + binding.write(flat_src, indices=env_ids) def set_fixed_tendon_position_limit_mask( self, @@ -1253,17 +2770,46 @@ def set_fixed_tendon_position_limit_mask( fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set fixed tendon position limits into internal buffers using masks. + """Set fixed-tendon position limits over selected env / tendon masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_LIMIT`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limit: Fixed tendon position limits. Shape is (num_instances, num_fixed_tendons) - with dtype wp.vec2f. - fixed_tendon_mask: Tendon mask. If None, then all fixed tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + limit: Fixed-tendon position limits ``[lower, upper]`` [m]. + Shape is (num_instances, num_fixed_tendons, 2) with dtype wp.float32. + fixed_tendon_mask: Fixed-tendon mask. If None, all fixed tendons are updated. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(limit, (self._num_instances, self._nft()), wp.vec2f, "limit") - if self._data._fixed_tendon_pos_limits is not None: - self._set_target_into_buffer_mask(self._data._fixed_tendon_pos_limits, limit, env_mask, fixed_tendon_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + tendon_mask_wp = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + self.assert_shape_and_dtype(limit, (self._num_instances, self._num_fixed_tendons, 2), wp.float32, "limit") + wp.launch( + shared_kernels.write_joint_position_limit_to_buffer_mask, + dim=(self._num_instances, self._num_fixed_tendons), + inputs=[limit, env_mask_wp, tendon_mask_wp], + outputs=[self._data._fixed_tendon_pos_limits.data], + device=self._device, + ) + flat_src = wp.array( + ptr=self._data._fixed_tendon_pos_limits.data.ptr, + shape=(self._num_instances, self._num_fixed_tendons, 2), + dtype=wp.float32, + device=self._device, + copy=False, + ) + binding = self._get_binding(TT.FIXED_TENDON_LIMIT) + binding.write(flat_src, mask=env_mask_wp) def set_fixed_tendon_rest_length_index( self, @@ -1272,18 +2818,40 @@ def set_fixed_tendon_rest_length_index( fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set fixed tendon rest length into internal buffers using indices. + """Set fixed-tendon rest lengths over selected env / tendon indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_REST_LENGTH`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - rest_length: Fixed tendon rest length. Shape is (len(env_ids), len(fixed_tendon_ids)). - fixed_tendon_ids: The tendon indices. Defaults to None (all fixed tendons). - env_ids: Environment indices. If None, then all indices are used. + rest_length: Fixed-tendon rest lengths [m]. May be a scalar + :class:`float` (broadcast), or shape + (len(env_ids), len(fixed_tendon_ids)) with dtype wp.float32. + fixed_tendon_ids: Fixed-tendon indices. Defaults to None (all fixed tendons). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - t = len(fixed_tendon_ids) if fixed_tendon_ids else self._nft() - self.assert_shape_and_dtype(rest_length, (n, t), wp.float32, "rest_length") - if self._data._fixed_tendon_rest_length is not None: - self._set_target_into_buffer(self._data._fixed_tendon_rest_length, rest_length, env_ids, fixed_tendon_ids) + env_ids = self._resolve_env_ids(env_ids) + tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_ids) + shape = (env_ids.shape[0], tendon_ids.shape[0]) + rest_length = self._broadcast_scalar_to_2d(rest_length, shape) + self.assert_shape_and_dtype(rest_length, shape, wp.float32, "rest_length") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=shape, + inputs=[rest_length, env_ids, tendon_ids], + outputs=[self._data._fixed_tendon_rest_length.data], + device=self._device, + ) + binding = self._get_binding(TT.FIXED_TENDON_REST_LENGTH) + binding.write(self._data._fixed_tendon_rest_length.data, indices=env_ids) def set_fixed_tendon_rest_length_mask( self, @@ -1292,18 +2860,42 @@ def set_fixed_tendon_rest_length_mask( fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set fixed tendon rest length into internal buffers using masks. + """Set fixed-tendon rest lengths over selected env / tendon masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_REST_LENGTH`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - rest_length: Fixed tendon rest length. Shape is (num_instances, num_fixed_tendons). - fixed_tendon_mask: Tendon mask. If None, then all fixed tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + rest_length: Fixed-tendon rest lengths [m]. May be a scalar + :class:`float` (broadcast), or shape + (num_instances, num_fixed_tendons) with dtype wp.float32. + fixed_tendon_mask: Fixed-tendon mask. If None, all fixed tendons are updated. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(rest_length, (self._num_instances, self._nft()), wp.float32, "rest_length") - if self._data._fixed_tendon_rest_length is not None: - self._set_target_into_buffer_mask( - self._data._fixed_tendon_rest_length, rest_length, env_mask, fixed_tendon_mask - ) + env_mask_wp = self._resolve_env_mask(env_mask) + tendon_mask_wp = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + shape = (self._num_instances, self._num_fixed_tendons) + rest_length = self._broadcast_scalar_to_2d(rest_length, shape) + self.assert_shape_and_dtype(rest_length, shape, wp.float32, "rest_length") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[rest_length, env_mask_wp, tendon_mask_wp], + outputs=[self._data._fixed_tendon_rest_length.data], + device=self._device, + ) + binding = self._get_binding(TT.FIXED_TENDON_REST_LENGTH) + binding.write(self._data._fixed_tendon_rest_length.data, mask=env_mask_wp) def set_fixed_tendon_offset_index( self, @@ -1312,18 +2904,40 @@ def set_fixed_tendon_offset_index( fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set fixed tendon offset into internal buffers using indices. + """Set fixed-tendon offsets over selected env / tendon indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_OFFSET`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - offset: Fixed tendon offset. Shape is (len(env_ids), len(fixed_tendon_ids)). - fixed_tendon_ids: The tendon indices. Defaults to None (all fixed tendons). - env_ids: Environment indices. If None, then all indices are used. + offset: Fixed-tendon offsets [m]. May be a scalar :class:`float` + (broadcast), or shape (len(env_ids), len(fixed_tendon_ids)) + with dtype wp.float32. + fixed_tendon_ids: Fixed-tendon indices. Defaults to None (all fixed tendons). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - t = len(fixed_tendon_ids) if fixed_tendon_ids else self._nft() - self.assert_shape_and_dtype(offset, (n, t), wp.float32, "offset") - if self._data._fixed_tendon_offset is not None: - self._set_target_into_buffer(self._data._fixed_tendon_offset, offset, env_ids, fixed_tendon_ids) + env_ids = self._resolve_env_ids(env_ids) + tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_ids) + shape = (env_ids.shape[0], tendon_ids.shape[0]) + offset = self._broadcast_scalar_to_2d(offset, shape) + self.assert_shape_and_dtype(offset, shape, wp.float32, "offset") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=shape, + inputs=[offset, env_ids, tendon_ids], + outputs=[self._data._fixed_tendon_offset.data], + device=self._device, + ) + binding = self._get_binding(TT.FIXED_TENDON_OFFSET) + binding.write(self._data._fixed_tendon_offset.data, indices=env_ids) def set_fixed_tendon_offset_mask( self, @@ -1332,16 +2946,42 @@ def set_fixed_tendon_offset_mask( fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set fixed tendon offset into internal buffers using masks. + """Set fixed-tendon offsets over selected env / tendon masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``FIXED_TENDON_OFFSET`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - offset: Fixed tendon offset. Shape is (num_instances, num_fixed_tendons). - fixed_tendon_mask: Tendon mask. If None, then all fixed tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + offset: Fixed-tendon offsets [m]. May be a scalar :class:`float` + (broadcast), or shape (num_instances, num_fixed_tendons) with + dtype wp.float32. + fixed_tendon_mask: Fixed-tendon mask. If None, all fixed tendons are updated. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(offset, (self._num_instances, self._nft()), wp.float32, "offset") - if self._data._fixed_tendon_offset is not None: - self._set_target_into_buffer_mask(self._data._fixed_tendon_offset, offset, env_mask, fixed_tendon_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + tendon_mask_wp = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + shape = (self._num_instances, self._num_fixed_tendons) + offset = self._broadcast_scalar_to_2d(offset, shape) + self.assert_shape_and_dtype(offset, shape, wp.float32, "offset") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[offset, env_mask_wp, tendon_mask_wp], + outputs=[self._data._fixed_tendon_offset.data], + device=self._device, + ) + binding = self._get_binding(TT.FIXED_TENDON_OFFSET) + binding.write(self._data._fixed_tendon_offset.data, mask=env_mask_wp) def write_fixed_tendon_properties_to_sim_index( self, @@ -1349,25 +2989,45 @@ def write_fixed_tendon_properties_to_sim_index( fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write fixed tendon properties into the simulation using indices. + """Push the cached fixed-tendon properties to the simulation in a single batch. + + PhysX's equivalent calls ``root_view.set_fixed_tendon_properties`` with + all six property buffers at once (articulation.py:3078-3086). OVPhysX + does not expose a single wheel-side batch setter; instead, we write + each ``FIXED_TENDON_*`` binding individually from the matching + ``self._data._fixed_tendon_*`` buffer. + + .. note:: + Only env indices apply to the wheel write; ``fixed_tendon_ids`` is + accepted for API parity with PhysX but is unused (the wheel writes + all tendons of the selected envs). Args: - fixed_tendon_ids: The fixed tendon indices to write the properties for. Defaults to None - (all fixed tendons). - env_ids: Environment indices. If None, then all indices are used. + fixed_tendon_ids: Accepted for PhysX API parity; ignored. + env_ids: Environment indices. If None, all environments are written. """ - if self._nft() == 0: - return - for tt, buf in [ + env_ids = self._resolve_env_ids(env_ids) + for tt, buf in ( (TT.FIXED_TENDON_STIFFNESS, self._data._fixed_tendon_stiffness), (TT.FIXED_TENDON_DAMPING, self._data._fixed_tendon_damping), (TT.FIXED_TENDON_LIMIT_STIFFNESS, self._data._fixed_tendon_limit_stiffness), - (TT.FIXED_TENDON_LIMIT, self._data._fixed_tendon_pos_limits), (TT.FIXED_TENDON_REST_LENGTH, self._data._fixed_tendon_rest_length), (TT.FIXED_TENDON_OFFSET, self._data._fixed_tendon_offset), - ]: - if buf is not None: - self._write_flat_tensor(tt, buf, env_ids, fixed_tendon_ids) + ): + binding = self._get_binding(tt) + if binding is not None: + binding.write(buf.data, indices=env_ids) + # Position-limit binding consumes a flat (N, T, 2) float32 view. + binding = self._get_binding(TT.FIXED_TENDON_LIMIT) + if binding is not None: + flat_src = wp.array( + ptr=self._data._fixed_tendon_pos_limits.data.ptr, + shape=(self._num_instances, self._num_fixed_tendons, 2), + dtype=wp.float32, + device=self._device, + copy=False, + ) + binding.write(flat_src, indices=env_ids) def write_fixed_tendon_properties_to_sim_mask( self, @@ -1375,24 +3035,33 @@ def write_fixed_tendon_properties_to_sim_mask( fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Write fixed tendon properties into the simulation using masks. + """Mask variant of :meth:`write_fixed_tendon_properties_to_sim_index`. Args: - fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + fixed_tendon_mask: Accepted for PhysX API parity; ignored. + env_mask: Environment mask. If None, all environments are written. """ - if self._nft() == 0: - return - for tt, buf in [ + env_mask_wp = self._resolve_env_mask(env_mask) + for tt, buf in ( (TT.FIXED_TENDON_STIFFNESS, self._data._fixed_tendon_stiffness), (TT.FIXED_TENDON_DAMPING, self._data._fixed_tendon_damping), (TT.FIXED_TENDON_LIMIT_STIFFNESS, self._data._fixed_tendon_limit_stiffness), - (TT.FIXED_TENDON_LIMIT, self._data._fixed_tendon_pos_limits), (TT.FIXED_TENDON_REST_LENGTH, self._data._fixed_tendon_rest_length), (TT.FIXED_TENDON_OFFSET, self._data._fixed_tendon_offset), - ]: - if buf is not None: - self._write_flat_tensor_mask(tt, buf, env_mask, fixed_tendon_mask) + ): + binding = self._get_binding(tt) + if binding is not None: + binding.write(buf.data, mask=env_mask_wp) + binding = self._get_binding(TT.FIXED_TENDON_LIMIT) + if binding is not None: + flat_src = wp.array( + ptr=self._data._fixed_tendon_pos_limits.data.ptr, + shape=(self._num_instances, self._num_fixed_tendons, 2), + dtype=wp.float32, + device=self._device, + copy=False, + ) + binding.write(flat_src, mask=env_mask_wp) def set_spatial_tendon_stiffness_index( self, @@ -1401,18 +3070,41 @@ def set_spatial_tendon_stiffness_index( spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set spatial tendon stiffness into internal buffers using indices. + """Set spatial-tendon stiffness over selected env / tendon indices into the simulation. + + ``SPATIAL_TENDON_STIFFNESS`` is a sim-device binding on OVPhysX (PhysX + applies tendon properties without a CPU clone), so the write goes + directly from the sim-device buffer to the wheel binding. + + .. note:: + This method expects partial data. A scalar :class:`float` is + broadcast to ``(len(env_ids), len(spatial_tendon_ids))``. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - stiffness: Spatial tendon stiffness. Shape is (len(env_ids), len(spatial_tendon_ids)). - spatial_tendon_ids: The tendon indices. Defaults to None (all spatial tendons). - env_ids: Environment indices. If None, then all indices are used. + stiffness: Spatial-tendon stiffness [N/m]. Scalar :class:`float`, + or shape ``(len(env_ids), len(spatial_tendon_ids))`` with + dtype wp.float32. + spatial_tendon_ids: Spatial-tendon indices. Defaults to None (all spatial tendons). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - t = len(spatial_tendon_ids) if spatial_tendon_ids else self._nst() - self.assert_shape_and_dtype(stiffness, (n, t), wp.float32, "stiffness") - if self._data._spatial_tendon_stiffness is not None: - self._set_target_into_buffer(self._data._spatial_tendon_stiffness, stiffness, env_ids, spatial_tendon_ids) + env_ids = self._resolve_env_ids(env_ids) + tendon_ids = self._resolve_spatial_tendon_ids(spatial_tendon_ids) + stiffness = self._broadcast_scalar_to_2d(stiffness, (env_ids.shape[0], tendon_ids.shape[0])) + self.assert_shape_and_dtype(stiffness, (env_ids.shape[0], tendon_ids.shape[0]), wp.float32, "stiffness") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], tendon_ids.shape[0]), + inputs=[stiffness, env_ids, tendon_ids], + outputs=[self._data._spatial_tendon_stiffness.data], + device=self._device, + ) + binding = self._get_binding(TT.SPATIAL_TENDON_STIFFNESS) + binding.write(self._data._spatial_tendon_stiffness.data, indices=env_ids) def set_spatial_tendon_stiffness_mask( self, @@ -1421,18 +3113,42 @@ def set_spatial_tendon_stiffness_mask( spatial_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set spatial tendon stiffness into internal buffers using masks. + """Set spatial-tendon stiffness over selected env / tendon masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``SPATIAL_TENDON_STIFFNESS`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - stiffness: Spatial tendon stiffness. Shape is (num_instances, num_spatial_tendons). - spatial_tendon_mask: Tendon mask. If None, then all spatial tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + stiffness: Spatial-tendon stiffness [N/m]. May be a scalar + :class:`float` (broadcast), or shape + (num_instances, num_spatial_tendons) with dtype wp.float32. + spatial_tendon_mask: Spatial-tendon mask. If None, all spatial tendons are updated. + Shape is (num_spatial_tendons,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(stiffness, (self._num_instances, self._nst()), wp.float32, "stiffness") - if self._data._spatial_tendon_stiffness is not None: - self._set_target_into_buffer_mask( - self._data._spatial_tendon_stiffness, stiffness, env_mask, spatial_tendon_mask - ) + env_mask_wp = self._resolve_env_mask(env_mask) + tendon_mask_wp = self._resolve_spatial_tendon_mask(spatial_tendon_mask) + shape = (self._num_instances, self._num_spatial_tendons) + stiffness = self._broadcast_scalar_to_2d(stiffness, shape) + self.assert_shape_and_dtype(stiffness, shape, wp.float32, "stiffness") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[stiffness, env_mask_wp, tendon_mask_wp], + outputs=[self._data._spatial_tendon_stiffness.data], + device=self._device, + ) + binding = self._get_binding(TT.SPATIAL_TENDON_STIFFNESS) + binding.write(self._data._spatial_tendon_stiffness.data, mask=env_mask_wp) def set_spatial_tendon_damping_index( self, @@ -1441,18 +3157,38 @@ def set_spatial_tendon_damping_index( spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set spatial tendon damping into internal buffers using indices. + """Set spatial-tendon damping over selected env / tendon indices into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``SPATIAL_TENDON_DAMPING`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - damping: Spatial tendon damping. Shape is (len(env_ids), len(spatial_tendon_ids)). - spatial_tendon_ids: The tendon indices. Defaults to None (all spatial tendons). - env_ids: Environment indices. If None, then all indices are used. + damping: Spatial-tendon damping [N·s/m]. Shape is + (len(env_ids), len(spatial_tendon_ids)) with dtype wp.float32. + spatial_tendon_ids: Spatial-tendon indices. Defaults to None (all spatial tendons). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - t = len(spatial_tendon_ids) if spatial_tendon_ids else self._nst() - self.assert_shape_and_dtype(damping, (n, t), wp.float32, "damping") - if self._data._spatial_tendon_damping is not None: - self._set_target_into_buffer(self._data._spatial_tendon_damping, damping, env_ids, spatial_tendon_ids) + env_ids = self._resolve_env_ids(env_ids) + tendon_ids = self._resolve_spatial_tendon_ids(spatial_tendon_ids) + damping = self._broadcast_scalar_to_2d(damping, (env_ids.shape[0], tendon_ids.shape[0])) + self.assert_shape_and_dtype(damping, (env_ids.shape[0], tendon_ids.shape[0]), wp.float32, "damping") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], tendon_ids.shape[0]), + inputs=[damping, env_ids, tendon_ids], + outputs=[self._data._spatial_tendon_damping.data], + device=self._device, + ) + binding = self._get_binding(TT.SPATIAL_TENDON_DAMPING) + binding.write(self._data._spatial_tendon_damping.data, indices=env_ids) def set_spatial_tendon_damping_mask( self, @@ -1461,18 +3197,42 @@ def set_spatial_tendon_damping_mask( spatial_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set spatial tendon damping into internal buffers using masks. + """Set spatial-tendon damping over selected env / tendon masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``SPATIAL_TENDON_DAMPING`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - damping: Spatial tendon damping. Shape is (num_instances, num_spatial_tendons). - spatial_tendon_mask: Tendon mask. If None, then all spatial tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + damping: Spatial-tendon damping [N·s/m]. May be a scalar + :class:`float` (broadcast), or shape + (num_instances, num_spatial_tendons) with dtype wp.float32. + spatial_tendon_mask: Spatial-tendon mask. If None, all spatial tendons are updated. + Shape is (num_spatial_tendons,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(damping, (self._num_instances, self._nst()), wp.float32, "damping") - if self._data._spatial_tendon_damping is not None: - self._set_target_into_buffer_mask( - self._data._spatial_tendon_damping, damping, env_mask, spatial_tendon_mask - ) + env_mask_wp = self._resolve_env_mask(env_mask) + tendon_mask_wp = self._resolve_spatial_tendon_mask(spatial_tendon_mask) + shape = (self._num_instances, self._num_spatial_tendons) + damping = self._broadcast_scalar_to_2d(damping, shape) + self.assert_shape_and_dtype(damping, shape, wp.float32, "damping") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[damping, env_mask_wp, tendon_mask_wp], + outputs=[self._data._spatial_tendon_damping.data], + device=self._device, + ) + binding = self._get_binding(TT.SPATIAL_TENDON_DAMPING) + binding.write(self._data._spatial_tendon_damping.data, mask=env_mask_wp) def set_spatial_tendon_limit_stiffness_index( self, @@ -1481,20 +3241,42 @@ def set_spatial_tendon_limit_stiffness_index( spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set spatial tendon limit stiffness into internal buffers using indices. + """Set spatial-tendon limit stiffness over selected env / tendon indices into the simulation. + + ``SPATIAL_TENDON_LIMIT_STIFFNESS`` is a sim-device binding on OVPhysX; + the write goes directly from the sim-device buffer to the wheel. + + .. note:: + This method expects partial data. A scalar :class:`float` is + broadcast to ``(len(env_ids), len(spatial_tendon_ids))``. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limit_stiffness: Spatial tendon limit stiffness. Shape is (len(env_ids), len(spatial_tendon_ids)). - spatial_tendon_ids: The tendon indices. Defaults to None (all spatial tendons). - env_ids: Environment indices. If None, then all indices are used. + limit_stiffness: Spatial-tendon limit stiffness [N/m]. Scalar + :class:`float`, or shape ``(len(env_ids), len(spatial_tendon_ids))`` + with dtype wp.float32. + spatial_tendon_ids: Spatial-tendon indices. Defaults to None (all spatial tendons). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - t = len(spatial_tendon_ids) if spatial_tendon_ids else self._nst() - self.assert_shape_and_dtype(limit_stiffness, (n, t), wp.float32, "limit_stiffness") - if self._data._spatial_tendon_limit_stiffness is not None: - self._set_target_into_buffer( - self._data._spatial_tendon_limit_stiffness, limit_stiffness, env_ids, spatial_tendon_ids - ) + env_ids = self._resolve_env_ids(env_ids) + tendon_ids = self._resolve_spatial_tendon_ids(spatial_tendon_ids) + limit_stiffness = self._broadcast_scalar_to_2d(limit_stiffness, (env_ids.shape[0], tendon_ids.shape[0])) + self.assert_shape_and_dtype( + limit_stiffness, (env_ids.shape[0], tendon_ids.shape[0]), wp.float32, "limit_stiffness" + ) + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], tendon_ids.shape[0]), + inputs=[limit_stiffness, env_ids, tendon_ids], + outputs=[self._data._spatial_tendon_limit_stiffness.data], + device=self._device, + ) + binding = self._get_binding(TT.SPATIAL_TENDON_LIMIT_STIFFNESS) + binding.write(self._data._spatial_tendon_limit_stiffness.data, indices=env_ids) def set_spatial_tendon_limit_stiffness_mask( self, @@ -1503,18 +3285,42 @@ def set_spatial_tendon_limit_stiffness_mask( spatial_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set spatial tendon limit stiffness into internal buffers using masks. + """Set spatial-tendon limit stiffness over selected env / tendon masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``SPATIAL_TENDON_LIMIT_STIFFNESS`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - limit_stiffness: Spatial tendon limit stiffness. Shape is (num_instances, num_spatial_tendons). - spatial_tendon_mask: Tendon mask. If None, then all spatial tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + limit_stiffness: Spatial-tendon limit stiffness [N/m]. May be a + scalar :class:`float` (broadcast), or shape + (num_instances, num_spatial_tendons) with dtype wp.float32. + spatial_tendon_mask: Spatial-tendon mask. If None, all spatial tendons are updated. + Shape is (num_spatial_tendons,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(limit_stiffness, (self._num_instances, self._nst()), wp.float32, "limit_stiffness") - if self._data._spatial_tendon_limit_stiffness is not None: - self._set_target_into_buffer_mask( - self._data._spatial_tendon_limit_stiffness, limit_stiffness, env_mask, spatial_tendon_mask - ) + env_mask_wp = self._resolve_env_mask(env_mask) + tendon_mask_wp = self._resolve_spatial_tendon_mask(spatial_tendon_mask) + shape = (self._num_instances, self._num_spatial_tendons) + limit_stiffness = self._broadcast_scalar_to_2d(limit_stiffness, shape) + self.assert_shape_and_dtype(limit_stiffness, shape, wp.float32, "limit_stiffness") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[limit_stiffness, env_mask_wp, tendon_mask_wp], + outputs=[self._data._spatial_tendon_limit_stiffness.data], + device=self._device, + ) + binding = self._get_binding(TT.SPATIAL_TENDON_LIMIT_STIFFNESS) + binding.write(self._data._spatial_tendon_limit_stiffness.data, mask=env_mask_wp) def set_spatial_tendon_offset_index( self, @@ -1523,18 +3329,40 @@ def set_spatial_tendon_offset_index( spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Set spatial tendon offset into internal buffers using indices. + """Set spatial-tendon offsets over selected env / tendon indices into the simulation. + + ``SPATIAL_TENDON_OFFSET`` is a sim-device binding on OVPhysX; the + write goes directly from the sim-device buffer to the wheel. + + .. note:: + This method expects partial data. A scalar :class:`float` is + broadcast to ``(len(env_ids), len(spatial_tendon_ids))``. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - offset: Spatial tendon offset. Shape is (len(env_ids), len(spatial_tendon_ids)). - spatial_tendon_ids: The tendon indices. Defaults to None (all spatial tendons). - env_ids: Environment indices. If None, then all indices are used. + offset: Spatial-tendon offsets [m]. Scalar :class:`float`, or + shape ``(len(env_ids), len(spatial_tendon_ids))`` with + dtype wp.float32. + spatial_tendon_ids: Spatial-tendon indices. Defaults to None (all spatial tendons). + env_ids: Environment indices. Defaults to None (all environments). """ - n = self._n_envs_index(env_ids) - t = len(spatial_tendon_ids) if spatial_tendon_ids else self._nst() - self.assert_shape_and_dtype(offset, (n, t), wp.float32, "offset") - if self._data._spatial_tendon_offset is not None: - self._set_target_into_buffer(self._data._spatial_tendon_offset, offset, env_ids, spatial_tendon_ids) + env_ids = self._resolve_env_ids(env_ids) + tendon_ids = self._resolve_spatial_tendon_ids(spatial_tendon_ids) + offset = self._broadcast_scalar_to_2d(offset, (env_ids.shape[0], tendon_ids.shape[0])) + self.assert_shape_and_dtype(offset, (env_ids.shape[0], tendon_ids.shape[0]), wp.float32, "offset") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], tendon_ids.shape[0]), + inputs=[offset, env_ids, tendon_ids], + outputs=[self._data._spatial_tendon_offset.data], + device=self._device, + ) + binding = self._get_binding(TT.SPATIAL_TENDON_OFFSET) + binding.write(self._data._spatial_tendon_offset.data, indices=env_ids) def set_spatial_tendon_offset_mask( self, @@ -1543,16 +3371,42 @@ def set_spatial_tendon_offset_mask( spatial_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Set spatial tendon offset into internal buffers using masks. + """Set spatial-tendon offsets over selected env / tendon masks into the simulation. + + This is a CPU-only write routed through pinned-host staging because + ``SPATIAL_TENDON_OFFSET`` is a CPU-only OVPhysX binding. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. + Performance is similar for both. However, to allow graphed pipelines, the + mask method must be used. Args: - offset: Spatial tendon offset. Shape is (num_instances, num_spatial_tendons). - spatial_tendon_mask: Tendon mask. If None, then all spatial tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. + offset: Spatial-tendon offsets [m]. May be a scalar :class:`float` + (broadcast), or shape (num_instances, num_spatial_tendons) with + dtype wp.float32. + spatial_tendon_mask: Spatial-tendon mask. If None, all spatial tendons are updated. + Shape is (num_spatial_tendons,). + env_mask: Environment mask. If None, all instances are updated. + Shape is (num_instances,). """ - self.assert_shape_and_dtype(offset, (self._num_instances, self._nst()), wp.float32, "offset") - if self._data._spatial_tendon_offset is not None: - self._set_target_into_buffer_mask(self._data._spatial_tendon_offset, offset, env_mask, spatial_tendon_mask) + env_mask_wp = self._resolve_env_mask(env_mask) + tendon_mask_wp = self._resolve_spatial_tendon_mask(spatial_tendon_mask) + shape = (self._num_instances, self._num_spatial_tendons) + offset = self._broadcast_scalar_to_2d(offset, shape) + self.assert_shape_and_dtype(offset, shape, wp.float32, "offset") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=shape, + inputs=[offset, env_mask_wp, tendon_mask_wp], + outputs=[self._data._spatial_tendon_offset.data], + device=self._device, + ) + binding = self._get_binding(TT.SPATIAL_TENDON_OFFSET) + binding.write(self._data._spatial_tendon_offset.data, mask=env_mask_wp) def write_spatial_tendon_properties_to_sim_index( self, @@ -1560,23 +3414,28 @@ def write_spatial_tendon_properties_to_sim_index( spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write spatial tendon properties into the simulation using indices. + """Push the cached spatial-tendon properties to the simulation in a single batch. + + Mirrors :meth:`write_fixed_tendon_properties_to_sim_index` for + spatial tendons. Only the four wheel-supported tensor types are + written; ``ARTICULATION_SPATIAL_TENDON_LIMIT`` and + ``ARTICULATION_SPATIAL_TENDON_REST_LENGTH`` are forward-compat + stubs (see ``docs/superpowers/specs/2026-04-28-ovphysx-wheel-gaps-for-marco.md``). Args: - spatial_tendon_ids: The spatial tendon indices to write the properties for. Defaults to None - (all spatial tendons). - env_ids: Environment indices. If None, then all indices are used. + spatial_tendon_ids: Accepted for PhysX API parity; ignored. + env_ids: Environment indices. If None, all environments are written. """ - if self._nst() == 0: - return - for tt, buf in [ + env_ids = self._resolve_env_ids(env_ids) + for tt, buf in ( (TT.SPATIAL_TENDON_STIFFNESS, self._data._spatial_tendon_stiffness), (TT.SPATIAL_TENDON_DAMPING, self._data._spatial_tendon_damping), (TT.SPATIAL_TENDON_LIMIT_STIFFNESS, self._data._spatial_tendon_limit_stiffness), (TT.SPATIAL_TENDON_OFFSET, self._data._spatial_tendon_offset), - ]: - if buf is not None: - self._write_flat_tensor(tt, buf, env_ids, spatial_tendon_ids) + ): + binding = self._get_binding(tt) + if binding is not None: + binding.write(buf.data, indices=env_ids) def write_spatial_tendon_properties_to_sim_mask( self, @@ -1584,80 +3443,48 @@ def write_spatial_tendon_properties_to_sim_mask( spatial_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Write spatial tendon properties into the simulation using masks. - - Args: - spatial_tendon_mask: Spatial tendon mask. If None, then all spatial tendons are used. - env_mask: Environment mask. If None, then all the instances are updated. - """ - if self._nst() == 0: - return - for tt, buf in [ + """Mask variant of :meth:`write_spatial_tendon_properties_to_sim_index`.""" + env_mask_wp = self._resolve_env_mask(env_mask) + for tt, buf in ( (TT.SPATIAL_TENDON_STIFFNESS, self._data._spatial_tendon_stiffness), (TT.SPATIAL_TENDON_DAMPING, self._data._spatial_tendon_damping), (TT.SPATIAL_TENDON_LIMIT_STIFFNESS, self._data._spatial_tendon_limit_stiffness), (TT.SPATIAL_TENDON_OFFSET, self._data._spatial_tendon_offset), - ]: - if buf is not None: - self._write_flat_tensor_mask(tt, buf, env_mask, spatial_tendon_mask) - - """ - Deprecated methods. - """ - - def write_root_state_to_sim( - self, - root_state: torch.Tensor | wp.array, - env_ids: Sequence[int] | wp.array | None = None, - ) -> None: - """Deprecated in base class. Use :meth:`write_root_pose_to_sim_index` and - :meth:`write_root_velocity_to_sim_index` instead.""" - self._write_root_state(TT.ROOT_POSE, root_state[:, :7], env_ids) - self._write_root_state(TT.ROOT_VELOCITY, root_state[:, 7:], env_ids) - - def write_root_com_state_to_sim( - self, - root_state: torch.Tensor | wp.array, - env_ids: Sequence[int] | wp.array | None = None, - ) -> None: - """Deprecated in base class. Use :meth:`write_root_com_pose_to_sim_index` and - :meth:`write_root_com_velocity_to_sim_index` instead.""" - self._write_root_state(TT.ROOT_POSE, root_state[:, :7], env_ids) - self._write_root_state(TT.ROOT_VELOCITY, root_state[:, 7:], env_ids) - - def write_root_link_state_to_sim( - self, - root_state: torch.Tensor | wp.array, - env_ids: Sequence[int] | wp.array | None = None, - ) -> None: - """Deprecated in base class. Use :meth:`write_root_link_pose_to_sim_index` and - :meth:`write_root_link_velocity_to_sim_index` instead.""" - self._write_root_state(TT.ROOT_POSE, root_state[:, :7], env_ids) - self._write_root_state(TT.ROOT_VELOCITY, root_state[:, 7:], env_ids) - - def write_joint_state_to_sim( - self, - position: torch.Tensor | wp.array, - velocity: torch.Tensor | wp.array, - joint_ids: Sequence[int] | None = None, - env_ids: Sequence[int] | wp.array | None = None, - ) -> None: - """Deprecated in base class. Use :meth:`write_joint_position_to_sim_index` and - :meth:`write_joint_velocity_to_sim_index` instead.""" - self.write_joint_position_to_sim_index(position=position, joint_ids=joint_ids, env_ids=env_ids) - self.write_joint_velocity_to_sim_index(velocity=velocity, joint_ids=joint_ids, env_ids=env_ids) + ): + binding = self._get_binding(tt) + if binding is not None: + binding.write(buf.data, mask=env_mask_wp) """ Internal helper. """ def _initialize_impl(self) -> None: + """Initialize the articulation from the OVPhysX simulation backend. + + 1. Acquire the OvPhysxManager.PhysX instance and the simulation device. + 2. Validate the USD prim tree (exactly one ArticulationRootAPI under + ``cfg.prim_path``; surface a clear RuntimeError otherwise). + 3. Convert IsaacLab prim-path notation to the glob pattern ovphysx expects, + extending it to the articulation root prim if needed. + 4. Eagerly create the bindings the data container needs and read counts + (num_instances, num_bodies, num_joints) and names from binding metadata. + 5. Discover tendon counts and names via :meth:`_process_tendons`. + 6. Construct the :class:`ArticulationData` container. + 7. Allocate asset-side buffers via :meth:`_create_buffers`. + 8. Apply initial state via :meth:`_process_cfg`. + 9. Prime the data via ``update(0.0)``. + 10. Mark the data as ready. + """ + # Step 1: Acquire PhysX instance and device. physx_instance = OvPhysxManager.get_physx_instance() if physx_instance is None: raise RuntimeError("OvPhysxManager has not been initialized yet.") + self._ovphysx = physx_instance + self._device = OvPhysxManager.get_device() prim_path = self.cfg.prim_path - # Convert IsaacLab prim-path notation to the glob patterns ovphysx expects. + # Step 2: Convert IsaacLab prim-path notation to the glob patterns ovphysx expects. # IsaacLab uses two conventions: # /World/envs/env_.*/Robot -- regex dot-star for "any env index" # /World/envs/{ENV_REGEX_NS}/Robot -- explicit placeholder @@ -1665,6 +3492,7 @@ def _initialize_impl(self) -> None: pattern = re.sub(r"\{ENV_REGEX_NS\}", "*", prim_path) pattern = re.sub(r"\.\*", "*", pattern) # env_.* -> env_* + # Step 3: Validate the USD prim tree and extend pattern to the articulation root if needed. # The pattern above points to the ArticulationCfg prim (e.g. /World/envs/env_*/Robot). # However, PhysicsArticulationRootAPI may be on a CHILD prim (e.g. /Robot/torso) # rather than on the prim itself. create_tensor_binding() only matches prims that @@ -1672,50 +3500,65 @@ def _initialize_impl(self) -> None: # articulation root. Mirror the PhysX backend's discovery logic: find the first # matching prim in the USD stage, walk its subtree for the articulation root, and # append the relative suffix to the glob pattern. - from pxr import UsdPhysics - - from isaaclab.sim.utils.queries import find_first_matching_prim, get_all_matching_child_prims - stage = PhysicsManager._sim.stage - first_prim = find_first_matching_prim(prim_path, stage=stage) - if first_prim is None: - raise RuntimeError(f"OvPhysxManager: no prim found for path '{prim_path}'.") - first_prim_path = first_prim.GetPath().pathString - root_prims = get_all_matching_child_prims( - first_prim_path, - predicate=lambda p: p.HasAPI(UsdPhysics.ArticulationRootAPI), - traverse_instance_prims=False, - ) - if len(root_prims) == 0: - raise RuntimeError( - f"No prim with PhysicsArticulationRootAPI found under '{first_prim_path}'." - " Check that the articulation has 'PhysicsArticulationRootAPI' applied." - ) - if len(root_prims) > 1: - raise RuntimeError( - f"Multiple articulation roots found under '{first_prim_path}': {root_prims}." - " There must be exactly one articulation root per prim path." - ) - self._articulation_root_path = root_prims[0].GetPath().pathString - root_relative = self._articulation_root_path[len(first_prim_path) :] - if root_relative: - # e.g. first_prim_path=/World/envs/env_0/Robot, root_relative=/torso - # pattern becomes /World/envs/env_*/Robot/torso + if self.cfg.articulation_root_prim_path is not None: + # Mirror the PhysX backend (articulation.py:3572): when the user has + # explicitly specified the articulation-root subpath, append it to + # the prim path directly and skip auto-discovery. Validate that the + # resulting expression resolves to at least one prim in the USD + # stage; otherwise surface a RuntimeError before downstream binding + # creation fails opaquely. + root_relative = self.cfg.articulation_root_prim_path + self._articulation_root_path = prim_path + root_relative + if sim_utils.find_first_matching_prim(self._articulation_root_path, stage=stage) is None: + raise RuntimeError( + f"Failed to find articulation root prim at '{self._articulation_root_path}'." + " Check that ``cfg.articulation_root_prim_path`` points at a prim that exists" + " in the USD stage." + ) pattern = pattern + root_relative - logger.info("OvPhysxManager: articulation root at '%s' (pattern extended to '%s')", root_relative, pattern) + logger.info("OvPhysxManager: explicit articulation root '%s' (pattern '%s')", root_relative, pattern) + else: + first_prim = sim_utils.find_first_matching_prim(prim_path, stage=stage) + if first_prim is None: + raise RuntimeError(f"Failed to find prim for expression: '{prim_path}'.") + first_prim_path = first_prim.GetPath().pathString + + root_prims = sim_utils.get_all_matching_child_prims( + first_prim_path, + predicate=lambda p: p.HasAPI(UsdPhysics.ArticulationRootAPI), + traverse_instance_prims=False, + ) + if len(root_prims) == 0: + raise RuntimeError( + f"Failed to find an articulation root when resolving '{prim_path}'." + " Ensure the prim has 'USD ArticulationRootAPI' applied." + ) + if len(root_prims) > 1: + raise RuntimeError( + f"Failed to find a single articulation root when resolving '{prim_path}'." + f" Found multiple under '{first_prim_path}'." + ) + + self._articulation_root_path = root_prims[0].GetPath().pathString + root_relative = self._articulation_root_path[len(first_prim_path) :] + if root_relative: + # e.g. first_prim_path=/World/envs/env_0/Robot, root_relative=/torso + # pattern becomes /World/envs/env_*/Robot/torso + pattern = pattern + root_relative + logger.info( + "OvPhysxManager: articulation root at '%s' (pattern extended to '%s')", root_relative, pattern + ) - # Bindings are created lazily (on first access) to avoid allocating - # handles for tensor types the user never queries. Only the root-pose - # binding is created eagerly because we need it to read articulation - # metadata (joint count, body count, names, fixed-base flag). - self._bindings: dict[int, Any] = {} - self._physx_instance = physx_instance self._binding_pattern = pattern + # Step 4: Eagerly create every binding the data container reads at init, + # so failures surface here rather than as KeyError downstream. eager_types = [ TT.ROOT_POSE, TT.DOF_POSITION, + TT.DOF_VELOCITY, TT.DOF_STIFFNESS, TT.DOF_DAMPING, TT.DOF_LIMIT, @@ -1733,6 +3576,14 @@ def _initialize_impl(self) -> None: except Exception: logger.debug("Could not create tensor binding for type %s on pattern %s", tt, pattern) + if not self._bindings: + raise RuntimeError( + f"OVPhysX could not create any articulation bindings for pattern {pattern!r}. " + f"Check that prim_path={prim_path!r} matches at least one " + "UsdPhysics.ArticulationRootAPI prim." + ) + + # Read metadata from the first available binding. sample = next(iter(self._bindings.values())) self._num_instances = sample.count self._num_joints = sample.dof_count @@ -1741,26 +3592,40 @@ def _initialize_impl(self) -> None: self._joint_names = list(sample.dof_names) self._body_names = list(sample.body_names) - # Create data container. - self._data = ArticulationData(self._bindings, self._device, binding_getter=self._get_binding) - - # Discover tendon counts/names before buffer allocation so that - # _create_buffers can size the tendon property arrays. + # Step 5: Discover tendon counts and names before buffer allocation. self._process_tendons() + # Step 6: Construct the data container. + self._data = ArticulationData( + self._bindings, + self._device, + self._num_instances, + self._num_bodies, + self._num_joints, + self._num_fixed_tendons, + self._num_spatial_tendons, + self._body_names, + self._joint_names, + self._fixed_tendon_names, + self._spatial_tendon_names, + binding_getter=self._get_binding, + ) + + # Step 7: Allocate asset-side buffers. self._create_buffers() + # Step 8: Apply initial state from config. self._process_cfg() + + # Step 8b: Build actuator instances and write drive properties to PhysX. self._process_actuators_cfg() - self._validate_cfg() - self._log_articulation_info() - # Cache the effort binding and a stable float32 view of the applied_torque - # buffer for write_data_to_sim(). The binding's internal write cache - # (keyed on object identity) handles the fast path automatically. + # Step 8c: Cache effort / target bindings and write-views for write_data_to_sim(). + # The effort view aliases applied_torque so the binding gets the actuator output + # without an extra copy. self._effort_binding = self._get_binding(TT.DOF_ACTUATION_FORCE) if self._effort_binding is not None: - torque = self._data.applied_torque + torque = self._data._applied_torque shape = self._effort_binding.shape self._effort_write_view = wp.array( ptr=torque.ptr, @@ -1772,7 +3637,6 @@ def _initialize_impl(self) -> None: else: self._effort_write_view = None - # Cache position/velocity target bindings + views for one-shot writes. def _make_write_view(tt, buf): b = self._get_binding(tt) if b is None or buf is None: @@ -1781,30 +3645,67 @@ def _make_write_view(tt, buf): return b, v self._pos_target_binding, self._pos_target_write_view = _make_write_view( - TT.DOF_POSITION_TARGET, self._data.joint_pos_target + TT.DOF_POSITION_TARGET, self._data._joint_pos_target ) self._vel_target_binding, self._vel_target_write_view = _make_write_view( - TT.DOF_VELOCITY_TARGET, self._data.joint_vel_target + TT.DOF_VELOCITY_TARGET, self._data._joint_vel_target ) - # Let the articulation data know that it is fully instantiated and ready to use. - self.data.is_primed = True + # Step 9: Validate the resolved configuration (raises ValueError when + # default joint pos/vel falls outside the wheel-reported limits). + # Mirrors PhysX's ``_validate_cfg`` ordering — after actuator/tendon + # processing, before priming. + self._validate_cfg() - def _create_buffers(self) -> None: - self._data._create_buffers() + # Step 10: Prime the data by performing the first read. + self.update(0.0) - self._ALL_INDICES = wp.array(np.arange(self._num_instances, dtype=np.int32), device=self._device) + # Step 11: Mark data as ready. + self._data.is_primed = True + def _create_buffers(self) -> None: + """Allocate asset-side buffers (index/mask constants, wrench buf, pinned CPU staging).""" + N = self._num_instances + B = self._num_bodies + J = self._num_joints + FT = self._num_fixed_tendons + ST = self._num_spatial_tendons + device = self._device + + # Index constants. + self._ALL_INDICES = wp.array(np.arange(N, dtype=np.int32), device=device) + self._ALL_BODY_INDICES = wp.array(np.arange(B, dtype=np.int32), device=device) + self._ALL_JOINT_INDICES = wp.array(np.arange(J, dtype=np.int32), device=device) + self._ALL_FIXED_TENDON_INDICES = wp.array(np.arange(FT, dtype=np.int32), device=device) + self._ALL_SPATIAL_TENDON_INDICES = wp.array(np.arange(ST, dtype=np.int32), device=device) + + # All-true masks. + self._ALL_TRUE_ENV_MASK = wp.array(np.ones(N, dtype=bool), dtype=wp.bool, device=device) + self._ALL_TRUE_BODY_MASK = wp.array(np.ones(B, dtype=bool), dtype=wp.bool, device=device) + self._ALL_TRUE_JOINT_MASK = wp.array(np.ones(J, dtype=bool), dtype=wp.bool, device=device) + self._ALL_TRUE_FIXED_TENDON_MASK = wp.array(np.ones(FT, dtype=bool), dtype=wp.bool, device=device) + self._ALL_TRUE_SPATIAL_TENDON_MASK = wp.array(np.ones(ST, dtype=bool), dtype=wp.bool, device=device) + + # Wrench buffer (force, torque, position) per body, written by the + # ``_body_wrench_to_world`` kernel and consumed by the + # ``LINK_WRENCH`` binding which expects the 3D ``(N, B, 9)`` shape. + self._wrench_buf = wp.zeros((N, B, 9), dtype=wp.float32, device=device) + + # Wrench composers. self._instantaneous_wrench_composer = WrenchComposer(self) self._permanent_wrench_composer = WrenchComposer(self) - self._wrench_buf = wp.zeros((self._num_instances, self._num_bodies, 9), dtype=wp.float32, device=self._device) - # Joint-index arrays for each actuator (filled by _process_actuators_cfg). + # Wrench scratch buffer (used by _apply_external_wrenches, not yet allocated above). + # Joint-index arrays for each actuator (populated by _process_actuators_cfg). self._joint_ids_per_actuator: dict[str, list[int]] = {} - self._write_scratch: dict[int, wp.array] = {} + + # Pinned-host CPU staging for env ids/masks (PR #5329 pattern). + self._cpu_env_ids_all = wp.zeros(N, dtype=wp.int32, device="cpu", pinned=True) + wp.copy(self._cpu_env_ids_all, self._ALL_INDICES) + self._cpu_env_mask = wp.zeros(N, dtype=wp.bool, device="cpu", pinned=True) def _process_cfg(self) -> None: - """Process the articulation configuration (initial state, soft limits, etc.).""" + """Populate default state buffers from the config (mirrors RigidObject and Newton Articulation).""" cfg = self.cfg N = self._num_instances D = self._num_joints @@ -1815,102 +3716,46 @@ def _process_cfg(self) -> None: default_root_vel = tuple(cfg.init_state.lin_vel) + tuple(cfg.init_state.ang_vel) np_pose = np.tile(np.array(default_root_pose, dtype=np.float32), (N, 1)) np_vel = np.tile(np.array(default_root_vel, dtype=np.float32), (N, 1)) - wp.copy( - self._data._default_root_pose, - wp.from_numpy(np_pose, dtype=wp.transformf, device=dev), - ) - wp.copy( - self._data._default_root_vel, - wp.from_numpy(np_vel, dtype=wp.spatial_vectorf, device=dev), - ) + self._data.default_root_pose = wp.array(np_pose, dtype=wp.transformf, device=dev) + self._data.default_root_vel = wp.array(np_vel, dtype=wp.spatial_vectorf, device=dev) # Default joint positions / velocities from config patterns. + # cfg.init_state.joint_pos is a dict[str, float] where keys are regex patterns + # matching joint names. We expand this into a (N, D) buffer. self._resolve_joint_values(cfg.init_state.joint_pos, self._data._default_joint_pos) self._resolve_joint_values(cfg.init_state.joint_vel, self._data._default_joint_vel) - # Keep soft-limit computation on-device, matching the PhysX/Newton path. - wp.launch( - update_soft_joint_pos_limits, - dim=(N, D), - inputs=[self._data.joint_pos_limits, cfg.soft_joint_pos_limit_factor], - outputs=[self._data._soft_joint_pos_limits], - device=dev, - ) + # Compute soft joint position limits from the hard limits read from the binding + # (or zeros if no joints). This matches the PhysX/Newton path. + if D > 0: + wp.launch( + update_soft_joint_pos_limits, + dim=(N, D), + inputs=[self._data.joint_pos_limits, cfg.soft_joint_pos_limit_factor], + outputs=[self._data._soft_joint_pos_limits], + device=dev, + ) - def _invalidate_initialize_callback(self, event) -> None: - self._is_initialized = False + def _process_tendons(self) -> None: + """Discover tendon counts from binding metadata and names from USD. - def _process_actuators_cfg(self) -> None: - """Build actuator instances from the config and write drive properties to PhysX. + Tendon counts come from the ovphysx binding metadata. Tendon names are + recovered from the exported USD articulation subtree because ovphysx + exposes joint names/counts, but not the per-joint USD paths that the + PhysX backend can query directly. + """ + self._fixed_tendon_names = [] + self._spatial_tendon_names = [] - Mirrors what the legacy PhysX backend does in its own _process_actuators_cfg: - - For ImplicitActuator: write the configured stiffness / damping to the PhysX - drive so the solver uses exactly the values from the actuator config. - - For all explicit actuators: zero out PhysX stiffness / damping so the - USD-authored drive gains cannot interfere with the explicit torque path. - - For all actuators: write effort_limit_sim and velocity_limit_sim. + sample = next(iter(self._bindings.values())) + self._num_fixed_tendons = getattr(sample, "fixed_tendon_count", 0) + self._num_spatial_tendons = getattr(sample, "spatial_tendon_count", 0) - These writes happen via TensorBinding (GPU-resident) after warmup has - allocated the GPU buffers (MODEL_INIT fires post-warmup). - """ - from isaaclab.actuators import ImplicitActuator - - self.actuators: dict[str, ActuatorBase] = {} - self._has_implicit_actuators = False - for name, act_cfg in self.cfg.actuators.items(): - joint_ids, joint_names = self.find_joints(act_cfg.joint_names_expr) - if not joint_ids: - logger.warning("Actuator '%s': no joints matched '%s'", name, act_cfg.joint_names_expr) - continue - act_cfg_copy = act_cfg.copy() - act = act_cfg_copy.class_type( - act_cfg_copy, - joint_names=joint_names, - joint_ids=joint_ids, - num_envs=self._num_instances, - device=self._device, - ) - self.actuators[name] = act - self._joint_ids_per_actuator[name] = joint_ids - - # Write drive gains and limits to PhysX to match the actuator config. - # Without this, PhysX retains whatever stiffness/damping was authored in the - # USD file, which can produce large restoring forces if the USD gains differ - # from the actuator config (e.g. a position-controlled robot exported with - # non-zero drive stiffness but configured with ImplicitActuator(stiffness=0)). - jids = list(joint_ids) - if isinstance(act, ImplicitActuator): - self._has_implicit_actuators = True - stiffness = act.stiffness # torch (N, J) - damping = act.damping # torch (N, J) - else: - stiffness = wp.zeros((self._num_instances, len(jids)), dtype=wp.float32, device=self._device) - damping = wp.zeros((self._num_instances, len(jids)), dtype=wp.float32, device=self._device) - self.write_joint_stiffness_to_sim_index(stiffness=stiffness, joint_ids=jids) - self.write_joint_damping_to_sim_index(damping=damping, joint_ids=jids) - self.write_joint_effort_limit_to_sim_index(limits=act.effort_limit_sim, joint_ids=jids) - self.write_joint_velocity_limit_to_sim_index(limits=act.velocity_limit_sim, joint_ids=jids) - - def _process_tendons(self) -> None: - """Discover tendon counts from binding metadata and names from USD. - - Tendon counts come from the ovphysx binding metadata. Tendon names are - recovered from the exported USD articulation subtree because ovphysx - exposes joint names/counts, but not the per-joint USD paths that the - PhysX backend can query directly. - """ - self._fixed_tendon_names = [] - self._spatial_tendon_names = [] - - sample = next(iter(self._bindings.values())) - self._num_fixed_tendons = getattr(sample, "fixed_tendon_count", 0) - self._num_spatial_tendons = getattr(sample, "spatial_tendon_count", 0) - - if self._num_fixed_tendons > 0 or self._num_spatial_tendons > 0: - stage_path = OvPhysxManager._stage_path - if stage_path is not None: - try: - from pxr import Usd, UsdPhysics + if self._num_fixed_tendons > 0 or self._num_spatial_tendons > 0: + stage_path = OvPhysxManager._stage_path + if stage_path is not None: + try: + from pxr import Usd from isaaclab.sim.utils.queries import get_all_matching_child_prims @@ -1947,55 +3792,165 @@ def _process_tendons(self) -> None: except Exception: logger.debug("Could not parse USD stage for tendon names at %s", stage_path) - self._data._num_fixed_tendons = self._num_fixed_tendons - self._data._num_spatial_tendons = self._num_spatial_tendons - self._data.fixed_tendon_names = self._fixed_tendon_names - self._data.spatial_tendon_names = self._spatial_tendon_names + # Push the discovered names into the data container if it already exists. + # (During _initialize_impl the data container is created AFTER _process_tendons, + # so self._data may not be set yet; the ArticulationData constructor receives the + # counts/names directly.) + if hasattr(self, "_data"): + self._data._num_fixed_tendons = self._num_fixed_tendons + self._data._num_spatial_tendons = self._num_spatial_tendons + self._data.fixed_tendon_names = self._fixed_tendon_names + self._data.spatial_tendon_names = self._spatial_tendon_names - def _apply_external_wrenches(self) -> None: - """Compose and write external wrenches to the LINK_WRENCH binding. + def _get_binding(self, tensor_type: int): + """Return a cached TensorBinding, creating it on first access. - WrenchComposer accumulates forces/torques in body (link) frame. - The LINK_WRENCH binding expects world-frame [fx,fy,fz,tx,ty,tz,px,py,pz]. - We rotate the body-frame vectors to world frame using the link quaternion - and pack them into the [N, L, 9] tensor with application position = origin. + Bindings are lightweight handles (a pointer + shape metadata into + PhysX's shared GPU buffer). Creating one does NOT allocate new GPU + memory -- the underlying simulation buffers are allocated once by PhysX + regardless of how many bindings point into them. Still, we defer + creation so that tensor types the user never queries are never looked up. + + Args: + tensor_type: The TensorType constant identifying which simulation + buffer to bind (e.g. :attr:`~isaaclab_ovphysx.tensor_types.ROOT_POSE`). + + Returns: + A TensorBinding object, or ``None`` if the binding could not be created. """ - inst = self._instantaneous_wrench_composer - perm = self._permanent_wrench_composer - if not inst.active and not perm.active: - return - if inst.active: - if perm.active: - inst.add_forces_and_torques_index( - forces=perm.composed_force, - torques=perm.composed_torque, - body_ids=list(range(self._num_bodies)), - env_ids=list(range(self._num_instances)), - ) - force_b = inst.composed_force - torque_b = inst.composed_torque - else: - force_b = perm.composed_force - torque_b = perm.composed_torque + binding = self._bindings.get(tensor_type) + if binding is not None: + return binding + try: + binding = self._ovphysx.create_tensor_binding(pattern=self._binding_pattern, tensor_type=tensor_type) + self._bindings[tensor_type] = binding + return binding + except Exception: + logger.debug("Could not create tensor binding for type %s", tensor_type) + return None - poses = self._data.body_link_pose_w - wp.launch( - _body_wrench_to_world, - dim=(self._num_instances, self._num_bodies), - inputs=[force_b, torque_b, poses], - outputs=[self._wrench_buf], - device=self._device, - ) - wrench_binding = self._get_binding(TT.LINK_WRENCH) - if wrench_binding is not None: - wrench_binding.write(self._wrench_buf) - inst.reset() + def _resolve_joint_values(self, pattern_dict: dict[str, float], buffer: wp.array) -> None: + """Resolve a ``{pattern: value}`` dict into a per-joint buffer. - def _apply_actuator_model(self) -> None: - """Run the actuator model to compute torques from user targets. + Builds values on CPU then copies to buffer's device (GPU arrays' + ``.numpy()`` returns a read-only copy, not a writable view). + + Args: + pattern_dict: A mapping from regex pattern strings to scalar values. + Matches joint names returned by :attr:`joint_names`. + buffer: Target warp array of shape ``(num_instances, num_joints)`` + to populate. + """ + buf_np = buffer.numpy() + modified = False + for pattern, value in pattern_dict.items(): + for j, name in enumerate(self._joint_names): + if re.fullmatch(pattern, name): + buf_np[:, j] = value + modified = True + if modified: + wp.copy(buffer, wp.from_numpy(buf_np, dtype=buffer.dtype, device=str(buffer.device))) + + def _n_envs_index(self, env_ids) -> int: + """Return the number of environments from an ``env_ids`` argument.""" + if env_ids is None: + return self._num_instances + if isinstance(env_ids, (list, tuple)): + return len(env_ids) + return env_ids.shape[0] if hasattr(env_ids, "shape") else len(env_ids) + + def _nft(self) -> int: + """Return the number of fixed tendons (0 if none).""" + return self._num_fixed_tendons + + def _nst(self) -> int: + """Return the number of spatial tendons (0 if none).""" + return self._num_spatial_tendons + + """ + Internal simulation callbacks. + """ + + def _invalidate_initialize_callback(self, event) -> None: + """Invalidate the asset on simulation reset.""" + super()._invalidate_initialize_callback(event) + + """ + Internal helpers -- Actuators. + """ + + def _process_actuators_cfg(self) -> None: + """Build actuator instances from the config and write drive properties to PhysX. + + Mirrors the PhysX backend's ``_process_actuators_cfg``: - IsaacLab actuators are torch-based. We convert warp -> torch via - DLPack (zero-copy on GPU), run the actuator, then write results back. + * For :class:`~isaaclab.actuators.ImplicitActuator`: write the configured + stiffness/damping to the PhysX drive so the solver uses exactly those values. + * For all explicit actuators: zero out PhysX stiffness/damping so USD-authored + drive gains cannot interfere with the explicit torque path. + * For all actuators: write :attr:`~isaaclab.actuators.ActuatorBase.effort_limit_sim` + and :attr:`~isaaclab.actuators.ActuatorBase.velocity_limit_sim`. + """ + from isaaclab.actuators import ImplicitActuator + + self.actuators: dict[str, Any] = {} + self._has_implicit_actuators = False + for name, act_cfg in self.cfg.actuators.items(): + joint_ids, joint_names = self.find_joints(act_cfg.joint_names_expr) + if not joint_ids: + logger.warning("Actuator '%s': no joints matched '%s'", name, act_cfg.joint_names_expr) + continue + act_cfg_copy = act_cfg.copy() + # Seed the actuator with the wheel's already-correct DOF defaults + # (USD-authored ``physxJoint:maxJointVelocity`` etc. that the wheel + # parsed at scene-load). Without these the ActuatorBase constructor + # falls back to ``inf`` for unset cfg fields, and the + # ``write_joint_*_to_sim_index`` calls below then *overwrite* the + # wheel's correct values with ``inf``. Mirrors the PhysX backend. + act = act_cfg_copy.class_type( + act_cfg_copy, + joint_names=joint_names, + joint_ids=joint_ids, + num_envs=self._num_instances, + device=self._device, + stiffness=self._data.joint_stiffness.torch[:, joint_ids], + damping=self._data.joint_damping.torch[:, joint_ids], + armature=self._data.joint_armature.torch[:, joint_ids], + friction=self._data.joint_friction_coeff.torch[:, joint_ids], + dynamic_friction=self._data.joint_dynamic_friction_coeff.torch[:, joint_ids], + viscous_friction=self._data.joint_viscous_friction_coeff.torch[:, joint_ids], + effort_limit=self._data.joint_effort_limits.torch[:, joint_ids].clone(), + velocity_limit=self._data.joint_vel_limits.torch[:, joint_ids], + ) + self.actuators[name] = act + self._joint_ids_per_actuator[name] = joint_ids + + # Write drive gains and limits to PhysX to match the actuator config. + # Without this, PhysX retains whatever stiffness/damping was authored in the + # USD file, which can produce large restoring forces when the USD gains differ + # from the actuator config. + jids = list(joint_ids) + if isinstance(act, ImplicitActuator): + self._has_implicit_actuators = True + stiffness = act.stiffness # torch (N, J) + damping = act.damping # torch (N, J) + else: + stiffness = wp.zeros((self._num_instances, len(jids)), dtype=wp.float32, device=self._device) + damping = wp.zeros((self._num_instances, len(jids)), dtype=wp.float32, device=self._device) + self.write_joint_stiffness_to_sim_index(stiffness=stiffness, joint_ids=jids) + self.write_joint_damping_to_sim_index(damping=damping, joint_ids=jids) + self.write_joint_effort_limit_to_sim_index(limits=act.effort_limit_sim, joint_ids=jids) + self.write_joint_velocity_limit_to_sim_index(limits=act.velocity_limit_sim, joint_ids=jids) + + def _apply_actuator_model(self) -> None: + """Run the actuator model to compute joint torques from user-supplied targets. + + IsaacLab actuators are torch-based. The method converts Warp buffers to + torch via DLPack (zero-copy on GPU), runs each actuator's + :meth:`~isaaclab.actuators.ActuatorBase.compute` method, then writes the + computed effort back to the private ``_computed_torque`` / ``_applied_torque`` + buffers of the data container. :meth:`write_data_to_sim` then pushes + ``_applied_torque`` to the ``DOF_ACTUATION_FORCE`` binding in one shot. """ from isaaclab.utils.types import ArticulationActions @@ -2006,7 +3961,7 @@ def _apply_actuator_model(self) -> None: jids_t = jids if isinstance(jids, list) else list(jids) all_joints = len(jids_t) == self._num_joints - # warp -> torch (zero-copy on same device via DLPack) + # Warp -> torch (zero-copy on same device via DLPack). jp_target_full = self._data.joint_pos_target.torch jv_target_full = self._data.joint_vel_target.torch je_target_full = self._data.joint_effort_target.torch @@ -2037,640 +3992,380 @@ def _apply_actuator_model(self) -> None: ct[:, jids_t] = act.computed_effort at[:, jids_t] = act.applied_effort + """ + Internal helpers -- Debugging. + """ + def _validate_cfg(self) -> None: - pass + """Validate the configuration after processing. - def _log_articulation_info(self) -> None: - """Log information about the articulation. - - .. note:: We purposefully read the values from the simulator to ensure that the values are configured as - expected. - """ - from prettytable import PrettyTable - - def format_large_number(_, v: float) -> str: - if abs(v) >= 1e3: - return f"{v:.1e}" - return f"{v:.3f}" - - def format_limits(_, v: tuple[float, float]) -> str: - if abs(v[0]) >= 1e3 or abs(v[1]) >= 1e3: - return f"[{v[0]:.1e}, {v[1]:.1e}]" - return f"[{v[0]:.3f}, {v[1]:.3f}]" - - stiffnesses = self.data.joint_stiffness.warp.numpy()[0].tolist() - dampings = self.data.joint_damping.warp.numpy()[0].tolist() - armatures = self.data.joint_armature.warp.numpy()[0].tolist() - frictions = self.data.joint_friction_coeff.warp.numpy()[0].tolist() - pos_limits_np = self.data.joint_pos_limits.warp.numpy().reshape(self._num_instances, self._num_joints, 2) - position_limits = [tuple(pos_limits_np[0, j].tolist()) for j in range(self._num_joints)] - velocity_limits = self.data.joint_vel_limits.warp.numpy()[0].tolist() - effort_limits = self.data.joint_effort_limits.warp.numpy()[0].tolist() - - joint_table = PrettyTable() - joint_table.title = f"Simulation Joint Information (Prim path: {self.cfg.prim_path})" - joint_table.field_names = [ - "Index", - "Name", - "Stiffness", - "Damping", - "Armature", - "Friction", - "Position Limits", - "Velocity Limits", - "Effort Limits", - ] - joint_table.custom_format["Stiffness"] = format_large_number - joint_table.custom_format["Damping"] = format_large_number - joint_table.custom_format["Armature"] = format_large_number - joint_table.custom_format["Friction"] = format_large_number - joint_table.custom_format["Position Limits"] = format_limits - joint_table.custom_format["Velocity Limits"] = format_large_number - joint_table.custom_format["Effort Limits"] = format_large_number - joint_table.align["Name"] = "l" - - for index, name in enumerate(self.joint_names): - joint_table.add_row( - [ - index, - name, - stiffnesses[index], - dampings[index], - armatures[index], - frictions[index], - position_limits[index], - velocity_limits[index], - effort_limits[index], - ] - ) - logger.info(f"Simulation parameters for joints in {self.cfg.prim_path}:\n" + joint_table.get_string()) - - if self.num_fixed_tendons > 0: - ft_stiffnesses = self.data.fixed_tendon_stiffness.warp.numpy()[0].tolist() - ft_dampings = self.data.fixed_tendon_damping.warp.numpy()[0].tolist() - ft_limit_stiffnesses = self.data.fixed_tendon_limit_stiffness.warp.numpy()[0].tolist() - ft_limits_np = self.data.fixed_tendon_pos_limits.warp.numpy().reshape( - self._num_instances, self.num_fixed_tendons, 2 - ) - ft_limits = [tuple(ft_limits_np[0, t].tolist()) for t in range(self.num_fixed_tendons)] - ft_rest_lengths = self.data.fixed_tendon_rest_length.warp.numpy()[0].tolist() - ft_offsets = self.data.fixed_tendon_offset.warp.numpy()[0].tolist() - - tendon_table = PrettyTable() - tendon_table.title = f"Simulation Fixed Tendon Information (Prim path: {self.cfg.prim_path})" - tendon_table.field_names = [ - "Index", - "Stiffness", - "Damping", - "Limit Stiffness", - "Limits", - "Rest Length", - "Offset", - ] - tendon_table.custom_format["Stiffness"] = format_large_number - tendon_table.custom_format["Damping"] = format_large_number - tendon_table.custom_format["Limit Stiffness"] = format_large_number - tendon_table.custom_format["Limits"] = format_limits - tendon_table.custom_format["Rest Length"] = format_large_number - tendon_table.custom_format["Offset"] = format_large_number - for index in range(self.num_fixed_tendons): - tendon_table.add_row( - [ - index, - ft_stiffnesses[index], - ft_dampings[index], - ft_limit_stiffnesses[index], - ft_limits[index], - ft_rest_lengths[index], - ft_offsets[index], - ] - ) - logger.info( - f"Simulation parameters for fixed tendons in {self.cfg.prim_path}:\n" + tendon_table.get_string() - ) + Mirrors :meth:`isaaclab_physx.assets.Articulation._validate_cfg` (raises + ``ValueError`` with a per-joint message when any default joint position + is outside ``[lower, upper]`` or any default joint velocity exceeds the + per-joint max velocity). Reads come from :attr:`ArticulationData` + accessors instead of PhysX's ``root_view.get_dof_limits`` / + ``get_dof_max_velocities`` because OVPhysX's ``root_view`` is the + per-tensor-type bindings dict. - if self.num_spatial_tendons > 0: - st_stiffnesses = self.data.spatial_tendon_stiffness.warp.numpy()[0].tolist() - st_dampings = self.data.spatial_tendon_damping.warp.numpy()[0].tolist() - st_limit_stiffnesses = self.data.spatial_tendon_limit_stiffness.warp.numpy()[0].tolist() - st_offsets = self.data.spatial_tendon_offset.warp.numpy()[0].tolist() - - tendon_table = PrettyTable() - tendon_table.title = f"Simulation Spatial Tendon Information (Prim path: {self.cfg.prim_path})" - tendon_table.field_names = [ - "Index", - "Stiffness", - "Damping", - "Limit Stiffness", - "Offset", - ] - tendon_table.float_format = ".3" - for index in range(self.num_spatial_tendons): - tendon_table.add_row( - [ - index, - st_stiffnesses[index], - st_dampings[index], - st_limit_stiffnesses[index], - st_offsets[index], - ] - ) - logger.info( - f"Simulation parameters for spatial tendons in {self.cfg.prim_path}:\n" + tendon_table.get_string() - ) + .. note:: + Must be called only after :meth:`_create_buffers` / + :meth:`_process_cfg` / :meth:`_process_actuators_cfg`, otherwise + limits and defaults may not yet reflect the final values. + """ + # check that the default joint positions are within the limits + joint_pos_limits = self._data.joint_pos_limits.torch[0] # (num_joints, 2) + default_joint_pos = self._data.default_joint_pos.torch[0] # (num_joints,) + out_of_range = default_joint_pos < joint_pos_limits[:, 0] + out_of_range |= default_joint_pos > joint_pos_limits[:, 1] + violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) + if len(violated_indices) > 0: + msg = "The following joints have default positions out of the limits: \n" + for idx in violated_indices: + joint_name = self._data.joint_names[idx] + joint_limit = joint_pos_limits[idx] + joint_pos = default_joint_pos[idx] + msg += f"\t- '{joint_name}': {joint_pos:.3f} not in [{joint_limit[0]:.3f}, {joint_limit[1]:.3f}]\n" + raise ValueError(msg) + + # check that the default joint velocities are within the limits + joint_max_vel = self._data.joint_vel_limits.torch[0] # (num_joints,) + default_joint_vel = self._data.default_joint_vel.torch[0] # (num_joints,) + out_of_range = torch.abs(default_joint_vel) > joint_max_vel + violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) + if len(violated_indices) > 0: + msg = "The following joints have default velocities out of the limits: \n" + for idx in violated_indices: + joint_name = self._data.joint_names[idx] + joint_limit = [-joint_max_vel[idx], joint_max_vel[idx]] + joint_vel = default_joint_vel[idx] + msg += f"\t- '{joint_name}': {joint_vel:.3f} not in [{joint_limit[0]:.3f}, {joint_limit[1]:.3f}]\n" + raise ValueError(msg) - """ - Internal helpers -- Bindings. - """ + def _log_articulation_info(self) -> None: + pass - def _get_binding(self, tensor_type: int): - """Return a cached TensorBinding, creating it on first access. + def _resolve_env_ids(self, env_ids) -> wp.array: + """Resolve environment indices to a warp int32 array on ``self._device`` (mirrors PhysX). - Bindings are lightweight handles (a pointer + shape metadata into - PhysX's shared GPU buffer). Creating one does NOT allocate new GPU - memory -- the underlying simulation buffers are allocated once by PhysX - regardless of how many bindings point into them. Still, we defer - creation so that tensor types the user never queries are never looked up. + Tests sometimes hand us indices on CPU even when the sim runs on GPU; we move the + resolved array onto ``self._device`` so kernel launches don't fail on a device + mismatch. """ - binding = self._bindings.get(tensor_type) - if binding is not None: - return binding - try: - binding = self._physx_instance.create_tensor_binding(pattern=self._binding_pattern, tensor_type=tensor_type) - self._bindings[tensor_type] = binding - return binding - except Exception: - logger.debug("Could not create tensor binding for type %s", tensor_type) - return None + if env_ids is None or env_ids == slice(None): + return self._ALL_INDICES + if isinstance(env_ids, list): + return wp.array(env_ids, dtype=wp.int32, device=self._device) + if isinstance(env_ids, torch.Tensor): + return wp.from_torch(env_ids.to(torch.int32), dtype=wp.int32) + if isinstance(env_ids, wp.array) and str(env_ids.device) != self._device: + env_ids = wp.clone(env_ids, device=self._device) + return env_ids + + def _resolve_body_ids(self, body_ids) -> wp.array: + """Resolve body indices to a warp int32 array on ``self._device`` (mirrors PhysX).""" + if body_ids is None or body_ids == slice(None): + return self._ALL_BODY_INDICES + if isinstance(body_ids, list): + return wp.array(body_ids, dtype=wp.int32, device=self._device) + if isinstance(body_ids, torch.Tensor): + return wp.from_torch(body_ids.to(torch.int32), dtype=wp.int32) + if isinstance(body_ids, wp.array) and str(body_ids.device) != self._device: + body_ids = wp.clone(body_ids, device=self._device) + return body_ids + + def _resolve_joint_ids(self, joint_ids) -> wp.array: + """Resolve joint indices to a warp int32 array on ``self._device``.""" + if joint_ids is None or joint_ids == slice(None): + return self._ALL_JOINT_INDICES + if isinstance(joint_ids, list): + return wp.array(joint_ids, dtype=wp.int32, device=self._device) + if isinstance(joint_ids, torch.Tensor): + return wp.from_torch(joint_ids.to(torch.int32), dtype=wp.int32) + if isinstance(joint_ids, wp.array) and str(joint_ids.device) != self._device: + joint_ids = wp.clone(joint_ids, device=self._device) + return joint_ids + + def _resolve_fixed_tendon_ids(self, tendon_ids) -> wp.array: + """Resolve fixed-tendon indices to a warp int32 array on ``self._device``.""" + if tendon_ids is None or tendon_ids == slice(None): + return self._ALL_FIXED_TENDON_INDICES + if isinstance(tendon_ids, list): + return wp.array(tendon_ids, dtype=wp.int32, device=self._device) + if isinstance(tendon_ids, torch.Tensor): + return wp.from_torch(tendon_ids.to(torch.int32), dtype=wp.int32) + if isinstance(tendon_ids, wp.array) and str(tendon_ids.device) != self._device: + tendon_ids = wp.clone(tendon_ids, device=self._device) + return tendon_ids + + def _resolve_spatial_tendon_ids(self, tendon_ids) -> wp.array: + """Resolve spatial-tendon indices to a warp int32 array on ``self._device``.""" + if tendon_ids is None or tendon_ids == slice(None): + return self._ALL_SPATIAL_TENDON_INDICES + if isinstance(tendon_ids, list): + return wp.array(tendon_ids, dtype=wp.int32, device=self._device) + if isinstance(tendon_ids, torch.Tensor): + return wp.from_torch(tendon_ids.to(torch.int32), dtype=wp.int32) + if isinstance(tendon_ids, wp.array) and str(tendon_ids.device) != self._device: + tendon_ids = wp.clone(tendon_ids, device=self._device) + return tendon_ids + + def _broadcast_scalar_to_2d( + self, value: float | torch.Tensor | wp.array, shape: tuple[int, int] + ) -> torch.Tensor | wp.array: + """Broadcast a scalar :class:`float` to a ``(rows, cols)`` torch ``float32`` tensor. + + Tendon and joint setters accept ``float | torch.Tensor | wp.array``; the + underlying ``shared_kernels.write_2d_data_to_buffer_*`` kernels only + accept 2D arrays. This helper expands a Python float into a constant + tensor on :attr:`_device`; tensor / warp inputs are returned as-is. + + Mirrors the PhysX backend's ``isinstance(value, float)`` branching, + which dispatches to ``articulation_kernels.float_data_to_buffer_with_*``. + OVPhysX does not have those scalar kernels, so we materialize the + broadcast on the Python side. - """ - Internal helpers -- Write. - """ + Args: + value: Scalar float or 2D tensor / warp array. + shape: ``(rows, cols)`` target shape used when broadcasting a + scalar. - def _to_flat_f32(self, data, target_shape: tuple[int, ...] | None = None) -> wp.array | np.ndarray: - """Ensure data is a contiguous float32 tensor suitable for binding I/O. + Returns: + A 2D :class:`torch.Tensor` on ``self._device`` if *value* was a + float; otherwise *value* unchanged. + """ + if isinstance(value, float): + return torch.full(shape, value, dtype=torch.float32, device=self._device) + return value - State tensor bindings (positions, velocities, poses) live on the - simulation device (GPU in GPU mode). We always return data on - self._device so the binding device check passes. + def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array: + """Resolve an environment mask to a ``wp.bool`` array on ``self._device``. - For structured warp dtypes (transformf, spatial_vectorf, etc.) a - zero-copy flat float32 view is created instead of roundtripping - through CPU numpy. + OVPhysX (like Newton) uses the wheel's native ``binding.write(mask=...)`` path, + so the mask is preserved end-to-end -- no ``torch.nonzero`` conversion needed. + ``None`` returns the pre-allocated all-true mask. """ - dev = self._device - if isinstance(data, wp.array): - if str(data.device) != dev: - data = wp.clone(data, device=dev) - if data.dtype == wp.float32: - return data - # Structured dtype: zero-copy flat float32 view. - # transformf -> [N, 7], spatial_vectorf -> [N, 6], etc. - floats_per_elem = data.strides[0] // 4 - return wp.array( - ptr=data.ptr, - shape=(data.shape[0], floats_per_elem), - dtype=wp.float32, - device=dev, - copy=False, - ) - elif isinstance(data, torch.Tensor): - if data.is_cuda and dev.startswith("cuda"): - return wp.from_torch(data.detach().contiguous().float()) - np_data = data.detach().cpu().numpy().astype(np.float32) - return wp.from_numpy(np_data, dtype=wp.float32, device=dev) - elif isinstance(data, np.ndarray): - return wp.from_numpy(data.astype(np.float32), dtype=wp.float32, device=dev) - elif isinstance(data, (int, float)): - return wp.from_numpy(np.array(data, dtype=np.float32), dtype=wp.float32, device=dev) - return wp.from_numpy(np.asarray(data, dtype=np.float32), dtype=wp.float32, device=dev) - - def _as_gpu_f32_2d(self, data, cols: int) -> wp.array: - """View/convert data as 2D [rows, cols] float32 on self._device. - - For warp arrays with structured dtypes (transformf, spatial_vectorf), - creates a zero-copy flat float32 view. For torch/numpy, converts to - warp on the simulation device. + if env_mask is None: + return self._ALL_TRUE_ENV_MASK + if isinstance(env_mask, torch.Tensor): + return wp.from_torch(env_mask.to(torch.bool), dtype=wp.bool) + if isinstance(env_mask, wp.array) and str(env_mask.device) != self._device: + env_mask = wp.clone(env_mask, device=self._device) + return env_mask + + def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array: + """Resolve a body mask to a ``wp.bool`` array on ``self._device`` (Newton-style).""" + if body_mask is None: + return self._ALL_TRUE_BODY_MASK + if isinstance(body_mask, torch.Tensor): + return wp.from_torch(body_mask.to(torch.bool), dtype=wp.bool) + if isinstance(body_mask, wp.array) and str(body_mask.device) != self._device: + body_mask = wp.clone(body_mask, device=self._device) + return body_mask + + def _resolve_joint_mask(self, joint_mask: wp.array | None) -> wp.array: + """Resolve a joint mask to a ``wp.bool`` array on ``self._device``.""" + if joint_mask is None: + return self._ALL_TRUE_JOINT_MASK + if isinstance(joint_mask, torch.Tensor): + return wp.from_torch(joint_mask.to(torch.bool), dtype=wp.bool) + if isinstance(joint_mask, wp.array) and str(joint_mask.device) != self._device: + joint_mask = wp.clone(joint_mask, device=self._device) + return joint_mask + + def _resolve_fixed_tendon_mask(self, tendon_mask: wp.array | None) -> wp.array: + """Resolve a fixed-tendon mask to a ``wp.bool`` array on ``self._device``.""" + if tendon_mask is None: + return self._ALL_TRUE_FIXED_TENDON_MASK + if isinstance(tendon_mask, torch.Tensor): + return wp.from_torch(tendon_mask.to(torch.bool), dtype=wp.bool) + if isinstance(tendon_mask, wp.array) and str(tendon_mask.device) != self._device: + tendon_mask = wp.clone(tendon_mask, device=self._device) + return tendon_mask + + def _resolve_spatial_tendon_mask(self, tendon_mask: wp.array | None) -> wp.array: + """Resolve a spatial-tendon mask to a ``wp.bool`` array on ``self._device``.""" + if tendon_mask is None: + return self._ALL_TRUE_SPATIAL_TENDON_MASK + if isinstance(tendon_mask, torch.Tensor): + return wp.from_torch(tendon_mask.to(torch.bool), dtype=wp.bool) + if isinstance(tendon_mask, wp.array) and str(tendon_mask.device) != self._device: + tendon_mask = wp.clone(tendon_mask, device=self._device) + return tendon_mask + + def _get_cpu_env_mask(self, env_mask: wp.array) -> wp.array: + """Return a pinned-host CPU copy of *env_mask* for a CPU-only binding write. + + ``env_mask`` is normally on ``self._device``; the wheel's ``binding.write(mask=...)`` + requires the mask on the binding's device, which is CPU for mass / coms / inertia. + Reuses the pre-allocated ``_cpu_env_mask`` pinned buffer. """ - dev = self._device - if isinstance(data, wp.array): - if str(data.device) != dev: - data = wp.clone(data, device=dev) - if data.dtype == wp.float32 and data.ndim == 2: - return data - n = data.shape[0] - return wp.array( - ptr=data.ptr, - shape=(n, cols), - dtype=wp.float32, - device=dev, - copy=False, - ) - if isinstance(data, torch.Tensor) and data.is_cuda and dev.startswith("cuda"): - return wp.from_torch(data.detach().contiguous().float().reshape(-1, cols)) - np_data = self._to_cpu_numpy(data).reshape(-1, cols) - return wp.from_numpy(np_data, dtype=wp.float32, device=dev) - - def _get_write_scratch(self, tensor_type: int, binding) -> wp.array: - """Return a cached GPU scratch buffer for read-modify-write.""" - if not hasattr(self, "_write_scratch"): - self._write_scratch = {} - buf = self._write_scratch.get(tensor_type) - if buf is None: - buf = wp.zeros(binding.shape, dtype=wp.float32, device=self._device) - self._write_scratch[tensor_type] = buf - return buf - - def _write_root_state(self, tensor_type: int, data, env_ids=None, mask=None, _ids_gpu=None) -> None: - """GPU-native write for root pose [N,7] or velocity [N,6]. - - Three paths, fastest first: - - Full write (no env_ids, no mask): zero-copy DLPack. - - Indexed write with full-size data: zero-copy view + indices. - The binding API only copies the indexed rows from the full buffer, - so no read-modify-write is needed when data is already [N,...]. - - Indexed write with partial data [K,...]: scatter kernel into a GPU - scratch buffer, then write with indices. - - Masked write: data is always full [N,...], pass directly with mask. - - Args: - _ids_gpu: Pre-converted GPU warp int32 array of env indices. - When provided, skips the per-call GPU->CPU->GPU conversion - of env_ids. - """ - binding = self._get_binding(tensor_type) - if binding is None: - return - N, C = binding.shape - - if env_ids is None and _ids_gpu is None and mask is None: - binding.write(self._to_flat_f32(data)) - self._invalidate_root_caches(tensor_type) - return - - src = self._as_gpu_f32_2d(data, C) - - if env_ids is not None or _ids_gpu is not None: - if _ids_gpu is None: - _ids_gpu = self._env_ids_to_gpu_warp(env_ids) - K = _ids_gpu.shape[0] - if src.shape[0] == N: - binding.write(src, indices=_ids_gpu) - else: - scratch = self._get_write_scratch(tensor_type, binding) - binding.read(scratch) - wp.launch( - _scatter_rows_partial, - dim=(K, C), - inputs=[scratch, src, _ids_gpu], - device=self._device, - ) - binding.write(scratch, indices=_ids_gpu) - else: - mask_u8 = wp.from_numpy( - self._to_cpu_numpy(mask).astype(np.uint8), - device=self._device, - ) - binding.write(src, mask=mask_u8) - self._invalidate_root_caches(tensor_type) - - def _invalidate_root_caches(self, tensor_type: int) -> None: - """Force re-read from GPU on next property access after a binding write.""" - if tensor_type == TT.ROOT_POSE: - self.data._root_link_pose_w.timestamp = -1.0 - self.data._root_com_pose_w.timestamp = -1.0 - elif tensor_type == TT.ROOT_VELOCITY: - self.data._root_link_vel_w.timestamp = -1.0 - self.data._root_com_vel_w.timestamp = -1.0 - - def _write_flat_tensor(self, tensor_type: int, data, env_ids=None, joint_ids=None, _ids_gpu=None) -> None: - """Write a 2-D tensor to a binding, with optional env/joint index subsetting.""" - if isinstance(data, (int, float)): - return - binding = self._get_binding(tensor_type) - if binding is None: - return - from isaaclab_ovphysx.tensor_types import _CPU_ONLY_TYPES - - is_cpu_only = tensor_type in _CPU_ONLY_TYPES - - # CPU-only types or column scatter must go through numpy. - if is_cpu_only or joint_ids is not None: - target_device = "cpu" if is_cpu_only else self._device - np_data = self._to_cpu_numpy(data) - if joint_ids is not None: - if is_cpu_only: - full = np.zeros(binding.shape, dtype=np.float32) - binding.read(full) - else: - scratch = self._get_write_scratch(tensor_type, binding) - binding.read(scratch) - full = scratch.numpy() - jids = self._to_cpu_indices(joint_ids, np.intp) - if env_ids is not None: - eids = self._to_cpu_indices(env_ids, np.intp) - full[np.ix_(eids, jids)] = np_data.reshape(len(eids), len(jids), *np_data.shape[2:]) - else: - full[:, jids] = np_data.reshape(full.shape[0], len(jids), *np_data.shape[2:]) - binding.write(wp.from_numpy(full, dtype=wp.float32, device=target_device)) - elif env_ids is not None: - if is_cpu_only: - full = np.zeros(binding.shape, dtype=np.float32) - binding.read(full) - else: - scratch = self._get_write_scratch(tensor_type, binding) - binding.read(scratch) - full = scratch.numpy() - eids = self._to_cpu_indices(env_ids, np.intp) - full[eids] = np_data if np_data.shape[0] == len(eids) else np_data[eids] - flat = wp.from_numpy(full.astype(np.float32), dtype=wp.float32, device=target_device) - idx = _ids_gpu if _ids_gpu is not None else self._env_ids_to_gpu_warp(env_ids) - binding.write(flat, indices=idx) - else: - binding.write(wp.from_numpy(np_data.astype(np.float32), dtype=wp.float32, device=target_device)) - return - - # GPU path: data stays on device. - if env_ids is None and _ids_gpu is None: - binding.write(self._to_flat_f32(data)) - return - - N, C = binding.shape[0], binding.shape[1] - src = self._as_gpu_f32_2d(data, C) - if _ids_gpu is None: - _ids_gpu = self._env_ids_to_gpu_warp(env_ids) - K = _ids_gpu.shape[0] - if src.shape[0] == N: - binding.write(src, indices=_ids_gpu) - else: - scratch = self._get_write_scratch(tensor_type, binding) - binding.read(scratch) - wp.launch( - _scatter_rows_partial, - dim=(K, C), - inputs=[scratch, src, _ids_gpu], - device=self._device, - ) - binding.write(scratch, indices=_ids_gpu) - - def _write_flat_tensor_mask(self, tensor_type: int, data, env_mask=None, joint_mask=None) -> None: - """Write a 2-D tensor to a binding, with optional env/joint mask subsetting.""" - if isinstance(data, (int, float)): - return - binding = self._get_binding(tensor_type) - if binding is None: - return - from isaaclab_ovphysx.tensor_types import _CPU_ONLY_TYPES - - is_cpu_only = tensor_type in _CPU_ONLY_TYPES - - # CPU-only types or column-mask scatter must go through numpy. - if is_cpu_only or joint_mask is not None: - target_device = "cpu" if is_cpu_only else self._device - np_data = self._to_cpu_numpy(data) - if joint_mask is not None: - # GPU bindings cannot read into numpy directly; read into GPU - # scratch first, then pull to CPU for column scatter. - if is_cpu_only: - full = np.zeros(binding.shape, dtype=np.float32) - binding.read(full) - else: - scratch = self._get_write_scratch(tensor_type, binding) - binding.read(scratch) - full = scratch.numpy() - jmask = self._to_cpu_numpy(joint_mask).astype(bool) - cols = np.where(jmask)[0] - if env_mask is not None: - emask = self._to_cpu_numpy(env_mask).astype(bool) - rows = np.where(emask)[0] - full[rows[:, None], cols] = np_data[rows[:, None], cols] - else: - full[:, cols] = np_data[:, cols] - binding.write(wp.from_numpy(full.astype(np.float32), dtype=wp.float32, device=target_device)) - elif env_mask is not None: - flat = wp.from_numpy(np_data.astype(np.float32), dtype=wp.float32, device=target_device) - mask_u8 = wp.from_numpy( - self._to_cpu_numpy(env_mask).astype(np.uint8), - device=target_device, - ) - binding.write(flat, mask=mask_u8) - else: - binding.write(wp.from_numpy(np_data.astype(np.float32), dtype=wp.float32, device=target_device)) - return + wp.copy(self._cpu_env_mask, env_mask) + return self._cpu_env_mask - # GPU path: data stays on device. - if env_mask is None: - binding.write(self._to_flat_f32(data)) - return + def _get_cpu_env_ids(self, env_ids: wp.array | torch.Tensor) -> wp.array: + """Return CPU int32 indices, using the pre-allocated pinned ``_cpu_env_ids_all`` + fast path when *env_ids* matches ``_ALL_INDICES`` (PR #5329 pattern). + """ + if isinstance(env_ids, torch.Tensor): + env_ids = wp.from_torch(env_ids, dtype=wp.int32) + if env_ids.ptr == self._ALL_INDICES.ptr: + return self._cpu_env_ids_all + return wp.clone(env_ids, device="cpu") - # Data is full [N, D], the binding API selects rows via the mask. - mask_u8 = wp.from_numpy( - self._to_cpu_numpy(env_mask).astype(np.uint8), - device=self._device, + """ + Deprecated methods. + """ + + def write_root_state_to_sim( + self, + root_state: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Deprecated; use :meth:`write_root_link_pose_to_sim_index` and + :meth:`write_root_com_velocity_to_sim_index` instead. + + Args: + root_state: Root state [m, m, m, qw, qx, qy, qz, m/s, m/s, m/s, rad/s, rad/s, rad/s]. + Shape is (len(env_ids), 13) with dtype wp.float32. + env_ids: Environment indices. Defaults to None (all environments). + """ + warnings.warn( + "The function 'write_root_state_to_sim' will be deprecated in a future release. Please" + " use 'write_root_link_pose_to_sim_index' and 'write_root_com_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, ) - binding.write(self._to_flat_f32(data), mask=mask_u8) + self.write_root_link_pose_to_sim_index(root_pose=root_state[:, :7], env_ids=env_ids) + self.write_root_com_velocity_to_sim_index(root_velocity=root_state[:, 7:], env_ids=env_ids) - def _write_friction_column(self, data, env_ids=None, joint_ids=None) -> None: - """Write static friction coefficient into column 0 of DOF_FRICTION_PROPERTIES [N,D,3].""" - binding = self._get_binding(TT.DOF_FRICTION_PROPERTIES) - if binding is None: - return - full = np.zeros(binding.shape, dtype=np.float32) - binding.read(full) - if isinstance(data, (int, float)): - if env_ids is not None and joint_ids is not None: - eids = self._to_cpu_numpy(env_ids).astype(np.intp) - jids = self._to_cpu_indices(joint_ids, np.intp) - full[np.ix_(eids, jids, [0])] = data - elif env_ids is not None: - eids = self._to_cpu_numpy(env_ids).astype(np.intp) - full[eids, :, 0] = data - elif joint_ids is not None: - jids = self._to_cpu_indices(joint_ids, np.intp) - full[:, jids, 0] = data - else: - full[..., 0] = data - binding.write(wp.from_numpy(full.astype(np.float32), dtype=wp.float32, device="cpu")) - return - np_data = self._to_cpu_numpy(data) - if env_ids is not None and joint_ids is not None: - eids = self._to_cpu_numpy(env_ids).astype(np.intp) - jids = self._to_cpu_indices(joint_ids, np.intp) - full[np.ix_(eids, jids, [0])] = np_data.reshape(len(eids), len(jids), 1) - elif env_ids is not None: - eids = self._to_cpu_numpy(env_ids).astype(np.intp) - full[eids, :, 0] = np_data.reshape(len(eids), -1) - elif joint_ids is not None: - jids = self._to_cpu_indices(joint_ids, np.intp) - full[:, jids, 0] = np_data.reshape(full.shape[0], len(jids)) - else: - full[..., 0] = np_data.reshape(full.shape[0], full.shape[1]) - binding.write(wp.from_numpy(full.astype(np.float32), dtype=wp.float32, device="cpu")) + def write_root_com_state_to_sim( + self, + root_state: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Deprecated; use :meth:`write_root_com_pose_to_sim_index` and + :meth:`write_root_com_velocity_to_sim_index` instead. - def _write_friction_column_mask(self, data, env_mask=None, joint_mask=None) -> None: - """Write static friction coefficient via mask into column 0 of DOF_FRICTION_PROPERTIES.""" - binding = self._get_binding(TT.DOF_FRICTION_PROPERTIES) - if binding is None: - return - full = np.zeros(binding.shape, dtype=np.float32) - binding.read(full) - if isinstance(data, (int, float)): - new_col = np.full((full.shape[0], full.shape[1]), data, dtype=np.float32) - else: - new_col = self._to_cpu_numpy(data).reshape(full.shape[0], full.shape[1]) - if env_mask is not None: - emask = self._to_cpu_numpy(env_mask).astype(bool) - if joint_mask is not None: - jmask = self._to_cpu_numpy(joint_mask).astype(bool) - rows = np.where(emask)[0] - cols = np.where(jmask)[0] - full[rows[:, None], cols, 0] = new_col[rows[:, None], cols] - else: - full[emask, :, 0] = new_col[emask] - elif joint_mask is not None: - jmask = self._to_cpu_numpy(joint_mask).astype(bool) - full[:, jmask, 0] = new_col[:, jmask] - else: - full[..., 0] = new_col - binding.write(wp.from_numpy(full.astype(np.float32), dtype=wp.float32, device="cpu")) - - def _write_joint_subset(self, tensor_type: int, buffer: wp.array, joint_ids: list[int]) -> None: - """Write a full-width joint buffer into the simulation for an actuator's joints.""" - binding = self._get_binding(tensor_type) - if binding is None: - return - if not hasattr(self, "_write_dltensor_cache"): - self._write_dltensor_cache = {} - cache_key = (tensor_type, buffer.ptr) - cached = self._write_dltensor_cache.get(cache_key) - if cached is None: - flat = self._to_flat_f32(buffer) - from ovphysx._dlpack_utils import acquire_dltensor - - dl, keepalive = acquire_dltensor(flat) - self._write_dltensor_cache[cache_key] = (dl, keepalive, flat) - cached = self._write_dltensor_cache[cache_key] - binding.write(cached[0]) - - @staticmethod - def _to_cpu_numpy(data) -> np.ndarray: - """Convert data (warp, torch, numpy, scalar) to a CPU numpy array.""" - if isinstance(data, wp.array): - return data.numpy().astype(np.float32) - if isinstance(data, torch.Tensor): - return data.detach().cpu().numpy().astype(np.float32) - return np.asarray(data, dtype=np.float32) - - @staticmethod - def _to_cpu_indices(data, dtype=np.int32) -> np.ndarray: - """Convert index array (warp, torch, list, numpy) to CPU numpy int array.""" - if isinstance(data, torch.Tensor): - return data.detach().cpu().numpy().astype(dtype) - if isinstance(data, wp.array): - return data.numpy().astype(dtype) - return np.asarray(data, dtype=dtype) - - def _env_ids_to_gpu_warp(self, env_ids) -> wp.array: - """Convert env_ids to a GPU int32 warp array, with single-entry caching. - - The cache avoids repeated GPU -> CPU -> GPU round-trips when the same - ``env_ids`` object is passed to multiple binding writes in a single step - (e.g. reset writes root_pose, root_vel, joint_pos, joint_vel). A new - object identity (``id()``) or shape change invalidates the cache. - """ - if hasattr(env_ids, "data_ptr"): - key = (env_ids.data_ptr(), env_ids.shape[0]) - elif isinstance(env_ids, wp.array): - key = (env_ids.ptr, env_ids.shape[0]) - else: - key = None - - if key is not None and hasattr(self, "_ids_cache_key") and self._ids_cache_key == key: - return self._ids_cache_val - - result = wp.array(self._to_cpu_indices(env_ids, np.int32), device=self._device) - if key is not None: - self._ids_cache_key = key - self._ids_cache_val = result - return result - - def _set_target_into_buffer(self, buffer: wp.array, data, env_ids=None, joint_ids=None) -> None: - """Set user-provided target data into a warp command buffer. - - For the common case (no index subset), this uses wp.copy to stay on - the simulation device. Subset writes (specific env_ids or joint_ids) - fall back to CPU because warp does not support scatter indexing. - """ - # Fast path: all-joints shortcut. When joint_ids covers every joint - # and env_ids is None, the subset is equivalent to a full copy. - if joint_ids is not None and env_ids is None: - n_joints = buffer.shape[1] if len(buffer.shape) > 1 else 1 - if hasattr(joint_ids, "__len__") and len(joint_ids) == n_joints: - joint_ids = None - if env_ids is None and joint_ids is None: - src = self._to_flat_f32(data) - if isinstance(src, np.ndarray): - src = wp.from_numpy(src, dtype=wp.float32, device=buffer.device) - wp.copy(buffer, src) - else: - np_data = self._to_cpu_numpy(data) - buf_np = buffer.numpy() - env_idx = self._to_cpu_numpy(env_ids).astype(np.intp) if env_ids is not None else None - jnt_idx = self._to_cpu_numpy(joint_ids).astype(np.intp) if joint_ids is not None else None - if env_idx is not None and jnt_idx is not None: - buf_np[np.ix_(env_idx, jnt_idx)] = np_data - elif env_idx is not None: - buf_np[env_idx] = np_data - else: - buf_np[:, jnt_idx] = np_data - wp.copy(buffer, wp.from_numpy(buf_np, dtype=wp.float32, device=buffer.device)) + Args: + root_state: Root CoM state [m, m, m, qw, qx, qy, qz, m/s, m/s, m/s, rad/s, rad/s, rad/s]. + Shape is (len(env_ids), 13) with dtype wp.float32. + env_ids: Environment indices. Defaults to None (all environments). + """ + warnings.warn( + "The function 'write_root_com_state_to_sim' will be deprecated in a future release. Please" + " use 'write_root_com_pose_to_sim_index' and 'write_root_com_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + self.write_root_com_pose_to_sim_index(root_pose=root_state[:, :7], env_ids=env_ids) + self.write_root_com_velocity_to_sim_index(root_velocity=root_state[:, 7:], env_ids=env_ids) - def _set_target_into_buffer_mask(self, buffer: wp.array, data, env_mask=None, joint_mask=None) -> None: - """Set user-provided target data into a warp command buffer using masks.""" - if env_mask is None: - src = self._to_flat_f32(data) - if isinstance(src, np.ndarray): - src = wp.from_numpy(src, dtype=wp.float32, device=buffer.device) - wp.copy(buffer, src) - else: - np_data = self._to_cpu_numpy(data) - buf_np = buffer.numpy() - mask_np = self._to_cpu_numpy(env_mask).astype(bool) - buf_np[mask_np] = np_data[mask_np] - wp.copy(buffer, wp.from_numpy(buf_np, dtype=wp.float32, device=buffer.device)) + def write_root_link_state_to_sim( + self, + root_state: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Deprecated; use :meth:`write_root_link_pose_to_sim_index` and + :meth:`write_root_link_velocity_to_sim_index` instead. - """ - Internal helpers -- Utilities. - """ + Args: + root_state: Root link state [m, m, m, qw, qx, qy, qz, m/s, m/s, m/s, rad/s, rad/s, rad/s]. + Shape is (len(env_ids), 13) with dtype wp.float32. + env_ids: Environment indices. Defaults to None (all environments). + """ + warnings.warn( + "The function 'write_root_link_state_to_sim' will be deprecated in a future release. Please" + " use 'write_root_link_pose_to_sim_index' and 'write_root_link_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + self.write_root_link_pose_to_sim_index(root_pose=root_state[:, :7], env_ids=env_ids) + self.write_root_link_velocity_to_sim_index(root_velocity=root_state[:, 7:], env_ids=env_ids) - def _n_envs_index(self, env_ids): - """Return the number of environments from an env_ids argument.""" - if env_ids is None: - return self._num_instances - if isinstance(env_ids, (list, tuple)): - return len(env_ids) - return env_ids.shape[0] if hasattr(env_ids, "shape") else len(env_ids) + def write_joint_state_to_sim( + self, + position: torch.Tensor | wp.array, + velocity: torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Deprecated combined joint-state write; use :meth:`write_joint_position_to_sim_index` + and :meth:`write_joint_velocity_to_sim_index` instead. - def _nft(self): - """Return the number of fixed tendons (0 if none).""" - return getattr(self, "_num_fixed_tendons", 0) + Args: + position: Joint positions [m or rad, depending on joint type]. Shape is + (len(env_ids), len(joint_ids)) with dtype wp.float32. + velocity: Joint velocities [m/s or rad/s, depending on joint type]. Shape is + (len(env_ids), len(joint_ids)) with dtype wp.float32. + joint_ids: Joint indices. Defaults to None (all joints). + env_ids: Environment indices. Defaults to None (all environments). + """ + warnings.warn( + "write_joint_state_to_sim is deprecated; use write_joint_position_to_sim_index" + " and write_joint_velocity_to_sim_index instead.", + DeprecationWarning, + stacklevel=2, + ) + self.write_joint_position_to_sim_index(position=position, joint_ids=joint_ids, env_ids=env_ids) + self.write_joint_velocity_to_sim_index(velocity=velocity, joint_ids=joint_ids, env_ids=env_ids) - def _nst(self): - """Return the number of spatial tendons (0 if none).""" - return getattr(self, "_num_spatial_tendons", 0) + def write_joint_friction_coefficient_to_sim( + self, + joint_friction_coeff: torch.Tensor | wp.array | float, + joint_dynamic_friction_coeff: torch.Tensor | wp.array | float | None = None, + joint_viscous_friction_coeff: torch.Tensor | wp.array | float | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Deprecated, same as :meth:`write_joint_friction_coefficient_to_sim_index`.""" + warnings.warn( + "The function 'write_joint_friction_coefficient_to_sim' will be deprecated in a future release. Please" + " use 'write_joint_friction_coefficient_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + self.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff=joint_friction_coeff, + joint_dynamic_friction_coeff=joint_dynamic_friction_coeff, + joint_viscous_friction_coeff=joint_viscous_friction_coeff, + joint_ids=joint_ids, + env_ids=env_ids, + ) - def _resolve_joint_values(self, pattern_dict: dict[str, float], buffer: wp.array) -> None: - """Resolve a {pattern: value} dict into a per-joint buffer. + def write_joint_dynamic_friction_coefficient_to_sim( + self, + joint_dynamic_friction_coeff: torch.Tensor | wp.array | float, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Deprecated, same as :meth:`write_joint_dynamic_friction_coefficient_to_sim_index`.""" + warnings.warn( + "The function 'write_joint_dynamic_friction_coefficient_to_sim' will be deprecated in a future release. " + "Please use 'write_joint_dynamic_friction_coefficient_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + self.write_joint_dynamic_friction_coefficient_to_sim_index( + joint_dynamic_friction_coeff=joint_dynamic_friction_coeff, + joint_ids=joint_ids, + env_ids=env_ids, + ) - Builds values on CPU then copies to buffer's device (GPU arrays' - .numpy() returns a read-only copy, not a writable view). - """ - buf_np = buffer.numpy() - modified = False - for pattern, value in pattern_dict.items(): - for j, name in enumerate(self._joint_names): - if re.fullmatch(pattern, name): - buf_np[:, j] = value - modified = True - if modified: - wp.copy(buffer, wp.from_numpy(buf_np, dtype=buffer.dtype, device=str(buffer.device))) + def write_joint_viscous_friction_coefficient_to_sim( + self, + joint_viscous_friction_coeff: torch.Tensor | wp.array | float, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Deprecated, same as :meth:`write_joint_viscous_friction_coefficient_to_sim_index`.""" + warnings.warn( + "The function 'write_joint_viscous_friction_coefficient_to_sim' will be deprecated in a future release. " + "Please use 'write_joint_viscous_friction_coefficient_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + self.write_joint_viscous_friction_coefficient_to_sim_index( + joint_viscous_friction_coeff=joint_viscous_friction_coeff, + joint_ids=joint_ids, + env_ids=env_ids, + ) diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation_data.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation_data.py index 5de3057e56cd..3ada3b36e63b 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation_data.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation_data.py @@ -3,11 +3,22 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Articulation data backed by ovphysx TensorBindingsAPI.""" +"""OVPhysX-backed ArticulationData implementation. + +Mirrors the post-refactor :class:`~isaaclab_ovphysx.assets.RigidObjectData` +shape with the API surface coming from +:class:`isaaclab_newton.assets.ArticulationData`: eager buffer allocation +in ``_create_buffers``; pinned-host CPU staging for CPU-only bindings via +``_binding_read`` (the PR #5329 pattern); ``ProxyArray`` returns from every +public property; timestamp-cached pull semantics through +``_read_*_binding``. +""" from __future__ import annotations +import logging import warnings +from collections.abc import Callable from typing import Any import numpy as np @@ -25,106 +36,118 @@ _projected_gravity, _world_vel_to_body_ang, _world_vel_to_body_lin, + concat_body_pose_and_vel_to_state, + concat_root_pose_and_vel_to_state, + get_body_com_pose_from_body_link_pose, + get_body_link_vel_from_body_com_vel, + vec13f, ) -from .kernels import _compose_body_com_poses, _fd_joint_acc - +from .kernels import _fd_joint_acc -class ArticulationData(BaseArticulationData): - """Data container for an articulation backed by ovphysx tensor bindings. +logger = logging.getLogger(__name__) - This class contains the data for an articulation in the simulation. The data includes the state of - the root rigid body, the state of all the bodies in the articulation, and the joint state. The data is - stored in the simulation world frame unless otherwise specified. - An articulation is comprised of multiple rigid bodies or links. For a rigid body, there are two frames - of reference that are used: +class ArticulationData(BaseArticulationData): + """Data container for an articulation backed by OVPhysX tensor bindings. - - Actor frame: The frame of reference of the rigid body prim. This typically corresponds to the Xform prim - with the rigid body schema. - - Center of mass frame: The frame of reference of the center of mass of the rigid body. + Reads simulation state on demand through ``ovphysx`` ``TensorBinding`` + objects keyed by :data:`~isaaclab_ovphysx.tensor_types.ARTICULATION_*`. + Buffers are allocated eagerly in :meth:`_create_buffers` and timestamped + so each binding is read at most once per simulation step. - Depending on the settings, the two frames may not coincide with each other. In the robotics sense, the actor frame - can be interpreted as the link frame. + CPU-only bindings (``BODY_MASS``, ``BODY_COM_POSE``, + ``BODY_INERTIA``, every ``DOF_*`` property except state) route through + pinned-host staging buffers via :meth:`_binding_read`, matching the + PR #5329 PhysX pattern. - Uses ovphysx :class:`TensorBinding` objects to lazily read simulation state into warp - arrays. Writes happen via the :class:`Articulation` class. + Properties return :class:`~isaaclab.utils.warp.ProxyArray` for time-varying + state and raw :class:`wp.array` for one-shot config buffers (defaults). """ __backend_name__: str = "ovphysx" """The name of the backend for the articulation data.""" - def __init__(self, bindings: dict[int, Any], device: str, binding_getter=None): - """Initialize the articulation data. + def __init__( + self, + bindings: dict[int, Any], + device: str, + num_instances: int, + num_bodies: int, + num_joints: int, + num_fixed_tendons: int, + num_spatial_tendons: int, + body_names: list[str], + joint_names: list[str], + fixed_tendon_names: list[str], + spatial_tendon_names: list[str], + binding_getter: Callable[[int], Any] | None = None, + ) -> None: + """Initialize the articulation data container. Args: - bindings: Mapping from ovphysx tensor type constant to a - live TensorBinding for this articulation. - device: The compute device (``"cpu"`` or ``"cuda:N"``). + bindings: Dictionary of OVPhysX :class:`TensorBinding` objects keyed + by :class:`isaaclab_ovphysx.tensor_types.TensorType`. + device: Simulation device string (e.g., ``"cuda:0"`` or ``"cpu"``). + num_instances: Number of articulation instances. + num_bodies: Number of bodies per articulation. + num_joints: Number of degrees of freedom per articulation. + num_fixed_tendons: Number of fixed tendons per articulation. + num_spatial_tendons: Number of spatial tendons per articulation. + body_names: Ordered list of body names. + joint_names: Ordered list of joint names. + fixed_tendon_names: Ordered list of fixed tendon names. + spatial_tendon_names: Ordered list of spatial tendon names. binding_getter: Optional callable(tensor_type) -> TensorBinding that lazily creates bindings on first access. When provided, - ``_get_binding()`` delegates to this instead of only checking + :meth:`_get_binding` delegates to this instead of only checking the static ``bindings`` dict. """ super().__init__(root_view=None, device=device) self._bindings = bindings self._binding_getter = binding_getter + # Counts and names are PLAIN INSTANCE ATTRIBUTES (not @property), + # mirroring the post-audit RigidObjectData demotion. + self.num_instances = num_instances + self.num_bodies = num_bodies + self.num_joints = num_joints + self.num_fixed_tendons = num_fixed_tendons + self.num_spatial_tendons = num_spatial_tendons + self.body_names = body_names + self.joint_names = joint_names + self.fixed_tendon_names = fixed_tendon_names + self.spatial_tendon_names = spatial_tendon_names + # Internal aliases used throughout _create_buffers and property bodies. + self._num_instances = num_instances + self._num_bodies = num_bodies + self._num_joints = num_joints + self._num_fixed_tendons = num_fixed_tendons + self._num_spatial_tendons = num_spatial_tendons + # Simulation timestamp drives the cache validity contract. self._sim_timestamp: float = 0.0 - self._is_primed = False - - # Metadata from an arbitrary articulation binding. - sample = next(iter(bindings.values())) - self._num_instances = sample.count - self._num_joints = sample.dof_count - self._num_bodies = sample.body_count - self._is_fixed_base = sample.is_fixed_base - - self.body_names = list(sample.body_names) - self.joint_names = list(sample.dof_names) - self.fixed_tendon_names: list[str] = [] - self.spatial_tendon_names: list[str] = [] - - self._num_fixed_tendons = 0 - self._num_spatial_tendons = 0 - - # Initialize parametric gravity and forward vectors (matching PhysX/Newton pattern). - # Guard against None sim context (e.g. mock/test environments). + self._is_primed: bool = False + # Pinned-host staging buffers for CPU-only bindings (keyed by tensor_type). + self._cpu_staging_buffers: dict[int, wp.array] = {} + # Scratch buffers for _get_read_view cache (keyed by (tensor_type, ptr)). + self._read_scratch: dict = {} + # Initialize gravity and forward constants (matching PhysX/Newton pattern). + gravity = (0.0, 0.0, -9.81) from isaaclab.physics import PhysicsManager - gravity = (0.0, 0.0, -9.81) if PhysicsManager._sim is not None and hasattr(PhysicsManager._sim, "cfg"): gravity = PhysicsManager._sim.cfg.gravity gravity_np = np.array(gravity, dtype=np.float32) - gravity_mag = np.linalg.norm(gravity_np) + gravity_mag = float(np.linalg.norm(gravity_np)) if gravity_mag == 0.0: gravity_dir = np.array([0.0, 0.0, -1.0], dtype=np.float32) else: gravity_dir = gravity_np / gravity_mag - gravity_dir_tiled = np.tile(gravity_dir, (self._num_instances, 1)) - forward_tiled = np.tile(np.array([1.0, 0.0, 0.0], dtype=np.float32), (self._num_instances, 1)) - + gravity_dir_tiled = np.tile(gravity_dir, (num_instances, 1)) + forward_tiled = np.tile(np.array([1.0, 0.0, 0.0], dtype=np.float32), (num_instances, 1)) self.GRAVITY_VEC_W = ProxyArray(wp.from_numpy(gravity_dir_tiled, dtype=wp.vec3f, device=device)) self.FORWARD_VEC_B = ProxyArray(wp.from_numpy(forward_tiled, dtype=wp.vec3f, device=device)) - - def update(self, dt: float) -> None: - """Update the data for the articulation. - - Args: - dt: The time step for the update [s]. This must be a positive value. - """ - self._sim_timestamp += dt - - # Finite-difference joint acceleration from velocity. - if dt > 0.0 and self._previous_joint_vel is not None: - cur_vel = self.joint_vel - wp.launch( - _fd_joint_acc, - dim=(self._num_instances, self._num_joints), - inputs=[cur_vel, self._previous_joint_vel, 1.0 / dt], - outputs=[self._joint_acc.data], - device=self.device, - ) - self._joint_acc.timestamp = self._sim_timestamp + # Allocate every TimestampedBufferWarp and pinned CPU staging buffer. + self._create_buffers() @property def is_primed(self) -> bool: @@ -148,21 +171,31 @@ def is_primed(self, value: bool) -> None: raise ValueError("The articulation data is already primed.") self._is_primed = True - """ - Names. - """ - - body_names: list[str] = None - """Body names in the order parsed by the simulation view.""" - - joint_names: list[str] = None - """Joint names in the order parsed by the simulation view.""" + def update(self, dt: float) -> None: + """Advance the simulation timestamp; trigger FD derivations. - fixed_tendon_names: list[str] = None - """Fixed tendon names in the order parsed by the simulation view.""" + Called by :class:`~isaaclab_ovphysx.assets.Articulation` each step. - spatial_tendon_names: list[str] = None - """Spatial tendon names in the order parsed by the simulation view.""" + Args: + dt: Simulation time-step [s]. + """ + self._sim_timestamp += dt + if not self._is_primed: + return + # FD joint acceleration from velocity. + if dt > 0.0 and self._previous_joint_vel is not None: + cur_vel_buf = self._joint_vel_buf + # Ensure joint vel buffer is fresh before differencing. + self._read_binding_into_buf(TT.DOF_VELOCITY, cur_vel_buf) + wp.launch( + _fd_joint_acc, + dim=(self._num_instances, self._num_joints), + inputs=[cur_vel_buf.data, self._previous_joint_vel, 1.0 / dt], + outputs=[self._joint_acc.data], + device=self.device, + ) + self._joint_acc.timestamp = self._sim_timestamp + wp.copy(self._previous_joint_vel, cur_vel_buf.data) """ Defaults - Initial state. @@ -170,10 +203,12 @@ def is_primed(self, value: bool) -> None: @property def default_root_pose(self) -> ProxyArray: - """Default root pose ``[pos, quat]`` in the local environment frame. + """Default root pose ``[pos, quat]`` in local environment frame [m, -]. - The position and quaternion are of the articulation root's actor frame. - Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). + Shape is (num_instances,), dtype = wp.transformf. + In torch this resolves to (num_instances, 7). + + Populated from :attr:`ArticulationCfg.init_state` during initialisation. """ if self._default_root_pose_ta is None: self._default_root_pose_ta = ProxyArray(self._default_root_pose) @@ -184,21 +219,23 @@ def default_root_pose(self, value: wp.array) -> None: """Set the default root pose. Args: - value: The default root pose. Shape is (num_instances, 7). + value: The default root pose, shape (num_instances, 7). Raises: ValueError: If the articulation data is already primed. """ - if self.is_primed: + if self._is_primed: raise ValueError("The articulation data is already primed.") self._default_root_pose.assign(value) @property def default_root_vel(self) -> ProxyArray: - """Default root velocity ``[lin_vel, ang_vel]`` in the local environment frame. + """Default root velocity ``[lin_vel, ang_vel]`` in local environment frame [m/s, rad/s]. + + Shape is (num_instances,), dtype = wp.spatial_vectorf. + In torch this resolves to (num_instances, 6). - The linear and angular velocities are of the articulation root's center of mass frame. - Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). + Populated from :attr:`ArticulationCfg.init_state` during initialisation. """ if self._default_root_vel_ta is None: self._default_root_vel_ta = ProxyArray(self._default_root_vel) @@ -209,12 +246,12 @@ def default_root_vel(self, value: wp.array) -> None: """Set the default root velocity. Args: - value: The default root velocity. Shape is (num_instances, 6). + value: The default root velocity, shape (num_instances, 6). Raises: ValueError: If the articulation data is already primed. """ - if self.is_primed: + if self._is_primed: raise ValueError("The articulation data is already primed.") self._default_root_vel.assign(value) @@ -222,10 +259,7 @@ def default_root_vel(self, value: wp.array) -> None: def default_joint_pos(self) -> ProxyArray: """Default joint positions of all joints [m or rad, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). - - This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. + Shape is (num_instances, num_joints), dtype = wp.float32. """ if self._default_joint_pos_ta is None: self._default_joint_pos_ta = ProxyArray(self._default_joint_pos) @@ -236,12 +270,12 @@ def default_joint_pos(self, value: wp.array) -> None: """Set the default joint positions. Args: - value: The default joint positions. Shape is (num_instances, num_joints). + value: The default joint positions, shape (num_instances, num_joints). Raises: ValueError: If the articulation data is already primed. """ - if self.is_primed: + if self._is_primed: raise ValueError("The articulation data is already primed.") self._default_joint_pos.assign(value) @@ -249,10 +283,7 @@ def default_joint_pos(self, value: wp.array) -> None: def default_joint_vel(self) -> ProxyArray: """Default joint velocities of all joints [m/s or rad/s, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). - - This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. + Shape is (num_instances, num_joints), dtype = wp.float32. """ if self._default_joint_vel_ta is None: self._default_joint_vel_ta = ProxyArray(self._default_joint_vel) @@ -263,12 +294,12 @@ def default_joint_vel(self, value: wp.array) -> None: """Set the default joint velocities. Args: - value: The default joint velocities. Shape is (num_instances, num_joints). + value: The default joint velocities, shape (num_instances, num_joints). Raises: ValueError: If the articulation data is already primed. """ - if self.is_primed: + if self._is_primed: raise ValueError("The articulation data is already primed.") self._default_joint_vel.assign(value) @@ -280,12 +311,7 @@ def default_joint_vel(self, value: wp.array) -> None: def joint_pos_target(self) -> ProxyArray: """Joint position targets commanded by the user [m or rad, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). - - For an implicit actuator model, the targets are directly set into the simulation. - For an explicit actuator model, the targets are used to compute the joint torques - (see :attr:`applied_torque`), which are then set into the simulation. + Shape is (num_instances, num_joints), dtype = wp.float32. """ if self._joint_pos_target_ta is None: self._joint_pos_target_ta = ProxyArray(self._joint_pos_target) @@ -295,12 +321,7 @@ def joint_pos_target(self) -> ProxyArray: def joint_vel_target(self) -> ProxyArray: """Joint velocity targets commanded by the user [m/s or rad/s, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). - - For an implicit actuator model, the targets are directly set into the simulation. - For an explicit actuator model, the targets are used to compute the joint torques - (see :attr:`applied_torque`), which are then set into the simulation. + Shape is (num_instances, num_joints), dtype = wp.float32. """ if self._joint_vel_target_ta is None: self._joint_vel_target_ta = ProxyArray(self._joint_vel_target) @@ -310,12 +331,7 @@ def joint_vel_target(self) -> ProxyArray: def joint_effort_target(self) -> ProxyArray: """Joint effort targets commanded by the user [N or N*m, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). - - For an implicit actuator model, the targets are directly set into the simulation. - For an explicit actuator model, the targets are used to compute the joint torques - (see :attr:`applied_torque`), which are then set into the simulation. + Shape is (num_instances, num_joints), dtype = wp.float32. """ if self._joint_effort_target_ta is None: self._joint_effort_target_ta = ProxyArray(self._joint_effort_target) @@ -329,12 +345,7 @@ def joint_effort_target(self) -> ProxyArray: def computed_torque(self) -> ProxyArray: """Joint torques computed from the actuator model (before clipping) [N*m]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). - - This quantity is the raw torque output from the actuator model, before any clipping is applied. - It is exposed for users who want to inspect the computations inside the actuator model. - For instance, to penalize the learning agent for a difference between the computed and applied torques. + Shape is (num_instances, num_joints), dtype = wp.float32. """ if self._computed_torque_ta is None: self._computed_torque_ta = ProxyArray(self._computed_torque) @@ -344,11 +355,7 @@ def computed_torque(self) -> ProxyArray: def applied_torque(self) -> ProxyArray: """Joint torques applied from the actuator model (after clipping) [N*m]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). - - These torques are set into the simulation, after clipping the :attr:`computed_torque` based on the - actuator model. + Shape is (num_instances, num_joints), dtype = wp.float32. """ if self._applied_torque_ta is None: self._applied_torque_ta = ProxyArray(self._applied_torque) @@ -360,85 +367,132 @@ def applied_torque(self) -> ProxyArray: @property def joint_stiffness(self) -> ProxyArray: - """Joint stiffness provided to the simulation. + """Joint stiffness provided to the simulation [N*m/rad or N/m, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + Shape is (num_instances, num_joints), dtype = wp.float32. - In the case of explicit actuators, the value for the corresponding joints is zero. + Routed through pinned-host staging because ``DOF_STIFFNESS`` is a + CPU-only OVPhysX binding. """ + self._read_scalar_binding(TT.DOF_STIFFNESS, self._joint_stiffness) if self._joint_stiffness_ta is None: - self._joint_stiffness_ta = ProxyArray(self._joint_stiffness) + self._joint_stiffness_ta = ProxyArray(self._joint_stiffness.data) return self._joint_stiffness_ta @property def joint_damping(self) -> ProxyArray: - """Joint damping provided to the simulation. + """Joint damping provided to the simulation [N*m*s/rad or N*s/m, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + Shape is (num_instances, num_joints), dtype = wp.float32. - In the case of explicit actuators, the value for the corresponding joints is zero. + Routed through pinned-host staging because ``DOF_DAMPING`` is a + CPU-only OVPhysX binding. """ + self._read_scalar_binding(TT.DOF_DAMPING, self._joint_damping) if self._joint_damping_ta is None: - self._joint_damping_ta = ProxyArray(self._joint_damping) + self._joint_damping_ta = ProxyArray(self._joint_damping.data) return self._joint_damping_ta @property def joint_armature(self) -> ProxyArray: - """Joint armature provided to the simulation. + """Joint armature provided to the simulation [kg*m^2]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + Shape is (num_instances, num_joints), dtype = wp.float32. + + Routed through pinned-host staging because ``DOF_ARMATURE`` is a + CPU-only OVPhysX binding. """ + self._read_scalar_binding(TT.DOF_ARMATURE, self._joint_armature) if self._joint_armature_ta is None: - self._joint_armature_ta = ProxyArray(self._joint_armature) + self._joint_armature_ta = ProxyArray(self._joint_armature.data) return self._joint_armature_ta @property def joint_friction_coeff(self) -> ProxyArray: - """Joint static friction coefficient provided to the simulation. + """Joint static friction coefficient [dimensionless]. + + Shape is (num_instances, num_joints), dtype = wp.float32. + Component ``[..., 0]`` of the ``DOF_FRICTION_PROPERTIES`` binding. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + Routed through pinned-host staging because ``DOF_FRICTION_PROPERTIES`` + is a CPU-only OVPhysX binding. """ + self._read_scalar_binding(TT.DOF_FRICTION_PROPERTIES, self._joint_friction_props_buf) if self._joint_friction_coeff_ta is None: self._joint_friction_coeff_ta = ProxyArray(self._joint_friction_coeff) return self._joint_friction_coeff_ta + @property + def joint_dynamic_friction_coeff(self) -> ProxyArray: + """Joint dynamic friction coefficient [dimensionless]. + + Shape is (num_instances, num_joints), dtype = wp.float32. + Component ``[..., 1]`` of the ``DOF_FRICTION_PROPERTIES`` binding. + + Routed through pinned-host staging because ``DOF_FRICTION_PROPERTIES`` + is a CPU-only OVPhysX binding. + """ + self._read_scalar_binding(TT.DOF_FRICTION_PROPERTIES, self._joint_friction_props_buf) + if self._joint_dynamic_friction_coeff_ta is None: + self._joint_dynamic_friction_coeff_ta = ProxyArray(self._joint_dynamic_friction_coeff) + return self._joint_dynamic_friction_coeff_ta + + @property + def joint_viscous_friction_coeff(self) -> ProxyArray: + """Joint viscous friction coefficient [N*m*s/rad or N*s/m, depending on joint type]. + + Shape is (num_instances, num_joints), dtype = wp.float32. + Component ``[..., 2]`` of the ``DOF_FRICTION_PROPERTIES`` binding. + + Routed through pinned-host staging because ``DOF_FRICTION_PROPERTIES`` + is a CPU-only OVPhysX binding. + """ + self._read_scalar_binding(TT.DOF_FRICTION_PROPERTIES, self._joint_friction_props_buf) + if self._joint_viscous_friction_coeff_ta is None: + self._joint_viscous_friction_coeff_ta = ProxyArray(self._joint_viscous_friction_coeff) + return self._joint_viscous_friction_coeff_ta + @property def joint_pos_limits(self) -> ProxyArray: - """Joint position limits provided to the simulation. + """Joint position limits provided to the simulation [m or rad, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.vec2f. In torch this resolves to - (num_instances, num_joints, 2). + Shape is (num_instances, num_joints), dtype = wp.vec2f. + In torch this resolves to (num_instances, num_joints, 2). - The limits are in the order :math:`[lower, upper]`. + The limits are in the order :math:`[lower, upper]`. Routed through + pinned-host staging because ``DOF_LIMIT`` is a CPU-only OVPhysX binding. """ + self._read_scalar_binding(TT.DOF_LIMIT, self._joint_pos_limits) if self._joint_pos_limits_ta is None: - self._joint_pos_limits_ta = ProxyArray(self._joint_pos_limits) + self._joint_pos_limits_ta = ProxyArray(self._joint_pos_limits.data) return self._joint_pos_limits_ta @property def joint_vel_limits(self) -> ProxyArray: """Joint maximum velocity provided to the simulation [m/s or rad/s, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + Shape is (num_instances, num_joints), dtype = wp.float32. + + Routed through pinned-host staging because ``DOF_MAX_VELOCITY`` is a + CPU-only OVPhysX binding. """ + self._read_scalar_binding(TT.DOF_MAX_VELOCITY, self._joint_vel_limits) if self._joint_vel_limits_ta is None: - self._joint_vel_limits_ta = ProxyArray(self._joint_vel_limits) + self._joint_vel_limits_ta = ProxyArray(self._joint_vel_limits.data) return self._joint_vel_limits_ta @property def joint_effort_limits(self) -> ProxyArray: """Joint maximum effort provided to the simulation [N or N*m, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + Shape is (num_instances, num_joints), dtype = wp.float32. + + Routed through pinned-host staging because ``DOF_MAX_FORCE`` is a + CPU-only OVPhysX binding. """ + self._read_scalar_binding(TT.DOF_MAX_FORCE, self._joint_effort_limits) if self._joint_effort_limits_ta is None: - self._joint_effort_limits_ta = ProxyArray(self._joint_effort_limits) + self._joint_effort_limits_ta = ProxyArray(self._joint_effort_limits.data) return self._joint_effort_limits_ta """ @@ -447,25 +501,12 @@ def joint_effort_limits(self) -> ProxyArray: @property def soft_joint_pos_limits(self) -> ProxyArray: - r"""Soft joint position limits for all joints. - - Shape is (num_instances, num_joints), dtype = wp.vec2f. In torch this resolves to - (num_instances, num_joints, 2). - - The limits are in the order :math:`[lower, upper]`. The soft joint position limits are computed as - a sub-region of the :attr:`joint_pos_limits` based on the - :attr:`~isaaclab.assets.ArticulationCfg.soft_joint_pos_limit_factor` parameter. - - Consider the joint position limits :math:`[lower, upper]` and the soft joint position limits - :math:`[soft\_lower, soft\_upper]`. The soft joint position limits are computed as: - - .. math:: + r"""Soft joint position limits for all joints [m or rad, depending on joint type]. - soft\_lower = (lower + upper) / 2 - factor * (upper - lower) / 2 - soft\_upper = (lower + upper) / 2 + factor * (upper - lower) / 2 + Shape is (num_instances, num_joints), dtype = wp.vec2f. + In torch this resolves to (num_instances, num_joints, 2). - The soft joint position limits help specify a safety region around the joint limits. It isn't used by the - simulation, but is useful for learning agents to prevent the joint positions from violating the limits. + The limits are in the order :math:`[lower, upper]`. """ if self._soft_joint_pos_limits_ta is None: self._soft_joint_pos_limits_ta = ProxyArray(self._soft_joint_pos_limits) @@ -473,13 +514,9 @@ def soft_joint_pos_limits(self) -> ProxyArray: @property def soft_joint_vel_limits(self) -> ProxyArray: - """Soft joint velocity limits for all joints. - - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + """Soft joint velocity limits for all joints [m/s or rad/s, depending on joint type]. - These are obtained from the actuator model. It may differ from :attr:`joint_vel_limits` if the actuator model - has a variable velocity limit model. For instance, in a variable gear ratio actuator model. + Shape is (num_instances, num_joints), dtype = wp.float32. """ if self._soft_joint_vel_limits_ta is None: self._soft_joint_vel_limits_ta = ProxyArray(self._soft_joint_vel_limits) @@ -489,8 +526,7 @@ def soft_joint_vel_limits(self) -> ProxyArray: def gear_ratio(self) -> ProxyArray: """Gear ratio for relating motor torques to applied joint torques. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + Shape is (num_instances, num_joints), dtype = wp.float32. """ if self._gear_ratio_ta is None: self._gear_ratio_ta = ProxyArray(self._gear_ratio) @@ -502,68 +538,84 @@ def gear_ratio(self) -> ProxyArray: @property def fixed_tendon_stiffness(self) -> ProxyArray: - """Fixed tendon stiffness provided to the simulation. + """Fixed-tendon stiffness gains [N*m/rad]. - Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to - (num_instances, num_fixed_tendons). + Shape is (num_instances, num_fixed_tendons), dtype = ``wp.float32``. + + Routed through pinned-host staging (CPU-only binding). """ + self._read_scalar_binding(TT.FIXED_TENDON_STIFFNESS, self._fixed_tendon_stiffness) if self._fixed_tendon_stiffness_ta is None: - self._fixed_tendon_stiffness_ta = ProxyArray(self._fixed_tendon_stiffness) + self._fixed_tendon_stiffness_ta = ProxyArray(self._fixed_tendon_stiffness.data) return self._fixed_tendon_stiffness_ta @property def fixed_tendon_damping(self) -> ProxyArray: - """Fixed tendon damping provided to the simulation. + """Fixed-tendon damping coefficients [N*m*s/rad]. + + Shape is (num_instances, num_fixed_tendons), dtype = ``wp.float32``. - Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to - (num_instances, num_fixed_tendons). + Routed through pinned-host staging (CPU-only binding). """ + self._read_scalar_binding(TT.FIXED_TENDON_DAMPING, self._fixed_tendon_damping) if self._fixed_tendon_damping_ta is None: - self._fixed_tendon_damping_ta = ProxyArray(self._fixed_tendon_damping) + self._fixed_tendon_damping_ta = ProxyArray(self._fixed_tendon_damping.data) return self._fixed_tendon_damping_ta @property def fixed_tendon_limit_stiffness(self) -> ProxyArray: - """Fixed tendon limit stiffness provided to the simulation. + """Fixed-tendon limit stiffness [N*m/rad]. - Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to - (num_instances, num_fixed_tendons). + Shape is (num_instances, num_fixed_tendons), dtype = ``wp.float32``. + + Routed through pinned-host staging (CPU-only binding). """ + self._read_scalar_binding(TT.FIXED_TENDON_LIMIT_STIFFNESS, self._fixed_tendon_limit_stiffness) if self._fixed_tendon_limit_stiffness_ta is None: - self._fixed_tendon_limit_stiffness_ta = ProxyArray(self._fixed_tendon_limit_stiffness) + self._fixed_tendon_limit_stiffness_ta = ProxyArray(self._fixed_tendon_limit_stiffness.data) return self._fixed_tendon_limit_stiffness_ta @property def fixed_tendon_rest_length(self) -> ProxyArray: - """Fixed tendon rest length provided to the simulation. + """Fixed-tendon rest lengths [m]. + + Shape is (num_instances, num_fixed_tendons), dtype = ``wp.float32``. - Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to - (num_instances, num_fixed_tendons). + Routed through pinned-host staging (CPU-only binding). """ + self._read_scalar_binding(TT.FIXED_TENDON_REST_LENGTH, self._fixed_tendon_rest_length) if self._fixed_tendon_rest_length_ta is None: - self._fixed_tendon_rest_length_ta = ProxyArray(self._fixed_tendon_rest_length) + self._fixed_tendon_rest_length_ta = ProxyArray(self._fixed_tendon_rest_length.data) return self._fixed_tendon_rest_length_ta @property def fixed_tendon_offset(self) -> ProxyArray: - """Fixed tendon offset provided to the simulation. + """Fixed-tendon offsets [m]. + + Shape is (num_instances, num_fixed_tendons), dtype = ``wp.float32``. - Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to - (num_instances, num_fixed_tendons). + Routed through pinned-host staging (CPU-only binding). """ + self._read_scalar_binding(TT.FIXED_TENDON_OFFSET, self._fixed_tendon_offset) if self._fixed_tendon_offset_ta is None: - self._fixed_tendon_offset_ta = ProxyArray(self._fixed_tendon_offset) + self._fixed_tendon_offset_ta = ProxyArray(self._fixed_tendon_offset.data) return self._fixed_tendon_offset_ta @property def fixed_tendon_pos_limits(self) -> ProxyArray: - """Fixed tendon position limits provided to the simulation. + """Fixed tendon position limits provided to the simulation [m or rad]. - Shape is (num_instances, num_fixed_tendons), dtype = wp.vec2f. In torch this resolves to - (num_instances, num_fixed_tendons, 2). + Shape is (num_instances, num_fixed_tendons), dtype = ``wp.vec2f``. + In torch this resolves to (num_instances, num_fixed_tendons, 2). + + .. deprecated:: + Use :attr:`fixed_tendon_limit` (shape ``(N, T, 2)``, dtype + ``wp.float32``) instead. This alias is kept for backwards + compatibility and reads the same underlying data. """ + self._read_scalar_binding(TT.FIXED_TENDON_LIMIT, self._fixed_tendon_pos_limits) if self._fixed_tendon_pos_limits_ta is None: - self._fixed_tendon_pos_limits_ta = ProxyArray(self._fixed_tendon_pos_limits) + self._fixed_tendon_pos_limits_ta = ProxyArray(self._fixed_tendon_pos_limits.data) return self._fixed_tendon_pos_limits_ta """ @@ -572,46 +624,54 @@ def fixed_tendon_pos_limits(self) -> ProxyArray: @property def spatial_tendon_stiffness(self) -> ProxyArray: - """Spatial tendon stiffness provided to the simulation. + """Spatial-tendon stiffness gains [N/m]. + + Shape is (num_instances, num_spatial_tendons), dtype = ``wp.float32``. - Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to - (num_instances, num_spatial_tendons). + Routed through pinned-host staging (CPU-only binding). """ + self._read_scalar_binding(TT.SPATIAL_TENDON_STIFFNESS, self._spatial_tendon_stiffness) if self._spatial_tendon_stiffness_ta is None: - self._spatial_tendon_stiffness_ta = ProxyArray(self._spatial_tendon_stiffness) + self._spatial_tendon_stiffness_ta = ProxyArray(self._spatial_tendon_stiffness.data) return self._spatial_tendon_stiffness_ta @property def spatial_tendon_damping(self) -> ProxyArray: - """Spatial tendon damping provided to the simulation. + """Spatial-tendon damping coefficients [N*s/m]. + + Shape is (num_instances, num_spatial_tendons), dtype = ``wp.float32``. - Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to - (num_instances, num_spatial_tendons). + Routed through pinned-host staging (CPU-only binding). """ + self._read_scalar_binding(TT.SPATIAL_TENDON_DAMPING, self._spatial_tendon_damping) if self._spatial_tendon_damping_ta is None: - self._spatial_tendon_damping_ta = ProxyArray(self._spatial_tendon_damping) + self._spatial_tendon_damping_ta = ProxyArray(self._spatial_tendon_damping.data) return self._spatial_tendon_damping_ta @property def spatial_tendon_limit_stiffness(self) -> ProxyArray: - """Spatial tendon limit stiffness provided to the simulation. + """Spatial-tendon limit stiffness [N/m]. - Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to - (num_instances, num_spatial_tendons). + Shape is (num_instances, num_spatial_tendons), dtype = ``wp.float32``. + + Routed through pinned-host staging (CPU-only binding). """ + self._read_scalar_binding(TT.SPATIAL_TENDON_LIMIT_STIFFNESS, self._spatial_tendon_limit_stiffness) if self._spatial_tendon_limit_stiffness_ta is None: - self._spatial_tendon_limit_stiffness_ta = ProxyArray(self._spatial_tendon_limit_stiffness) + self._spatial_tendon_limit_stiffness_ta = ProxyArray(self._spatial_tendon_limit_stiffness.data) return self._spatial_tendon_limit_stiffness_ta @property def spatial_tendon_offset(self) -> ProxyArray: - """Spatial tendon offset provided to the simulation. + """Spatial-tendon offsets [m]. + + Shape is (num_instances, num_spatial_tendons), dtype = ``wp.float32``. - Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to - (num_instances, num_spatial_tendons). + Routed through pinned-host staging (CPU-only binding). """ + self._read_scalar_binding(TT.SPATIAL_TENDON_OFFSET, self._spatial_tendon_offset) if self._spatial_tendon_offset_ta is None: - self._spatial_tendon_offset_ta = ProxyArray(self._spatial_tendon_offset) + self._spatial_tendon_offset_ta = ProxyArray(self._spatial_tendon_offset.data) return self._spatial_tendon_offset_ta """ @@ -620,8 +680,10 @@ def spatial_tendon_offset(self) -> ProxyArray: @property def root_link_pose_w(self) -> ProxyArray: - """Root link pose ``[pos, quat]`` in simulation world frame. - Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). + """Root link pose ``[pos, quat]`` in simulation world frame [m, -]. + + Shape is (num_instances,), dtype = wp.transformf. + In torch this resolves to (num_instances, 7). This quantity is the pose of the articulation root's actor frame relative to the world. The orientation is provided in (x, y, z, w) format. @@ -631,10 +693,21 @@ def root_link_pose_w(self) -> ProxyArray: self._root_link_pose_w_ta = ProxyArray(self._root_link_pose_w.data) return self._root_link_pose_w_ta + @property + def root_pose_w(self) -> ProxyArray: + """Alias for :attr:`root_link_pose_w` matching Newton's convention. + + Shape is (num_instances,), dtype = wp.transformf. + In torch this resolves to (num_instances, 7). + """ + return self.root_link_pose_w + @property def root_link_vel_w(self) -> ProxyArray: - """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. - Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). + """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame [m/s, rad/s]. + + Shape is (num_instances,), dtype = wp.spatial_vectorf. + In torch this resolves to (num_instances, 6). This quantity contains the linear and angular velocities of the articulation root's actor frame relative to the world. @@ -645,7 +718,7 @@ def root_link_vel_w(self) -> ProxyArray: if self._root_link_vel_w.timestamp < self._sim_timestamp: wp.launch( _copy_first_body, - dim=self._num_instances, + dim=self.num_instances, inputs=[self._body_link_vel_w.data], outputs=[self._root_link_vel_w.data], device=self.device, @@ -657,8 +730,10 @@ def root_link_vel_w(self) -> ProxyArray: @property def root_com_pose_w(self) -> ProxyArray: - """Root center of mass pose ``[pos, quat]`` in simulation world frame. - Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). + """Root center of mass pose ``[pos, quat]`` in simulation world frame [m, -]. + + Shape is (num_instances,), dtype = wp.transformf. + In torch this resolves to (num_instances, 7). This quantity is the pose of the articulation root's center of mass frame relative to the world. The orientation is provided in (x, y, z, w) format. @@ -666,7 +741,7 @@ def root_com_pose_w(self) -> ProxyArray: if self._root_com_pose_w.timestamp < self._sim_timestamp: wp.launch( _compose_root_com_pose, - dim=self._num_instances, + dim=self.num_instances, inputs=[self.root_link_pose_w, self.body_com_pose_b], outputs=[self._root_com_pose_w.data], device=self.device, @@ -678,8 +753,10 @@ def root_com_pose_w(self) -> ProxyArray: @property def root_com_vel_w(self) -> ProxyArray: - """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. - Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). + """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame [m/s, rad/s]. + + Shape is (num_instances,), dtype = wp.spatial_vectorf. + In torch this resolves to (num_instances, 6). This quantity contains the linear and angular velocities of the articulation root's center of mass frame relative to the world. @@ -695,33 +772,39 @@ def root_com_vel_w(self) -> ProxyArray: @property def body_mass(self) -> ProxyArray: - """Body mass in the world frame [kg]. + """Body masses [kg]. + + Shape is (num_instances, num_bodies), dtype = ``wp.float32``. - Shape is (num_instances, num_bodies), dtype = wp.float32. In torch this resolves to - (num_instances, num_bodies). + Routed through pinned-host staging because the underlying OVPhysX + binding is CPU-only (``ARTICULATION_BODY_MASS``). """ + self._read_scalar_binding(TT.BODY_MASS, self._body_mass) if self._body_mass_ta is None: - self._body_mass_ta = ProxyArray(self._body_mass) + self._body_mass_ta = ProxyArray(self._body_mass.data) return self._body_mass_ta @property def body_inertia(self) -> ProxyArray: - """Flattened body inertia in the world frame [kg*m^2]. + """Body inertia tensors [kg*m^2]. - Shape is (num_instances, num_bodies, 9), dtype = wp.float32. In torch this resolves to - (num_instances, num_bodies, 9). + Shape is (num_instances, num_bodies, 9), dtype = ``wp.float32``; the + trailing 9 is the row-major 3×3 inertia tensor. - Stored as a flattened 3x3 inertia matrix per body. + Routed through pinned-host staging (``ARTICULATION_BODY_INERTIA`` is + a CPU-only binding). """ + self._read_scalar_binding(TT.BODY_INERTIA, self._body_inertia) if self._body_inertia_ta is None: - self._body_inertia_ta = ProxyArray(self._body_inertia) + self._body_inertia_ta = ProxyArray(self._body_inertia.data) return self._body_inertia_ta @property def body_link_pose_w(self) -> ProxyArray: - """Body link pose ``[pos, quat]`` in simulation world frame. - Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to - (num_instances, num_bodies, 7). + """Body link pose ``[pos, quat]`` in simulation world frame [m, -]. + + Shape is (num_instances, num_bodies), dtype = wp.transformf. + In torch this resolves to (num_instances, num_bodies, 7). This quantity is the pose of the articulation links' actor frame relative to the world. The orientation is provided in (x, y, z, w) format. @@ -732,61 +815,81 @@ def body_link_pose_w(self) -> ProxyArray: return self._body_link_pose_w_ta @property - def body_link_vel_w(self) -> ProxyArray: - """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. - Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to - (num_instances, num_bodies, 6). + def body_com_vel_w(self) -> ProxyArray: + """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame [m/s, rad/s]. - This quantity contains the linear and angular velocities of the articulation links' actor frame - relative to the world. + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. + In torch this resolves to (num_instances, num_bodies, 6). """ - self._read_spatial_vector_binding(TT.LINK_VELOCITY, self._body_link_vel_w) + self._read_spatial_vector_binding(TT.LINK_VELOCITY, self._body_com_vel_w) + if self._body_com_vel_w_ta is None: + self._body_com_vel_w_ta = ProxyArray(self._body_com_vel_w.data) + return self._body_com_vel_w_ta + + @property + def body_link_vel_w(self) -> ProxyArray: + """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame [m/s, rad/s]. + + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. + In torch this resolves to (num_instances, num_bodies, 6). + + Derived from :attr:`body_com_vel_w` and :attr:`body_com_pose_b` via + :func:`~isaaclab_ovphysx.assets.kernels.get_body_link_vel_from_body_com_vel`. + """ + if self._body_link_vel_w.timestamp >= self._sim_timestamp: + if self._body_link_vel_w_ta is None: + self._body_link_vel_w_ta = ProxyArray(self._body_link_vel_w.data) + return self._body_link_vel_w_ta + _ = self.body_com_vel_w + _ = self.body_link_pose_w + _ = self.body_com_pose_b + wp.launch( + get_body_link_vel_from_body_com_vel, + dim=(self.num_instances, self.num_bodies), + inputs=[self._body_com_vel_w.data, self._body_link_pose_w.data, self._body_com_pose_b.data], + outputs=[self._body_link_vel_w.data], + device=self.device, + ) + self._body_link_vel_w.timestamp = self._sim_timestamp if self._body_link_vel_w_ta is None: self._body_link_vel_w_ta = ProxyArray(self._body_link_vel_w.data) return self._body_link_vel_w_ta @property def body_com_pose_w(self) -> ProxyArray: - """Body center of mass pose ``[pos, quat]`` in simulation world frame. - Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to - (num_instances, num_bodies, 7). + """Body center of mass pose ``[pos, quat]`` in simulation world frame [m, -]. - This quantity is the pose of the center of mass frame of the articulation links relative to the world. + Shape is (num_instances, num_bodies), dtype = wp.transformf. + In torch this resolves to (num_instances, num_bodies, 7). + + Derived from :attr:`body_link_pose_w` and :attr:`body_com_pose_b` via + :func:`~isaaclab_ovphysx.assets.kernels.get_body_com_pose_from_body_link_pose`. The orientation is provided in (x, y, z, w) format. """ - if self._body_com_pose_w.timestamp < self._sim_timestamp: - wp.launch( - _compose_body_com_poses, - dim=(self._num_instances, self._num_bodies), - inputs=[self.body_link_pose_w, self.body_com_pose_b], - outputs=[self._body_com_pose_w.data], - device=self.device, - ) - self._body_com_pose_w.timestamp = self._sim_timestamp + if self._body_com_pose_w.timestamp >= self._sim_timestamp: + if self._body_com_pose_w_ta is None: + self._body_com_pose_w_ta = ProxyArray(self._body_com_pose_w.data) + return self._body_com_pose_w_ta + _ = self.body_link_pose_w + _ = self.body_com_pose_b + wp.launch( + get_body_com_pose_from_body_link_pose, + dim=(self.num_instances, self.num_bodies), + inputs=[self._body_link_pose_w.data, self._body_com_pose_b.data], + outputs=[self._body_com_pose_w.data], + device=self.device, + ) + self._body_com_pose_w.timestamp = self._sim_timestamp if self._body_com_pose_w_ta is None: self._body_com_pose_w_ta = ProxyArray(self._body_com_pose_w.data) return self._body_com_pose_w_ta - @property - def body_com_vel_w(self) -> ProxyArray: - """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. - Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to - (num_instances, num_bodies, 6). - - This quantity contains the linear and angular velocities of the articulation links' center of mass frame - relative to the world. - - .. note:: - This is currently approximated using the link velocity. A proper COM velocity derivation - accounting for the COM offset is not yet implemented. - """ - return self.body_link_vel_w - @property def body_com_acc_w(self) -> ProxyArray: """Acceleration of all bodies center of mass ``[lin_acc, ang_acc]`` [m/s^2, rad/s^2]. - Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to - (num_instances, num_bodies, 6). + + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. + In torch this resolves to (num_instances, num_bodies, 6). All values are relative to the world. """ @@ -797,9 +900,10 @@ def body_com_acc_w(self) -> ProxyArray: @property def body_com_pose_b(self) -> ProxyArray: - """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. - Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to - (num_instances, num_bodies, 7). + """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames [m, -]. + + Shape is (num_instances, num_bodies), dtype = wp.transformf. + In torch this resolves to (num_instances, num_bodies, 7). This quantity is the pose of the center of mass frame of the rigid body relative to the body's link frame. The orientation is provided in (x, y, z, w) format. @@ -809,6 +913,23 @@ def body_com_pose_b(self) -> ProxyArray: self._body_com_pose_b_ta = ProxyArray(self._body_com_pose_b.data) return self._body_com_pose_b_ta + @property + def body_incoming_joint_wrench_b(self) -> ProxyArray: + """Incoming joint wrenches on each body in the body frame [N, N*m]. + + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. + In torch this resolves to (num_instances, num_bodies, 6). + + All body reaction wrenches are provided including the root body to the world of an articulation. + """ + self._read_spatial_vector_binding( + TT.LINK_INCOMING_JOINT_FORCE, + self._body_incoming_joint_wrench_buf, + ) + if self._body_incoming_joint_wrench_b_ta is None: + self._body_incoming_joint_wrench_b_ta = ProxyArray(self._body_incoming_joint_wrench_buf.data) + return self._body_incoming_joint_wrench_b_ta + """ Joint state properties. """ @@ -817,8 +938,7 @@ def body_com_pose_b(self) -> ProxyArray: def joint_pos(self) -> ProxyArray: """Joint positions of all joints [m or rad, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + Shape is (num_instances, num_joints), dtype = wp.float32. """ self._read_binding_into_buf(TT.DOF_POSITION, self._joint_pos_buf) if self._joint_pos_ta is None: @@ -829,8 +949,7 @@ def joint_pos(self) -> ProxyArray: def joint_vel(self) -> ProxyArray: """Joint velocities of all joints [m/s or rad/s, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + Shape is (num_instances, num_joints), dtype = wp.float32. """ self._read_binding_into_buf(TT.DOF_VELOCITY, self._joint_vel_buf) if self._joint_vel_ta is None: @@ -841,8 +960,7 @@ def joint_vel(self) -> ProxyArray: def joint_acc(self) -> ProxyArray: """Joint acceleration of all joints [m/s^2 or rad/s^2, depending on joint type]. - Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to - (num_instances, num_joints). + Shape is (num_instances, num_joints), dtype = wp.float32. .. note:: This quantity is computed via finite differencing of joint velocities. @@ -858,12 +976,13 @@ def joint_acc(self) -> ProxyArray: @property def projected_gravity_b(self) -> ProxyArray: """Projection of the gravity direction on base frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ if self._projected_gravity_b.timestamp < self._sim_timestamp: wp.launch( _projected_gravity, - dim=self._num_instances, + dim=self.num_instances, inputs=[self.GRAVITY_VEC_W, self.root_link_pose_w], outputs=[self._projected_gravity_b.data], device=self.device, @@ -875,7 +994,8 @@ def projected_gravity_b(self) -> ProxyArray: @property def heading_w(self) -> ProxyArray: - """Yaw heading of the base frame (in radians). + """Yaw heading of the base frame (in radians) [rad]. + Shape is (num_instances,), dtype = wp.float32. .. note:: @@ -885,7 +1005,7 @@ def heading_w(self) -> ProxyArray: if self._heading_w.timestamp < self._sim_timestamp: wp.launch( _compute_heading, - dim=self._num_instances, + dim=self.num_instances, inputs=[self.FORWARD_VEC_B, self.root_link_pose_w], outputs=[self._heading_w.data], device=self.device, @@ -898,6 +1018,7 @@ def heading_w(self) -> ProxyArray: @property def root_link_lin_vel_b(self) -> ProxyArray: """Root link linear velocity in base frame [m/s]. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the articulation root's actor frame with respect to its actor frame. @@ -905,7 +1026,7 @@ def root_link_lin_vel_b(self) -> ProxyArray: if self._root_link_lin_vel_b.timestamp < self._sim_timestamp: wp.launch( _world_vel_to_body_lin, - dim=self._num_instances, + dim=self.num_instances, inputs=[self.root_link_pose_w, self.root_link_vel_w], outputs=[self._root_link_lin_vel_b.data], device=self.device, @@ -918,6 +1039,7 @@ def root_link_lin_vel_b(self) -> ProxyArray: @property def root_link_ang_vel_b(self) -> ProxyArray: """Root link angular velocity in base frame [rad/s]. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the articulation root's actor frame with respect to its actor frame. @@ -925,7 +1047,7 @@ def root_link_ang_vel_b(self) -> ProxyArray: if self._root_link_ang_vel_b.timestamp < self._sim_timestamp: wp.launch( _world_vel_to_body_ang, - dim=self._num_instances, + dim=self.num_instances, inputs=[self.root_link_pose_w, self.root_link_vel_w], outputs=[self._root_link_ang_vel_b.data], device=self.device, @@ -938,6 +1060,7 @@ def root_link_ang_vel_b(self) -> ProxyArray: @property def root_com_lin_vel_b(self) -> ProxyArray: """Root center of mass linear velocity in base frame [m/s]. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the articulation root's center of mass frame @@ -946,7 +1069,7 @@ def root_com_lin_vel_b(self) -> ProxyArray: if self._root_com_lin_vel_b.timestamp < self._sim_timestamp: wp.launch( _world_vel_to_body_lin, - dim=self._num_instances, + dim=self.num_instances, inputs=[self.root_link_pose_w, self.root_com_vel_w], outputs=[self._root_com_lin_vel_b.data], device=self.device, @@ -959,6 +1082,7 @@ def root_com_lin_vel_b(self) -> ProxyArray: @property def root_com_ang_vel_b(self) -> ProxyArray: """Root center of mass angular velocity in base frame [rad/s]. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the articulation root's center of mass frame @@ -967,7 +1091,7 @@ def root_com_ang_vel_b(self) -> ProxyArray: if self._root_com_ang_vel_b.timestamp < self._sim_timestamp: wp.launch( _world_vel_to_body_ang, - dim=self._num_instances, + dim=self.num_instances, inputs=[self.root_link_pose_w, self.root_com_vel_w], outputs=[self._root_com_ang_vel_b.data], device=self.device, @@ -983,10 +1107,9 @@ def root_com_ang_vel_b(self) -> ProxyArray: @property def root_link_pos_w(self) -> ProxyArray: - """Root link position in simulation world frame. - Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + """Root link position in simulation world frame [m]. - This quantity is the position of the actor frame of the root rigid body relative to the world. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ parent = self.root_link_pose_w if self._root_link_pos_w_ta is None: @@ -996,9 +1119,8 @@ def root_link_pos_w(self) -> ProxyArray: @property def root_link_quat_w(self) -> ProxyArray: """Root link orientation (x, y, z, w) in simulation world frame. - Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). - This quantity is the orientation of the actor frame of the root rigid body. + Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). """ parent = self.root_link_pose_w if self._root_link_quat_w_ta is None: @@ -1007,10 +1129,9 @@ def root_link_quat_w(self) -> ProxyArray: @property def root_link_lin_vel_w(self) -> ProxyArray: - """Root linear velocity in simulation world frame [m/s]. - Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + """Root link linear velocity in simulation world frame [m/s]. - This quantity is the linear velocity of the root rigid body's actor frame relative to the world. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ parent = self.root_link_vel_w if self._root_link_lin_vel_w_ta is None: @@ -1020,9 +1141,8 @@ def root_link_lin_vel_w(self) -> ProxyArray: @property def root_link_ang_vel_w(self) -> ProxyArray: """Root link angular velocity in simulation world frame [rad/s]. - Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). - This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ parent = self.root_link_vel_w if self._root_link_ang_vel_w_ta is None: @@ -1031,10 +1151,9 @@ def root_link_ang_vel_w(self) -> ProxyArray: @property def root_com_pos_w(self) -> ProxyArray: - """Root center of mass position in simulation world frame. - Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + """Root center of mass position in simulation world frame [m]. - This quantity is the position of the center of mass frame of the root rigid body relative to the world. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ parent = self.root_com_pose_w if self._root_com_pos_w_ta is None: @@ -1044,9 +1163,8 @@ def root_com_pos_w(self) -> ProxyArray: @property def root_com_quat_w(self) -> ProxyArray: """Root center of mass orientation (x, y, z, w) in simulation world frame. - Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). - This quantity is the orientation of the principal axes of inertia of the root rigid body relative to the world. + Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). """ parent = self.root_com_pose_w if self._root_com_quat_w_ta is None: @@ -1056,9 +1174,8 @@ def root_com_quat_w(self) -> ProxyArray: @property def root_com_lin_vel_w(self) -> ProxyArray: """Root center of mass linear velocity in simulation world frame [m/s]. - Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). - This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ parent = self.root_com_vel_w if self._root_com_lin_vel_w_ta is None: @@ -1068,9 +1185,8 @@ def root_com_lin_vel_w(self) -> ProxyArray: @property def root_com_ang_vel_w(self) -> ProxyArray: """Root center of mass angular velocity in simulation world frame [rad/s]. - Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). - This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ parent = self.root_com_vel_w if self._root_com_ang_vel_w_ta is None: @@ -1079,11 +1195,10 @@ def root_com_ang_vel_w(self) -> ProxyArray: @property def body_link_pos_w(self) -> ProxyArray: - """Positions of all bodies in simulation world frame. - Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to - (num_instances, num_bodies, 3). + """Positions of all bodies in simulation world frame [m]. - This quantity is the position of the articulation bodies' actor frame relative to the world. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. + In torch this resolves to (num_instances, num_bodies, 3). """ parent = self.body_link_pose_w if self._body_link_pos_w_ta is None: @@ -1093,10 +1208,9 @@ def body_link_pos_w(self) -> ProxyArray: @property def body_link_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of all bodies in simulation world frame. - Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to - (num_instances, num_bodies, 4). - This quantity is the orientation of the articulation bodies' actor frame relative to the world. + Shape is (num_instances, num_bodies), dtype = wp.quatf. + In torch this resolves to (num_instances, num_bodies, 4). """ parent = self.body_link_pose_w if self._body_link_quat_w_ta is None: @@ -1106,10 +1220,9 @@ def body_link_quat_w(self) -> ProxyArray: @property def body_link_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame [m/s]. - Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to - (num_instances, num_bodies, 3). - This quantity is the linear velocity of the articulation bodies' actor frame relative to the world. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. + In torch this resolves to (num_instances, num_bodies, 3). """ parent = self.body_link_vel_w if self._body_link_lin_vel_w_ta is None: @@ -1119,10 +1232,9 @@ def body_link_lin_vel_w(self) -> ProxyArray: @property def body_link_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame [rad/s]. - Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to - (num_instances, num_bodies, 3). - This quantity is the angular velocity of the articulation bodies' actor frame relative to the world. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. + In torch this resolves to (num_instances, num_bodies, 3). """ parent = self.body_link_vel_w if self._body_link_ang_vel_w_ta is None: @@ -1131,11 +1243,10 @@ def body_link_ang_vel_w(self) -> ProxyArray: @property def body_com_pos_w(self) -> ProxyArray: - """Positions of all bodies' center of mass in simulation world frame. - Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to - (num_instances, num_bodies, 3). + """Positions of all bodies' center of mass in simulation world frame [m]. - This quantity is the position of the articulation bodies' center of mass frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. + In torch this resolves to (num_instances, num_bodies, 3). """ parent = self.body_com_pose_w if self._body_com_pos_w_ta is None: @@ -1145,10 +1256,9 @@ def body_com_pos_w(self) -> ProxyArray: @property def body_com_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. - Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to - (num_instances, num_bodies, 4). - This quantity is the orientation of the articulation bodies' principal axes of inertia. + Shape is (num_instances, num_bodies), dtype = wp.quatf. + In torch this resolves to (num_instances, num_bodies, 4). """ parent = self.body_com_pose_w if self._body_com_quat_w_ta is None: @@ -1158,10 +1268,9 @@ def body_com_quat_w(self) -> ProxyArray: @property def body_com_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame [m/s]. - Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to - (num_instances, num_bodies, 3). - This quantity is the linear velocity of the articulation bodies' center of mass frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. + In torch this resolves to (num_instances, num_bodies, 3). """ parent = self.body_com_vel_w if self._body_com_lin_vel_w_ta is None: @@ -1171,10 +1280,9 @@ def body_com_lin_vel_w(self) -> ProxyArray: @property def body_com_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame [rad/s]. - Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to - (num_instances, num_bodies, 3). - This quantity is the angular velocity of the articulation bodies' center of mass frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. + In torch this resolves to (num_instances, num_bodies, 3). """ parent = self.body_com_vel_w if self._body_com_ang_vel_w_ta is None: @@ -1184,10 +1292,9 @@ def body_com_ang_vel_w(self) -> ProxyArray: @property def body_com_lin_acc_w(self) -> ProxyArray: """Linear acceleration of all bodies in simulation world frame [m/s^2]. - Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to - (num_instances, num_bodies, 3). - This quantity is the linear acceleration of the articulation bodies' center of mass frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. + In torch this resolves to (num_instances, num_bodies, 3). """ parent = self.body_com_acc_w if self._body_com_lin_acc_w_ta is None: @@ -1197,10 +1304,9 @@ def body_com_lin_acc_w(self) -> ProxyArray: @property def body_com_ang_acc_w(self) -> ProxyArray: """Angular acceleration of all bodies in simulation world frame [rad/s^2]. - Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to - (num_instances, num_bodies, 3). - This quantity is the angular acceleration of the articulation bodies' center of mass frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. + In torch this resolves to (num_instances, num_bodies, 3). """ parent = self.body_com_acc_w if self._body_com_ang_acc_w_ta is None: @@ -1209,11 +1315,10 @@ def body_com_ang_acc_w(self) -> ProxyArray: @property def body_com_pos_b(self) -> ProxyArray: - """Center of mass position of all of the bodies in their respective link frames. - Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to - (num_instances, num_bodies, 3). + """Center of mass position of all of the bodies in their respective link frames [m]. - This quantity is the center of mass location relative to its body's link frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. + In torch this resolves to (num_instances, num_bodies, 3). """ parent = self.body_com_pose_b if self._body_com_pos_b_ta is None: @@ -1224,10 +1329,9 @@ def body_com_pos_b(self) -> ProxyArray: def body_com_quat_b(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their respective link frames. - Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to - (num_instances, num_bodies, 4). - This quantity is the orientation of the principal axes of inertia relative to its body's link frame. + Shape is (num_instances, num_bodies), dtype = wp.quatf. + In torch this resolves to (num_instances, num_bodies, 4). """ parent = self.body_com_pose_b if self._body_com_quat_b_ta is None: @@ -1235,95 +1339,12 @@ def body_com_quat_b(self) -> ProxyArray: return self._body_com_quat_b_ta """ - Deprecated in base class (required by ABC for backward compatibility). - """ - - @property - def default_root_state(self) -> ProxyArray: - """Deprecated. Use :attr:`default_root_pose` and :attr:`default_root_vel` instead.""" - warnings.warn( - "default_root_state is deprecated. Use default_root_pose and default_root_vel.", - DeprecationWarning, - stacklevel=2, - ) - if self._root_state_w_buf is None: - self._root_state_w_buf = wp.zeros( - self._num_instances, dtype=wp.types.vector(13, wp.float32), device=self.device - ) - if self._default_root_state_ta is None: - self._default_root_state_ta = ProxyArray(self._root_state_w_buf) - return self._default_root_state_ta - - @property - def root_state_w(self) -> ProxyArray: - """Deprecated. Use :attr:`root_link_pose_w` and :attr:`root_com_vel_w` instead.""" - warnings.warn( - "root_state_w is deprecated. Use root_link_pose_w and root_com_vel_w.", - DeprecationWarning, - stacklevel=2, - ) - return self.root_link_pose_w - - @property - def root_link_state_w(self) -> ProxyArray: - """Deprecated. Use :attr:`root_link_pose_w` and :attr:`root_link_vel_w` instead.""" - warnings.warn( - "root_link_state_w is deprecated. Use root_link_pose_w and root_link_vel_w.", - DeprecationWarning, - stacklevel=2, - ) - return self.root_link_pose_w - - @property - def root_com_state_w(self) -> ProxyArray: - """Deprecated. Use :attr:`root_com_pose_w` and :attr:`root_com_vel_w` instead.""" - warnings.warn( - "root_com_state_w is deprecated. Use root_com_pose_w and root_com_vel_w.", - DeprecationWarning, - stacklevel=2, - ) - return self.root_com_pose_w - - @property - def body_state_w(self) -> ProxyArray: - """Deprecated. Use :attr:`body_link_pose_w` and :attr:`body_com_vel_w` instead.""" - warnings.warn( - "body_state_w is deprecated. Use body_link_pose_w and body_com_vel_w.", - DeprecationWarning, - stacklevel=2, - ) - return self.body_link_pose_w - - @property - def body_link_state_w(self) -> ProxyArray: - """Deprecated. Use :attr:`body_link_pose_w` and :attr:`body_link_vel_w` instead.""" - warnings.warn( - "body_link_state_w is deprecated. Use body_link_pose_w and body_link_vel_w.", - DeprecationWarning, - stacklevel=2, - ) - return self.body_link_pose_w - - @property - def body_com_state_w(self) -> ProxyArray: - """Deprecated. Use :attr:`body_com_pose_w` and :attr:`body_com_vel_w` instead.""" - warnings.warn( - "body_com_state_w is deprecated. Use body_com_pose_w and body_com_vel_w.", - DeprecationWarning, - stacklevel=2, - ) - return self.body_com_pose_w - - """ - Internal helper. + Internal helpers. """ def _create_buffers(self) -> None: # noqa: C901 + """Eagerly allocate every TimestampedBuffer and pinned CPU staging buffer.""" super()._create_buffers() - # Scratch buffers for _read_binding_into_* methods, allocated lazily - # on first use and reused every subsequent step to avoid per-step - # allocation overhead on the hot RL path. - self._read_scratch: dict = {} N = self._num_instances D = self._num_joints @@ -1349,20 +1370,26 @@ def _create_buffers(self) -> None: # noqa: C901 self._joint_acc = TimestampedBuffer((N, D), dev, wp.float32) self._previous_joint_vel = wp.zeros((N, D), dtype=wp.float32, device=dev) - # -- Joint properties - self._joint_stiffness = wp.zeros((N, D), dtype=wp.float32, device=dev) - self._joint_damping = wp.zeros((N, D), dtype=wp.float32, device=dev) - self._joint_armature = wp.zeros((N, D), dtype=wp.float32, device=dev) - self._joint_friction_coeff = wp.zeros((N, D), dtype=wp.float32, device=dev) - self._joint_pos_limits = wp.zeros((N, D), dtype=wp.vec2f, device=dev) - self._joint_vel_limits = wp.zeros((N, D), dtype=wp.float32, device=dev) - self._joint_effort_limits = wp.zeros((N, D), dtype=wp.float32, device=dev) - - # -- Body properties - self._body_mass = wp.zeros((N, L), dtype=wp.float32, device=dev) - self._body_inertia = wp.zeros((N, L, 9), dtype=wp.float32, device=dev) - - # -- Soft limits / custom properties + # -- Joint properties (CPU-only; timestamped so they can be re-read after writes) + self._joint_stiffness = TimestampedBuffer((N, D), dev, wp.float32) + self._joint_damping = TimestampedBuffer((N, D), dev, wp.float32) + self._joint_armature = TimestampedBuffer((N, D), dev, wp.float32) + self._joint_pos_limits = TimestampedBuffer((N, D), dev, wp.vec2f) + self._joint_vel_limits = TimestampedBuffer((N, D), dev, wp.float32) + self._joint_effort_limits = TimestampedBuffer((N, D), dev, wp.float32) + # Friction: single (N, D, 3) TimestampedBuffer; per-component views are created lazily. + self._joint_friction_props_buf = TimestampedBuffer((N, D, 3), dev, wp.float32) + # These are strided wp.array views into _joint_friction_props_buf.data; created in + # _pin_proxy_arrays after the buffer exists. + self._joint_friction_coeff: wp.array | None = None + self._joint_dynamic_friction_coeff: wp.array | None = None + self._joint_viscous_friction_coeff: wp.array | None = None + + # -- Body properties (CPU-only; read once at init, re-read via _read_scalar_binding) + self._body_mass = TimestampedBuffer((N, L), dev, wp.float32) + self._body_inertia = TimestampedBuffer((N, L, 9), dev, wp.float32) + + # -- Soft limits / custom joint properties self._soft_joint_pos_limits = wp.zeros((N, D), dtype=wp.vec2f, device=dev) self._soft_joint_vel_limits = wp.zeros((N, D), dtype=wp.float32, device=dev) self._gear_ratio = wp.ones((N, D), dtype=wp.float32, device=dev) @@ -1388,48 +1415,180 @@ def _create_buffers(self) -> None: # noqa: C901 self._root_com_lin_vel_b = TimestampedBuffer(N, dev, wp.vec3f) self._root_com_ang_vel_b = TimestampedBuffer(N, dev, wp.vec3f) - # -- Deprecated combined state buffers - self._root_state_w_buf = None - self._root_link_state_w_buf = None - self._root_com_state_w_buf = None - self._body_state_w_buf = None - self._body_link_state_w_buf = None - self._body_com_state_w_buf = None - - # -- Tendon property buffers - T_fix = getattr(self, "_num_fixed_tendons", 0) - T_spa = getattr(self, "_num_spatial_tendons", 0) + # -- Deprecated combined state buffers (TimestampedBuffer; lazily filled on first access) + self._root_state_w_buf = TimestampedBuffer(N, dev, vec13f) + self._root_link_state_w_buf = TimestampedBuffer(N, dev, vec13f) + self._root_com_state_w_buf = TimestampedBuffer(N, dev, vec13f) + self._default_root_state_buf = wp.zeros(N, dtype=vec13f, device=dev) + # -- Deprecated body combined state buffers (TimestampedBuffer; lazily filled on first access) + self._body_state_w_buf = TimestampedBuffer((N, L), dev, vec13f) + self._body_link_state_w_buf = TimestampedBuffer((N, L), dev, vec13f) + self._body_com_state_w_buf = TimestampedBuffer((N, L), dev, vec13f) + + # -- Tendon property buffers (always allocated; empty shape when T==0 so + # properties never return None). Routed through _read_scalar_binding. + T_fix = self._num_fixed_tendons + T_spa = self._num_spatial_tendons + self._fixed_tendon_stiffness = TimestampedBuffer((N, T_fix), dev, wp.float32) + self._fixed_tendon_damping = TimestampedBuffer((N, T_fix), dev, wp.float32) + self._fixed_tendon_limit_stiffness = TimestampedBuffer((N, T_fix), dev, wp.float32) + self._fixed_tendon_rest_length = TimestampedBuffer((N, T_fix), dev, wp.float32) + self._fixed_tendon_offset = TimestampedBuffer((N, T_fix), dev, wp.float32) + # Legacy alias kept for any internal callers that used the old vec2f buffer. + self._fixed_tendon_pos_limits = TimestampedBuffer((N, T_fix), dev, wp.vec2f) + self._spatial_tendon_stiffness = TimestampedBuffer((N, T_spa), dev, wp.float32) + self._spatial_tendon_damping = TimestampedBuffer((N, T_spa), dev, wp.float32) + self._spatial_tendon_limit_stiffness = TimestampedBuffer((N, T_spa), dev, wp.float32) + self._spatial_tendon_offset = TimestampedBuffer((N, T_spa), dev, wp.float32) + + # -- CPU staging buffers for CPU-only bindings. + # Pre-allocate all of them so there is no per-step allocation on the hot path. + # These are keyed by tensor_type in self._cpu_staging_buffers; _binding_read + # selects the right one at read time. The sizes must match the binding shapes + # (flat float32). On a GPU sim the buffers are pinned-host (page-locked) so + # the wheel can dispatch async copies; on a CPU sim the staging copy is + # functionally redundant but the buffer must still exist for the write + # helpers, so we allocate unpinned and pay only the intra-CPU memcpy. + pinned = dev != "cpu" + self._cpu_body_mass = wp.zeros((N, L), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_body_coms = wp.zeros((N, L, 7), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_body_inertia = wp.zeros((N, L, 9), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_joint_stiffness = wp.zeros((N, D), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_joint_damping = wp.zeros((N, D), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_joint_position_limit = wp.zeros((N, D, 2), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_joint_velocity_limit = wp.zeros((N, D), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_joint_effort_limit = wp.zeros((N, D), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_joint_armature = wp.zeros((N, D), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_joint_friction_coeff = wp.zeros((N, D), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_joint_dynamic_friction_coeff = wp.zeros((N, D), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_joint_viscous_friction_coeff = wp.zeros((N, D), dtype=wp.float32, device="cpu", pinned=pinned) if T_fix > 0: - self._fixed_tendon_stiffness = wp.zeros((N, T_fix), dtype=wp.float32, device=dev) - self._fixed_tendon_damping = wp.zeros((N, T_fix), dtype=wp.float32, device=dev) - self._fixed_tendon_limit_stiffness = wp.zeros((N, T_fix), dtype=wp.float32, device=dev) - self._fixed_tendon_rest_length = wp.zeros((N, T_fix), dtype=wp.float32, device=dev) - self._fixed_tendon_offset = wp.zeros((N, T_fix), dtype=wp.float32, device=dev) - self._fixed_tendon_pos_limits = wp.zeros((N, T_fix), dtype=wp.vec2f, device=dev) - else: - self._fixed_tendon_stiffness = None - self._fixed_tendon_damping = None - self._fixed_tendon_limit_stiffness = None - self._fixed_tendon_rest_length = None - self._fixed_tendon_offset = None - self._fixed_tendon_pos_limits = None + self._cpu_fixed_tendon_stiffness = wp.zeros((N, T_fix), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_fixed_tendon_damping = wp.zeros((N, T_fix), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_fixed_tendon_limit_stiffness = wp.zeros((N, T_fix), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_fixed_tendon_rest_length = wp.zeros((N, T_fix), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_fixed_tendon_offset = wp.zeros((N, T_fix), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_fixed_tendon_pos_limits = wp.zeros((N, T_fix, 2), dtype=wp.float32, device="cpu", pinned=pinned) if T_spa > 0: - self._spatial_tendon_stiffness = wp.zeros((N, T_spa), dtype=wp.float32, device=dev) - self._spatial_tendon_damping = wp.zeros((N, T_spa), dtype=wp.float32, device=dev) - self._spatial_tendon_limit_stiffness = wp.zeros((N, T_spa), dtype=wp.float32, device=dev) - self._spatial_tendon_offset = wp.zeros((N, T_spa), dtype=wp.float32, device=dev) - else: - self._spatial_tendon_stiffness = None - self._spatial_tendon_damping = None - self._spatial_tendon_limit_stiffness = None - self._spatial_tendon_offset = None + self._cpu_spatial_tendon_stiffness = wp.zeros((N, T_spa), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_spatial_tendon_damping = wp.zeros((N, T_spa), dtype=wp.float32, device="cpu", pinned=pinned) + self._cpu_spatial_tendon_limit_stiffness = wp.zeros( + (N, T_spa), dtype=wp.float32, device="cpu", pinned=pinned + ) + self._cpu_spatial_tendon_offset = wp.zeros((N, T_spa), dtype=wp.float32, device="cpu", pinned=pinned) - # Read initial joint properties from bindings + # Read initial joint/body properties from bindings (one-time CPU reads). self._read_initial_properties() - - # Initialize ProxyArray wrappers (lazily created on first access) + # Initialize ProxyArray wrappers (lazily created on first property access). self._pin_proxy_arrays() + def _binding_read(self, tensor_type: int, binding: Any, dst: wp.array) -> None: + """Read *binding* into *dst*, staging through a pinned-host buffer for CPU-only bindings. + + For GPU-resident state bindings (pose, velocity, etc.) the read goes directly + into the destination array. For CPU-only property bindings (mass, COM, limits, + stiffness, …) the wheel writes into a pinned-host staging buffer first, then + :func:`wp.copy` moves the data to the simulation device asynchronously. + + Args: + tensor_type: TensorType key identifying the binding. + binding: OVPhysX TensorBinding whose ``read`` method is called. + dst: Destination :class:`wp.array` on the simulation device. + """ + if tensor_type not in TT._CPU_ONLY_TYPES or self.device == "cpu": + binding.read(dst) + return + # Route through a lazily-allocated pinned-host staging buffer. + staging = self._cpu_staging_buffers.get(tensor_type) + if staging is None: + staging = wp.zeros(binding.shape, dtype=wp.float32, device="cpu", pinned=True) + self._cpu_staging_buffers[tensor_type] = staging + binding.read(staging) + # Build a flat float32 view of dst matching the binding's flat shape. + if dst.dtype == wp.float32: + view = dst + else: + view = wp.array( + ptr=dst.ptr, + shape=binding.shape, + dtype=wp.float32, + device=str(dst.device), + copy=False, + ) + wp.copy(view, staging) + + def _binding_write( + self, + tensor_type: int, + binding: Any, + src: wp.array, + *, + indices: wp.array | None = None, + mask: wp.array | None = None, + ) -> None: + """Write *src* to *binding*, staging through pinned-host buffers for CPU-only bindings. + + Args: + tensor_type: TensorType key identifying the binding. + binding: OVPhysX TensorBinding whose ``write`` method is called. + src: Source :class:`wp.array` on the simulation device. + indices: Optional environment indices for partial writes. + mask: Optional boolean mask for partial writes. + """ + if tensor_type not in TT._CPU_ONLY_TYPES or self.device == "cpu": + binding.write(src, indices=indices, mask=mask) + return + # Stage through a pinned-host buffer. + staging = self._cpu_staging_buffers.get(tensor_type) + if staging is None: + staging = wp.zeros(binding.shape, dtype=wp.float32, device="cpu", pinned=True) + self._cpu_staging_buffers[tensor_type] = staging + if src.dtype == wp.float32: + src_view = src + else: + src_view = wp.array( + ptr=src.ptr, + shape=binding.shape, + dtype=wp.float32, + device=str(src.device), + copy=False, + ) + wp.copy(staging, src_view) + binding.write(staging, indices=indices, mask=mask) + + def _stage_to_pinned_cpu(self, tensor_type: int, role: str, src: wp.array) -> wp.array: + """Copy *src* into a lazily-allocated pinned-host :class:`wp.array`. + + Keyed on *(tensor_type, role)* so the same pair always reuses the same + buffer, avoiding per-call allocation on the hot path. + + Args: + tensor_type: TensorType identifying the binding. + role: Disambiguating string when the same tensor_type may serve + multiple purposes (e.g. ``"read"`` vs ``"write"``). + src: Source array on the simulation device. + + Returns: + Pinned-host wp.array containing a copy of *src*. + """ + key = (tensor_type, role) + staging = self._cpu_staging_buffers.get(key) # type: ignore[call-overload] + if staging is None: + if src.dtype == wp.float32: + shape = src.shape + else: + # Flatten to float32 shape matching the element byte size. + elem_floats = src.dtype.size // 4 + shape = src.shape + (elem_floats,) + staging = wp.zeros(shape, dtype=wp.float32, device="cpu", pinned=True) + self._cpu_staging_buffers[key] = staging # type: ignore[index] + if src.dtype == wp.float32: + wp.copy(staging, src) + else: + flat_src = wp.array(ptr=src.ptr, shape=staging.shape, dtype=wp.float32, device=str(src.device), copy=False) + wp.copy(staging, flat_src) + return staging + def _read_initial_properties(self) -> None: """Read static/initial joint and body properties from ovphysx bindings. @@ -1439,7 +1598,6 @@ def _read_initial_properties(self) -> None: simulation device. """ - # Property reads always use CPU numpy (property tensors are host-side). def _read_cpu(tensor_type): binding = self._get_binding(tensor_type) if binding is None: @@ -1448,75 +1606,97 @@ def _read_cpu(tensor_type): binding.read(np_buf) return np_buf - for tt, dst in [ + # Joint scalar properties — write to .data since buffers are now TimestampedBuffer. + for tt, buf in [ (TT.DOF_STIFFNESS, self._joint_stiffness), (TT.DOF_DAMPING, self._joint_damping), (TT.DOF_ARMATURE, self._joint_armature), (TT.DOF_MAX_VELOCITY, self._joint_vel_limits), (TT.DOF_MAX_FORCE, self._joint_effort_limits), - (TT.BODY_MASS, self._body_mass), ]: np_buf = _read_cpu(tt) if np_buf is not None: - wp.copy(dst, wp.from_numpy(np_buf, dtype=wp.float32, device=self.device)) + wp.copy(buf.data, wp.from_numpy(np_buf, dtype=wp.float32, device=self.device)) + buf.timestamp = self._sim_timestamp + + # Body mass (now a TimestampedBuffer). + np_buf = _read_cpu(TT.BODY_MASS) + if np_buf is not None: + wp.copy(self._body_mass.data, wp.from_numpy(np_buf, dtype=wp.float32, device=self.device)) + self._body_mass.timestamp = self._sim_timestamp - # Joint position limits: [N, D, 2] -> (N, D) wp.vec2f + # Joint position limits: [N, D, 2] -> (N, D) wp.vec2f stored in TimestampedBuffer.data np_lim = _read_cpu(TT.DOF_LIMIT) if np_lim is not None: - self._joint_pos_limits = wp.from_numpy( + src = wp.from_numpy( np_lim.reshape(self._num_instances, self._num_joints, 2), dtype=wp.vec2f, device=self.device ) + wp.copy(self._joint_pos_limits.data, src) + self._joint_pos_limits.timestamp = self._sim_timestamp - # Body inertia: [N, L, 9] + # Body inertia (now a TimestampedBuffer): [N, L, 9] np_iner = _read_cpu(TT.BODY_INERTIA) if np_iner is not None: - self._body_inertia = wp.from_numpy(np_iner, dtype=wp.float32, device=self.device) + wp.copy( + self._body_inertia.data, + wp.from_numpy(np_iner, dtype=wp.float32, device=self.device), + ) + self._body_inertia.timestamp = self._sim_timestamp - # Friction: [N, D, 3] -> extract static friction (column 0) + # Friction: [N, D, 3] -> load directly into the combined TimestampedBuffer. + # The strided per-component views (_joint_friction_coeff/dynamic/viscous) are + # created later in _pin_proxy_arrays, so we write to the combined buffer here. np_fric = _read_cpu(TT.DOF_FRICTION_PROPERTIES) if np_fric is not None: - self._joint_friction_coeff = wp.from_numpy(np_fric[..., 0].copy(), dtype=wp.float32, device=self.device) + fric_contiguous = np.ascontiguousarray(np_fric.reshape(self._num_instances, self._num_joints, 3)) + wp.copy( + self._joint_friction_props_buf.data, + wp.from_numpy(fric_contiguous, dtype=wp.float32, device=self.device), + ) + self._joint_friction_props_buf.timestamp = self._sim_timestamp - # Fixed tendon properties (CPU-side, read once) - T_fix = getattr(self, "_num_fixed_tendons", 0) + # Fixed tendon properties. PhysX exposes tendons on the simulation + # device (no ``device="cpu"`` clone in its ``set_fixed_tendon_properties`` + # call); the OVPhysX wheel mirrors that, so we read directly into the + # sim-device buffer rather than via a numpy round-trip. + T_fix = self._num_fixed_tendons if T_fix > 0: - for tt, dst in [ + for tt, buf in [ (TT.FIXED_TENDON_STIFFNESS, self._fixed_tendon_stiffness), (TT.FIXED_TENDON_DAMPING, self._fixed_tendon_damping), (TT.FIXED_TENDON_LIMIT_STIFFNESS, self._fixed_tendon_limit_stiffness), (TT.FIXED_TENDON_REST_LENGTH, self._fixed_tendon_rest_length), (TT.FIXED_TENDON_OFFSET, self._fixed_tendon_offset), ]: - np_buf = _read_cpu(tt) - if np_buf is not None and dst is not None: - wp.copy(dst, wp.from_numpy(np_buf, dtype=wp.float32, device=self.device)) - # Fixed tendon limits: [N, T, 2] -> (N, T) wp.vec2f - np_tlim = _read_cpu(TT.FIXED_TENDON_LIMIT) - if np_tlim is not None and self._fixed_tendon_pos_limits is not None: - self._fixed_tendon_pos_limits = wp.from_numpy( - np_tlim.reshape(self._num_instances, T_fix, 2), dtype=wp.vec2f, device=self.device - ) - - # Spatial tendon properties (CPU-side, read once) - T_spa = getattr(self, "_num_spatial_tendons", 0) + binding = self._get_binding(tt) + if binding is not None: + self._binding_read(tt, binding, buf.data) + buf.timestamp = self._sim_timestamp + binding = self._get_binding(TT.FIXED_TENDON_LIMIT) + if binding is not None: + self._binding_read(TT.FIXED_TENDON_LIMIT, binding, self._fixed_tendon_pos_limits.data) + self._fixed_tendon_pos_limits.timestamp = self._sim_timestamp + + # Spatial tendon properties (sim-device, see fixed-tendon comment above). + T_spa = self._num_spatial_tendons if T_spa > 0: - for tt, dst in [ + for tt, buf in [ (TT.SPATIAL_TENDON_STIFFNESS, self._spatial_tendon_stiffness), (TT.SPATIAL_TENDON_DAMPING, self._spatial_tendon_damping), (TT.SPATIAL_TENDON_LIMIT_STIFFNESS, self._spatial_tendon_limit_stiffness), (TT.SPATIAL_TENDON_OFFSET, self._spatial_tendon_offset), ]: - np_buf = _read_cpu(tt) - if np_buf is not None and dst is not None: - wp.copy(dst, wp.from_numpy(np_buf, dtype=wp.float32, device=self.device)) + binding = self._get_binding(tt) + if binding is not None: + self._binding_read(tt, binding, buf.data) + buf.timestamp = self._sim_timestamp def _pin_proxy_arrays(self) -> None: """Create pinned ProxyArray wrappers for all data buffers. - This is called once from :meth:`_create_buffers` during initialization. + Called once from :meth:`_create_buffers` during initialization. All ``_ta`` fields are lazily populated on first property access. """ - # -- Pinned ProxyArray cache (one per read property, lazily created on first access) # Defaults self._default_root_pose_ta: ProxyArray | None = None self._default_root_vel_ta: ProxyArray | None = None @@ -1534,6 +1714,8 @@ def _pin_proxy_arrays(self) -> None: self._joint_damping_ta: ProxyArray | None = None self._joint_armature_ta: ProxyArray | None = None self._joint_friction_coeff_ta: ProxyArray | None = None + self._joint_dynamic_friction_coeff_ta: ProxyArray | None = None + self._joint_viscous_friction_coeff_ta: ProxyArray | None = None self._joint_pos_limits_ta: ProxyArray | None = None self._joint_vel_limits_ta: ProxyArray | None = None self._joint_effort_limits_ta: ProxyArray | None = None @@ -1562,6 +1744,7 @@ def _pin_proxy_arrays(self) -> None: self._body_link_pose_w_ta: ProxyArray | None = None self._body_link_vel_w_ta: ProxyArray | None = None self._body_com_pose_w_ta: ProxyArray | None = None + self._body_com_vel_w_ta: ProxyArray | None = None self._body_com_acc_w_ta: ProxyArray | None = None self._body_com_pose_b_ta: ProxyArray | None = None # Body properties @@ -1605,13 +1788,75 @@ def _pin_proxy_arrays(self) -> None: self._body_com_quat_b_ta: ProxyArray | None = None # Deprecated state-concat properties self._default_root_state_ta: ProxyArray | None = None + self._root_state_w_ta: ProxyArray | None = None + self._root_link_state_w_ta: ProxyArray | None = None + self._root_com_state_w_ta: ProxyArray | None = None + # Deprecated body state-concat properties + self._body_state_w_ta: ProxyArray | None = None + self._body_link_state_w_ta: ProxyArray | None = None + self._body_com_state_w_ta: ProxyArray | None = None + + # Create strided wp.array views into _joint_friction_props_buf.data so that + # each friction component is accessible without copying data. The combined + # buffer has shape (N, D, 3) and contiguous float32 storage, so component k + # lives at byte offset k*4 with strides (D*3*4, 3*4). + N = self._num_instances + D = self._num_joints + _fp = self._joint_friction_props_buf.data + _float_bytes = 4 # sizeof(float32) + _stride_row = D * 3 * _float_bytes # bytes between rows + _stride_col = 3 * _float_bytes # bytes between columns (elements) + _dev = str(_fp.device) + self._joint_friction_coeff = wp.array( + ptr=_fp.ptr, + shape=(N, D), + strides=(_stride_row, _stride_col), + dtype=wp.float32, + device=_dev, + copy=False, + ) + self._joint_dynamic_friction_coeff = wp.array( + ptr=_fp.ptr + _float_bytes, + shape=(N, D), + strides=(_stride_row, _stride_col), + dtype=wp.float32, + device=_dev, + copy=False, + ) + self._joint_viscous_friction_coeff = wp.array( + ptr=_fp.ptr + 2 * _float_bytes, + shape=(N, D), + strides=(_stride_row, _stride_col), + dtype=wp.float32, + device=_dev, + copy=False, + ) - """ - Internal helpers -- Bindings. - """ + def _invalidate_initialize_callback(self, event) -> None: + """Invalidate cached buffers when the simulation is reinitialized. + + Args: + event: Simulation event (unused). + """ + self._is_primed = False + self._sim_timestamp = 0.0 + # Reset every TimestampedBuffer timestamp so the next property access + # triggers a fresh pull from the binding. + for attr_name in dir(self): + if attr_name.startswith("_") and not attr_name.startswith("__"): + val = getattr(self, attr_name, None) + if isinstance(val, TimestampedBuffer): + val.timestamp = -1.0 def _get_binding(self, tensor_type: int): - """Return a binding, lazily creating it if a binding_getter was provided.""" + """Return a binding, lazily creating it if a binding_getter was provided. + + Args: + tensor_type: TensorType key. + + Returns: + The TensorBinding, or ``None`` if not available. + """ b = self._bindings.get(tensor_type) if b is not None: return b @@ -1622,26 +1867,6 @@ def _get_binding(self, tensor_type: int): return b return None - def _get_read_scratch(self, tensor_type: int) -> wp.array | None: - """Return a pre-allocated flat float32 scratch buffer for a binding. - - Allocated once on first use, then reused every step. CPU-only - bindings (body properties, DOF properties) get CPU scratch; GPU - bindings get GPU scratch. wp.copy handles cross-device transfer - when the destination buffer lives on a different device. - """ - if tensor_type in self._read_scratch: - return self._read_scratch[tensor_type] - binding = self._get_binding(tensor_type) - if binding is None: - return None - from isaaclab_ovphysx.tensor_types import _CPU_ONLY_TYPES - - dev = "cpu" if tensor_type in _CPU_ONLY_TYPES else self.device - buf = wp.zeros(binding.shape, dtype=wp.float32, device=dev) - self._read_scratch[tensor_type] = buf - return buf - def _get_read_view(self, tensor_type: int, wp_array: wp.array, floats_per_elem: int = 0) -> wp.array | None: """Return a stable float32 view of a warp buffer for reading from a binding. @@ -1651,6 +1876,16 @@ def _get_read_view(self, tensor_type: int, wp_array: wp.array, floats_per_elem: The returned view is cached so that ``binding.read(view)`` sees the same object on every call, enabling the binding's internal read cache. + + Args: + tensor_type: TensorType key. + wp_array: Destination warp array. + floats_per_elem: Number of float32 elements per logical element + (e.g. 7 for transformf, 6 for spatial_vectorf). Pass 0 to + return the array as-is. + + Returns: + Float32 view suitable for ``binding.read()``, or ``None``. """ if not hasattr(self, "_read_view_cache"): self._read_view_cache = {} @@ -1678,18 +1913,13 @@ def _get_read_view(self, tensor_type: int, wp_array: wp.array, floats_per_elem: self._read_view_cache[cache_key] = view return view - def _read_binding_into_flat(self, tensor_type: int, wp_array: wp.array) -> None: - """Read a flat binding (no structured dtype) into an existing warp array. + def _read_binding_into_buf(self, tensor_type: int, buf: TimestampedBuffer) -> None: + """Read from an ovphysx binding into a :class:`TimestampedBuffer`, skipping if fresh. - Reads directly into the target array -- no scratch buffer, no extra copy. + Args: + tensor_type: TensorType key. + buf: Timestamped buffer to refresh. """ - binding = self._get_binding(tensor_type) - if binding is None: - return - binding.read(wp_array) - - def _read_binding_into_buf(self, tensor_type: int, buf: TimestampedBuffer) -> None: - """Read from an ovphysx binding into a TimestampedBuffer, skipping if fresh.""" if buf.timestamp >= self._sim_timestamp: return view = self._get_read_view(tensor_type, buf.data) @@ -1699,17 +1929,34 @@ def _read_binding_into_buf(self, tensor_type: int, buf: TimestampedBuffer) -> No buf.timestamp = self._sim_timestamp def _read_transform_binding(self, tensor_type: int, buf: TimestampedBuffer) -> None: - """Read a pose binding (float32 view of transformf buffer), skipping if fresh.""" + """Read a pose binding (float32 view of transformf buffer), skipping if fresh. + + CPU-only bindings (e.g. ``BODY_COM_POSE``) are routed through a + pinned-host staging buffer via :meth:`_binding_read` so the wheel's + device-match requirement is satisfied even on a GPU sim. + + Args: + tensor_type: TensorType key. + buf: Timestamped :class:`wp.transformf` buffer to refresh. + """ if buf.timestamp >= self._sim_timestamp: return + binding = self._get_binding(tensor_type) + if binding is None: + return view = self._get_read_view(tensor_type, buf.data, 7) if view is None: return - self._get_binding(tensor_type).read(view) + self._binding_read(tensor_type, binding, view) buf.timestamp = self._sim_timestamp def _read_spatial_vector_binding(self, tensor_type: int, buf: TimestampedBuffer) -> None: - """Read a velocity binding (float32 view of spatial_vectorf buffer), skipping if fresh.""" + """Read a velocity binding (float32 view of spatial_vectorf buffer), skipping if fresh. + + Args: + tensor_type: TensorType key. + buf: Timestamped :class:`wp.spatial_vectorf` buffer to refresh. + """ if buf.timestamp >= self._sim_timestamp: return view = self._get_read_view(tensor_type, buf.data, 6) @@ -1718,11 +1965,37 @@ def _read_spatial_vector_binding(self, tensor_type: int, buf: TimestampedBuffer) self._get_binding(tensor_type).read(view) buf.timestamp = self._sim_timestamp - """ - Internal helpers -- Extraction. - """ + def _read_scalar_binding(self, tensor_type: int, buf: TimestampedBuffer) -> None: + """Refresh a scalar or flat float32 buffer from the matching binding if stale. + + Identical timestamp-gating contract as :meth:`_read_transform_binding` + but without a structured-dtype reinterpret cast. CPU-only bindings + (e.g. ``DOF_STIFFNESS``, ``DOF_LIMIT``) are routed through a + pre-allocated pinned-host staging buffer via :meth:`_binding_read` so + the wheel's device-match requirement is satisfied even on a GPU sim. + + Args: + tensor_type: TensorType key identifying the binding. + buf: Timestamped buffer whose :attr:`~TimestampedBuffer.data` field + will be refreshed. + """ + if buf.timestamp >= self._sim_timestamp: + return + binding = self._get_binding(tensor_type) + if binding is None: + return + self._binding_read(tensor_type, binding, buf.data) + buf.timestamp = self._sim_timestamp def _get_pos_from_transform(self, transform: wp.array) -> wp.array: + """Return a position view aliased into a transform array. + + Args: + transform: Source transform array. + + Returns: + vec3f view into the position component. + """ return wp.array( ptr=transform.ptr, shape=transform.shape, @@ -1732,6 +2005,14 @@ def _get_pos_from_transform(self, transform: wp.array) -> wp.array: ) def _get_quat_from_transform(self, transform: wp.array) -> wp.array: + """Return a quaternion view aliased into a transform array. + + Args: + transform: Source transform array. + + Returns: + quatf view into the quaternion component (offset 3 floats = 12 bytes). + """ return wp.array( ptr=transform.ptr + 3 * 4, shape=transform.shape, @@ -1741,6 +2022,14 @@ def _get_quat_from_transform(self, transform: wp.array) -> wp.array: ) def _get_lin_vel_from_spatial_vector(self, sv: wp.array) -> wp.array: + """Return a linear velocity view aliased into a spatial vector array. + + Args: + sv: Source spatial vector array. + + Returns: + vec3f view into the linear velocity component. + """ return wp.array( ptr=sv.ptr, shape=sv.shape, @@ -1750,6 +2039,14 @@ def _get_lin_vel_from_spatial_vector(self, sv: wp.array) -> wp.array: ) def _get_ang_vel_from_spatial_vector(self, sv: wp.array) -> wp.array: + """Return an angular velocity view aliased into a spatial vector array. + + Args: + sv: Source spatial vector array. + + Returns: + vec3f view into the angular velocity component (offset 3 floats = 12 bytes). + """ return wp.array( ptr=sv.ptr + 3 * 4, shape=sv.shape, @@ -1757,3 +2054,197 @@ def _get_ang_vel_from_spatial_vector(self, sv: wp.array) -> wp.array: strides=sv.strides, device=self.device, ) + + """ + Deprecated properties. + """ + + @property + def default_root_state(self) -> ProxyArray: + """Deprecated. Use :attr:`default_root_pose` and :attr:`default_root_vel` instead. + + Shape is (num_instances,), dtype = ``vec13f``. In torch this resolves to (num_instances, 13). + """ + warnings.warn( + "default_root_state is deprecated. Use default_root_pose and default_root_vel.", + DeprecationWarning, + stacklevel=2, + ) + wp.launch( + concat_root_pose_and_vel_to_state, + dim=self.num_instances, + inputs=[self._default_root_pose, self._default_root_vel], + outputs=[self._default_root_state_buf], + device=self.device, + ) + if self._default_root_state_ta is None: + self._default_root_state_ta = ProxyArray(self._default_root_state_buf) + return self._default_root_state_ta + + @property + def root_state_w(self) -> ProxyArray: + """Deprecated. Use :attr:`root_link_pose_w` and :attr:`root_com_vel_w` instead. + + Shape is (num_instances,), dtype = ``vec13f``. In torch this resolves to (num_instances, 13). + """ + warnings.warn( + "The `root_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_link_pose_w` and " + "`root_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._root_state_w_buf.timestamp < self._sim_timestamp: + wp.launch( + concat_root_pose_and_vel_to_state, + dim=self.num_instances, + inputs=[self.root_link_pose_w, self.root_com_vel_w], + outputs=[self._root_state_w_buf.data], + device=self.device, + ) + self._root_state_w_buf.timestamp = self._sim_timestamp + if self._root_state_w_ta is None: + self._root_state_w_ta = ProxyArray(self._root_state_w_buf.data) + return self._root_state_w_ta + + @property + def root_link_state_w(self) -> ProxyArray: + """Deprecated. Use :attr:`root_link_pose_w` and :attr:`root_link_vel_w` instead. + + Shape is (num_instances,), dtype = ``vec13f``. In torch this resolves to (num_instances, 13). + """ + warnings.warn( + "The `root_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_link_pose_w` and " + "`root_link_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._root_link_state_w_buf.timestamp < self._sim_timestamp: + wp.launch( + concat_root_pose_and_vel_to_state, + dim=self.num_instances, + inputs=[self.root_link_pose_w, self.root_link_vel_w], + outputs=[self._root_link_state_w_buf.data], + device=self.device, + ) + self._root_link_state_w_buf.timestamp = self._sim_timestamp + if self._root_link_state_w_ta is None: + self._root_link_state_w_ta = ProxyArray(self._root_link_state_w_buf.data) + return self._root_link_state_w_ta + + @property + def root_com_state_w(self) -> ProxyArray: + """Deprecated. Use :attr:`root_com_pose_w` and :attr:`root_com_vel_w` instead. + + Shape is (num_instances,), dtype = ``vec13f``. In torch this resolves to (num_instances, 13). + """ + warnings.warn( + "The `root_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_com_pose_w` and " + "`root_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._root_com_state_w_buf.timestamp < self._sim_timestamp: + wp.launch( + concat_root_pose_and_vel_to_state, + dim=self.num_instances, + inputs=[self.root_com_pose_w, self.root_com_vel_w], + outputs=[self._root_com_state_w_buf.data], + device=self.device, + ) + self._root_com_state_w_buf.timestamp = self._sim_timestamp + if self._root_com_state_w_ta is None: + self._root_com_state_w_ta = ProxyArray(self._root_com_state_w_buf.data) + return self._root_com_state_w_ta + + @property + def body_state_w(self) -> ProxyArray: + """Deprecated. Use :attr:`body_link_pose_w` and :attr:`body_com_vel_w` instead. + + Shape is (num_instances, num_bodies), dtype = ``vec13f``. + In torch this resolves to (num_instances, num_bodies, 13). + """ + warnings.warn( + "The `body_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " + "`body_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_state_w_buf.timestamp >= self._sim_timestamp: + if self._body_state_w_ta is None: + self._body_state_w_ta = ProxyArray(self._body_state_w_buf.data) + return self._body_state_w_ta + _ = self.body_link_pose_w + _ = self.body_com_vel_w + wp.launch( + concat_body_pose_and_vel_to_state, + dim=(self.num_instances, self.num_bodies), + inputs=[self._body_link_pose_w.data, self._body_com_vel_w.data], + outputs=[self._body_state_w_buf.data], + device=self.device, + ) + self._body_state_w_buf.timestamp = self._sim_timestamp + if self._body_state_w_ta is None: + self._body_state_w_ta = ProxyArray(self._body_state_w_buf.data) + return self._body_state_w_ta + + @property + def body_link_state_w(self) -> ProxyArray: + """Deprecated. Use :attr:`body_link_pose_w` and :attr:`body_link_vel_w` instead. + + Shape is (num_instances, num_bodies), dtype = ``vec13f``. + In torch this resolves to (num_instances, num_bodies, 13). + """ + warnings.warn( + "The `body_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " + "`body_link_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_link_state_w_buf.timestamp >= self._sim_timestamp: + if self._body_link_state_w_ta is None: + self._body_link_state_w_ta = ProxyArray(self._body_link_state_w_buf.data) + return self._body_link_state_w_ta + _ = self.body_link_pose_w + _ = self.body_link_vel_w + wp.launch( + concat_body_pose_and_vel_to_state, + dim=(self.num_instances, self.num_bodies), + inputs=[self._body_link_pose_w.data, self._body_link_vel_w.data], + outputs=[self._body_link_state_w_buf.data], + device=self.device, + ) + self._body_link_state_w_buf.timestamp = self._sim_timestamp + if self._body_link_state_w_ta is None: + self._body_link_state_w_ta = ProxyArray(self._body_link_state_w_buf.data) + return self._body_link_state_w_ta + + @property + def body_com_state_w(self) -> ProxyArray: + """Deprecated. Use :attr:`body_com_pose_w` and :attr:`body_com_vel_w` instead. + + Shape is (num_instances, num_bodies), dtype = ``vec13f``. + In torch this resolves to (num_instances, num_bodies, 13). + """ + warnings.warn( + "The `body_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_com_pose_w` and " + "`body_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_com_state_w_buf.timestamp >= self._sim_timestamp: + if self._body_com_state_w_ta is None: + self._body_com_state_w_ta = ProxyArray(self._body_com_state_w_buf.data) + return self._body_com_state_w_ta + _ = self.body_com_pose_w + _ = self.body_com_vel_w + wp.launch( + concat_body_pose_and_vel_to_state, + dim=(self.num_instances, self.num_bodies), + inputs=[self._body_com_pose_w.data, self._body_com_vel_w.data], + outputs=[self._body_com_state_w_buf.data], + device=self.device, + ) + self._body_com_state_w_buf.timestamp = self._sim_timestamp + if self._body_com_state_w_ta is None: + self._body_com_state_w_ta = ProxyArray(self._body_com_state_w_buf.data) + return self._body_com_state_w_ta diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/kernels.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/kernels.py index cc9faf15753a..1f11a6d74789 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/kernels.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/kernels.py @@ -3,7 +3,25 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Warp kernels for the ovphysx articulation.""" +"""Warp kernels for the OVPhysX Articulation backend. + +Mirrors the structure of :mod:`isaaclab_physx.assets.articulation.kernels` for +the kernels OVPhysX exposes today. The following PhysX kernels are intentionally +absent because the corresponding write paths are not yet plumbed through the +OVPhysX Python API: + +* ``write_joint_vel_data`` +* ``write_joint_limit_data_to_buffer`` +* ``float_data_to_buffer_with_indices`` +* ``update_default_joint_values`` +* ``update_targets`` +* ``update_actuator_state_model`` +* ``extract_friction_properties`` + +OVPhysX-only data-layer kernels (``_fd_joint_acc``, ``_compose_body_com_poses``) +support the timestamped-buffer / pull-on-demand model used by +:class:`isaaclab_ovphysx.assets.ArticulationData`; they have no PhysX equivalent. +""" import warp as wp @@ -13,7 +31,18 @@ def compute_soft_joint_pos_limits_func( joint_pos_limits: wp.vec2f, soft_limit_factor: wp.float32, ): - """Compute soft joint position limits from hard limits.""" + """Compute the soft joint position limits. + + Args: + joint_pos_limits: Hard joint position limits as ``(lower, upper)`` ``[m or rad, + depending on joint type]``. + soft_limit_factor: Scale factor in ``[0, 1]`` shrinking the soft range around + the midpoint of the hard range; ``1.0`` makes the soft limits equal the + hard limits, smaller values create a tighter window. + + Returns: + The soft joint position limits as ``(lower, upper)``. + """ joint_pos_mean = (joint_pos_limits[0] + joint_pos_limits[1]) / 2.0 joint_pos_range = joint_pos_limits[1] - joint_pos_limits[0] return wp.vec2f( @@ -28,8 +57,136 @@ def update_soft_joint_pos_limits( soft_limit_factor: wp.float32, soft_joint_pos_limits: wp.array2d(dtype=wp.vec2f), ): - """Update soft joint position limits from hard limits and a scale factor.""" + """Update soft joint position limits from hard limits and a soft limit factor. + + Soft limits provide a safety margin before reaching the hard joint position + limits. See :func:`compute_soft_joint_pos_limits_func` for the per-joint + formula. + + Args: + joint_pos_limits: Hard joint position limits as vec2f ``(lower, upper)`` + ``[m or rad, depending on joint type]``. Shape is ``(num_envs, num_joints)``. + soft_limit_factor: Scale factor in ``[0, 1]``. ``1.0`` makes the soft + limits equal the hard limits; smaller values create a tighter window. + soft_joint_pos_limits: Output array. Shape is ``(num_envs, num_joints)``. + """ + i, j = wp.tid() + soft_joint_pos_limits[i, j] = compute_soft_joint_pos_limits_func(joint_pos_limits[i, j], soft_limit_factor) + + +@wp.kernel +def clamp_default_joint_pos_and_update_soft_limits_index( + joint_pos_limits: wp.array2d(dtype=wp.vec2f), + env_ids: wp.array(dtype=wp.int32), + joint_ids: wp.array(dtype=wp.int32), + soft_limit_factor: wp.float32, + default_joint_pos: wp.array2d(dtype=wp.float32), + soft_joint_pos_limits: wp.array2d(dtype=wp.vec2f), + clamped_count: wp.array(dtype=wp.int32), +): + """Clamp default joint positions to new limits and refresh soft limits over (env_ids × joint_ids). + + Mirrors PhysX's :func:`isaaclab_physx.assets.articulation.kernels.write_joint_limit_data_to_buffer` + side-effects, minus the limit-write itself (the existing + :func:`shared_kernels.write_joint_position_limit_to_buffer_index` launch handles that). + + For each ``(i, j)`` thread the kernel: + + * Clamps :attr:`default_joint_pos` ``[env_ids[i], joint_ids[j]]`` if it falls outside + the new limits, atomically incrementing :attr:`clamped_count`. + * Recomputes :attr:`soft_joint_pos_limits` ``[env_ids[i], joint_ids[j]]`` from the new + hard limits and :attr:`soft_limit_factor`. + """ + i, j = wp.tid() + e = env_ids[i] + k = joint_ids[j] + lo = joint_pos_limits[e, k][0] + hi = joint_pos_limits[e, k][1] + if (default_joint_pos[e, k] < lo) or (default_joint_pos[e, k] > hi): + wp.atomic_add(clamped_count, 0, 1) + default_joint_pos[e, k] = wp.clamp(default_joint_pos[e, k], lo, hi) + soft_joint_pos_limits[e, k] = compute_soft_joint_pos_limits_func(joint_pos_limits[e, k], soft_limit_factor) + + +@wp.kernel +def write_joint_friction_data_to_buffer_index( + in_static: wp.array2d(dtype=wp.float32), + in_dynamic: wp.array2d(dtype=wp.float32), + in_viscous: wp.array2d(dtype=wp.float32), + env_ids: wp.array(dtype=wp.int32), + joint_ids: wp.array(dtype=wp.int32), + out_buffer: wp.array3d(dtype=wp.float32), +): + """Conditionally update the static / dynamic / viscous slots of the friction buffer. + + Mirrors :func:`isaaclab_physx.assets.articulation.kernels.write_joint_friction_data_to_buffer`: + each of the three input arrays is optional (``None`` translates to a null pointer + which evaluates ``False`` inside the kernel), so callers can update any subset + of the friction components without disturbing the others. + + Args: + in_static: Static friction coefficients, or ``None`` to leave that + component unchanged. Shape is ``(num_selected_envs, num_selected_joints)``. + in_dynamic: Dynamic friction coefficients, or ``None``. + in_viscous: Viscous friction coefficients, or ``None``. + env_ids: Environment indices to write. + joint_ids: Joint indices to write. + out_buffer: The combined ``(num_envs, num_joints, 3)`` friction buffer. + """ + i, j = wp.tid() + if in_static: + out_buffer[env_ids[i], joint_ids[j], 0] = in_static[i, j] + if in_dynamic: + out_buffer[env_ids[i], joint_ids[j], 1] = in_dynamic[i, j] + if in_viscous: + out_buffer[env_ids[i], joint_ids[j], 2] = in_viscous[i, j] + + +@wp.kernel +def write_joint_friction_data_to_buffer_mask( + in_static: wp.array2d(dtype=wp.float32), + in_dynamic: wp.array2d(dtype=wp.float32), + in_viscous: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), + joint_mask: wp.array(dtype=wp.bool), + out_buffer: wp.array3d(dtype=wp.float32), +): + """Mask variant of :func:`write_joint_friction_data_to_buffer_index`.""" + i, j = wp.tid() + if not env_mask[i] or not joint_mask[j]: + return + if in_static: + out_buffer[i, j, 0] = in_static[i, j] + if in_dynamic: + out_buffer[i, j, 1] = in_dynamic[i, j] + if in_viscous: + out_buffer[i, j, 2] = in_viscous[i, j] + + +@wp.kernel +def clamp_default_joint_pos_and_update_soft_limits_mask( + joint_pos_limits: wp.array2d(dtype=wp.vec2f), + env_mask: wp.array(dtype=wp.bool), + joint_mask: wp.array(dtype=wp.bool), + soft_limit_factor: wp.float32, + default_joint_pos: wp.array2d(dtype=wp.float32), + soft_joint_pos_limits: wp.array2d(dtype=wp.vec2f), + clamped_count: wp.array(dtype=wp.int32), +): + """Mask variant of :func:`clamp_default_joint_pos_and_update_soft_limits_index`. + + Iterates the full ``(num_envs, num_joints)`` grid and applies the clamp / + soft-limit refresh only where both :paramref:`env_mask` and :paramref:`joint_mask` + are ``True``. + """ i, j = wp.tid() + if not env_mask[i] or not joint_mask[j]: + return + lo = joint_pos_limits[i, j][0] + hi = joint_pos_limits[i, j][1] + if (default_joint_pos[i, j] < lo) or (default_joint_pos[i, j] > hi): + wp.atomic_add(clamped_count, 0, 1) + default_joint_pos[i, j] = wp.clamp(default_joint_pos[i, j], lo, hi) soft_joint_pos_limits[i, j] = compute_soft_joint_pos_limits_func(joint_pos_limits[i, j], soft_limit_factor) @@ -47,11 +204,18 @@ def _fd_joint_acc( ): """Compute joint acceleration via finite differencing and update previous velocity. + Diverges from PhysX's :func:`get_joint_acc_from_joint_vel` in taking the + inverse time step rather than ``dt`` itself; the multiply-by-reciprocal + avoids per-element division inside the kernel. + Args: - cur_vel: Current joint velocities. Shape is (num_envs, num_joints). - prev_vel: Previous joint velocities (updated in-place). Shape is (num_envs, num_joints). - inv_dt: Inverse time step (1/dt) [1/s]. - out: Output joint accelerations. Shape is (num_envs, num_joints). + cur_vel: Current joint velocities ``[m/s or rad/s, depending on joint type]``. + Shape is ``(num_envs, num_joints)``. + prev_vel: Previous joint velocities (updated in-place). Same shape and units + as :paramref:`cur_vel`. + inv_dt: Inverse time step ``1 / dt`` ``[1/s]``. + out: Output joint accelerations ``[m/s^2 or rad/s^2, depending on joint type]``. + Shape is ``(num_envs, num_joints)``. """ i, j = wp.tid() out[i, j] = (cur_vel[i, j] - prev_vel[i, j]) * inv_dt diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/kernels.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/kernels.py index cf49c8362636..9d1b6ca1b8aa 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/kernels.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/kernels.py @@ -1106,3 +1106,271 @@ def write_body_com_pose_to_buffer_mask( i, j = wp.tid() if env_mask[i] and body_mask[j]: out_data[i, j] = in_data[i, j] + + +""" +Articulation-only kernels (used by isaaclab_ovphysx.assets.articulation). +""" + + +@wp.kernel +def _copy_first_body( + body_vel: wp.array(dtype=wp.spatial_vectorf, ndim=2), + root_vel: wp.array(dtype=wp.spatial_vectorf), +): + """Copy the first body's spatial velocity to the root velocity buffer. + + For single rigid-body assets, index 0 is always the root body. This + kernel extracts that slice without allocating an intermediate buffer. + + Args: + body_vel: Body spatial velocities ``[m/s, rad/s]``. Shape is + ``(num_envs, num_bodies)`` with dtype ``wp.spatial_vectorf``. + root_vel: Output root spatial velocities ``[m/s, rad/s]``. Shape is + ``(num_envs,)`` with dtype ``wp.spatial_vectorf``. + """ + i = wp.tid() + root_vel[i] = body_vel[i, 0] + + +@wp.kernel +def _compose_root_com_pose( + link_pose: wp.array(dtype=wp.transformf), + com_pose_b: wp.array(dtype=wp.transformf, ndim=2), + com_pose_w: wp.array(dtype=wp.transformf), +): + """Compose root link pose with the body-frame COM offset to get the world-frame COM pose. + + Implements the forward transform: + + ``com_pose_w = link_pose_w * com_pose_b[0]`` + + where ``*`` denotes ``wp.transform_multiply``. Only the first body + (index ``0``) is used; for articulations this is the base link body. + + Args: + link_pose: Root link poses in world frame ``[m, -]``. Shape is + ``(num_envs,)`` with dtype ``wp.transformf``. + com_pose_b: Body-frame COM offsets ``[m, -]`` from the + ``RIGID_BODY_COM_POSE`` binding. Shape is ``(num_envs, num_bodies)`` + with dtype ``wp.transformf``. + com_pose_w: Output world-frame root COM poses ``[m, -]``. Shape is + ``(num_envs,)`` with dtype ``wp.transformf``. + """ + i = wp.tid() + com_pose_w[i] = wp.transform_multiply(link_pose[i], com_pose_b[i, 0]) + + +@wp.kernel +def _projected_gravity( + gravity_vec_w: wp.array(dtype=wp.vec3f), + root_pose: wp.array(dtype=wp.transformf), + out: wp.array(dtype=wp.vec3f), +): + """Project the world-frame gravity direction into the root body frame. + + Applies the inverse of the root orientation quaternion to the world-frame + gravity vector, yielding the gravity direction expressed in the body frame. + The magnitude is preserved (unit vector in, unit vector out if input is a + unit vector). + + Args: + gravity_vec_w: Gravity direction per instance in world frame ``[-]`` + (typically the normalised ``(0, 0, -1)`` gravitational acceleration + direction). Shape is ``(num_envs,)`` with dtype ``wp.vec3f``. + root_pose: Root link poses in world frame ``[m, -]``. Only the + rotation component is used. Shape is ``(num_envs,)`` with dtype + ``wp.transformf``. + out: Output gravity direction in body frame ``[-]``. Shape is + ``(num_envs,)`` with dtype ``wp.vec3f``. + """ + i = wp.tid() + q = wp.transform_get_rotation(root_pose[i]) + out[i] = wp.quat_rotate_inv(q, gravity_vec_w[i]) + + +@wp.kernel +def _compute_heading( + forward_vec_b: wp.array(dtype=wp.vec3f), + root_pose: wp.array(dtype=wp.transformf), + out: wp.array(dtype=wp.float32), +): + """Compute the yaw heading angle by rotating a body-frame forward vector to world frame. + + Rotates ``forward_vec_b`` by the root orientation quaternion and then computes the + heading as ``atan2(forward_w.y, forward_w.x)`` ``[rad]``, i.e. the signed angle + from the world X-axis to the projected forward direction in the XY plane. + + Args: + forward_vec_b: Forward direction in body frame per instance ``[-]``. + Shape is ``(num_envs,)`` with dtype ``wp.vec3f``. + root_pose: Root link poses in world frame ``[m, -]``. Only the rotation + component is used. Shape is ``(num_envs,)`` with dtype ``wp.transformf``. + out: Output heading angles ``[rad]`` in ``[-π, π]``. Shape is + ``(num_envs,)`` with dtype ``wp.float32``. + """ + i = wp.tid() + q = wp.transform_get_rotation(root_pose[i]) + forward = wp.quat_rotate(q, forward_vec_b[i]) + out[i] = wp.atan2(forward[1], forward[0]) + + +@wp.kernel +def _world_vel_to_body_lin( + root_pose: wp.array(dtype=wp.transformf), + vel_w: wp.array(dtype=wp.spatial_vectorf), + out: wp.array(dtype=wp.vec3f), +): + """Rotate the world-frame linear velocity component into the root body frame. + + Extracts the linear velocity from the top three components of the spatial + velocity vector (``wp.spatial_top``) and rotates it by the inverse of the + root orientation quaternion. + + Args: + root_pose: Root link poses in world frame ``[m, -]``. Only the rotation + component is used. Shape is ``(num_envs,)`` with dtype ``wp.transformf``. + vel_w: Root spatial velocities in world frame ``[m/s, rad/s]``. + Shape is ``(num_envs,)`` with dtype ``wp.spatial_vectorf``. + out: Output linear velocity in body frame ``[m/s]``. Shape is + ``(num_envs,)`` with dtype ``wp.vec3f``. + """ + i = wp.tid() + q = wp.transform_get_rotation(root_pose[i]) + lin = wp.spatial_top(vel_w[i]) + out[i] = wp.quat_rotate_inv(q, lin) + + +@wp.kernel +def _world_vel_to_body_ang( + root_pose: wp.array(dtype=wp.transformf), + vel_w: wp.array(dtype=wp.spatial_vectorf), + out: wp.array(dtype=wp.vec3f), +): + """Rotate the world-frame angular velocity component into the root body frame. + + Extracts the angular velocity from the bottom three components of the spatial + velocity vector (``wp.spatial_bottom``) and rotates it by the inverse of the + root orientation quaternion. + + Args: + root_pose: Root link poses in world frame ``[m, -]``. Only the rotation + component is used. Shape is ``(num_envs,)`` with dtype ``wp.transformf``. + vel_w: Root spatial velocities in world frame ``[m/s, rad/s]``. + Shape is ``(num_envs,)`` with dtype ``wp.spatial_vectorf``. + out: Output angular velocity in body frame ``[rad/s]``. Shape is + ``(num_envs,)`` with dtype ``wp.vec3f``. + """ + i = wp.tid() + q = wp.transform_get_rotation(root_pose[i]) + ang = wp.spatial_bottom(vel_w[i]) + out[i] = wp.quat_rotate_inv(q, ang) + + +@wp.kernel +def write_joint_position_limit_to_buffer_index( + in_data: wp.array3d(dtype=wp.float32), + env_ids: wp.array(dtype=wp.int32), + joint_ids: wp.array(dtype=wp.int32), + out_data: wp.array(dtype=wp.vec2f, ndim=2), +): + """Write joint position-limit data to a vec2f buffer at specified indices. + + This kernel copies ``[lower, upper]`` limit pairs from a partial float32 input + array into the output ``wp.vec2f`` buffer at the specified environment and joint + indices. + + Args: + in_data: Input array containing limit pairs ``[lower, upper]`` [m or rad]. + Shape is (num_selected_envs, num_selected_joints, 2). + env_ids: Input array of environment indices to write to. + Shape is (num_selected_envs,). + joint_ids: Input array of joint indices to write to. + Shape is (num_selected_joints,). + out_data: Output array where limit data is written. Shape is + (num_envs, num_joints) with dtype ``wp.vec2f``. + """ + i, j = wp.tid() + out_data[env_ids[i], joint_ids[j]] = wp.vec2f(in_data[i, j, 0], in_data[i, j, 1]) + + +@wp.kernel +def write_joint_position_limit_to_buffer_mask( + in_data: wp.array3d(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), + joint_mask: wp.array(dtype=wp.bool), + out_data: wp.array(dtype=wp.vec2f, ndim=2), +): + """Mask-scatter joint position-limit data into the vec2f cache buffer. + + Copies ``[lower, upper]`` limit pairs where both ``env_mask[i]`` and + ``joint_mask[j]`` are True. + + Args: + in_data: Input array containing limit pairs ``[lower, upper]`` [m or rad]. + Shape is (num_envs, num_joints, 2). + env_mask: Boolean environment mask. Shape is (num_envs,). + joint_mask: Boolean joint mask. Shape is (num_joints,). + out_data: Output array where limit data is written. Shape is + (num_envs, num_joints) with dtype ``wp.vec2f``. + """ + i, j = wp.tid() + if env_mask[i] and joint_mask[j]: + out_data[i, j] = wp.vec2f(in_data[i, j, 0], in_data[i, j, 1]) + + +@wp.kernel +def write_joint_friction_to_buffer_index( + in_data: wp.array2d(dtype=wp.float32), + env_ids: wp.array(dtype=wp.int32), + joint_ids: wp.array(dtype=wp.int32), + out_data: wp.array3d(dtype=wp.float32), +): + """Write joint friction coefficient to all three components of the friction buffer. + + Broadcasts a single friction value into the static (index 0), dynamic (index 1), + and viscous (index 2) components of the ``(N, D, 3)`` friction properties buffer + at the specified environment and joint indices. + + Args: + in_data: Input friction coefficients [dimensionless]. Shape is + (num_selected_envs, num_selected_joints). + env_ids: Input array of environment indices to write to. + Shape is (num_selected_envs,). + joint_ids: Input array of joint indices to write to. + Shape is (num_selected_joints,). + out_data: Output friction properties buffer. Shape is (num_envs, num_joints, 3). + """ + i, j = wp.tid() + val = in_data[i, j] + out_data[env_ids[i], joint_ids[j], 0] = val + out_data[env_ids[i], joint_ids[j], 1] = val + out_data[env_ids[i], joint_ids[j], 2] = val + + +@wp.kernel +def write_joint_friction_to_buffer_mask( + in_data: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), + joint_mask: wp.array(dtype=wp.bool), + out_data: wp.array3d(dtype=wp.float32), +): + """Mask-scatter joint friction coefficient into all three components of the friction buffer. + + Broadcasts a single friction value into the static (index 0), dynamic (index 1), + and viscous (index 2) components where both ``env_mask[i]`` and ``joint_mask[j]`` + are True. + + Args: + in_data: Input friction coefficients [dimensionless]. Shape is + (num_envs, num_joints). + env_mask: Boolean environment mask. Shape is (num_envs,). + joint_mask: Boolean joint mask. Shape is (num_joints,). + out_data: Output friction properties buffer. Shape is (num_envs, num_joints, 3). + """ + i, j = wp.tid() + if env_mask[i] and joint_mask[j]: + val = in_data[i, j] + out_data[i, j, 0] = val + out_data[i, j, 1] = val + out_data[i, j, 2] = val diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/tensor_types.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/tensor_types.py index 41afe07cf09c..9e7df44aef9a 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/tensor_types.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/tensor_types.py @@ -365,6 +365,11 @@ # 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. +# +# Tendon tensor types are NOT in this set: PhysX exposes tendons on the +# simulation device (its ``set_fixed_tendon_properties`` takes ``data.warp`` +# without a ``device="cpu"`` clone, unlike ``set_dof_stiffnesses``), and the +# OVPhysX wheel mirrors that — tendon bindings are GPU-resident on a GPU sim. _CPU_ONLY_TYPES_CANDIDATES: tuple = ( DOF_STIFFNESS, DOF_DAMPING, @@ -378,16 +383,6 @@ 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, diff --git a/source/isaaclab_ovphysx/test/assets/test_articulation.py b/source/isaaclab_ovphysx/test/assets/test_articulation.py index 52998a2ec5f6..ad17be4eb7ab 100644 --- a/source/isaaclab_ovphysx/test/assets/test_articulation.py +++ b/source/isaaclab_ovphysx/test/assets/test_articulation.py @@ -3,114 +3,2444 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Unit tests for ovphysx articulation helpers.""" +# ignore private usage of variables warning +# pyright: reportPrivateUsage=none + + +"""Real-backend tests for the OVPhysX Articulation. + +Mirrors :mod:`isaaclab_physx.test.assets.test_articulation` 1-to-1: same set +of test functions, names, parametrizations, and assertions. + +OVPhysX runs kitless under ``./scripts/run_ovphysx.sh`` so there is no +``AppLauncher`` boot — :class:`~isaaclab.sim.SimulationContext` is driven +directly via ``build_simulation_context(sim_cfg=SimulationCfg(physics=OvPhysxCfg(), ...))`` +which works because :func:`isaaclab.app.has_kit` returns False in this +environment. + +PhysX-specific ``cube_object.root_view.set_X(...)`` / ``get_X(...)`` calls are +adapted to OVPhysX by going through the backend's per-tensor-type binding +dictionary (``cube_object._bindings`` / :meth:`~isaaclab_ovphysx.assets.Articulation._get_binding`) +and the public setters (:meth:`set_masses_index`, :meth:`set_coms_index`, +:meth:`set_inertias_index`). Reads use the data-class properties +(``cube_object.data.body_mass``, ``body_inertia``, ``body_com_pose_b``). + +Process-global device lock +-------------------------- + +``ovphysx<=0.3.7`` binds device mode (CPU vs GPU) at the C++ layer on the +first ``ovphysx.PhysX(device=...)`` call and cannot release/swap it without a +process restart. :class:`~isaaclab_ovphysx.physics.OvPhysxManager` tracks +this on ``_locked_device`` and raises :exc:`RuntimeError` if a later +:class:`SimulationContext` requests a different device. The +``_ovphysx_skip_other_device`` autouse fixture below preempts that error in +parametrized tests by ``pytest.skip``-ing on the unlocked device, so the +session finishes cleanly when only one device is exercised. + +CI note +------- +Because the lock is process-global, full coverage requires **two separate +``./scripts/run_ovphysx.sh -m pytest`` invocations** -- once with ``-k 'cpu'`` +and once with ``-k 'cuda:0'``. Tracked as gap G5 in +``docs/superpowers/specs/2026-04-28-ovphysx-wheel-gaps-for-marco.md``; until +the wheel exposes a way to reset Carbonite device state, this is the supported +pattern. +""" from __future__ import annotations -from types import SimpleNamespace +import sys +import numpy as np import pytest +import torch import warp as wp +from isaaclab_ovphysx import tensor_types as TT +from isaaclab_ovphysx.assets import Articulation +from isaaclab_ovphysx.physics import OvPhysxCfg -from pxr import Sdf, Usd, UsdPhysics - -# The CI isaaclab_ov* pattern unintentionally collects isaaclab_ovphysx tests, -# but the ovphysx wheel is not installed in that environment. Skip gracefully -# so the isaaclab_ov CI pipeline is not blocked by an unrelated dependency. -pytest.importorskip("ovphysx.types", reason="ovphysx wheel not installed") +import isaaclab.sim as sim_utils +import isaaclab.utils.math as math_utils +import isaaclab.utils.string as string_utils +from isaaclab.actuators import ActuatorBase, IdealPDActuatorCfg, ImplicitActuatorCfg +from isaaclab.assets import ArticulationCfg +from isaaclab.envs.mdp.terminations import joint_effort_out_of_limit +from isaaclab.managers import SceneEntityCfg +from isaaclab.sim import SimulationCfg, build_simulation_context +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.version import get_isaac_sim_version, has_kit -from isaaclab_ovphysx.assets.articulation.articulation import Articulation # noqa: E402 -from isaaclab_ovphysx.physics import OvPhysxManager # noqa: E402 -from isaaclab_ovphysx.test.mock_interfaces.views import MockOvPhysxBindingSet # noqa: E402 +## +# Pre-defined configs +## +from isaaclab_assets import ANYMAL_C_CFG, FRANKA_PANDA_CFG, SHADOW_HAND_CFG # isort:skip wp.init() -def _define_tendon_joint(stage: Usd.Stage, path: str, schema_name: str) -> None: - """Define a revolute joint prim with a tendon schema marker.""" - joint = UsdPhysics.RevoluteJoint.Define(stage, path) - schemas = Sdf.TokenListOp() - schemas.explicitItems = [schema_name] - joint.GetPrim().SetMetadata("apiSchemas", schemas) - - -def _make_articulation_root_stage(tmp_path) -> str: - """Create a stage with one relevant articulation subtree and unrelated joints elsewhere.""" - stage = Usd.Stage.CreateInMemory() - stage.DefinePrim("/World", "Xform") - stage.DefinePrim("/World/envs", "Xform") - stage.DefinePrim("/World/envs/env_0", "Xform") - stage.DefinePrim("/World/envs/env_0/Robot", "Xform") - stage.DefinePrim("/World/envs/env_0/Robot/root", "Xform") - stage.DefinePrim("/World/unrelated", "Xform") - - _define_tendon_joint( - stage, - "/World/envs/env_0/Robot/root/fixed_joint", - "PhysxTendonAxisRootAPI:inst0", - ) - _define_tendon_joint( - stage, - "/World/envs/env_0/Robot/root/spatial_joint", - "PhysxTendonAttachmentRootAPI:inst0", - ) - _define_tendon_joint( - stage, - "/World/unrelated/unrelated_fixed_joint", - "PhysxTendonAxisRootAPI:inst0", - ) - _define_tendon_joint( - stage, - "/World/unrelated/unrelated_spatial_joint", - "PhysxTendonAttachmentLeafAPI:inst0", - ) - - stage_path = tmp_path / "scene.usda" - stage.Export(str(stage_path)) - return str(stage_path) - - -def _make_articulation_shell() -> Articulation: - """Create a minimal ovphysx articulation shell for tendon processing tests.""" - articulation = object.__new__(Articulation) - bindings = MockOvPhysxBindingSet( - num_instances=1, - num_joints=2, - num_bodies=2, - num_fixed_tendons=1, - num_spatial_tendons=1, - ) - object.__setattr__(articulation, "_bindings", bindings.bindings) - object.__setattr__(articulation, "_articulation_root_path", "/World/envs/env_0/Robot/root") - object.__setattr__(articulation, "_initialize_handle", None) - object.__setattr__(articulation, "_invalidate_initialize_handle", None) - object.__setattr__(articulation, "_prim_deletion_handle", None) - object.__setattr__(articulation, "_debug_vis_handle", None) - object.__setattr__( - articulation, - "_data", - SimpleNamespace( - _num_fixed_tendons=0, - _num_spatial_tendons=0, - fixed_tendon_names=[], - spatial_tendon_names=[], - ), - ) - return articulation - - -def test_process_tendons_scopes_to_articulation_root(tmp_path): - """Tendon discovery should ignore joints that live outside the current articulation subtree.""" - articulation = _make_articulation_shell() - stage_path = _make_articulation_root_stage(tmp_path) - old_stage_path = OvPhysxManager._stage_path - OvPhysxManager._stage_path = stage_path - try: - articulation._process_tendons() - finally: - OvPhysxManager._stage_path = old_stage_path - - assert articulation.fixed_tendon_names == ["fixed_joint"] - assert articulation.spatial_tendon_names == ["spatial_joint"] - assert articulation._data.fixed_tendon_names == ["fixed_joint"] - assert articulation._data.spatial_tendon_names == ["spatial_joint"] +_OMNI_PHYSX_SCHEMAS_GAP_REASON = ( + "Schema-level fixed-joint creation in :mod:`isaaclab.sim.schemas` imports " + "``omni.physx.scripts.utils``, which is a Kit-only module not shipped by " + "the ovphysx wheel. See " + "docs/superpowers/specs/2026-04-28-ovphysx-wheel-gaps-for-marco.md." +) + +_MATERIAL_GAP_REASON = ( + "Requires a ``RIGID_BODY_MATERIAL`` TensorType (or a view-helper) on the " + "ovphysx wheel side. ``Articulation.root_view`` is a per-tensor-type " + "bindings dict on OVPhysX, so ``root_view.get_material_properties()`` / " + "``set_material_properties()`` / ``max_shapes`` are not available. See " + "docs/superpowers/specs/2026-04-28-ovphysx-wheel-gaps-for-marco.md." +) + +_FK_ON_DEMAND_GAP_REASON = ( + "GPU-only: PhysX's data-class getters call " + "``SimulationView.update_articulations_kinematic`` before reading link " + "transforms (see ``isaaclab_physx.assets.articulation_data:735``), so body " + "poses reflect new joint positions immediately after " + "``write_joint_position_to_sim_*`` without a sim step. The OVPhysX wheel's " + "``ovphysx.PhysX`` class does not expose an equivalent FK-on-demand API " + "(omni.physics.tensors has it; the OVPhysX wrapper does not surface it), " + "so body-pose bindings remain stale on a GPU sim until the next ``step``. " + "On a CPU sim the same write path happens to update the bindings synchronously " + "(no async stream involved), so the test xpasses there — hence ``strict=False``. " + "See docs/superpowers/specs/2026-04-28-ovphysx-wheel-gaps-for-marco.md." +) + + +def _read_binding_to_torch(articulation: Articulation, tensor_type: int, device: str | torch.device) -> torch.Tensor: + """Read an OVPhysX TensorBinding into a torch tensor on *device*. + + Test-side adapter for the verbatim PhysX mirror. PhysX cross-checks the + data class against the simulation via ``articulation.root_view.get_X()`` + accessors; on OVPhysX, ``root_view`` is a per-tensor-type bindings dict + (no view-level getters), so we read the binding directly into a CPU + numpy buffer (CPU-only types) and move the result to *device*. + """ + binding = articulation.root_view[tensor_type] + np_buf = np.zeros(binding.shape, dtype=np.float32) + binding.read(np_buf) + return torch.from_numpy(np_buf).to(device) + + +# Session-locked device. Set on the first parametrized test that runs and +# never reassigned -- ovphysx's process-global device lock means subsequent +# tests on the other device must skip. +_LOCKED_DEVICE: list[str | None] = [None] + + +@pytest.fixture(autouse=True) +def _ovphysx_skip_other_device(request): + """Skip tests whose ``device`` parameter mismatches the session-locked device. + + ``ovphysx<=0.3.7`` locks the process-global device mode on the first + ``ovphysx.PhysX(device=...)`` call, so any test parametrized to a different + device after the first ``sim.reset()`` would hit + :exc:`ovphysx.types.PhysXDeviceError`. We detect the locked device on the + first encounter and skip subsequent tests on the other device with a clear + message so the run finishes cleanly rather than producing spurious failures. + """ + callspec = getattr(request.node, "callspec", None) + device = callspec.params.get("device") if callspec is not None else None + if device is None: + # Test does not parametrize on device (e.g. test_warmup_attach_stage_not_called_for_cpu). + return + locked = _LOCKED_DEVICE[0] + if locked is None: + _LOCKED_DEVICE[0] = device + return + if device != locked: + pytest.skip( + f"ovphysx process-global device lock is held by '{locked}'; cannot run '{device}' " + "tests in the same session. Run pytest twice (once per device) for full coverage." + ) + + +def _ovphysx_sim_context(device: str, **kwargs): + """Wrapper around :func:`build_simulation_context` that injects OVPhysX cfg. + + PhysX tests pass ``device=device`` directly and let + :func:`build_simulation_context` build a default :class:`SimulationCfg`. + OVPhysX needs ``physics=OvPhysxCfg()`` set on the cfg so the manager + dispatches to OVPhysX rather than PhysX, so we build the cfg here and + pass it through. ``gravity_enabled`` is consumed locally (it is ignored + by ``build_simulation_context`` once a ``sim_cfg`` is provided). + ``add_ground_plane``, ``auto_add_lighting``, and other kwargs continue + to flow through ``build_simulation_context`` as before. + """ + dt = kwargs.pop("dt", 1.0 / 60.0) + gravity_enabled = kwargs.pop("gravity_enabled", True) + gravity = (0.0, 0.0, -9.81) if gravity_enabled else (0.0, 0.0, 0.0) + sim_cfg = SimulationCfg(physics=OvPhysxCfg(), device=device, dt=dt, gravity=gravity) + return build_simulation_context(device=device, sim_cfg=sim_cfg, **kwargs) + + +def generate_articulation_cfg( + articulation_type: str, + stiffness: float | None = 10.0, + damping: float | None = 2.0, + velocity_limit: float | None = None, + effort_limit: float | None = None, + velocity_limit_sim: float | None = None, + effort_limit_sim: float | None = None, +) -> ArticulationCfg: + """Generate an articulation configuration. + + Args: + articulation_type: Type of articulation to generate. + It should be one of: "humanoid", "panda", "anymal", "shadow_hand", "single_joint_implicit", + "single_joint_explicit". + stiffness: Stiffness value for the articulation's actuators. Only currently used for "humanoid". + Defaults to 10.0. + damping: Damping value for the articulation's actuators. Only currently used for "humanoid". + Defaults to 2.0. + velocity_limit: Velocity limit for the actuators. Only currently used for "single_joint_implicit" + and "single_joint_explicit". + effort_limit: Effort limit for the actuators. Only currently used for "single_joint_implicit" + and "single_joint_explicit". + velocity_limit_sim: Velocity limit for the actuators (set into the simulation). + Only currently used for "single_joint_implicit" and "single_joint_explicit". + effort_limit_sim: Effort limit for the actuators (set into the simulation). + Only currently used for "single_joint_implicit" and "single_joint_explicit". + + Returns: + The articulation configuration for the requested articulation type. + + """ + if articulation_type == "humanoid": + articulation_cfg = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/IsaacSim/Humanoid/humanoid_instanceable.usd" + ), + init_state=ArticulationCfg.InitialStateCfg(pos=(0.0, 0.0, 1.34)), + actuators={"body": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=stiffness, damping=damping)}, + ) + elif articulation_type == "panda": + articulation_cfg = FRANKA_PANDA_CFG + elif articulation_type == "anymal": + articulation_cfg = ANYMAL_C_CFG + elif articulation_type == "shadow_hand": + articulation_cfg = SHADOW_HAND_CFG + elif articulation_type == "single_joint_implicit": + articulation_cfg = ArticulationCfg( + # we set 80.0 default for max force because default in USD is 10e10 which makes testing annoying. + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/IsaacSim/SimpleArticulation/revolute_articulation.usd", + joint_drive_props=sim_utils.JointDrivePropertiesCfg(max_effort=80.0, max_velocity=5.0), + ), + actuators={ + "joint": ImplicitActuatorCfg( + joint_names_expr=[".*"], + effort_limit_sim=effort_limit_sim, + velocity_limit_sim=velocity_limit_sim, + effort_limit=effort_limit, + velocity_limit=velocity_limit, + stiffness=2000.0, + damping=100.0, + ), + }, + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.0, 0.0, 0.0), + joint_pos=({"RevoluteJoint": 1.5708}), + rot=(0.7071081, 0, 0, 0.7071055), + ), + ) + elif articulation_type == "single_joint_explicit": + # we set 80.0 default for max force because default in USD is 10e10 which makes testing annoying. + articulation_cfg = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/IsaacSim/SimpleArticulation/revolute_articulation.usd", + joint_drive_props=sim_utils.JointDrivePropertiesCfg(max_effort=80.0, max_velocity=5.0), + ), + actuators={ + "joint": IdealPDActuatorCfg( + joint_names_expr=[".*"], + effort_limit_sim=effort_limit_sim, + velocity_limit_sim=velocity_limit_sim, + effort_limit=effort_limit, + velocity_limit=velocity_limit, + stiffness=0.0, + damping=10.0, + ), + }, + ) + elif articulation_type == "spatial_tendon_test_asset": + # we set 80.0 default for max force because default in USD is 10e10 which makes testing annoying. + articulation_cfg = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/IsaacLab/Tests/spatial_tendons.usd", + ), + actuators={ + "joint": ImplicitActuatorCfg( + joint_names_expr=[".*"], + stiffness=2000.0, + damping=100.0, + ), + }, + ) + else: + raise ValueError( + f"Invalid articulation type: {articulation_type}, valid options are 'humanoid', 'panda', 'anymal'," + " 'shadow_hand', 'single_joint_implicit', 'single_joint_explicit' or 'spatial_tendon_test_asset'." + ) + + return articulation_cfg + + +def generate_articulation( + articulation_cfg: ArticulationCfg, num_articulations: int, device: str +) -> tuple[Articulation, torch.tensor]: + """Generate an articulation from a configuration. + + Handles the creation of the articulation, the environment prims and the articulation's environment + translations + + Args: + articulation_cfg: Articulation configuration. + num_articulations: Number of articulations to generate. + device: Device to use for the tensors. + + Returns: + The articulation and environment translations. + + """ + # Generate translations of 2.5 m in x for each articulation + translations = torch.zeros(num_articulations, 3, device=device) + translations[:, 0] = torch.arange(num_articulations) * 2.5 + + # Create Top-level Xforms, one for each articulation + for i in range(num_articulations): + sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=translations[i][:3]) + articulation = Articulation(articulation_cfg.replace(prim_path="/World/Env_.*/Robot")) + + return articulation, translations + + +@pytest.fixture +def sim(request): + """Create simulation context with the specified device.""" + device = request.getfixturevalue("device") + if "gravity_enabled" in request.fixturenames: + gravity_enabled = request.getfixturevalue("gravity_enabled") + else: + gravity_enabled = True # default to gravity enabled + if "add_ground_plane" in request.fixturenames: + add_ground_plane = request.getfixturevalue("add_ground_plane") + else: + add_ground_plane = False # default to no ground plane + with _ovphysx_sim_context( + device=device, auto_add_lighting=True, gravity_enabled=gravity_enabled, add_ground_plane=add_ground_plane + ) as sim: + sim._app_control_on_stop_handle = None + yield sim + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.isaacsim_ci +def test_initialization_floating_base_non_root(sim, num_articulations, device, add_ground_plane): + """Test initialization for a floating-base with articulation root on a rigid body. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is not fixed base + 3. All buffers have correct shapes + 4. The articulation can be simulated + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + articulation_cfg = generate_articulation_cfg(articulation_type="humanoid", stiffness=0.0, damping=0.0) + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + + # Check if articulation is initialized + assert articulation.is_initialized + # Check that is fixed base + assert not articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 21) + + # Cross-check binding shapes against cached counts. PhysX does this via + # ``root_view.max_dofs == shared_metatype.dof_count``; on OVPhysX + # ``root_view`` is the per-tensor-type bindings dict, so the equivalent + # invariant is that each per-DOF / per-link binding's shape agrees with + # the count cached on the asset. + for tt in (TT.DOF_POSITION, TT.DOF_VELOCITY, TT.DOF_STIFFNESS): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_joints + for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_bodies + # Body-name ordering check is degenerate on OVPhysX: ``body_names`` is + # sourced from binding metadata (``sample.body_names``), so the PhysX + # ``link_paths[0]`` round-trip is a no-op here and is omitted. + # -- actuator type + for actuator_name, actuator in articulation.actuators.items(): + is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) + assert actuator.is_implicit_model == is_implicit_model_cfg + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.isaacsim_ci +def test_initialization_floating_base(sim, num_articulations, device, add_ground_plane): + """Test initialization for a floating-base with articulation root on provided prim path. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is not fixed base + 3. All buffers have correct shapes + 4. The articulation can be simulated + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + articulation_cfg = generate_articulation_cfg(articulation_type="anymal", stiffness=0.0, damping=0.0) + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that floating base + assert not articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 12) + assert articulation.data.body_mass.torch.shape == (num_articulations, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (num_articulations, articulation.num_bodies, 9) + + # Cross-check binding shapes against cached counts. PhysX does this via + # ``root_view.max_dofs == shared_metatype.dof_count``; on OVPhysX + # ``root_view`` is the per-tensor-type bindings dict, so the equivalent + # invariant is that each per-DOF / per-link binding's shape agrees with + # the count cached on the asset. + for tt in (TT.DOF_POSITION, TT.DOF_VELOCITY, TT.DOF_STIFFNESS): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_joints + for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_bodies + # Body-name ordering check is degenerate on OVPhysX: ``body_names`` is + # sourced from binding metadata (``sample.body_names``), so the PhysX + # ``link_paths[0]`` round-trip is a no-op here and is omitted. + # -- actuator type + for actuator_name, actuator in articulation.actuators.items(): + is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) + assert actuator.is_implicit_model == is_implicit_model_cfg + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_initialization_fixed_base(sim, num_articulations, device): + """Test initialization for fixed base. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is fixed base + 3. All buffers have correct shapes + 4. The articulation maintains its default state + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + articulation_cfg = generate_articulation_cfg(articulation_type="panda") + articulation, translations = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that fixed base + assert articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 9) + assert articulation.data.body_mass.torch.shape == (num_articulations, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (num_articulations, articulation.num_bodies, 9) + + # Cross-check binding shapes against cached counts. PhysX does this via + # ``root_view.max_dofs == shared_metatype.dof_count``; on OVPhysX + # ``root_view`` is the per-tensor-type bindings dict, so the equivalent + # invariant is that each per-DOF / per-link binding's shape agrees with + # the count cached on the asset. + for tt in (TT.DOF_POSITION, TT.DOF_VELOCITY, TT.DOF_STIFFNESS): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_joints + for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_bodies + # Body-name ordering check is degenerate on OVPhysX: ``body_names`` is + # sourced from binding metadata (``sample.body_names``), so the PhysX + # ``link_paths[0]`` round-trip is a no-op here and is omitted. + # -- actuator type + for actuator_name, actuator in articulation.actuators.items(): + is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) + assert actuator.is_implicit_model == is_implicit_model_cfg + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + # check that the root is at the correct state - its default state as it is fixed base + default_root_pose = articulation.data.default_root_pose.torch.clone() + default_root_vel = articulation.data.default_root_vel.torch.clone() + default_root_pose[:, :3] = default_root_pose[:, :3] + translations + + torch.testing.assert_close(articulation.data.root_link_pose_w.torch, default_root_pose) + torch.testing.assert_close(articulation.data.root_com_vel_w.torch, default_root_vel) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.isaacsim_ci +def test_initialization_fixed_base_single_joint(sim, num_articulations, device, add_ground_plane): + """Test initialization for fixed base articulation with a single joint. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is fixed base + 3. All buffers have correct shapes + 4. The articulation maintains its default state + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + articulation_cfg = generate_articulation_cfg(articulation_type="single_joint_implicit") + articulation, translations = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that fixed base + assert articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 1) + assert articulation.data.body_mass.torch.shape == (num_articulations, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (num_articulations, articulation.num_bodies, 9) + + # Cross-check binding shapes against cached counts. PhysX does this via + # ``root_view.max_dofs == shared_metatype.dof_count``; on OVPhysX + # ``root_view`` is the per-tensor-type bindings dict, so the equivalent + # invariant is that each per-DOF / per-link binding's shape agrees with + # the count cached on the asset. + for tt in (TT.DOF_POSITION, TT.DOF_VELOCITY, TT.DOF_STIFFNESS): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_joints + for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_bodies + # Body-name ordering check is degenerate on OVPhysX: ``body_names`` is + # sourced from binding metadata (``sample.body_names``), so the PhysX + # ``link_paths[0]`` round-trip is a no-op here and is omitted. + # -- actuator type + for actuator_name, actuator in articulation.actuators.items(): + is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) + assert actuator.is_implicit_model == is_implicit_model_cfg + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + # check that the root is at the correct state - its default state as it is fixed base + default_root_pose = articulation.data.default_root_pose.torch.clone() + default_root_vel = articulation.data.default_root_vel.torch.clone() + default_root_pose[:, :3] = default_root_pose[:, :3] + translations + + torch.testing.assert_close(articulation.data.root_link_pose_w.torch, default_root_pose) + torch.testing.assert_close(articulation.data.root_com_vel_w.torch, default_root_vel) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_initialization_hand_with_tendons(sim, num_articulations, device): + """Test initialization for fixed base articulated hand with tendons. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is fixed base + 3. All buffers have correct shapes + 4. The articulation can be simulated + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + articulation_cfg = generate_articulation_cfg(articulation_type="shadow_hand") + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that fixed base + assert articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 24) + assert articulation.data.body_mass.torch.shape == (num_articulations, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (num_articulations, articulation.num_bodies, 9) + + # Cross-check binding shapes against cached counts. See the equivalent + # block in test_initialization_fixed_base_single_joint for why the verbatim + # PhysX ``root_view.max_dofs == shared_metatype.dof_count`` identity is + # replaced with binding-shape checks on OVPhysX. + for tt in (TT.DOF_POSITION, TT.DOF_VELOCITY, TT.DOF_STIFFNESS): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_joints + for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_bodies + # -- actuator type + for actuator_name, actuator in articulation.actuators.items(): + is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) + assert actuator.is_implicit_model == is_implicit_model_cfg + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.isaacsim_ci +@pytest.mark.xfail(reason=_OMNI_PHYSX_SCHEMAS_GAP_REASON, strict=False) +def test_initialization_floating_base_made_fixed_base(sim, num_articulations, device, add_ground_plane): + """Test initialization for a floating-base articulation made fixed-base using schema properties. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is fixed base after modification + 3. All buffers have correct shapes + 4. The articulation maintains its default state + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type="anymal").copy() + # Fix root link by making it kinematic + articulation_cfg.spawn.articulation_props.fix_root_link = True + articulation, translations = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that is fixed base + assert articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 12) + + # Cross-check binding shapes against cached counts. PhysX does this via + # ``root_view.max_dofs == shared_metatype.dof_count``; on OVPhysX + # ``root_view`` is the per-tensor-type bindings dict, so the equivalent + # invariant is that each per-DOF / per-link binding's shape agrees with + # the count cached on the asset. + for tt in (TT.DOF_POSITION, TT.DOF_VELOCITY, TT.DOF_STIFFNESS): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_joints + for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_bodies + # Body-name ordering check is degenerate on OVPhysX: ``body_names`` is + # sourced from binding metadata (``sample.body_names``), so the PhysX + # ``link_paths[0]`` round-trip is a no-op here and is omitted. + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + # check that the root is at the correct state - its default state as it is fixed base + default_root_pose = articulation.data.default_root_pose.torch.clone() + default_root_vel = articulation.data.default_root_vel.torch.clone() + default_root_pose[:, :3] = default_root_pose[:, :3] + translations + + torch.testing.assert_close(articulation.data.root_link_pose_w.torch, default_root_pose) + torch.testing.assert_close(articulation.data.root_com_vel_w.torch, default_root_vel) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.isaacsim_ci +def test_initialization_fixed_base_made_floating_base(sim, num_articulations, device, add_ground_plane): + """Test initialization for fixed base made floating-base using schema properties. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is floating base after modification + 3. All buffers have correct shapes + 4. The articulation can be simulated + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type="panda") + # Unfix root link by making it non-kinematic + articulation_cfg.spawn.articulation_props.fix_root_link = False + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that is floating base + assert not articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 9) + + # Cross-check binding shapes against cached counts. PhysX does this via + # ``root_view.max_dofs == shared_metatype.dof_count``; on OVPhysX + # ``root_view`` is the per-tensor-type bindings dict, so the equivalent + # invariant is that each per-DOF / per-link binding's shape agrees with + # the count cached on the asset. + for tt in (TT.DOF_POSITION, TT.DOF_VELOCITY, TT.DOF_STIFFNESS): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_joints + for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): + if tt in articulation.root_view: + assert articulation.root_view[tt].shape[1] == articulation.num_bodies + # Body-name ordering check is degenerate on OVPhysX: ``body_names`` is + # sourced from binding metadata (``sample.body_names``), so the PhysX + # ``link_paths[0]`` round-trip is a no-op here and is omitted. + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.isaacsim_ci +def test_out_of_range_default_joint_pos(sim, num_articulations, device, add_ground_plane): + """Test that the default joint position from configuration is out of range. + + This test verifies that: + 1. The articulation fails to initialize when joint positions are out of range + 2. The error is properly handled + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + # Create articulation + articulation_cfg = generate_articulation_cfg(articulation_type="panda").copy() + articulation_cfg.init_state.joint_pos = { + "panda_joint1": 10.0, + "panda_joint[2, 4]": -20.0, + } + + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + with pytest.raises(ValueError): + sim.reset() + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_out_of_range_default_joint_vel(sim, device): + """Test that the default joint velocity from configuration is out of range. + + This test verifies that: + 1. The articulation fails to initialize when joint velocities are out of range + 2. The error is properly handled + """ + articulation_cfg = FRANKA_PANDA_CFG.replace(prim_path="/World/Robot") + articulation_cfg.init_state.joint_vel = { + "panda_joint1": 100.0, + "panda_joint[2, 4]": -60.0, + } + articulation = Articulation(articulation_cfg) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + with pytest.raises(ValueError): + sim.reset() + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.isaacsim_ci +def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane): + """Test write_joint_limits_to_sim API and when default pos falls outside of the new limits. + + This test verifies that: + 1. Joint limits can be set correctly + 2. Default positions are preserved when setting new limits + 3. Joint limits can be set with indexing + 4. Invalid joint positions are properly handled + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + # Create articulation + articulation_cfg = generate_articulation_cfg(articulation_type="panda") + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device) + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + + # Get current default joint pos + default_joint_pos = articulation._data.default_joint_pos.torch.clone() + + # Set new joint limits + limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) + limits[..., 0] = (torch.rand(num_articulations, articulation.num_joints, device=device) + 5.0) * -1.0 + limits[..., 1] = torch.rand(num_articulations, articulation.num_joints, device=device) + 5.0 + articulation.write_joint_position_limit_to_sim_index(limits=limits) + + # Check new limits are in place + torch.testing.assert_close(articulation._data.joint_pos_limits.torch, limits) + torch.testing.assert_close(articulation._data.default_joint_pos.torch, default_joint_pos) + + # Set new joint limits with indexing + env_ids = torch.arange(1, device=device, dtype=torch.int32) + joint_ids = torch.arange(2, device=device, dtype=torch.int32) + limits = torch.zeros(env_ids.shape[0], joint_ids.shape[0], 2, device=device) + limits[..., 0] = (torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) + 5.0) * -1.0 + limits[..., 1] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) + 5.0 + articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) + + # Check new limits are in place + torch.testing.assert_close(articulation._data.joint_pos_limits.torch[env_ids][:, joint_ids], limits) + torch.testing.assert_close(articulation._data.default_joint_pos.torch, default_joint_pos) + + # Set new joint limits that invalidate default joint pos + limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) + limits[..., 0] = torch.rand(num_articulations, articulation.num_joints, device=device) * -0.1 + limits[..., 1] = torch.rand(num_articulations, articulation.num_joints, device=device) * 0.1 + articulation.write_joint_position_limit_to_sim_index(limits=limits) + + # Check if all values are within the bounds + default_joint_pos_torch = articulation._data.default_joint_pos.torch + within_bounds = (default_joint_pos_torch >= limits[..., 0]) & (default_joint_pos_torch <= limits[..., 1]) + assert torch.all(within_bounds) + + # Set new joint limits that invalidate default joint pos with indexing + limits = torch.zeros(env_ids.shape[0], joint_ids.shape[0], 2, device=device) + limits[..., 0] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) * -0.1 + limits[..., 1] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) * 0.1 + articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) + + # Check if all values are within the bounds + default_joint_pos_torch = articulation._data.default_joint_pos.torch + within_bounds = (default_joint_pos_torch[env_ids][:, joint_ids] >= limits[..., 0]) & ( + default_joint_pos_torch[env_ids][:, joint_ids] <= limits[..., 1] + ) + assert torch.all(within_bounds) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +def test_joint_effort_limits(sim, num_articulations, device, add_ground_plane): + """Validate joint effort limits via joint_effort_out_of_limit().""" + # Create articulation + articulation_cfg = generate_articulation_cfg(articulation_type="panda") + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device) + + # Minimal env wrapper exposing scene["robot"] + class _Env: + def __init__(self, art): + self.scene = {"robot": art} + + env = _Env(articulation) + robot_all = SceneEntityCfg(name="robot") + + sim.reset() + assert articulation.is_initialized + + # Case A: no clipping → should NOT terminate + articulation._data.computed_torque.torch.zero_() + articulation._data.applied_torque.torch.zero_() + out = joint_effort_out_of_limit(env, robot_all) # [N] + assert torch.all(~out) + + # Case B: simulate clipping → should terminate + articulation._data.computed_torque.torch.fill_(100.0) # pretend controller commanded 100 + articulation._data.applied_torque.torch.fill_(50.0) # pretend actuator clipped to 50 + out = joint_effort_out_of_limit(env, robot_all) # [N] + assert torch.all(out) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_external_force_buffer(sim, num_articulations, device): + """Test if external force buffer correctly updates in the force value is zero case. + + This test verifies that: + 1. External forces can be applied correctly + 2. Force buffers are updated properly + 3. Zero forces are handled correctly + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type="anymal") + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + + # play the simulator + sim.reset() + + # find bodies to apply the force + body_ids, _ = articulation.find_bodies("base") + + # reset root state + articulation.write_root_pose_to_sim_index(root_pose=articulation.data.default_root_pose.torch.clone()) + articulation.write_root_velocity_to_sim_index(root_velocity=articulation.data.default_root_vel.torch.clone()) + + # reset dof state + joint_pos, joint_vel = ( + articulation.data.default_joint_pos.torch, + articulation.data.default_joint_vel.torch, + ) + articulation.write_joint_position_to_sim_index(position=joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) + + # reset articulation + articulation.reset() + + # perform simulation + for step in range(5): + # initiate force tensor + external_wrench_b = torch.zeros(articulation.num_instances, len(body_ids), 6, device=sim.device) + + if step == 0 or step == 3: + # set a non-zero force + force = 1 + else: + # set a zero force + force = 0 + + # set force value + external_wrench_b[:, :, 0] = force + external_wrench_b[:, :, 3] = force + + # apply force + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + body_ids=body_ids, + ) + + # check if the articulation's force and torque buffers are correctly updated + for i in range(num_articulations): + assert articulation.permanent_wrench_composer.composed_force.torch[i, 0, 0].item() == force + assert articulation.permanent_wrench_composer.composed_torque.torch[i, 0, 0].item() == force + + # Check if the instantaneous wrench is correctly added to the permanent wrench + articulation.instantaneous_wrench_composer.add_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + body_ids=body_ids, + ) + + # apply action to the articulation + articulation.set_joint_position_target_index(target=articulation.data.default_joint_pos.torch.clone()) + articulation.write_data_to_sim() + + # perform step + sim.step() + + # update buffers + articulation.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_external_force_on_single_body(sim, num_articulations, device): + """Test application of external force on the base of the articulation. + + This test verifies that: + 1. External forces can be applied to specific bodies + 2. The forces affect the articulation's motion correctly + 3. The articulation responds to the forces as expected + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type="anymal") + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + # Play the simulator + sim.reset() + + # Find bodies to apply the force + body_ids, _ = articulation.find_bodies("base") + # Sample a large force + external_wrench_b = torch.zeros(articulation.num_instances, len(body_ids), 6, device=sim.device) + external_wrench_b[..., 1] = 1000.0 + + # Now we are ready! + for _ in range(5): + # reset root state + articulation.write_root_pose_to_sim_index(root_pose=articulation.data.default_root_pose.torch.clone()) + articulation.write_root_velocity_to_sim_index(root_velocity=articulation.data.default_root_vel.torch.clone()) + # reset dof state + joint_pos, joint_vel = ( + articulation.data.default_joint_pos.torch, + articulation.data.default_joint_vel.torch, + ) + articulation.write_joint_position_to_sim_index(position=joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) + # reset articulation + articulation.reset() + # apply force + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], body_ids=body_ids + ) + # perform simulation + for _ in range(100): + # apply action to the articulation + articulation.set_joint_position_target_index(target=articulation.data.default_joint_pos.torch.clone()) + articulation.write_data_to_sim() + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + # check condition that the articulations have fallen down + for i in range(num_articulations): + assert articulation.data.root_pos_w.torch[i, 2].item() < 0.2 + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_external_force_on_single_body_at_position(sim, num_articulations, device): + """Test application of external force on the base of the articulation at a given position. + + This test verifies that: + 1. External forces can be applied to specific bodies at a given position + 2. External forces can be applied to specific bodies in the global frame + 3. External forces are calculated and composed correctly + 4. The forces affect the articulation's motion correctly + 5. The articulation responds to the forces as expected + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type="anymal") + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + # Play the simulator + sim.reset() + + # Find bodies to apply the force + body_ids, _ = articulation.find_bodies("base") + # Sample a large force + external_wrench_b = torch.zeros(articulation.num_instances, len(body_ids), 6, device=sim.device) + external_wrench_b[..., 2] = 500.0 + external_wrench_positions_b = torch.zeros(articulation.num_instances, len(body_ids), 3, device=sim.device) + external_wrench_positions_b[..., 1] = 1.0 + + desired_force = torch.zeros(articulation.num_instances, len(body_ids), 3, device=sim.device) + desired_force[..., 2] = 1000.0 + desired_torque = torch.zeros(articulation.num_instances, len(body_ids), 3, device=sim.device) + desired_torque[..., 0] = 1000.0 + + # Now we are ready! + for i in range(5): + # reset root state + root_pose = articulation.data.default_root_pose.torch.clone() + root_pose[0, 0] = 2.5 # space them apart by 2.5m + + articulation.write_root_pose_to_sim_index(root_pose=root_pose) + articulation.write_root_velocity_to_sim_index(root_velocity=articulation.data.default_root_vel.torch.clone()) + # reset dof state + joint_pos, joint_vel = ( + articulation.data.default_joint_pos.torch, + articulation.data.default_joint_vel.torch, + ) + articulation.write_joint_position_to_sim_index(position=joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) + # reset articulation + articulation.reset() + # apply force + is_global = False + + if i % 2 == 0: + body_com_pos_w = articulation.data.body_com_pos_w.torch[:, body_ids, :3] + # is_global = True + external_wrench_positions_b[..., 0] = 0.0 + external_wrench_positions_b[..., 1] = 1.0 + external_wrench_positions_b[..., 2] = 0.0 + external_wrench_positions_b += body_com_pos_w + else: + external_wrench_positions_b[..., 0] = 0.0 + external_wrench_positions_b[..., 1] = 1.0 + external_wrench_positions_b[..., 2] = 0.0 + + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=external_wrench_positions_b, + body_ids=body_ids, + is_global=is_global, + ) + articulation.permanent_wrench_composer.add_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=external_wrench_positions_b, + body_ids=body_ids, + is_global=is_global, + ) + # perform simulation + for _ in range(100): + # apply action to the articulation + articulation.set_joint_position_target_index(target=articulation.data.default_joint_pos.torch.clone()) + articulation.write_data_to_sim() + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + # check condition that the articulations have fallen down + for i in range(num_articulations): + assert articulation.data.root_pos_w.torch[i, 2].item() < 0.2 + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_external_force_on_multiple_bodies(sim, num_articulations, device): + """Test application of external force on the legs of the articulation. + + This test verifies that: + 1. External forces can be applied to multiple bodies + 2. The forces affect the articulation's motion correctly + 3. The articulation responds to the forces as expected + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type="anymal") + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + + # Play the simulator + sim.reset() + + # Find bodies to apply the force + body_ids, _ = articulation.find_bodies(".*_SHANK") + # Sample a large force + external_wrench_b = torch.zeros(articulation.num_instances, len(body_ids), 6, device=sim.device) + external_wrench_b[..., 1] = 100.0 + + # Now we are ready! + for _ in range(5): + # reset root state + articulation.write_root_pose_to_sim_index(root_pose=articulation.data.default_root_pose.torch.clone()) + articulation.write_root_velocity_to_sim_index(root_velocity=articulation.data.default_root_vel.torch.clone()) + # reset dof state + joint_pos, joint_vel = ( + articulation.data.default_joint_pos.torch, + articulation.data.default_joint_vel.torch, + ) + articulation.write_joint_position_to_sim_index(position=joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) + # reset articulation + articulation.reset() + # apply force + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], body_ids=body_ids + ) + # perform simulation + for _ in range(100): + # apply action to the articulation + articulation.set_joint_position_target_index(target=articulation.data.default_joint_pos.torch.clone()) + articulation.write_data_to_sim() + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + # check condition + for i in range(num_articulations): + # since there is a moment applied on the articulation, the articulation should rotate + assert articulation.data.root_ang_vel_w.torch[i, 2].item() > 0.1 + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_external_force_on_multiple_bodies_at_position(sim, num_articulations, device): + """Test application of external force on the legs of the articulation at a given position. + + This test verifies that: + 1. External forces can be applied to multiple bodies at a given position + 2. External forces can be applied to multiple bodies in the global frame + 3. External forces are calculated and composed correctly + 4. The forces affect the articulation's motion correctly + 5. The articulation responds to the forces as expected + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type="anymal") + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + + # Play the simulator + sim.reset() + + # Find bodies to apply the force + body_ids, _ = articulation.find_bodies(".*_SHANK") + # Sample a large force + external_wrench_b = torch.zeros(articulation.num_instances, len(body_ids), 6, device=sim.device) + external_wrench_b[..., 2] = 500.0 + external_wrench_positions_b = torch.zeros(articulation.num_instances, len(body_ids), 3, device=sim.device) + external_wrench_positions_b[..., 1] = 1.0 + + desired_force = torch.zeros(articulation.num_instances, len(body_ids), 3, device=sim.device) + desired_force[..., 2] = 1000.0 + desired_torque = torch.zeros(articulation.num_instances, len(body_ids), 3, device=sim.device) + desired_torque[..., 0] = 1000.0 + + # Now we are ready! + for i in range(5): + # reset root state + articulation.write_root_pose_to_sim_index(root_pose=articulation.data.default_root_pose.torch.clone()) + articulation.write_root_velocity_to_sim_index(root_velocity=articulation.data.default_root_vel.torch.clone()) + # reset dof state + joint_pos, joint_vel = ( + articulation.data.default_joint_pos.torch, + articulation.data.default_joint_vel.torch, + ) + articulation.write_joint_position_to_sim_index(position=joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) + # reset articulation + articulation.reset() + + is_global = False + if i % 2 == 0: + body_com_pos_w = articulation.data.body_com_pos_w.torch[:, body_ids, :3] + is_global = True + external_wrench_positions_b[..., 0] = 0.0 + external_wrench_positions_b[..., 1] = 1.0 + external_wrench_positions_b[..., 2] = 0.0 + external_wrench_positions_b += body_com_pos_w + else: + external_wrench_positions_b[..., 0] = 0.0 + external_wrench_positions_b[..., 1] = 1.0 + external_wrench_positions_b[..., 2] = 0.0 + + # apply force + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=external_wrench_positions_b, + body_ids=body_ids, + is_global=is_global, + ) + articulation.permanent_wrench_composer.add_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=external_wrench_positions_b, + body_ids=body_ids, + is_global=is_global, + ) + # perform simulation + for _ in range(100): + # apply action to the articulation + articulation.set_joint_position_target_index(target=articulation.data.default_joint_pos.torch.clone()) + articulation.write_data_to_sim() + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + # check condition + for i in range(num_articulations): + # since there is a moment applied on the articulation, the articulation should rotate + assert torch.abs(articulation.data.root_ang_vel_w.torch[i, 2]).item() > 0.1 + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_loading_gains_from_usd(sim, num_articulations, device): + """Test that gains are loaded from USD file if actuator model has them as None. + + This test verifies that: + 1. Gains are loaded correctly from USD file + 2. Default gains are applied when not specified + 3. The gains match the expected values + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type="humanoid", stiffness=None, damping=None) + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + + # Play sim + sim.reset() + + # Expected gains + # -- Stiffness values + expected_stiffness = { + ".*_waist.*": 20.0, + ".*_upper_arm.*": 10.0, + "pelvis": 10.0, + ".*_lower_arm": 2.0, + ".*_thigh:0": 10.0, + ".*_thigh:1": 20.0, + ".*_thigh:2": 10.0, + ".*_shin": 5.0, + ".*_foot.*": 2.0, + } + indices_list, _, values_list = string_utils.resolve_matching_names_values( + expected_stiffness, articulation.joint_names + ) + expected_stiffness = torch.zeros(articulation.num_instances, articulation.num_joints, device=articulation.device) + expected_stiffness[:, indices_list] = torch.tensor(values_list, device=articulation.device) + # -- Damping values + expected_damping = { + ".*_waist.*": 5.0, + ".*_upper_arm.*": 5.0, + "pelvis": 5.0, + ".*_lower_arm": 1.0, + ".*_thigh:0": 5.0, + ".*_thigh:1": 5.0, + ".*_thigh:2": 5.0, + ".*_shin": 0.1, + ".*_foot.*": 1.0, + } + indices_list, _, values_list = string_utils.resolve_matching_names_values( + expected_damping, articulation.joint_names + ) + expected_damping = torch.zeros_like(expected_stiffness) + expected_damping[:, indices_list] = torch.tensor(values_list, device=articulation.device) + + # Check that gains are loaded from USD file + torch.testing.assert_close(articulation.actuators["body"].stiffness, expected_stiffness) + torch.testing.assert_close(articulation.actuators["body"].damping, expected_damping) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.isaacsim_ci +def test_setting_gains_from_cfg(sim, num_articulations, device, add_ground_plane): + """Test that gains are loaded from the configuration correctly. + + This test verifies that: + 1. Gains are loaded correctly from configuration + 2. The gains match the expected values + 3. The gains are applied correctly to the actuators + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type="humanoid") + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=sim.device + ) + + # Play sim + sim.reset() + + # Expected gains + expected_stiffness = torch.full( + (articulation.num_instances, articulation.num_joints), 10.0, device=articulation.device + ) + expected_damping = torch.full_like(expected_stiffness, 2.0) + + # Check that gains are loaded from USD file + torch.testing.assert_close(articulation.actuators["body"].stiffness, expected_stiffness) + torch.testing.assert_close(articulation.actuators["body"].damping, expected_damping) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_setting_gains_from_cfg_dict(sim, num_articulations, device): + """Test that gains are loaded from the configuration dictionary correctly. + + This test verifies that: + 1. Gains are loaded correctly from configuration dictionary + 2. The gains match the expected values + 3. The gains are applied correctly to the actuators + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type="humanoid") + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=sim.device + ) + # Play sim + sim.reset() + + # Expected gains + expected_stiffness = torch.full( + (articulation.num_instances, articulation.num_joints), 10.0, device=articulation.device + ) + expected_damping = torch.full_like(expected_stiffness, 2.0) + + # Check that gains are loaded from USD file + torch.testing.assert_close(articulation.actuators["body"].stiffness, expected_stiffness) + torch.testing.assert_close(articulation.actuators["body"].damping, expected_damping) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("vel_limit_sim", [1e5, None]) +@pytest.mark.parametrize("vel_limit", [1e2, None]) +@pytest.mark.parametrize("add_ground_plane", [False]) +@pytest.mark.isaacsim_ci +def test_setting_velocity_limit_implicit(sim, num_articulations, device, vel_limit_sim, vel_limit, add_ground_plane): + """Test setting of velocity limit for implicit actuators. + + This test verifies that: + 1. Velocity limits can be set correctly for implicit actuators + 2. The limits are applied correctly to the simulation + 3. The limits are handled correctly when both sim and non-sim limits are set + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + vel_limit_sim: The velocity limit to set in simulation + vel_limit: The velocity limit to set in actuator + """ + # create simulation + articulation_cfg = generate_articulation_cfg( + articulation_type="single_joint_implicit", + velocity_limit_sim=vel_limit_sim, + velocity_limit=vel_limit, + ) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, + num_articulations=num_articulations, + device=device, + ) + # Play sim + if vel_limit_sim is not None and vel_limit is not None: + with pytest.raises(ValueError): + sim.reset() + return + sim.reset() + + # read the values set into the simulation + physx_vel_limit = _read_binding_to_torch(articulation, TT.DOF_MAX_VELOCITY, device) + # check data buffer + torch.testing.assert_close(articulation.data.joint_velocity_limits.torch, physx_vel_limit) + # check actuator has simulation velocity limit + torch.testing.assert_close(articulation.actuators["joint"].velocity_limit_sim, physx_vel_limit) + # check that both values match for velocity limit + torch.testing.assert_close( + articulation.actuators["joint"].velocity_limit_sim, + articulation.actuators["joint"].velocity_limit, + ) + + if vel_limit_sim is None: + # Case 2: both velocity limit and velocity limit sim are not set + # This is the case where the velocity limit keeps its USD default value + # Case 3: velocity limit sim is not set but velocity limit is set + # For backwards compatibility, we do not set velocity limit to simulation + # Thus, both default to USD default value. + limit = articulation_cfg.spawn.joint_drive_props.max_velocity + else: + # Case 4: only velocity limit sim is set + # In this case, the velocity limit is set to the USD value + limit = vel_limit_sim + + # check max velocity is what we set + expected_velocity_limit = torch.full_like(physx_vel_limit, limit) + torch.testing.assert_close(physx_vel_limit, expected_velocity_limit) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("vel_limit_sim", [1e5, None]) +@pytest.mark.parametrize("vel_limit", [1e2, None]) +@pytest.mark.isaacsim_ci +def test_setting_velocity_limit_explicit(sim, num_articulations, device, vel_limit_sim, vel_limit): + """Test setting of velocity limit for explicit actuators.""" + articulation_cfg = generate_articulation_cfg( + articulation_type="single_joint_explicit", + velocity_limit_sim=vel_limit_sim, + velocity_limit=vel_limit, + ) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, + num_articulations=num_articulations, + device=device, + ) + # Play sim + sim.reset() + + # collect limit init values + physx_vel_limit = _read_binding_to_torch(articulation, TT.DOF_MAX_VELOCITY, device) + actuator_vel_limit = articulation.actuators["joint"].velocity_limit + actuator_vel_limit_sim = articulation.actuators["joint"].velocity_limit_sim + + # check data buffer for joint_velocity_limits_sim + torch.testing.assert_close(articulation.data.joint_velocity_limits.torch, physx_vel_limit) + # check actuator velocity_limit_sim is set to physx + torch.testing.assert_close(actuator_vel_limit_sim, physx_vel_limit) + + if vel_limit is not None: + expected_actuator_vel_limit = torch.full( + (articulation.num_instances, articulation.num_joints), + vel_limit, + device=articulation.device, + ) + # check actuator is set + torch.testing.assert_close(actuator_vel_limit, expected_actuator_vel_limit) + # check physx is not velocity_limit + assert not torch.allclose(actuator_vel_limit, physx_vel_limit) + else: + # check actuator velocity_limit is the same as the PhysX default + torch.testing.assert_close(actuator_vel_limit, physx_vel_limit) + + # simulation velocity limit is set to USD value unless user overrides + if vel_limit_sim is not None: + limit = vel_limit_sim + else: + limit = articulation_cfg.spawn.joint_drive_props.max_velocity + # check physx is set to expected value + expected_vel_limit = torch.full_like(physx_vel_limit, limit) + torch.testing.assert_close(physx_vel_limit, expected_vel_limit) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("effort_limit_sim", [1e5, None]) +@pytest.mark.parametrize("effort_limit", [1e2, 80.0, None]) +@pytest.mark.isaacsim_ci +def test_setting_effort_limit_implicit(sim, num_articulations, device, effort_limit_sim, effort_limit): + """Test setting of effort limit for implicit actuators. + + This test verifies the effort limit resolution logic for actuator models implemented in :class:`ActuatorBase`: + - Case 1: If USD value == actuator config value: values match correctly + - Case 2: If USD value != actuator config value: actuator config value is used + - Case 3: If actuator config value is None: USD value is used as default + """ + articulation_cfg = generate_articulation_cfg( + articulation_type="single_joint_implicit", + effort_limit_sim=effort_limit_sim, + effort_limit=effort_limit, + ) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, + num_articulations=num_articulations, + device=device, + ) + # Play sim + if effort_limit_sim is not None and effort_limit is not None: + with pytest.raises(ValueError): + sim.reset() + return + sim.reset() + + # obtain the physx effort limits + physx_effort_limit = _read_binding_to_torch(articulation, TT.DOF_MAX_FORCE, device) + + # check that the two are equivalent + torch.testing.assert_close( + articulation.actuators["joint"].effort_limit_sim, + articulation.actuators["joint"].effort_limit, + ) + torch.testing.assert_close(articulation.actuators["joint"].effort_limit_sim, physx_effort_limit) + + # decide the limit based on what is set + if effort_limit_sim is None and effort_limit is None: + limit = articulation_cfg.spawn.joint_drive_props.max_effort + elif effort_limit_sim is not None and effort_limit is None: + limit = effort_limit_sim + elif effort_limit_sim is None and effort_limit is not None: + limit = effort_limit + + # check that the max force is what we set + expected_effort_limit = torch.full_like(physx_effort_limit, limit) + torch.testing.assert_close(physx_effort_limit, expected_effort_limit) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("effort_limit_sim", [1e5, None]) +@pytest.mark.parametrize("effort_limit", [80.0, 1e2, None]) +@pytest.mark.isaacsim_ci +def test_setting_effort_limit_explicit(sim, num_articulations, device, effort_limit_sim, effort_limit): + """Test setting of effort limit for explicit actuators. + + This test verifies the effort limit resolution logic for actuator models implemented in :class:`ActuatorBase`: + - Case 1: If USD value == actuator config value: values match correctly + - Case 2: If USD value != actuator config value: actuator config value is used + - Case 3: If actuator config value is None: USD value is used as default + + """ + + articulation_cfg = generate_articulation_cfg( + articulation_type="single_joint_explicit", + effort_limit_sim=effort_limit_sim, + effort_limit=effort_limit, + ) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, + num_articulations=num_articulations, + device=device, + ) + # Play sim + sim.reset() + + # usd default effort limit is set to 80 + usd_default_effort_limit = 80.0 + + # collect limit init values + physx_effort_limit = _read_binding_to_torch(articulation, TT.DOF_MAX_FORCE, device) + actuator_effort_limit = articulation.actuators["joint"].effort_limit + actuator_effort_limit_sim = articulation.actuators["joint"].effort_limit_sim + + # check actuator effort_limit_sim is set to physx + torch.testing.assert_close(actuator_effort_limit_sim, physx_effort_limit) + + if effort_limit is not None: + expected_actuator_effort_limit = torch.full_like(actuator_effort_limit, effort_limit) + # check actuator is set + torch.testing.assert_close(actuator_effort_limit, expected_actuator_effort_limit) + + # check physx effort limit does not match the one explicit actuator has + assert not (torch.allclose(actuator_effort_limit, physx_effort_limit)) + else: + # When effort_limit is None, actuator should use USD default values + expected_actuator_effort_limit = torch.full_like(physx_effort_limit, usd_default_effort_limit) + torch.testing.assert_close(actuator_effort_limit, expected_actuator_effort_limit) + + # when using explicit actuators, the limits are set to high unless user overrides + if effort_limit_sim is not None: + limit = effort_limit_sim + else: + limit = ActuatorBase._DEFAULT_MAX_EFFORT_SIM # type: ignore + # check physx internal value matches the expected sim value + expected_effort_limit = torch.full_like(physx_effort_limit, limit) + torch.testing.assert_close(actuator_effort_limit_sim, expected_effort_limit) + torch.testing.assert_close(physx_effort_limit, expected_effort_limit) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_reset(sim, num_articulations, device): + """Test that reset method works properly.""" + articulation_cfg = generate_articulation_cfg(articulation_type="humanoid") + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=device + ) + + # Play the simulator + sim.reset() + + # Now we are ready! + # reset articulation + articulation.reset() + + # Reset should zero external forces and torques + assert not articulation._instantaneous_wrench_composer.active + assert not articulation._permanent_wrench_composer.active + assert torch.count_nonzero(articulation._instantaneous_wrench_composer.composed_force.torch) == 0 + assert torch.count_nonzero(articulation._instantaneous_wrench_composer.composed_torque.torch) == 0 + assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_force.torch) == 0 + assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_torque.torch) == 0 + + if num_articulations > 1: + num_bodies = articulation.num_bodies + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=torch.ones((num_articulations, num_bodies, 3), device=device), + torques=torch.ones((num_articulations, num_bodies, 3), device=device), + ) + articulation.instantaneous_wrench_composer.add_forces_and_torques_index( + forces=torch.ones((num_articulations, num_bodies, 3), device=device), + torques=torch.ones((num_articulations, num_bodies, 3), device=device), + ) + articulation.reset(env_ids=torch.tensor([0], device=device)) + assert articulation._instantaneous_wrench_composer.active + assert articulation._permanent_wrench_composer.active + assert torch.count_nonzero(articulation._instantaneous_wrench_composer.composed_force.torch) == num_bodies * 3 + assert torch.count_nonzero(articulation._instantaneous_wrench_composer.composed_torque.torch) == num_bodies * 3 + assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_force.torch) == num_bodies * 3 + assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_torque.torch) == num_bodies * 3 + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.isaacsim_ci +def test_apply_joint_command(sim, num_articulations, device, add_ground_plane): + """Test applying of joint position target functions correctly for a robotic arm.""" + articulation_cfg = generate_articulation_cfg(articulation_type="panda") + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=device + ) + + # Play the simulator + sim.reset() + + for _ in range(100): + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + # reset dof state + joint_pos = articulation.data.default_joint_pos.torch.clone() + joint_pos[:, 3] = 0.0 + + # apply action to the articulation + articulation.set_joint_position_target_index(target=joint_pos) + articulation.write_data_to_sim() + + for _ in range(100): + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + # Check that current joint position is not the same as default joint position, meaning + # the articulation moved. We can't check that it reached its desired joint position as the gains + # are not properly tuned + assert not torch.allclose(articulation.data.joint_pos.torch, joint_pos) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("with_offset", [True, False]) +@pytest.mark.isaacsim_ci +def test_body_root_state(sim, num_articulations, device, with_offset): + """Test for reading the `body_state_w` property. + + This test verifies that: + 1. Body states can be read correctly + 2. States are correct with and without offsets + 3. States are consistent across different devices + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + with_offset: Whether to test with offset + """ + sim._app_control_on_stop_handle = None + articulation_cfg = generate_articulation_cfg(articulation_type="single_joint_implicit") + articulation, env_pos = generate_articulation(articulation_cfg, num_articulations, device) + env_idx = torch.tensor([x for x in range(num_articulations)], device=device, dtype=torch.int32) + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10, "Possible reference leak for articulation" + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized, "Articulation is not initialized" + # Check that fixed base + assert articulation.is_fixed_base, "Articulation is not a fixed base" + + # Resolve body indices by name (ordering may differ across physics backends) + root_idx = articulation.body_names.index("CenterPivot") + arm_idx = articulation.body_names.index("Arm") + + # change center of mass offset from link frame + if with_offset: + offset = [0.5, 0.0, 0.0] + else: + offset = [0.0, 0.0, 0.0] + + # create com offsets — apply offset to the Arm body + num_bodies = articulation.num_bodies + com = _read_binding_to_torch(articulation, TT.BODY_COM_POSE, device) + link_offset = [1.0, 0.0, 0.0] # the offset from CenterPivot to Arm frames + new_com = torch.tensor(offset, device=device).repeat(num_articulations, 1, 1) + com[:, arm_idx, :3] = new_com.squeeze(-2) + # PhysX uses ``root_view.set_coms``; OVPhysX wraps the wheel + # ``BODY_COM_POSE`` write in :meth:`set_coms_index` (wp.transformf contract). + articulation.set_coms_index( + coms=wp.from_torch(com.contiguous(), dtype=wp.transformf), + env_ids=wp.from_torch(env_idx, dtype=wp.int32), + ) + + # check they are set + torch.testing.assert_close(_read_binding_to_torch(articulation, TT.BODY_COM_POSE, device), com) + + for i in range(50): + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + # get state properties + root_link_pose_w = articulation.data.root_link_pose_w.torch + root_link_vel_w = articulation.data.root_link_vel_w.torch + root_com_pose_w = articulation.data.root_com_pose_w.torch + root_com_vel_w = articulation.data.root_com_vel_w.torch + body_link_pose_w = articulation.data.body_link_pose_w.torch + body_link_vel_w = articulation.data.body_link_vel_w.torch + body_com_pose_w = articulation.data.body_com_pose_w.torch + body_com_vel_w = articulation.data.body_com_vel_w.torch + + if with_offset: + # get joint state + joint_pos = articulation.data.joint_pos.torch.unsqueeze(-1) + joint_vel = articulation.data.joint_vel.torch.unsqueeze(-1) + + # LINK state + # angular velocity should be the same for both COM and link frames + torch.testing.assert_close(root_com_vel_w[..., 3:], root_link_vel_w[..., 3:]) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) + + # lin_vel arm + lin_vel_gt = torch.zeros(num_articulations, num_bodies, 3, device=device) + vx = -(link_offset[0]) * joint_vel * torch.sin(joint_pos) + vy = torch.zeros(num_articulations, 1, 1, device=device) + vz = (link_offset[0]) * joint_vel * torch.cos(joint_pos) + lin_vel_gt[:, arm_idx, :] = torch.cat([vx, vy, vz], dim=-1).squeeze(-2) + + # linear velocity of root link should be zero + torch.testing.assert_close(lin_vel_gt[:, root_idx, :], root_link_vel_w[..., :3], atol=1e-3, rtol=1e-1) + # linear velocity of pendulum link should be + torch.testing.assert_close(lin_vel_gt, body_link_vel_w[..., :3], atol=1e-3, rtol=1e-1) + + # ang_vel + torch.testing.assert_close(root_com_vel_w[..., 3:], root_link_vel_w[..., 3:]) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) + + # COM state + # position and orientation shouldn't match for the _state_com_w but everything else will + pos_gt = torch.zeros(num_articulations, num_bodies, 3, device=device) + px = (link_offset[0] + offset[0]) * torch.cos(joint_pos) + py = torch.zeros(num_articulations, 1, 1, device=device) + pz = (link_offset[0] + offset[0]) * torch.sin(joint_pos) + pos_gt[:, arm_idx, :] = torch.cat([px, py, pz], dim=-1).squeeze(-2) + pos_gt += env_pos.unsqueeze(-2).repeat(1, num_bodies, 1) + torch.testing.assert_close(pos_gt[:, root_idx, :], root_com_pose_w[..., :3], atol=1e-3, rtol=1e-1) + torch.testing.assert_close(pos_gt, body_com_pose_w[..., :3], atol=1e-3, rtol=1e-1) + + # orientation + com_quat_b = articulation.data.body_com_quat_b.torch + com_quat_w = math_utils.quat_mul(body_link_pose_w[..., 3:], com_quat_b) + torch.testing.assert_close(com_quat_w, body_com_pose_w[..., 3:]) + torch.testing.assert_close(com_quat_w[:, root_idx, :], root_com_pose_w[..., 3:]) + + # angular velocity should be the same for both COM and link frames + torch.testing.assert_close(root_com_vel_w[..., 3:], root_link_vel_w[..., 3:]) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) + else: + # single joint center of masses are at link frames so they will be the same + torch.testing.assert_close(root_link_pose_w, root_com_pose_w) + torch.testing.assert_close(root_com_vel_w, root_link_vel_w) + torch.testing.assert_close(body_link_pose_w, body_com_pose_w) + torch.testing.assert_close(body_com_vel_w, body_link_vel_w) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("with_offset", [True, False]) +@pytest.mark.parametrize("state_location", ["com", "link"]) +@pytest.mark.parametrize("gravity_enabled", [False]) +@pytest.mark.isaacsim_ci +def test_write_root_state(sim, num_articulations, device, with_offset, state_location, gravity_enabled): + """Test the setters for root_state using both the link frame and center of mass as reference frame. + + This test verifies that: + 1. Root states can be written correctly + 2. States are correct with and without offsets + 3. States can be written for both COM and link frames + 4. States are consistent across different devices + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + with_offset: Whether to test with offset + state_location: Whether to test COM or link frame + """ + sim._app_control_on_stop_handle = None + articulation_cfg = generate_articulation_cfg(articulation_type="anymal") + articulation, env_pos = generate_articulation(articulation_cfg, num_articulations, device) + env_idx = torch.tensor([x for x in range(num_articulations)], device=device, dtype=torch.int32) + + # Play sim + sim.reset() + + # change center of mass offset from link frame + if with_offset: + offset = torch.tensor([1.0, 0.0, 0.0]).repeat(num_articulations, 1, 1) + else: + offset = torch.tensor([0.0, 0.0, 0.0]).repeat(num_articulations, 1, 1) + + # create com offsets + com = _read_binding_to_torch(articulation, TT.BODY_COM_POSE, device) + new_com = offset.to(device) + com[:, 0, :3] = new_com.squeeze(-2) + # See test_body_root_state for the PhysX → OVPhysX setter substitution. + articulation.set_coms_index( + coms=wp.from_torch(com.contiguous(), dtype=wp.transformf), + env_ids=wp.from_torch(env_idx, dtype=wp.int32), + ) + + # check they are set + torch.testing.assert_close(_read_binding_to_torch(articulation, TT.BODY_COM_POSE, device), com) + + rand_state = torch.zeros(num_articulations, 13, device=device) + rand_state[..., :7] = articulation.data.default_root_pose.torch + rand_state[..., :3] += env_pos + # make quaternion a unit vector + rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) + + env_idx = env_idx.to(device) + for i in range(10): + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + if state_location == "com": + if i % 2 == 0: + articulation.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7]) + articulation.write_root_com_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) + else: + articulation.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_idx) + articulation.write_root_com_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_idx) + elif state_location == "link": + if i % 2 == 0: + articulation.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7]) + articulation.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) + else: + articulation.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_idx) + articulation.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_idx) + + if state_location == "com": + torch.testing.assert_close(rand_state[..., :7], articulation.data.root_com_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], articulation.data.root_com_vel_w.torch) + elif state_location == "link": + torch.testing.assert_close(rand_state[..., :7], articulation.data.root_link_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], articulation.data.root_link_vel_w.torch) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_body_incoming_joint_wrench_b_single_joint(sim, num_articulations, device): + """Test the data.body_incoming_joint_wrench_b buffer is populated correctly and statically correct for single joint. + + This test verifies that: + 1. The body incoming joint wrench buffer has correct shape + 2. The wrench values are statically correct for a single joint + 3. The wrench values match expected values from gravity and external forces + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + articulation_cfg = generate_articulation_cfg(articulation_type="single_joint_implicit") + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=device + ) + + # Play the simulator + sim.reset() + + # Resolve body indices by name (ordering may differ across physics backends) + arm_idx = articulation.body_names.index("Arm") + root_idx = articulation.body_names.index("CenterPivot") + # apply external force + external_force_vector_b = torch.zeros((num_articulations, articulation.num_bodies, 3), device=device) + external_force_vector_b[:, arm_idx, 1] = 10.0 # 10 N in Y direction + external_torque_vector_b = torch.zeros((num_articulations, articulation.num_bodies, 3), device=device) + external_torque_vector_b[:, arm_idx, 2] = 10.0 # 10 Nm in z direction + + # apply action to the articulation + joint_pos = torch.ones_like(articulation.data.joint_pos.torch) * 1.5708 / 2.0 + articulation.write_joint_position_to_sim_index( + position=torch.ones_like(articulation.data.joint_pos.torch), + ) + articulation.write_joint_velocity_to_sim_index( + velocity=torch.zeros_like(articulation.data.joint_vel.torch), + ) + articulation.set_joint_position_target_index(target=joint_pos) + articulation.write_data_to_sim() + for _ in range(50): + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_force_vector_b, torques=external_torque_vector_b + ) + articulation.write_data_to_sim() + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + # check shape + assert articulation.data.body_incoming_joint_wrench_b.torch.shape == ( + num_articulations, + articulation.num_bodies, + 6, + ) + + # calculate expected static + mass = articulation.data.body_mass.torch.to("cpu") + pos_w = articulation.data.body_pos_w.torch + quat_w = articulation.data.body_quat_w.torch + + mass_link2 = mass[:, arm_idx].view(num_articulations, -1) + gravity = torch.tensor(sim.cfg.gravity, device="cpu").repeat(num_articulations, 1).view((num_articulations, 3)) + + # NOTE: the com and link pose for single joint are colocated + weight_vector_w = mass_link2 * gravity + # expected wrench from link mass and external wrench + # PhysX reports the incoming joint wrench as the force FROM body0 ONTO body1 (body1's frame). + # The USD asset defines body0=CenterPivot, body1=Arm, so the wrench is the constraint/support + # force from CenterPivot onto Arm, expressed in Arm's frame. + # In static equilibrium this equals -(gravity + external forces on Arm). + total_force_w = weight_vector_w.to(device) + math_utils.quat_apply( + quat_w[:, arm_idx, :], external_force_vector_b[:, arm_idx, :] + ) + total_torque_w = torch.cross( + pos_w[:, arm_idx, :].to(device) - pos_w[:, root_idx, :].to(device), + total_force_w, + dim=-1, + ) + math_utils.quat_apply(quat_w[:, arm_idx, :], external_torque_vector_b[:, arm_idx, :]) + expected_wrench = torch.zeros((num_articulations, 6), device=device) + expected_wrench[:, :3] = math_utils.quat_apply( + math_utils.quat_conjugate(quat_w[:, arm_idx, :]), + -total_force_w, + ) + expected_wrench[:, 3:] = math_utils.quat_apply( + math_utils.quat_conjugate(quat_w[:, arm_idx, :]), + -total_torque_w, + ) + + # check value of last joint wrench + torch.testing.assert_close( + expected_wrench, + articulation.data.body_incoming_joint_wrench_b.torch[:, arm_idx, :].squeeze(1), + atol=1e-2, + rtol=1e-3, + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_setting_articulation_root_prim_path(sim, device): + """Test that the articulation root prim path can be set explicitly.""" + sim._app_control_on_stop_handle = None + # Create articulation + articulation_cfg = generate_articulation_cfg(articulation_type="humanoid") + articulation_cfg.articulation_root_prim_path = "/torso" + articulation, _ = generate_articulation(articulation_cfg, 1, device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation._is_initialized + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_setting_invalid_articulation_root_prim_path(sim, device): + """Test that the articulation root prim path can be set explicitly.""" + sim._app_control_on_stop_handle = None + # Create articulation + articulation_cfg = generate_articulation_cfg(articulation_type="humanoid") + articulation_cfg.articulation_root_prim_path = "/non_existing_prim_path" + articulation, _ = generate_articulation(articulation_cfg, 1, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + with pytest.raises(RuntimeError): + sim.reset() + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("gravity_enabled", [False]) +@pytest.mark.isaacsim_ci +@pytest.mark.xfail(reason=_FK_ON_DEMAND_GAP_REASON, strict=False) +def test_write_joint_state_data_consistency(sim, num_articulations, device, gravity_enabled): + """Test the setters for root_state using both the link frame and center of mass as reference frame. + + This test verifies that after write_joint_state_to_sim operations: + 1. state, com_state, link_state value consistency + 2. body_pose, link + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + sim._app_control_on_stop_handle = None + articulation_cfg = generate_articulation_cfg(articulation_type="anymal") + articulation, env_pos = generate_articulation(articulation_cfg, num_articulations, device) + env_idx = torch.tensor([x for x in range(num_articulations)]) + + # Play sim + sim.reset() + + limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) + limits[..., 0] = (torch.rand(num_articulations, articulation.num_joints, device=device) + 5.0) * -1.0 + limits[..., 1] = torch.rand(num_articulations, articulation.num_joints, device=device) + 5.0 + articulation.write_joint_position_limit_to_sim_index(limits=limits) + + from torch.distributions import Uniform + + joint_pos_limits = articulation.data.joint_pos_limits.torch + joint_vel_limits = articulation.data.joint_vel_limits.torch + pos_dist = Uniform(joint_pos_limits[..., 0], joint_pos_limits[..., 1]) + vel_dist = Uniform(-joint_vel_limits, joint_vel_limits) + + original_body_link_pose_w = articulation.data.body_link_pose_w.torch.clone() + original_body_com_vel_w = articulation.data.body_com_vel_w.torch.clone() + + rand_joint_pos = pos_dist.sample() + rand_joint_vel = vel_dist.sample() + + articulation.write_joint_position_to_sim_index(position=rand_joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=rand_joint_vel) + # make sure valued updated + body_link_pose_w = articulation.data.body_link_pose_w.torch + body_com_vel_w = articulation.data.body_com_vel_w.torch + original_body_states = torch.cat([original_body_link_pose_w, original_body_com_vel_w], dim=-1) + body_state_w = torch.cat([body_link_pose_w, body_com_vel_w], dim=-1) + assert torch.count_nonzero(original_body_states[:, 1:] != body_state_w[:, 1:]) > ( + len(original_body_states[:, 1:]) / 2 + ) + # validate body - link consistency + body_link_vel_w = articulation.data.body_link_vel_w.torch + torch.testing.assert_close(body_link_pose_w, articulation.data.body_link_pose_w.torch) + # skip lin_vel because it differs from link frame, this should be fine because we are only checking + # if velocity update is triggered, which can be determined by comparing angular velocity + torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) + + # validate link - com conistency + body_com_pos_b = articulation.data.body_com_pos_b.torch + body_com_quat_b = articulation.data.body_com_quat_b.torch + expected_com_pos, expected_com_quat = math_utils.combine_frame_transforms( + body_link_pose_w[..., :3].view(-1, 3), + body_link_pose_w[..., 3:].view(-1, 4), + body_com_pos_b.view(-1, 3), + body_com_quat_b.view(-1, 4), + ) + body_com_pos_w = articulation.data.body_com_pos_w.torch + body_com_quat_w = articulation.data.body_com_quat_w.torch + torch.testing.assert_close(expected_com_pos.view(len(env_idx), -1, 3), body_com_pos_w) + torch.testing.assert_close(expected_com_quat.view(len(env_idx), -1, 4), body_com_quat_w) + + # validate body - com consistency + body_com_lin_vel_w = articulation.data.body_com_lin_vel_w.torch + body_com_ang_vel_w = articulation.data.body_com_ang_vel_w.torch + torch.testing.assert_close(body_com_vel_w[..., :3], body_com_lin_vel_w) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_com_ang_vel_w) + + # validate pos_w, quat_w, pos_b, quat_b is consistent with pose_w and pose_b + expected_com_pose_w = torch.cat((body_com_pos_w, body_com_quat_w), dim=2) + expected_com_pose_b = torch.cat((body_com_pos_b, body_com_quat_b), dim=2) + body_pos_w = articulation.data.body_pos_w.torch + body_quat_w = articulation.data.body_quat_w.torch + expected_body_pose_w = torch.cat((body_pos_w, body_quat_w), dim=2) + body_link_pos_w = articulation.data.body_link_pos_w.torch + body_link_quat_w = articulation.data.body_link_quat_w.torch + expected_body_link_pose_w = torch.cat((body_link_pos_w, body_link_quat_w), dim=2) + body_com_pose_w = articulation.data.body_com_pose_w.torch + body_com_pose_b = articulation.data.body_com_pose_b.torch + body_pose_w = articulation.data.body_pose_w.torch + body_link_pose_w_fresh = articulation.data.body_link_pose_w.torch + torch.testing.assert_close(body_com_pose_w, expected_com_pose_w) + torch.testing.assert_close(body_com_pose_b, expected_com_pose_b) + torch.testing.assert_close(body_pose_w, expected_body_pose_w) + torch.testing.assert_close(body_link_pose_w_fresh, expected_body_link_pose_w) + + # validate pose_w is consistent with individual properties + body_vel_w = articulation.data.body_vel_w.torch + body_com_vel_w_fresh = articulation.data.body_com_vel_w.torch + torch.testing.assert_close(body_pose_w, body_link_pose_w) + torch.testing.assert_close(body_vel_w, body_com_vel_w) + torch.testing.assert_close(body_link_pose_w_fresh, body_link_pose_w) + torch.testing.assert_close(body_com_pose_w, articulation.data.body_com_pose_w.torch) + torch.testing.assert_close(body_vel_w, body_com_vel_w_fresh) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_spatial_tendons(sim, num_articulations, device): + """Test spatial tendons apis. + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation has spatial tendons + 3. All buffers have correct shapes + 4. The articulation can be simulated + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + # skip test if Isaac Sim version is less than 5.0 + if has_kit() and get_isaac_sim_version().major < 5: + pytest.skip("Spatial tendons are not supported in Isaac Sim < 5.0. Please update to Isaac Sim 5.0 or later.") + return + articulation_cfg = generate_articulation_cfg(articulation_type="spatial_tendon_test_asset") + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that fixed base + assert articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 3) + assert articulation.data.body_mass.torch.shape == (num_articulations, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (num_articulations, articulation.num_bodies, 9) + assert articulation.num_spatial_tendons == 1 + + articulation.set_spatial_tendon_stiffness_index(stiffness=10.0) + articulation.set_spatial_tendon_limit_stiffness_index(limit_stiffness=10.0) + articulation.set_spatial_tendon_damping_index(damping=10.0) + articulation.set_spatial_tendon_offset_index(offset=10.0) + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_write_joint_frictions_to_sim(sim, num_articulations, device, add_ground_plane): + """Test applying of joint position target functions correctly for a robotic arm.""" + articulation_cfg = generate_articulation_cfg(articulation_type="panda") + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=device + ) + + # Play the simulator + sim.reset() + + for _ in range(100): + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + # apply action to the articulation + dynamic_friction = torch.rand(num_articulations, articulation.num_joints, device=device) + viscous_friction = torch.rand(num_articulations, articulation.num_joints, device=device) + friction = torch.rand(num_articulations, articulation.num_joints, device=device) + + # Guarantee that the dynamic friction is not greater than the static friction + dynamic_friction = torch.min(dynamic_friction, friction) + + # The static friction must be set first to be sure the dynamic friction is not greater than static + # when both are set. + articulation.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff=friction, + joint_dynamic_friction_coeff=dynamic_friction, + joint_viscous_friction_coeff=viscous_friction, + ) + articulation.write_data_to_sim() + + for _ in range(100): + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + friction_props_from_sim = _read_binding_to_torch(articulation, TT.DOF_FRICTION_PROPERTIES, "cpu") + joint_friction_coeff_sim = friction_props_from_sim[:, :, 0] + joint_dynamic_friction_coeff_sim = friction_props_from_sim[:, :, 1] + joint_viscous_friction_coeff_sim = friction_props_from_sim[:, :, 2] + assert torch.allclose(joint_dynamic_friction_coeff_sim, dynamic_friction.cpu()) + assert torch.allclose(joint_viscous_friction_coeff_sim, viscous_friction.cpu()) + assert torch.allclose(joint_friction_coeff_sim, friction.cpu()) + + # For Isaac Sim >= 5.0: also test the combined API that can set dynamic and viscous via + # write_joint_friction_coefficient_to_sim; reset the sim to isolate this path. + if has_kit() and get_isaac_sim_version().major >= 5: + # Reset simulator to ensure a clean state for the alternative API path + sim.reset() + + # Warm up a few steps to populate buffers + for _ in range(100): + sim.step() + articulation.update(sim.cfg.dt) + + # New random coefficients + dynamic_friction_2 = torch.rand(num_articulations, articulation.num_joints, device=device) + viscous_friction_2 = torch.rand(num_articulations, articulation.num_joints, device=device) + friction_2 = torch.rand(num_articulations, articulation.num_joints, device=device) + + # Guarantee that the dynamic friction is not greater than the static friction + dynamic_friction_2 = torch.min(dynamic_friction_2, friction_2) + + # Use the combined setter to write all three at once + articulation.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff=friction_2, + joint_dynamic_friction_coeff=dynamic_friction_2, + joint_viscous_friction_coeff=viscous_friction_2, + ) + articulation.write_data_to_sim() + + # Step to let sim ingest new params and refresh data buffers + for _ in range(100): + sim.step() + articulation.update(sim.cfg.dt) + + friction_props_from_sim_2 = _read_binding_to_torch(articulation, TT.DOF_FRICTION_PROPERTIES, "cpu") + joint_friction_coeff_sim_2 = friction_props_from_sim_2[:, :, 0] + friction_dynamic_coef_sim_2 = friction_props_from_sim_2[:, :, 1] + friction_viscous_coeff_sim_2 = friction_props_from_sim_2[:, :, 2] + + # Validate values propagated + assert torch.allclose(friction_viscous_coeff_sim_2, viscous_friction_2.cpu()) + assert torch.allclose(friction_dynamic_coef_sim_2, dynamic_friction_2.cpu()) + assert torch.allclose(joint_friction_coeff_sim_2, friction_2.cpu()) + + +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("articulation_type", ["panda"]) +@pytest.mark.isaacsim_ci +@pytest.mark.xfail(reason=_MATERIAL_GAP_REASON, strict=False) +def test_set_material_properties(sim, num_articulations, device, add_ground_plane, articulation_type): + """Test getting and setting material properties (friction/restitution) of articulation shapes.""" + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=device + ) + + # Play the simulator + sim.reset() + + # Get number of shapes from the articulation + max_shapes = articulation.root_view.max_shapes + + # Generate random material properties: (static_friction, dynamic_friction, restitution) + materials = torch.empty(num_articulations, max_shapes, 3, device="cpu").uniform_(0.0, 1.0) + # Ensure dynamic friction <= static friction + materials[..., 1] = torch.min(materials[..., 0], materials[..., 1]) + + # Set material properties via the PhysX view-level API + env_ids = torch.arange(num_articulations, dtype=torch.int32) + articulation.root_view.set_material_properties( + wp.from_torch(materials, dtype=wp.float32), wp.from_torch(env_ids, dtype=wp.int32) + ) + + # Simulate physics + sim.step() + articulation.update(sim.cfg.dt) + + # Get material properties from simulation + materials_check = wp.to_torch(articulation.root_view.get_material_properties()) + + # Check if material properties are set correctly + torch.testing.assert_close(materials_check, materials) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "--maxfail=1"]) diff --git a/source/isaaclab_ovphysx/test/assets/test_articulation_helpers.py b/source/isaaclab_ovphysx/test/assets/test_articulation_helpers.py new file mode 100644 index 000000000000..9be6185022da --- /dev/null +++ b/source/isaaclab_ovphysx/test/assets/test_articulation_helpers.py @@ -0,0 +1,142 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""OVPhysX-only unit tests for articulation helpers. + +These tests cover OVPhysX-specific scaffolding (USD tendon-scope resolution, +mock binding-set shape contracts) that has no PhysX equivalent and therefore +does not appear in the PhysX-mirrored ``test_articulation.py``. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import warp as wp + +from pxr import Sdf, Usd, UsdPhysics + +# The CI isaaclab_ov* pattern unintentionally collects isaaclab_ovphysx tests, +# but the ovphysx wheel is not installed in that environment. Skip gracefully +# so the isaaclab_ov CI pipeline is not blocked by an unrelated dependency. +pytest.importorskip("ovphysx.types", reason="ovphysx wheel not installed") + +from isaaclab_ovphysx.assets.articulation.articulation import Articulation # noqa: E402 +from isaaclab_ovphysx.physics import OvPhysxManager # noqa: E402 +from isaaclab_ovphysx.test.mock_interfaces.views import MockOvPhysxBindingSet # noqa: E402 + +wp.init() + + +def _define_tendon_joint(stage: Usd.Stage, path: str, schema_name: str) -> None: + """Define a revolute joint prim with a tendon schema marker.""" + joint = UsdPhysics.RevoluteJoint.Define(stage, path) + schemas = Sdf.TokenListOp() + schemas.explicitItems = [schema_name] + joint.GetPrim().SetMetadata("apiSchemas", schemas) + + +def _make_articulation_root_stage(tmp_path) -> str: + """Create a stage with one relevant articulation subtree and unrelated joints elsewhere.""" + stage = Usd.Stage.CreateInMemory() + stage.DefinePrim("/World", "Xform") + stage.DefinePrim("/World/envs", "Xform") + stage.DefinePrim("/World/envs/env_0", "Xform") + stage.DefinePrim("/World/envs/env_0/Robot", "Xform") + stage.DefinePrim("/World/envs/env_0/Robot/root", "Xform") + stage.DefinePrim("/World/unrelated", "Xform") + + _define_tendon_joint( + stage, + "/World/envs/env_0/Robot/root/fixed_joint", + "PhysxTendonAxisRootAPI:inst0", + ) + _define_tendon_joint( + stage, + "/World/envs/env_0/Robot/root/spatial_joint", + "PhysxTendonAttachmentRootAPI:inst0", + ) + _define_tendon_joint( + stage, + "/World/unrelated/unrelated_fixed_joint", + "PhysxTendonAxisRootAPI:inst0", + ) + _define_tendon_joint( + stage, + "/World/unrelated/unrelated_spatial_joint", + "PhysxTendonAttachmentLeafAPI:inst0", + ) + + stage_path = tmp_path / "scene.usda" + stage.Export(str(stage_path)) + return str(stage_path) + + +def _make_articulation_shell() -> Articulation: + """Create a minimal ovphysx articulation shell for tendon processing tests.""" + articulation = object.__new__(Articulation) + bindings = MockOvPhysxBindingSet( + num_instances=1, + num_joints=2, + num_bodies=2, + num_fixed_tendons=1, + num_spatial_tendons=1, + ) + object.__setattr__(articulation, "_bindings", bindings.bindings) + object.__setattr__(articulation, "_articulation_root_path", "/World/envs/env_0/Robot/root") + object.__setattr__(articulation, "_initialize_handle", None) + object.__setattr__(articulation, "_invalidate_initialize_handle", None) + object.__setattr__(articulation, "_prim_deletion_handle", None) + object.__setattr__(articulation, "_debug_vis_handle", None) + object.__setattr__( + articulation, + "_data", + SimpleNamespace( + _num_fixed_tendons=0, + _num_spatial_tendons=0, + fixed_tendon_names=[], + spatial_tendon_names=[], + ), + ) + return articulation + + +def test_process_tendons_scopes_to_articulation_root(tmp_path): + """Tendon discovery should ignore joints that live outside the current articulation subtree.""" + articulation = _make_articulation_shell() + stage_path = _make_articulation_root_stage(tmp_path) + old_stage_path = OvPhysxManager._stage_path + OvPhysxManager._stage_path = stage_path + try: + articulation._process_tendons() + finally: + OvPhysxManager._stage_path = old_stage_path + + assert articulation.fixed_tendon_names == ["fixed_joint"] + assert articulation.spatial_tendon_names == ["spatial_joint"] + assert articulation._data.fixed_tendon_names == ["fixed_joint"] + assert articulation._data.spatial_tendon_names == ["spatial_joint"] + + +def test_mock_binding_set_rigid_object_shapes(): + pytest.importorskip("isaaclab_ovphysx.tensor_types").RIGID_BODY_POSE # gates on wheel + from isaaclab_ovphysx import tensor_types as TT + from isaaclab_ovphysx.test.mock_interfaces.views import MockOvPhysxBindingSet + + bindings = MockOvPhysxBindingSet( + num_instances=4, + num_joints=0, + num_bodies=1, + asset_kind="rigid_object", + ) + assert bindings.bindings[TT.RIGID_BODY_POSE].shape == (4, 7) + assert bindings.bindings[TT.RIGID_BODY_VELOCITY].shape == (4, 6) + assert bindings.bindings[TT.RIGID_BODY_WRENCH].shape == (4, 9) + assert bindings.bindings[TT.RIGID_BODY_MASS].shape == (4,) + assert bindings.bindings[TT.RIGID_BODY_INERTIA].shape == (4, 9) + # Articulation-only bindings must be absent + assert TT.DOF_POSITION not in bindings.bindings + assert TT.LINK_WRENCH not in bindings.bindings