diff --git a/.github/actions/run-package-tests/action.yml b/.github/actions/run-package-tests/action.yml index ce3472fd6021..a283bb1eb8bd 100644 --- a/.github/actions/run-package-tests/action.yml +++ b/.github/actions/run-package-tests/action.yml @@ -44,7 +44,7 @@ inputs: test-k-expr: description: >- Global pytest -k expression applied inside every per-file pytest run - spawned by tools/conftest.py (combined with device-split selectors). + spawned by tools/conftest.py. default: '' required: false test-jobs: diff --git a/.github/actions/run-tests/action.yml b/.github/actions/run-tests/action.yml index 9f04d670b191..04e4efa706ab 100644 --- a/.github/actions/run-tests/action.yml +++ b/.github/actions/run-tests/action.yml @@ -44,7 +44,7 @@ inputs: test-k-expr: description: >- Global pytest -k expression applied inside every per-file pytest run - spawned by tools/conftest.py (combined with device-split selectors). + spawned by tools/conftest.py. Unlike pytest-options, this reaches the individual test processes, so it can deselect parametrized cases (e.g. "not ovphysx"). default: '' diff --git a/.github/workflows/tools-tests.yml b/.github/workflows/tools-tests.yml index 4913e366e2fb..3062204d6dd9 100644 --- a/.github/workflows/tools-tests.yml +++ b/.github/workflows/tools-tests.yml @@ -60,7 +60,7 @@ jobs: - name: Run tools tests env: - # The tests import the modules under test (crash_journal, _device_split) directly. + # The tests import the modules under test (crash_journal) directly. PYTHONPATH: tools # Files are listed explicitly rather than collected from tools/: test_settings.py is a # configuration module, not a test, and tools/test/ needs the full Isaac Lab install. @@ -68,7 +68,7 @@ jobs: # --noconftest keeps tools/conftest.py out of the session. That file is the CI test # orchestrator - its pytest_sessionstart scans source/ and scripts/ and runs the whole # suite, so loading it here would ignore the files named below. - run: python3 -m pytest tools/test_crash_journal.py tools/test_device_split.py tools/test_file_scheduler.py + run: python3 -m pytest tools/test_crash_journal.py tools/test_file_scheduler.py .github/actions/_lib/test_registry_credential_fallback.py -v --noconftest - uses: astral-sh/setup-uv@v6 diff --git a/docs/source/concepts/ovphysx.rst b/docs/source/concepts/ovphysx.rst index 12a6cd0b2188..6307062dabf7 100644 --- a/docs/source/concepts/ovphysx.rst +++ b/docs/source/concepts/ovphysx.rst @@ -149,7 +149,7 @@ Then run a small backend smoke test: .. code-block:: bash - uv run --extra ovphysx --extra test python -m pytest source/isaaclab_ov/test/assets/test_rigid_object.py::test_initialization -k cpu + uv run --extra ovphysx --extra test python -m pytest source/isaaclab_ov/test/assets/test_rigid_object.py::test_rigid_object_real_ovphysx_seams -k cpu To try a task that declares an OvPhysX physics preset, use the same preset CLI syntax as the other backends: diff --git a/pyproject.toml b/pyproject.toml index 75fb748dbcc7..92e9fd2e4aef 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -360,7 +360,6 @@ ignore-words-list = "aheared,collet,followng,haa,publically,rouines,slq,collapsa markers = [ "isaacsim_ci: mark test to run in isaacsim ci", - "device_split: re-invoke this file once per device (CPU and GPU) in CI due to process-global device locks (e.g., ovphysx<=0.3.7 gap G5)", "windows_ci: mark test to run on Windows platforms in CI", "arm_ci: mark test to run on ARM platforms in CI (e.g. NVIDIA DGX Spark)", "unit: test exercises isolated logic and does not launch the simulator", diff --git a/source/isaaclab/changelog.d/antoiner-asset-tests-redesign.skip b/source/isaaclab/changelog.d/antoiner-asset-tests-redesign.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab/isaaclab/test/utils/actuator_equivalence.py b/source/isaaclab/isaaclab/test/utils/actuator_equivalence.py index 7c83b950123d..46b183d81994 100644 --- a/source/isaaclab/isaaclab/test/utils/actuator_equivalence.py +++ b/source/isaaclab/isaaclab/test/utils/actuator_equivalence.py @@ -12,9 +12,7 @@ * :class:`EquivalenceAssertionsMixin` with the trajectory/telemetry ``test_*_match`` oracles, * dummy TorchScript checkpoint factories for the neural actuator tests, -* mock scene/env plumbing for driving ``randomize_actuator_gains``, -* :class:`ActuatorStateResetBase` with the per-env actuator state reset - scenario. +* mock scene/env plumbing for driving ``randomize_actuator_gains``. This module must stay free of backend packages (``isaaclab_newton``, ``isaaclab_physx``) and of ``isaaclab_assets``; everything backend-specific @@ -25,11 +23,8 @@ import tempfile import torch -import warp as wp -from ... import sim as sim_utils from ...actuators import DCMotorCfg, DelayedPDActuatorCfg, IdealPDActuatorCfg, ImplicitActuatorCfg -from ...sim import SimulationCfg, build_simulation_context # --------------------------------------------------------------------------- # Actuator configurations under test @@ -44,17 +39,6 @@ ), } -DC_MOTOR_ACTUATORS = { - "legs": DCMotorCfg( - joint_names_expr=[".*HAA", ".*HFE", ".*KFE"], - saturation_effort=120.0, - actuator_effort_limit=80.0, - actuator_velocity_limit=7.5, - stiffness={".*": 40.0}, - damping={".*": 5.0}, - ), -} - MIXED_ACTUATORS = { "hips": IdealPDActuatorCfg( joint_names_expr=[".*HAA"], @@ -313,142 +297,3 @@ def build_dr_term(env, asset_name, joint_ids=None): asset_cfg = cfg.params["asset_cfg"] asset_cfg.finalize(env.device) return term, asset_cfg - - -# --------------------------------------------------------------------------- -# Per-env reset: actuator state isolation -# --------------------------------------------------------------------------- - - -class ActuatorStateResetBase: - """Reset must clear the actuator state buffers for the requested envs only. - - Inspects ``adapter.actuators[i].state.delay_state.num_pushes`` directly: - - * After warmup, ``num_pushes > 0`` for every DOF (buffer was populated). - * After ``articulation.reset(env_ids=[0])``, the entries for env 0's DOFs - must be ``0`` and the entries for env 1's DOFs must remain ``> 0``. - - Done independently on Lab and Newton paths. Backend-specific twins mix - this into a ``unittest.TestCase`` and provide :meth:`_make_sim_cfg`, - :meth:`_make_articulation`, and :meth:`_get_adapter`. - """ - - RESET_ENV: int = 0 - UNCHANGED_ENV: int = 1 - NUM_ENVS: int = 2 - DT: float = 1.0 / 120.0 - TARGET_OFFSET: float = 0.1 # [rad] added to initial joint positions - RESET_WARMUP_STEPS: int = 3 - - def _make_sim_cfg(self, use_newton_actuators: bool) -> SimulationCfg: - """Return the backend simulation config for the run.""" - raise NotImplementedError - - def _make_articulation(self): - """Construct the backend articulation (DelayedPD on all joints) at ``/World/Env_.*/Robot``.""" - raise NotImplementedError - - def _get_adapter(self, articulation): - """Return the Newton actuator adapter that owns ``articulation``'s actuators.""" - raise NotImplementedError - - def _build_and_warm(self, *, use_newton_actuators: bool): - ctx = build_simulation_context( - device="cuda:0", - gravity_enabled=True, - add_ground_plane=True, - sim_cfg=self._make_sim_cfg(use_newton_actuators), - ) - sim = ctx.__enter__() - sim._app_control_on_stop_handle = None - for i in range(self.NUM_ENVS): - sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0)) - articulation = self._make_articulation() - sim.reset() - - init_pos = wp.to_torch(articulation.data.joint_pos).clone() - target_pos = init_pos + self.TARGET_OFFSET - target_vel = torch.zeros_like(init_pos) - articulation.set_joint_position_target_index(target=target_pos) - articulation.set_joint_velocity_target_index(target=target_vel) - for _ in range(self.RESET_WARMUP_STEPS): - articulation.write_data_to_sim() - sim.step() - articulation.update(self.DT) - return ctx, sim, articulation - - def test_newton_state_reset_isolated_to_reset_env(self): - """Newton: ``num_pushes`` zeroes for env 0's DOFs only after reset of [0].""" - ctx, sim, articulation = self._build_and_warm(use_newton_actuators=True) - try: - adapter = self._get_adapter(articulation) - self.assertIsNotNone(adapter) - # Find a DelayedPD actuator (it's the only one with delay_state). - stateful_pairs = [ - (act, st) - for act, st in zip(adapter.actuators, adapter._states_a) - if st is not None and getattr(st, "delay_state", None) is not None - ] - self.assertGreater(len(stateful_pairs), 0, "expected at least one DelayedPD actuator with delay_state") - - for act, state in stateful_pairs: - pushes_before = state.delay_state.num_pushes.numpy() - self.assertTrue( - (pushes_before > 0).all(), - "expected non-zero num_pushes for all DOFs after warmup", - ) - - articulation.reset(env_ids=torch.tensor([self.RESET_ENV], device=articulation.device, dtype=torch.long)) - - # Map each entry of ``act.indices`` to its env via the adapter's - # per-env DOF count. On the Newton backend the adapter is model-wide - # (includes free-joint DOFs on floating-base articulations); on - # PhysX it is per-articulation — ``adapter.num_joints`` is the - # correct stride in both cases. - for act, state in stateful_pairs: - pushes_after = state.delay_state.num_pushes.numpy() - indices_np = act.indices.numpy() - for i, global_dof in enumerate(indices_np): - env = int(global_dof) // adapter.num_joints - if env == self.RESET_ENV: - self.assertEqual( - int(pushes_after[i]), - 0, - f"DOF {i} (env {env}) should be reset to 0, got {pushes_after[i]}", - ) - else: - self.assertGreater( - int(pushes_after[i]), - 0, - f"DOF {i} (env {env}) was NOT in reset env_ids but num_pushes is 0", - ) - finally: - ctx.__exit__(None, None, None) - - def test_lab_state_reset_isolated_to_reset_env(self): - """Reset environments accept fresh commands while the remaining environments retain their history.""" - ctx, sim, articulation = self._build_and_warm(use_newton_actuators=False) - try: - commands = articulation.actuators.target_command - old_target = commands.position.torch.clone() - # Episode reset samples the configured lag; construction alone leaves actuator lag at zero. - articulation.reset() - commands.set_position_index(value=old_target) - articulation.write_data_to_sim() - new_target = old_target + 0.02 - articulation.reset(env_ids=torch.tensor([self.RESET_ENV], device=articulation.device)) - commands.set_position_index(value=new_target) - articulation.write_data_to_sim() - expected_target = old_target.clone() - expected_target[self.RESET_ENV] = new_target[self.RESET_ENV] - for actuator in articulation.actuators.values(): - joints = actuator.joint_indices - demand = ( - actuator.stiffness * (expected_target[:, joints] - articulation.data.joint_pos.torch[:, joints]) - - actuator.damping * articulation.data.joint_vel.torch[:, joints] - ) - torch.testing.assert_close(actuator.computed_effort, demand) - - finally: - ctx.__exit__(None, None, None) diff --git a/source/isaaclab/test/actuators/test_newton_adapter.py b/source/isaaclab/test/actuators/test_newton_adapter.py index 21fac108251a..055da4b5807e 100644 --- a/source/isaaclab/test/actuators/test_newton_adapter.py +++ b/source/isaaclab/test/actuators/test_newton_adapter.py @@ -3,20 +3,23 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Tests for the PhysX-side Newton actuator adapter.""" +"""Tests for the PhysX-side Newton actuator adapter and its host runtime.""" from pathlib import Path from types import SimpleNamespace +from unittest.mock import Mock import numpy as np import pytest import torch +import warp as wp from newton.actuators import ClampingDCMotor, ClampingMaxEffort, ClampingPositionBased, DrivePD from pxr import Usd, UsdGeom, UsdPhysics from isaaclab.actuators import ActuatorBaseCfg, DCMotor, DCMotorCfg, DelayedPDActuatorCfg, RemotizedPDActuatorCfg from isaaclab.actuators.newton import NewtonActuatorAdapter +from isaaclab.actuators.newton.physx_runtime import PhysxActuatorRuntime from isaaclab.sim import SimulationCfg, SimulationContext from isaaclab.sim.schemas.schemas_actuators import ( author_actuator_prims, @@ -59,7 +62,8 @@ def __reduce__(self): return _write_pickle_marker, (self.marker_path,) -def _make_actuator_stage() -> Usd.Stage: +def _make_robot_stage() -> Usd.Stage: + """Create a robot with one revolute joint per entry of ``_JOINT_NAMES`` and no actuators.""" stage = Usd.Stage.CreateInMemory() UsdGeom.Xform.Define(stage, "/World/Robot") @@ -70,7 +74,11 @@ def _make_actuator_stage() -> Usd.Stage: joints = [UsdPhysics.RevoluteJoint.Define(stage, f"/World/Robot/{name}") for name in _JOINT_NAMES] for joint, body in zip(joints, bodies, strict=True): joint.CreateBody1Rel().SetTargets([body.GetPath()]) + return stage + +def _make_actuator_stage() -> Usd.Stage: + stage = _make_robot_stage() author_actuator_prims( stage, "/World/Robot", @@ -183,6 +191,35 @@ def test_from_usd_groups_by_structure_and_preserves_per_dof_values(): } == {(10.0, 20.0), (11.0, 21.0)} +def test_is_stateful_only_with_delay_state(): + """Report the adapter as stateful when an actuator keeps a delay buffer, and stateless otherwise.""" + stateless_stage = _make_robot_stage() + author_actuator_prims( + stateless_stage, + "/World/Robot", + { + "dc": DCMotorCfg( + joint_names_expr=[".*"], + stiffness=1.0, + damping=0.1, + actuator_effort_limit=2.0, + actuator_velocity_limit=3.0, + saturation_effort=4.0, + ) + }, + ) + for stage, expected in ((_make_actuator_stage(), True), (stateless_stage, False)): + adapter = NewtonActuatorAdapter.from_usd( + stage=stage, + joint_names=_JOINT_NAMES, + num_envs=2, + num_joints=len(_JOINT_NAMES), + device="cpu", + articulation_prim_path="/World/Robot", + ) + assert adapter.is_stateful is expected + + @pytest.mark.parametrize( ("configured_limit", "expected_limits"), [(None, (71.0, 72.0)), ({"pd_a": 5.0}, (5.0, 0.0))], @@ -311,3 +348,86 @@ def test_define_actuator_properties_authors_every_spawned_articulation(monkeypat for robot_path in robot_paths: actuators = [prim for prim in stage.Traverse() if prim.GetTypeName() == "NewtonActuator"] assert any(str(prim.GetPath()).startswith(robot_path + "/") for prim in actuators), robot_path + + +# --------------------------------------------------------------------------- +# CUDA-graph handling of the host-side PhysX actuator runtime +# --------------------------------------------------------------------------- + + +def _runtime() -> PhysxActuatorRuntime: + """Create a runtime for a CUDA articulation with a mocked logger.""" + return PhysxActuatorRuntime(SimpleNamespace(device="cuda:0"), logger=Mock()) + + +def _failing_capture(fail_at: int) -> type: + """Return a ``wp.ScopedCapture`` double whose capture number ``fail_at`` (0-based) raises on entry.""" + + class _Capture: + count = 0 + + def __init__(self, *args, **kwargs) -> None: + self.index = type(self).count + type(self).count += 1 + self.graph = object() + + def __enter__(self): + if self.index >= fail_at: + raise RuntimeError("capture unavailable") + return self + + def __exit__(self, *exc_info) -> bool: + return False + + return _Capture + + +def test_graph_capture_failure_restores_adapter_state_and_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) -> None: + """A failed second capture discards the graphs and restores the state swapped by the first capture.""" + state_a, state_b = object(), object() + runtime = _runtime() + runtime.adapter = SimpleNamespace(_states_a=state_a, _states_b=state_b) + + def _swap_adapter_state(*args, **kwargs) -> None: + runtime.adapter._states_a, runtime.adapter._states_b = runtime.adapter._states_b, runtime.adapter._states_a + + monkeypatch.setattr(wp, "ScopedCapture", _failing_capture(fail_at=1)) + monkeypatch.setattr(runtime, "_run_native_actuator_kernels", _swap_adapter_state) + + runtime._capture_native_actuator_graphs(SimpleNamespace(), 0.01) + + assert runtime.native_actuator_graphs == () + assert runtime.adapter._states_a is state_a + assert runtime.adapter._states_b is state_b + runtime._logger.warning.assert_called_once() + + +def test_compute_runs_eagerly_after_graph_capture_failure(monkeypatch: pytest.MonkeyPatch) -> None: + """A graphable adapter whose capture fails still computes the current command eagerly.""" + capture = _failing_capture(fail_at=0) + runtime = _runtime() + runtime.adapter = SimpleNamespace(is_stateful=False, is_all_graphable=True, _states_a=object(), _states_b=object()) + eager_compute = Mock() + monkeypatch.setattr(wp, "get_device", lambda device: SimpleNamespace(is_cuda=True, is_capturing=False)) + monkeypatch.setattr(wp, "ScopedCapture", capture) + monkeypatch.setattr(runtime, "_run_native_actuator_kernels", eager_compute) + collection = SimpleNamespace() + + runtime.compute(collection, 0.01) + runtime.compute(collection, 0.01) + + # The failed capture is not retried, and every step falls back to the eager kernels. + assert capture.count == 1 + assert runtime.native_actuator_graphs == () + assert eager_compute.call_count == 2 + eager_compute.assert_called_with(collection, 0.01) + + +def test_stateful_actuator_rejects_outer_cuda_capture(monkeypatch: pytest.MonkeyPatch) -> None: + """Stateful adapters cannot safely mutate their buffers inside an outer CUDA capture.""" + runtime = _runtime() + runtime.adapter = SimpleNamespace(is_stateful=True) + monkeypatch.setattr(wp, "get_device", lambda device: SimpleNamespace(is_cuda=True, is_capturing=True)) + + with pytest.raises(RuntimeError, match="stateful Newton actuators cannot run inside an outer CUDA graph capture"): + runtime.compute(SimpleNamespace(), 0.01) diff --git a/source/isaaclab/test/assets/_articulation_iface_test_utils.py b/source/isaaclab/test/assets/_articulation_iface_test_utils.py deleted file mode 100644 index d852c5f13e35..000000000000 --- a/source/isaaclab/test/assets/_articulation_iface_test_utils.py +++ /dev/null @@ -1,504 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -# ignore private usage of variables warning -# pyright: reportPrivateUsage=none - -"""Shared mocked articulation backend factories for interface tests.""" - -from unittest.mock import MagicMock - -import _iface_test_boot # noqa: F401 (starts the runtime) - -import numpy as np -import torch -import warp as wp - -from isaaclab.assets.articulation.articulation_cfg import ArticulationCfg -from isaaclab.utils.wrench_composer import WrenchComposer - -BACKENDS: list[str] = [] -BACKEND_UNAVAILABLE_REASONS: dict[str, str] = {} - -try: - from isaaclab_physx.assets.articulation.articulation import Articulation as PhysXArticulation - from isaaclab_physx.assets.articulation.articulation_data import ArticulationData as PhysXArticulationData - from isaaclab_physx.physics import PhysxManager as SimulationManager - from isaaclab_physx.physics.physx_manager import PhysxSceneDataBackend - from isaaclab_physx.test.fixtures.views import MockArticulationViewWarp as PhysXMockArticulationViewWarp -except ImportError as error: - BACKEND_UNAVAILABLE_REASONS["physx"] = f"{type(error).__name__}: {error}" -else: - # PhysX data classes need gravity even though interface tests do not create a physics scene. - _mock_physics_sim_view = MagicMock() - _mock_physics_sim_view.get_gravity.return_value = (0.0, 0.0, -9.81) - SimulationManager.get_physics_sim_view = MagicMock(return_value=_mock_physics_sim_view) - SimulationManager._scene_data_backend = PhysxSceneDataBackend() - - BACKENDS.append("physx") - -try: - from isaaclab_newton.assets.articulation.articulation import Articulation as NewtonArticulation - from isaaclab_newton.assets.articulation.articulation_data import ArticulationData as NewtonArticulationData - from isaaclab_newton.test.fixtures.views import MockNewtonArticulationView as NewtonMockArticulationView -except ImportError as error: - BACKEND_UNAVAILABLE_REASONS["newton"] = f"{type(error).__name__}: {error}" -else: - BACKENDS.append("newton") - -try: - import ovphysx # noqa: F401 - - from isaaclab_ov.assets.articulation.articulation import Articulation as OvPhysxArticulation - from isaaclab_ov.assets.articulation.articulation_data import ArticulationData as OvPhysxArticulationData - from isaaclab_ov.test.fixtures.views import MockOvPhysxBindingSet - from isaaclab_ov.physics.ovphysx_manager import OvPhysxManager, OvPhysxSceneDataBackend -except ImportError as error: - BACKEND_UNAVAILABLE_REASONS["ovphysx"] = f"{type(error).__name__}: {error}" -else: - # Writers bump the scene-data transform version that ``initialize()`` would normally create. - OvPhysxManager._scene_data_backend = OvPhysxSceneDataBackend() - BACKENDS.append("ovphysx") - - -def create_physx_articulation( - num_instances: int = 2, - num_joints: int = 6, - num_bodies: int = 7, - num_fixed_tendons: int = 0, - num_spatial_tendons: int = 0, - device: str = "cuda:0", - is_fixed_base: bool = False, - joint_ordering: tuple[str, ...] | None = None, - body_ordering: tuple[str, ...] | None = None, -): - """Create a test Articulation instance with mocked dependencies.""" - joint_names = [f"joint_{i}" for i in range(num_joints)] - body_names = [f"body_{i}" for i in range(num_bodies)] - fixed_tendon_names = [f"fixed_tendon_{i}" for i in range(num_fixed_tendons)] - spatial_tendon_names = [f"spatial_tendon_{i}" for i in range(num_spatial_tendons)] - - articulation = object.__new__(PhysXArticulation) - - articulation.cfg = ArticulationCfg( - prim_path="/World/Robot", - soft_joint_pos_limit_factor=1.0, - actuators={}, - joint_ordering=joint_ordering, - body_ordering=body_ordering, - ) - object.__setattr__(articulation, "_sim_cfg", None) - object.__setattr__(articulation, "_fixed_tendon_target_dirty", False) - - # Create PhysX mock view - mock_view = PhysXMockArticulationViewWarp( - count=num_instances, - num_links=num_bodies, - num_dofs=num_joints, - device=device, - max_fixed_tendons=num_fixed_tendons, - max_spatial_tendons=num_spatial_tendons, - ) - mock_view.set_random_mock_data() - mock_view._noop_setters = True - - # Set up the mock view's metatype for accessing names/counts - mock_metatype = MagicMock() - mock_metatype.fixed_base = is_fixed_base - mock_metatype.dof_count = num_joints - mock_metatype.link_count = num_bodies - mock_metatype.dof_names = joint_names - mock_metatype.link_names = body_names - object.__setattr__(mock_view, "_shared_metatype", mock_metatype) - - object.__setattr__(articulation, "_root_view", mock_view) - object.__setattr__(articulation, "_device", device) - articulation._clamped_default_count = wp.zeros(1, dtype=wp.int32, device=device) - - # We can't call the initialize method here, because we don't have a good mock for the actuators yet. - # We need to set the _data attribute manually. - - # Create ArticulationData instance (SimulationManager already mocked at module level) - data = PhysXArticulationData(mock_view, device) - object.__setattr__(articulation, "_data", data) - - # Set tendon names on articulation and data - object.__setattr__(articulation, "_fixed_tendon_names", fixed_tendon_names) - object.__setattr__(articulation, "_spatial_tendon_names", spatial_tendon_names) - data.fixed_tendon_names = fixed_tendon_names - data.spatial_tendon_names = spatial_tendon_names - - # Create wrench composers (pass articulation which has num_instances, num_bodies, device properties) - mock_inst_wrench = WrenchComposer(articulation, supports_world_at_com=True) - mock_perm_wrench = WrenchComposer(articulation, supports_world_at_com=True) - object.__setattr__(articulation, "_instantaneous_wrench_composer", mock_inst_wrench) - object.__setattr__(articulation, "_permanent_wrench_composer", mock_perm_wrench) - - # Prevent __del__ / _clear_callbacks from raising AttributeError - 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) - - # Set up other required attributes - object.__setattr__(articulation, "_ALL_INDICES", wp.array(np.arange(num_instances, dtype=np.int32), device=device)) - object.__setattr__( - articulation, "_ALL_BODY_INDICES", wp.array(np.arange(num_bodies, dtype=np.int32), device=device) - ) - object.__setattr__( - articulation, "_ALL_JOINT_INDICES", wp.array(np.arange(num_joints, dtype=np.int32), device=device) - ) - - # Tendon index arrays - object.__setattr__( - articulation, - "_ALL_FIXED_TENDON_INDICES", - wp.array(np.arange(num_fixed_tendons, dtype=np.int32), device=device), - ) - object.__setattr__( - articulation, - "_ALL_SPATIAL_TENDON_INDICES", - wp.array(np.arange(num_spatial_tendons, dtype=np.int32), device=device), - ) - - # Warp arrays for set_external_force_and_torque - object.__setattr__( - articulation, "_ALL_INDICES_WP", wp.array(np.arange(num_instances, dtype=np.int32), device=device) - ) - object.__setattr__( - articulation, "_ALL_BODY_INDICES_WP", wp.array(np.arange(num_bodies, dtype=np.int32), device=device) - ) - - object.__setattr__(articulation, "_joint_pos_target_backend", None) - object.__setattr__(articulation, "_joint_vel_target_backend", None) - object.__setattr__(articulation, "_joint_effort_target_backend", None) - object.__setattr__(articulation, "_body_wrench_force_backend", None) - object.__setattr__(articulation, "_body_wrench_torque_backend", None) - - articulation._resolve_and_install_ordering_maps() - articulation._ordering_configure_backend_staging() - - # Cached .view(wp.float32) wrappers - object.__setattr__(articulation, "_root_link_pose_w_f32", None) - object.__setattr__(articulation, "_root_com_vel_w_f32", None) - object.__setattr__(articulation, "_root_link_vel_w_f32", None) - object.__setattr__(articulation, "_inst_wrench_force_f32", None) - object.__setattr__(articulation, "_inst_wrench_torque_f32", None) - object.__setattr__(articulation, "_perm_wrench_force_f32", None) - object.__setattr__(articulation, "_perm_wrench_torque_f32", None) - - # Pre-allocated pinned CPU buffers for PhysX TensorAPI writes - N, J, B = num_instances, num_joints, num_bodies - object.__setattr__(articulation, "_sim_env_ids", wp.empty(N, dtype=wp.int32, device=device)) - object.__setattr__(articulation, "_sim_env_ids_views", {}) - cpu_env_ids = wp.array(np.arange(N, dtype=np.int32), device="cpu") - object.__setattr__(articulation, "_cpu_env_ids_all", cpu_env_ids) - object.__setattr__(articulation, "_cpu_env_ids", wp.empty(N, dtype=wp.int32, device="cpu", pinned=True)) - object.__setattr__(articulation, "_cpu_env_ids_views", {}) - object.__setattr__(articulation, "_cpu_joint_stiffness", wp.zeros((N, J), dtype=wp.float32, device="cpu")) - object.__setattr__(articulation, "_cpu_joint_damping", wp.zeros((N, J), dtype=wp.float32, device="cpu")) - object.__setattr__(articulation, "_cpu_joint_pos_limits", wp.zeros((N, J, 2), dtype=wp.float32, device="cpu")) - object.__setattr__(articulation, "_cpu_joint_vel_limits", wp.zeros((N, J), dtype=wp.float32, device="cpu")) - object.__setattr__(articulation, "_cpu_joint_effort_limits", wp.zeros((N, J), dtype=wp.float32, device="cpu")) - object.__setattr__(articulation, "_cpu_joint_armature", wp.zeros((N, J), dtype=wp.float32, device="cpu")) - object.__setattr__(articulation, "_cpu_joint_friction_props", wp.zeros((N, J, 3), dtype=wp.float32, device="cpu")) - object.__setattr__(articulation, "_cpu_body_mass", wp.zeros((N, B), dtype=wp.float32, device="cpu")) - object.__setattr__(articulation, "_cpu_body_coms", wp.zeros((N, B, 7), dtype=wp.float32, device="cpu")) - object.__setattr__(articulation, "_cpu_body_inertia", wp.zeros((N, B, 9), dtype=wp.float32, device="cpu")) - - articulation._process_actuators_cfg() - - return articulation, mock_view - - -def create_ovphysx_articulation( - num_instances: int = 2, - num_joints: int = 6, - num_bodies: int = 7, - num_fixed_tendons: int = 0, - num_spatial_tendons: int = 0, - device: str = "cuda:0", - is_fixed_base: bool = False, - joint_ordering: tuple[str, ...] | None = None, - body_ordering: tuple[str, ...] | None = None, -): - """Create a test OvPhysX Articulation instance with mocked tensor bindings.""" - joint_names = [f"joint_{i}" for i in range(num_joints)] - body_names = [f"body_{i}" for i in range(num_bodies)] - - articulation = object.__new__(OvPhysxArticulation) - - articulation.cfg = ArticulationCfg( - prim_path="/World/Robot", - soft_joint_pos_limit_factor=1.0, - actuators={}, - joint_ordering=joint_ordering, - body_ordering=body_ordering, - ) - object.__setattr__(articulation, "_sim_cfg", None) - object.__setattr__(articulation, "_fixed_tendon_target_dirty", False) - - # Create mock binding set - mock_bindings = MockOvPhysxBindingSet( - num_instances=num_instances, - num_joints=num_joints, - num_bodies=num_bodies, - is_fixed_base=is_fixed_base, - joint_names=joint_names, - body_names=body_names, - num_fixed_tendons=num_fixed_tendons, - num_spatial_tendons=num_spatial_tendons, - ) - mock_bindings.set_random_data() - - fixed_tendon_names = [f"fixed_tendon_{i}" for i in range(num_fixed_tendons)] - spatial_tendon_names = [f"spatial_tendon_{i}" for i in range(num_spatial_tendons)] - - object.__setattr__(articulation, "_device", device) - object.__setattr__(articulation, "_ovphysx", MagicMock()) - object.__setattr__(articulation, "_root_view", mock_bindings.view) - object.__setattr__(articulation, "_bindings", mock_bindings.bindings) - object.__setattr__(articulation, "_num_instances", num_instances) - object.__setattr__(articulation, "_num_joints", num_joints) - object.__setattr__(articulation, "_num_bodies", num_bodies) - object.__setattr__(articulation, "_is_fixed_base", is_fixed_base) - object.__setattr__(articulation, "_joint_names", joint_names) - object.__setattr__(articulation, "_body_names", body_names) - object.__setattr__(articulation, "_fixed_tendon_names", fixed_tendon_names) - object.__setattr__(articulation, "_spatial_tendon_names", spatial_tendon_names) - object.__setattr__(articulation, "_num_fixed_tendons", num_fixed_tendons) - object.__setattr__(articulation, "_num_spatial_tendons", num_spatial_tendons) - - # Create ArticulationData; counts come from the view, names are set after. - data = OvPhysxArticulationData(mock_bindings.view, device) - 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._is_fixed_base = is_fixed_base - 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._resolve_and_install_ordering_maps() - articulation._create_buffers() - - # Wrench composers - mock_inst_wrench = WrenchComposer(articulation, supports_world_at_com=True) - mock_perm_wrench = WrenchComposer(articulation, supports_world_at_com=True) - object.__setattr__(articulation, "_instantaneous_wrench_composer", mock_inst_wrench) - object.__setattr__(articulation, "_permanent_wrench_composer", mock_perm_wrench) - # Prevent __del__ / _clear_callbacks from raising - 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) - articulation._process_actuators_cfg() - - from isaaclab_ov import tensor_types as TT - - object.__setattr__(articulation, "_can_write_effort", articulation._get_binding(TT.DOF_ACTUATION_FORCE) is not None) - object.__setattr__(articulation, "_can_write_pos_target", articulation._get_binding(TT.DOF_POSITION_TARGET) is not None) - object.__setattr__(articulation, "_can_write_vel_target", articulation._get_binding(TT.DOF_VELOCITY_TARGET) is not None) - - return articulation, mock_bindings - - -def create_newton_articulation( - num_instances: int = 2, - num_joints: int = 6, - num_bodies: int = 7, - num_fixed_tendons: int = 0, - num_spatial_tendons: int = 0, - device: str = "cuda:0", - is_fixed_base: bool = False, - joint_ordering: tuple[str, ...] | None = None, - body_ordering: tuple[str, ...] | None = None, -): - """Create a test Newton Articulation instance with mocked dependencies.""" - import isaaclab_newton.assets.articulation.articulation_data as newton_data_module - - joint_names = [f"joint_{i}" for i in range(num_joints)] - body_names = [f"body_{i}" for i in range(num_bodies)] - fixed_tendon_names = [f"fixed_tendon_{i}" for i in range(num_fixed_tendons)] - - # Create Newton mock view - mock_view = NewtonMockArticulationView( - num_instances=num_instances, - num_bodies=num_bodies, - num_joints=num_joints, - device=device, - is_fixed_base=is_fixed_base, - joint_names=joint_names, - body_names=body_names, - tendon_names=fixed_tendon_names, - ) - mock_view.set_random_mock_data() - mock_view._noop_setters = True - - # Mock NewtonManager (aliased as SimulationManager in Newton modules) - mock_model = MagicMock() - mock_model.world_count = num_instances - mock_model.gravity = wp.array( - np.tile(np.array([[0.0, 0.0, -9.81]], dtype=np.float32), (num_instances + 1, 1)), - dtype=wp.vec3f, - device=device, - ) - # Sizes consumed by NewtonArticulationData's lazy task-space buffers. - # Model-wide counts equal the per-articulation counts because the mock contains only - # this homogeneous articulation batch. - mock_model.articulation_count = num_instances - mock_model.max_joints_per_articulation = num_bodies - total_dofs = num_joints + (0 if is_fixed_base else 6) - mock_model.max_dofs_per_articulation = total_dofs - mock_model.joint_dof_count = num_instances * total_dofs - mock_model.body_count = num_instances * num_bodies - mock_view.model = mock_model - mock_state = MagicMock() - mock_control = MagicMock() - - mock_manager = MagicMock() - mock_manager.get_model.return_value = mock_model - mock_manager.get_state_0.return_value = mock_state - mock_manager.get_state_1.return_value = mock_state - mock_manager.get_control.return_value = mock_control - - # Patch SimulationManager in the Newton data module - original_sim_manager = newton_data_module.SimulationManager - newton_data_module.SimulationManager = mock_manager - - try: - data = NewtonArticulationData(mock_view, device) - finally: - newton_data_module.SimulationManager = original_sim_manager - mock_view._tendon_count = num_fixed_tendons - - # Create Articulation shell (bypass __init__) - articulation = object.__new__(NewtonArticulation) - - articulation.cfg = ArticulationCfg( - prim_path="/World/Robot", - soft_joint_pos_limit_factor=1.0, - actuators={}, - joint_ordering=joint_ordering, - body_ordering=body_ordering, - ) - object.__setattr__(articulation, "_sim_cfg", None) - object.__setattr__(articulation, "_fixed_tendon_target_dirty", False) - - object.__setattr__(articulation, "_root_view", mock_view) - object.__setattr__(articulation, "_device", device) - object.__setattr__(articulation, "_data", data) - object.__setattr__(articulation, "_test_simulation_manager", mock_manager) - articulation._clamped_default_count = wp.zeros(1, dtype=wp.int32, device=device) - # the solver builds this adapter; the shell has no model, so it stays absent - object.__setattr__(articulation, "_fixed_tendon_control", None) - - # Newton supports fixed tendons but not spatial tendons. - object.__setattr__(articulation, "_fixed_tendon_names", fixed_tendon_names) - object.__setattr__(articulation, "_spatial_tendon_names", []) - data.fixed_tendon_names = fixed_tendon_names - data.spatial_tendon_names = [] - - # Wrench composers - mock_inst_wrench = WrenchComposer(articulation) - mock_perm_wrench = WrenchComposer(articulation) - object.__setattr__(articulation, "_instantaneous_wrench_composer", mock_inst_wrench) - object.__setattr__(articulation, "_permanent_wrench_composer", mock_perm_wrench) - - # Prevent __del__ / _clear_callbacks from raising AttributeError - 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) - - # Newton uses wp.array for indices (not torch) - object.__setattr__(articulation, "_ALL_INDICES", wp.array(np.arange(num_instances, dtype=np.int32), device=device)) - object.__setattr__( - articulation, "_ALL_BODY_INDICES", wp.array(np.arange(num_bodies, dtype=np.int32), device=device) - ) - object.__setattr__( - articulation, "_ALL_JOINT_INDICES", wp.array(np.arange(num_joints, dtype=np.int32), device=device) - ) - - # Newton uses wp.bool masks - object.__setattr__(articulation, "_ALL_ENV_MASK", wp.ones((num_instances,), dtype=wp.bool, device=device)) - object.__setattr__(articulation, "_ALL_BODY_MASK", wp.ones((num_bodies,), dtype=wp.bool, device=device)) - object.__setattr__(articulation, "_ALL_JOINT_MASK", wp.ones((num_joints,), dtype=wp.bool, device=device)) - - articulation._resolve_and_install_ordering_maps() - articulation._ordering_configure_backend_staging() - - # Tendon arrays - object.__setattr__( - articulation, - "_ALL_FIXED_TENDON_INDICES", - wp.array(np.arange(num_fixed_tendons, dtype=np.int32), device=device), - ) - object.__setattr__( - articulation, "_ALL_FIXED_TENDON_MASK", wp.ones((num_fixed_tendons,), dtype=wp.bool, device=device) - ) - object.__setattr__( - articulation, "_ALL_SPATIAL_TENDON_INDICES", wp.array(np.array([], dtype=np.int32), device=device) - ) - object.__setattr__(articulation, "_ALL_SPATIAL_TENDON_MASK", wp.ones((0,), dtype=wp.bool, device=device)) - - articulation._process_actuators_cfg() - - return articulation, mock_view - - - -def get_articulation( - backend: str, - num_instances: int = 2, - num_joints: int = 6, - num_bodies: int = 7, - num_fixed_tendons: int = 0, - num_spatial_tendons: int = 0, - device: str = "cuda:0", - is_fixed_base: bool = False, - joint_ordering: tuple[str, ...] | None = None, - body_ordering: tuple[str, ...] | None = None, -): - if backend == "physx": - return create_physx_articulation( - num_instances, - num_joints, - num_bodies, - num_fixed_tendons, - num_spatial_tendons, - device, - is_fixed_base=is_fixed_base, - joint_ordering=joint_ordering, - body_ordering=body_ordering, - ) - elif backend == "ovphysx": - return create_ovphysx_articulation( - num_instances, - num_joints, - num_bodies, - num_fixed_tendons, - num_spatial_tendons, - device, - is_fixed_base=is_fixed_base, - joint_ordering=joint_ordering, - body_ordering=body_ordering, - ) - elif backend == "newton": - return create_newton_articulation( - num_instances, - num_joints, - num_bodies, - num_fixed_tendons, - num_spatial_tendons, - device, - is_fixed_base=is_fixed_base, - joint_ordering=joint_ordering, - body_ordering=body_ordering, - ) - else: - raise ValueError(f"Invalid backend: {backend}") diff --git a/source/isaaclab/test/assets/_iface_test_boot.py b/source/isaaclab/test/assets/_iface_test_boot.py deleted file mode 100644 index d93f02b84c02..000000000000 --- a/source/isaaclab/test/assets/_iface_test_boot.py +++ /dev/null @@ -1,30 +0,0 @@ -# Copyright (c) 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 - -"""Shared Kit and kitless bootstrap for backend interface tests.""" - -import os -import sys -from unittest.mock import MagicMock - -_kitless = "ovphysx" in os.environ.get("LD_PRELOAD", "") or ( - os.environ.get("LD_PRELOAD", "") == "" and "EXP_PATH" not in os.environ -) - -if not _kitless: - from isaaclab.test.utils import launch_test_simulation - - launch_test_simulation() -else: - # ``omni`` is a real namespace package in OvPhysX kitless runs. Install - # missing submodules in both ``sys.modules`` and the namespace attributes. - import omni as _omni - - for _mod in ("physics", "physics.tensors", "physx", "timeline", "usd"): - _stub = MagicMock() - sys.modules[f"omni.{_mod}"] = _stub - setattr(_omni, _mod.split(".", 1)[0], _stub) - for _mod in ("isaacsim.core", "isaacsim.core.simulation_manager"): - sys.modules.setdefault(_mod, MagicMock()) diff --git a/source/isaaclab/test/assets/_rigid_object_collection_iface_test_utils.py b/source/isaaclab/test/assets/_rigid_object_collection_iface_test_utils.py deleted file mode 100644 index e7593e8c6b8e..000000000000 --- a/source/isaaclab/test/assets/_rigid_object_collection_iface_test_utils.py +++ /dev/null @@ -1,296 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -# ignore private usage of variables warning -# pyright: reportPrivateUsage=none - -"""Shared mocked rigid-object-collection backend factories for interface tests.""" - -from unittest.mock import MagicMock - -import _iface_test_boot # noqa: F401 (starts the runtime) - -import numpy as np -import warp as wp - -from isaaclab.assets.rigid_object.rigid_object_cfg import RigidObjectCfg -from isaaclab.assets.rigid_object_collection.rigid_object_collection_cfg import RigidObjectCollectionCfg -from isaaclab.utils.wrench_composer import WrenchComposer - -BACKENDS: list[str] = [] - -try: - from isaaclab_physx.assets.rigid_object_collection.rigid_object_collection import ( - RigidObjectCollection as PhysXRigidObjectCollection, - ) - from isaaclab_physx.assets.rigid_object_collection.rigid_object_collection_data import ( - RigidObjectCollectionData as PhysXRigidObjectCollectionData, - ) - from isaaclab_physx.physics import PhysxManager as SimulationManager - from isaaclab_physx.physics.physx_manager import PhysxSceneDataBackend - from isaaclab_physx.test.fixtures.views import MockRigidBodyViewWarp as PhysXMockRigidBodyViewWarp -except ImportError: - pass -else: - # PhysX data classes need gravity even though interface tests do not create a physics scene. - _mock_physics_sim_view = MagicMock() - _mock_physics_sim_view.get_gravity.return_value = (0.0, 0.0, -9.81) - SimulationManager.get_physics_sim_view = MagicMock(return_value=_mock_physics_sim_view) - SimulationManager._scene_data_backend = PhysxSceneDataBackend() - - BACKENDS.append("physx") - -try: - from isaaclab_newton.assets.rigid_object_collection.rigid_object_collection import ( - RigidObjectCollection as NewtonRigidObjectCollection, - ) - from isaaclab_newton.assets.rigid_object_collection.rigid_object_collection_data import ( - RigidObjectCollectionData as NewtonRigidObjectCollectionData, - ) - from isaaclab_newton.test.fixtures.views import MockNewtonCollectionView as NewtonMockCollectionView -except ImportError: - pass -else: - BACKENDS.append("newton") - -try: - import ovphysx # noqa: F401 - - from isaaclab_ov.assets.rigid_object_collection.rigid_object_collection import ( - RigidObjectCollection as OvPhysxRigidObjectCollection, - ) - from isaaclab_ov.assets.rigid_object_collection.rigid_object_collection_data import ( - RigidObjectCollectionData as OvPhysxRigidObjectCollectionData, - ) - from isaaclab_ov.test.fixtures.views import MockOvPhysxBindingSet - from isaaclab_ov.physics.ovphysx_manager import OvPhysxManager, OvPhysxSceneDataBackend -except ImportError: - pass -else: - if hasattr(OvPhysxRigidObjectCollection, "_create_buffers"): - # Writers bump the scene-data transform version that ``initialize()`` would normally create. - OvPhysxManager._scene_data_backend = OvPhysxSceneDataBackend() - BACKENDS.append("ovphysx") - - -def create_physx_rigid_object_collection( - num_instances: int = 2, - num_bodies: int = 3, - device: str = "cuda:0", -): - """Create a test RigidObjectCollection instance with mocked dependencies.""" - collection = object.__new__(PhysXRigidObjectCollection) - - rigid_objects = {f"object_{i}": RigidObjectCfg(prim_path=f"/World/Object_{i}") for i in range(num_bodies)} - collection.cfg = RigidObjectCollectionCfg(rigid_objects=rigid_objects) - - # View count = num_instances * num_bodies (body-major view order) - mock_view = PhysXMockRigidBodyViewWarp( - count=num_instances * num_bodies, - device=device, - ) - mock_view.set_random_mock_data() - mock_view._noop_setters = True - - object.__setattr__(collection, "_root_view", mock_view) - object.__setattr__(collection, "_device", device) - object.__setattr__(collection, "_num_bodies", num_bodies) - object.__setattr__(collection, "_num_instances", num_instances) - object.__setattr__(collection, "_body_names_list", [f"object_{i}" for i in range(num_bodies)]) - - # Create RigidObjectCollectionData instance - data = PhysXRigidObjectCollectionData(mock_view, num_bodies, device) - object.__setattr__(collection, "_data", data) - data.body_names = [f"object_{i}" for i in range(num_bodies)] - - # Create wrench composers - mock_inst_wrench = WrenchComposer(collection, supports_world_at_com=True) - mock_perm_wrench = WrenchComposer(collection, supports_world_at_com=True) - object.__setattr__(collection, "_instantaneous_wrench_composer", mock_inst_wrench) - object.__setattr__(collection, "_permanent_wrench_composer", mock_perm_wrench) - - # Prevent __del__ / _clear_callbacks from raising AttributeError - object.__setattr__(collection, "_initialize_handle", None) - object.__setattr__(collection, "_invalidate_initialize_handle", None) - object.__setattr__(collection, "_prim_deletion_handle", None) - object.__setattr__(collection, "_debug_vis_handle", None) - - # Set up index arrays - object.__setattr__( - collection, "_ALL_ENV_INDICES", wp.array(np.arange(num_instances, dtype=np.int32), device=device) - ) - object.__setattr__(collection, "_ALL_BODY_INDICES", wp.array(np.arange(num_bodies, dtype=np.int32), device=device)) - num_view_ids = num_instances * num_bodies - all_view_ids = wp.array(np.arange(num_view_ids, dtype=np.int32), device=device) - object.__setattr__(collection, "_ALL_VIEW_INDICES", all_view_ids) - object.__setattr__(collection, "_sim_view_ids", wp.empty(num_view_ids, dtype=wp.int32, device=device)) - object.__setattr__(collection, "_sim_view_ids_views", {}) - cpu_all_view_ids = wp.empty(num_view_ids, dtype=wp.int32, device="cpu", pinned=True) - wp.copy(cpu_all_view_ids, all_view_ids) - object.__setattr__(collection, "_cpu_all_view_ids", cpu_all_view_ids) - object.__setattr__(collection, "_cpu_view_ids", wp.empty(num_view_ids, dtype=wp.int32, device="cpu", pinned=True)) - object.__setattr__(collection, "_cpu_view_ids_views", {}) - - return collection, mock_view - - -def create_newton_rigid_object_collection( - num_instances: int = 2, - num_bodies: int = 3, - device: str = "cuda:0", -): - """Create a test Newton RigidObjectCollection instance with mocked dependencies.""" - import isaaclab_newton.assets.rigid_object_collection.rigid_object_collection as newton_coll_module - import isaaclab_newton.assets.rigid_object_collection.rigid_object_collection_data as newton_data_module - - body_names = [f"object_{i}" for i in range(num_bodies)] - - # Create collection-specific mock view with (N, B) root shapes - mock_view = NewtonMockCollectionView( - num_envs=num_instances, - num_bodies=num_bodies, - device=device, - body_names=body_names, - ) - mock_view.set_random_mock_data() - mock_view._noop_setters = True - - # Mock NewtonManager (aliased as SimulationManager in Newton modules) - mock_model = MagicMock() - mock_model.world_count = num_instances - mock_model.gravity = wp.array( - np.tile(np.array([[0.0, 0.0, -9.81]], dtype=np.float32), (num_instances + 1, 1)), - dtype=wp.vec3f, - device=device, - ) - mock_state = MagicMock() - mock_control = MagicMock() - - mock_manager = MagicMock() - mock_manager.get_model.return_value = mock_model - mock_manager.get_state_0.return_value = mock_state - mock_manager.get_state_1.return_value = mock_state - mock_manager.get_control.return_value = mock_control - - # Patch SimulationManager in both data and collection modules - original_data_manager = newton_data_module.SimulationManager - original_coll_manager = newton_coll_module.SimulationManager - newton_data_module.SimulationManager = mock_manager - newton_coll_module.SimulationManager = mock_manager - - try: - data = NewtonRigidObjectCollectionData(mock_view, num_bodies, device) - finally: - newton_data_module.SimulationManager = original_data_manager - newton_coll_module.SimulationManager = original_coll_manager - - # Create collection shell (bypass __init__) - collection = object.__new__(NewtonRigidObjectCollection) - - rigid_objects = {f"object_{i}": RigidObjectCfg(prim_path=f"/World/Object_{i}") for i in range(num_bodies)} - collection.cfg = RigidObjectCollectionCfg(rigid_objects=rigid_objects) - - object.__setattr__(collection, "_root_view", mock_view) - object.__setattr__(collection, "_device", device) - object.__setattr__(collection, "_num_bodies", num_bodies) - object.__setattr__(collection, "_num_instances", num_instances) - object.__setattr__(collection, "_body_names_list", body_names) - object.__setattr__(collection, "_data", data) - data.body_names = body_names - - # Wrench composers (Newton-specific) - mock_inst_wrench = WrenchComposer(collection) - mock_perm_wrench = WrenchComposer(collection) - object.__setattr__(collection, "_instantaneous_wrench_composer", mock_inst_wrench) - object.__setattr__(collection, "_permanent_wrench_composer", mock_perm_wrench) - - # Prevent __del__ / _clear_callbacks from raising AttributeError - object.__setattr__(collection, "_initialize_handle", None) - object.__setattr__(collection, "_invalidate_initialize_handle", None) - object.__setattr__(collection, "_prim_deletion_handle", None) - object.__setattr__(collection, "_debug_vis_handle", None) - - # Index arrays (warp) - object.__setattr__( - collection, "_ALL_ENV_INDICES", wp.array(np.arange(num_instances, dtype=np.int32), device=device) - ) - object.__setattr__(collection, "_ALL_BODY_INDICES", wp.array(np.arange(num_bodies, dtype=np.int32), device=device)) - object.__setattr__(collection, "_ALL_ENV_MASK", wp.ones((num_instances,), dtype=wp.bool, device=device)) - object.__setattr__(collection, "_ALL_BODY_MASK", wp.ones((num_bodies,), dtype=wp.bool, device=device)) - - return collection, mock_view - - -def create_ovphysx_rigid_object_collection( - num_instances: int = 2, - num_bodies: int = 3, - device: str = "cuda:0", -): - """Create a test OVPhysX RigidObjectCollection instance with mocked tensor bindings.""" - body_names = [f"object_{i}" for i in range(num_bodies)] - - collection = object.__new__(OvPhysxRigidObjectCollection) - - rigid_objects = {f"object_{i}": RigidObjectCfg(prim_path=f"/World/Object_{i}") for i in range(num_bodies)} - collection.cfg = RigidObjectCollectionCfg(rigid_objects=rigid_objects) - - # Use articulation-mode bindings with num_joints=0 to get (N, B, ...) shaped tensors. - mock_bindings = MockOvPhysxBindingSet( - num_instances=num_instances, - num_joints=0, - num_bodies=num_bodies, - body_names=body_names, - asset_kind="articulation", - ) - mock_bindings.set_random_data() - - object.__setattr__(collection, "_device", device) - object.__setattr__(collection, "_ovphysx", MagicMock()) - object.__setattr__(collection, "_root_view", mock_bindings.view) - object.__setattr__(collection, "_bindings", mock_bindings.bindings) - object.__setattr__(collection, "_num_instances", num_instances) - object.__setattr__(collection, "_num_bodies", num_bodies) - object.__setattr__(collection, "_body_names_list", body_names) - - # Create RigidObjectCollectionData - data = OvPhysxRigidObjectCollectionData(mock_bindings.view, num_bodies, device) - data.num_instances = num_instances - data.num_bodies = num_bodies - data._is_primed = True - object.__setattr__(collection, "_data", data) - - # Allocate the buffers that RigidObjectCollection normally allocates in _initialize_impl. - collection._create_buffers() - - # Use production wrench composers for interface coverage. - mock_inst_wrench = WrenchComposer(collection, supports_world_at_com=True) - mock_perm_wrench = WrenchComposer(collection, supports_world_at_com=True) - object.__setattr__(collection, "_instantaneous_wrench_composer", mock_inst_wrench) - object.__setattr__(collection, "_permanent_wrench_composer", mock_perm_wrench) - - # Prevent __del__ / _clear_callbacks from raising - object.__setattr__(collection, "_initialize_handle", None) - object.__setattr__(collection, "_invalidate_initialize_handle", None) - object.__setattr__(collection, "_prim_deletion_handle", None) - object.__setattr__(collection, "_debug_vis_handle", None) - - return collection, mock_bindings - - - -def get_rigid_object_collection( - backend: str, - num_instances: int = 2, - num_bodies: int = 3, - device: str = "cuda:0", -): - if backend == "physx": - return create_physx_rigid_object_collection(num_instances, num_bodies, device) - elif backend == "ovphysx": - return create_ovphysx_rigid_object_collection(num_instances, num_bodies, device) - elif backend == "newton": - return create_newton_rigid_object_collection(num_instances, num_bodies, device) - else: - raise ValueError(f"Invalid backend: {backend}") diff --git a/source/isaaclab/test/assets/_rigid_object_iface_test_utils.py b/source/isaaclab/test/assets/_rigid_object_iface_test_utils.py deleted file mode 100644 index b58eae0b64c6..000000000000 --- a/source/isaaclab/test/assets/_rigid_object_iface_test_utils.py +++ /dev/null @@ -1,278 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -# ignore private usage of variables warning -# pyright: reportPrivateUsage=none - -"""Shared mocked rigid-object backend factories for interface tests.""" - -from unittest.mock import MagicMock - -import _iface_test_boot # noqa: F401 (starts the runtime) - -import numpy as np -import warp as wp - -from isaaclab.assets.rigid_object.rigid_object_cfg import RigidObjectCfg -from isaaclab.utils.wrench_composer import WrenchComposer - -BACKENDS: list[str] = [] - -try: - from isaaclab_physx.assets.rigid_object.rigid_object import RigidObject as PhysXRigidObject - from isaaclab_physx.assets.rigid_object.rigid_object_data import RigidObjectData as PhysXRigidObjectData - from isaaclab_physx.physics import PhysxManager as SimulationManager - from isaaclab_physx.physics.physx_manager import PhysxSceneDataBackend - from isaaclab_physx.test.fixtures.views import MockRigidBodyViewWarp as PhysXMockRigidBodyViewWarp -except ImportError: - pass -else: - # PhysX data classes need gravity even though interface tests do not create a physics scene. - _mock_physics_sim_view = MagicMock() - _mock_physics_sim_view.get_gravity.return_value = (0.0, 0.0, -9.81) - SimulationManager.get_physics_sim_view = MagicMock(return_value=_mock_physics_sim_view) - SimulationManager._scene_data_backend = PhysxSceneDataBackend() - - BACKENDS.append("physx") - -try: - from isaaclab_newton.assets.rigid_object.rigid_object import RigidObject as NewtonRigidObject - from isaaclab_newton.assets.rigid_object.rigid_object_data import RigidObjectData as NewtonRigidObjectData - from isaaclab_newton.test.fixtures.views import MockNewtonArticulationView as NewtonMockArticulationView -except ImportError: - pass -else: - BACKENDS.append("newton") - -try: - import ovphysx # noqa: F401 - - from isaaclab_ov.assets.rigid_object.rigid_object import RigidObject as OvPhysxRigidObject - from isaaclab_ov.assets.rigid_object.rigid_object_data import RigidObjectData as OvPhysxRigidObjectData - from isaaclab_ov.test.fixtures.views import MockOvPhysxBindingSet - from isaaclab_ov.physics.ovphysx_manager import OvPhysxManager, OvPhysxSceneDataBackend -except ImportError: - pass -else: - # Writers bump the scene-data transform version that ``initialize()`` would normally create. - OvPhysxManager._scene_data_backend = OvPhysxSceneDataBackend() - BACKENDS.append("ovphysx") - - -def create_physx_rigid_object( - num_instances: int = 2, - device: str = "cuda:0", -): - """Create a test RigidObject instance with mocked dependencies.""" - body_names = ["body_0"] - - rigid_object = object.__new__(PhysXRigidObject) - - rigid_object.cfg = RigidObjectCfg(prim_path="/World/Object") - - # Create PhysX mock view - mock_view = PhysXMockRigidBodyViewWarp( - count=num_instances, - device=device, - ) - mock_view.set_random_mock_data() - mock_view._noop_setters = True - - object.__setattr__(rigid_object, "_root_view", mock_view) - object.__setattr__(rigid_object, "_device", device) - - # Create RigidObjectData instance (SimulationManager already mocked at module level) - data = PhysXRigidObjectData(mock_view, device) - object.__setattr__(rigid_object, "_data", data) - - # Set body names on data - data.body_names = body_names - - # Create wrench composers - mock_inst_wrench = WrenchComposer(rigid_object, supports_world_at_com=True) - mock_perm_wrench = WrenchComposer(rigid_object, supports_world_at_com=True) - object.__setattr__(rigid_object, "_instantaneous_wrench_composer", mock_inst_wrench) - object.__setattr__(rigid_object, "_permanent_wrench_composer", mock_perm_wrench) - - # Prevent __del__ / _clear_callbacks from raising AttributeError - object.__setattr__(rigid_object, "_initialize_handle", None) - object.__setattr__(rigid_object, "_invalidate_initialize_handle", None) - object.__setattr__(rigid_object, "_prim_deletion_handle", None) - object.__setattr__(rigid_object, "_debug_vis_handle", None) - - # Set up index arrays (warp arrays for rigid object) - object.__setattr__(rigid_object, "_ALL_INDICES", wp.array(np.arange(num_instances, dtype=np.int32), device=device)) - object.__setattr__(rigid_object, "_ALL_BODY_INDICES", wp.array(np.array([0], dtype=np.int32), device=device)) - - # Cached .view(wp.float32) wrappers - object.__setattr__(rigid_object, "_root_link_pose_w_f32", None) - object.__setattr__(rigid_object, "_root_com_vel_w_f32", None) - - # Pre-allocated pinned CPU buffers for PhysX TensorAPI writes - N, B = num_instances, 1 # rigid object has 1 body - object.__setattr__(rigid_object, "_sim_env_ids", wp.empty(N, dtype=wp.int32, device=device)) - object.__setattr__(rigid_object, "_sim_env_ids_views", {}) - cpu_env_ids = wp.array(np.arange(N, dtype=np.int32), device="cpu") - object.__setattr__(rigid_object, "_cpu_env_ids_all", cpu_env_ids) - object.__setattr__(rigid_object, "_cpu_env_ids", wp.empty(N, dtype=wp.int32, device="cpu", pinned=True)) - object.__setattr__(rigid_object, "_cpu_env_ids_views", {}) - object.__setattr__(rigid_object, "_cpu_body_mass", wp.zeros((N, B), dtype=wp.float32, device="cpu")) - object.__setattr__(rigid_object, "_cpu_body_coms", wp.zeros((N, B, 7), dtype=wp.float32, device="cpu")) - object.__setattr__(rigid_object, "_cpu_body_inertia", wp.zeros((N, B, 9), dtype=wp.float32, device="cpu")) - - return rigid_object, mock_view - - -def create_newton_rigid_object( - num_instances: int = 2, - device: str = "cuda:0", -): - """Create a test Newton RigidObject instance with mocked dependencies.""" - import isaaclab_newton.assets.rigid_object.rigid_object_data as newton_data_module - - body_names = ["body_0"] - - # Create Newton mock view (uses ArticulationView with num_bodies=1 for rigid objects) - mock_view = NewtonMockArticulationView( - num_instances=num_instances, - num_bodies=1, - num_joints=0, - device=device, - is_fixed_base=False, - joint_names=[], - body_names=body_names, - ) - mock_view.set_random_mock_data() - mock_view._noop_setters = True - - # Mock NewtonManager (aliased as SimulationManager in Newton modules) - mock_model = MagicMock() - mock_model.world_count = num_instances - mock_model.gravity = wp.array( - np.tile(np.array([[0.0, 0.0, -9.81]], dtype=np.float32), (num_instances + 1, 1)), - dtype=wp.vec3f, - device=device, - ) - mock_state = MagicMock() - mock_control = MagicMock() - - mock_manager = MagicMock() - mock_manager.get_model.return_value = mock_model - mock_manager.get_state_0.return_value = mock_state - mock_manager.get_state_1.return_value = mock_state - mock_manager.get_control.return_value = mock_control - - # Patch SimulationManager in the Newton data module - original_sim_manager = newton_data_module.SimulationManager - newton_data_module.SimulationManager = mock_manager - - try: - data = NewtonRigidObjectData(mock_view, device) - finally: - newton_data_module.SimulationManager = original_sim_manager - - # Create RigidObject shell (bypass __init__) - rigid_object = object.__new__(NewtonRigidObject) - - rigid_object.cfg = RigidObjectCfg(prim_path="/World/Object") - - object.__setattr__(rigid_object, "_root_view", mock_view) - object.__setattr__(rigid_object, "_device", device) - object.__setattr__(rigid_object, "_data", data) - - # Wrench composers - mock_inst_wrench = WrenchComposer(rigid_object) - mock_perm_wrench = WrenchComposer(rigid_object) - object.__setattr__(rigid_object, "_instantaneous_wrench_composer", mock_inst_wrench) - object.__setattr__(rigid_object, "_permanent_wrench_composer", mock_perm_wrench) - - # Prevent __del__ / _clear_callbacks from raising AttributeError - object.__setattr__(rigid_object, "_initialize_handle", None) - object.__setattr__(rigid_object, "_invalidate_initialize_handle", None) - object.__setattr__(rigid_object, "_prim_deletion_handle", None) - object.__setattr__(rigid_object, "_debug_vis_handle", None) - - # Newton uses wp.array for indices - object.__setattr__(rigid_object, "_ALL_INDICES", wp.array(np.arange(num_instances, dtype=np.int32), device=device)) - object.__setattr__(rigid_object, "_ALL_BODY_INDICES", wp.array(np.array([0], dtype=np.int32), device=device)) - - # Newton uses wp.bool masks - object.__setattr__(rigid_object, "_ALL_ENV_MASK", wp.ones((num_instances,), dtype=wp.bool, device=device)) - object.__setattr__(rigid_object, "_ALL_BODY_MASK", wp.ones((1,), dtype=wp.bool, device=device)) - - return rigid_object, mock_view - - -def create_ovphysx_rigid_object( - num_instances: int = 2, - device: str = "cuda:0", -): - """Create a test OvPhysX RigidObject instance with mocked tensor bindings.""" - body_names = ["base_link"] - - obj = object.__new__(OvPhysxRigidObject) - - obj.cfg = RigidObjectCfg(prim_path="/World/object") - - # Create mock binding set - mock_bindings = MockOvPhysxBindingSet( - num_instances=num_instances, - num_joints=0, - num_bodies=1, - body_names=body_names, - asset_kind="rigid_object", - ) - mock_bindings.set_random_data() - - object.__setattr__(obj, "_device", device) - object.__setattr__(obj, "_ovphysx", MagicMock()) - object.__setattr__(obj, "_root_view", mock_bindings.view) - object.__setattr__(obj, "_bindings", mock_bindings.bindings) - object.__setattr__(obj, "_num_instances", num_instances) - object.__setattr__(obj, "_num_bodies", 1) - object.__setattr__(obj, "_body_names", body_names) - - # Create RigidObjectData - data = OvPhysxRigidObjectData(mock_bindings.view, device) - data.num_instances = num_instances - data.num_bodies = 1 - data._is_primed = True - object.__setattr__(obj, "_data", data) - - # Build the buffers RigidObject normally allocates in _initialize_impl - # (_ALL_INDICES, _ALL_*_MASK, pinned CPU staging buffers, wrench buf). - # _create_buffers also instantiates real WrenchComposers; those get - # replaced with mocks just below. - obj._create_buffers() - - # Use production wrench composers for interface coverage. - mock_inst_wrench = WrenchComposer(obj, supports_world_at_com=True) - mock_perm_wrench = WrenchComposer(obj, supports_world_at_com=True) - object.__setattr__(obj, "_instantaneous_wrench_composer", mock_inst_wrench) - object.__setattr__(obj, "_permanent_wrench_composer", mock_perm_wrench) - - # Prevent __del__ / _clear_callbacks from raising - object.__setattr__(obj, "_initialize_handle", None) - object.__setattr__(obj, "_invalidate_initialize_handle", None) - object.__setattr__(obj, "_prim_deletion_handle", None) - object.__setattr__(obj, "_debug_vis_handle", None) - - return obj, mock_bindings - - - -def get_rigid_object( - backend: str, - num_instances: int = 2, - device: str = "cuda:0", -): - if backend == "physx": - return create_physx_rigid_object(num_instances, device) - elif backend == "ovphysx": - return create_ovphysx_rigid_object(num_instances, device) - elif backend == "newton": - return create_newton_rigid_object(num_instances, device) - else: - raise ValueError(f"Invalid backend: {backend}") diff --git a/source/isaaclab/test/assets/contract/__init__.py b/source/isaaclab/test/assets/contract/__init__.py new file mode 100644 index 000000000000..0257e2759638 --- /dev/null +++ b/source/isaaclab/test/assets/contract/__init__.py @@ -0,0 +1,6 @@ +# 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 + +"""Shared asset contract test infrastructure.""" diff --git a/source/isaaclab/test/assets/contract/articulation_factory.py b/source/isaaclab/test/assets/contract/articulation_factory.py new file mode 100644 index 000000000000..0a5ed0b754b1 --- /dev/null +++ b/source/isaaclab/test/assets/contract/articulation_factory.py @@ -0,0 +1,269 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# ignore private usage of variables warning +# pyright: reportPrivateUsage=none + +"""Mocked articulation backends for the contract tests. + +Each factory bypasses ``__init__`` and sets the attributes that ``_initialize_impl`` would create, so the real +articulation and data classes run against mocked views. They return the articulation and the raw backend view or +binding set. +""" + +from unittest.mock import MagicMock + +import numpy as np +import pytest +import warp as wp + +from isaaclab.assets.articulation.articulation_cfg import ArticulationCfg + +from .backends import ( + AVAILABLE, + finish_shell, + indices, + install_physx_recording_setters, + newton_manager, + patch_ovphysx_manager, + patch_physx_manager, +) + +if "physx" in AVAILABLE: + import isaaclab_physx.assets as physx_assets + from isaaclab_physx.test.fixtures.views import MockArticulationViewWarp + +if "newton" in AVAILABLE: + import isaaclab_newton.assets as newton_assets + import isaaclab_newton.assets.articulation.articulation_data as newton_articulation_data + from isaaclab_newton.test.fixtures.views import MockNewtonArticulationView + +if "ovphysx" in AVAILABLE: + import isaaclab_ov.assets as ovphysx_assets + from isaaclab_ov import tensor_types as TT + from isaaclab_ov.test.fixtures.views import MockOvPhysxBindingSet + +_PHYSX_STORAGE = { + "set_root_transforms": "_root_transforms", + "set_root_velocities": "_root_velocities", + "set_dof_positions": "_dof_positions", + "set_dof_velocities": "_dof_velocities", + "set_dof_position_targets": "_dof_position_targets", + "set_dof_velocity_targets": "_dof_velocity_targets", + "set_dof_actuation_forces": "_dof_actuation_forces", + "set_dof_limits": "_dof_limits", + "set_dof_stiffnesses": "_dof_stiffnesses", + "set_dof_dampings": "_dof_dampings", + "set_dof_max_forces": "_dof_max_forces", + "set_dof_max_velocities": "_dof_max_velocities", + "set_dof_armatures": "_dof_armatures", + "set_dof_friction_coefficients": "_dof_friction_coefficients", + "set_dof_friction_properties": "_dof_friction_properties", + "set_masses": "_masses", + "set_coms": "_coms", + "set_inertias": "_inertias", +} + + +def _names(prefix: str, count: int) -> list[str]: + return [f"{prefix}_{i}" for i in range(count)] + + +def _shell(cls, joint_ordering, body_ordering): + """Create an uninitialized articulation of ``cls`` with a minimal configuration.""" + articulation = object.__new__(cls) + articulation.cfg = ArticulationCfg( + prim_path="/World/Robot", + soft_joint_pos_limit_factor=1.0, + actuators={}, + joint_ordering=joint_ordering, + body_ordering=body_ordering, + ) + articulation._sim_cfg = None + articulation._fixed_tendon_target_dirty = False + return articulation + + +def _set_tendon_names(articulation, fixed: list[str], spatial: list[str]) -> None: + articulation._fixed_tendon_names = articulation.data.fixed_tendon_names = fixed + articulation._spatial_tendon_names = articulation.data.spatial_tendon_names = spatial + + +def _physx_articulation(N, J, B, FT, ST, device, is_fixed_base, joint_ordering, body_ordering, monkeypatch): + patch_physx_manager(monkeypatch) + view = MockArticulationViewWarp( + count=N, num_links=B, num_dofs=J, device=device, max_fixed_tendons=FT, max_spatial_tendons=ST + ) + view.set_random_mock_data() + install_physx_recording_setters(view, _PHYSX_STORAGE) + view._shared_metatype = MagicMock( + fixed_base=is_fixed_base, dof_count=J, link_count=B, dof_names=_names("joint", J), link_names=_names("body", B) + ) + + articulation = _shell(physx_assets.Articulation, joint_ordering, body_ordering) + articulation._root_view = view + articulation._device = device + articulation._clamped_default_count = wp.zeros(1, dtype=wp.int32, device=device) + articulation._data = physx_assets.ArticulationData(view, device) + _set_tendon_names(articulation, _names("fixed_tendon", FT), _names("spatial_tendon", ST)) + finish_shell(articulation) + articulation._ALL_INDICES = articulation._ALL_INDICES_WP = indices(N, device) + articulation._ALL_BODY_INDICES = articulation._ALL_BODY_INDICES_WP = indices(B, device) + articulation._ALL_JOINT_INDICES = indices(J, device) + articulation._ALL_FIXED_TENDON_INDICES = indices(FT, device) + articulation._ALL_SPATIAL_TENDON_INDICES = indices(ST, device) + for name in ("joint_pos_target", "joint_vel_target", "joint_effort_target", "body_wrench_force"): + setattr(articulation, f"_{name}_backend", None) + articulation._body_wrench_torque_backend = None + articulation._resolve_and_install_ordering_maps() + articulation._ordering_configure_backend_staging() + # Cached ``.view(wp.float32)`` wrappers. + for name in ("root_link_pose_w", "root_com_vel_w", "root_link_vel_w", "inst_wrench_force", "inst_wrench_torque"): + setattr(articulation, f"_{name}_f32", None) + articulation._perm_wrench_force_f32 = articulation._perm_wrench_torque_f32 = None + # Pinned CPU staging buffers for PhysX TensorAPI writes. + articulation._sim_env_ids = wp.empty(N, dtype=wp.int32, device=device) + articulation._sim_env_ids_views = {} + articulation._cpu_env_ids_all = indices(N, "cpu") + articulation._cpu_env_ids = wp.empty(N, dtype=wp.int32, device="cpu", pinned=wp.is_cuda_available()) + articulation._cpu_env_ids_views = {} + for name, shape in { + "joint_stiffness": (N, J), + "joint_damping": (N, J), + "joint_pos_limits": (N, J, 2), + "joint_vel_limits": (N, J), + "joint_effort_limits": (N, J), + "joint_armature": (N, J), + "joint_friction_props": (N, J, 3), + "body_mass": (N, B), + "body_coms": (N, B, 7), + "body_inertia": (N, B, 9), + }.items(): + setattr(articulation, f"_cpu_{name}", wp.zeros(shape, dtype=wp.float32, device="cpu")) + articulation._process_actuators_cfg() + return articulation, view + + +def _ovphysx_articulation(N, J, B, FT, ST, device, is_fixed_base, joint_ordering, body_ordering, monkeypatch): + patch_ovphysx_manager(monkeypatch) + joint_names, body_names = _names("joint", J), _names("body", B) + bindings = MockOvPhysxBindingSet( + num_instances=N, + num_joints=J, + num_bodies=B, + is_fixed_base=is_fixed_base, + joint_names=joint_names, + body_names=body_names, + num_fixed_tendons=FT, + num_spatial_tendons=ST, + ) + bindings.set_random_data() + + articulation = _shell(ovphysx_assets.Articulation, joint_ordering, body_ordering) + articulation._device = device + articulation._ovphysx = MagicMock() + articulation._root_view = bindings.view + articulation._bindings = bindings.bindings + articulation._num_instances = N + articulation._num_joints = J + articulation._num_bodies = B + articulation._is_fixed_base = is_fixed_base + articulation._joint_names = joint_names + articulation._body_names = body_names + articulation._num_fixed_tendons = FT + articulation._num_spatial_tendons = ST + # Counts come from the view; names are set on the data afterwards. + articulation._data = ovphysx_assets.ArticulationData(bindings.view, device) + articulation._data.body_names = body_names + articulation._data.joint_names = joint_names + articulation._data._is_fixed_base = is_fixed_base + _set_tendon_names(articulation, _names("fixed_tendon", FT), _names("spatial_tendon", ST)) + articulation._resolve_and_install_ordering_maps() + articulation._create_buffers() + finish_shell(articulation) + articulation._process_actuators_cfg() + articulation._can_write_effort = articulation._get_binding(TT.DOF_ACTUATION_FORCE) is not None + articulation._can_write_pos_target = articulation._get_binding(TT.DOF_POSITION_TARGET) is not None + articulation._can_write_vel_target = articulation._get_binding(TT.DOF_VELOCITY_TARGET) is not None + return articulation, bindings + + +def _newton_articulation(N, J, B, FT, ST, device, is_fixed_base, joint_ordering, body_ordering, monkeypatch): + # Newton supports fixed tendons but not spatial tendons. + fixed_tendon_names = _names("fixed_tendon", FT) + view = MockNewtonArticulationView( + num_instances=N, + num_bodies=B, + num_joints=J, + num_tendons=FT, + device=device, + is_fixed_base=is_fixed_base, + joint_names=_names("joint", J), + body_names=_names("body", B), + tendon_names=fixed_tendon_names, + ) + view.set_random_mock_data() + # MuJoCo tendon properties the data binds when the articulation has fixed tendons, with distinct values. + tendon_values = np.arange(1, N * FT + 1, dtype=np.float32).reshape(N, 1, FT) + view._attributes["mujoco.tendon_stiffness"] = wp.array(tendon_values, dtype=wp.float32, device=device) + view._attributes["mujoco.tendon_damping"] = wp.array(tendon_values + 0.5, dtype=wp.float32, device=device) + tendon_range = np.stack((-tendon_values, tendon_values), axis=-1) + view._attributes["mujoco.tendon_range"] = wp.array(tendon_range, dtype=wp.vec2f, device=device) + + manager = newton_manager(N, device) + # Sizes of the lazy task-space buffers; the mock model holds only this homogeneous articulation batch. + model = view.model = manager.get_model.return_value + total_dofs = J + (0 if is_fixed_base else 6) + model.articulation_count = N + model.max_joints_per_articulation = B + model.max_dofs_per_articulation = total_dofs + model.joint_dof_count = N * total_dofs + model.body_count = N * B + monkeypatch.setattr(newton_articulation_data, "SimulationManager", manager, raising=False) + + articulation = _shell(newton_assets.Articulation, joint_ordering, body_ordering) + articulation._root_view = view + articulation._device = device + articulation._clamped_default_count = wp.zeros(1, dtype=wp.int32, device=device) + articulation._data = newton_assets.ArticulationData(view, device) + articulation._test_simulation_manager = manager + # The solver builds this adapter; the shell has no model, so it stays absent. + articulation._fixed_tendon_control = None + _set_tendon_names(articulation, fixed_tendon_names, []) + finish_shell(articulation, supports_world_at_com=False) + articulation._ALL_INDICES = indices(N, device) + articulation._ALL_BODY_INDICES = indices(B, device) + articulation._ALL_JOINT_INDICES = indices(J, device) + articulation._ALL_ENV_MASK = wp.ones((N,), dtype=wp.bool, device=device) + articulation._ALL_BODY_MASK = wp.ones((B,), dtype=wp.bool, device=device) + articulation._ALL_JOINT_MASK = wp.ones((J,), dtype=wp.bool, device=device) + articulation._resolve_and_install_ordering_maps() + articulation._ordering_configure_backend_staging() + articulation._ALL_FIXED_TENDON_INDICES = indices(FT, device) + articulation._ALL_FIXED_TENDON_MASK = wp.ones((FT,), dtype=wp.bool, device=device) + articulation._ALL_SPATIAL_TENDON_INDICES = indices(0, device) + articulation._ALL_SPATIAL_TENDON_MASK = wp.ones((0,), dtype=wp.bool, device=device) + articulation._process_actuators_cfg() + return articulation, view + + +def get_articulation( + backend: str, + num_instances: int = 2, + num_joints: int = 6, + num_bodies: int = 7, + num_fixed_tendons: int = 0, + num_spatial_tendons: int = 0, + device: str = "cuda:0", + is_fixed_base: bool = False, + joint_ordering: tuple[str, ...] | None = None, + body_ordering: tuple[str, ...] | None = None, + *, + monkeypatch: pytest.MonkeyPatch, +): + """Create a mocked articulation of one backend and return it with its backend view or bindings.""" + factory = {"physx": _physx_articulation, "ovphysx": _ovphysx_articulation, "newton": _newton_articulation}[backend] + counts = (num_instances, num_joints, num_bodies, num_fixed_tendons, num_spatial_tendons) + return factory(*counts, device, is_fixed_base, joint_ordering, body_ordering, monkeypatch) diff --git a/source/isaaclab/test/assets/contract/backends.py b/source/isaaclab/test/assets/contract/backends.py new file mode 100644 index 000000000000..bbd5bdfda8de --- /dev/null +++ b/source/isaaclab/test/assets/contract/backends.py @@ -0,0 +1,189 @@ +# 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 + +"""Backend bootstrap, availability, and manager doubles shared by the asset contract tests. + +Importing this module launches Kit when the process runs inside Isaac Sim. Otherwise it stubs the Kit-only modules +that the PhysX and OVPhysX asset classes import, so the real asset and data classes run against mocked views. +""" + +import os +import sys +from collections.abc import Callable +from importlib.machinery import ModuleSpec +from importlib.util import find_spec +from unittest.mock import MagicMock + +import numpy as np +import pytest +import warp as wp + +from isaaclab.utils.wrench_composer import WrenchComposer + +if "ovphysx" not in os.environ.get("LD_PRELOAD", "") and (os.environ.get("LD_PRELOAD") or "EXP_PATH" in os.environ): + from isaaclab.test.utils import launch_test_simulation + + launch_test_simulation() +else: + import omni # noqa: F401 # real namespace package; the stubs below become its attributes + + for _name, _is_package in ( + ("carb", False), + ("usdrt", True), + ("omni.kit", True), + ("omni.kit.app", False), + ("omni.physics", True), + ("omni.physics.tensors", False), + ("omni.physx", False), + ("omni.timeline", False), + ("omni.usd", False), + ("isaacsim", True), + ("isaacsim.core", True), + ("isaacsim.core.simulation_manager", False), + ): + if _name in sys.modules: + continue + _stub = MagicMock() + _stub.__spec__ = ModuleSpec(_name, loader=None, is_package=_is_package) + if _is_package: + _stub.__path__ = [] + sys.modules[_name] = _stub + if "." in _name: + _parent, _attribute = _name.rsplit(".", 1) + setattr(sys.modules[_parent], _attribute, _stub) + sys.modules["omni.kit.app"].get_app.return_value = None + + +def _unavailable_reason(*modules: str, cuda: bool = False) -> str | None: + """Return why a backend cannot run in this process, or None when it can.""" + for module in modules: + if find_spec(module) is None: + return f"missing module: {module}" + if cuda and not wp.is_cuda_available(): + # The mocked OVPhysX bindings allocate pinned host staging buffers even for CPU tensors. + return "requires a CUDA runtime" + return None + + +UNAVAILABLE = { + name: reason + for name, reason in { + "physx": _unavailable_reason("isaaclab_physx"), + "newton": _unavailable_reason("isaaclab_newton"), + "ovphysx": _unavailable_reason("ovphysx", "isaaclab_ov", cuda=True), + }.items() + if reason is not None +} +"""Backend name to the reason it is skipped in this process.""" + +AVAILABLE = [name for name in ("physx", "newton", "ovphysx") if name not in UNAVAILABLE] +"""Backends that run in this process.""" + + +def backends(*names: str) -> list: + """Return pytest parameters for the named backends (default: all), skipping the unavailable ones.""" + return [ + pytest.param(name, id=name, marks=pytest.mark.skipif(name in UNAVAILABLE, reason=UNAVAILABLE.get(name, ""))) + for name in names or ("physx", "newton", "ovphysx") + ] + + +def requires(name: str) -> pytest.MarkDecorator: + """Skip a backend-specific test when that backend is unavailable.""" + return pytest.mark.skipif(name in UNAVAILABLE, reason=f"{name}: {UNAVAILABLE.get(name)}") + + +def patch_physx_manager(monkeypatch: pytest.MonkeyPatch) -> None: + """Give PhysX data classes gravity and a scene-data backend without creating a physics scene.""" + from isaaclab_physx.physics import PhysxManager + from isaaclab_physx.physics.physx_manager import PhysxSceneDataBackend + + physics_sim_view = MagicMock() + physics_sim_view.get_gravity.return_value = (0.0, 0.0, -9.81) + monkeypatch.setattr(PhysxManager, "get_physics_sim_view", MagicMock(return_value=physics_sim_view), raising=False) + # Writers bump the scene-data transform version that ``initialize()`` would normally create. + monkeypatch.setattr(PhysxManager, "_scene_data_backend", PhysxSceneDataBackend(), raising=False) + + +def patch_ovphysx_manager(monkeypatch: pytest.MonkeyPatch) -> None: + """Create the OVPhysX scene-data backend whose transform version the writers bump.""" + from isaaclab_ov.physics.ovphysx_manager import OvPhysxManager, OvPhysxSceneDataBackend + + monkeypatch.setattr(OvPhysxManager, "_scene_data_backend", OvPhysxSceneDataBackend(), raising=False) + + +def install_physx_recording_setters(mock_view, storage_by_method: dict[str, str]) -> None: + """Replace PhysX view setters with TensorAPI-faithful ones that apply the selected rows to view storage. + + The PhysX TensorAPI receives full-size data and applies only the rows named by ``indices``. The fixture view + setters either drop writes or expect compact data, so the contract installs setters that require full-size data + and leave every unselected row untouched. + + Args: + mock_view: PhysX fixture view whose setters are replaced. + storage_by_method: Setter name to the view attribute that stores its values. + """ + + def make_setter(storage_name: str) -> Callable[[wp.array, wp.array | None], None]: + def setter(values: wp.array, indices: wp.array | None = None) -> None: + values_np = values.numpy() + stored = getattr(mock_view, storage_name, None) + if stored is None: + stored = wp.zeros((mock_view._count, *values_np.shape[1:]), dtype=wp.float32, device=values.device) + setattr(mock_view, storage_name, stored) + stored_np = stored.numpy() + if indices is None: + stored_np[...] = values_np.reshape(stored_np.shape) + else: + if values_np.size != stored_np.size: + raise ValueError(f"{storage_name}: expected full-size data of shape {stored_np.shape}") + rows = indices.numpy().astype(np.int64) + stored_np[rows] = values_np.reshape(stored_np.shape)[rows] + if stored.device.is_cuda: + stored.assign(wp.array(stored_np, dtype=stored.dtype, device=stored.device)) + + return setter + + mock_view._noop_setters = False + for method_name, storage_name in storage_by_method.items(): + setattr(mock_view, method_name, make_setter(storage_name)) + + +def read_backend_joint_state(backend: str, art, raw_backend) -> tuple[np.ndarray, np.ndarray]: + """Return joint position and velocity from backend-order storage.""" + if backend == "physx": + return raw_backend._dof_positions.numpy().copy(), raw_backend._dof_velocities.numpy().copy() + if backend == "ovphysx": + from isaaclab_ov import tensor_types as TT + + position = np.asarray(raw_backend.bindings[TT.DOF_POSITION]._data).copy() + velocity = np.asarray(raw_backend.bindings[TT.DOF_VELOCITY]._data).copy() + return position, velocity + return art.data._sim_bind_joint_pos.numpy().copy(), art.data._sim_bind_joint_vel.numpy().copy() + + +def finish_shell(asset, *, supports_world_at_com: bool = True) -> None: + """Attach wrench composers and clear the callback handles that ``__del__`` releases.""" + asset._instantaneous_wrench_composer = WrenchComposer(asset, supports_world_at_com=supports_world_at_com) + asset._permanent_wrench_composer = WrenchComposer(asset, supports_world_at_com=supports_world_at_com) + asset._initialize_handle = None + asset._invalidate_initialize_handle = None + asset._prim_deletion_handle = None + asset._debug_vis_handle = None + + +def indices(count: int, device: str) -> wp.array: + """Return ``[0, count)`` as an int32 Warp array.""" + return wp.array(np.arange(count, dtype=np.int32), device=device) + + +def newton_manager(num_instances: int, device: str) -> MagicMock: + """Return a NewtonManager double with a gravity-carrying model.""" + model = MagicMock(world_count=num_instances) + model.gravity = wp.array(np.tile([[0.0, 0.0, -9.81]], (num_instances + 1, 1)), dtype=wp.vec3f, device=device) + manager = MagicMock() + manager.get_model.return_value = model + manager.get_state_0.return_value = manager.get_state_1.return_value = MagicMock() + return manager diff --git a/source/isaaclab/test/assets/contract/conftest.py b/source/isaaclab/test/assets/contract/conftest.py new file mode 100644 index 000000000000..16e2a0e7acce --- /dev/null +++ b/source/isaaclab/test/assets/contract/conftest.py @@ -0,0 +1,45 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Backend and device parameters shared by the asset contracts.""" + +import pytest +import warp as wp + +from isaaclab.test.utils import DeviceScope, test_devices + +from .backends import backends + + +@pytest.fixture(params=backends()) +def backend(request: pytest.FixtureRequest) -> str: + """Run each shared contract on every backend, skipping the ones unavailable in this process.""" + return request.param + + +@pytest.fixture(params=test_devices(DeviceScope.CPU_AND_DEFAULT_CUDA)) +def device(request: pytest.FixtureRequest) -> str: + """Exercise CPU buffers and the CUDA staging paths of PhysX and OVPhysX.""" + return request.param + + +def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None: + """Deselect Newton CUDA cases: Newton's asset code has no CUDA-specific path, so the CPU case covers it.""" + deselected = [] + for item in items: + params = getattr(item, "callspec", None) and item.callspec.params + if params and params.get("backend") == "newton" and params.get("device", "cpu") != "cpu": + deselected.append(item) + if deselected: + config.hook.pytest_deselected(items=deselected) + items[:] = [item for item in items if item not in deselected] + + +@pytest.fixture(autouse=True) +def _default_warp_device(request: pytest.FixtureRequest): + """Allocate unqualified Warp arrays on the case's device, or on CPU when the case has none.""" + callspec = getattr(request.node, "callspec", None) + with wp.ScopedDevice(callspec.params.get("device", "cpu") if callspec else "cpu"): + yield diff --git a/source/isaaclab/test/assets/contract/rigid_object_factory.py b/source/isaaclab/test/assets/contract/rigid_object_factory.py new file mode 100644 index 000000000000..b5a5530327f0 --- /dev/null +++ b/source/isaaclab/test/assets/contract/rigid_object_factory.py @@ -0,0 +1,243 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# ignore private usage of variables warning +# pyright: reportPrivateUsage=none + +"""Mocked rigid-object and rigid-object-collection backends for the contract tests. + +Each factory bypasses ``__init__`` and sets the attributes that ``_initialize_impl`` would create, so the real asset +and data classes run against mocked views. They return the asset and the raw backend view or binding set. +""" + +from unittest.mock import MagicMock + +import pytest +import warp as wp + +from isaaclab.assets.rigid_object.rigid_object_cfg import RigidObjectCfg +from isaaclab.assets.rigid_object_collection.rigid_object_collection_cfg import RigidObjectCollectionCfg + +from .backends import ( + AVAILABLE, + finish_shell, + indices, + install_physx_recording_setters, + newton_manager, + patch_ovphysx_manager, + patch_physx_manager, +) + +if "physx" in AVAILABLE: + import isaaclab_physx.assets as physx_assets + from isaaclab_physx.test.fixtures.views import MockRigidBodyViewWarp + +if "newton" in AVAILABLE: + import isaaclab_newton.assets as newton_assets + import isaaclab_newton.assets.rigid_object.rigid_object_data as newton_rigid_object_data + import isaaclab_newton.assets.rigid_object_collection.rigid_object_collection as newton_collection + import isaaclab_newton.assets.rigid_object_collection.rigid_object_collection_data as newton_collection_data + from isaaclab_newton.test.fixtures.views import MockNewtonArticulationView, MockNewtonCollectionView + +if "ovphysx" in AVAILABLE: + import isaaclab_ov.assets as ovphysx_assets + from isaaclab_ov.test.fixtures.views import MockOvPhysxBindingSet + +_PHYSX_STORAGE = { + "set_transforms": "_transforms", + "set_velocities": "_velocities", + "set_masses": "_masses", + "set_coms": "_coms", + "set_inertias": "_inertias", +} + + +def _physx_rigid_object(num_instances: int, device: str, monkeypatch: pytest.MonkeyPatch): + patch_physx_manager(monkeypatch) + view = MockRigidBodyViewWarp(count=num_instances, device=device) + view.set_random_mock_data() + install_physx_recording_setters(view, _PHYSX_STORAGE) + + obj = object.__new__(physx_assets.RigidObject) + obj.cfg = RigidObjectCfg(prim_path="/World/Object") + obj._root_view = view + obj._device = device + obj._data = physx_assets.RigidObjectData(view, device) + obj._data.body_names = ["body_0"] + finish_shell(obj) + obj._ALL_INDICES = indices(num_instances, device) + obj._ALL_BODY_INDICES = indices(1, device) + obj._root_link_pose_w_f32 = None + obj._root_com_vel_w_f32 = None + # Pinned CPU staging buffers for PhysX TensorAPI writes. + pinned = wp.is_cuda_available() + obj._sim_env_ids = wp.empty(num_instances, dtype=wp.int32, device=device) + obj._sim_env_ids_views = {} + obj._cpu_env_ids_all = indices(num_instances, "cpu") + obj._cpu_env_ids = wp.empty(num_instances, dtype=wp.int32, device="cpu", pinned=pinned) + obj._cpu_env_ids_views = {} + obj._cpu_body_mass = wp.zeros((num_instances, 1), dtype=wp.float32, device="cpu") + obj._cpu_body_coms = wp.zeros((num_instances, 1, 7), dtype=wp.float32, device="cpu") + obj._cpu_body_inertia = wp.zeros((num_instances, 1, 9), dtype=wp.float32, device="cpu") + return obj, view + + +def _newton_rigid_object(num_instances: int, device: str, monkeypatch: pytest.MonkeyPatch): + view = MockNewtonArticulationView( + num_instances=num_instances, + num_bodies=1, + num_joints=0, + device=device, + is_fixed_base=False, + joint_names=[], + body_names=["body_0"], + ) + view.set_random_mock_data() + monkeypatch.setattr(newton_rigid_object_data, "SimulationManager", newton_manager(num_instances, device)) + + obj = object.__new__(newton_assets.RigidObject) + obj.cfg = RigidObjectCfg(prim_path="/World/Object") + obj._root_view = view + obj._device = device + obj._data = newton_assets.RigidObjectData(view, device) + finish_shell(obj, supports_world_at_com=False) + obj._ALL_INDICES = indices(num_instances, device) + obj._ALL_BODY_INDICES = indices(1, device) + obj._ALL_ENV_MASK = wp.ones((num_instances,), dtype=wp.bool, device=device) + obj._ALL_BODY_MASK = wp.ones((1,), dtype=wp.bool, device=device) + return obj, view + + +def _ovphysx_rigid_object(num_instances: int, device: str, monkeypatch: pytest.MonkeyPatch): + patch_ovphysx_manager(monkeypatch) + bindings = MockOvPhysxBindingSet( + num_instances=num_instances, num_joints=0, num_bodies=1, body_names=["body_0"], asset_kind="rigid_object" + ) + bindings.set_random_data() + + obj = object.__new__(ovphysx_assets.RigidObject) + obj.cfg = RigidObjectCfg(prim_path="/World/Object") + obj._device = device + obj._ovphysx = MagicMock() + obj._root_view = bindings.view + obj._bindings = bindings.bindings + obj._num_instances = num_instances + obj._num_bodies = 1 + obj._body_names = ["body_0"] + obj._data = ovphysx_assets.RigidObjectData(bindings.view, device) + obj._data.num_instances = num_instances + obj._data.num_bodies = 1 + obj._data._is_primed = True + obj._create_buffers() + finish_shell(obj) + return obj, bindings + + +def _collection_cfg(body_names: list[str]) -> RigidObjectCollectionCfg: + return RigidObjectCollectionCfg( + rigid_objects={name: RigidObjectCfg(prim_path=f"/World/{name}") for name in body_names} + ) + + +def _physx_collection(num_instances: int, num_bodies: int, device: str, monkeypatch: pytest.MonkeyPatch): + patch_physx_manager(monkeypatch) + body_names = [f"object_{i}" for i in range(num_bodies)] + # PhysX collection views are body-major: one view entry per (body, environment). + num_view_ids = num_instances * num_bodies + view = MockRigidBodyViewWarp(count=num_view_ids, device=device) + view.set_random_mock_data() + install_physx_recording_setters(view, _PHYSX_STORAGE) + + collection = object.__new__(physx_assets.RigidObjectCollection) + collection.cfg = _collection_cfg(body_names) + collection._root_view = view + collection._device = device + collection._num_bodies = num_bodies + collection._num_instances = num_instances + collection._body_names_list = body_names + collection._data = physx_assets.RigidObjectCollectionData(view, num_bodies, device) + collection._data.body_names = body_names + finish_shell(collection) + collection._ALL_ENV_INDICES = indices(num_instances, device) + collection._ALL_BODY_INDICES = indices(num_bodies, device) + collection._ALL_VIEW_INDICES = indices(num_view_ids, device) + pinned = wp.is_cuda_available() + collection._sim_view_ids = wp.empty(num_view_ids, dtype=wp.int32, device=device) + collection._sim_view_ids_views = {} + collection._cpu_all_view_ids = wp.empty(num_view_ids, dtype=wp.int32, device="cpu", pinned=pinned) + wp.copy(collection._cpu_all_view_ids, collection._ALL_VIEW_INDICES) + collection._cpu_view_ids = wp.empty(num_view_ids, dtype=wp.int32, device="cpu", pinned=pinned) + collection._cpu_view_ids_views = {} + return collection, view + + +def _newton_collection(num_instances: int, num_bodies: int, device: str, monkeypatch: pytest.MonkeyPatch): + body_names = [f"object_{i}" for i in range(num_bodies)] + view = MockNewtonCollectionView(num_envs=num_instances, num_bodies=num_bodies, device=device, body_names=body_names) + view.set_random_mock_data() + manager = newton_manager(num_instances, device) + monkeypatch.setattr(newton_collection_data, "SimulationManager", manager) + monkeypatch.setattr(newton_collection, "SimulationManager", manager) + + collection = object.__new__(newton_assets.RigidObjectCollection) + collection.cfg = _collection_cfg(body_names) + collection._root_view = view + collection._device = device + collection._num_bodies = num_bodies + collection._num_instances = num_instances + collection._body_names_list = body_names + collection._data = newton_assets.RigidObjectCollectionData(view, num_bodies, device) + collection._data.body_names = body_names + finish_shell(collection, supports_world_at_com=False) + collection._ALL_ENV_INDICES = indices(num_instances, device) + collection._ALL_BODY_INDICES = indices(num_bodies, device) + collection._ALL_ENV_MASK = wp.ones((num_instances,), dtype=wp.bool, device=device) + collection._ALL_BODY_MASK = wp.ones((num_bodies,), dtype=wp.bool, device=device) + return collection, view + + +def _ovphysx_collection(num_instances: int, num_bodies: int, device: str, monkeypatch: pytest.MonkeyPatch): + patch_ovphysx_manager(monkeypatch) + body_names = [f"object_{i}" for i in range(num_bodies)] + # Articulation-mode bindings without joints give the (N, B, ...) tensors of a collection. + bindings = MockOvPhysxBindingSet( + num_instances=num_instances, + num_joints=0, + num_bodies=num_bodies, + body_names=body_names, + asset_kind="articulation", + ) + bindings.set_random_data() + + collection = object.__new__(ovphysx_assets.RigidObjectCollection) + collection.cfg = _collection_cfg(body_names) + collection._device = device + collection._ovphysx = MagicMock() + collection._root_view = bindings.view + collection._bindings = bindings.bindings + collection._num_instances = num_instances + collection._num_bodies = num_bodies + collection._body_names_list = body_names + collection._data = ovphysx_assets.RigidObjectCollectionData(bindings.view, num_bodies, device) + collection._data.num_instances = num_instances + collection._data.num_bodies = num_bodies + collection._data._is_primed = True + collection._create_buffers() + finish_shell(collection) + return collection, bindings + + +def get_rigid_object(backend: str, num_instances: int = 2, device: str = "cpu", *, monkeypatch: pytest.MonkeyPatch): + """Create a mocked single-body rigid object of the given backend.""" + factory = {"physx": _physx_rigid_object, "newton": _newton_rigid_object, "ovphysx": _ovphysx_rigid_object} + return factory[backend](num_instances, device, monkeypatch) + + +def get_rigid_object_collection( + backend: str, num_instances: int = 2, num_bodies: int = 3, device: str = "cpu", *, monkeypatch: pytest.MonkeyPatch +): + """Create a mocked rigid-object collection of the given backend.""" + factory = {"physx": _physx_collection, "newton": _newton_collection, "ovphysx": _ovphysx_collection} + return factory[backend](num_instances, num_bodies, device, monkeypatch) diff --git a/source/isaaclab/test/assets/contract/test_articulation_contract.py b/source/isaaclab/test/assets/contract/test_articulation_contract.py new file mode 100644 index 000000000000..83f141148791 --- /dev/null +++ b/source/isaaclab/test/assets/contract/test_articulation_contract.py @@ -0,0 +1,1779 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# ignore private usage of variables warning +# pyright: reportPrivateUsage=none + +"""Shared articulation contracts across the production backends. + +Checks that every articulation backend provides the data, writer, and joint/body ordering behavior the base +articulation class advertises. The backends run on mocked views, so these cases need neither Isaac Sim nor a GPU +simulation. Pure bookkeeping and ordering run on CPU only; getters and writers also run on CUDA for the backends that +stage through CPU-pinned buffers there. +""" + +import math +import warnings +from operator import attrgetter +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +import torch +import warp as wp + +from isaaclab.assets.articulation.base_articulation_data import BaseArticulationData +from isaaclab.utils import replace +from isaaclab.utils.warp import ProxyArray + +from .articulation_factory import get_articulation +from .backends import backends, read_backend_joint_state, requires + +pytestmark = pytest.mark.integration + +# Distinct instance, joint, body, and tendon counts so any swapped axis shows up. +N, J, B, FT, ST = 2, 6, 7, 3, 2 +_AXIS_COUNTS = {"joint": J, "body": B, "fixed_tendon": FT, "spatial_tendon": ST} + + +def _articulation(backend: str, device: str, monkeypatch: pytest.MonkeyPatch, **kwargs): + """Create the shared mocked articulation; keyword arguments override the default counts.""" + counts = dict(num_instances=N, num_joints=J, num_bodies=B, num_fixed_tendons=FT, num_spatial_tendons=ST) + return get_articulation(backend, device=device, monkeypatch=monkeypatch, **(counts | kwargs)) + + +def _mask(total: int, selected, device: str) -> wp.array: + """Return a boolean Warp mask with the selected indices set.""" + mask = np.zeros(total, dtype=bool) + mask[list(selected)] = True + return wp.array(mask, dtype=wp.bool, device=device) + + +def _payload(shape: tuple[int, ...], trailing: int, dtype, device: str, offset: float = 0.0) -> torch.Tensor: + """Return valid writer data whose entries differ per element. + + Transforms carry a fixed 90-degree rotation about Z, and ``vec2f`` limits a ``[-v, v]`` pair. + """ + values = torch.arange(1, math.prod(shape) + 1, dtype=torch.float32, device=device).reshape(shape) + offset + 0.25 + if dtype == wp.vec2f: + return torch.stack((-values, values), dim=-1) + if trailing: + values = values.unsqueeze(-1) + torch.arange(trailing, dtype=torch.float32, device=device) / 10.0 + if dtype == wp.transformf: + values[..., 3:] = torch.tensor([0.0, 0.0, 0.5**0.5, 0.5**0.5], device=device) + return values + + +def _assert_reads_back(value: ProxyArray, expected: torch.Tensor, name: str) -> None: + """Assert a data getter returns the values a writer just wrote.""" + torch.testing.assert_close(value.torch, expected, atol=1e-5, rtol=1e-5, msg=lambda msg: f"{name}: {msg}") + + +def _prime_timestamped_properties(data, buffer_names: list[str]) -> list: + """Read lazy public properties and stamp their backing buffers as current.""" + buffers = [] + for name in buffer_names: + getattr(data, name.removeprefix("_").removesuffix("_buf")) + buffer = getattr(data, name) + assert buffer is not None, name + buffer.timestamp = data._sim_timestamp + buffers.append((name, buffer)) + return buffers + + +def _assert_buffers_stale(data, buffers) -> None: + for name, buffer in buffers: + assert buffer.timestamp < data._sim_timestamp, name + + +def _ignore_newton_model_changes(backend: str, monkeypatch: pytest.MonkeyPatch) -> None: + """The mocked Newton asset has no solver to receive model-change notifications.""" + if backend == "newton": + from isaaclab_newton.physics import NewtonManager + + monkeypatch.setattr(NewtonManager, "add_model_change", MagicMock()) + + +# --------------------------------------------------------------------------- +# Counts, names, finders, and selectors +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("is_fixed_base", [False, True], ids=["floating", "fixed"]) +def test_counts_and_names(monkeypatch, backend, is_fixed_base): + art, _ = _articulation(backend, "cpu", monkeypatch, is_fixed_base=is_fixed_base) + # Newton has no spatial tendons. + num_spatial = 0 if backend == "newton" else ST + + assert isinstance(art.data, BaseArticulationData) + assert art.num_instances == N + assert art.is_fixed_base is is_fixed_base + for names, count, expected in ( + (art.joint_names, art.num_joints, J), + (art.body_names, art.num_bodies, B), + (art.fixed_tendon_names, art.num_fixed_tendons, FT), + (art.spatial_tendon_names, art.num_spatial_tendons, num_spatial), + ): + assert count == expected + assert isinstance(names, list) and len(names) == expected + assert all(isinstance(name, str) for name in names) + + +def test_finders(monkeypatch, backend): + """Finders return fresh lists or a cached index proxy, honor pattern order, and reject unmatched patterns.""" + art, _ = _articulation(backend, "cpu", monkeypatch) + finders = {"find_bodies": B, "find_joints": J, "find_fixed_tendons": FT} + if backend != "newton": + finders["find_spatial_tendons"] = ST + + for finder_name, count in finders.items(): + finder = getattr(art, finder_name) + indices, names = finder(".*") + proxy, proxy_names = finder(".*", as_proxy=True) + assert indices == list(range(count)) and len(names) == count + assert all(isinstance(name, str) for name in names) + assert indices == proxy.torch.tolist() and names == proxy_names + assert proxy is finder(".*", as_proxy=True)[0] + assert proxy.dtype == wp.int32 and str(proxy.device) == art.device + assert finder(names[0]) == ([0], [names[0]]) + # Mutating a returned list must not corrupt the cache. + indices.clear() + names.append("corrupted") + assert finder(".*") == (list(range(count)), proxy_names) + assert finder([proxy_names[1], proxy_names[0]], preserve_order=True) == ([1, 0], proxy_names[1::-1]) + with pytest.raises(ValueError): + finder("nonexistent_xyz") + + +def test_resolve_ids_accept_tensor_views(monkeypatch, backend): + """Selector resolution honors tensor views, and environment slices alias the cached indices without copies.""" + art, _ = _articulation(backend, "cpu", monkeypatch, num_instances=4, num_joints=4, num_bodies=4) + ids = torch.arange(4, dtype=torch.int32) + for resolve in (art._resolve_env_ids, art._resolve_joint_ids, art._resolve_body_ids): + assert resolve(ids).shape[0] == 4 + assert resolve(ids[:2]).shape[0] == 2 + + cached = wp.to_torch(art._ALL_INDICES) + for selection in (slice(None), slice(1, None, 2), slice(0, 0)): + resolved = wp.to_torch(art._resolve_env_ids(selection)) + torch.testing.assert_close(resolved, cached[selection]) + assert resolved.data_ptr() == cached[selection].data_ptr() + assert resolved.stride() == cached[selection].stride() + + +# --------------------------------------------------------------------------- +# Fixed-tendon targets +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("backend", backends("physx", "ovphysx")) +def test_commanding_a_tendon_target_schedules_the_write(monkeypatch, backend): + art, _ = _articulation(backend, "cpu", monkeypatch, num_fixed_tendons=2) + art.set_fixed_tendon_position_target_index(target=torch.zeros((N, 2))) + assert art._fixed_tendon_target_dirty is True + + +@requires("newton") +def test_newton_tendon_target_requires_a_tendon_actuator(monkeypatch): + """Without a MuJoCo tendon actuator Newton cannot command targets, and an untransmitting solver says so.""" + from isaaclab_newton.physics import NewtonManager + + art, _ = _articulation("newton", "cpu", monkeypatch, num_fixed_tendons=2) + with pytest.raises(RuntimeError, match="no MuJoCo tendon actuator"): + art.set_fixed_tendon_position_target_index(target=torch.zeros((N, 2))) + with pytest.raises(NotImplementedError, match="does not drive fixed tendons"): + NewtonManager.create_fixed_tendon_control(art) + + +# --------------------------------------------------------------------------- +# Data properties +# --------------------------------------------------------------------------- + + +def _properties(suffix: tuple[int, ...], dtype, names: str) -> list[tuple[str, tuple[int, ...], type]]: + return [(name, suffix, dtype) for name in names.split()] + + +# (property, shape after the instance axis, dtype) +_DATA_PROPERTIES = [ + *_properties((), wp.transformf, "root_link_pose_w root_com_pose_w default_root_pose"), + *_properties((), wp.spatial_vectorf, "root_link_vel_w root_com_vel_w default_root_vel"), + *_properties((), wp.quatf, "root_link_quat_w root_com_quat_w"), + *_properties( + (), + wp.vec3f, + """root_link_pos_w root_link_lin_vel_w root_link_ang_vel_w root_com_pos_w root_com_lin_vel_w + root_com_ang_vel_w projected_gravity_b root_link_lin_vel_b root_link_ang_vel_b root_com_lin_vel_b + root_com_ang_vel_b""", + ), + *_properties((), wp.float32, "heading_w"), + *_properties((B,), wp.transformf, "body_link_pose_w body_com_pose_w body_com_pose_b"), + *_properties((B,), wp.spatial_vectorf, "body_link_vel_w body_com_vel_w body_com_acc_w"), + *_properties((B,), wp.quatf, "body_link_quat_w body_com_quat_w body_com_quat_b"), + *_properties( + (B,), wp.vec3f, "body_link_pos_w body_link_lin_vel_w body_link_ang_vel_w body_com_pos_w body_com_pos_b" + ), + *_properties((B,), wp.float32, "body_mass"), + *_properties((B, 9), wp.float32, "body_inertia"), + *_properties( + (J,), + wp.float32, + """joint_pos joint_vel joint_acc joint_stiffness joint_damping joint_armature joint_friction_coeff + joint_vel_limits joint_effort_limits default_joint_pos default_joint_vel""", + ), + *_properties((J,), wp.vec2f, "joint_pos_limits soft_joint_pos_limits"), + *_properties((FT,), wp.float32, "fixed_tendon_stiffness fixed_tendon_damping"), + *_properties((FT,), wp.vec2f, "fixed_tendon_pos_limits"), +] +# Tendon properties that Newton does not implement. +_PHYSX_FAMILY_TENDON_PROPERTIES = [ + *_properties((FT,), wp.float32, "fixed_tendon_limit_stiffness fixed_tendon_rest_length fixed_tendon_offset"), + *_properties( + (ST,), + wp.float32, + "spatial_tendon_stiffness spatial_tendon_damping spatial_tendon_limit_stiffness spatial_tendon_offset", + ), +] +_ALIASES = { + "root_pose_w": "root_link_pose_w", + "root_pos_w": "root_link_pos_w", + "root_quat_w": "root_link_quat_w", + "root_vel_w": "root_com_vel_w", + "root_lin_vel_w": "root_com_lin_vel_w", + "root_ang_vel_w": "root_com_ang_vel_w", + "body_pose_w": "body_link_pose_w", + "body_pos_w": "body_link_pos_w", + "body_quat_w": "body_link_quat_w", + "body_vel_w": "body_com_vel_w", + "body_lin_vel_w": "body_com_lin_vel_w", + "body_ang_vel_w": "body_com_ang_vel_w", + "joint_limits": "joint_pos_limits", + "joint_friction": "joint_friction_coeff", +} + + +def test_data_properties(monkeypatch, backend, device): + """Every data property is a ProxyArray of the advertised shape and dtype; components and aliases share data.""" + art, _ = _articulation(backend, device, monkeypatch) + if backend == "newton": + # Native selections may stride over other articulations in each world. + for name in ("root_link_pose_w", "root_com_vel_w", "body_link_pose_w", "body_com_vel_w"): + binding = getattr(art.data, "_sim_bind_" + name) + view = wp.empty((2 * binding.shape[0], *binding.shape[1:]), dtype=binding.dtype, device=device)[::2] + view.assign(binding) + setattr(art.data, "_sim_bind_" + name, view) + art.data._pin_proxy_arrays() + art.data.update(dt=0.01) + if backend == "newton": + # Derived Newton quantities stay unallocated until read. + for name in ("root_link_vel_w", "body_link_vel_w", "body_com_pose_b", "root_com_pose_w", "body_com_pose_w"): + assert getattr(art.data, "_" + name).data is None, name + for name in ("projected_gravity_b", "heading_w"): + assert getattr(art.data, "_" + name).data is None, name + + properties = _DATA_PROPERTIES + (_PHYSX_FAMILY_TENDON_PROPERTIES if backend != "newton" else []) + for name, suffix, dtype in properties: + if backend == "newton" and name == "body_com_pose_b": + # Newton stores a position-only COM and warns that the pose appends a unit quaternion. + with pytest.warns(UserWarning, match="unit quaternion"): + value = getattr(art.data, name) + else: + value = getattr(art.data, name) + shape = (N, *suffix) + if backend == "physx" and name == "fixed_tendon_pos_limits": + # Known inconsistency: PhysX exposes (N, T, 2) float32 although its docstring and OVPhysX + # advertise (N, T) vec2f. Pin the current layout so a change on either side is noticed. + shape, dtype = (*shape, 2), wp.float32 + assert isinstance(value, ProxyArray), name + assert (value.shape, value.dtype) == (shape, dtype), name + + for frame in ("root_link", "root_com", "body_link", "body_com"): + for quantity, components in (("pose", ("pos", "quat")), ("vel", ("lin_vel", "ang_vel"))): + packed = getattr(art.data, f"{frame}_{quantity}_w").torch + for component, expected in zip(components, (packed[..., :3], packed[..., 3:]), strict=True): + view = getattr(art.data, f"{frame}_{component}_w").torch + torch.testing.assert_close(view, expected) + assert view.data_ptr() == expected.data_ptr() + assert view.stride() == expected.stride() + + # Random mock state makes link and COM quantities differ, so a retargeted alias fails. + for alias, canonical in _ALIASES.items(): + alias_value, canonical_value = getattr(art.data, alias), getattr(art.data, canonical) + assert (alias_value.shape, alias_value.dtype) == (canonical_value.shape, canonical_value.dtype), alias + assert torch.equal(alias_value.torch, canonical_value.torch), alias + + +def test_actuator_backed_data_properties(monkeypatch, backend): + """Soft velocity limits alias the actuator buffer without warning; deprecated command properties warn and alias.""" + art, _ = _articulation(backend, "cpu", monkeypatch, num_joints=4) + actuators = art.actuators + soft_joint_vel_limits = torch.arange(1, N * 4 + 1, dtype=torch.float32).reshape(N, 4) / 2.0 + wp.copy(actuators._soft_joint_vel_limits, wp.from_torch(soft_joint_vel_limits)) + + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + value = art.data.soft_joint_vel_limits + assert value.warp.ptr == art.data.soft_joint_vel_limits.warp.ptr == actuators._soft_joint_vel_limits.ptr + assert (value.shape, value.dtype) == ((N, 4), wp.float32) + torch.testing.assert_close(value.torch, soft_joint_vel_limits, rtol=0.0, atol=0.0) + + aliases = { + "joint_pos_target": actuators.target_command.position, + "joint_vel_target": actuators.target_command.velocity, + "joint_effort_target": actuators.target_command.effort, + "computed_torque": actuators.computed_effort, + "applied_torque": actuators.applied_effort, + } + for name, collection_buffer in aliases.items(): + with pytest.warns(DeprecationWarning, match=name): + assert getattr(art.data, name) is collection_buffer, name + + +@pytest.mark.parametrize("backend", backends("physx", "ovphysx")) +def test_dynamics_buffers_allocate_on_read_and_reuse_across_refreshes(monkeypatch, backend, device): + """Optional dynamics allocate independently on demand and retain stable public views.""" + art, _ = _articulation(backend, device, monkeypatch, num_joints=3, num_bodies=4) + names = ("body_com_jacobian_w", "mass_matrix", "gravity_compensation_forces") + art.data.update(dt=0.01) + for index, name in enumerate(names): + for unread in names[index:]: + assert getattr(art.data, "_" + unread).data is None + wrapper = getattr(art.data, name) + assert wrapper is getattr(art.data, "_" + name).data + art.data.update(dt=0.01) + assert getattr(art.data, name) is wrapper + assert art.data._body_link_jacobian_w is None + wrapper = art.data.body_link_jacobian_w + art.data.update(dt=0.01) + assert art.data.body_link_jacobian_w is wrapper + + +# --------------------------------------------------------------------------- +# Cache invalidation +# --------------------------------------------------------------------------- + +_BODY_ORDERINGS = pytest.mark.parametrize( + "body_ordering", [None, ("body_0", "body_3", "body_2", "body_1")], ids=["backend_order", "reordered"] +) +_ROOT_BODY_FRAME_VELOCITIES = ["_root_link_lin_vel_b", "_root_link_ang_vel_b", "_root_com_lin_vel_b"] +_ROOT_BODY_FRAME_VELOCITIES += ["_root_com_ang_vel_b"] + + +@_BODY_ORDERINGS +def test_root_writes_invalidate_dependent_caches(monkeypatch, backend, body_ordering): + art, _ = _articulation(backend, "cpu", monkeypatch, num_joints=3, num_bodies=4, body_ordering=body_ordering) + art.data.update(dt=0.01) + pose_dependents = ["_root_link_vel_w", "_body_link_vel_w", "_projected_gravity_b", "_heading_w"] + if backend != "newton": + pose_dependents.append("_body_com_vel_w") + + buffers = _prime_timestamped_properties(art.data, pose_dependents + _ROOT_BODY_FRAME_VELOCITIES) + if backend == "ovphysx" and art.data._body_com_vel_w_backend is not None: + buffers.append(("_body_com_vel_w_backend", art.data._body_com_vel_w_backend)) + art.write_root_link_pose_to_sim_index(root_pose=_payload((N,), 7, wp.transformf, "cpu")) + _assert_buffers_stale(art.data, buffers) + + buffers = _prime_timestamped_properties(art.data, _ROOT_BODY_FRAME_VELOCITIES) + art.write_root_com_velocity_to_sim_index(root_velocity=_payload((N,), 6, wp.spatial_vectorf, "cpu")) + _assert_buffers_stale(art.data, buffers) + + +@_BODY_ORDERINGS +def test_set_coms_invalidates_same_timestamp_dependents(monkeypatch, backend, body_ordering): + _ignore_newton_model_changes(backend, monkeypatch) + art, _ = _articulation(backend, "cpu", monkeypatch, num_joints=3, num_bodies=4, body_ordering=body_ordering) + art.data.update(dt=0.01) + names = ["_root_com_pose_w", "_root_link_vel_w", "_body_com_pose_w", "_body_link_vel_w"] + names += _ROOT_BODY_FRAME_VELOCITIES + if backend == "newton": + names.append("_body_com_pose_b") + else: + names += ["_root_com_vel_w", "_body_com_vel_w"] + if backend == "physx": + # COM changes also move the COM Jacobian and the generalized dynamics. + names += ["_body_com_jacobian_w", "_mass_matrix", "_gravity_compensation_forces"] + for frame in ("", "_link", "_com"): + for owner in ("root", "body"): + names.append(f"_{owner}{frame}_state_w{'_buf' if backend == 'ovphysx' else ''}") + if backend == "newton": + coms = wp.zeros((N, 4), dtype=wp.vec3f, device="cpu") + else: + coms = wp.from_torch(_payload((N, 4), 7, wp.transformf, "cpu"), dtype=wp.transformf) + + for set_coms in (art.set_coms_index, art.set_coms_mask): + buffers = _prime_timestamped_properties(art.data, names) + if backend == "ovphysx" and art.data._body_com_vel_w_backend is not None: + art.data._body_com_vel_w_backend.timestamp = art.data._sim_timestamp + buffers.append(("_body_com_vel_w_backend", art.data._body_com_vel_w_backend)) + set_coms(coms=coms) + _assert_buffers_stale(art.data, buffers) + + +@requires("physx") +def test_physx_com_cache_survives_com_and_joint_writes(monkeypatch): + """A COM write refreshes the PhysX COM cache directly, and a joint write leaves it valid.""" + art, view = _articulation("physx", "cpu", monkeypatch, num_joints=3, num_bodies=4) + view.get_coms = MagicMock(wraps=view.get_coms) + + art.set_coms_index(coms=wp.zeros((N, 4), dtype=wp.transformf, device="cpu"), full_data=True) + art.data.body_com_pose_b + art.write_joint_position_to_sim_index(position=torch.zeros((N, 3)), full_data=True) + art.data.body_com_pose_b + + assert view.get_coms.call_count == 0 + + +# --------------------------------------------------------------------------- +# Writers +# --------------------------------------------------------------------------- + +# (writer, keyword, selected item axis, trailing size, dtype, read-back getter) +_WRITERS = [ + ("write_root_pose_to_sim", "root_pose", None, 7, wp.transformf, "data.root_link_pose_w"), + ("write_root_link_pose_to_sim", "root_pose", None, 7, wp.transformf, "data.root_link_pose_w"), + ("write_root_com_pose_to_sim", "root_pose", None, 7, wp.transformf, "data.root_com_pose_w"), + ("write_root_velocity_to_sim", "root_velocity", None, 6, wp.spatial_vectorf, "data.root_com_vel_w"), + ("write_root_link_velocity_to_sim", "root_velocity", None, 6, wp.spatial_vectorf, "data.root_link_vel_w"), + ("write_root_com_velocity_to_sim", "root_velocity", None, 6, wp.spatial_vectorf, "data.root_com_vel_w"), + ("write_joint_position_to_sim", "position", "joint", 0, wp.float32, "data.joint_pos"), + ("write_joint_velocity_to_sim", "velocity", "joint", 0, wp.float32, "data.joint_vel"), + ("write_joint_stiffness_to_sim", "stiffness", "joint", 0, wp.float32, "data.joint_stiffness"), + ("write_joint_damping_to_sim", "damping", "joint", 0, wp.float32, "data.joint_damping"), + ("write_joint_position_limit_to_sim", "limits", "joint", 2, wp.vec2f, "data.joint_pos_limits"), + ("write_joint_velocity_limit_to_sim", "limits", "joint", 0, wp.float32, "data.joint_vel_limits"), + ("write_joint_effort_limit_to_sim", "limits", "joint", 0, wp.float32, "data.joint_effort_limits"), + ("write_joint_armature_to_sim", "armature", "joint", 0, wp.float32, "data.joint_armature"), + ( + "write_joint_friction_coefficient_to_sim", + "joint_friction_coeff", + "joint", + 0, + wp.float32, + "data.joint_friction_coeff", + ), + ("set_joint_position_target", "target", "joint", 0, wp.float32, "actuators.target_command.position"), + ("set_joint_velocity_target", "target", "joint", 0, wp.float32, "actuators.target_command.velocity"), + ("set_joint_effort_target", "target", "joint", 0, wp.float32, "actuators.target_command.effort"), + ("set_masses", "masses", "body", 0, wp.float32, "data.body_mass"), + ("set_coms", "coms", "body", 7, wp.transformf, "data.body_com_pose_b"), + ("set_inertias", "inertias", "body", 9, wp.float32, "data.body_inertia"), + *[ + (f"set_fixed_tendon_{name}", name, "fixed_tendon", 0, wp.float32, f"data.fixed_tendon_{name}") + for name in ("stiffness", "damping", "limit_stiffness", "rest_length", "offset") + ], + ("set_fixed_tendon_position_limit", "limit", "fixed_tendon", 2, wp.vec2f, "data.fixed_tendon_pos_limits"), + *[ + (f"set_spatial_tendon_{name}", name, "spatial_tendon", 0, wp.float32, f"data.spatial_tendon_{name}") + for name in ("stiffness", "damping", "limit_stiffness", "offset") + ], +] +# Writers that also accept one scalar for every selected entry; limit pairs reject it. +_SCALAR_WRITERS = { + "write_joint_stiffness_to_sim", + "write_joint_damping_to_sim", + "write_joint_position_limit_to_sim", + "write_joint_velocity_limit_to_sim", + "write_joint_effort_limit_to_sim", + "write_joint_armature_to_sim", + *[writer for writer, _, axis, *_ in _WRITERS if axis in ("fixed_tendon", "spatial_tendon")], +} +# Tendon setters that Newton implements. +_NEWTON_TENDON_WRITERS = {"set_fixed_tendon_stiffness", "set_fixed_tendon_damping", "set_fixed_tendon_position_limit"} + + +@pytest.mark.parametrize("selection", ["index", "mask"]) +@pytest.mark.parametrize("writer, kwarg, axis, trailing, dtype, getter", _WRITERS, ids=[w[0] for w in _WRITERS]) +def test_writer(monkeypatch, backend, device, selection, writer, kwarg, axis, trailing, dtype, getter): + """Writers accept warp and torch data, write only the selected entries, and reject mismatched shapes.""" + is_tendon = axis in ("fixed_tendon", "spatial_tendon") + if backend == "newton" and is_tendon and writer not in _NEWTON_TENDON_WRITERS: + pytest.skip("Newton implements only fixed-tendon stiffness, damping, and position limits") + if backend == "newton" and writer == "set_coms": + # Newton stores the COM as a position only. + trailing, dtype, getter = 3, wp.vec3f, "data.body_com_pos_b" + _ignore_newton_model_changes(backend, monkeypatch) + art, _ = _articulation(backend, device, monkeypatch) + art.data.update(dt=0.01) + method = getattr(art, f"{writer}_{selection}") + shape = (N,) if axis is None else (N, _AXIS_COUNTS[axis]) + + def read() -> ProxyArray: + if is_tendon: + # Tendon setters stage values that the property push sends to the simulation. + getattr(art, f"write_{axis}_properties_to_sim_index")() + return attrgetter(getter)(art) + + # Full warp data. + expected = _payload(shape, trailing, dtype, device) + method(**{kwarg: wp.from_torch(expected.contiguous(), dtype=dtype)}) + _assert_reads_back(read(), expected, writer) + + # Partial torch data for environment 1 and, for per-item writers, items 1 and 0 in that order. + values = _payload(shape, trailing, dtype, device, offset=100.0) + items = [1, 0] + if selection == "index": + kwargs = {kwarg: values[1:], "env_ids": torch.tensor([1], dtype=torch.int32, device=device)} + if axis is not None: + kwargs.update({kwarg: values[1:, items], f"{axis}_ids": items}) + else: + kwargs = {kwarg: values, "env_mask": _mask(N, [1], device)} + if axis is not None: + kwargs[f"{axis}_mask"] = _mask(shape[1], items, device) + method(**kwargs) + if axis is None: + expected[1] = values[1] + else: + expected[1, items] = values[1, items] + _assert_reads_back(read(), expected, f"{writer} partial") + + if writer in _SCALAR_WRITERS: + if dtype == wp.vec2f: + with pytest.raises((ValueError, TypeError)): + method(**{kwarg: 1.0}) + else: + method(**{kwarg: 1.0}) + _assert_reads_back(read(), torch.ones_like(expected), f"{writer} scalar") + + if backend == "newton" and selection == "mask" and is_tendon: + return # pinned by test_newton_fixed_tendon_mask_setters_reject_extra_environments + extra_env = _payload((N + 1, *shape[1:]), trailing, dtype, device) + with pytest.raises((AssertionError, RuntimeError)): + method(**{kwarg: extra_env}) + with pytest.raises((AssertionError, RuntimeError)): + method(**{kwarg: wp.from_torch(extra_env.contiguous(), dtype=dtype)}) + + +@requires("newton") +@pytest.mark.xfail( + raises=pytest.fail.Exception, + strict=True, + reason="Newton fixed-tendon mask setters slice full data without validating its shape", +) +def test_newton_fixed_tendon_mask_setters_reject_extra_environments(monkeypatch): + art, _ = _articulation("newton", "cpu", monkeypatch) + with pytest.raises((AssertionError, RuntimeError)): + art.set_fixed_tendon_stiffness_mask(stiffness=torch.ones((N + 1, FT))) + + +# quantity -> (writer, keyword, trailing size). Joint partial writes are covered by the ordering contract. +_PARTIAL_WRITES = { + "root_pose": ("write_root_link_pose_to_sim", "root_pose", 7), + "root_velocity": ("write_root_com_velocity_to_sim", "root_velocity", 6), + "mass": ("set_masses", "masses", 1), +} + + +def _read_backend_rows(backend: str, art, raw_backend, quantity: str) -> torch.Tensor: + """Read one articulation quantity from backend storage as one row per environment.""" + if backend == "ovphysx": + from isaaclab_ov import tensor_types as TT + + binding = {"root_pose": TT.ROOT_POSE, "root_velocity": TT.ROOT_VELOCITY, "mass": TT.BODY_MASS}[quantity] + values = torch.as_tensor(raw_backend.bindings[binding]._data) + elif quantity == "mass": + values = wp.to_torch(raw_backend.get_masses() if backend == "physx" else art.data._sim_bind_body_mass) + else: + getter = raw_backend.get_root_transforms if quantity == "root_pose" else raw_backend.get_root_velocities + values = wp.to_torch(getter() if backend == "physx" else getter(None)) + return values.reshape(art.num_instances, -1).cpu().clone() + + +@pytest.mark.parametrize("selection", ["index", "mask"]) +@pytest.mark.parametrize("quantity", _PARTIAL_WRITES) +def test_partial_write_preserves_unselected_backend_rows(monkeypatch, backend, selection, quantity): + art, raw_backend = _articulation(backend, "cpu", monkeypatch, num_joints=3, num_bodies=4) + writer, kwarg, trailing = _PARTIAL_WRITES[quantity] + width = 4 if quantity == "mass" else trailing + # Literal, per-element distinct payload; poses keep an identity rotation. + values = 100.0 * torch.arange(1, N + 1, dtype=torch.float32).unsqueeze(-1) + torch.arange(width) + if quantity == "root_pose": + values[:, 3:6] = 0.0 + values[:, 6] = 1.0 + before = _read_backend_rows(backend, art, raw_backend, quantity) + # Select the second environment (and the second body) so a writer that ignores the selection cannot pass. + env, body = 1, 1 + if selection == "index": + kwargs = {kwarg: values[env:], "env_ids": torch.tensor([env], dtype=torch.int32)} + if quantity == "mass": + kwargs.update({kwarg: values[env:, body : body + 1], "body_ids": [body]}) + else: + kwargs = {kwarg: values, "env_mask": _mask(N, [env], "cpu")} + if quantity == "mass": + kwargs["body_mask"] = _mask(4, [body], "cpu") + + getattr(art, f"{writer}_{selection}")(**kwargs) + + expected = before.clone() + if quantity == "mass": + expected[env, body] = values[env, body] + else: + expected[env] = values[env] + torch.testing.assert_close(_read_backend_rows(backend, art, raw_backend, quantity), expected, rtol=0.0, atol=0.0) + + +def test_partial_joint_state_write_follows_int64_selectors_and_restarts_acceleration(monkeypatch, backend): + """A fused joint-state write honors unsorted int64 selectors and restarts the finite-difference baseline.""" + art, raw_backend = _articulation(backend, "cpu", monkeypatch, num_joints=3, num_bodies=2) + art.data.update(dt=0.01) + expected_joint_pos = art.data.joint_pos.torch.clone() + expected_joint_vel = art.data.joint_vel.torch.clone() + # Newton reads body poses through FK bindings, so only its body velocity is a timestamped cache. + cached_properties = ("body_link_vel_w",) if backend == "newton" else ("body_link_pose_w", "body_com_vel_w") + caches = _prime_timestamped_properties(art.data, [f"_{name}" for name in cached_properties]) + # Seed the finite-difference state so an untouched cell shows whether the write leaked into it. + previous_joint_vel = wp.to_torch(art.data._previous_joint_vel) + previous_joint_vel.fill_(3.0) + wp.to_torch(art.data._joint_acc.data).fill_(4.0) + art.data._joint_acc.timestamp = -1.0 + position = torch.tensor([[0.1, -0.1]]) + velocity = torch.tensor([[0.2, -0.2]]) + + # int64 selectors with the joints in reverse order; the payload follows the selector order. + art.write_joint_state_to_sim_index( + position=position[:, [1, 0]], + velocity=velocity[:, [1, 0]], + env_ids=torch.tensor([1], dtype=torch.int64), + joint_ids=torch.tensor([2, 0], dtype=torch.int64), + skip_forward=True, + ) + + expected_joint_pos[1, [0, 2]] = position[0] + expected_joint_vel[1, [0, 2]] = velocity[0] + expected_previous_joint_vel = torch.full_like(previous_joint_vel, 3.0) + expected_previous_joint_vel[1, [0, 2]] = velocity[0] + expected_joint_acc = torch.full_like(previous_joint_vel, 4.0) + expected_joint_acc[1, [0, 2]] = 0.0 + torch.testing.assert_close(art.data.joint_pos.torch, expected_joint_pos) + torch.testing.assert_close(art.data.joint_vel.torch, expected_joint_vel) + torch.testing.assert_close(previous_joint_vel, expected_previous_joint_vel) + torch.testing.assert_close(wp.to_torch(art.data._joint_acc.data), expected_joint_acc) + if backend == "ovphysx": + # OVPhysX stamps the reset acceleration, so a read before the next step returns it unchanged. + assert art.data._joint_acc.timestamp == art.data._sim_timestamp + torch.testing.assert_close(art.data.joint_acc.torch, expected_joint_acc) + else: + torch.testing.assert_close(art.data.joint_acc.torch[1, [0, 2]], torch.zeros(2)) + backend_joint_pos, backend_joint_vel = read_backend_joint_state(backend, art, raw_backend) + torch.testing.assert_close(torch.from_numpy(backend_joint_pos), expected_joint_pos) + torch.testing.assert_close(torch.from_numpy(backend_joint_vel), expected_joint_vel) + # ``skip_forward`` leaves the body caches to the caller; a regular write invalidates them. + for name, buffer in caches: + assert buffer.timestamp == art.data._sim_timestamp, name + art.write_joint_state_to_sim_index(position=position, velocity=velocity, env_ids=[1], joint_ids=[0, 2]) + _assert_buffers_stale(art.data, caches) + + +def test_deprecated_writers_forward_to_public_writers(monkeypatch, backend): + """Deprecated joint-state, friction, and position-target writers warn and write what the public writers would.""" + art, raw_backend = _articulation(backend, "cpu", monkeypatch, num_joints=4, num_bodies=2) + art.data.update(dt=0.01) + expected_joint_pos = art.data.joint_pos.torch.clone() + expected_joint_vel = art.data.joint_vel.torch.clone() + + with pytest.warns(DeprecationWarning, match="write_joint_state_to_sim"): + art.write_joint_state_to_sim( + position=torch.tensor([[0.1]]), velocity=torch.tensor([[0.2]]), joint_ids=[1], env_ids=[1] + ) + expected_joint_pos[1, 1] = 0.1 + expected_joint_vel[1, 1] = 0.2 + torch.testing.assert_close(art.data.joint_pos.torch, expected_joint_pos) + torch.testing.assert_close(art.data.joint_vel.torch, expected_joint_vel) + backend_joint_pos, backend_joint_vel = read_backend_joint_state(backend, art, raw_backend) + torch.testing.assert_close(torch.from_numpy(backend_joint_pos), expected_joint_pos) + torch.testing.assert_close(torch.from_numpy(backend_joint_vel), expected_joint_vel) + + for writer_name, value in (("write_joint_friction_coefficient_to_sim", 0.5), ("write_joint_friction_to_sim", 0.25)): + friction = torch.full((N, 4), value) + with pytest.warns(DeprecationWarning): + getattr(art, writer_name)(friction) + _assert_reads_back(art.data.joint_friction_coeff, friction, writer_name) + + target = torch.tensor([[11.0, 12.0], [21.0, 22.0]]) + with pytest.warns(DeprecationWarning): + art.set_joint_position_target(target, joint_ids=slice(1, 3)) + expected_target = torch.zeros((N, 4)) + expected_target[:, 1:3] = target + torch.testing.assert_close(art.actuators.target_command.position.torch, expected_target) + + +@pytest.mark.parametrize("selection", ["index", "mask"]) +@pytest.mark.parametrize("kind", ["fixed", "spatial"]) +def test_write_tendon_properties_to_sim_selects_envs(monkeypatch, backend, device, selection, kind): + """Pushing tendon properties writes the selected environments to the backend.""" + if backend == "newton" and kind == "spatial": + pytest.skip("Newton does not support spatial tendons") + _ignore_newton_model_changes(backend, monkeypatch) + num_instances = 3 + art, raw_backend = _articulation(backend, device, monkeypatch, num_instances=num_instances, num_joints=2) + art.data.update(dt=0.01) + env_writes = [] + if backend == "physx": + + def capture(*args, indices=None, **kwargs): + env_writes.append(indices.numpy().tolist()) + + setattr(raw_backend, f"set_{kind}_tendon_properties", MagicMock(side_effect=capture)) + elif backend == "ovphysx": + from isaaclab_ov import tensor_types as TT + + prefix = "FIXED_TENDON" if kind == "fixed" else "SPATIAL_TENDON" + tendon_types = {getattr(TT, name) for name in dir(TT) if name.startswith(prefix)} + set_attribute = art._root_view.set_attribute + + def capture(name, values, *, indices=None, mask=None): + if name in tendon_types: + if mask is not None: + env_writes.append(np.flatnonzero(mask.numpy()).tolist()) + else: + env_writes.append(list(range(num_instances)) if indices is None else indices.numpy().tolist()) + set_attribute(name, values, indices=indices, mask=mask) + + object.__setattr__(art._root_view, "set_attribute", capture) + write = getattr(art, f"write_{kind}_tendon_properties_to_sim_{selection}") + + for selected_envs in (list(range(num_instances)), [0, 2]): + env_writes.clear() + if backend == "newton": + # Newton copies staged rows into the solver-model bindings; stage a value no environment holds yet. + art.set_fixed_tendon_stiffness_index(stiffness=float(10 * len(selected_envs))) + sim_before = wp.to_torch(art.data._sim_bind_fixed_tendon_stiffness).clone() + partial = len(selected_envs) < num_instances + if selection == "index": + write(env_ids=torch.tensor(selected_envs, device=device) if partial else None) + else: + write(env_mask=_mask(num_instances, selected_envs, device) if partial else None) + if backend == "newton": + sim_after = wp.to_torch(art.data._sim_bind_fixed_tendon_stiffness) + env_writes.append([env for env in range(num_instances) if not torch.equal(sim_after[env], sim_before[env])]) + assert env_writes, "no tendon properties were written" + assert all(envs == selected_envs for envs in env_writes), env_writes + + +# --------------------------------------------------------------------------- +# Joint and body ordering +# +# A configured ordering exposes joints and bodies under public names while the backend keeps its own order. These +# cases seed backend-order storage, act through the public API, and compare against independently derived +# expectations. Cyclic orderings are preferred because, unlike reversals, they are not their own inverse, so a swapped +# map fails. +# --------------------------------------------------------------------------- + + +def _names(prefix: str, count: int) -> tuple[str, ...]: + return tuple(f"{prefix}_{index}" for index in range(count)) + + +def _ordering(mode: str, names: tuple[str, ...], keep_first: bool = False) -> tuple[str, ...] | None: + """Return the public ordering of backend names for a mode, optionally keeping a fixed root first.""" + if keep_first: + rest = _ordering(mode, names[1:]) + return None if rest is None else (names[0], *rest) + return {"none": None, "identity": names, "reversed": names[::-1], "cyclic": (*names[-1:], *names[:-1])}[mode] + + +def _user_to_backend(ordering, count: int) -> np.ndarray: + """Return public-to-backend indices for an optional ordering map.""" + return np.arange(count) if ordering is None else np.asarray(ordering.user_to_backend_indices, dtype=np.int64) + + +def _backend_to_user(public_names: tuple[str, ...] | None, backend_names: tuple[str, ...]) -> np.ndarray: + """Derive backend-to-public indices independently from the ordering map.""" + if public_names is None: + return np.arange(len(backend_names)) + return np.asarray([public_names.index(name) for name in backend_names]) + + +def _install_cyclic_ordering(art) -> tuple[np.ndarray, np.ndarray]: + """Install cyclic joint and body orderings on a constructed articulation, keeping a fixed root body first.""" + art.cfg = replace( + art.cfg, + joint_ordering=_ordering("cyclic", tuple(art.backend_joint_names)), + body_ordering=_ordering("cyclic", tuple(art.backend_body_names), keep_first=art.is_fixed_base), + ) + # Re-resolve and re-stage the maps as backend initialization does after a config change. + art._resolve_and_install_ordering_maps() + art._ordering_configure_backend_staging() + return _user_to_backend(art.joint_ordering, art.num_joints), _user_to_backend(art.body_ordering, art.num_bodies) + + +def _to_torch(array) -> torch.Tensor: + """Return a CPU clone of a ProxyArray, Warp array, or NumPy array.""" + if isinstance(array, np.ndarray): + return torch.from_numpy(array.copy()) + if hasattr(array, "torch"): + return array.torch.detach().cpu().clone() + return wp.to_torch(array).detach().cpu().clone() + + +def _assert_proxy_close(actual, expected: torch.Tensor) -> None: + torch.testing.assert_close(actual.torch.cpu(), expected, rtol=1.0e-5, atol=1.0e-5) + + +def _make_body_state(num_instances: int, num_bodies: int) -> tuple[np.ndarray, ...]: + """Return distinct backend-order root pose, root velocity, link pose, COM, COM velocity, and acceleration.""" + root_pose = np.zeros((num_instances, 7), dtype=np.float32) + root_pose[:, 0] = 10.0 + np.arange(num_instances) + root_pose[:, 6] = 1.0 + root_vel = np.zeros((num_instances, 6), dtype=np.float32) + root_vel[:, 0] = 3.0 + np.arange(num_instances) + root_vel[:, 5] = 7.0 + np.arange(num_instances) + env, body = np.meshgrid(np.arange(num_instances), np.arange(num_bodies), indexing="ij") + link_pose = np.zeros((num_instances, num_bodies, 7), dtype=np.float32) + link_pose[..., 0], link_pose[..., 1], link_pose[..., 6] = 10.0 * body, env, 1.0 + com_pose_b = np.zeros((num_instances, num_bodies, 7), dtype=np.float32) + com_pose_b[..., 0], com_pose_b[..., 6] = body + 1.0, 1.0 + body_com_vel = np.zeros((num_instances, num_bodies, 6), dtype=np.float32) + body_com_vel[..., 0], body_com_vel[..., 5] = 20.0 + body, 30.0 + body + body_acc = np.zeros((num_instances, num_bodies, 6), dtype=np.float32) + body_acc[..., 0], body_acc[..., 3] = 100.0 + body, 200.0 + body + return root_pose, root_vel, link_pose, com_pose_b, body_com_vel, body_acc + + +def _set_body_state(backend: str, art, raw_backend, state: tuple[np.ndarray, ...]) -> None: + """Write backend-order body state from :func:`_make_body_state` into the backend mock.""" + root_pose, root_vel, link_pose, com_pose_b, body_com_vel, body_acc = state + device = art.device + if backend == "physx": + raw_backend._root_transforms = wp.array(root_pose, dtype=wp.float32, device=device) + raw_backend._root_velocities = wp.array(root_vel, dtype=wp.float32, device=device) + raw_backend._link_transforms = wp.array(link_pose, dtype=wp.float32, device=device) + raw_backend._link_velocities = wp.array(body_com_vel, dtype=wp.float32, device=device) + raw_backend._link_accelerations = wp.array(body_acc, dtype=wp.float32, device=device) + raw_backend._coms = wp.array(com_pose_b, dtype=wp.float32, device="cpu") + elif backend == "ovphysx": + from isaaclab_ov import tensor_types as TT + + for binding, value in ( + (TT.ROOT_POSE, root_pose), + (TT.ROOT_VELOCITY, root_vel), + (TT.LINK_POSE, link_pose), + (TT.LINK_VELOCITY, body_com_vel), + (TT.LINK_ACCELERATION, body_acc), + (TT.BODY_COM_POSE, com_pose_b), + ): + raw_backend.bindings[binding]._data = value.copy() + _invalidate_com_caches(art) + else: + root_pose_wp = wp.array(root_pose[:, None], dtype=wp.transformf, device=device) + root_vel_wp = wp.array(root_vel[:, None], dtype=wp.spatial_vectorf, device=device) + link_pose_wp = wp.array(link_pose[:, None], dtype=wp.transformf, device=device) + body_com_vel_wp = wp.array(body_com_vel[:, None], dtype=wp.spatial_vectorf, device=device) + body_com_pos_wp = wp.array(com_pose_b[:, None, :, :3], dtype=wp.vec3f, device=device) + raw_backend.set_mock_root_transforms(root_pose_wp) + raw_backend.set_mock_root_velocities(root_vel_wp) + raw_backend.set_mock_link_transforms(link_pose_wp) + raw_backend.set_mock_link_velocities(body_com_vel_wp) + raw_backend.set_mock_coms(body_com_pos_wp) + art.data._sim_bind_root_link_pose_w.assign(root_pose_wp[:, 0]) + art.data._sim_bind_root_com_vel_w.assign(root_vel_wp[:, 0]) + art.data._sim_bind_body_link_pose_w.assign(link_pose_wp[:, 0]) + art.data._sim_bind_body_com_vel_w.assign(body_com_vel_wp[:, 0]) + art.data._sim_bind_body_com_pos_b.assign(body_com_pos_wp[:, 0]) + art.data._previous_body_com_vel.zero_() + + +def _invalidate_com_caches(art) -> None: + art.data._body_com_pose_b.timestamp = -1.0 + if art.data._body_com_pose_b_backend is not None: + art.data._body_com_pose_b_backend.timestamp = -1.0 + + +def _seed_backend_coms(backend: str, art, raw_backend) -> np.ndarray: + """Seed distinct backend-order PhysX or OVPhysX COM transforms and return them.""" + value = 100.0 * np.arange(art.num_instances)[:, None] + 10.0 * np.arange(art.num_bodies) + coms = np.zeros((art.num_instances, art.num_bodies, 7), dtype=np.float32) + coms[..., :3] = value[..., None] + np.arange(1.0, 4.0) + coms[..., 6] = 1.0 + if backend == "physx": + raw_backend._coms = wp.array(coms, dtype=wp.float32, device="cpu") + else: + from isaaclab_ov import tensor_types as TT + + raw_backend.bindings[TT.BODY_COM_POSE]._data = coms.copy() + _invalidate_com_caches(art) + return coms + + +def _read_backend_coms(backend: str, art, raw_backend) -> np.ndarray: + if backend == "physx": + coms = raw_backend.get_coms() + else: + from isaaclab_ov import tensor_types as TT + + coms = art.root_view.get_attribute(TT.BODY_COM_POSE) + return coms.numpy().reshape(art.num_instances, art.num_bodies, 7).copy() + + +def _com_payload(env_ids: list[int], body_ids: list[int], base: float) -> np.ndarray: + """Return distinct identity-rotation COM transforms for the selected environments and bodies.""" + value = base + 100.0 * np.asarray(env_ids)[:, None] + 10.0 * np.asarray(body_ids) + coms = np.zeros((len(env_ids), len(body_ids), 7), dtype=np.float32) + coms[..., :3] = value[..., None] + np.arange(1.0, 4.0) + coms[..., 6] = 1.0 + return coms + + +def _backend_properties(backend: str, art, raw_backend) -> dict[str, torch.Tensor]: + """Return backend-order joint and body properties.""" + if backend == "physx": + friction = _to_torch(raw_backend.get_dof_friction_properties()) + properties = { + "stiffness": raw_backend.get_dof_stiffnesses(), + "damping": raw_backend.get_dof_dampings(), + "armature": raw_backend.get_dof_armatures(), + "position_limits": raw_backend.get_dof_limits(), + "velocity_limits": raw_backend.get_dof_max_velocities(), + "effort_limits": raw_backend.get_dof_max_forces(), + "mass": raw_backend.get_masses(), + "inertia": raw_backend.get_inertias(), + "com": raw_backend.get_coms(), + } + elif backend == "ovphysx": + from isaaclab_ov import tensor_types as TT + + friction = _to_torch(raw_backend.bindings[TT.DOF_FRICTION_PROPERTIES]._data) + bindings = { + "stiffness": TT.DOF_STIFFNESS, + "damping": TT.DOF_DAMPING, + "armature": TT.DOF_ARMATURE, + "position_limits": TT.DOF_LIMIT, + "velocity_limits": TT.DOF_MAX_VELOCITY, + "effort_limits": TT.DOF_MAX_FORCE, + "mass": TT.BODY_MASS, + "inertia": TT.BODY_INERTIA, + "com": TT.BODY_COM_POSE, + } + properties = {name: raw_backend.bindings[binding]._data for name, binding in bindings.items()} + else: + data = art.data + position_limits = torch.stack( + (_to_torch(data._sim_bind_joint_pos_limits_lower), _to_torch(data._sim_bind_joint_pos_limits_upper)), -1 + ) + return { + "stiffness": _to_torch(data._sim_bind_joint_stiffness_sim), + "damping": _to_torch(data._sim_bind_joint_damping_sim), + "armature": _to_torch(data._sim_bind_joint_armature), + "position_limits": position_limits, + "velocity_limits": _to_torch(data._sim_bind_joint_vel_limits_sim), + "effort_limits": _to_torch(data._sim_bind_joint_effort_limits_sim), + "friction": _to_torch(data._sim_bind_joint_friction_coeff), + "viscous_friction": _to_torch(data._sim_bind_joint_viscous_friction_coeff), + "mass": _to_torch(data._sim_bind_body_mass), + "inertia": _to_torch(data._sim_bind_body_inertia), + "com": _to_torch(data._sim_bind_body_com_pos_b), + } + properties = {name: _to_torch(value) for name, value in properties.items()} + properties["com"] = properties["com"].reshape(art.num_instances, art.num_bodies, 7) + properties["inertia"] = properties["inertia"].reshape(art.num_instances, art.num_bodies, 9) + properties.update(friction=friction[..., 0], dynamic_friction=friction[..., 1], viscous_friction=friction[..., 2]) + return properties + + +def _write_joint_state(art, selection: str, operation: str, env_ids, joint_ids, position, velocity) -> None: + """Write public joint state of the selected cells; the payloads have one row per selected environment.""" + values = {"position": position, "velocity": velocity} + values = {name: value for name, value in values.items() if operation in (name, "state")} + writer = getattr(art, f"write_joint_{operation}_to_sim_{selection}") + if selection == "index": + values = {name: torch.tensor(value, dtype=torch.float32, device=art.device) for name, value in values.items()} + env_ids_wp = wp.array(env_ids, dtype=wp.int32, device=art.device) + writer(**values, env_ids=env_ids_wp, joint_ids=wp.array(joint_ids, dtype=wp.int32, device=art.device)) + return + full = {} + for name, value in values.items(): + full[name] = torch.zeros((art.num_instances, art.num_joints), device=art.device) + full[name][np.ix_(env_ids, joint_ids)] = torch.tensor(value, dtype=torch.float32, device=art.device) + env_mask = _mask(art.num_instances, env_ids, art.device) + writer(**full, env_mask=env_mask, joint_mask=_mask(art.num_joints, joint_ids, art.device)) + + +def _seed_joint_state(backend: str, art, raw_backend, position: np.ndarray, velocity: np.ndarray) -> None: + """Seed backend-order joint state and invalidate the public caches.""" + if backend == "physx": + raw_backend._dof_positions = wp.array(position, dtype=wp.float32, device=art.device) + raw_backend._dof_velocities = wp.array(velocity, dtype=wp.float32, device=art.device) + art.data._joint_pos.timestamp = -1.0 + art.data._joint_vel.timestamp = -1.0 + elif backend == "ovphysx": + _set_ovphysx_joint_state(raw_backend, position, velocity) + for name in ("_joint_pos_buf", "_joint_vel_buf", "_joint_pos_backend", "_joint_vel_backend"): + if getattr(art.data, name, None) is not None: + getattr(art.data, name).timestamp = -1.0 + else: + art.data._sim_bind_joint_pos.assign(wp.array(position, dtype=wp.float32, device=art.device)) + art.data._sim_bind_joint_vel.assign(wp.array(velocity, dtype=wp.float32, device=art.device)) + # Raw binding writes bypass the solver step, so mirror its publish into the public-order shadows. + art.data._refresh_user_order_joint_state() + + +def _set_ovphysx_joint_state(raw_backend, position: np.ndarray, velocity: np.ndarray) -> None: + from isaaclab_ov import tensor_types as TT + + raw_backend.bindings[TT.DOF_POSITION]._data = position.copy() + raw_backend.bindings[TT.DOF_VELOCITY]._data = velocity.copy() + + +_INITIAL_POSITION = np.asarray([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]], dtype=np.float32) +_INITIAL_VELOCITY = _INITIAL_POSITION + 100.0 + + +def _ordered_joint_articulation(backend: str, monkeypatch, mode: str = "cyclic"): + """Create a two-environment, three-joint articulation with seeded backend joint state.""" + art, raw_backend = get_articulation( + backend, 2, 3, 2, device="cpu", joint_ordering=_ordering(mode, _names("joint", 3)), monkeypatch=monkeypatch + ) + _seed_joint_state(backend, art, raw_backend, _INITIAL_POSITION, _INITIAL_VELOCITY) + return art, raw_backend + + +# --------------------------------------------------------------------------- +# Ordering resolution and public reads +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("mode", ["none", "identity", "cyclic"]) +def test_configured_ordering_resolves_public_names(monkeypatch, backend, mode): + """Default and backend-order configurations normalize to no ordering; others map public names to backend.""" + joint_names, body_names = _names("joint", 3), _names("body", 4) + joint_ordering, body_ordering = _ordering(mode, joint_names), _ordering(mode, body_names) + art, _ = get_articulation( + backend, + 2, + 3, + 4, + device="cpu", + joint_ordering=joint_ordering, + body_ordering=body_ordering, + monkeypatch=monkeypatch, + ) + + # The asset ordering properties delegate to the single source on data. + assert art.joint_ordering is art.data.joint_ordering + assert art.body_ordering is art.data.body_ordering + if mode != "cyclic": + assert art.joint_ordering is None and art.body_ordering is None + assert (art.joint_names, art.body_names) == (list(joint_names), list(body_names)) + return + assert (tuple(art.joint_names), tuple(art.body_names)) == (joint_ordering, body_ordering) + np.testing.assert_array_equal( + art.joint_ordering.backend_to_user_indices, _backend_to_user(joint_ordering, joint_names) + ) + np.testing.assert_array_equal( + art.body_ordering.backend_to_user_indices, _backend_to_user(body_ordering, body_names) + ) + + +def test_body_ordering_reorders_public_body_state(monkeypatch, backend): + """Body quantities follow the public body order while root quantities are unchanged.""" + identity_art, identity_raw = get_articulation(backend, 2, 1, 3, device="cpu", monkeypatch=monkeypatch) + ordered_art, ordered_raw = get_articulation(backend, 2, 1, 3, device="cpu", monkeypatch=monkeypatch) + state = _make_body_state(2, 3) + _set_body_state(backend, identity_art, identity_raw, state) + _set_body_state(backend, ordered_art, ordered_raw, state) + _, user_to_backend = _install_cyclic_ordering(ordered_art) + identity_art.data.update(dt=0.01) + ordered_art.data.update(dt=0.01) + + for name in ("body_com_pose_b", "body_com_acc_w", "body_link_vel_w", "body_com_pose_w"): + _assert_proxy_close( + getattr(ordered_art.data, name), _to_torch(getattr(identity_art.data, name))[:, user_to_backend] + ) + for name in ("root_com_pose_w", "root_link_vel_w"): + _assert_proxy_close(getattr(ordered_art.data, name), _to_torch(getattr(identity_art.data, name))) + + +@pytest.mark.parametrize("backend", backends("physx", "newton")) +@pytest.mark.parametrize("is_fixed_base", [False, True], ids=["floating", "fixed"]) +def test_ordering_reorders_public_dynamics_quantities(monkeypatch, backend, is_fixed_base): + num_instances, num_joints, num_bodies = 2, 3, 4 + identity_art, identity_raw = get_articulation( + backend, + num_instances, + num_joints, + num_bodies, + device="cpu", + is_fixed_base=is_fixed_base, + monkeypatch=monkeypatch, + ) + identity_manager = getattr(identity_art, "_test_simulation_manager", None) + ordered_art, ordered_raw = get_articulation( + backend, + num_instances, + num_joints, + num_bodies, + device="cpu", + is_fixed_base=is_fixed_base, + monkeypatch=monkeypatch, + ) + num_base_dofs = identity_art.num_base_dofs + num_dofs = num_joints + num_base_dofs + # A fixed root has no Jacobian row. + num_jacobi_bodies = num_bodies - (1 if num_base_dofs == 0 else 0) + jacobian = np.arange(num_instances * num_jacobi_bodies * 6 * num_dofs, dtype=np.float32) + jacobian = jacobian.reshape(num_instances, num_jacobi_bodies, 6, num_dofs) + mass_matrix = np.arange(num_instances * num_dofs**2, dtype=np.float32).reshape(num_instances, num_dofs, num_dofs) + gravity = np.arange(num_instances * num_dofs, dtype=np.float32).reshape(num_instances, num_dofs) + for art, raw_backend in ((identity_art, identity_raw), (ordered_art, ordered_raw)): + _set_body_state(backend, art, raw_backend, _make_body_state(num_instances, num_bodies)) + if backend == "physx": + raw_backend.set_mock_jacobians(wp.array(jacobian, dtype=wp.float32, device="cpu")) + raw_backend.set_mock_generalized_mass_matrices(wp.array(mass_matrix, dtype=wp.float32, device="cpu")) + raw_backend.set_mock_gravity_compensation_forces(wp.array(gravity, dtype=wp.float32, device="cpu")) + else: + model_jacobian = np.zeros((num_instances, num_bodies, 6, num_dofs), dtype=np.float32) + model_jacobian[:, num_bodies - num_jacobi_bodies :] = jacobian + raw_backend.set_mock_jacobians(wp.array(model_jacobian, dtype=wp.float32, device="cpu")) + raw_backend.set_mock_mass_matrices(wp.array(mass_matrix, dtype=wp.float32, device="cpu")) + if backend == "physx": + # Buffers allocated before the ordering is installed must follow it too. + for name in ("body_com_jacobian_w", "mass_matrix", "gravity_compensation_forces"): + getattr(ordered_art.data, name) + joint_user_to_backend, body_user_to_backend = _install_cyclic_ordering(ordered_art) + identity_art.data.update(dt=0.01) + ordered_art.data.update(dt=0.01) + if backend == "newton": + # Both mocks patch the same Newton module; read the shared model through the identity manager. + import isaaclab_newton.assets.articulation.articulation_data as newton_data_module + + monkeypatch.setattr(newton_data_module, "SimulationManager", identity_manager) + + dofs = np.concatenate((np.arange(num_base_dofs), num_base_dofs + joint_user_to_backend)) + bodies = body_user_to_backend if num_base_dofs else body_user_to_backend[body_user_to_backend != 0] - 1 + for name in ("body_com_jacobian_w", "body_link_jacobian_w"): + expected = _to_torch(getattr(identity_art.data, name))[:, bodies][:, :, :, dofs] + _assert_proxy_close(getattr(ordered_art.data, name), expected) + _assert_proxy_close(ordered_art.data.mass_matrix, _to_torch(identity_art.data.mass_matrix)[:, dofs][:, :, dofs]) + if backend == "physx": + expected = _to_torch(identity_art.data.gravity_compensation_forces)[:, dofs] + _assert_proxy_close(ordered_art.data.gravity_compensation_forces, expected) + + +@requires("ovphysx") +def test_ovphysx_ordered_reads_refresh_after_reset(monkeypatch): + """Ordered OVPhysX pose, velocity, and friction shadows re-read the backend after invalidation.""" + from isaaclab_ov import tensor_types as TT + + art, raw_backend = get_articulation( + "ovphysx", + 2, + 3, + 3, + device="cpu", + joint_ordering=_ordering("cyclic", _names("joint", 3)), + body_ordering=_ordering("cyclic", _names("body", 3)), + monkeypatch=monkeypatch, + ) + state = list(_make_body_state(2, 3)) + # Zero COM offsets make the link and COM velocities equal. + state[3][..., :3] = 0.0 + _set_body_state("ovphysx", art, raw_backend, state) + body_user_to_backend = _user_to_backend(art.body_ordering, 3) + joint_user_to_backend = _user_to_backend(art.joint_ordering, 3) + art.data.update(dt=0.01) + for name in ("body_link_pose_w", "body_com_vel_w", "body_link_vel_w", "root_link_vel_w"): + getattr(art.data, name).torch.clone() + + link_pose = state[2].copy() + link_pose[..., 0] += 1000.0 + velocity = state[4].copy() + velocity[..., 0] += 500.0 + raw_backend.bindings[TT.LINK_POSE]._data = link_pose + raw_backend.bindings[TT.LINK_VELOCITY]._data = velocity + art.data._reset_pose() + art.data._reset_velocity() + + _assert_proxy_close(art.data.body_link_pose_w, torch.from_numpy(link_pose[:, body_user_to_backend])) + _assert_proxy_close(art.data.body_com_vel_w, torch.from_numpy(velocity[:, body_user_to_backend])) + _assert_proxy_close(art.data.body_link_vel_w, torch.from_numpy(velocity[:, body_user_to_backend])) + _assert_proxy_close(art.data.root_link_vel_w, torch.from_numpy(velocity[:, 0])) + + # A fresh native read must gather too; initialization and writer caches can hide a missing gather. + friction = 100.0 + np.arange(2 * 3 * 3, dtype=np.float32).reshape(2, 3, 3) + raw_backend.bindings[TT.DOF_FRICTION_PROPERTIES]._data = friction + art.data._joint_friction_props_buf.timestamp = -1 + art.data._joint_friction_props_backend.timestamp = -1 + for component, name in enumerate(("friction_coeff", "dynamic_friction_coeff", "viscous_friction_coeff")): + expected = torch.from_numpy(friction[:, joint_user_to_backend, component]) + _assert_proxy_close(getattr(art.data, f"joint_{name}"), expected) + + +@requires("ovphysx") +def test_ovphysx_ordered_reads_gather_once_per_step(monkeypatch): + """Ordered root and body velocities share one backend read, and joint state is gathered once per step.""" + from isaaclab_ov import tensor_types as TT + + art, raw_backend = get_articulation( + "ovphysx", + 2, + 3, + 3, + device="cpu", + joint_ordering=("joint_2", "joint_1", "joint_0"), + body_ordering=("body_2", "body_1", "body_0"), + monkeypatch=monkeypatch, + ) + _set_body_state("ovphysx", art, raw_backend, _make_body_state(2, 3)) + art.data.update(dt=0.01) + art.data._binding_read = MagicMock(wraps=art.data._binding_read) + + for _ in range(2): + for name in ("body_com_vel_w", "root_link_vel_w", "joint_pos", "joint_vel"): + getattr(art.data, name).torch.clone() + + reads = [call.args[0] for call in art.data._binding_read.call_args_list] + assert [reads.count(binding) for binding in (TT.LINK_VELOCITY, TT.DOF_POSITION, TT.DOF_VELOCITY)] == [1, 1, 1] + + +@requires("ovphysx") +def test_ovphysx_joint_acceleration_differences_public_order_velocities(monkeypatch): + art, raw_backend = _ordered_joint_articulation("ovphysx", monkeypatch) + user_to_backend = _user_to_backend(art.joint_ordering, 3) + first = np.asarray([[1.0, 2.0, 4.0], [10.0, 20.0, 40.0]], dtype=np.float32) + second = first + np.asarray([[3.0, 5.0, 7.0], [11.0, 13.0, 17.0]], dtype=np.float32) + _set_ovphysx_joint_state(raw_backend, _INITIAL_POSITION, first) + art.data.update(0.1) + art.data.joint_acc.torch.clone() + _set_ovphysx_joint_state(raw_backend, _INITIAL_POSITION, second) + art.data.update(0.1) + + torch.testing.assert_close(art.data.joint_vel.torch, torch.from_numpy(second[:, user_to_backend])) + torch.testing.assert_close(art.data.joint_acc.torch, torch.from_numpy((second - first)[:, user_to_backend] / 0.1)) + + +# --------------------------------------------------------------------------- +# Property writes +# --------------------------------------------------------------------------- + +# (writer, keyword, axis, backend property, public getter, trailing size) +_PROPERTY_WRITERS = [ + ("write_joint_stiffness_to_sim", "stiffness", "joint", "stiffness", "joint_stiffness", 0), + ("write_joint_damping_to_sim", "damping", "joint", "damping", "joint_damping", 0), + ("write_joint_armature_to_sim", "armature", "joint", "armature", "joint_armature", 0), + ("write_joint_position_limit_to_sim", "limits", "joint", "position_limits", "joint_pos_limits", 2), + ("write_joint_velocity_limit_to_sim", "limits", "joint", "velocity_limits", "joint_vel_limits", 0), + ("write_joint_effort_limit_to_sim", "limits", "joint", "effort_limits", "joint_effort_limits", 0), + ("write_joint_friction_coefficient_to_sim", "joint_friction_coeff", "joint", "friction", "joint_friction_coeff", 0), + *[ + (f"write_joint_{kind}_friction_coefficient_to_sim", f"joint_{kind}_friction_coeff", "joint") + + (f"{kind}_friction", f"joint_{kind}_friction_coeff", 0) + for kind in ("dynamic", "viscous") + ], + ("set_masses", "masses", "body", "mass", "body_mass", 0), + ("set_inertias", "inertias", "body", "inertia", "body_inertia", 9), + ("set_coms", "coms", "body", "com", "body_com_pose_b", 7), +] + + +def _property_payload(shape: tuple[int, ...], trailing: int, offset: float) -> torch.Tensor: + """Return distinct property values; limits are ``[-v, v]`` pairs and COMs carry an identity rotation.""" + values = offset + torch.arange(np.prod(shape), dtype=torch.float32).reshape(shape) + if trailing == 2: + return torch.stack((-values, values), dim=-1) + if trailing: + values = values.unsqueeze(-1) + torch.arange(trailing, dtype=torch.float32) / 10.0 + if trailing == 7: + values[..., 3:] = torch.tensor([0.0, 0.0, 0.0, 1.0]) + return values + + +@pytest.mark.parametrize("mode", ["none", "cyclic"]) +@pytest.mark.parametrize("selection", ["index", "mask"]) +def test_property_writes_route_to_backend_and_preserve_other_rows(monkeypatch, backend, selection, mode): + """Partial public property writes reach the matching backend joints and bodies and keep every other entry.""" + num_instances, count = 2, 4 + if backend == "newton": + from isaaclab_newton.physics import NewtonManager + + monkeypatch.setattr(NewtonManager, "add_model_change", MagicMock()) + art, raw_backend = get_articulation( + backend, + num_instances, + count, + count, + device="cpu", + joint_ordering=_ordering(mode, _names("joint", count)), + body_ordering=_ordering(mode, _names("body", count)), + monkeypatch=monkeypatch, + ) + user_to_backend = {"joint": _user_to_backend(art.joint_ordering, count)} + user_to_backend["body"] = _user_to_backend(art.body_ordering, count) + expected = _backend_properties(backend, art, raw_backend) + + for case, (writer, kwarg, axis, name, getter, trailing) in enumerate(_PROPERTY_WRITERS): + if name not in expected: + continue # Newton has no dynamic joint friction. + if backend == "newton" and name == "com": + # Newton stores the COM as a position only. + getter, trailing = "body_com_pos_b", 3 + # A write to every environment and two items, then one to a single environment and two unsorted items. + for env_ids, items in (([0, 1], [0, 2]), ([1], [3, 1])): + soft_limits_before = art.data.soft_joint_pos_limits.torch.clone() + values = _property_payload( + (num_instances, count), trailing, offset=100.0 * (case + 1) + 10.0 * len(env_ids) + ) + selected = values[np.ix_(env_ids, items)] + if selection == "index": + kwargs = {kwarg: selected, "env_ids": env_ids, f"{axis}_ids": items} + else: + kwargs = {kwarg: values, "env_mask": _mask(num_instances, env_ids, "cpu")} + kwargs[f"{axis}_mask"] = _mask(count, items, "cpu") + getattr(art, f"{writer}_{selection}")(**kwargs) + + expected[name][np.ix_(env_ids, user_to_backend[axis][items])] = selected + backend_values = _backend_properties(backend, art, raw_backend)[name] + torch.testing.assert_close(backend_values, expected[name], rtol=0.0, atol=0.0, msg=f"{writer}: {mode}") + _assert_proxy_close(getattr(art.data, getter), expected[name][:, user_to_backend[axis]]) + if name == "position_limits": + # The soft limits of the written joints follow (factor 1.0); the others are untouched. + soft_limits = art.data.soft_joint_pos_limits.torch + torch.testing.assert_close(soft_limits[np.ix_(env_ids, items)], selected) + untouched = [joint for joint in range(count) if joint not in items] + torch.testing.assert_close(soft_limits[:, untouched], soft_limits_before[:, untouched]) + + +@pytest.mark.parametrize("backend", backends("physx", "ovphysx")) +@pytest.mark.parametrize("mode", ["none", "cyclic"]) +def test_duplicate_com_selectors_preserve_omitted_rows(monkeypatch, backend, mode): + """Duplicate body or environment selectors write their last value and keep every omitted backend row.""" + art, raw_backend = get_articulation( + backend, 2, 1, 4, device="cpu", body_ordering=_ordering(mode, _names("body", 4)), monkeypatch=monkeypatch + ) + expected = _seed_backend_coms(backend, art, raw_backend) + user_to_backend = _user_to_backend(art.body_ordering, 4) + if backend == "physx": + raw_backend.get_coms = MagicMock(wraps=raw_backend.get_coms) + + # Duplicate body IDs omit body 3. + body_ids = [0, 1, 1, 2] + payload = _com_payload([0, 1], body_ids, 3000.0) + art.set_coms_index(coms=wp.array(payload, dtype=wp.transformf), body_ids=wp.array(body_ids, dtype=wp.int32)) + public = art.data.body_com_pose_b.torch.numpy().copy() + if backend == "physx": + # The cold partial write seeds untouched bodies with one backend read; the public read reuses it. + assert raw_backend.get_coms.call_count == 1 + for offset, body in enumerate(body_ids): + expected[:, user_to_backend[body]] = payload[:, offset] + np.testing.assert_array_equal(_read_backend_coms(backend, art, raw_backend), expected) + np.testing.assert_array_equal(public, expected[:, user_to_backend]) + + # From cold caches, duplicate environment IDs omit environment 0 and must not mark the global COM caches valid. + _invalidate_com_caches(art) + env_ids = [1, 1] + payload = _com_payload(env_ids, [0, 1, 2, 3], 4000.0) + art.set_coms_index(coms=wp.array(payload, dtype=wp.transformf), env_ids=wp.array(env_ids, dtype=wp.int32)) + assert art.data._body_com_pose_b.timestamp < 0.0 + backend_staging = art.data._body_com_pose_b_backend + assert backend_staging is None or backend_staging.timestamp < 0.0 + expected[1, user_to_backend] = payload[-1] + np.testing.assert_array_equal(_read_backend_coms(backend, art, raw_backend), expected) + np.testing.assert_array_equal(art.data.body_com_pose_b.torch.numpy(), expected[:, user_to_backend]) + + +@pytest.mark.parametrize("backend", backends("physx", "ovphysx")) +@pytest.mark.parametrize("mode", ["none", "cyclic"]) +def test_static_com_cache_does_not_follow_sim_timestamp(monkeypatch, backend, mode): + """The COM is read from the backend once; later steps keep serving the cached public and root COM poses.""" + art, raw_backend = get_articulation( + backend, 2, 1, 4, device="cpu", body_ordering=_ordering(mode, _names("body", 4)), monkeypatch=monkeypatch + ) + root_pose = np.zeros((2, 7), dtype=np.float32) + root_pose[:, :3] = [[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]] + root_pose[:, 6] = 1.0 + if backend == "physx": + raw_backend._root_transforms = wp.array(root_pose, dtype=wp.float32) + raw_backend.get_coms = read_spy = MagicMock(wraps=raw_backend.get_coms) + else: + from isaaclab_ov import tensor_types as TT + + raw_backend.bindings[TT.ROOT_POSE]._data = root_pose.copy() + com_binding = raw_backend.bindings[TT.BODY_COM_POSE] + com_binding.read = read_spy = MagicMock(wraps=com_binding.read) + coms = _seed_backend_coms(backend, art, raw_backend) + expected_root = root_pose.copy() + expected_root[:, :3] += coms[:, 0, :3] + + for _ in range(3): + np.testing.assert_array_equal( + art.data.body_com_pose_b.torch.numpy(), coms[:, _user_to_backend(art.body_ordering, 4)] + ) + np.testing.assert_array_equal(art.data.root_com_pose_w.torch.numpy(), expected_root) + assert read_spy.call_count == 1 + art.data.update(dt=0.01) + + +# --------------------------------------------------------------------------- +# Root writes +# --------------------------------------------------------------------------- + + +def _read_backend_root_state(backend: str, art, raw_backend) -> tuple[np.ndarray, np.ndarray]: + """Read the backend root link pose and COM velocity.""" + if backend == "physx": + return raw_backend._root_transforms.numpy(), art.data.root_com_vel_w.warp.numpy() + if backend == "ovphysx": + from isaaclab_ov import tensor_types as TT + + return raw_backend.bindings[TT.ROOT_POSE]._data, raw_backend.bindings[TT.ROOT_VELOCITY]._data + return art.data._sim_bind_root_link_pose_w.numpy(), art.data._sim_bind_root_com_vel_w.numpy() + + +@pytest.mark.parametrize("operation", ["com_pose", "link_velocity"]) +@pytest.mark.parametrize("selection", ["index", "mask"]) +def test_floating_root_writers_use_the_root_body_com_after_body_reordering(monkeypatch, backend, selection, operation): + """Root COM conversions use the backend root body, whichever body is first in public order.""" + results = [] + for body_ordering in (_names("body", 4), _ordering("cyclic", _names("body", 4))): + art, raw_backend = get_articulation( + backend, 2, 1, 4, device="cpu", body_ordering=body_ordering, monkeypatch=monkeypatch + ) + state = list(_make_body_state(2, 4)) + state[3][..., :3] = 1.0 + np.arange(4)[:, None] + np.arange(3) + _set_body_state(backend, art, raw_backend, state) + pose = wp.array( + [[4.0, 5.0, 6.0, 0.0, 0.0, 0.0, 1.0], [14.0, 15.0, 16.0, 0.0, 0.0, 0.0, 1.0]], dtype=wp.transformf + ) + velocity = wp.array( + [[1.0, 2.0, 3.0, 0.4, 0.5, 0.6], [11.0, 12.0, 13.0, 1.4, 1.5, 1.6]], dtype=wp.spatial_vectorf + ) + if selection == "index": + selector, pose, velocity = {"env_ids": wp.array([0], dtype=wp.int32)}, pose[:1], velocity[:1] + else: + selector = {"env_mask": _mask(2, [0], "cpu")} + if operation == "com_pose": + getattr(art, f"write_root_com_pose_to_sim_{selection}")(root_pose=pose, **selector) + else: + getattr(art, f"write_root_link_pose_to_sim_{selection}")(root_pose=pose, **selector) + getattr(art, f"write_root_link_velocity_to_sim_{selection}")(root_velocity=velocity, **selector) + results.append(_read_backend_root_state(backend, art, raw_backend)[0 if operation == "com_pose" else 1][0]) + + assert art.body_names[0] != "body_0" + np.testing.assert_allclose(results[1], results[0], atol=1e-6) + + +# --------------------------------------------------------------------------- +# Joint state writes +# --------------------------------------------------------------------------- + + +# Each backend selects the joint map separately in every ordered writer, so every ordered (selection, operation) pair +# has one row; the selected environments and joints alternate across rows. +@pytest.mark.parametrize( + "mode, selection, operation, env_ids, joint_ids", + [ + ("none", "index", "state", [0, 1], [0]), + ("none", "mask", "position", [1], [0, 1, 2]), + ("cyclic", "index", "position", [1], [0, 1, 2]), + ("cyclic", "index", "velocity", [0, 1], [0]), + ("cyclic", "index", "state", [1], [2, 0, 1]), + ("cyclic", "mask", "position", [0, 1], [0]), + ("cyclic", "mask", "velocity", [1], [0, 1, 2]), + ("cyclic", "mask", "state", [0, 1], [2]), + ], +) +def test_partial_joint_write_preserves_backend_rows( + monkeypatch, backend, mode, selection, operation, env_ids, joint_ids +): + art, raw_backend = _ordered_joint_articulation(backend, monkeypatch, mode) + user_to_backend = _user_to_backend(art.joint_ordering, 3) + expected = {"position": _INITIAL_POSITION.copy(), "velocity": _INITIAL_VELOCITY.copy()} + + for base in (900.0, 1900.0): + payload = base + 10.0 * np.asarray(env_ids, dtype=np.float32)[:, None] + np.asarray(joint_ids) + _write_joint_state(art, selection, operation, env_ids, joint_ids, payload, payload + 0.5) + cells = np.ix_(env_ids, user_to_backend[joint_ids]) + if operation in ("position", "state"): + expected["position"][cells] = payload + if operation in ("velocity", "state"): + expected["velocity"][cells] = payload + 0.5 + position, velocity = read_backend_joint_state(backend, art, raw_backend) + np.testing.assert_array_equal(position, expected["position"]) + np.testing.assert_array_equal(velocity, expected["velocity"]) + + np.testing.assert_array_equal(art.data.joint_pos.torch.numpy(), expected["position"][:, user_to_backend]) + np.testing.assert_array_equal(art.data.joint_vel.torch.numpy(), expected["velocity"][:, user_to_backend]) + + +@pytest.mark.parametrize("backend", backends("physx", "ovphysx")) +@pytest.mark.parametrize("operation", ["position", "velocity", "state"]) +def test_partial_joint_write_after_a_step_preserves_newer_backend_rows(monkeypatch, backend, operation): + """Partial writes restage the backend rows once the simulation advances, also for duplicate full-length selectors. + + Newton writes the selected cells straight into the simulation buffers and has no row staging that could go stale. + """ + art, raw_backend = _ordered_joint_articulation(backend, monkeypatch) + user_to_backend = _user_to_backend(art.joint_ordering, 3) + art.data.joint_pos, art.data.joint_vel + _write_joint_state(art, "index", operation, [1], [0], np.full((1, 1), 901.0), np.full((1, 1), 902.0)) + + # The simulation advances and changes every backend row. + art.data._sim_timestamp += 0.01 + newer = {"position": _INITIAL_POSITION + 1000.0, "velocity": _INITIAL_VELOCITY + 1000.0} + if backend == "physx": + raw_backend._dof_positions.assign(wp.array(newer["position"], dtype=wp.float32)) + raw_backend._dof_velocities.assign(wp.array(newer["velocity"], dtype=wp.float32)) + else: + _set_ovphysx_joint_state(raw_backend, newer["position"], newer["velocity"]) + + # A duplicate full-length selector omits joint 2, whose newer backend value must survive. + joint_ids = [0, 1, 1] + payload = np.asarray([[911.0, 912.0, 913.0]], dtype=np.float32) + _write_joint_state(art, "index", operation, [1], joint_ids, payload, payload + 0.5) + for name, value in (("position", payload), ("velocity", payload + 0.5)): + if operation in (name, "state"): + for offset, joint in enumerate(joint_ids): + newer[name][1, user_to_backend[joint]] = value[0, offset] + position, velocity = read_backend_joint_state(backend, art, raw_backend) + np.testing.assert_array_equal(position, newer["position"]) + np.testing.assert_array_equal(velocity, newer["velocity"]) + + +@requires("physx") +@pytest.mark.parametrize("operation", ["position", "velocity", "state"]) +@pytest.mark.parametrize("selector", ["partial_mask", "default_mask", "explicit_full_index"]) +def test_physx_joint_writes_read_the_backend_only_for_partial_selectors(monkeypatch, operation, selector): + """A provably complete write skips the staging read; partial and opaque selectors read each axis once.""" + art, raw_backend = _ordered_joint_articulation("physx", monkeypatch) + raw_backend.get_dof_positions = MagicMock(wraps=raw_backend.get_dof_positions) + raw_backend.get_dof_velocities = MagicMock(wraps=raw_backend.get_dof_velocities) + + if selector == "partial_mask": + _write_joint_state(art, "mask", operation, [1], [0], np.full((1, 1), 901.0), np.full((1, 1), 902.0)) + else: + values = {"position": torch.tensor(_INITIAL_POSITION), "velocity": torch.tensor(_INITIAL_VELOCITY)} + values = {name: value for name, value in values.items() if operation in (name, "state")} + if selector == "default_mask": + getattr(art, f"write_joint_{operation}_to_sim_mask")(**values) + else: + joint_ids = wp.array([0, 1, 2], dtype=wp.int32) + getattr(art, f"write_joint_{operation}_to_sim_index")(**values, joint_ids=joint_ids) + + reads = int(selector != "default_mask") + assert raw_backend.get_dof_positions.call_count == (reads if operation in ("position", "state") else 0) + assert raw_backend.get_dof_velocities.call_count == (reads if operation in ("velocity", "state") else 0) + + +# --------------------------------------------------------------------------- +# Commands, wrenches, and configuration +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("mode", ["none", "cyclic"]) +@pytest.mark.parametrize("is_fixed_base", [False, True], ids=["floating", "fixed"]) +def test_external_wrenches_are_written_in_backend_body_order(monkeypatch, backend, mode, is_fixed_base): + """Public-order body wrenches reach each backend in backend body order and the requested frame.""" + num_instances, num_bodies = 2, 4 + backend_body_names = _names("body", num_bodies) + body_ordering = _ordering(mode, backend_body_names, keep_first=is_fixed_base) + art, raw_backend = get_articulation( + backend, + num_instances, + 1, + num_bodies, + device="cpu", + is_fixed_base=is_fixed_base, + body_ordering=body_ordering, + monkeypatch=monkeypatch, + ) + # Distinct quarter-turn rotations about Z and nonzero positions, in backend body order. + poses = np.zeros((num_instances, num_bodies, 7), dtype=np.float32) + poses[..., 0] = np.arange(num_bodies) + 10.0 + poses[..., 5] = np.sin(np.arange(num_bodies) * np.pi / 4) + poses[..., 6] = np.cos(np.arange(num_bodies) * np.pi / 4) + captured = {} + if backend == "physx": + + def capture_wrench(*, force_data, torque_data, position_data, indices, is_global): + assert position_data is None + captured.update(is_global=is_global, force=force_data.numpy().copy(), torque=torque_data.numpy().copy()) + + raw_backend.apply_forces_and_torques_at_position = capture_wrench + elif backend == "newton": + poses_wp = wp.array(poses[:, None], dtype=wp.transformf) + raw_backend.set_mock_link_transforms(poses_wp) + art.data._sim_bind_body_link_pose_w.assign(poses_wp[:, 0]) + art.data._refresh_user_order_body_state() + else: + from isaaclab_ov import tensor_types as TT + + raw_backend.bindings[TT.LINK_POSE]._data = poses + art.data._reset_pose() + + forces = np.arange(num_instances * num_bodies * 3, dtype=np.float32).reshape(num_instances, num_bodies, 3) + torques = forces + 100.0 + backend_to_user = _backend_to_user(body_ordering, backend_body_names) + composer = art.instantaneous_wrench_composer + for is_global in (False, True): + composer.set_forces_and_torques_index( + forces=wp.array(forces, dtype=wp.vec3f), torques=wp.array(torques, dtype=wp.vec3f), is_global=is_global + ) + with patch.object(composer, "compose_to_body_frame", wraps=composer.compose_to_body_frame) as compose: + art.write_data_to_sim() + assert compose.call_count == int(is_global and backend == "newton") + + expected_force, expected_torque = forces[:, backend_to_user], torques[:, backend_to_user] + if backend == "physx": + # PhysX receives the frame flag and rotates the wrench itself. + assert captured["is_global"] is is_global + force, torque = captured["force"], captured["torque"] + else: + if not is_global: + # The 0/90/180/270-degree body rotations, independent of the backend quaternion transform. + rotations = np.asarray( + [ + [[1, 0, 0], [0, 1, 0]], + [[0, -1, 0], [1, 0, 0]], + [[-1, 0, 0], [0, -1, 0]], + [[0, 1, 0], [-1, 0, 0]], + ], + dtype=np.float32, + ) + rotations = np.concatenate((rotations, np.tile([[[0, 0, 1]]], (4, 1, 1))), axis=1) + expected_force = np.einsum("bij,nbj->nbi", rotations, expected_force) + expected_torque = np.einsum("bij,nbj->nbi", rotations, expected_torque) + if backend == "newton": + wrench = art.data._sim_bind_body_external_wrench.numpy() + else: + wrench = raw_backend.bindings[TT.LINK_WRENCH]._data + # OVPhysX packs the link positions after the wrench. + np.testing.assert_array_equal(wrench[..., 6:9], poses[..., :3]) + force, torque = wrench[..., :3], wrench[..., 3:6] + np.testing.assert_allclose(force.reshape(expected_force.shape), expected_force, atol=1e-4) + np.testing.assert_allclose(torque.reshape(expected_torque.shape), expected_torque, atol=1e-4) + + +@requires("ovphysx") +def test_ovphysx_actuator_commands_are_written_in_backend_order(monkeypatch): + """Implicit position, velocity, and effort commands reach their matching backend joints.""" + from isaaclab_ov import tensor_types as TT + + art, raw_backend = _ordered_joint_articulation("ovphysx", monkeypatch) + backend_to_user = _backend_to_user(tuple(art.joint_names), _names("joint", 3)) + position = np.arange(6, dtype=np.float32).reshape(2, 3) + # Seed the actuator collection's submitted command buffers (no actuator groups recompute them). + art.actuators._joint_pos_target.assign(wp.array(position, dtype=wp.float32)) + art.actuators._joint_vel_target.assign(wp.array(position + 100.0, dtype=wp.float32)) + art.actuators._joint_effort_target_sim.assign(wp.array(position + 200.0, dtype=wp.float32)) + object.__setattr__(art, "_has_implicit_actuators", True) + object.__setattr__(art, "_can_write_effort", True) + + art.write_data_to_sim() + + for binding, offset in ( + (TT.DOF_POSITION_TARGET, 0.0), + (TT.DOF_VELOCITY_TARGET, 100.0), + (TT.DOF_ACTUATION_FORCE, 200.0), + ): + np.testing.assert_array_equal(raw_backend.bindings[binding]._data, (position + offset)[:, backend_to_user]) + + +@requires("ovphysx") +def test_ovphysx_partial_effort_target_write_keeps_commanded_targets(monkeypatch): + """A partial effort-target write keeps the other joints' commanded targets, not the last applied effort. + + The applied effort that ``write_data_to_sim`` pushes can differ from the commanded target (for example once explicit + actuators clip it), so it must not share the staging that a partial target write reuses for unselected joints. + """ + from isaaclab_ov import tensor_types as TT + + art, raw_backend = _ordered_joint_articulation("ovphysx", monkeypatch) + object.__setattr__(art, "_can_write_effort", True) + backend_to_user = _backend_to_user(tuple(art.joint_names), _names("joint", 3)) + targets = np.asarray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32) + art.set_joint_effort_target_mask(target=torch.tensor(targets)) + # An actuator model applies an effort that differs from the commanded target. + art.actuators._joint_effort_target_sim.assign(wp.array(targets + 100.0, dtype=wp.float32)) + art.write_data_to_sim() + np.testing.assert_array_equal( + raw_backend.bindings[TT.DOF_ACTUATION_FORCE]._data, (targets + 100.0)[:, backend_to_user] + ) + + art.set_joint_effort_target_index(target=torch.tensor([[901.0], [902.0]]), joint_ids=[1], env_ids=[0, 1]) + + targets[:, 1] = [901.0, 902.0] + np.testing.assert_array_equal(raw_backend.bindings[TT.DOF_ACTUATION_FORCE]._data, targets[:, backend_to_user]) + + +@requires("ovphysx") +@pytest.mark.parametrize("selector_kind", ["torch", "warp"]) +def test_ovphysx_int64_effort_target_selector_reaches_binding(monkeypatch, selector_kind): + """Indexed effort-target environment selectors reach the OVPhysX binding as int32 Warp indices.""" + from isaaclab_ov import tensor_types as TT + + art, raw_backend = get_articulation("ovphysx", 2, 3, 2, device="cpu", monkeypatch=monkeypatch) + object.__setattr__(art, "_can_write_effort", True) + expected = raw_backend.bindings[TT.DOF_ACTUATION_FORCE]._data.copy() + set_attribute = art._root_view.set_attribute + captured = [] + + def strict_set_attribute(name, values, *, indices=None, mask=None): + if name == TT.DOF_ACTUATION_FORCE: + assert isinstance(indices, wp.array) and indices.dtype == wp.int32 and str(indices.device) == art.device + captured.append(indices.numpy().tolist()) + set_attribute(name, values, indices=indices, mask=mask) + + monkeypatch.setattr(art._root_view, "set_attribute", strict_set_attribute, raising=False) + env_ids = torch.tensor([1], dtype=torch.int64) + if selector_kind == "warp": + env_ids = wp.from_torch(env_ids, dtype=wp.int64) + + art.set_joint_effort_target_index(target=torch.tensor([[7.0]]), env_ids=env_ids, joint_ids=torch.tensor([2])) + + expected[1] = [0.0, 0.0, 7.0] + assert captured == [[1]] + np.testing.assert_array_equal(raw_backend.bindings[TT.DOF_ACTUATION_FORCE]._data, expected) + + +@requires("ovphysx") +def test_ovphysx_configured_defaults_use_public_joint_names(monkeypatch): + art, _ = get_articulation( + "ovphysx", 1, 3, 2, device="cpu", joint_ordering=("joint_2", "joint_1", "joint_0"), monkeypatch=monkeypatch + ) + patterns = {"joint_0": 10.0, "joint_1": 20.0, "joint_2": 30.0} + + art._resolve_joint_values(patterns, art.data._default_joint_pos) + + expected = [[patterns[name] for name in art.joint_names]] + np.testing.assert_array_equal(art.data.default_joint_pos.warp.numpy(), expected) + + +@requires("physx") +def test_physx_newton_actuator_forces_are_written_in_backend_order(monkeypatch): + art, raw_backend = get_articulation( + "physx", + 2, + 4, + 2, + device="cpu", + joint_ordering=_ordering("cyclic", _names("joint", 4)), + monkeypatch=monkeypatch, + ) + forces = np.arange(8, dtype=np.float32).reshape(2, 4) + object.__setattr__(art, "_joint_effort_target_backend", wp.zeros_like(art.data.joint_effort_target.warp)) + object.__setattr__(art, "_physx_actuator_wrapper", MagicMock(joint_f_2d=wp.array(forces, dtype=wp.float32))) + object.__setattr__(art, "_has_newton_actuators", True) + object.__setattr__(art, "_has_implicit_actuators", False) + written = [] + raw_backend.set_dof_actuation_forces = lambda values, indices: written.append(values.numpy().copy()) + + art.write_data_to_sim() + + np.testing.assert_allclose(written[-1], forces[:, _backend_to_user(tuple(art.joint_names), _names("joint", 4))]) + + +@requires("physx") +def test_physx_validate_cfg_reports_velocity_limits_in_public_joint_order(monkeypatch): + """Pair public default velocities with limits for the same named joint.""" + art, raw_backend = get_articulation( + "physx", 1, 3, 2, device="cpu", joint_ordering=("joint_2", "joint_1", "joint_0"), monkeypatch=monkeypatch + ) + backend_velocity_limits = np.asarray([[1.0, 100.0, 100.0]], dtype=np.float32) + raw_backend.set_mock_dof_max_velocities(wp.array(backend_velocity_limits, dtype=wp.float32)) + art.data._joint_vel_limits.assign(wp.array(backend_velocity_limits[:, ::-1].copy(), dtype=wp.float32)) + art.data._joint_pos_limits.assign(wp.array(np.tile((-100.0, 100.0), (1, 3, 1)), dtype=wp.vec2f)) + art.data._default_joint_pos.zero_() + art.data._default_joint_vel.assign(wp.array([[50.0, 50.0, 2.0]], dtype=wp.float32)) + + with pytest.raises(ValueError, match=r"'joint_0': 2\.000 not in \[-1\.000, 1\.000\]"): + art._validate_cfg() diff --git a/source/isaaclab/test/assets/contract/test_rigid_object_contract.py b/source/isaaclab/test/assets/contract/test_rigid_object_contract.py new file mode 100644 index 000000000000..07922735d759 --- /dev/null +++ b/source/isaaclab/test/assets/contract/test_rigid_object_contract.py @@ -0,0 +1,451 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# ignore private usage of variables warning +# pyright: reportPrivateUsage=none + +"""Shared rigid-object and rigid-object-collection contracts across the production backends. + +Both asset families run on mocked backend views, so these cases need neither Isaac Sim nor a GPU simulation. A rigid +object exposes root quantities of shape ``(N,)`` and body quantities of shape ``(N, 1)``; a collection exposes body +quantities of shape ``(N, B)``. Bookkeeping runs on CPU; getters and writers run on every test device because PhysX +stages writes through pinned CPU buffers on CUDA. +""" + +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +import torch +import warp as wp + +from isaaclab.assets import BaseRigidObjectCollectionData, BaseRigidObjectData +from isaaclab.utils.warp import ProxyArray + +from .rigid_object_factory import get_rigid_object, get_rigid_object_collection + +pytestmark = pytest.mark.integration + +# Distinct instance and body counts make swapped axes visible. +_NUM_INSTANCES, _NUM_BODIES = 2, 3 +_ASSETS = ["object", "collection"] + + +def _create(asset: str, backend: str, monkeypatch: pytest.MonkeyPatch, device: str = "cpu", num_instances: int = 2): + """Create a mocked asset and return it with its raw backend view, its top-level prefix, and its body count.""" + if asset == "object": + obj, raw = get_rigid_object(backend, num_instances, device, monkeypatch=monkeypatch) + return obj, raw, "root", 1 + obj, raw = get_rigid_object_collection(backend, num_instances, _NUM_BODIES, device, monkeypatch=monkeypatch) + return obj, raw, "body", _NUM_BODIES + + +def _payload(shape: tuple[int, ...], trailing: tuple[int, ...], offset: float = 0.0, device: str = "cpu"): + """Return per-element distinct values; 7-wide transforms carry a 90-degree rotation about Z.""" + data = torch.arange(np.prod(shape + trailing), dtype=torch.float32, device=device).reshape(shape + trailing) + data += 1.25 + offset + if trailing == (7,): + data[..., 3:5] = 0.0 + data[..., 5:7] = 2.0**-0.5 + return data + + +def _select_last(values: torch.Tensor, selection: str, per_body: bool) -> tuple[torch.Tensor, dict]: + """Return the payload and selectors that write the last environment (and last body) of ``values``. + + Selecting the last rather than the first element keeps identity or transposed routing from passing. + """ + device = values.device + if selection == "index": + selected = values[-1:, -1:] if per_body else values[-1:] + selectors = {"env_ids": torch.tensor([values.shape[0] - 1], dtype=torch.int32, device=device)} + if per_body: + selectors["body_ids"] = torch.tensor([values.shape[1] - 1], dtype=torch.int32, device=device) + return selected, selectors + masks = [wp.array(np.arange(n) == n - 1, dtype=wp.bool, device=str(device)) for n in values.shape[:2]] + selectors = {"env_mask": masks[0]} + if per_body: + selectors["body_mask"] = masks[1] + return values, selectors + + +# --------------------------------------------------------------------------- +# Bookkeeping, finders, and data properties +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("asset", _ASSETS) +def test_counts_names_finders_and_index_resolution(monkeypatch, backend, asset): + obj, _, _, num_bodies = _create(asset, backend, monkeypatch, num_instances=4) + + assert obj.num_instances == 4 + assert obj.num_bodies == num_bodies + assert len(obj.body_names) == num_bodies and all(isinstance(name, str) for name in obj.body_names) + assert isinstance(obj.data, BaseRigidObjectData if asset == "object" else BaseRigidObjectCollectionData) + + # The legacy return mode is a list for rigid objects and an int32 tensor for collections. + indices, names = obj.find_bodies(".*") + proxy, proxy_names = obj.find_bodies(".*", as_proxy=True) + assert isinstance(indices, list if asset == "object" else torch.Tensor) + assert list(indices) == proxy.torch.tolist() == list(range(num_bodies)) + assert names == proxy_names == obj.body_names + assert proxy is obj.find_bodies(".*", as_proxy=True)[0] + assert proxy.dtype == wp.int32 and str(proxy.device) == obj.device + assert obj.find_bodies(obj.body_names[-1])[1] == [obj.body_names[-1]] + if asset == "collection": + with pytest.warns(DeprecationWarning): + assert obj.find_objects(".*", as_proxy=True)[0] is proxy + assert obj._resolve_body_ids(torch.arange(3, dtype=torch.int32)[:2]).shape[0] == 2 + + # Environment selections resolve to strided views of the cached indices without copying. + env_ids = torch.arange(4, dtype=torch.int32) + assert obj._resolve_env_ids(env_ids[:2]).shape[0] == 2 + cached = wp.to_torch(obj._ALL_INDICES if asset == "object" else obj._ALL_ENV_INDICES) + for selection in (slice(None), slice(1, None, 2), slice(0, 0)): + resolved = wp.to_torch(obj._resolve_env_ids(selection)) + torch.testing.assert_close(resolved, cached[selection]) + assert resolved.data_ptr() == cached[selection].data_ptr() + assert resolved.stride() == cached[selection].stride() + + +def _expected_properties(top: str) -> dict[str, type]: + """Return the public data properties and their dtypes for an asset whose top-level prefix is ``top``.""" + frame_quantities = { + "pose_w": wp.transformf, + "vel_w": wp.spatial_vectorf, + "pos_w": wp.vec3f, + "quat_w": wp.quatf, + "lin_vel_w": wp.vec3f, + "ang_vel_w": wp.vec3f, + } + properties = { + f"{prefix}_{frame}_{quantity}": dtype + for prefix in {top, "body"} + for frame in ("link", "com") + for quantity, dtype in frame_quantities.items() + } + properties |= {f"{top}_{frame}_{axis}_vel_b": wp.vec3f for frame in ("link", "com") for axis in ("lin", "ang")} + return properties | { + "projected_gravity_b": wp.vec3f, + "heading_w": wp.float32, + "body_com_acc_w": wp.spatial_vectorf, + "body_com_lin_acc_w": wp.vec3f, + "body_com_ang_acc_w": wp.vec3f, + "body_com_pose_b": wp.transformf, + "body_com_pos_b": wp.vec3f, + "body_com_quat_b": wp.quatf, + "body_mass": wp.float32, + "body_inertia": wp.float32, + f"default_{top}_pose": wp.transformf, + f"default_{top}_vel": wp.spatial_vectorf, + } + + +def _expected_aliases(top: str) -> dict[str, str]: + """Return the shorthand data properties and the canonical properties they alias.""" + aliases = { + "body_acc_w": "body_com_acc_w", + "body_lin_acc_w": "body_com_lin_acc_w", + "body_ang_acc_w": "body_com_ang_acc_w", + "com_pos_b": "body_com_pos_b", + "com_quat_b": "body_com_quat_b", + } + for prefix in {top, "body"}: + aliases |= {f"{prefix}_{quantity}_w": f"{prefix}_link_{quantity}_w" for quantity in ("pose", "pos", "quat")} + aliases |= { + f"{prefix}_{quantity}_w": f"{prefix}_com_{quantity}_w" for quantity in ("vel", "lin_vel", "ang_vel") + } + return aliases + + +@pytest.mark.parametrize("asset", _ASSETS) +def test_data_properties_shapes_aliases_and_views(monkeypatch, backend, device, asset): + obj, _, top, num_bodies = _create(asset, backend, monkeypatch, device) + data = obj.data + data.update(dt=0.01) + if backend == "newton" and asset == "object": + # Derived quantities stay unallocated until first read. + for name in ("root_link_vel_w", "root_com_pose_w", "body_com_pose_b", "projected_gravity_b", "heading_w"): + assert getattr(data, "_" + name).data is None + + for name, dtype in _expected_properties(top).items(): + value = getattr(data, name) + per_instance = asset == "object" and not name.startswith("body") + shape = (_NUM_INSTANCES,) if per_instance else (_NUM_INSTANCES, num_bodies) + assert isinstance(value, ProxyArray), name + assert (value.shape, value.dtype) == (shape + ((9,) if name == "body_inertia" else ()), dtype), name + + # Random mock state makes link and COM quantities differ, so a retargeted alias fails. + for alias, canonical in _expected_aliases(top).items(): + alias_value, canonical_value = getattr(data, alias), getattr(data, canonical) + assert (alias_value.shape, alias_value.dtype) == (canonical_value.shape, canonical_value.dtype), alias + assert torch.equal(alias_value.torch, canonical_value.torch), alias + + # Position, orientation, and velocity components are zero-copy views of the packed pose and velocity. + for frame in {f"{top}_link", f"{top}_com", "body_link", "body_com"}: + for quantity, components in (("pose", ("pos", "quat")), ("vel", ("lin_vel", "ang_vel"))): + packed = getattr(data, f"{frame}_{quantity}_w").torch + for component, expected in zip(components, (packed[..., :3], packed[..., 3:]), strict=True): + view = getattr(data, f"{frame}_{component}_w").torch + torch.testing.assert_close(view, expected) + assert (view.data_ptr(), view.stride()) == (expected.data_ptr(), expected.stride()), component + + +@pytest.mark.parametrize("trigger", ["pose", "velocity", "coms_index", "coms_mask"]) +@pytest.mark.parametrize("asset", _ASSETS) +def test_writes_invalidate_dependent_caches(monkeypatch, backend, asset, trigger): + obj, _, top, num_bodies = _create(asset, backend, monkeypatch) + data = obj.data + data.update(dt=0.01) + body_frame_velocities = [f"{top}_{frame}_{axis}_vel_b" for frame in ("link", "com") for axis in ("lin", "ang")] + stale = { + "pose": [f"{top}_link_vel_w", "projected_gravity_b", "heading_w", *body_frame_velocities], + "velocity": body_frame_velocities, + }.get(trigger) + if stale is None: + states = [f"{top}_state_w", f"{top}_link_state_w", f"{top}_com_state_w"] + stale = [f"{top}_com_pose_w", f"{top}_link_vel_w", *body_frame_velocities, *states] + stale.append("body_com_pose_b" if backend == "newton" else f"{top}_com_vel_w") + # Read each property so lazy caches allocate, then mark it current. + buffers = {} + for name in stale: + getattr(data, name) + buffers[name] = getattr(data, "_" + name) + buffers[name].timestamp = data._sim_timestamp + body_velocity = data.body_link_vel_w + + shape = (_NUM_INSTANCES,) if asset == "object" else (_NUM_INSTANCES, num_bodies) + pose_kwarg, velocity_kwarg = ( + ("root_pose", "root_velocity") if asset == "object" else ("body_poses", "body_velocities") + ) + if trigger == "pose": + getattr(obj, f"write_{top}_link_pose_to_sim_index")(**{pose_kwarg: _payload(shape, (7,))}) + elif trigger == "velocity": + # Zero angular velocity makes link and COM velocities equal. + velocity = torch.zeros(shape + (6,)) + velocity[..., :3] = _payload(shape, (3,)) + getattr(obj, f"write_{top}_com_velocity_to_sim_index")(**{velocity_kwarg: velocity}) + body_velocity_expected = velocity if asset == "collection" else velocity.unsqueeze(1) + torch.testing.assert_close(data.body_link_vel_w.torch, body_velocity_expected) + assert data.body_link_vel_w is body_velocity + else: + layout = (3,) if backend == "newton" else (7,) + coms = _payload((_NUM_INSTANCES, num_bodies), layout) + if backend == "newton": + from isaaclab_newton.physics import NewtonManager + + monkeypatch.setattr(NewtonManager, "add_model_change", MagicMock()) + getattr(obj, f"set_{trigger}")(coms=coms) + + for name, buffer in buffers.items(): + assert buffer.timestamp < data._sim_timestamp, name + + +# --------------------------------------------------------------------------- +# Writers +# --------------------------------------------------------------------------- + +_POSE, _VELOCITY = ((7,), wp.transformf), ((6,), wp.spatial_vectorf) +# written quantity -> (read-back getter with ``{top}`` as ``root`` or ``body``, (per-element shape, warp dtype)) +_WRITERS = { + "link_pose": ("{top}_link_pose_w", _POSE), + "com_pose": ("{top}_com_pose_w", _POSE), + "link_velocity": ("{top}_link_vel_w", _VELOCITY), + "com_velocity": ("{top}_com_vel_w", _VELOCITY), + "masses": ("body_mass", ((), wp.float32)), + "inertias": ("body_inertia", ((9,), wp.float32)), + "coms": ("body_com_pose_b", _POSE), +} + + +def _writer(obj, asset: str, quantity: str, selection: str): + """Return the index or mask writer of a quantity and its payload keyword.""" + if quantity in ("masses", "inertias", "coms"): + return getattr(obj, f"set_{quantity}_{selection}"), quantity + top = "root" if asset == "object" else "body" + kind = "pose" if quantity.endswith("pose") else "velocity" + kwarg = f"root_{kind}" if asset == "object" else {"pose": "body_poses", "velocity": "body_velocities"}[kind] + return getattr(obj, f"write_{top}_{quantity}_to_sim_{selection}"), kwarg + + +@pytest.mark.parametrize("selection", ["index", "mask"]) +@pytest.mark.parametrize("quantity", _WRITERS) +@pytest.mark.parametrize("asset", _ASSETS) +def test_writer_reads_back_full_and_partial_writes(monkeypatch, backend, device, asset, quantity, selection): + obj, _, top, num_bodies = _create(asset, backend, monkeypatch, device) + obj.data.update(dt=0.01) + getter, (trailing, wp_dtype) = _WRITERS[quantity] + if quantity == "coms" and backend == "newton": + # Newton stores the center of mass as a position only. + getter, (trailing, wp_dtype) = "body_com_pos_b", ((3,), wp.vec3f) + getter = getter.format(top=top) + method, kwarg = _writer(obj, asset, quantity, selection) + per_body = asset == "collection" or quantity in ("masses", "inertias", "coms") + shape = (_NUM_INSTANCES, num_bodies) if per_body else (_NUM_INSTANCES,) + + # A torch write of every element reads back through the matching getter. + full = _payload(shape, trailing, device=device) + method(**{kwarg: full}) + torch.testing.assert_close(getattr(obj.data, getter).torch, full, atol=1e-5, rtol=1e-5) + + # A warp write selecting the last environment (and last body) leaves every other element untouched. + update = _payload(shape, trailing, offset=100.0, device=device) + selected, selectors = _select_last(update, selection, per_body) + method(**{kwarg: wp.from_torch(selected.contiguous(), dtype=wp_dtype)}, **selectors) + expected = full.clone() + last = (-1, -1) if per_body else -1 + expected[last] = update[last] + torch.testing.assert_close(getattr(obj.data, getter).torch, expected, atol=1e-5, rtol=1e-5) + + # Data with an extra environment is rejected for torch and warp inputs. + bad = _payload((_NUM_INSTANCES + 1, *shape[1:]), trailing, device=device) + for value in (bad, wp.from_torch(bad.contiguous(), dtype=wp_dtype)): + with pytest.raises((AssertionError, RuntimeError)): + method(**{kwarg: value}) + + +@pytest.mark.parametrize("asset", _ASSETS) +def test_pose_and_velocity_writer_aliases_forward_to_canonical_writers(monkeypatch, backend, asset): + obj, _, top, num_bodies = _create(asset, backend, monkeypatch) + shape = (_NUM_INSTANCES,) if asset == "object" else (_NUM_INSTANCES, num_bodies) + for quantity, getter, trailing in (("pose", f"{top}_link_pose_w", (7,)), ("velocity", f"{top}_com_vel_w", (6,))): + for offset, selection in enumerate(("index", "mask")): + values = _payload(shape, trailing, offset=10.0 * offset) + method, kwarg = _writer(obj, asset, quantity, selection) + method(**{kwarg: values}) + torch.testing.assert_close(getattr(obj.data, getter).torch, values, atol=1e-5, rtol=1e-5) + + +# backend-stored quantity -> (written quantity, per-element shape) +_PARTIAL_WRITES = {"pose": ("link_pose", (7,)), "velocity": ("com_velocity", (6,)), "mass": ("masses", ())} + + +def _read_backend(backend: str, asset: str, raw, quantity: str, num_bodies: int) -> torch.Tensor: + """Read one quantity from backend storage as an ``(env, body, -1)`` tensor.""" + if backend == "physx": + values = {"pose": raw.get_transforms, "velocity": raw.get_velocities, "mass": raw.get_masses}[quantity]() + values = wp.to_torch(values) + if asset == "collection": + # PhysX collection views are body-major. + values = values.reshape(num_bodies, _NUM_INSTANCES, -1).transpose(0, 1) + elif backend == "newton": + if quantity == "mass": + values = wp.to_torch(raw.get_attribute("body_mass", None)) + else: + values = wp.to_torch({"pose": raw.get_root_transforms, "velocity": raw.get_root_velocities}[quantity](None)) + else: + from isaaclab_ov import tensor_types as TT + + bindings = { + "object": {"pose": TT.RIGID_BODY_POSE, "velocity": TT.RIGID_BODY_VELOCITY, "mass": TT.RIGID_BODY_MASS}, + "collection": {"pose": TT.LINK_POSE, "velocity": TT.LINK_VELOCITY, "mass": TT.BODY_MASS}, + } + values = torch.as_tensor(raw.bindings[bindings[asset][quantity]]._data) + return values.reshape(_NUM_INSTANCES, num_bodies, -1).cpu().clone() + + +@pytest.mark.parametrize("selection", ["index", "mask"]) +@pytest.mark.parametrize("quantity", _PARTIAL_WRITES) +@pytest.mark.parametrize("asset", _ASSETS) +def test_partial_write_preserves_unselected_backend_elements(monkeypatch, request, backend, asset, quantity, selection): + if backend == "ovphysx" and asset == "collection" and quantity != "mass": + # Product bug: OVPhysX pushes whole environment rows from a pose/velocity buffer that a partial write does + # not refresh first, so the unselected bodies of the selected environment receive stale (here zero) state. + request.applymarker( + pytest.mark.xfail( + raises=AssertionError, + strict=True, + reason="OVPhysX partial body writes overwrite unselected bodies of the selected environment", + ) + ) + obj, raw, _, num_bodies = _create(asset, backend, monkeypatch) + written, trailing = _PARTIAL_WRITES[quantity] + per_body = asset == "collection" or quantity == "mass" + shape = (_NUM_INSTANCES, num_bodies) if per_body else (_NUM_INSTANCES,) + values = _payload(shape, trailing, offset=100.0) + before = _read_backend(backend, asset, raw, quantity, num_bodies) + + selected, selectors = _select_last(values, selection, per_body) + method, kwarg = _writer(obj, asset, written, selection) + method(**{kwarg: selected}, **selectors) + + expected = before.clone() + expected[-1, -1] = values.reshape(_NUM_INSTANCES, num_bodies, -1)[-1, -1] + torch.testing.assert_close(_read_backend(backend, asset, raw, quantity, num_bodies), expected, rtol=0, atol=0) + + +# --------------------------------------------------------------------------- +# Family-specific behavior +# --------------------------------------------------------------------------- + + +def test_rigid_object_external_wrench_frames(monkeypatch, backend): + """Forward local and world wrenches through the real writer in each backend's frame.""" + obj, raw, _, _ = _create("object", backend, monkeypatch) + # A known 90-degree rotation about Z keeps the local-frame expectation independent of production. + root_pose = torch.tensor( + [[1.0, 2.0, 3.0, 0.0, 0.0, 2.0**-0.5, 2.0**-0.5], [4.0, 5.0, 6.0, 0.0, 0.0, 2.0**-0.5, 2.0**-0.5]] + ) + obj.write_root_link_pose_to_sim_index(root_pose=root_pose) + if backend == "newton": + # Seed the body pose the native forward kinematics would publish. + obj.data._sim_bind_body_link_pose_w.assign(wp.from_torch(root_pose, dtype=wp.transformf)) + if backend == "physx": + raw.apply_forces_and_torques_at_position = MagicMock() + composer = obj.permanent_wrench_composer + forces = torch.arange(1.0, 7.0).reshape(2, 1, 3) + torques = forces + 10.0 + + for is_global in (False, True): + composer.reset() + composer.set_forces_and_torques_index(forces=forces, torques=torques, is_global=is_global) + with patch.object(composer, "compose_to_body_frame", wraps=composer.compose_to_body_frame) as compose: + obj.write_data_to_sim() + assert compose.call_count == int(is_global and backend == "newton") + + expected_force, expected_torque = forces, torques + if backend == "physx": + call = raw.apply_forces_and_torques_at_position.call_args.kwargs + assert call["is_global"] is is_global + assert call["position_data"] is None + actual_force, actual_torque = call["force_data"].numpy(), call["torque_data"].numpy() + else: + if not is_global: + # A 90-degree Z rotation maps body (x, y, z) to world (-y, x, z). + expected_force = torch.stack((-forces[..., 1], forces[..., 0], forces[..., 2]), dim=-1) + expected_torque = torch.stack((-torques[..., 1], torques[..., 0], torques[..., 2]), dim=-1) + if backend == "newton": + packed = obj.data._sim_bind_body_external_wrench.numpy() + else: + from isaaclab_ov import tensor_types as TT + + packed = raw.bindings[TT.RIGID_BODY_WRENCH]._data.reshape(2, 1, 9) + np.testing.assert_allclose(packed[..., 6:9], root_pose[:, None, :3].numpy()) + actual_force, actual_torque = packed[..., :3], packed[..., 3:6] + np.testing.assert_allclose(actual_force.reshape(2, 1, 3), expected_force.numpy(), atol=1e-5, rtol=1e-5) + np.testing.assert_allclose(actual_torque.reshape(2, 1, 3), expected_torque.numpy(), atol=1e-5, rtol=1e-5) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_collection_reshape_data_to_view_3d(monkeypatch, backend, device, dtype): + """Environment-major data becomes body-major view data, keeping the input dtype and array type.""" + obj, _, _, num_bodies = _create("collection", backend, monkeypatch, device) + data = torch.arange(_NUM_INSTANCES * num_bodies * 4, dtype=dtype, device=device).reshape( + _NUM_INSTANCES, num_bodies, 4 + ) + expected = data.permute(1, 0, 2).reshape(num_bodies * _NUM_INSTANCES, 4) + + view = obj.reshape_data_to_view_3d(data, 4, device=device) + assert isinstance(view, torch.Tensor) and view.is_contiguous() + assert (view.dtype, view.device) == (dtype, data.device) + torch.testing.assert_close(view, expected) + + if dtype == torch.float32: + warp_view = obj.reshape_data_to_view_3d(wp.from_torch(data, dtype=wp.float32), 4, device=device) + assert isinstance(warp_view, wp.array) and warp_view.dtype == wp.float32 and str(warp_view.device) == device + torch.testing.assert_close(wp.to_torch(warp_view), expected) + if device != "cpu": + # Torch inputs move to the requested device. + torch.testing.assert_close(obj.reshape_data_to_view_3d(data, 4, device="cpu"), expected.cpu()) diff --git a/source/isaaclab/test/assets/test_articulation_iface.py b/source/isaaclab/test/assets/test_articulation_iface.py deleted file mode 100644 index 7c8b169e2249..000000000000 --- a/source/isaaclab/test/assets/test_articulation_iface.py +++ /dev/null @@ -1,1417 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -# ignore private usage of variables warning -# pyright: reportPrivateUsage=none - -""" -Checks that the articulation interfaces are consistent across backends, and are providing the exact same data as what -the base articulation class advertises. All articulation interfaces need to comply with the same interface contract. - -The setup is a bit convoluted so that we can run these tests without requiring Isaac Sim or GPU simulation. -""" - -import math -import warnings -from unittest.mock import MagicMock - -import numpy as np -import pytest -import torch -import warp as wp -from _articulation_iface_test_utils import BACKEND_UNAVAILABLE_REASONS, BACKENDS, get_articulation - -from isaaclab.test.utils import DeviceScope, test_devices - -pytestmark = pytest.mark.integration - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _check_proxy_array(arr, *, expected_shape: tuple, expected_dtype: type, name: str): - """Assert that `arr` is a ProxyArray with the expected shape and dtype.""" - from isaaclab.utils.warp import ProxyArray - - assert isinstance(arr, ProxyArray), f"{name}: expected ProxyArray, got {type(arr)}" - assert arr.shape == expected_shape, f"{name}: expected shape {expected_shape}, got {arr.shape}" - assert arr.dtype == expected_dtype, f"{name}: expected dtype {expected_dtype}, got {arr.dtype}" - - -# Common parametrize decorators. Pure bookkeeping (counts, names, finders, aliases) runs on CPU only; -# getters and writers keep every test device because PhysX stages through CPU-pinned buffers on CUDA. -_backends = pytest.mark.parametrize("backend", BACKENDS, indirect=False) -_devices = pytest.mark.parametrize("device", test_devices(DeviceScope.CPU_AND_DEFAULT_CUDA)) -# One fixture with distinct instance, joint, and body counts so any swapped axis shows up. -_NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES = 2, 6, 7 -_index_resolution_backends = pytest.mark.parametrize( - "backend", [backend for backend in ("physx", "newton") if backend in BACKENDS], indirect=False -) -_production_backends = pytest.mark.parametrize( - "backend", [backend for backend in ("physx", "newton", "ovphysx") if backend in BACKENDS], indirect=False -) - - -# --------------------------------------------------------------------------- -# Tests: Index resolution helpers -# --------------------------------------------------------------------------- - - -class TestArticulationIndexResolution: - """Test backend-specific index resolution helpers.""" - - @_production_backends - @_devices - def test_resolve_env_ids_handles_tensor_view_shape(self, backend, device): - art, _ = get_articulation(backend, num_instances=4, device=device) - - env_ids = torch.arange(4, dtype=torch.int32, device=device) - resolved_full = art._resolve_env_ids(env_ids) - resolved_view = art._resolve_env_ids(env_ids[:2]) - - assert resolved_full.shape[0] == 4 - assert resolved_view.shape[0] == 2 - # Native indexed writers can use a strided view of cached IDs without copying or uploading them. - cached = wp.to_torch(art._ALL_INDICES) - for selection in (slice(None), slice(1, None, 2), slice(0, 0)): - resolved = wp.to_torch(art._resolve_env_ids(selection)) - torch.testing.assert_close(resolved, cached[selection]) - assert resolved.data_ptr() == cached[selection].data_ptr() - assert resolved.stride() == cached[selection].stride() - - @_index_resolution_backends - def test_resolve_joint_ids_handles_tensor_view_shape(self, backend): - art, _ = get_articulation(backend, num_joints=4, device="cpu") - - joint_ids = torch.arange(4, dtype=torch.int32, device="cpu") - resolved_full = art._resolve_joint_ids(joint_ids) - resolved_view = art._resolve_joint_ids(joint_ids[:2]) - - assert resolved_full.shape[0] == 4 - assert resolved_view.shape[0] == 2 - - @_index_resolution_backends - def test_resolve_body_ids_handles_tensor_view_shape(self, backend): - art, _ = get_articulation(backend, num_bodies=4, device="cpu") - - body_ids = torch.arange(4, dtype=torch.int32, device="cpu") - resolved_full = art._resolve_body_ids(body_ids) - resolved_view = art._resolve_body_ids(body_ids[:2]) - - assert resolved_full.shape[0] == 4 - assert resolved_view.shape[0] == 2 - - -# --------------------------------------------------------------------------- -# Tests: Articulation properties and finders -# --------------------------------------------------------------------------- - - -class TestArticulationProperties: - """Test that articulation properties return the correct types/values.""" - - @_backends - @pytest.mark.parametrize("is_fixed_base", [False, True], ids=["floating", "fixed"]) - def test_articulation_counts_names_and_finders(self, backend, is_fixed_base): - from isaaclab.assets.articulation.base_articulation_data import BaseArticulationData - - art, _ = get_articulation( - backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device="cpu", is_fixed_base=is_fixed_base - ) - - assert art.num_instances == _NUM_INSTANCES - assert art.num_joints == _NUM_JOINTS - assert art.num_bodies == _NUM_BODIES - assert art.is_fixed_base is is_fixed_base - for names, count in ((art.joint_names, _NUM_JOINTS), (art.body_names, _NUM_BODIES)): - assert isinstance(names, list) - assert len(names) == count - assert all(isinstance(n, str) for n in names) - assert isinstance(art.data, BaseArticulationData) - - -class TestArticulationFinderReturnModes: - """Test finder return modes on production articulation backends.""" - - @_production_backends - @pytest.mark.parametrize( - "finder_name, names_attr, count", - [ - ("find_bodies", "body_names", 4), - ("find_joints", "joint_names", 3), - ("find_fixed_tendons", "fixed_tendon_names", 2), - ("find_spatial_tendons", "spatial_tendon_names", 5), - ], - ) - def test_finder_returns_legacy_list_or_cached_proxy(self, backend, finder_name, names_attr, count): - art, _ = get_articulation( - backend, - num_instances=2, - num_joints=3, - num_bodies=4, - num_fixed_tendons=2, - num_spatial_tendons=5, - device="cpu", - ) - if backend == "newton" and finder_name == "find_spatial_tendons": - pytest.skip("Newton does not support spatial tendons.") - finder = getattr(art, finder_name) - - indices, names = finder(".*") - proxy, proxy_names = finder(".*", as_proxy=True) - - assert isinstance(indices, list) and isinstance(names, list) - assert len(indices) == count - assert len(names) == count - assert all(isinstance(i, int) for i in indices) - assert all(isinstance(n, str) for n in names) - assert indices == proxy.torch.tolist() - assert names == proxy_names - assert proxy is finder(".*", as_proxy=True)[0] - assert proxy.dtype == wp.int32 - assert str(proxy.device) == art.device - - first_name = getattr(art, names_attr)[0] - assert finder(first_name) == ([0], [first_name]) - - -class TestFixedTendonTargetScheduling: - """Commanding a tendon target is what schedules the write, not merely having tendons.""" - - @staticmethod - def _articulation(backend): - art, _ = get_articulation( - backend, - num_instances=2, - num_joints=2, - num_bodies=2, - num_fixed_tendons=2, - num_spatial_tendons=0, - device="cpu", - ) - return art - - @_production_backends - def test_commanding_a_target_schedules_the_write(self, backend): - """The target setter marks the write pending; a static offset write keeps its own contract.""" - if backend == "newton": - pytest.skip("Newton needs a MuJoCo tendon actuator, covered by the delegation tests below.") - art = self._articulation(backend) - art.set_fixed_tendon_position_target_index(target=torch.zeros((2, 2), dtype=torch.float32)) - assert art._fixed_tendon_target_dirty is True - - def test_newton_reports_a_missing_tendon_actuator(self): - """Without a MuJoCo tendon actuator the solver has no adapter, so commanding must say so.""" - if "newton" not in BACKENDS: - pytest.skip(BACKEND_UNAVAILABLE_REASONS.get("newton", "newton backend unavailable")) - art = self._articulation("newton") - with pytest.raises(RuntimeError, match="no MuJoCo tendon actuator"): - art.set_fixed_tendon_position_target_index(target=torch.zeros((2, 2), dtype=torch.float32)) - - def test_a_solver_without_tendon_transmission_refuses_to_build_an_adapter(self): - """A solver with no tendon transmission says so, rather than returning nothing.""" - if "newton" not in BACKENDS: - pytest.skip(BACKEND_UNAVAILABLE_REASONS.get("newton", "newton backend unavailable")) - from isaaclab_newton.physics import NewtonManager - - art = self._articulation("newton") - with pytest.raises(NotImplementedError, match="does not drive fixed tendons"): - NewtonManager.create_fixed_tendon_control(art) - - -# --------------------------------------------------------------------------- -# Tests: resolve_matching_names caching behavior -# --------------------------------------------------------------------------- - - -class TestResolveMatchingNamesCache: - """Test that resolve_matching_names caching returns correct, isolated results.""" - - @_backends - def test_unmatched_regex_raises(self, backend): - """ValueError from resolve_matching_names propagates correctly.""" - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device="cpu") - with pytest.raises(ValueError): - art.find_bodies("nonexistent_body_xyz") - with pytest.raises(ValueError): - art.find_joints("nonexistent_joint_xyz") - - @_backends - def test_mutating_result_does_not_corrupt_cache(self, backend): - """Mutating returned lists must not affect future cached results.""" - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device="cpu") - - for finder, expected_len in [("find_bodies", _NUM_BODIES), ("find_joints", _NUM_JOINTS)]: - idx1, names1 = getattr(art, finder)(".*") - assert len(idx1) == expected_len - - idx1.clear() - names1.append("corrupted") - - idx2, names2 = getattr(art, finder)(".*") - assert len(idx2) == expected_len - assert "corrupted" not in names2 - - @_backends - def test_find_with_preserve_order(self, backend): - """A list of patterns resolves each name; preserve_order=True keeps the pattern order.""" - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device="cpu") - idx_fwd, names_fwd = art.find_joints(["joint_1", "joint_0"], preserve_order=True) - assert names_fwd == ["joint_1", "joint_0"] - - idx_rev, names_rev = art.find_joints(["joint_0", "joint_1"], preserve_order=True) - assert names_rev == ["joint_0", "joint_1"] - - -# --------------------------------------------------------------------------- -# Tests: ArticulationData property contract -# --------------------------------------------------------------------------- - -# (property, shape kind, dtype). Shape kinds: "N" = (num_instances,), "NB" = (num_instances, num_bodies), -# "NB9" = (num_instances, num_bodies, 9), "NJ" = (num_instances, num_joints). -_ARTICULATION_DATA_PROPERTIES = [ - # root state - ("root_link_pose_w", "N", wp.transformf), - ("root_link_vel_w", "N", wp.spatial_vectorf), - ("root_com_pose_w", "N", wp.transformf), - ("root_com_vel_w", "N", wp.spatial_vectorf), - ("root_link_pos_w", "N", wp.vec3f), - ("root_link_quat_w", "N", wp.quatf), - ("root_link_lin_vel_w", "N", wp.vec3f), - ("root_link_ang_vel_w", "N", wp.vec3f), - ("root_com_pos_w", "N", wp.vec3f), - ("root_com_quat_w", "N", wp.quatf), - ("root_com_lin_vel_w", "N", wp.vec3f), - ("root_com_ang_vel_w", "N", wp.vec3f), - # derived - ("projected_gravity_b", "N", wp.vec3f), - ("heading_w", "N", wp.float32), - ("root_link_lin_vel_b", "N", wp.vec3f), - ("root_link_ang_vel_b", "N", wp.vec3f), - ("root_com_lin_vel_b", "N", wp.vec3f), - ("root_com_ang_vel_b", "N", wp.vec3f), - # body state - ("body_link_pose_w", "NB", wp.transformf), - ("body_link_vel_w", "NB", wp.spatial_vectorf), - ("body_com_pose_w", "NB", wp.transformf), - ("body_com_vel_w", "NB", wp.spatial_vectorf), - ("body_com_acc_w", "NB", wp.spatial_vectorf), - ("body_com_pose_b", "NB", wp.transformf), - ("body_mass", "NB", wp.float32), - ("body_inertia", "NB9", wp.float32), - ("body_link_pos_w", "NB", wp.vec3f), - ("body_link_quat_w", "NB", wp.quatf), - ("body_link_lin_vel_w", "NB", wp.vec3f), - ("body_link_ang_vel_w", "NB", wp.vec3f), - ("body_com_pos_w", "NB", wp.vec3f), - ("body_com_quat_w", "NB", wp.quatf), - ("body_com_pos_b", "NB", wp.vec3f), - ("body_com_quat_b", "NB", wp.quatf), - # joint state and properties - ("joint_pos", "NJ", wp.float32), - ("joint_vel", "NJ", wp.float32), - ("joint_acc", "NJ", wp.float32), - ("joint_stiffness", "NJ", wp.float32), - ("joint_damping", "NJ", wp.float32), - ("joint_armature", "NJ", wp.float32), - ("joint_friction_coeff", "NJ", wp.float32), - ("joint_pos_limits", "NJ", wp.vec2f), - ("joint_vel_limits", "NJ", wp.float32), - ("joint_effort_limits", "NJ", wp.float32), - ("soft_joint_pos_limits", "NJ", wp.vec2f), - # defaults and command targets - ("default_root_pose", "N", wp.transformf), - ("default_root_vel", "N", wp.spatial_vectorf), - ("default_joint_pos", "NJ", wp.float32), - ("default_joint_vel", "NJ", wp.float32), - ("joint_pos_target", "NJ", wp.float32), - ("joint_vel_target", "NJ", wp.float32), - ("joint_effort_target", "NJ", wp.float32), -] - - -class TestArticulationDataProperties: - """Test that every data property is a ProxyArray with the advertised shape and dtype.""" - - @_backends - @_devices - def test_articulation_data_property_contract(self, backend, device): - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device) - if backend == "newton": - # Native selections may stride over other articulations in each world. - for name in ("root_link_pose_w", "root_com_vel_w", "body_link_pose_w", "body_com_vel_w"): - binding = getattr(art.data, "_sim_bind_" + name) - storage = wp.empty((2 * binding.shape[0], *binding.shape[1:]), dtype=binding.dtype, device=device) - view = storage[::2] - view.assign(binding) - setattr(art.data, "_sim_bind_" + name, view) - art.data._pin_proxy_arrays() - art.data.update(dt=0.01) - if backend == "newton": - names = ("root_link_vel_w", "body_link_vel_w", "body_com_pose_b", "root_com_pose_w") - for name in (*names, "body_com_pose_w", "projected_gravity_b", "heading_w"): - assert getattr(art.data, "_" + name).data is None - shapes = { - "N": (_NUM_INSTANCES,), - "NB": (_NUM_INSTANCES, _NUM_BODIES), - "NB9": (_NUM_INSTANCES, _NUM_BODIES, 9), - "NJ": (_NUM_INSTANCES, _NUM_JOINTS), - } - for name, shape_kind, dtype in _ARTICULATION_DATA_PROPERTIES: - if backend == "newton" and name == "body_com_pose_b": - # Newton stores a position-only COM and warns that the pose appends a unit quaternion. - with pytest.warns(UserWarning, match="unit quaternion"): - value = getattr(art.data, name) - else: - value = getattr(art.data, name) - _check_proxy_array(value, expected_shape=shapes[shape_kind], expected_dtype=dtype, name=name) - - for frame in ("root_link", "root_com", "body_link", "body_com"): - for quantity, components in (("pose", ("pos", "quat")), ("vel", ("lin_vel", "ang_vel"))): - packed = getattr(art.data, f"{frame}_{quantity}_w").torch - for component, expected in zip(components, (packed[..., :3], packed[..., 3:]), strict=True): - view = getattr(art.data, f"{frame}_{component}_w").torch - torch.testing.assert_close(view, expected) - assert view.data_ptr() == expected.data_ptr() - assert view.stride() == expected.stride() - - @pytest.mark.skipif("physx" not in BACKENDS, reason="PhysX backend unavailable") - def test_physx_set_coms_index_updates_body_com_pose_b_cache(self): - art, view = get_articulation("physx", num_instances=2, num_joints=3, num_bodies=4, device="cpu") - - num_get_coms_calls = 0 - get_coms = view.get_coms - - def counted_get_coms(): - nonlocal num_get_coms_calls - num_get_coms_calls += 1 - return get_coms() - - view.get_coms = counted_get_coms - - coms = wp.zeros((art.num_instances, art.num_bodies), dtype=wp.transformf, device="cpu") - art.set_coms_index(coms=coms, full_data=True) - art.data.body_com_pose_b - - assert num_get_coms_calls == 0 - - @pytest.mark.skipif("physx" not in BACKENDS, reason="PhysX backend unavailable") - def test_physx_joint_position_write_preserves_body_com_pose_b_cache(self): - art, view = get_articulation("physx", num_instances=2, num_joints=3, num_bodies=4, device="cpu") - - num_get_coms_calls = 0 - get_coms = view.get_coms - - def counted_get_coms(): - nonlocal num_get_coms_calls - num_get_coms_calls += 1 - return get_coms() - - view.get_coms = counted_get_coms - - art.data.update(dt=0.01) - art.data.body_com_pose_b - assert num_get_coms_calls == 1 - - joint_pos = torch.zeros((art.num_instances, art.num_joints), device="cpu") - art.write_joint_position_to_sim_index(position=joint_pos, full_data=True) - art.data.body_com_pose_b - - assert num_get_coms_calls == 1 - - @pytest.mark.parametrize("backend", [name for name in ("physx", "ovphysx") if name in BACKENDS]) - @_devices - def test_dynamics_buffers_allocate_on_read_and_reuse_across_refreshes(self, backend, device): - """Optional dynamics allocate independently on demand and retain stable public views.""" - art, _ = get_articulation(backend, num_instances=2, num_joints=3, num_bodies=4, device=device) - names = ("body_com_jacobian_w", "mass_matrix", "gravity_compensation_forces") - art.data.update(dt=0.01) - for index, name in enumerate(names): - for unread in names[index:]: - assert getattr(art.data, "_" + unread).data is None - wrapper = getattr(art.data, name) - assert wrapper is getattr(art.data, "_" + name).data - art.data.update(dt=0.01) - assert getattr(art.data, name) is wrapper - assert art.data._body_link_jacobian_w is None - wrapper = art.data.body_link_jacobian_w - art.data.update(dt=0.01) - assert art.data.body_link_jacobian_w is wrapper - - @_production_backends - def test_actuator_compatibility_projections_are_stable(self, backend): - num_instances, num_joints = 2, 4 - art, _ = get_articulation(backend, num_instances, num_joints, 5, device="cpu") - soft_joint_vel_limits = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]], dtype=torch.float32) - wp.copy(art.actuators._soft_joint_vel_limits, wp.from_torch(soft_joint_vel_limits)) - - with warnings.catch_warnings(record=True) as caught_warnings: - warnings.simplefilter("always") - soft_joint_vel_limits_data = art.data.soft_joint_vel_limits - soft_joint_vel_limits_repeat = art.data.soft_joint_vel_limits - - _check_proxy_array( - soft_joint_vel_limits_data, - expected_shape=(num_instances, num_joints), - expected_dtype=wp.float32, - name="soft_joint_vel_limits", - ) - torch.testing.assert_close(soft_joint_vel_limits_data.torch, soft_joint_vel_limits, rtol=0.0, atol=0.0) - assert soft_joint_vel_limits_data.warp.ptr == soft_joint_vel_limits_repeat.warp.ptr - assert soft_joint_vel_limits_data.warp.ptr == art.actuators._soft_joint_vel_limits.ptr - assert not [warning for warning in caught_warnings if warning.category is DeprecationWarning] - - -# --------------------------------------------------------------------------- -# Tests: Alias/shorthand properties -# --------------------------------------------------------------------------- - -_ARTICULATION_ALIASES = [ - ("root_pose_w", "root_link_pose_w"), - ("root_pos_w", "root_link_pos_w"), - ("root_quat_w", "root_link_quat_w"), - ("root_vel_w", "root_com_vel_w"), - ("root_lin_vel_w", "root_com_lin_vel_w"), - ("root_ang_vel_w", "root_com_ang_vel_w"), - ("body_pose_w", "body_link_pose_w"), - ("body_pos_w", "body_link_pos_w"), - ("body_quat_w", "body_link_quat_w"), - ("body_vel_w", "body_com_vel_w"), - ("body_lin_vel_w", "body_com_lin_vel_w"), - ("body_ang_vel_w", "body_com_ang_vel_w"), - ("joint_limits", "joint_pos_limits"), - ("joint_friction", "joint_friction_coeff"), -] - - -class TestArticulationDataAliases: - """Test that alias properties return the values of their canonical counterparts.""" - - @_backends - def test_aliases_match_canonical_values(self, backend): - # Random mock state makes link and COM quantities differ, so a retargeted alias fails. - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device="cpu") - art.data.update(dt=0.01) - d = art.data - - for alias, canonical in _ARTICULATION_ALIASES: - alias_value, canonical_value = getattr(d, alias), getattr(d, canonical) - assert alias_value.shape == canonical_value.shape, alias - assert alias_value.dtype == canonical_value.dtype, alias - assert torch.equal(alias_value.torch, canonical_value.torch), alias - - -# --------------------------------------------------------------------------- -# Writer/setter test helpers -# --------------------------------------------------------------------------- - -# Map warp structured dtypes to their torch trailing dimension size. -_WP_DTYPE_TO_TRAILING = { - wp.transformf: 7, - wp.spatial_vectorf: 6, - wp.vec2f: 2, - wp.float32: 0, # no trailing dimension -} - - -def _make_data_torch(shape: tuple, device: str, wp_dtype=wp.float32) -> torch.Tensor: - """Create valid torch test data for a given warp dtype. - - For transformf shapes, appends a trailing dim of 7 and sets quat w=1. - For spatial_vectorf, appends trailing 6. - For vec2f, appends trailing 2 with [-1, 1]. - For float32, no trailing dim. - """ - trailing = _WP_DTYPE_TO_TRAILING[wp_dtype] - if trailing: - full_shape = (*shape, trailing) - else: - full_shape = shape - data = torch.zeros(full_shape, device=device, dtype=torch.float32) - if wp_dtype == wp.transformf: - data[..., 6] = 1.0 # identity quat w - elif wp_dtype == wp.vec2f: - data[..., 0] = -1.0 - data[..., 1] = 1.0 - elif wp_dtype == wp.float32: - data.fill_(1.0) - return data - - -def _make_data_warp(shape: tuple, device: str, wp_dtype=wp.float32) -> wp.array: - """Create valid warp test data for a given warp dtype. - - Warp structured types collapse the trailing dim into the dtype, - so a (N,) transformf array is equivalent to (N, 7) float32 in torch. - """ - t = _make_data_torch(shape, device, wp_dtype) - if wp_dtype == wp.float32: - return wp.from_torch(t, dtype=wp.float32) - # For structured types, the torch tensor has the trailing dim; convert to warp - return wp.from_torch(t.contiguous(), dtype=wp_dtype) - - -def _make_payload_torch(shape: tuple, device: str, wp_dtype=wp.float32) -> torch.Tensor: - """Create valid torch data whose entries differ per element, for writer read-back checks. - - Transforms get distinct positions and a fixed 90-degree rotation about Z; ``vec2f`` limits get - distinct ``[-v, v]`` pairs. - """ - count = math.prod(shape) - values = torch.arange(1, count + 1, dtype=torch.float32, device=device).reshape(shape) + 0.25 - if wp_dtype == wp.float32: - return values - if wp_dtype == wp.vec2f: - return torch.stack((-values, values), dim=-1) - if wp_dtype == wp.spatial_vectorf: - return values.unsqueeze(-1) + torch.arange(6, dtype=torch.float32, device=device) / 10.0 - if wp_dtype == wp.transformf: - data = torch.zeros((*shape, 7), dtype=torch.float32, device=device) - data[..., :3] = values.unsqueeze(-1) + torch.arange(3, dtype=torch.float32, device=device) / 10.0 - data[..., 5] = data[..., 6] = 2.0**-0.5 - return data - raise ValueError(f"Unsupported payload dtype: {wp_dtype}") - - -def _make_payload_warp(shape: tuple, device: str, wp_dtype=wp.float32) -> wp.array: - """Create the per-element distinct payload of :func:`_make_payload_torch` as a warp array.""" - return wp.from_torch(_make_payload_torch(shape, device, wp_dtype).contiguous(), dtype=wp_dtype) - - -def _assert_reads_back(value, expected: torch.Tensor, name: str) -> None: - """Assert a data getter (ProxyArray) returns the values a writer just wrote.""" - torch.testing.assert_close(value.torch, expected, atol=1e-5, rtol=1e-5, msg=lambda msg: f"{name}: {msg}") - - -def _make_com_data(backend: str, shape: tuple[int, ...], device: str) -> wp.array: - """Create backend-compatible center-of-mass test data.""" - if backend == "newton": - return wp.zeros(shape, dtype=wp.vec3f, device=device) - return _make_data_warp(shape, device, wp.transformf) - - -def _prime_timestamped_properties(data, property_buffer_pairs: list[tuple[str, str]]): - """Prime public lazy properties and return their concrete timestamped buffers.""" - buffers = [] - for property_name, buffer_name in property_buffer_pairs: - getattr(data, property_name) - buffer = getattr(data, buffer_name) - assert buffer is not None, buffer_name - buffer.timestamp = data._sim_timestamp - buffers.append((buffer_name, buffer)) - return buffers - - -def _assert_buffers_stale(data, buffers) -> None: - for name, buffer in buffers: - assert buffer.timestamp < data._sim_timestamp, name - - -def _make_bad_data_torch(shape: tuple, device: str, wp_dtype=wp.float32) -> torch.Tensor: - """Create torch data with wrong leading shape for negative testing. - - Adds +1 to the first dimension so the shape doesn't match. - """ - bad_shape = (shape[0] + 1,) + shape[1:] - return _make_data_torch(bad_shape, device, wp_dtype) - - -def _make_bad_data_warp(shape: tuple, device: str, wp_dtype=wp.float32) -> wp.array: - """Create warp data with wrong leading shape for negative testing.""" - bad_shape = (shape[0] + 1,) + shape[1:] - return _make_data_warp(bad_shape, device, wp_dtype) - - -def _make_env_mask(num_instances: int, device: str, partial: bool) -> wp.array | None: - """Create an env_mask: None for all envs, or a partial bool mask.""" - if not partial: - return None - mask_np = np.zeros(num_instances, dtype=bool) - mask_np[0] = True - return wp.array(mask_np, dtype=wp.bool, device=device) - - -def _make_env_ids(device: str, subset: bool) -> torch.Tensor | None: - """Create env_ids: None for all envs, or [0] for a subset.""" - if not subset: - return None - return torch.tensor([0], dtype=torch.int32, device=device) - - -def _make_item_mask(total: int, selected: list[int], device: str) -> wp.array: - """Create a bool warp mask with True at `selected` indices, False elsewhere.""" - mask_np = np.zeros(total, dtype=bool) - for i in selected: - mask_np[i] = True - return wp.array(mask_np, dtype=wp.bool, device=device) - - -# --------------------------------------------------------------------------- -# Tests: Root writers — torch/warp × index/mask × all/subset × negative -# --------------------------------------------------------------------------- - -# writer suffix -> data getter that reads the written quantity back -_ROOT_POSE_METHODS = { - "root_pose": "root_link_pose_w", - "root_link_pose": "root_link_pose_w", - "root_com_pose": "root_com_pose_w", -} -_ROOT_VEL_METHODS = { - "root_velocity": "root_com_vel_w", - "root_link_velocity": "root_link_vel_w", - "root_com_velocity": "root_com_vel_w", -} - - -class TestArticulationWritersRoot: - """Test root pose/velocity writers with all input combinations.""" - - @_production_backends - @pytest.mark.parametrize( - "body_ordering", - [None, ("body_0", "body_3", "body_2", "body_1")], - ) - def test_pose_write_invalidates_pose_dependent_caches(self, backend, body_ordering): - art, _ = get_articulation( - backend, - num_instances=2, - num_joints=3, - num_bodies=4, - device="cpu", - body_ordering=body_ordering, - ) - art.data.update(dt=0.01) - buffers = _prime_timestamped_properties( - art.data, - [ - ("root_link_vel_w", "_root_link_vel_w"), - ("body_link_vel_w", "_body_link_vel_w"), - *([("body_com_vel_w", "_body_com_vel_w")] if backend != "newton" else []), - ("projected_gravity_b", "_projected_gravity_b"), - ("heading_w", "_heading_w"), - ("root_link_lin_vel_b", "_root_link_lin_vel_b"), - ("root_link_ang_vel_b", "_root_link_ang_vel_b"), - ("root_com_lin_vel_b", "_root_com_lin_vel_b"), - ("root_com_ang_vel_b", "_root_com_ang_vel_b"), - ], - ) - if backend == "ovphysx" and art.data._body_com_vel_w_backend is not None: - buffers.append(("_body_com_vel_w_backend", art.data._body_com_vel_w_backend)) - root_pose = _make_data_warp((art.num_instances,), "cpu", wp.transformf) - art.write_root_link_pose_to_sim_index(root_pose=root_pose) - _assert_buffers_stale(art.data, buffers) - - @_production_backends - @pytest.mark.parametrize( - "body_ordering", - [None, ("body_0", "body_3", "body_2", "body_1")], - ) - def test_velocity_write_invalidates_body_frame_caches(self, backend, body_ordering): - art, _ = get_articulation( - backend, - num_instances=2, - num_joints=3, - num_bodies=4, - device="cpu", - body_ordering=body_ordering, - ) - art.data.update(dt=0.01) - buffers = _prime_timestamped_properties( - art.data, - [ - ("root_link_lin_vel_b", "_root_link_lin_vel_b"), - ("root_link_ang_vel_b", "_root_link_ang_vel_b"), - ("root_com_lin_vel_b", "_root_com_lin_vel_b"), - ("root_com_ang_vel_b", "_root_com_ang_vel_b"), - ], - ) - root_velocity = _make_data_warp((art.num_instances,), "cpu", wp.spatial_vectorf) - art.write_root_com_velocity_to_sim_index(root_velocity=root_velocity) - _assert_buffers_stale(art.data, buffers) - - @_production_backends - @pytest.mark.parametrize("setter_kind", ["index", "mask"]) - @pytest.mark.parametrize( - "body_ordering", - [None, ("body_0", "body_3", "body_2", "body_1")], - ) - def test_set_coms_invalidates_same_timestamp_dependents(self, backend, setter_kind, body_ordering): - art, _ = get_articulation( - backend, - num_instances=2, - num_joints=3, - num_bodies=4, - device="cpu", - body_ordering=body_ordering, - ) - art.data.update(dt=0.01) - common_pairs = [ - ("root_com_pose_w", "_root_com_pose_w"), - ("root_link_vel_w", "_root_link_vel_w"), - ("body_com_pose_w", "_body_com_pose_w"), - ("body_link_vel_w", "_body_link_vel_w"), - ("root_link_lin_vel_b", "_root_link_lin_vel_b"), - ("root_link_ang_vel_b", "_root_link_ang_vel_b"), - ("root_com_lin_vel_b", "_root_com_lin_vel_b"), - ("root_com_ang_vel_b", "_root_com_ang_vel_b"), - ] - if backend != "newton": - common_pairs += [ - ("root_com_vel_w", "_root_com_vel_w"), - ("body_com_vel_w", "_body_com_vel_w"), - ] - if backend == "physx": - # COM changes also move the COM Jacobian and the generalized dynamics. - common_pairs += [ - ("body_com_jacobian_w", "_body_com_jacobian_w"), - ("mass_matrix", "_mass_matrix"), - ("gravity_compensation_forces", "_gravity_compensation_forces"), - ] - state_buffer_suffix = "_buf" if backend == "ovphysx" else "" - common_pairs += [ - ("root_state_w", f"_root_state_w{state_buffer_suffix}"), - ("root_link_state_w", f"_root_link_state_w{state_buffer_suffix}"), - ("root_com_state_w", f"_root_com_state_w{state_buffer_suffix}"), - ("body_state_w", f"_body_state_w{state_buffer_suffix}"), - ("body_link_state_w", f"_body_link_state_w{state_buffer_suffix}"), - ("body_com_state_w", f"_body_com_state_w{state_buffer_suffix}"), - ] - # Prime public properties before resolving private buffers so Newton allocates lazy caches. - buffers = _prime_timestamped_properties(art.data, common_pairs) - if backend == "newton": - buffers += _prime_timestamped_properties(art.data, [("body_com_pose_b", "_body_com_pose_b")]) - if backend == "ovphysx" and art.data._body_com_vel_w_backend is not None: - art.data._body_com_vel_w_backend.timestamp = art.data._sim_timestamp - buffers.append(("_body_com_vel_w_backend", art.data._body_com_vel_w_backend)) - coms = _make_com_data(backend, (art.num_instances, art.num_bodies), "cpu") - - def set_coms() -> None: - if setter_kind == "index": - kwargs = {} if backend != "physx" else {"full_data": True} - art.set_coms_index(coms=coms, **kwargs) - else: - art.set_coms_mask(coms=coms) - - if backend == "newton": - from unittest.mock import patch - - from isaaclab_newton.physics import NewtonManager - - with patch.object(NewtonManager, "add_model_change"): - set_coms() - else: - set_coms() - _assert_buffers_stale(art.data, buffers) - - # -- index variants -- - - @_backends - @_devices - @pytest.mark.parametrize("method_suffix", _ROOT_POSE_METHODS) - def test_write_root_pose_to_sim_index(self, backend, device, method_suffix): - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device) - num_instances = _NUM_INSTANCES - art.data.update(dt=0.01) - method = getattr(art, f"write_{method_suffix}_to_sim_index") - - # torch, all envs - method(root_pose=_make_data_torch((num_instances,), device, wp.transformf)) - # torch, subset - method(root_pose=_make_data_torch((1,), device, wp.transformf), env_ids=_make_env_ids(device, True)) - # warp, all envs: the matching getter reads the written poses back - method(root_pose=_make_payload_warp((num_instances,), device, wp.transformf)) - getter = _ROOT_POSE_METHODS[method_suffix] - _assert_reads_back( - getattr(art.data, getter), _make_payload_torch((num_instances,), device, wp.transformf), getter - ) - # warp, subset - method(root_pose=_make_data_warp((1,), device, wp.transformf), env_ids=_make_env_ids(device, True)) - # negative: bad torch shape - with pytest.raises((AssertionError, RuntimeError)): - method(root_pose=_make_bad_data_torch((num_instances,), device, wp.transformf)) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(root_pose=_make_bad_data_warp((num_instances,), device, wp.transformf)) - - @_backends - @_devices - @pytest.mark.parametrize("method_suffix", _ROOT_VEL_METHODS) - def test_write_root_velocity_to_sim_index(self, backend, device, method_suffix): - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device) - num_instances = _NUM_INSTANCES - art.data.update(dt=0.01) - method = getattr(art, f"write_{method_suffix}_to_sim_index") - - # torch, all envs - method(root_velocity=_make_data_torch((num_instances,), device, wp.spatial_vectorf)) - # torch, subset - method(root_velocity=_make_data_torch((1,), device, wp.spatial_vectorf), env_ids=_make_env_ids(device, True)) - # warp, all envs: the matching getter reads the written velocities back - method(root_velocity=_make_payload_warp((num_instances,), device, wp.spatial_vectorf)) - getter = _ROOT_VEL_METHODS[method_suffix] - _assert_reads_back( - getattr(art.data, getter), _make_payload_torch((num_instances,), device, wp.spatial_vectorf), getter - ) - # warp, subset - method(root_velocity=_make_data_warp((1,), device, wp.spatial_vectorf), env_ids=_make_env_ids(device, True)) - # negative: bad torch shape - with pytest.raises((AssertionError, RuntimeError)): - method(root_velocity=_make_bad_data_torch((num_instances,), device, wp.spatial_vectorf)) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(root_velocity=_make_bad_data_warp((num_instances,), device, wp.spatial_vectorf)) - - # -- mask variants -- - - @_backends - @_devices - @pytest.mark.parametrize("method_suffix", _ROOT_POSE_METHODS) - def test_write_root_pose_to_sim_mask(self, backend, device, method_suffix): - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device) - num_instances = _NUM_INSTANCES - art.data.update(dt=0.01) - method = getattr(art, f"write_{method_suffix}_to_sim_mask") - - # torch, no mask (all) - method(root_pose=_make_data_torch((num_instances,), device, wp.transformf)) - # torch, partial mask - method( - root_pose=_make_data_torch((num_instances,), device, wp.transformf), - env_mask=_make_env_mask(num_instances, device, True), - ) - # warp, no mask: the matching getter reads the written poses back - method(root_pose=_make_payload_warp((num_instances,), device, wp.transformf)) - getter = _ROOT_POSE_METHODS[method_suffix] - _assert_reads_back( - getattr(art.data, getter), _make_payload_torch((num_instances,), device, wp.transformf), getter - ) - # warp, partial mask - method( - root_pose=_make_data_warp((num_instances,), device, wp.transformf), - env_mask=_make_env_mask(num_instances, device, True), - ) - # negative: bad torch shape - with pytest.raises((AssertionError, RuntimeError)): - method(root_pose=_make_bad_data_torch((num_instances,), device, wp.transformf)) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(root_pose=_make_bad_data_warp((num_instances,), device, wp.transformf)) - - @_backends - @_devices - @pytest.mark.parametrize("method_suffix", _ROOT_VEL_METHODS) - def test_write_root_velocity_to_sim_mask(self, backend, device, method_suffix): - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device) - num_instances = _NUM_INSTANCES - art.data.update(dt=0.01) - method = getattr(art, f"write_{method_suffix}_to_sim_mask") - - # torch, no mask - method(root_velocity=_make_data_torch((num_instances,), device, wp.spatial_vectorf)) - # torch, partial mask - method( - root_velocity=_make_data_torch((num_instances,), device, wp.spatial_vectorf), - env_mask=_make_env_mask(num_instances, device, True), - ) - # warp, no mask: the matching getter reads the written velocities back - method(root_velocity=_make_payload_warp((num_instances,), device, wp.spatial_vectorf)) - getter = _ROOT_VEL_METHODS[method_suffix] - _assert_reads_back( - getattr(art.data, getter), _make_payload_torch((num_instances,), device, wp.spatial_vectorf), getter - ) - # warp, partial mask - method( - root_velocity=_make_data_warp((num_instances,), device, wp.spatial_vectorf), - env_mask=_make_env_mask(num_instances, device, True), - ) - # negative: bad torch shape - with pytest.raises((AssertionError, RuntimeError)): - method(root_velocity=_make_bad_data_torch((num_instances,), device, wp.spatial_vectorf)) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(root_velocity=_make_bad_data_warp((num_instances,), device, wp.spatial_vectorf)) - - -@pytest.mark.parametrize("backend", BACKENDS) -def test_deprecated_joint_friction_writers(backend): - """The deprecated joint friction writers warn and forward to the index writer on every backend.""" - num_instances, num_joints = 2, 4 - art, _ = get_articulation(backend, num_instances, num_joints, 5, device="cpu") - for writer_name, value in (("write_joint_friction_coefficient_to_sim", 0.5), ("write_joint_friction_to_sim", 0.25)): - friction = torch.full((num_instances, num_joints), value) - with pytest.warns(DeprecationWarning): - getattr(art, writer_name)(friction) - _assert_reads_back(art.data.joint_friction_coeff, friction, writer_name) - - -# --------------------------------------------------------------------------- -# Tests: Joint writers — torch/warp × index/mask × all/subset × negative -# --------------------------------------------------------------------------- - -# (method_name, kwarg_name, wp_dtype, accepts_float, read-back getter) -_JOINT_METHODS = [ - ("write_joint_position_to_sim", "position", wp.float32, False, "joint_pos"), - ("write_joint_velocity_to_sim", "velocity", wp.float32, False, "joint_vel"), - ("write_joint_stiffness_to_sim", "stiffness", wp.float32, True, "joint_stiffness"), - ("write_joint_damping_to_sim", "damping", wp.float32, True, "joint_damping"), - ("write_joint_position_limit_to_sim", "limits", wp.vec2f, True, "joint_pos_limits"), - ("write_joint_velocity_limit_to_sim", "limits", wp.float32, True, "joint_vel_limits"), - ("write_joint_effort_limit_to_sim", "limits", wp.float32, True, "joint_effort_limits"), - ("write_joint_armature_to_sim", "armature", wp.float32, True, "joint_armature"), - ("write_joint_friction_coefficient_to_sim", "joint_friction_coeff", wp.float32, False, "joint_friction_coeff"), - ("set_joint_position_target", "target", wp.float32, False, "actuators.target_command.position"), - ("set_joint_velocity_target", "target", wp.float32, False, "actuators.target_command.velocity"), - ("set_joint_effort_target", "target", wp.float32, False, "actuators.target_command.effort"), -] - - -def _read_joint_writer_target(art, getter: str): - """Resolve a read-back getter: a data property, or a dotted path on the articulation.""" - if "." not in getter: - return getattr(art.data, getter) - value = art - for attr in getter.split("."): - value = getattr(value, attr) - return value - - -class TestArticulationWritersJoint: - """Test joint writers/setters with all input combinations.""" - - @_backends - @_devices - @pytest.mark.parametrize( - "method_base, kwarg, wp_dtype, accepts_float, getter", - _JOINT_METHODS, - ids=[m[0] for m in _JOINT_METHODS], - ) - def test_joint_writer_index(self, backend, device, method_base, kwarg, wp_dtype, accepts_float, getter): - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device) - num_instances, num_joints = _NUM_INSTANCES, _NUM_JOINTS - art.data.update(dt=0.01) - method = getattr(art, f"{method_base}_index") - sub_j = min(2, num_joints) - sub_joint_ids = list(range(sub_j)) - - # torch, all envs + all joints - method(**{kwarg: _make_data_torch((num_instances, num_joints), device, wp_dtype)}) - # torch, subset envs + subset joints - method( - **{ - kwarg: _make_data_torch((1, sub_j), device, wp_dtype), - "joint_ids": sub_joint_ids, - "env_ids": _make_env_ids(device, True), - } - ) - # warp, all envs + all joints: the matching getter reads the written values back - method(**{kwarg: _make_payload_warp((num_instances, num_joints), device, wp_dtype)}) - _assert_reads_back( - _read_joint_writer_target(art, getter), - _make_payload_torch((num_instances, num_joints), device, wp_dtype), - getter, - ) - # warp, subset - method( - **{ - kwarg: _make_data_warp((1, sub_j), device, wp_dtype), - "joint_ids": sub_joint_ids, - "env_ids": _make_env_ids(device, True), - } - ) - # float scalar (only for accepts_float methods, and NOT for vec2f position_limit) - if accepts_float and wp_dtype != wp.vec2f: - method(**{kwarg: 1.0}) - # float scalar for vec2f position_limit should raise ValueError - if accepts_float and wp_dtype == wp.vec2f: - with pytest.raises((ValueError, TypeError)): - method(**{kwarg: 1.0}) - # negative: bad torch shape - with pytest.raises((AssertionError, RuntimeError)): - method(**{kwarg: _make_bad_data_torch((num_instances, num_joints), device, wp_dtype)}) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(**{kwarg: _make_bad_data_warp((num_instances, num_joints), device, wp_dtype)}) - - @_backends - @_devices - @pytest.mark.parametrize( - "method_base, kwarg, wp_dtype, accepts_float, getter", - _JOINT_METHODS, - ids=[m[0] for m in _JOINT_METHODS], - ) - def test_joint_writer_mask(self, backend, device, method_base, kwarg, wp_dtype, accepts_float, getter): - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device) - num_instances, num_joints = _NUM_INSTANCES, _NUM_JOINTS - art.data.update(dt=0.01) - method = getattr(art, f"{method_base}_mask") - sub_joint_sel = list(range(min(2, num_joints))) - - # torch, no mask - method(**{kwarg: _make_data_torch((num_instances, num_joints), device, wp_dtype)}) - # torch, partial env_mask + joint_mask - method( - **{ - kwarg: _make_data_torch((num_instances, num_joints), device, wp_dtype), - "joint_mask": _make_item_mask(num_joints, sub_joint_sel, device), - "env_mask": _make_env_mask(num_instances, device, True), - } - ) - # warp, no mask: the matching getter reads the written values back - method(**{kwarg: _make_payload_warp((num_instances, num_joints), device, wp_dtype)}) - _assert_reads_back( - _read_joint_writer_target(art, getter), - _make_payload_torch((num_instances, num_joints), device, wp_dtype), - getter, - ) - # warp, partial env_mask + joint_mask - method( - **{ - kwarg: _make_data_warp((num_instances, num_joints), device, wp_dtype), - "joint_mask": _make_item_mask(num_joints, sub_joint_sel, device), - "env_mask": _make_env_mask(num_instances, device, True), - } - ) - # float scalar (only for accepts_float methods, and NOT for vec2f) - if accepts_float and wp_dtype != wp.vec2f: - method(**{kwarg: 1.0}) - # negative: bad torch shape - with pytest.raises((AssertionError, RuntimeError)): - method(**{kwarg: _make_bad_data_torch((num_instances, num_joints), device, wp_dtype)}) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(**{kwarg: _make_bad_data_warp((num_instances, num_joints), device, wp_dtype)}) - - -# --------------------------------------------------------------------------- -# Tests: Body writers — torch/warp × index/mask × all/subset × negative -# --------------------------------------------------------------------------- - -# (method_name, kwarg_name, read-back getter) -_BODY_METHODS = [ - ("set_masses", "masses", "body_mass"), - ("set_coms", "coms", "body_com_pose_b"), - ("set_inertias", "inertias", "body_inertia"), -] - - -def _body_writer_layout(backend: str, method_base: str) -> tuple[type, int, str]: - """Return the warp dtype, torch trailing size, and read-back getter of a body property writer.""" - if method_base == "set_masses": - return wp.float32, 0, "body_mass" - if method_base == "set_inertias": - return wp.float32, 9, "body_inertia" - if backend == "newton": - # Newton stores the COM as a position only. - return wp.vec3f, 3, "body_com_pos_b" - return wp.transformf, 7, "body_com_pose_b" - - -def _make_body_torch(shape: tuple[int, int], device: str, wp_dtype: type, trailing: int, payload: bool = False): - """Create body-property torch data: valid ones by default, or per-element distinct values.""" - full_shape = (*shape, trailing) if trailing else shape - if payload: - data = torch.arange(1, math.prod(full_shape) + 1, dtype=torch.float32, device=device).reshape(full_shape) - else: - data = torch.ones(full_shape, device=device, dtype=torch.float32) - if wp_dtype == wp.transformf: - data[..., 3:6] = 0.0 - data[..., 6] = 1.0 - if not payload: - data[..., :3] = 0.0 - return data - - -def _make_body_warp(shape: tuple[int, int], device: str, wp_dtype: type, trailing: int, payload: bool = False): - """Create body-property data as a warp array (structured dtypes collapse the trailing dim).""" - t = _make_body_torch(shape, device, wp_dtype, trailing, payload).contiguous() - return wp.from_torch(t, dtype=wp.float32 if wp_dtype == wp.float32 else wp_dtype) - - -class TestArticulationWritersBody: - """Test body property writers/setters with all input combinations.""" - - @_backends - @_devices - @pytest.mark.parametrize("method_base, kwarg", [m[:2] for m in _BODY_METHODS], ids=[m[0] for m in _BODY_METHODS]) - def test_body_writer_index(self, backend, device, method_base, kwarg): - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device) - num_instances, num_bodies = _NUM_INSTANCES, _NUM_BODIES - wp_dtype, trailing, getter = _body_writer_layout(backend, method_base) - art.data.update(dt=0.01) - method = getattr(art, f"{method_base}_index") - sub_b = min(2, num_bodies) - sub_body_ids = list(range(sub_b)) - - # torch, all envs + all bodies - method(**{kwarg: _make_body_torch((num_instances, num_bodies), device, wp_dtype, trailing)}) - # torch, subset - method( - **{ - kwarg: _make_body_torch((1, sub_b), device, wp_dtype, trailing), - "body_ids": sub_body_ids, - "env_ids": _make_env_ids(device, True), - } - ) - # warp, all envs + all bodies: the matching getter reads the written values back - method(**{kwarg: _make_body_warp((num_instances, num_bodies), device, wp_dtype, trailing, payload=True)}) - expected = _make_body_torch((num_instances, num_bodies), device, wp_dtype, trailing, payload=True) - _assert_reads_back(getattr(art.data, getter), expected, getter) - # warp, subset - method( - **{ - kwarg: _make_body_warp((1, sub_b), device, wp_dtype, trailing), - "body_ids": sub_body_ids, - "env_ids": _make_env_ids(device, True), - } - ) - # negative: bad torch shape (extra env) - with pytest.raises((AssertionError, RuntimeError)): - method(**{kwarg: _make_body_torch((num_instances + 1, num_bodies), device, wp_dtype, trailing)}) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(**{kwarg: _make_body_warp((num_instances + 1, num_bodies), device, wp_dtype, trailing)}) - - @_backends - @_devices - @pytest.mark.parametrize("method_base, kwarg", [m[:2] for m in _BODY_METHODS], ids=[m[0] for m in _BODY_METHODS]) - def test_body_writer_mask(self, backend, device, method_base, kwarg): - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device) - num_instances, num_bodies = _NUM_INSTANCES, _NUM_BODIES - wp_dtype, trailing, getter = _body_writer_layout(backend, method_base) - art.data.update(dt=0.01) - method = getattr(art, f"{method_base}_mask") - sub_body_sel = list(range(min(2, num_bodies))) - - # torch, no mask - method(**{kwarg: _make_body_torch((num_instances, num_bodies), device, wp_dtype, trailing)}) - # torch, partial env_mask + body_mask - method( - **{ - kwarg: _make_body_torch((num_instances, num_bodies), device, wp_dtype, trailing), - "body_mask": _make_item_mask(num_bodies, sub_body_sel, device), - "env_mask": _make_env_mask(num_instances, device, True), - } - ) - # warp, no mask: the matching getter reads the written values back - method(**{kwarg: _make_body_warp((num_instances, num_bodies), device, wp_dtype, trailing, payload=True)}) - expected = _make_body_torch((num_instances, num_bodies), device, wp_dtype, trailing, payload=True) - _assert_reads_back(getattr(art.data, getter), expected, getter) - # warp, partial env_mask + body_mask - method( - **{ - kwarg: _make_body_warp((num_instances, num_bodies), device, wp_dtype, trailing), - "body_mask": _make_item_mask(num_bodies, sub_body_sel, device), - "env_mask": _make_env_mask(num_instances, device, True), - } - ) - # negative: bad torch shape - with pytest.raises((AssertionError, RuntimeError)): - method(**{kwarg: _make_body_torch((num_instances + 1, num_bodies), device, wp_dtype, trailing)}) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(**{kwarg: _make_body_warp((num_instances + 1, num_bodies), device, wp_dtype, trailing)}) - - -# --------------------------------------------------------------------------- -# Tendon tests — counts, names, data, writers -# --------------------------------------------------------------------------- - -# Newton's fixed-tendon solver integration is covered in isaaclab_newton/test/assets/test_articulation.py. -# This fixture matrix also requires spatial tendons and PhysX-specific properties. -_tendon_backends = pytest.mark.parametrize("backend", [b for b in BACKENDS if b != "newton"], indirect=False) - -# The first row covers the zero-spatial-tendon edge; the second has distinct counts on every axis. -_tendon_dims = pytest.mark.parametrize( - "num_instances, num_joints, num_bodies, num_fixed_tendons, num_spatial_tendons", - [ - (1, 2, 2, 1, 0), # fixed only - (2, 6, 7, 3, 2), # both types - ], -) - -# (property, tendon kind, dtype) -_TENDON_DATA_PROPERTIES = [ - ("fixed_tendon_stiffness", "fixed", wp.float32), - ("fixed_tendon_damping", "fixed", wp.float32), - ("fixed_tendon_limit_stiffness", "fixed", wp.float32), - ("fixed_tendon_rest_length", "fixed", wp.float32), - ("fixed_tendon_offset", "fixed", wp.float32), - ("fixed_tendon_pos_limits", "fixed", wp.vec2f), - ("spatial_tendon_stiffness", "spatial", wp.float32), - ("spatial_tendon_damping", "spatial", wp.float32), - ("spatial_tendon_limit_stiffness", "spatial", wp.float32), - ("spatial_tendon_offset", "spatial", wp.float32), -] - - -class TestArticulationTendons: - """Test tendon counts, names, and data properties.""" - - @_tendon_backends - @_tendon_dims - @_devices - def test_articulation_tendon_contract( - self, backend, num_instances, num_joints, num_bodies, num_fixed_tendons, num_spatial_tendons, device - ): - art, _ = get_articulation( - backend, num_instances, num_joints, num_bodies, num_fixed_tendons, num_spatial_tendons, device - ) - counts = {"fixed": num_fixed_tendons, "spatial": num_spatial_tendons} - assert art.num_fixed_tendons == num_fixed_tendons - assert art.num_spatial_tendons == num_spatial_tendons - for names, count in ( - (art.fixed_tendon_names, num_fixed_tendons), - (art.spatial_tendon_names, num_spatial_tendons), - ): - assert isinstance(names, list) - assert len(names) == count - assert all(isinstance(n, str) for n in names) - - art.data.update(dt=0.01) - for name, kind, dtype in _TENDON_DATA_PROPERTIES: - if counts[kind] == 0: - continue - expected_shape = (num_instances, counts[kind]) - if backend == "physx" and name == "fixed_tendon_pos_limits": - # Known inconsistency: PhysX exposes (N, T, 2) float32 although its docstring and OVPhysX - # advertise (N, T) vec2f. Pin the current layout so a change on either side is noticed. - expected_shape, dtype = (*expected_shape, 2), wp.float32 - _check_proxy_array(getattr(art.data, name), expected_shape=expected_shape, expected_dtype=dtype, name=name) - - -# (method_name, kwarg_name, read-back data property) -_FIXED_TENDON_METHODS = [ - ("set_fixed_tendon_stiffness", "stiffness", "fixed_tendon_stiffness"), - ("set_fixed_tendon_damping", "damping", "fixed_tendon_damping"), - ("set_fixed_tendon_limit_stiffness", "limit_stiffness", "fixed_tendon_limit_stiffness"), - ("set_fixed_tendon_rest_length", "rest_length", "fixed_tendon_rest_length"), - ("set_fixed_tendon_offset", "offset", "fixed_tendon_offset"), - ("set_fixed_tendon_position_limit", "limit", "fixed_tendon_pos_limits"), -] -_SPATIAL_TENDON_METHODS = [ - ("set_spatial_tendon_stiffness", "stiffness", "spatial_tendon_stiffness"), - ("set_spatial_tendon_damping", "damping", "spatial_tendon_damping"), - ("set_spatial_tendon_limit_stiffness", "limit_stiffness", "spatial_tendon_limit_stiffness"), - ("set_spatial_tendon_offset", "offset", "spatial_tendon_offset"), -] -_TENDON_METHODS = [("fixed", *m) for m in _FIXED_TENDON_METHODS] + [("spatial", *m) for m in _SPATIAL_TENDON_METHODS] - - -class TestArticulationWritersTendon: - """Test tendon writers/setters with all input combinations.""" - - @_tendon_backends - @_tendon_dims - @_devices - @pytest.mark.parametrize("selection", ["index", "mask"]) - @pytest.mark.parametrize("kind, method_base, kwarg, getter", _TENDON_METHODS, ids=[m[1] for m in _TENDON_METHODS]) - def test_tendon_writer( - self, - backend, - num_instances, - num_joints, - num_bodies, - num_fixed_tendons, - num_spatial_tendons, - device, - selection, - kind, - method_base, - kwarg, - getter, - ): - num_tendons = num_fixed_tendons if kind == "fixed" else num_spatial_tendons - if num_tendons == 0: - pytest.skip(f"No {kind} tendons configured") - art, _ = get_articulation( - backend, num_instances, num_joints, num_bodies, num_fixed_tendons, num_spatial_tendons, device - ) - art.data.update(dt=0.01) - method = getattr(art, f"{method_base}_{selection}") - wp_dtype = wp.vec2f if getter == "fixed_tendon_pos_limits" else wp.float32 - full_shape = (num_instances, num_tendons) - sub_t = min(2, num_tendons) - sub_tendons = list(range(sub_t)) - if selection == "index": - subset = {f"{kind}_tendon_ids": sub_tendons, "env_ids": _make_env_ids(device, True)} - subset_shape = (1, sub_t) - else: - subset = { - f"{kind}_tendon_mask": _make_item_mask(num_tendons, sub_tendons, device), - "env_mask": _make_env_mask(num_instances, device, True), - } - subset_shape = full_shape - - # torch, all envs + all tendons - method(**{kwarg: _make_data_torch(full_shape, device, wp_dtype)}) - # torch, subset - method(**{kwarg: _make_data_torch(subset_shape, device, wp_dtype)}, **subset) - # warp, all envs + all tendons: the matching data property reads the written values back - method(**{kwarg: _make_payload_warp(full_shape, device, wp_dtype)}) - _assert_reads_back(getattr(art.data, getter), _make_payload_torch(full_shape, device, wp_dtype), getter) - # warp, subset - method(**{kwarg: _make_data_warp(subset_shape, device, wp_dtype)}, **subset) - # Position limits require a lower/upper pair; other properties accept scalar values. - if wp_dtype == wp.vec2f: - with pytest.raises((ValueError, TypeError)): - method(**{kwarg: 1.0}) - else: - method(**{kwarg: 1.0}) - # negative: bad torch shape - with pytest.raises((AssertionError, RuntimeError)): - method(**{kwarg: _make_bad_data_torch(full_shape, device, wp_dtype)}) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(**{kwarg: _make_bad_data_warp(full_shape, device, wp_dtype)}) - - @_tendon_backends - @_devices - @pytest.mark.parametrize("selection", ["index", "mask"]) - @pytest.mark.parametrize("kind", ["fixed", "spatial"]) - def test_write_tendon_properties_to_sim_selects_envs(self, backend, device, selection, kind): - """Pushing tendon properties writes the selected environments to the backend.""" - num_instances = 3 - art, raw_backend = get_articulation(backend, num_instances, 2, 2, 2, 2, device) - art.data.update(dt=0.01) - env_writes = [] - if backend == "physx": - - def capture(*args, indices=None, **kwargs): - env_writes.append(indices.numpy().tolist()) - - setattr(raw_backend, f"set_{kind}_tendon_properties", MagicMock(side_effect=capture)) - else: - from isaaclab_ov import tensor_types as TT - - prefix = "FIXED_TENDON" if kind == "fixed" else "SPATIAL_TENDON" - tendon_types = {getattr(TT, name) for name in dir(TT) if name.startswith(prefix)} - set_attribute = art._root_view.set_attribute - - def capture(name, values, *, indices=None, mask=None): - if name in tendon_types: - if mask is not None: - env_writes.append(np.flatnonzero(mask.numpy()).tolist()) - else: - env_writes.append(list(range(num_instances)) if indices is None else indices.numpy().tolist()) - set_attribute(name, values, indices=indices, mask=mask) - - object.__setattr__(art._root_view, "set_attribute", capture) - write = getattr(art, f"write_{kind}_tendon_properties_to_sim_{selection}") - - for selected_envs in (list(range(num_instances)), [0, 2]): - env_writes.clear() - if selection == "index": - write( - env_ids=None if len(selected_envs) == num_instances else torch.tensor(selected_envs, device=device) - ) - else: - mask = None - if len(selected_envs) < num_instances: - mask = wp.array([i in selected_envs for i in range(num_instances)], dtype=wp.bool, device=device) - write(env_mask=mask) - assert env_writes, "no tendon properties were written" - assert all(envs == selected_envs for envs in env_writes), env_writes diff --git a/source/isaaclab/test/assets/test_articulation_ordering_iface.py b/source/isaaclab/test/assets/test_articulation_ordering_iface.py deleted file mode 100644 index 7fdbad251e57..000000000000 --- a/source/isaaclab/test/assets/test_articulation_ordering_iface.py +++ /dev/null @@ -1,2419 +0,0 @@ -# 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 - -# pyright: reportPrivateUsage=none - -"""Mocked cross-backend articulation ordering interface tests.""" - -from unittest.mock import MagicMock, patch - -import numpy as np -import pytest -import torch -import warp as wp -from _articulation_iface_test_utils import BACKEND_UNAVAILABLE_REASONS, BACKENDS, get_articulation -from _pytest.mark.structures import ParameterSet - -from isaaclab.utils import replace -from isaaclab.utils.buffers import TimestampedBuffer - - -def _make_body_ordering_backend_data(num_instances: int, num_bodies: int) -> tuple[np.ndarray, ...]: - """Create deterministic backend-order body data with identity rotations.""" - root_pose = np.zeros((num_instances, 7), dtype=np.float32) - root_pose[:, 6] = 1.0 - root_vel = np.zeros((num_instances, 6), dtype=np.float32) - link_pose = np.zeros((num_instances, num_bodies, 7), dtype=np.float32) - com_pose_b = np.zeros((num_instances, num_bodies, 7), dtype=np.float32) - body_com_vel = np.zeros((num_instances, num_bodies, 6), dtype=np.float32) - body_acc = np.zeros((num_instances, num_bodies, 6), dtype=np.float32) - - for env_index in range(num_instances): - root_pose[env_index, :3] = (10.0 + env_index, 0.0, 0.0) - root_vel[env_index, 0] = 3.0 + env_index - root_vel[env_index, 5] = 7.0 + env_index - for body_index in range(num_bodies): - link_pose[env_index, body_index, :3] = (10.0 * body_index, float(env_index), 0.0) - link_pose[env_index, body_index, 6] = 1.0 - com_pose_b[env_index, body_index, :3] = (float(body_index + 1), 0.0, 0.0) - com_pose_b[env_index, body_index, 6] = 1.0 - body_com_vel[env_index, body_index, 0] = 20.0 + body_index - body_com_vel[env_index, body_index, 5] = 30.0 + body_index - body_acc[env_index, body_index, 0] = 100.0 + body_index - body_acc[env_index, body_index, 3] = 200.0 + body_index - - return root_pose, root_vel, link_pose, com_pose_b, body_com_vel, body_acc - - -def _install_test_body_ordering(art, mode: str = "reversed") -> np.ndarray: - """Install a reversed or cyclic body ordering that preserves a fixed root body.""" - names = art.backend_body_names - movable = names[1:] if art.is_fixed_base else names - if mode == "reversed": - movable = tuple(reversed(movable)) - elif mode == "cyclic": - movable = (movable[-1], *movable[:-1]) - else: - raise ValueError(f"Unsupported body ordering mode: {mode}") - body_ordering = (names[0], *movable) if art.is_fixed_base else tuple(movable) - art.cfg = replace(art.cfg, body_ordering=body_ordering) - # Re-resolve and re-stage the ordering maps exactly as backend initialization - # does after a config change installs a new ordering on an already-initialized - # articulation. - art._resolve_and_install_ordering_maps() - art._ordering_configure_backend_staging() - return np.asarray(art.body_ordering.user_to_backend_indices, dtype=np.int64) - - -def _install_test_joint_ordering(art, mode: str = "reversed") -> np.ndarray: - """Install a reversed or cyclic public joint ordering on an already constructed articulation.""" - art.cfg = replace(art.cfg, joint_ordering=_joint_ordering_for_mode(mode, art.num_joints)) - # Re-resolve and re-stage the ordering maps exactly as backend initialization - # does after a config change installs a new ordering on an already-initialized - # articulation. - art._resolve_and_install_ordering_maps() - art._ordering_configure_backend_staging() - return np.asarray(art.joint_ordering.user_to_backend_indices, dtype=np.int64) - - -def _joint_ordering_for_mode(mode: str, num_joints: int) -> tuple[str, ...] | None: - """Return the configured public joint ordering for a parity-test mode.""" - backend_names = tuple(f"joint_{index}" for index in range(num_joints)) - if mode == "none": - return None - if mode == "identity": - return backend_names - if mode == "reversed": - return tuple(reversed(backend_names)) - if mode == "cyclic": - return (backend_names[-1], *backend_names[:-1]) - raise ValueError(f"Unsupported joint ordering mode: {mode}") - - -def _body_ordering_for_mode(mode: str, num_bodies: int) -> tuple[str, ...] | None: - """Return the configured public body ordering for a parity-test mode.""" - backend_names = tuple(f"body_{index}" for index in range(num_bodies)) - if mode == "none": - return None - if mode == "identity": - return backend_names - if mode == "reversed": - return tuple(reversed(backend_names)) - if mode == "cyclic": - return (backend_names[-1], *backend_names[:-1]) - raise ValueError(f"Unsupported body ordering mode: {mode}") - - -def _ordering_shadow_shape(buffer: wp.array | TimestampedBuffer) -> tuple[int, ...]: - """Return the allocation shape for a raw or timestamped ordering shadow.""" - if isinstance(buffer, TimestampedBuffer): - buffer = buffer.data - return tuple(buffer.shape) - - -def _ordering_user_to_backend(ordering, item_count: int) -> np.ndarray: - """Return public-to-backend indices for an optional ordering map.""" - if ordering is None: - return np.arange(item_count, dtype=np.int64) - return np.asarray(ordering.user_to_backend_indices, dtype=np.int64) - - -def _expected_backend_to_user( - public_names: tuple[str, ...] | None, - backend_names: tuple[str, ...], -) -> np.ndarray: - """Derive backend-to-public indices independently from the ordering map.""" - if public_names is None: - return np.arange(len(backend_names), dtype=np.int64) - return np.asarray([public_names.index(name) for name in backend_names], dtype=np.int64) - - -def _make_backend_com_poses(num_instances: int, num_bodies: int) -> np.ndarray: - """Create distinct deterministic body COM transforms in backend order.""" - coms = np.zeros((num_instances, num_bodies, 7), dtype=np.float32) - for env_index in range(num_instances): - for body_index in range(num_bodies): - value = float(100 * env_index + 10 * body_index) - coms[env_index, body_index, :3] = (value + 1.0, value + 2.0, value + 3.0) - coms[env_index, body_index, 6] = 1.0 - return coms - - -def _seed_backend_com_poses(backend: str, art, raw_backend, coms: np.ndarray) -> None: - """Seed backend-order COM transforms and invalidate both public and ordered caches.""" - if backend == "physx": - raw_backend._coms = wp.array(coms, dtype=wp.float32, device="cpu") - elif backend == "ovphysx": - from isaaclab_ov import tensor_types as TT - - raw_backend.bindings[TT.BODY_COM_POSE]._data = coms.copy() - else: - raise AssertionError(f"Unsupported backend for COM parity test: {backend}") - - art.data._body_com_pose_b.timestamp = -1.0 - backend_staging = art.data._body_com_pose_b_backend - if backend_staging is not None: - backend_staging.timestamp = -1.0 - - -def _read_backend_com_poses(backend: str, art, raw_backend) -> np.ndarray: - """Read complete backend-order body COM transforms through the backend API.""" - if backend == "physx": - return raw_backend.get_coms().numpy().reshape(art.num_instances, art.num_bodies, 7).copy() - if backend == "ovphysx": - from isaaclab_ov import tensor_types as TT - - return ( - art.root_view.get_attribute(TT.BODY_COM_POSE).numpy().reshape(art.num_instances, art.num_bodies, 7).copy() - ) - raise AssertionError(f"Unsupported backend for COM parity test: {backend}") - - -def _make_selected_com_poses(env_ids: list[int], body_ids: list[int], base_value: float) -> np.ndarray: - """Create distinct deterministic public-order COM transforms for a selected rectangle.""" - coms = np.zeros((len(env_ids), len(body_ids), 7), dtype=np.float32) - for env_offset, env_id in enumerate(env_ids): - for body_offset, body_id in enumerate(body_ids): - value = base_value + float(100 * env_id + 10 * body_id) - coms[env_offset, body_offset, :3] = (value + 1.0, value + 2.0, value + 3.0) - coms[env_offset, body_offset, 6] = 1.0 - return coms - - -def _seed_backend_joint_state( - backend: str, - art, - raw_backend, - position: np.ndarray, - velocity: np.ndarray, -) -> None: - """Seed deterministic joint state directly in backend-order storage.""" - if backend == "physx": - raw_backend._noop_setters = False - raw_backend._dof_positions = wp.array(position, dtype=wp.float32, device=art.device) - raw_backend._dof_velocities = wp.array(velocity, dtype=wp.float32, device=art.device) - art.data._joint_pos.timestamp = -1.0 - art.data._joint_vel.timestamp = -1.0 - return - if backend == "ovphysx": - from isaaclab_ov import tensor_types as TT - - raw_backend.bindings[TT.DOF_POSITION]._data = position.copy() - raw_backend.bindings[TT.DOF_VELOCITY]._data = velocity.copy() - for buffer_name in ("_joint_pos_buf", "_joint_vel_buf", "_joint_pos_backend", "_joint_vel_backend"): - buffer = getattr(art.data, buffer_name, None) - if buffer is not None: - buffer.timestamp = -1.0 - return - if backend == "newton": - art.data._sim_bind_joint_pos.assign(wp.array(position, dtype=wp.float32, device=art.device)) - art.data._sim_bind_joint_vel.assign(wp.array(velocity, dtype=wp.float32, device=art.device)) - # Raw sim-bind writes bypass the solver step, so mirror the post-step - # publish that keeps the public-order shadows coherent under ordering. - art.data._refresh_user_order_joint_state() - return - raise AssertionError(f"Unsupported backend for joint-state parity test: {backend}") - - -def _mutate_backend_joint_state( - backend: str, - art, - raw_backend, - position: np.ndarray, - velocity: np.ndarray, -) -> None: - """Mutate backend-order joint state without invalidating staging timestamps.""" - if backend == "physx": - raw_backend._dof_positions.assign(wp.array(position, dtype=wp.float32, device=art.device)) - raw_backend._dof_velocities.assign(wp.array(velocity, dtype=wp.float32, device=art.device)) - return - if backend == "ovphysx": - from isaaclab_ov import tensor_types as TT - - raw_backend.bindings[TT.DOF_POSITION]._data = position.copy() - raw_backend.bindings[TT.DOF_VELOCITY]._data = velocity.copy() - return - raise AssertionError(f"Unsupported backend for stale joint-state parity test: {backend}") - - -def _read_backend_joint_state(backend: str, art, raw_backend) -> tuple[np.ndarray, np.ndarray]: - """Return joint position and velocity from backend-order storage.""" - if backend == "physx": - return raw_backend._dof_positions.numpy().copy(), raw_backend._dof_velocities.numpy().copy() - if backend == "ovphysx": - from isaaclab_ov import tensor_types as TT - - position = np.asarray(raw_backend.bindings[TT.DOF_POSITION]._data).copy() - velocity = np.asarray(raw_backend.bindings[TT.DOF_VELOCITY]._data).copy() - return position, velocity - if backend == "newton": - return art.data._sim_bind_joint_pos.numpy().copy(), art.data._sim_bind_joint_vel.numpy().copy() - raise AssertionError(f"Unsupported backend for joint-state parity test: {backend}") - - -def _coverage_ids(coverage: str, num_instances: int, num_items: int) -> tuple[list[int], list[int]]: - """Return environment and item IDs for a write-coverage mode.""" - if coverage == "one_env_one_item": - return [1], [0] - if coverage == "all_envs_one_item": - return list(range(num_instances)), [0] - if coverage == "one_env_all_items": - return [1], list(range(num_items)) - raise ValueError(f"Unsupported write coverage: {coverage}") - - -def _write_selected_joint_state( - art, - selection: str, - operation: str, - coverage: str, - position_value: float, - velocity_value: float, -) -> tuple[list[int], list[int], np.ndarray, np.ndarray]: - """Write selected public joint state through an index or mask API.""" - if operation not in ("position", "velocity", "state"): - raise ValueError(f"Unsupported joint-state operation: {operation}") - - env_ids, joint_ids = _coverage_ids(coverage, art.num_instances, art.num_joints) - payload_offsets = 10.0 * np.asarray(env_ids, dtype=np.float32)[:, None] - payload_offsets = payload_offsets + np.asarray(joint_ids, dtype=np.float32)[None, :] - position_payload = np.asarray(position_value + payload_offsets, dtype=np.float32) - velocity_payload = np.asarray(velocity_value + payload_offsets, dtype=np.float32) - if selection == "index": - env_ids_wp = wp.array(env_ids, dtype=wp.int32, device=art.device) - joint_ids_wp = wp.array(joint_ids, dtype=wp.int32, device=art.device) - position = torch.tensor(position_payload, dtype=torch.float32, device=art.device) - velocity = torch.tensor(velocity_payload, dtype=torch.float32, device=art.device) - if operation == "position": - art.write_joint_position_to_sim_index( - position=position, - env_ids=env_ids_wp, - joint_ids=joint_ids_wp, - ) - elif operation == "velocity": - art.write_joint_velocity_to_sim_index( - velocity=velocity, - env_ids=env_ids_wp, - joint_ids=joint_ids_wp, - ) - else: - art.write_joint_state_to_sim_index( - position=position, - velocity=velocity, - env_ids=env_ids_wp, - joint_ids=joint_ids_wp, - ) - return env_ids, joint_ids, position_payload, velocity_payload - - if selection == "mask": - env_mask = wp.array( - [env_index in env_ids for env_index in range(art.num_instances)], - dtype=wp.bool, - device=art.device, - ) - joint_mask = wp.array( - [joint_index in joint_ids for joint_index in range(art.num_joints)], - dtype=wp.bool, - device=art.device, - ) - position = torch.zeros((art.num_instances, art.num_joints), dtype=torch.float32, device=art.device) - velocity = torch.zeros((art.num_instances, art.num_joints), dtype=torch.float32, device=art.device) - env_index = torch.tensor(env_ids, dtype=torch.long, device=art.device) - joint_index = torch.tensor(joint_ids, dtype=torch.long, device=art.device) - position[env_index[:, None], joint_index] = torch.tensor( - position_payload, dtype=torch.float32, device=art.device - ) - velocity[env_index[:, None], joint_index] = torch.tensor( - velocity_payload, dtype=torch.float32, device=art.device - ) - if operation == "position": - art.write_joint_position_to_sim_mask( - position=position, - env_mask=env_mask, - joint_mask=joint_mask, - ) - elif operation == "velocity": - art.write_joint_velocity_to_sim_mask( - velocity=velocity, - env_mask=env_mask, - joint_mask=joint_mask, - ) - else: - art.write_joint_state_to_sim_mask( - position=position, - velocity=velocity, - env_mask=env_mask, - joint_mask=joint_mask, - ) - return env_ids, joint_ids, position_payload, velocity_payload - - raise ValueError(f"Unsupported joint selection mode: {selection}") - - -def _exercise_partial_joint_write( - backend: str, - ordering_mode: str, - selection: str, - operation: str, - coverage: str, -) -> None: - """Assert partial public writes preserve complete backend joint rows.""" - num_instances = 2 - num_joints = 3 - joint_ordering = _joint_ordering_for_mode(ordering_mode, num_joints) - art, raw_backend = get_articulation( - backend, - num_instances=num_instances, - num_joints=num_joints, - num_bodies=2, - device="cpu", - joint_ordering=joint_ordering, - ) - initial_position = np.asarray([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]], dtype=np.float32) - initial_velocity = np.asarray([[110.0, 120.0, 130.0], [140.0, 150.0, 160.0]], dtype=np.float32) - _seed_backend_joint_state(backend, art, raw_backend, initial_position, initial_velocity) - - user_to_backend = _ordering_user_to_backend(art.joint_ordering, num_joints) - expected_position = initial_position.copy() - expected_velocity = initial_velocity.copy() - - for position_value, velocity_value in ((901.0, 902.0), (903.0, 904.0)): - env_ids, joint_ids, position_payload, velocity_payload = _write_selected_joint_state( - art, - selection, - operation, - coverage, - position_value, - velocity_value, - ) - backend_joint_ids = user_to_backend[joint_ids] - selected_cells = np.ix_(env_ids, backend_joint_ids) - if operation in ("position", "state"): - expected_position[selected_cells] = position_payload - if operation in ("velocity", "state"): - expected_velocity[selected_cells] = velocity_payload - - backend_position, backend_velocity = _read_backend_joint_state(backend, art, raw_backend) - np.testing.assert_array_equal(backend_position, expected_position) - np.testing.assert_array_equal(backend_velocity, expected_velocity) - - public_position = art.data.joint_pos.torch.detach().cpu().numpy() - public_velocity = art.data.joint_vel.torch.detach().cpu().numpy() - np.testing.assert_array_equal(public_position, expected_position[:, user_to_backend]) - np.testing.assert_array_equal(public_velocity, expected_velocity[:, user_to_backend]) - - -def _exercise_stale_partial_joint_write(backend: str, operation: str) -> None: - """Assert partial writes refresh complete rows after the simulation timestamp advances.""" - num_instances = 2 - num_joints = 3 - art, raw_backend = get_articulation( - backend, - num_instances=num_instances, - num_joints=num_joints, - num_bodies=2, - device="cpu", - joint_ordering=_joint_ordering_for_mode("reversed", num_joints), - ) - initial_position = np.asarray([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]], dtype=np.float32) - initial_velocity = np.asarray([[110.0, 120.0, 130.0], [140.0, 150.0, 160.0]], dtype=np.float32) - _seed_backend_joint_state(backend, art, raw_backend, initial_position, initial_velocity) - user_to_backend = np.asarray(art.joint_ordering.user_to_backend_indices, dtype=np.int64) - - expected_position = initial_position.copy() - expected_velocity = initial_velocity.copy() - public_position_staging = art.data._joint_pos if backend == "physx" else art.data._joint_pos_buf - public_velocity_staging = art.data._joint_vel if backend == "physx" else art.data._joint_vel_buf - public_staging = [] - backend_staging = [] - if operation in ("position", "state"): - _ = art.data.joint_pos - public_staging.append(("public position", public_position_staging)) - backend_staging.append(("backend position", art.data._joint_pos_backend)) - if operation in ("velocity", "state"): - _ = art.data.joint_vel - public_staging.append(("public velocity", public_velocity_staging)) - backend_staging.append(("backend velocity", art.data._joint_vel_backend)) - # After a read, the public-order buffer is current on every backend. Whether the backend-order - # staging image is also current depends on the backend: PhysX gathers ordered reads straight - # from the tensor-API view and stages the backend image only on the write path, so it is not - # refreshed by a read; OVPhysX still restages it on read. The backend-agnostic guarantee is that - # the public buffer is current -- the write path's re-staging is exercised by the final assert. - for staging_name, staging_buffer in public_staging: - assert staging_buffer is not None - assert staging_buffer.timestamp == art.data._sim_timestamp, f"{backend} {staging_name} staging is not current" - for _, staging_buffer in backend_staging: - assert staging_buffer is not None - staging_buffers = public_staging + backend_staging - - for write_index, (position_value, velocity_value) in enumerate(((901.0, 902.0), (903.0, 904.0))): - if write_index == 1: - for staging_name, staging_buffer in staging_buffers: - assert staging_buffer is not None - assert staging_buffer.timestamp >= 0.0, f"{backend} {staging_name} staging was invalidated" - assert staging_buffer.timestamp < art.data._sim_timestamp - env_ids, joint_ids, position_payload, velocity_payload = _write_selected_joint_state( - art, - "index", - operation, - "one_env_one_item", - position_value, - velocity_value, - ) - backend_joint_ids = user_to_backend[joint_ids] - selected_cells = np.ix_(env_ids, backend_joint_ids) - if operation in ("position", "state"): - expected_position[selected_cells] = position_payload - if operation in ("velocity", "state"): - expected_velocity[selected_cells] = velocity_payload - - backend_position, backend_velocity = _read_backend_joint_state(backend, art, raw_backend) - np.testing.assert_array_equal(backend_position, expected_position) - np.testing.assert_array_equal(backend_velocity, expected_velocity) - - if write_index == 0: - selected_mask = np.zeros((num_instances, num_joints), dtype=bool) - selected_mask[selected_cells] = True - art.data._sim_timestamp += 0.01 - if operation in ("position", "state"): - newer_position = 1000.0 + np.arange(num_instances * num_joints, dtype=np.float32).reshape( - num_instances, num_joints - ) - expected_position[~selected_mask] = newer_position[~selected_mask] - if operation in ("velocity", "state"): - newer_velocity = 2000.0 + np.arange(num_instances * num_joints, dtype=np.float32).reshape( - num_instances, num_joints - ) - expected_velocity[~selected_mask] = newer_velocity[~selected_mask] - _mutate_backend_joint_state(backend, art, raw_backend, expected_position, expected_velocity) - - -def _exercise_duplicate_full_length_joint_write(backend: str, operation: str) -> None: - """Assert duplicate full-length selectors preserve omitted backend joints.""" - num_instances = 2 - num_joints = 3 - art, raw_backend = get_articulation( - backend, - num_instances=num_instances, - num_joints=num_joints, - num_bodies=2, - device="cpu", - joint_ordering=_joint_ordering_for_mode("reversed", num_joints), - ) - initial_position = np.asarray([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]], dtype=np.float32) - initial_velocity = np.asarray([[110.0, 120.0, 130.0], [140.0, 150.0, 160.0]], dtype=np.float32) - _seed_backend_joint_state(backend, art, raw_backend, initial_position, initial_velocity) - - public_staging = [] - backend_staging = [] - if operation in ("position", "state"): - _ = art.data.joint_pos - public_staging.append(art.data._joint_pos if backend == "physx" else art.data._joint_pos_buf) - backend_staging.append(art.data._joint_pos_backend) - if operation in ("velocity", "state"): - _ = art.data.joint_vel - public_staging.append(art.data._joint_vel if backend == "physx" else art.data._joint_vel_buf) - backend_staging.append(art.data._joint_vel_backend) - # A read makes the public buffer current on every backend; the backend-order staging image is - # only guaranteed current on backends that restage on read (OVPhysX), not PhysX, which stages it - # lazily on the write path. Preservation of the omitted rows is verified by the final assert. - for staging_buffer in public_staging: - assert staging_buffer is not None - assert staging_buffer.timestamp == art.data._sim_timestamp - for staging_buffer in backend_staging: - assert staging_buffer is not None - - art.data._sim_timestamp += 0.01 - newer_position = 1000.0 + np.arange(num_instances * num_joints, dtype=np.float32).reshape(num_instances, num_joints) - newer_velocity = 2000.0 + np.arange(num_instances * num_joints, dtype=np.float32).reshape(num_instances, num_joints) - _mutate_backend_joint_state(backend, art, raw_backend, newer_position, newer_velocity) - - env_ids = [1] - joint_ids = [0, 1, 1] - position_payload = np.asarray([[901.0, 902.0, 902.0]], dtype=np.float32) - velocity_payload = np.asarray([[1901.0, 1902.0, 1902.0]], dtype=np.float32) - env_ids_wp = wp.array(env_ids, dtype=wp.int32, device=art.device) - joint_ids_wp = wp.array(joint_ids, dtype=wp.int32, device=art.device) - position = torch.tensor(position_payload, dtype=torch.float32, device=art.device) - velocity = torch.tensor(velocity_payload, dtype=torch.float32, device=art.device) - if operation == "position": - art.write_joint_position_to_sim_index(position=position, env_ids=env_ids_wp, joint_ids=joint_ids_wp) - elif operation == "velocity": - art.write_joint_velocity_to_sim_index(velocity=velocity, env_ids=env_ids_wp, joint_ids=joint_ids_wp) - else: - art.write_joint_state_to_sim_index( - position=position, - velocity=velocity, - env_ids=env_ids_wp, - joint_ids=joint_ids_wp, - ) - - user_to_backend = np.asarray(art.joint_ordering.user_to_backend_indices, dtype=np.int64) - expected_position = newer_position.copy() - expected_velocity = newer_velocity.copy() - for payload_index, joint_id in enumerate(joint_ids): - backend_joint_id = user_to_backend[joint_id] - if operation in ("position", "state"): - expected_position[env_ids[0], backend_joint_id] = position_payload[0, payload_index] - if operation in ("velocity", "state"): - expected_velocity[env_ids[0], backend_joint_id] = velocity_payload[0, payload_index] - - backend_position, backend_velocity = _read_backend_joint_state(backend, art, raw_backend) - np.testing.assert_array_equal(backend_position, expected_position) - np.testing.assert_array_equal(backend_velocity, expected_velocity) - omitted_backend_joint = user_to_backend[2] - duplicate_backend_joint = user_to_backend[1] - if operation in ("position", "state"): - assert backend_position[env_ids[0], omitted_backend_joint] == newer_position[env_ids[0], omitted_backend_joint] - assert backend_position[env_ids[0], duplicate_backend_joint] == position_payload[0, 1] - if operation in ("velocity", "state"): - assert backend_velocity[env_ids[0], omitted_backend_joint] == newer_velocity[env_ids[0], omitted_backend_joint] - assert backend_velocity[env_ids[0], duplicate_backend_joint] == velocity_payload[0, 1] - - -def _make_dynamics_ordering_backend_data( - num_instances: int, num_joints: int, num_jacobi_bodies: int, num_base_dofs: int -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """Create deterministic backend-order dynamics data.""" - num_dofs = num_joints + num_base_dofs - jacobian = np.arange(num_instances * num_jacobi_bodies * 6 * num_dofs, dtype=np.float32).reshape( - num_instances, num_jacobi_bodies, 6, num_dofs - ) - mass_matrix = np.arange(num_instances * num_dofs * num_dofs, dtype=np.float32).reshape( - num_instances, num_dofs, num_dofs - ) - gravity = np.arange(num_instances * num_dofs, dtype=np.float32).reshape(num_instances, num_dofs) - return jacobian, mass_matrix, gravity - - -def _generalized_dof_user_to_backend(num_base_dofs: int, joint_user_to_backend: np.ndarray) -> np.ndarray: - """Return generalized DoF user-to-backend indices including the floating base prefix.""" - return np.concatenate( - ( - np.arange(num_base_dofs, dtype=np.int64), - num_base_dofs + joint_user_to_backend, - ) - ) - - -def _jacobian_body_user_to_backend(art, body_user_to_backend: np.ndarray) -> np.ndarray: - """Return Jacobian body-axis user-to-backend indices with the fixed root excluded.""" - if art.num_base_dofs == 0: - return np.asarray([backend_id - 1 for backend_id in body_user_to_backend if backend_id != 0], dtype=np.int64) - return body_user_to_backend - - -def _set_dynamics_ordering_backend_data( - backend: str, - art, - raw_backend, - jacobian: np.ndarray, - mass_matrix: np.ndarray, - gravity: np.ndarray, -) -> None: - """Write deterministic backend-order dynamics data into the backend mock.""" - if backend == "physx": - raw_backend.set_mock_jacobians(wp.array(jacobian, dtype=wp.float32, device=art.device)) - raw_backend.set_mock_generalized_mass_matrices(wp.array(mass_matrix, dtype=wp.float32, device=art.device)) - raw_backend.set_mock_gravity_compensation_forces(wp.array(gravity, dtype=wp.float32, device=art.device)) - elif backend == "newton": - if art.is_fixed_base: - model_jacobian = np.zeros((art.num_instances, art.num_bodies, 6, art.num_joints), dtype=np.float32) - model_jacobian[:, 1:] = jacobian - else: - model_jacobian = jacobian - raw_backend.set_mock_jacobians(wp.array(model_jacobian, dtype=wp.float32, device=art.device)) - raw_backend.set_mock_mass_matrices(wp.array(mass_matrix, dtype=wp.float32, device=art.device)) - else: - raise AssertionError(f"Unsupported backend for dynamics-ordering test: {backend}") - - -def _set_body_ordering_backend_data( - backend: str, - art, - raw_backend, - root_pose: np.ndarray, - root_vel: np.ndarray, - link_pose: np.ndarray, - com_pose_b: np.ndarray, - body_com_vel: np.ndarray, - body_acc: np.ndarray, -) -> None: - """Write deterministic backend-order body state into the backend mock.""" - if backend == "physx": - raw_backend._root_transforms = wp.array(root_pose, dtype=wp.float32, device=art.device) - raw_backend._root_velocities = wp.array(root_vel, dtype=wp.float32, device=art.device) - raw_backend._link_transforms = wp.array(link_pose, dtype=wp.float32, device=art.device) - raw_backend._link_velocities = wp.array(body_com_vel, dtype=wp.float32, device=art.device) - raw_backend._link_accelerations = wp.array(body_acc, dtype=wp.float32, device=art.device) - raw_backend._coms = wp.array(com_pose_b, dtype=wp.float32, device="cpu") - elif backend == "ovphysx": - from isaaclab_ov import tensor_types as TT - - raw_backend.bindings[TT.ROOT_POSE]._data = root_pose.copy() - raw_backend.bindings[TT.ROOT_VELOCITY]._data = root_vel.copy() - raw_backend.bindings[TT.LINK_POSE]._data = link_pose.copy() - raw_backend.bindings[TT.LINK_VELOCITY]._data = body_com_vel.copy() - raw_backend.bindings[TT.LINK_ACCELERATION]._data = body_acc.copy() - raw_backend.bindings[TT.BODY_COM_POSE]._data = com_pose_b.copy() - art.data._body_com_pose_b.timestamp = -1.0 - if art.data._body_com_pose_b_backend is not None: - art.data._body_com_pose_b_backend.timestamp = -1.0 - elif backend == "newton": - root_pose_wp = wp.array(root_pose[:, None, :], dtype=wp.transformf, device=art.device) - root_vel_wp = wp.array(root_vel[:, None, :], dtype=wp.spatial_vectorf, device=art.device) - link_pose_wp = wp.array(link_pose[:, None, :, :], dtype=wp.transformf, device=art.device) - body_com_vel_wp = wp.array(body_com_vel[:, None, :, :], dtype=wp.spatial_vectorf, device=art.device) - body_com_pos_wp = wp.array(com_pose_b[:, None, :, :3], dtype=wp.vec3f, device=art.device) - raw_backend.set_mock_root_transforms(root_pose_wp) - raw_backend.set_mock_root_velocities(root_vel_wp) - raw_backend.set_mock_link_transforms(link_pose_wp) - raw_backend.set_mock_link_velocities(body_com_vel_wp) - raw_backend.set_mock_coms(body_com_pos_wp) - art.data._sim_bind_root_link_pose_w.assign(root_pose_wp[:, 0]) - art.data._sim_bind_root_com_vel_w.assign(root_vel_wp[:, 0]) - art.data._sim_bind_body_link_pose_w.assign(link_pose_wp[:, 0]) - art.data._sim_bind_body_com_vel_w.assign(body_com_vel_wp[:, 0]) - art.data._sim_bind_body_com_pos_b.assign(body_com_pos_wp[:, 0]) - art.data._previous_body_com_vel.assign( - wp.zeros((art.num_instances, art.num_bodies), dtype=wp.spatial_vectorf, device=art.device) - ) - else: - raise AssertionError(f"Unsupported backend for body-ordering test: {backend}") - - -def _set_rotated_body_poses(backend: str, art, raw_backend) -> None: - """Use distinct quarter-turn Z rotations and nonzero positions for wrench packing.""" - poses = np.zeros((art.num_instances, art.num_bodies, 7), dtype=np.float32) - poses[..., 0] = np.arange(art.num_bodies) + 10.0 - half_angles = np.arange(art.num_bodies) * (np.pi / 4) - poses[..., 5] = np.sin(half_angles) - poses[..., 6] = np.cos(half_angles) - if backend == "newton": - poses_wp = wp.array(poses[:, None], dtype=wp.transformf, device=art.device) - raw_backend.set_mock_link_transforms(poses_wp) - art.data._sim_bind_body_link_pose_w.assign(poses_wp[:, 0]) - art.data._refresh_user_order_body_state() - elif backend == "ovphysx": - from isaaclab_ov import tensor_types as TT - - raw_backend.bindings[TT.LINK_POSE]._data = poses - art.data._reset_pose() - - -def _read_backend_wrench(backend: str, art, raw_backend, captured: dict) -> tuple[np.ndarray, np.ndarray]: - """Return force and torque from the concrete backend write target.""" - if backend == "physx": - return captured["force"], captured["torque"] - if backend == "newton": - wrench = art.data._sim_bind_body_external_wrench.numpy() - return wrench[..., :3], wrench[..., 3:6] - from isaaclab_ov import tensor_types as TT - - wrench = raw_backend.bindings[TT.LINK_WRENCH]._data - return wrench[..., :3], wrench[..., 3:6] - - -def _seed_root_writer_backend_com(backend: str, art, raw_backend, backend_coms: np.ndarray) -> None: - """Seed root-writer COM inputs in backend order and verify the public view.""" - if backend in ("physx", "ovphysx"): - _seed_backend_com_poses(backend, art, raw_backend, backend_coms) - backend_actual = art.data._backend_body_com_pose_b.numpy() - public_actual = art.data.body_com_pose_b.warp.numpy() - expected_public = backend_coms[:, _ordering_user_to_backend(art.body_ordering, art.num_bodies)] - elif backend == "newton": - backend_positions = backend_coms[..., :3] - art.data._sim_bind_body_com_pos_b.assign(wp.array(backend_positions, dtype=wp.vec3f, device=art.device)) - if art.body_ordering is not None: - art.data._sync_user_ordered_body_property_buffers() - backend_actual = art.data._sim_bind_body_com_pos_b.numpy() - public_actual = art.data.body_com_pos_b.warp.numpy() - expected_public = backend_positions[:, _ordering_user_to_backend(art.body_ordering, art.num_bodies)] - backend_coms = backend_positions - else: - raise AssertionError(f"Unsupported backend for root-writer parity test: {backend}") - - np.testing.assert_array_equal(backend_actual, backend_coms) - np.testing.assert_array_equal(public_actual, expected_public) - - -def _write_root_writer_command( - art, - selection: str, - operation: str, - root_com_pose: np.ndarray, - root_link_velocity: np.ndarray, -) -> None: - """Write one root command through an index or mask API.""" - if selection == "index": - env_ids = wp.array([0], dtype=wp.int32, device=art.device) - if operation == "com_pose": - art.write_root_com_pose_to_sim_index( - root_pose=wp.array(root_com_pose[:1], dtype=wp.transformf, device=art.device), - env_ids=env_ids, - ) - else: - art.write_root_link_pose_to_sim_index( - root_pose=wp.array(root_com_pose[:1], dtype=wp.transformf, device=art.device), - env_ids=env_ids, - ) - art.write_root_link_velocity_to_sim_index( - root_velocity=wp.array(root_link_velocity[:1], dtype=wp.spatial_vectorf, device=art.device), - env_ids=env_ids, - ) - elif selection == "mask": - env_mask = wp.array([True, False], dtype=wp.bool, device=art.device) - if operation == "com_pose": - art.write_root_com_pose_to_sim_mask( - root_pose=wp.array(root_com_pose, dtype=wp.transformf, device=art.device), - env_mask=env_mask, - ) - else: - art.write_root_link_pose_to_sim_mask( - root_pose=wp.array(root_com_pose, dtype=wp.transformf, device=art.device), - env_mask=env_mask, - ) - art.write_root_link_velocity_to_sim_mask( - root_velocity=wp.array(root_link_velocity, dtype=wp.spatial_vectorf, device=art.device), - env_mask=env_mask, - ) - else: - raise AssertionError(f"Unsupported root-writer selection: {selection}") - - -def _read_backend_root_state(backend: str, art, raw_backend) -> tuple[np.ndarray, np.ndarray]: - """Read backend root link pose and COM velocity after a writer call.""" - if backend == "physx": - return ( - raw_backend._root_transforms.numpy().copy(), - art.data.root_com_vel_w.warp.numpy().copy(), - ) - if backend == "ovphysx": - from isaaclab_ov import tensor_types as TT - - return ( - raw_backend.bindings[TT.ROOT_POSE]._data.copy(), - raw_backend.bindings[TT.ROOT_VELOCITY]._data.copy(), - ) - if backend == "newton": - return ( - art.data._sim_bind_root_link_pose_w.numpy().copy(), - art.data._sim_bind_root_com_vel_w.numpy().copy(), - ) - raise AssertionError(f"Unsupported backend for root-writer parity test: {backend}") - - -def _clone_proxy_tensor(arr) -> torch.Tensor: - """Return a CPU clone of a proxy array's torch view.""" - return arr.torch.detach().cpu().clone() - - -def _assert_proxy_close(actual, expected: torch.Tensor) -> None: - """Assert a proxy array is close to a CPU tensor.""" - torch.testing.assert_close(actual.torch.cpu(), expected, rtol=1.0e-5, atol=1.0e-5) - - -def _clone_backend_tensor(array) -> torch.Tensor: - """Return a CPU tensor clone from a backend Warp or NumPy array.""" - if isinstance(array, np.ndarray): - return torch.from_numpy(array.copy()) - return wp.to_torch(array).detach().cpu().clone() - - -def _get_backend_joint_property_tensors(backend: str, art, raw_backend) -> dict[str, torch.Tensor]: - """Return backend-order joint property values for ordering parity checks.""" - if backend == "physx": - friction_properties = _clone_backend_tensor(raw_backend.get_dof_friction_properties()) - return { - "stiffness": _clone_backend_tensor(raw_backend.get_dof_stiffnesses()), - "damping": _clone_backend_tensor(raw_backend.get_dof_dampings()), - "armature": _clone_backend_tensor(raw_backend.get_dof_armatures()), - "position_limits": _clone_backend_tensor(raw_backend.get_dof_limits()), - "velocity_limits": _clone_backend_tensor(raw_backend.get_dof_max_velocities()), - "effort_limits": _clone_backend_tensor(raw_backend.get_dof_max_forces()), - "friction": friction_properties[..., 0], - "dynamic_friction": friction_properties[..., 1], - "viscous_friction": friction_properties[..., 2], - } - if backend == "ovphysx": - from isaaclab_ov import tensor_types as TT - - friction_properties = _clone_backend_tensor(raw_backend.bindings[TT.DOF_FRICTION_PROPERTIES]._data) - return { - "stiffness": _clone_backend_tensor(raw_backend.bindings[TT.DOF_STIFFNESS]._data), - "damping": _clone_backend_tensor(raw_backend.bindings[TT.DOF_DAMPING]._data), - "armature": _clone_backend_tensor(raw_backend.bindings[TT.DOF_ARMATURE]._data), - "position_limits": _clone_backend_tensor(raw_backend.bindings[TT.DOF_LIMIT]._data), - "velocity_limits": _clone_backend_tensor(raw_backend.bindings[TT.DOF_MAX_VELOCITY]._data), - "effort_limits": _clone_backend_tensor(raw_backend.bindings[TT.DOF_MAX_FORCE]._data), - "friction": friction_properties[..., 0], - "dynamic_friction": friction_properties[..., 1], - "viscous_friction": friction_properties[..., 2], - } - if backend == "newton": - return { - "stiffness": _clone_backend_tensor(art.data._sim_bind_joint_stiffness_sim), - "damping": _clone_backend_tensor(art.data._sim_bind_joint_damping_sim), - "armature": _clone_backend_tensor(art.data._sim_bind_joint_armature), - "position_limits": torch.stack( - ( - _clone_backend_tensor(art.data._sim_bind_joint_pos_limits_lower), - _clone_backend_tensor(art.data._sim_bind_joint_pos_limits_upper), - ), - dim=-1, - ), - "velocity_limits": _clone_backend_tensor(art.data._sim_bind_joint_vel_limits_sim), - "effort_limits": _clone_backend_tensor(art.data._sim_bind_joint_effort_limits_sim), - "friction": _clone_backend_tensor(art.data._sim_bind_joint_friction_coeff), - "viscous_friction": _clone_backend_tensor(art.data._sim_bind_joint_viscous_friction_coeff), - } - raise AssertionError(f"Unsupported backend for joint-property ordering test: {backend}") - - -def _get_backend_body_property_tensors(backend: str, art, raw_backend) -> dict[str, torch.Tensor]: - """Return backend-order body property values for ordering parity checks.""" - if backend == "physx": - return { - "mass": _clone_backend_tensor(raw_backend.get_masses()), - "inertia": _clone_backend_tensor(raw_backend.get_inertias()), - } - if backend == "ovphysx": - from isaaclab_ov import tensor_types as TT - - return { - "mass": _clone_backend_tensor(raw_backend.bindings[TT.BODY_MASS]._data), - "inertia": _clone_backend_tensor(raw_backend.bindings[TT.BODY_INERTIA]._data), - } - if backend == "newton": - return { - "mass": _clone_backend_tensor(art.data._sim_bind_body_mass), - "inertia": _clone_backend_tensor(art.data._sim_bind_body_inertia), - } - raise AssertionError(f"Unsupported backend for body-property ordering test: {backend}") - - -def _backend_param(backend: str, *values, **kwargs) -> ParameterSet: - """Build a backend parameter that skips unavailable plugins at collection time.""" - marks = list(kwargs.pop("marks", ())) - marks.append(pytest.mark.skipif(backend not in BACKENDS, reason=_backend_unavailable_reason(backend))) - return pytest.param(backend, *values, marks=marks, **kwargs) - - -def _backend_unavailable_reason(backend: str) -> str: - """Describe why an optional articulation backend could not be imported.""" - label = {"physx": "PhysX", "ovphysx": "OVPhysX", "newton": "Newton"}.get(backend, backend) - reason = f"{label} backend is not available" - if detail := BACKEND_UNAVAILABLE_REASONS.get(backend): - reason = f"{reason}: {detail}" - return reason - - -_requires_physx = pytest.mark.skipif("physx" not in BACKENDS, reason=_backend_unavailable_reason("physx")) -_requires_ovphysx = pytest.mark.skipif("ovphysx" not in BACKENDS, reason=_backend_unavailable_reason("ovphysx")) -_requires_newton = pytest.mark.skipif("newton" not in BACKENDS, reason=_backend_unavailable_reason("newton")) -_all_backends = pytest.mark.parametrize( - "backend", [_backend_param(backend) for backend in ("physx", "ovphysx", "newton")], indirect=False -) -_physx_ovphysx_backends = pytest.mark.parametrize( - "backend", [_backend_param(backend) for backend in ("physx", "ovphysx")], indirect=False -) -_dynamics_ordering_backends = pytest.mark.parametrize( - "backend", [_backend_param(backend) for backend in ("physx", "newton")], indirect=False -) -_non_mock_backends = _all_backends - - -class TestArticulationDataBodyState: - """Test data properties for all body states.""" - - @_non_mock_backends - @pytest.mark.parametrize("num_instances, num_joints, num_bodies", [(2, 1, 3)]) - @pytest.mark.parametrize("device", ["cpu"]) - def test_reversed_body_ordering_reorders_public_body_quantities( - self, backend, num_instances, num_joints, num_bodies, device - ): - identity_art, identity_raw = get_articulation(backend, num_instances, num_joints, num_bodies, device=device) - ordered_art, ordered_raw = get_articulation(backend, num_instances, num_joints, num_bodies, device=device) - raw_data = _make_body_ordering_backend_data(num_instances, num_bodies) - _set_body_ordering_backend_data(backend, identity_art, identity_raw, *raw_data) - _set_body_ordering_backend_data(backend, ordered_art, ordered_raw, *raw_data) - user_to_backend = _install_test_body_ordering(ordered_art) - - identity_art.data.update(dt=0.01) - ordered_art.data.update(dt=0.01) - - identity_body_com_pose_b = _clone_proxy_tensor(identity_art.data.body_com_pose_b) - identity_body_com_acc_w = _clone_proxy_tensor(identity_art.data.body_com_acc_w) - identity_body_link_vel_w = _clone_proxy_tensor(identity_art.data.body_link_vel_w) - identity_body_com_pose_w = _clone_proxy_tensor(identity_art.data.body_com_pose_w) - identity_root_com_pose_w = _clone_proxy_tensor(identity_art.data.root_com_pose_w) - identity_root_link_vel_w = _clone_proxy_tensor(identity_art.data.root_link_vel_w) - - _assert_proxy_close(ordered_art.data.body_com_pose_b, identity_body_com_pose_b[:, user_to_backend]) - _assert_proxy_close(ordered_art.data.body_com_acc_w, identity_body_com_acc_w[:, user_to_backend]) - _assert_proxy_close(ordered_art.data.body_link_vel_w, identity_body_link_vel_w[:, user_to_backend]) - _assert_proxy_close(ordered_art.data.body_com_pose_w, identity_body_com_pose_w[:, user_to_backend]) - _assert_proxy_close(ordered_art.data.root_com_pose_w, identity_root_com_pose_w) - _assert_proxy_close(ordered_art.data.root_link_vel_w, identity_root_link_vel_w) - - @_non_mock_backends - @pytest.mark.parametrize("ordering_mode", ["reversed", "cyclic"]) - @pytest.mark.parametrize("num_instances, num_joints, num_bodies", [(2, 1, 3)]) - @pytest.mark.parametrize("device", ["cpu"]) - def test_body_ordering_reorders_public_body_properties( - self, backend, ordering_mode, num_instances, num_joints, num_bodies, device - ): - """Expose mass and every inertia component under the matching public body name.""" - body_ordering = _body_ordering_for_mode(ordering_mode, num_bodies) - art, raw_backend = get_articulation( - backend, - num_instances, - num_joints, - num_bodies, - device=device, - body_ordering=body_ordering, - ) - user_to_backend = np.asarray(art.body_ordering.user_to_backend_indices, dtype=np.int64) - backend_properties = _get_backend_body_property_tensors(backend, art, raw_backend) - - _assert_proxy_close(art.data.body_mass, backend_properties["mass"][:, user_to_backend]) - _assert_proxy_close( - art.data.body_inertia, - backend_properties["inertia"][:, user_to_backend], - ) - - @_requires_ovphysx - def test_ovphysx_reversed_body_ordering_rereads_backend_shadows_after_reset(self): - """Refresh OVPhysX backend shadow buffers after same-step pose/velocity invalidation.""" - num_instances = 2 - num_joints = 1 - num_bodies = 3 - art, raw_backend = get_articulation("ovphysx", num_instances, num_joints, num_bodies, device="cpu") - raw_data = _make_body_ordering_backend_data(num_instances, num_bodies) - _set_body_ordering_backend_data("ovphysx", art, raw_backend, *raw_data) - user_to_backend = _install_test_body_ordering(art) - art.data.update(dt=0.01) - - _ = _clone_proxy_tensor(art.data.body_link_pose_w) - _ = _clone_proxy_tensor(art.data.body_com_vel_w) - - from isaaclab_ov import tensor_types as TT - - next_link_pose = raw_data[2].copy() - next_link_pose[..., 0] += 1000.0 - next_body_com_vel = raw_data[4].copy() - next_body_com_vel[..., 0] += 500.0 - raw_backend.bindings[TT.LINK_POSE]._data = next_link_pose - raw_backend.bindings[TT.LINK_VELOCITY]._data = next_body_com_vel - - art.data._reset_pose() - art.data._reset_velocity() - - expected_link_pose = torch.from_numpy(next_link_pose[:, user_to_backend]) - expected_body_com_vel = torch.from_numpy(next_body_com_vel[:, user_to_backend]) - _assert_proxy_close(art.data.body_link_pose_w, expected_link_pose) - _assert_proxy_close(art.data.body_com_vel_w, expected_body_com_vel) - - @_requires_ovphysx - def test_ovphysx_reversed_body_ordering_rereads_all_velocity_shadows_after_reset(self): - """Refresh every OVPhysX velocity shadow after a same-step reset.""" - art, raw_backend = get_articulation( - "ovphysx", 2, 1, 3, device="cpu", body_ordering=("body_2", "body_1", "body_0") - ) - from isaaclab_ov import tensor_types as TT - - raw_data = list(_make_body_ordering_backend_data(2, 3)) - raw_data[3][..., :3] = 0.0 - _set_body_ordering_backend_data("ovphysx", art, raw_backend, *raw_data) - user_to_backend = np.asarray(art.body_ordering.user_to_backend_indices, dtype=np.int64) - art.data.update(0.01) - art.data.body_com_vel_w.torch.clone() - art.data.body_link_vel_w.torch.clone() - art.data.root_link_vel_w.torch.clone() - - next_velocity = raw_data[4].copy() - next_velocity[..., 0] += 500.0 - raw_backend.bindings[TT.LINK_VELOCITY]._data = next_velocity - art.data._reset_velocity() - - expected = torch.from_numpy(next_velocity[:, user_to_backend]) - _assert_proxy_close(art.data.body_com_vel_w, expected) - _assert_proxy_close(art.data.body_link_vel_w, expected) - _assert_proxy_close(art.data.root_link_vel_w, torch.from_numpy(next_velocity[:, 0])) - - @_requires_ovphysx - def test_ovphysx_ordered_root_and_body_velocity_share_one_backend_read(self) -> None: - """Read the shared backend link-velocity binding once per timestamp.""" - art, raw_backend = get_articulation( - "ovphysx", - 2, - 1, - 3, - device="cpu", - body_ordering=("body_2", "body_1", "body_0"), - ) - raw_data = _make_body_ordering_backend_data(2, 3) - _set_body_ordering_backend_data("ovphysx", art, raw_backend, *raw_data) - art.data.update(0.01) - binding_read = MagicMock(wraps=art.data._binding_read) - art.data._binding_read = binding_read - - _ = art.data.body_com_vel_w - _ = art.data.root_link_vel_w - - from isaaclab_ov import tensor_types as TT - - link_velocity_reads = [call for call in binding_read.call_args_list if call.args[0] == TT.LINK_VELOCITY] - assert len(link_velocity_reads) == 1 - - @_dynamics_ordering_backends - @pytest.mark.parametrize("num_instances, num_joints, num_bodies", [(2, 3, 4)]) - @pytest.mark.parametrize("device", ["cpu"]) - @pytest.mark.parametrize("is_fixed_base", [False, True], ids=["floating", "fixed"]) - def test_ordering_reorders_public_dynamics_quantities( - self, backend, num_instances, num_joints, num_bodies, device, is_fixed_base - ): - identity_art, identity_raw = get_articulation( - backend, num_instances, num_joints, num_bodies, device=device, is_fixed_base=is_fixed_base - ) - ordered_art, ordered_raw = get_articulation( - backend, num_instances, num_joints, num_bodies, device=device, is_fixed_base=is_fixed_base - ) - raw_body_data = _make_body_ordering_backend_data(num_instances, num_bodies) - _set_body_ordering_backend_data(backend, identity_art, identity_raw, *raw_body_data) - _set_body_ordering_backend_data(backend, ordered_art, ordered_raw, *raw_body_data) - - num_base_dofs = identity_art.num_base_dofs - num_jacobi_bodies = num_bodies - (1 if num_base_dofs == 0 else 0) - raw_dynamics_data = _make_dynamics_ordering_backend_data( - num_instances, num_joints, num_jacobi_bodies, num_base_dofs - ) - _set_dynamics_ordering_backend_data(backend, identity_art, identity_raw, *raw_dynamics_data) - _set_dynamics_ordering_backend_data(backend, ordered_art, ordered_raw, *raw_dynamics_data) - if backend != "newton": - for name in ("body_com_jacobian_w", "mass_matrix", "gravity_compensation_forces"): - getattr(ordered_art.data, name) - # Cyclic orderings are not involutions, so a user_to_backend/backend_to_user swap shows up. - body_user_to_backend = _install_test_body_ordering(ordered_art, "cyclic") - joint_user_to_backend = _install_test_joint_ordering(ordered_art, "cyclic") - - identity_art.data.update(dt=0.01) - ordered_art.data.update(dt=0.01) - - if backend == "newton": - import isaaclab_newton.assets.articulation.articulation_data as newton_data_module - - original_simulation_manager = newton_data_module.SimulationManager - newton_data_module.SimulationManager = identity_art._test_simulation_manager - else: - newton_data_module = None - original_simulation_manager = None - - try: - jacobian_body_indices = _jacobian_body_user_to_backend(ordered_art, body_user_to_backend) - dof_indices = _generalized_dof_user_to_backend(num_base_dofs, joint_user_to_backend) - identity_body_com_jacobian_w = _clone_proxy_tensor(identity_art.data.body_com_jacobian_w) - identity_body_link_jacobian_w = _clone_proxy_tensor(identity_art.data.body_link_jacobian_w) - identity_mass_matrix = _clone_proxy_tensor(identity_art.data.mass_matrix) - - expected_body_com_jacobian_w = identity_body_com_jacobian_w[:, jacobian_body_indices][:, :, :, dof_indices] - expected_body_link_jacobian_w = identity_body_link_jacobian_w[:, jacobian_body_indices][ - :, :, :, dof_indices - ] - expected_mass_matrix = identity_mass_matrix[:, dof_indices][:, :, dof_indices] - - _assert_proxy_close(ordered_art.data.body_com_jacobian_w, expected_body_com_jacobian_w) - _assert_proxy_close(ordered_art.data.body_link_jacobian_w, expected_body_link_jacobian_w) - _assert_proxy_close(ordered_art.data.mass_matrix, expected_mass_matrix) - - if backend == "physx": - identity_gravity_compensation_forces = _clone_proxy_tensor( - identity_art.data.gravity_compensation_forces - ) - expected_gravity_compensation_forces = identity_gravity_compensation_forces[:, dof_indices] - _assert_proxy_close( - ordered_art.data.gravity_compensation_forces, - expected_gravity_compensation_forces, - ) - finally: - if newton_data_module is not None: - newton_data_module.SimulationManager = original_simulation_manager - - -class TestArticulationOrderingAllocation: - """Test that ordering-only shadows exist only for nonidentity mappings.""" - - _DATA_SHADOWS = { - "physx": ( - "_joint_pos_backend", - "_joint_vel_backend", - "_joint_stiffness_backend", - "_joint_damping_backend", - "_joint_armature_backend", - "_joint_pos_limits_backend", - "_joint_vel_limits_backend", - "_joint_effort_limits_backend", - "_joint_friction_props_backend", - "_joint_friction_props_user", - "_body_com_pose_b_backend", - "_body_mass_backend", - "_body_inertia_backend", - ), - "ovphysx": ( - "_body_link_pose_w_backend", - "_body_com_pose_b_backend", - "_body_com_vel_w_backend", - "_body_com_acc_w_backend", - "_joint_pos_backend", - "_joint_vel_backend", - "_joint_stiffness_backend", - "_joint_damping_backend", - "_joint_armature_backend", - "_joint_pos_limits_backend", - "_joint_vel_limits_backend", - "_joint_effort_limits_backend", - "_joint_friction_props_backend", - "_body_mass_backend", - "_body_inertia_backend", - ), - "newton": ( - "_joint_pos_user", - "_joint_vel_user", - "_joint_stiffness_user", - "_joint_damping_user", - "_joint_armature_user", - "_joint_friction_coeff_user", - "_joint_viscous_friction_user", - "_joint_pos_limits_lower_user", - "_joint_pos_limits_upper_user", - "_joint_vel_limits_user", - "_joint_effort_limits_user", - "_body_link_pose_w_user", - "_body_com_vel_w_user", - "_body_mass_user", - "_body_inertia_user", - "_body_com_pos_b_user", - ), - } - _ARTICULATION_SHADOWS = { - "physx": ( - "_joint_pos_target_backend", - "_joint_vel_target_backend", - "_joint_effort_target_backend", - "_body_wrench_force_backend", - "_body_wrench_torque_backend", - ), - "ovphysx": ( - "_joint_pos_target_backend", - "_joint_vel_target_backend", - "_joint_effort_target_backend", - ), - "newton": (), - } - - @_all_backends - @pytest.mark.parametrize("ordering_mode", ["none", "identity", "reversed"]) - def test_backends_allocate_shadows_only_for_nonidentity_ordering(self, backend: str, ordering_mode: str) -> None: - num_instances = 2 - num_joints = 3 - num_bodies = 4 - art, _ = get_articulation( - backend, - num_instances, - num_joints, - num_bodies, - device="cpu", - joint_ordering=_joint_ordering_for_mode(ordering_mode, num_joints), - body_ordering=_body_ordering_for_mode(ordering_mode, num_bodies), - ) - - expected_ordering = ordering_mode == "reversed" - assert (art.joint_ordering is not None) is expected_ordering - assert (art.body_ordering is not None) is expected_ordering - if not expected_ordering: - # Default and explicit backend orders keep the backend names as the public names. - assert art.joint_names == list(art.backend_joint_names) - assert art.body_names == list(art.backend_body_names) - # The asset ordering properties delegate to the single source on data. - assert art.joint_ordering is art.data.joint_ordering - assert art.body_ordering is art.data.body_ordering - - for owner, field_names in ( - (art.data, self._DATA_SHADOWS[backend]), - (art, self._ARTICULATION_SHADOWS[backend]), - ): - for field_name in field_names: - shadow = getattr(owner, field_name) - if not expected_ordering: - assert shadow is None, f"{backend} {ordering_mode} unexpectedly allocated {field_name}" - continue - assert shadow is not None, f"{backend} reversed ordering did not allocate {field_name}" - item_count = num_bodies if "_body_" in field_name else num_joints - assert _ordering_shadow_shape(shadow)[:2] == (num_instances, item_count) - - -class TestArticulationOrderingComWrites: - """Test coherent partial COM writes and step-independent static COM caches.""" - - @pytest.mark.parametrize("ordering_mode", ["none", "reversed"]) - @pytest.mark.parametrize("selection", ["index", "mask"]) - # ``one_env_one_item`` is dropped: env selection is always explicit for both index and mask - # writers, so it exercises the same branch as ``all_envs_one_item`` (both partial-body). The - # retained pair keeps one partial-body and one full-body write. - @pytest.mark.parametrize( - "coverage", - ["all_envs_one_item", "one_env_all_items"], - ) - @_requires_ovphysx - def test_ovphysx_partial_com_write_preserves_backend_rows( - self, - ordering_mode: str, - selection: str, - coverage: str, - ) -> None: - num_instances = 2 - num_bodies = 4 - art, raw_backend = get_articulation( - "ovphysx", - num_instances=num_instances, - num_joints=1, - num_bodies=num_bodies, - device="cpu", - body_ordering=_body_ordering_for_mode(ordering_mode, num_bodies), - ) - backend_seed = _make_backend_com_poses(num_instances, num_bodies) - _seed_backend_com_poses("ovphysx", art, raw_backend, backend_seed) - np.testing.assert_array_equal(_read_backend_com_poses("ovphysx", art, raw_backend), backend_seed) - - user_to_backend = _ordering_user_to_backend(art.body_ordering, num_bodies) - - first_env_ids, first_body_ids = _coverage_ids(coverage, num_instances, num_bodies) - if len(first_body_ids) == num_bodies: - second_env_ids = [0] - second_body_ids = first_body_ids - else: - second_env_ids = first_env_ids - second_body_ids = [1] - write_selections = ( - (first_env_ids, first_body_ids, 1000.0), - (second_env_ids, second_body_ids, 2000.0), - ) - if ordering_mode == "reversed": - for _, body_ids, _ in write_selections: - if len(body_ids) < num_bodies: - assert all(user_to_backend[body_id] != body_id for body_id in body_ids) - - expected_backend = backend_seed.copy() - first_backend_cells = None - first_payload = None - backend_after = None - for write_index, (env_ids, body_ids, base_value) in enumerate(write_selections): - payload = _make_selected_com_poses(env_ids, body_ids, base_value) - expected_before = expected_backend.copy() - if selection == "index": - art.set_coms_index( - coms=wp.array(payload, dtype=wp.transformf, device=art.device), - env_ids=wp.array(env_ids, dtype=wp.int32, device=art.device), - body_ids=wp.array(body_ids, dtype=wp.int32, device=art.device), - ) - else: - full_coms = np.zeros_like(backend_seed) - full_coms[np.ix_(env_ids, body_ids)] = payload - env_mask = wp.array( - [env_index in env_ids for env_index in range(num_instances)], - dtype=wp.bool, - device=art.device, - ) - body_mask = None - if len(body_ids) < num_bodies: - body_mask = wp.array( - [body_index in body_ids for body_index in range(num_bodies)], - dtype=wp.bool, - device=art.device, - ) - art.set_coms_mask( - coms=wp.array(full_coms, dtype=wp.transformf, device=art.device), - env_mask=env_mask, - body_mask=body_mask, - ) - - backend_body_ids = user_to_backend[np.asarray(body_ids, dtype=np.int64)] - selected_cells = np.ix_(env_ids, backend_body_ids) - selected_mask = np.zeros((num_instances, num_bodies), dtype=bool) - selected_mask[selected_cells] = True - expected_backend[selected_cells] = payload - - backend_after = _read_backend_com_poses("ovphysx", art, raw_backend) - np.testing.assert_array_equal(backend_after, expected_backend) - np.testing.assert_array_equal(backend_after[~selected_mask], expected_before[~selected_mask]) - public_coms = art.data.body_com_pose_b.torch.detach().cpu().numpy() - np.testing.assert_array_equal(public_coms, expected_backend[:, user_to_backend]) - - if write_index == 0: - first_backend_cells = selected_cells - first_payload = payload.copy() - - assert backend_after is not None - assert first_backend_cells is not None - assert first_payload is not None - np.testing.assert_array_equal(backend_after[first_backend_cells], first_payload) - - @pytest.mark.parametrize("ordering_mode", ["none", "reversed"]) - @pytest.mark.parametrize("selection", ["index", "mask"]) - @_requires_newton - def test_newton_partial_com_write_routes_to_backend_bodies(self, ordering_mode: str, selection: str) -> None: - """Route partial Newton COM writes to backend bodies and preserve the untouched rows. - - Newton stores the body COM as a position-only ``vec3`` and its ``set_coms`` deliberately - cannot set an orientation, so the 7-component COM parity harness used for PhysX/OVPhysX does - not apply. This test exercises the same reorder-routing and preservation contract natively in - ``vec3`` against ``_sim_bind_body_com_pos_b``. - """ - num_instances = 2 - num_bodies = 4 - art, _ = get_articulation( - "newton", - num_instances=num_instances, - num_joints=1, - num_bodies=num_bodies, - device="cpu", - body_ordering=_body_ordering_for_mode(ordering_mode, num_bodies), - ) - backend_seed = _make_backend_com_poses(num_instances, num_bodies)[..., :3].copy() - art.data._sim_bind_body_com_pos_b.assign(wp.array(backend_seed, dtype=wp.vec3f, device=art.device)) - if art.body_ordering is not None: - art.data._sync_user_ordered_body_property_buffers() - - user_to_backend = _ordering_user_to_backend(art.body_ordering, num_bodies) - env_ids = list(range(num_instances)) - public_body_ids = [0, 2] - omitted_public_body_ids = [1, 3] - payload = _make_selected_com_poses(env_ids, public_body_ids, 5000.0)[..., :3].copy() - if ordering_mode == "reversed": - assert all(user_to_backend[body_id] != body_id for body_id in public_body_ids) - - if selection == "index": - art.set_coms_index( - coms=wp.array(payload, dtype=wp.vec3f, device=art.device), - env_ids=wp.array(env_ids, dtype=wp.int32, device=art.device), - body_ids=wp.array(public_body_ids, dtype=wp.int32, device=art.device), - ) - else: - full_coms = np.zeros((num_instances, num_bodies, 3), dtype=np.float32) - full_coms[:, public_body_ids] = payload - art.set_coms_mask( - coms=wp.array(full_coms, dtype=wp.vec3f, device=art.device), - env_mask=wp.ones(num_instances, dtype=wp.bool, device=art.device), - body_mask=_make_item_mask(num_bodies, public_body_ids, art.device), - ) - - backend_body_ids = user_to_backend[np.asarray(public_body_ids, dtype=np.int64)] - expected_backend = backend_seed.copy() - expected_backend[np.ix_(env_ids, backend_body_ids)] = payload - - backend_after = art.data._sim_bind_body_com_pos_b.numpy() - np.testing.assert_array_equal(backend_after, expected_backend) - omitted_backend_body_ids = user_to_backend[np.asarray(omitted_public_body_ids, dtype=np.int64)] - np.testing.assert_array_equal( - backend_after[:, omitted_backend_body_ids], - backend_seed[:, omitted_backend_body_ids], - ) - public_after = art.data.body_com_pos_b.torch.detach().cpu().numpy() - np.testing.assert_array_equal(public_after, expected_backend[:, user_to_backend]) - - @_physx_ovphysx_backends - @pytest.mark.parametrize("ordering_mode", ["none", "reversed"]) - def test_duplicate_body_ids_preserve_omitted_backend_com(self, backend: str, ordering_mode: str) -> None: - num_instances = 2 - num_bodies = 4 - art, raw_backend = get_articulation( - backend, - num_instances=num_instances, - num_joints=1, - num_bodies=num_bodies, - device="cpu", - body_ordering=_body_ordering_for_mode(ordering_mode, num_bodies), - ) - if backend == "physx": - raw_backend._noop_setters = False - backend_seed = _make_backend_com_poses(num_instances, num_bodies) - _seed_backend_com_poses(backend, art, raw_backend, backend_seed) - np.testing.assert_array_equal(_read_backend_com_poses(backend, art, raw_backend), backend_seed) - if backend == "physx": - raw_backend.get_coms = MagicMock(wraps=raw_backend.get_coms) - - user_to_backend = _ordering_user_to_backend(art.body_ordering, num_bodies) - duplicate_body_ids = [0, 1, 1, 2] - omitted_public_body_id = 3 - env_ids = list(range(num_instances)) - payload = _make_selected_com_poses(env_ids, duplicate_body_ids, 3000.0) - - art.set_coms_index( - coms=wp.array(payload, dtype=wp.transformf, device=art.device), - body_ids=wp.array(duplicate_body_ids, dtype=wp.int32, device=art.device), - ) - public_after = art.data.body_com_pose_b.torch.detach().cpu().numpy() - if backend == "physx": - # The cold partial write seeds untouched bodies with one backend read; the public read reuses it. - assert raw_backend.get_coms.call_count == 1 - - expected_backend = backend_seed.copy() - for env_offset, env_id in enumerate(env_ids): - for body_offset, public_body_id in enumerate(duplicate_body_ids): - expected_backend[env_id, user_to_backend[public_body_id]] = payload[env_offset, body_offset] - backend_after = _read_backend_com_poses(backend, art, raw_backend) - np.testing.assert_array_equal(backend_after, expected_backend) - np.testing.assert_array_equal( - backend_after[:, user_to_backend[omitted_public_body_id]], - backend_seed[:, user_to_backend[omitted_public_body_id]], - ) - np.testing.assert_array_equal(public_after, expected_backend[:, user_to_backend]) - - @_physx_ovphysx_backends - @pytest.mark.parametrize("ordering_mode", ["none", "reversed"]) - def test_duplicate_env_ids_leave_global_com_cache_invalid(self, backend: str, ordering_mode: str) -> None: - num_instances = 2 - num_bodies = 4 - art, raw_backend = get_articulation( - backend, - num_instances=num_instances, - num_joints=1, - num_bodies=num_bodies, - device="cpu", - body_ordering=_body_ordering_for_mode(ordering_mode, num_bodies), - ) - if backend == "physx": - raw_backend._noop_setters = False - backend_seed = _make_backend_com_poses(num_instances, num_bodies) - _seed_backend_com_poses(backend, art, raw_backend, backend_seed) - np.testing.assert_array_equal(_read_backend_com_poses(backend, art, raw_backend), backend_seed) - - user_to_backend = _ordering_user_to_backend(art.body_ordering, num_bodies) - duplicate_env_ids = [1, 1] - omitted_env_id = 0 - public_body_ids = list(range(num_bodies)) - payload = _make_selected_com_poses(duplicate_env_ids, public_body_ids, 4000.0) - - art.set_coms_index( - coms=wp.array(payload, dtype=wp.transformf, device=art.device), - env_ids=wp.array(duplicate_env_ids, dtype=wp.int32, device=art.device), - ) - - expected_backend = backend_seed.copy() - for env_offset, env_id in enumerate(duplicate_env_ids): - expected_backend[env_id, user_to_backend] = payload[env_offset] - backend_after = _read_backend_com_poses(backend, art, raw_backend) - np.testing.assert_array_equal(backend_after, expected_backend) - np.testing.assert_array_equal(backend_after[omitted_env_id], backend_seed[omitted_env_id]) - - public_cache_was_invalid = art.data._body_com_pose_b.timestamp < 0.0 - backend_staging = art.data._body_com_pose_b_backend - backend_cache_was_invalid = backend_staging is None or backend_staging.timestamp < 0.0 - public_after = art.data.body_com_pose_b.torch.detach().cpu().numpy() - np.testing.assert_array_equal(public_after, expected_backend[:, user_to_backend]) - assert public_cache_was_invalid - assert backend_cache_was_invalid - - @_physx_ovphysx_backends - @pytest.mark.parametrize("ordering_mode", ["none", "reversed"]) - def test_static_com_cache_does_not_follow_sim_timestamp(self, backend: str, ordering_mode: str) -> None: - num_instances = 2 - num_bodies = 4 - art, raw_backend = get_articulation( - backend, - num_instances=num_instances, - num_joints=1, - num_bodies=num_bodies, - device="cpu", - body_ordering=_body_ordering_for_mode(ordering_mode, num_bodies), - ) - root_pose = np.zeros((num_instances, 7), dtype=np.float32) - root_pose[:, :3] = np.asarray([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]], dtype=np.float32) - root_pose[:, 6] = 1.0 - backend_seed = _make_backend_com_poses(num_instances, num_bodies) - if backend == "physx": - raw_backend._root_transforms = wp.array(root_pose, dtype=wp.float32, device=art.device) - else: - from isaaclab_ov import tensor_types as TT - - raw_backend.bindings[TT.ROOT_POSE]._data = root_pose.copy() - _seed_backend_com_poses(backend, art, raw_backend, backend_seed) - np.testing.assert_array_equal(_read_backend_com_poses(backend, art, raw_backend), backend_seed) - art.data._body_com_pose_b.timestamp = -1.0 - backend_staging = art.data._body_com_pose_b_backend - if backend_staging is not None: - backend_staging.timestamp = -1.0 - - user_to_backend = _ordering_user_to_backend(art.body_ordering, num_bodies) - expected_public = backend_seed[:, user_to_backend] - expected_root = root_pose.copy() - expected_root[:, :3] += backend_seed[:, 0, :3] - - if backend == "physx": - raw_backend.get_coms = MagicMock(wraps=raw_backend.get_coms) - read_spy = raw_backend.get_coms - else: - com_binding = raw_backend.bindings[TT.BODY_COM_POSE] - com_binding.read = MagicMock(wraps=com_binding.read) - read_spy = com_binding.read - - np.testing.assert_array_equal(art.data.body_com_pose_b.torch.detach().cpu().numpy(), expected_public) - np.testing.assert_array_equal(art.data.root_com_pose_w.torch.detach().cpu().numpy(), expected_root) - assert read_spy.call_count == 1 - - for _ in range(3): - art.data.update(dt=0.01) - np.testing.assert_array_equal(art.data.body_com_pose_b.torch.detach().cpu().numpy(), expected_public) - np.testing.assert_array_equal(art.data.root_com_pose_w.torch.detach().cpu().numpy(), expected_root) - assert read_spy.call_count == 1 - - -# --------------------------------------------------------------------------- -# Tests: ArticulationData joint state and properties -# --------------------------------------------------------------------------- - - -class TestArticulationOrderingRootWriteParity: - """Test root writers against backend-order COM offsets.""" - - @pytest.mark.parametrize("operation", ["com_pose", "link_velocity"]) - @pytest.mark.parametrize( - "backend, selection", - [ - _backend_param("physx", "index", id="physx-index"), - _backend_param("physx", "mask", id="physx-mask"), - _backend_param("ovphysx", "index", id="ovphysx-index"), - _backend_param("ovphysx", "mask", id="ovphysx-mask"), - _backend_param("newton", "index", id="newton-index"), - _backend_param("newton", "mask", id="newton-mask"), - ], - ) - def test_floating_root_writers_match_identity_after_body_reordering( - self, - backend: str, - selection: str, - operation: str, - ) -> None: - num_instances = 2 - num_bodies = 4 - backend_body_names = tuple(f"body_{index}" for index in range(num_bodies)) - ordered_body_names = _body_ordering_for_mode("cyclic", num_bodies) - identity_art, identity_raw = get_articulation( - backend, - num_instances=num_instances, - num_joints=1, - num_bodies=num_bodies, - device="cpu", - is_fixed_base=False, - body_ordering=backend_body_names, - ) - ordered_art, ordered_raw = get_articulation( - backend, - num_instances=num_instances, - num_joints=1, - num_bodies=num_bodies, - device="cpu", - is_fixed_base=False, - body_ordering=ordered_body_names, - ) - # The identity-configured reference normalizes to ``None`` at install time. - assert identity_art.body_ordering is None - assert ordered_art.body_ordering is not None - assert tuple(ordered_art.body_names) == ordered_body_names - np.testing.assert_array_equal( - ordered_art.body_ordering.backend_to_user_indices, - _expected_backend_to_user(ordered_body_names, backend_body_names), - ) - if backend == "physx": - identity_raw._noop_setters = False - ordered_raw._noop_setters = False - - backend_coms = _make_backend_com_poses(num_instances, num_bodies) - _seed_root_writer_backend_com(backend, identity_art, identity_raw, backend_coms) - _seed_root_writer_backend_com(backend, ordered_art, ordered_raw, backend_coms) - - root_com_pose = np.asarray( - [[4.0, 5.0, 6.0, 0.0, 0.0, 0.0, 1.0], [14.0, 15.0, 16.0, 0.0, 0.0, 0.0, 1.0]], - dtype=np.float32, - ) - root_link_velocity = np.asarray( - [[1.0, 2.0, 3.0, 0.4, 0.5, 0.6], [11.0, 12.0, 13.0, 1.4, 1.5, 1.6]], - dtype=np.float32, - ) - _write_root_writer_command(identity_art, selection, operation, root_com_pose, root_link_velocity) - _write_root_writer_command(ordered_art, selection, operation, root_com_pose, root_link_velocity) - - identity_root_link_pose, identity_root_com_velocity = _read_backend_root_state( - backend, identity_art, identity_raw - ) - ordered_root_link_pose, ordered_root_com_velocity = _read_backend_root_state(backend, ordered_art, ordered_raw) - if operation == "com_pose": - np.testing.assert_allclose(ordered_root_link_pose[0], identity_root_link_pose[0], atol=1e-6) - else: - np.testing.assert_allclose(ordered_root_com_velocity[0], identity_root_com_velocity[0], atol=1e-6) - - -class TestArticulationOrderingWriteParity: - """Test that partial joint writes preserve unselected backend state.""" - - # Each backend selects the joint map separately in every ordered writer (position, velocity, state x - # index, mask), so every ordered writer gets one row. The ordering is cyclic: an involution such as - # 'reversed' on three joints is its own inverse and cannot tell user_to_backend from backend_to_user. - # Unordered rows keep one index and one mask writer per backend; coverage alternates across rows. - @pytest.mark.parametrize( - ("backend", "ordering_mode", "selection", "operation", "coverage"), - [ - row - for backend in ("physx", "ovphysx", "newton") - for row in ( - _backend_param(backend, "none", "index", "state", "all_envs_one_item"), - _backend_param(backend, "none", "mask", "position", "one_env_all_items"), - _backend_param(backend, "cyclic", "index", "position", "one_env_all_items"), - _backend_param(backend, "cyclic", "index", "velocity", "all_envs_one_item"), - _backend_param(backend, "cyclic", "index", "state", "one_env_all_items"), - _backend_param(backend, "cyclic", "mask", "position", "all_envs_one_item"), - _backend_param(backend, "cyclic", "mask", "velocity", "one_env_all_items"), - _backend_param(backend, "cyclic", "mask", "state", "all_envs_one_item"), - ) - ], - ) - def test_partial_joint_write_preserves_backend_rows( - self, - backend: str, - ordering_mode: str, - selection: str, - operation: str, - coverage: str, - ) -> None: - _exercise_partial_joint_write(backend, ordering_mode, selection, operation, coverage) - - @pytest.mark.parametrize( - "operation, expected_position_reads, expected_velocity_reads", - [("position", 1, 0), ("velocity", 0, 1), ("state", 1, 1)], - ) - @_requires_physx - def test_physx_mask_partial_joint_write_reads_each_staging_axis_once( - self, - operation: str, - expected_position_reads: int, - expected_velocity_reads: int, - ) -> None: - art, raw_backend = get_articulation( - "physx", - num_instances=2, - num_joints=3, - num_bodies=2, - device="cpu", - joint_ordering=_joint_ordering_for_mode("reversed", 3), - ) - initial_position = np.asarray([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]], dtype=np.float32) - initial_velocity = np.asarray([[110.0, 120.0, 130.0], [140.0, 150.0, 160.0]], dtype=np.float32) - _seed_backend_joint_state("physx", art, raw_backend, initial_position, initial_velocity) - # Spy on the public backend tensor-API seam rather than a private refresh helper: a - # partial write must stage each touched axis from the backend view exactly once so the - # unselected rows survive, independent of how the data class names its internals. - position_read = MagicMock(wraps=raw_backend.get_dof_positions) - velocity_read = MagicMock(wraps=raw_backend.get_dof_velocities) - raw_backend.get_dof_positions = position_read - raw_backend.get_dof_velocities = velocity_read - - env_ids, joint_ids, position_payload, velocity_payload = _write_selected_joint_state( - art, - "mask", - operation, - "one_env_one_item", - 901.0, - 902.0, - ) - - assert position_read.call_count == expected_position_reads - assert velocity_read.call_count == expected_velocity_reads - # One representative preservation check (exhaustive coverage lives in - # test_partial_joint_write_preserves_backend_rows): the untouched rows keep - # their seeded values while the selected cell takes the payload. - user_to_backend = np.asarray(art.joint_ordering.user_to_backend_indices, dtype=np.int64) - selected_cells = np.ix_(env_ids, user_to_backend[joint_ids]) - if operation in ("position", "state"): - initial_position[selected_cells] = position_payload - backend_position, _ = _read_backend_joint_state("physx", art, raw_backend) - np.testing.assert_array_equal(backend_position, initial_position) - - @_physx_ovphysx_backends - @pytest.mark.parametrize("operation", ["position", "velocity", "state"]) - def test_duplicate_full_length_joint_selector_preserves_newer_backend_rows( - self, backend: str, operation: str - ) -> None: - _exercise_duplicate_full_length_joint_write(backend, operation) - - @pytest.mark.parametrize("selector", ["default_mask", "explicit_full_index"]) - @pytest.mark.parametrize("operation", ["position", "velocity", "state"]) - @_requires_physx - def test_physx_joint_selector_provenance_controls_staging_read(self, operation: str, selector: str) -> None: - art, raw_backend = get_articulation( - "physx", - num_instances=2, - num_joints=3, - num_bodies=2, - device="cpu", - joint_ordering=_joint_ordering_for_mode("reversed", 3), - ) - initial_position = np.asarray([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]], dtype=np.float32) - initial_velocity = np.asarray([[110.0, 120.0, 130.0], [140.0, 150.0, 160.0]], dtype=np.float32) - _seed_backend_joint_state("physx", art, raw_backend, initial_position, initial_velocity) - # Spy on the public backend tensor-API seam. A full-length write expressed as the default - # mask (joint_ids is None) is provably complete, so it must not read the backend to stage - # unselected rows; an explicit full index selector is opaque to the writer and - # conservatively triggers exactly one staging read per touched axis. - position_read = MagicMock(wraps=raw_backend.get_dof_positions) - velocity_read = MagicMock(wraps=raw_backend.get_dof_velocities) - raw_backend.get_dof_positions = position_read - raw_backend.get_dof_velocities = velocity_read - - position = torch.tensor(initial_position[:, ::-1].copy(), dtype=torch.float32, device=art.device) - velocity = torch.tensor(initial_velocity[:, ::-1].copy(), dtype=torch.float32, device=art.device) - if selector == "default_mask": - if operation == "position": - art.write_joint_position_to_sim_mask(position=position) - elif operation == "velocity": - art.write_joint_velocity_to_sim_mask(velocity=velocity) - else: - art.write_joint_state_to_sim_mask(position=position, velocity=velocity) - else: - joint_ids = wp.array([0, 1, 2], dtype=wp.int32, device=art.device) - if operation == "position": - art.write_joint_position_to_sim_index(position=position, joint_ids=joint_ids) - elif operation == "velocity": - art.write_joint_velocity_to_sim_index(velocity=velocity, joint_ids=joint_ids) - else: - art.write_joint_state_to_sim_index(position=position, velocity=velocity, joint_ids=joint_ids) - - should_read = selector == "explicit_full_index" - assert position_read.call_count == int(should_read and operation in ("position", "state")) - assert velocity_read.call_count == int(should_read and operation in ("velocity", "state")) - - # Newton is intentionally absent from the backend list. Its partial joint writers (e.g. - # write_joint_position_to_sim_index) scatter the selected cells straight into the live sim-bound - # buffer instead of staging and rewriting a full backend row, so there is no stale full-row - # staging image that a later partial write could resurrect. The hazard this test guards against - # is structural to the PhysX/OVPhysX read-modify-write staging path and cannot occur on Newton. - @_physx_ovphysx_backends - @pytest.mark.parametrize("operation", ["position", "velocity", "state"]) - def test_stale_partial_joint_write_preserves_newer_backend_rows(self, backend: str, operation: str) -> None: - _exercise_stale_partial_joint_write(backend, operation) - - @_requires_ovphysx - def test_ovphysx_partial_effort_target_write_preserves_unselected_backend_rows(self) -> None: - """A partial effort-target write must not leak the last applied torque onto unselected joints. - - ``write_data_to_sim`` pushes the processed effort command (which can differ - from the raw commanded target, e.g. once explicit actuators clip or otherwise transform it) - to the ``DOF_ACTUATION_FORCE`` binding. That push must use its own backend-order scratch - buffer rather than ``_joint_effort_target_backend``, because the latter is also the staging - that :meth:`~isaaclab_ov.assets.articulation.Articulation.set_joint_effort_target_index` - reuses for unselected joints on a partial write. If the two were the same buffer, a partial - write would resurrect stale applied torque for every joint it did not touch. - """ - from isaaclab_ov import tensor_types as TT - - num_instances = 2 - num_joints = 3 - art, raw_backend = get_articulation( - "ovphysx", - num_instances=num_instances, - num_joints=num_joints, - num_bodies=2, - device="cpu", - joint_ordering=_joint_ordering_for_mode("reversed", num_joints), - ) - object.__setattr__(art, "_can_write_effort", True) - user_to_backend = _ordering_user_to_backend(art.joint_ordering, num_joints) - - # Persist raw effort targets through the public setter, in backend order via the - # ``_joint_effort_target_backend`` mirror. - raw_targets = np.asarray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32) - art.set_joint_effort_target_mask(target=torch.tensor(raw_targets, dtype=torch.float32, device=art.device)) - - # Simulate an actuator model computing an applied torque that differs from the raw - # target, then run one full simulation step. - applied_torque = raw_targets + 100.0 - art.actuators._joint_effort_target_sim.assign(wp.array(applied_torque, dtype=wp.float32, device=art.device)) - art.write_data_to_sim() - - pushed_after_step = raw_backend.bindings[TT.DOF_ACTUATION_FORCE]._data.copy() - np.testing.assert_array_equal(pushed_after_step, applied_torque[:, user_to_backend]) - - # A partial write on a single joint must preserve the persisted raw targets -- not the - # applied torque from the last step -- for every joint it does not touch. - joint_ids = wp.array([1], dtype=wp.int32, device=art.device) - env_ids = wp.array([0, 1], dtype=wp.int32, device=art.device) - new_target = torch.tensor([[901.0], [902.0]], dtype=torch.float32, device=art.device) - art.set_joint_effort_target_index(target=new_target, joint_ids=joint_ids, env_ids=env_ids) - - expected_targets = raw_targets.copy() - expected_targets[:, 1] = new_target.cpu().numpy()[:, 0] - expected_backend = expected_targets[:, user_to_backend] - - pushed_after_partial_write = raw_backend.bindings[TT.DOF_ACTUATION_FORCE]._data.copy() - np.testing.assert_array_equal(pushed_after_partial_write, expected_backend) - - -class TestArticulationDataJointState: - """Test data properties for joint state and joint properties.""" - - @_requires_ovphysx - def test_ovphysx_joint_acceleration_differences_public_order_velocities(self): - """Finite-difference OVPhysX joint velocity entirely in public joint order.""" - art, raw_backend = get_articulation("ovphysx", 2, 3, 2, device="cpu") - user_to_backend = _install_test_joint_ordering(art) - from isaaclab_ov import tensor_types as TT - - first = np.asarray([[1.0, 2.0, 4.0], [10.0, 20.0, 40.0]], dtype=np.float32) - second = first + np.asarray([[3.0, 5.0, 7.0], [11.0, 13.0, 17.0]], dtype=np.float32) - raw_backend.bindings[TT.DOF_VELOCITY]._data = first - art.data.update(0.1) - art.data.joint_acc.torch.clone() - raw_backend.bindings[TT.DOF_VELOCITY]._data = second - art.data.update(0.1) - - torch.testing.assert_close(art.data.joint_vel.torch, torch.from_numpy(second[:, user_to_backend])) - torch.testing.assert_close( - art.data.joint_acc.torch, - torch.from_numpy((second - first)[:, user_to_backend] / 0.1), - ) - - @_requires_ovphysx - def test_ovphysx_ordered_joint_state_is_cached_per_sim_timestamp(self): - """Gather ordered OVPhysX joint position and velocity at most once per timestamp.""" - art, _ = get_articulation( - "ovphysx", - 2, - 3, - 2, - device="cpu", - joint_ordering=("joint_2", "joint_1", "joint_0"), - ) - art.data.update(0.01) - binding_read = MagicMock(wraps=art.data._binding_read) - art.data._binding_read = binding_read - from isaaclab_ov import tensor_types as TT - - for _ in range(2): - art.data.joint_pos.torch.clone() - art.data.joint_vel.torch.clone() - - read_types = [call.args[0] for call in binding_read.call_args_list] - assert read_types.count(TT.DOF_POSITION) == 1 - assert read_types.count(TT.DOF_VELOCITY) == 1 - - @_non_mock_backends - @pytest.mark.parametrize("num_instances, num_joints, num_bodies", [(2, 3, 2)]) - @pytest.mark.parametrize("device", ["cpu"]) - def test_joint_ordering_reorders_public_joint_properties( - self, backend, num_instances, num_joints, num_bodies, device - ): - """Expose every backend joint property under the matching public joint name.""" - joint_ordering = _joint_ordering_for_mode("cyclic", num_joints) - art, raw_backend = get_articulation( - backend, - num_instances, - num_joints, - num_bodies, - device=device, - joint_ordering=joint_ordering, - ) - user_to_backend = np.asarray(art.joint_ordering.user_to_backend_indices, dtype=np.int64) - backend_properties = _get_backend_joint_property_tensors(backend, art, raw_backend) - public_properties = { - "stiffness": art.data.joint_stiffness, - "damping": art.data.joint_damping, - "armature": art.data.joint_armature, - "position_limits": art.data.joint_pos_limits, - "velocity_limits": art.data.joint_vel_limits, - "effort_limits": art.data.joint_effort_limits, - "friction": art.data.joint_friction_coeff, - } - if backend in ("physx", "ovphysx", "newton"): - public_properties.update( - { - "viscous_friction": art.data.joint_viscous_friction_coeff, - } - ) - if backend in ("physx", "ovphysx"): - public_properties["dynamic_friction"] = art.data.joint_dynamic_friction_coeff - - for property_name, public_property in public_properties.items(): - _assert_proxy_close(public_property, backend_properties[property_name][:, user_to_backend]) - - if backend == "ovphysx": - from isaaclab_ov import tensor_types as TT - - # A fresh native read must gather too; initialization and writer caches can hide a missing gather. - friction = np.arange(num_instances * num_joints * 3, dtype=np.float32).reshape(num_instances, num_joints, 3) - friction += 100.0 - raw_backend.bindings[TT.DOF_FRICTION_PROPERTIES]._data = friction - art.data._joint_friction_props_buf.timestamp = -1 - art.data._joint_friction_props_backend.timestamp = -1 - for component, property_name in enumerate( - ("joint_friction_coeff", "joint_dynamic_friction_coeff", "joint_viscous_friction_coeff") - ): - _assert_proxy_close( - getattr(art.data, property_name), torch.from_numpy(friction[:, user_to_backend, component]) - ) - - -def _make_item_mask(total: int, selected: list[int], device: str) -> wp.array: - """Create a bool Warp mask with the selected indices set.""" - mask = np.zeros(total, dtype=bool) - mask[selected] = True - return wp.array(mask, dtype=wp.bool, device=device) - - -# --------------------------------------------------------------------------- -# Tests: Articulation operations -# --------------------------------------------------------------------------- - - -class TestArticulationOperations: - """Test cross-cutting articulation operations.""" - - @_non_mock_backends - def test_legacy_position_target_accepts_joint_slice(self, backend: str) -> None: - art, _ = get_articulation(backend, num_instances=2, num_joints=4, num_bodies=2, device="cpu") - target = torch.tensor([[11.0, 12.0], [21.0, 22.0]], dtype=torch.float32) - - with pytest.warns(DeprecationWarning): - art.set_joint_position_target(target, joint_ids=slice(1, 3)) - - expected = torch.zeros((2, 4), dtype=torch.float32) - expected[:, 1:3] = target - torch.testing.assert_close(art.actuators.target_command.position.torch, expected) - - @_non_mock_backends - @pytest.mark.parametrize("ordering_mode", ["none", "reversed", "cyclic"]) - @pytest.mark.parametrize("is_fixed_base", [False, True], ids=["floating", "fixed"]) - @pytest.mark.parametrize("device", ["cpu"]) - def test_external_wrenches_are_written_in_backend_body_order(self, backend, ordering_mode, is_fixed_base, device): - """Write public-order body wrenches to each backend in backend body order.""" - num_instances, num_joints, num_bodies = 2, 1, 4 - backend_body_names = tuple(f"body_{index}" for index in range(num_bodies)) - if ordering_mode == "none": - body_ordering = None - elif is_fixed_base: - if ordering_mode == "reversed": - body_ordering = (backend_body_names[0], *reversed(backend_body_names[1:])) - else: - body_ordering = (backend_body_names[0], backend_body_names[-1], *backend_body_names[1:-1]) - else: - body_ordering = _body_ordering_for_mode(ordering_mode, num_bodies) - art, raw_backend = get_articulation( - backend, - num_instances, - num_joints, - num_bodies, - device=device, - is_fixed_base=is_fixed_base, - body_ordering=body_ordering, - ) - _set_rotated_body_poses(backend, art, raw_backend) - captured = {} - if backend == "physx": - - def capture_wrench(*, force_data, torque_data, position_data, indices, is_global): - captured["is_global"] = is_global - assert position_data is None - captured["force"] = force_data.numpy().reshape(num_instances, num_bodies, 3).copy() - captured["torque"] = torque_data.numpy().reshape(num_instances, num_bodies, 3).copy() - - raw_backend.apply_forces_and_torques_at_position = capture_wrench - - forces = np.arange(num_instances * num_bodies * 3, dtype=np.float32).reshape(num_instances, num_bodies, 3) - torques = forces + 100.0 - backend_to_user = _expected_backend_to_user(body_ordering, backend_body_names) - composer = art.instantaneous_wrench_composer - for is_global in (False, True): - composer.set_forces_and_torques_index( - forces=wp.array(forces, dtype=wp.vec3f, device=device), - torques=wp.array(torques, dtype=wp.vec3f, device=device), - is_global=is_global, - ) - with patch.object(composer, "compose_to_body_frame", wraps=composer.compose_to_body_frame) as compose: - art.write_data_to_sim() - assert compose.call_count == int(is_global and backend == "newton") - - expected_force, expected_torque = forces[:, backend_to_user], torques[:, backend_to_user] - if backend == "physx": - assert captured["is_global"] is is_global - elif not is_global: - # Known 0/90/180/270-degree rotations, independently of the backend quaternion transform. - rotations = np.asarray( - [ - [[1, 0, 0], [0, 1, 0], [0, 0, 1]], - [[0, -1, 0], [1, 0, 0], [0, 0, 1]], - [[-1, 0, 0], [0, -1, 0], [0, 0, 1]], - [[0, 1, 0], [-1, 0, 0], [0, 0, 1]], - ], - dtype=np.float32, - ) - expected_force = np.einsum("bij,nbj->nbi", rotations, expected_force) - expected_torque = np.einsum("bij,nbj->nbi", rotations, expected_torque) - backend_force, backend_torque = _read_backend_wrench(backend, art, raw_backend, captured) - np.testing.assert_allclose(backend_force, expected_force, atol=1e-4) - np.testing.assert_allclose(backend_torque, expected_torque, atol=1e-4) - if backend == "ovphysx": - from isaaclab_ov import tensor_types as TT - - packed_positions = raw_backend.bindings[TT.LINK_WRENCH]._data[..., 6:9] - np.testing.assert_array_equal(packed_positions, raw_backend.bindings[TT.LINK_POSE]._data[..., :3]) - - @_requires_ovphysx - def test_ovphysx_configured_defaults_use_public_joint_names(self): - """Resolve configured OVPhysX defaults against public joint names.""" - art, _ = get_articulation( - "ovphysx", - 1, - 3, - 2, - device="cpu", - joint_ordering=("joint_2", "joint_1", "joint_0"), - ) - zeros = wp.zeros((1, 3), dtype=wp.float32, device="cpu") - art.data._default_joint_pos.assign(zeros) - art.data._default_joint_vel.assign(zeros) - patterns = {"joint_0": 10.0, "joint_1": 20.0, "joint_2": 30.0} - - art._resolve_joint_values(patterns, art.data._default_joint_pos) - - expected = np.asarray([[patterns[name] for name in art.joint_names]], dtype=np.float32) - np.testing.assert_array_equal(art.data.default_joint_pos.warp.numpy(), expected) - - @_requires_ovphysx - @pytest.mark.parametrize("ordering_mode", ["reversed", "cyclic"]) - def test_ovphysx_implicit_targets_are_written_in_backend_order(self, ordering_mode: str) -> None: - """Write implicit position and velocity targets under their matching backend joint names.""" - from isaaclab_ov import tensor_types as TT - - num_instances, num_joints = 2, 3 - backend_joint_names = tuple(f"joint_{index}" for index in range(num_joints)) - joint_ordering = _joint_ordering_for_mode(ordering_mode, num_joints) - art, raw_backend = get_articulation( - "ovphysx", - num_instances=num_instances, - num_joints=num_joints, - device="cpu", - joint_ordering=joint_ordering, - ) - position = np.arange(num_instances * num_joints, dtype=np.float32).reshape(num_instances, num_joints) - velocity = position + 100.0 - effort = position + 200.0 - # Seed the actuator collection's submitted command buffers (no actuator groups recompute them). - art.actuators._joint_pos_target.assign(wp.array(position, dtype=wp.float32, device=art.device)) - art.actuators._joint_vel_target.assign(wp.array(velocity, dtype=wp.float32, device=art.device)) - art.actuators._joint_effort_target_sim.assign(wp.array(effort, dtype=wp.float32, device=art.device)) - object.__setattr__(art, "_has_implicit_actuators", True) - object.__setattr__(art, "_can_write_effort", True) - - art.write_data_to_sim() - - backend_to_user = _expected_backend_to_user(joint_ordering, backend_joint_names) - np.testing.assert_array_equal( - raw_backend.bindings[TT.DOF_POSITION_TARGET]._data, - position[:, backend_to_user], - ) - np.testing.assert_array_equal( - raw_backend.bindings[TT.DOF_VELOCITY_TARGET]._data, - velocity[:, backend_to_user], - ) - np.testing.assert_array_equal( - raw_backend.bindings[TT.DOF_ACTUATION_FORCE]._data, - effort[:, backend_to_user], - ) - - @_requires_ovphysx - @pytest.mark.parametrize("selector_kind", ["torch", "warp"]) - def test_ovphysx_int64_effort_target_selector_reaches_binding( - self, selector_kind: str, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Convert indexed effort-target environment selectors for the OVPhysX binding.""" - from isaaclab_ov import tensor_types as TT - - art, raw_backend = get_articulation("ovphysx", 2, 3, 2, device="cpu") - object.__setattr__(art, "_can_write_effort", True) - expected = raw_backend.bindings[TT.DOF_ACTUATION_FORCE]._data.copy() - captured_indices = None - set_attribute = art._root_view.set_attribute - - def strict_set_attribute(name, values, *, indices=None, mask=None): - nonlocal captured_indices - if name == TT.DOF_ACTUATION_FORCE: - assert isinstance(indices, wp.array) - assert indices.dtype == wp.int32 - assert str(indices.device) == art.device - captured_indices = indices - set_attribute(name, values, indices=indices, mask=mask) - - monkeypatch.setattr(art._root_view, "set_attribute", strict_set_attribute) - env_ids_torch = torch.tensor([1], dtype=torch.int64) - env_ids = env_ids_torch if selector_kind == "torch" else wp.from_torch(env_ids_torch, dtype=wp.int64) - joint_ids = torch.tensor([2], dtype=torch.int64) - target = torch.tensor([[7.0]], dtype=torch.float32) - - art.set_joint_effort_target_index(target=target, env_ids=env_ids, joint_ids=joint_ids) - - expected[1] = 0.0 - expected[1, 2] = 7.0 - assert captured_indices is not None - np.testing.assert_array_equal(captured_indices.numpy(), [1]) - np.testing.assert_array_equal(raw_backend.bindings[TT.DOF_ACTUATION_FORCE]._data, expected) - - @_requires_physx - @pytest.mark.parametrize("ordering_mode", ["reversed", "cyclic"]) - def test_physx_newton_actuator_forces_are_written_in_backend_order(self, ordering_mode: str): - """Write Newton-actuator PhysX forces in backend joint order.""" - num_instances = 2 - num_joints = 4 - num_bodies = 2 - backend_joint_names = tuple(f"joint_{index}" for index in range(num_joints)) - joint_ordering = _joint_ordering_for_mode(ordering_mode, num_joints) - art, raw_backend = get_articulation( - "physx", - num_instances, - num_joints, - num_bodies, - device="cpu", - joint_ordering=joint_ordering, - ) - user_forces_np = np.arange(num_instances * num_joints, dtype=np.float32).reshape(num_instances, num_joints) - user_forces = wp.array(user_forces_np, dtype=wp.float32, device=art.device) - object.__setattr__(art, "_joint_effort_target_backend", wp.zeros_like(art.data.joint_effort_target.warp)) - wrapper = MagicMock() - wrapper.joint_f_2d = user_forces - object.__setattr__(art, "_physx_actuator_wrapper", wrapper) - object.__setattr__(art, "_has_newton_actuators", True) - object.__setattr__(art, "_has_implicit_actuators", False) - captured = {} - - def _capture_forces(forces, indices): - captured["forces"] = wp.clone(forces, device="cpu").numpy() - captured["indices"] = wp.clone(indices, device="cpu").numpy() - - raw_backend.set_dof_actuation_forces = _capture_forces - - art.write_data_to_sim() - - backend_to_user = _expected_backend_to_user(joint_ordering, backend_joint_names) - np.testing.assert_allclose(captured["forces"], user_forces_np[:, backend_to_user]) - - @_requires_physx - def test_physx_validate_cfg_reports_velocity_limits_in_public_joint_order(self): - """Pair public default velocities with limits for the same named joint.""" - art, raw_backend = get_articulation( - "physx", - 1, - 3, - 2, - device="cpu", - joint_ordering=("joint_2", "joint_1", "joint_0"), - ) - backend_velocity_limits = np.asarray([[1.0, 100.0, 100.0]], dtype=np.float32) - public_velocity_limits = backend_velocity_limits[:, [2, 1, 0]] - raw_backend.set_mock_dof_max_velocities(wp.array(backend_velocity_limits, dtype=wp.float32, device="cpu")) - art.data._joint_vel_limits.assign(wp.array(public_velocity_limits, dtype=wp.float32, device="cpu")) - art.data._joint_pos_limits.assign(wp.array(np.tile((-100.0, 100.0), (1, 3, 1)), dtype=wp.vec2f, device="cpu")) - art.data._default_joint_pos.assign(wp.zeros((1, 3), dtype=wp.float32, device="cpu")) - art.data._default_joint_vel.assign( - wp.array(np.asarray([[50.0, 50.0, 2.0]], dtype=np.float32), dtype=wp.float32, device="cpu") - ) - - with pytest.raises(ValueError, match=r"'joint_0': 2\.000 not in \[-1\.000, 1\.000\]"): - art._validate_cfg() - - -class TestArticulationWritersJoint: - """Test joint writers/setters with all input combinations.""" - - @_non_mock_backends - @pytest.mark.parametrize("ordering_mode", ["reversed", "cyclic"]) - @pytest.mark.parametrize("selection", ["index", "mask"]) - def test_reversed_joint_ordering_routes_property_writes_to_backend( - self, backend: str, ordering_mode: str, selection: str - ): - """Route partial public property writes to matching backend joints.""" - num_instances, num_joints, num_bodies = 2, 4, 2 - joint_ordering = _joint_ordering_for_mode(ordering_mode, num_joints) - art, raw_backend = get_articulation( - backend, - num_instances, - num_joints, - num_bodies, - device="cpu", - joint_ordering=joint_ordering, - ) - if backend == "physx": - raw_backend._noop_setters = False - - user_to_backend = np.asarray(art.joint_ordering.user_to_backend_indices, dtype=np.int64) - selected_envs = np.arange(num_instances, dtype=np.int32) - selected_joints = np.asarray([0, 2], dtype=np.int32) - property_cases = [ - ("write_joint_stiffness_to_sim", "stiffness", "stiffness"), - ("write_joint_damping_to_sim", "damping", "damping"), - ("write_joint_armature_to_sim", "armature", "armature"), - ("write_joint_position_limit_to_sim", "limits", "position_limits"), - ("write_joint_velocity_limit_to_sim", "limits", "velocity_limits"), - ("write_joint_effort_limit_to_sim", "limits", "effort_limits"), - ("write_joint_friction_coefficient_to_sim", "joint_friction_coeff", "friction"), - ] - if backend in ("physx", "ovphysx"): - property_cases.extend( - [ - ( - "write_joint_dynamic_friction_coefficient_to_sim", - "joint_dynamic_friction_coeff", - "dynamic_friction", - ), - ( - "write_joint_viscous_friction_coefficient_to_sim", - "joint_viscous_friction_coeff", - "viscous_friction", - ), - ] - ) - if backend == "newton": - property_cases.append( - ( - "write_joint_viscous_friction_coefficient_to_sim", - "joint_viscous_friction_coeff", - "viscous_friction", - ) - ) - - for case_index, (method_name, value_name, property_name) in enumerate(property_cases): - before = _get_backend_joint_property_tensors(backend, art, raw_backend)[property_name] - leading_shape = ( - (num_instances, selected_joints.size) if selection == "index" else (num_instances, num_joints) - ) - values = torch.arange(int(np.prod(leading_shape)), dtype=torch.float32).reshape(leading_shape) - values = values + 100.0 * (case_index + 1) - if property_name == "position_limits": - values = torch.stack((-values, values), dim=-1) - soft_limits_before = art.data.soft_joint_pos_limits.torch.clone() - - kwargs = {value_name: values} - if selection == "index": - kwargs.update( - { - "env_ids": wp.array(selected_envs, dtype=wp.int32, device=art.device), - "joint_ids": wp.array(selected_joints, dtype=wp.int32, device=art.device), - } - ) - else: - kwargs.update( - { - "env_mask": wp.ones(num_instances, dtype=wp.bool, device=art.device), - "joint_mask": _make_item_mask(num_joints, selected_joints.tolist(), art.device), - } - ) - - getattr(art, f"{method_name}_{selection}")(**kwargs) - - after = _get_backend_joint_property_tensors(backend, art, raw_backend)[property_name] - expected = before.clone() - backend_joint_ids = torch.as_tensor(user_to_backend[selected_joints], dtype=torch.long) - if selection == "index": - expected[:, backend_joint_ids] = values - else: - public_joint_ids = torch.as_tensor(selected_joints, dtype=torch.long) - expected[:, backend_joint_ids] = values[:, public_joint_ids] - torch.testing.assert_close(after, expected, rtol=0.0, atol=0.0) - - if property_name == "position_limits": - soft_limits_after = art.data.soft_joint_pos_limits.torch - public_joint_ids = torch.as_tensor(selected_joints, dtype=torch.long) - selected_limits = expected[:, backend_joint_ids] - center = selected_limits.mean(dim=-1) - half_range = 0.5 * (selected_limits[..., 1] - selected_limits[..., 0]) - factor = art.cfg.soft_joint_pos_limit_factor - expected_soft_limits = torch.stack( - (center - factor * half_range, center + factor * half_range), - dim=-1, - ) - torch.testing.assert_close( - soft_limits_after[:, public_joint_ids], - expected_soft_limits, - ) - unselected_joint_ids = torch.tensor([1, 3], dtype=torch.long) - torch.testing.assert_close( - soft_limits_after[:, unselected_joint_ids], - soft_limits_before[:, unselected_joint_ids], - ) - - -class TestArticulationWritersBody: - """Test body property writers/setters with all input combinations.""" - - @_non_mock_backends - @pytest.mark.parametrize("ordering_mode", ["reversed", "cyclic"]) - @pytest.mark.parametrize("selection", ["index", "mask"]) - def test_body_ordering_routes_property_writes_to_backend(self, backend: str, ordering_mode: str, selection: str): - """Route partial public property writes to matching backend bodies.""" - num_instances, num_joints, num_bodies = 2, 1, 4 - body_ordering = _body_ordering_for_mode(ordering_mode, num_bodies) - art, raw_backend = get_articulation( - backend, - num_instances, - num_joints, - num_bodies, - device="cpu", - body_ordering=body_ordering, - ) - if backend == "physx": - raw_backend._noop_setters = False - - user_to_backend = np.asarray(art.body_ordering.user_to_backend_indices, dtype=np.int64) - selected_envs = np.arange(num_instances, dtype=np.int32) - selected_bodies = np.asarray([0, 2], dtype=np.int32) - property_cases = [ - ("set_masses", "masses", "mass"), - ("set_inertias", "inertias", "inertia"), - ] - - for case_index, (method_name, value_name, property_name) in enumerate(property_cases): - before = _get_backend_body_property_tensors(backend, art, raw_backend)[property_name] - leading_shape = ( - (num_instances, selected_bodies.size) if selection == "index" else (num_instances, num_bodies) - ) - values = torch.arange(int(np.prod(leading_shape)), dtype=torch.float32).reshape(leading_shape) - values = values + 100.0 * (case_index + 1) - if property_name == "inertia": - component_offsets = torch.arange(9, dtype=torch.float32) / 10.0 - values = values.unsqueeze(-1) + component_offsets - - kwargs = {value_name: values} - if selection == "index": - kwargs.update( - { - "env_ids": wp.array(selected_envs, dtype=wp.int32, device=art.device), - "body_ids": wp.array(selected_bodies, dtype=wp.int32, device=art.device), - } - ) - else: - kwargs.update( - { - "env_mask": wp.ones(num_instances, dtype=wp.bool, device=art.device), - "body_mask": _make_item_mask(num_bodies, selected_bodies.tolist(), art.device), - } - ) - - getattr(art, f"{method_name}_{selection}")(**kwargs) - - after = _get_backend_body_property_tensors(backend, art, raw_backend)[property_name] - expected = before.clone() - backend_body_ids = torch.as_tensor(user_to_backend[selected_bodies], dtype=torch.long) - if selection == "index": - expected[:, backend_body_ids] = values - else: - public_body_ids = torch.as_tensor(selected_bodies, dtype=torch.long) - expected[:, backend_body_ids] = values[:, public_body_ids] - torch.testing.assert_close(after, expected, rtol=0.0, atol=0.0) diff --git a/source/isaaclab/test/assets/test_rigid_object_collection_iface.py b/source/isaaclab/test/assets/test_rigid_object_collection_iface.py deleted file mode 100644 index 35604f52d9f3..000000000000 --- a/source/isaaclab/test/assets/test_rigid_object_collection_iface.py +++ /dev/null @@ -1,834 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -# ignore private usage of variables warning -# pyright: reportPrivateUsage=none - -""" -Checks that the rigid object collection interfaces are consistent across backends, and are providing -the exact same data as what the base rigid object collection class advertises. All rigid object -collection interfaces need to comply with the same interface contract. - -The setup is a bit convoluted so that we can run these tests without requiring Isaac Sim or GPU simulation. -""" - -import math -from unittest.mock import patch - -import numpy as np -import pytest -import torch -import warp as wp -from _rigid_object_collection_iface_test_utils import BACKENDS, get_rigid_object_collection - -from isaaclab.test.utils import DeviceScope, test_devices - -pytestmark = pytest.mark.integration - -# Distinct instance and body counts make swapped axes visible. -_NUM_INSTANCES, _NUM_BODIES = 2, 3 - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _check_proxy_array(arr, *, expected_shape: tuple, expected_dtype: type, name: str): - """Assert that `arr` is a ProxyArray with the expected shape and dtype.""" - from isaaclab.utils.warp import ProxyArray - - assert isinstance(arr, ProxyArray), f"{name}: expected ProxyArray, got {type(arr)}" - assert arr.shape == expected_shape, f"{name}: expected shape {expected_shape}, got {arr.shape}" - assert arr.dtype == expected_dtype, f"{name}: expected dtype {expected_dtype}, got {arr.dtype}" - - -# Common parametrize decorators. Pure bookkeeping (counts, names, finders, aliases) runs on CPU only; -# getters and writers keep every test device because PhysX stages through CPU-pinned buffers on CUDA. -_backends = pytest.mark.parametrize("backend", BACKENDS, indirect=False) -_devices = pytest.mark.parametrize("device", test_devices(DeviceScope.CPU_AND_DEFAULT_CUDA)) -_index_resolution_backends = pytest.mark.parametrize( - "backend", [backend for backend in ("physx", "newton") if backend in BACKENDS], indirect=False -) -_reshape_3d_backends = pytest.mark.parametrize( - "backend", [backend for backend in ("physx", "newton", "ovphysx") if backend in BACKENDS], indirect=False -) -_production_backends = pytest.mark.parametrize( - "backend", [backend for backend in ("physx", "newton", "ovphysx") if backend in BACKENDS], indirect=False -) - - -# --------------------------------------------------------------------------- -# Writer/setter test helpers -# --------------------------------------------------------------------------- - -_WP_DTYPE_TO_TRAILING = { - wp.transformf: 7, - wp.spatial_vectorf: 6, - wp.vec2f: 2, - wp.float32: 0, -} - - -def _make_data_torch(shape: tuple, device: str, wp_dtype=wp.float32) -> torch.Tensor: - """Create valid torch test data for a given warp dtype.""" - trailing = _WP_DTYPE_TO_TRAILING[wp_dtype] - if trailing: - full_shape = (*shape, trailing) - else: - full_shape = shape - data = torch.zeros(full_shape, device=device, dtype=torch.float32) - if wp_dtype == wp.transformf: - data[..., 6] = 1.0 - elif wp_dtype == wp.float32: - data.fill_(1.0) - return data - - -def _make_data_warp(shape: tuple, device: str, wp_dtype=wp.float32) -> wp.array: - """Create valid warp test data for a given warp dtype.""" - t = _make_data_torch(shape, device, wp_dtype) - if wp_dtype == wp.float32: - return wp.from_torch(t, dtype=wp.float32) - return wp.from_torch(t.contiguous(), dtype=wp_dtype) - - -def _make_payload_torch(shape: tuple, device: str, wp_dtype=wp.float32) -> torch.Tensor: - """Create valid torch data whose entries differ per element, for writer read-back checks. - - Transforms get distinct positions and a fixed 90-degree rotation about Z. - """ - values = torch.arange(1, math.prod(shape) + 1, dtype=torch.float32, device=device).reshape(shape) + 0.25 - if wp_dtype == wp.spatial_vectorf: - return values.unsqueeze(-1) + torch.arange(6, dtype=torch.float32, device=device) / 10.0 - if wp_dtype == wp.transformf: - data = torch.zeros((*shape, 7), dtype=torch.float32, device=device) - data[..., :3] = values.unsqueeze(-1) + torch.arange(3, dtype=torch.float32, device=device) / 10.0 - data[..., 5] = data[..., 6] = 2.0**-0.5 - return data - raise ValueError(f"Unsupported payload dtype: {wp_dtype}") - - -def _make_payload_warp(shape: tuple, device: str, wp_dtype=wp.float32) -> wp.array: - """Create the per-element distinct payload of :func:`_make_payload_torch` as a warp array.""" - return wp.from_torch(_make_payload_torch(shape, device, wp_dtype).contiguous(), dtype=wp_dtype) - - -def _assert_reads_back(value, expected: torch.Tensor, name: str) -> None: - """Assert a data getter (ProxyArray) returns the values a writer just wrote.""" - torch.testing.assert_close(value.torch, expected, atol=1e-5, rtol=1e-5, msg=lambda msg: f"{name}: {msg}") - - -def _make_com_data(backend: str, shape: tuple[int, ...], device: str) -> wp.array: - """Create backend-compatible center-of-mass test data.""" - if backend == "newton": - return wp.zeros(shape, dtype=wp.vec3f, device=device) - return _make_data_warp(shape, device, wp.transformf) - - -def _prime_timestamped_properties(data, property_buffer_pairs: list[tuple[str, str]]): - """Prime public lazy properties and return their concrete timestamped buffers.""" - buffers = [] - for property_name, buffer_name in property_buffer_pairs: - getattr(data, property_name) - buffer = getattr(data, buffer_name) - assert buffer is not None, buffer_name - buffer.timestamp = data._sim_timestamp - buffers.append((buffer_name, buffer)) - return buffers - - -def _assert_buffers_stale(data, buffers) -> None: - for name, buffer in buffers: - assert buffer.timestamp < data._sim_timestamp, name - - -def _make_bad_data_torch(shape: tuple, device: str, wp_dtype=wp.float32) -> torch.Tensor: - """Create torch data with wrong leading shape for negative testing.""" - bad_shape = (shape[0] + 1,) + shape[1:] - return _make_data_torch(bad_shape, device, wp_dtype) - - -def _make_bad_data_warp(shape: tuple, device: str, wp_dtype=wp.float32) -> wp.array: - """Create warp data with wrong leading shape for negative testing.""" - bad_shape = (shape[0] + 1,) + shape[1:] - return _make_data_warp(bad_shape, device, wp_dtype) - - -def _make_env_mask(num_instances: int, device: str, partial: bool) -> wp.array | None: - """Create an env_mask: None for all envs, or a partial bool mask.""" - if not partial: - return None - mask_np = np.zeros(num_instances, dtype=bool) - mask_np[0] = True - return wp.array(mask_np, dtype=wp.bool, device=device) - - -def _make_env_ids(device: str, subset: bool) -> torch.Tensor | None: - """Create env_ids: None for all envs, or [0] for a subset.""" - if not subset: - return None - return torch.tensor([0], dtype=torch.int32, device=device) - - -def _make_body_ids(device: str, subset_ids: list[int] | None) -> torch.Tensor | None: - """Create body_ids: None for all bodies, or a list for a subset.""" - if subset_ids is None: - return None - return torch.tensor(subset_ids, dtype=torch.int32, device=device) - - -def _make_item_mask(total: int, selected: list[int], device: str) -> wp.array: - """Create a bool warp mask with True at `selected` indices, False elsewhere.""" - mask_np = np.zeros(total, dtype=bool) - for i in selected: - mask_np[i] = True - return wp.array(mask_np, dtype=wp.bool, device=device) - - -# --------------------------------------------------------------------------- -# Tests: Index resolution helpers -# --------------------------------------------------------------------------- - - -class TestCollectionIndexResolution: - """Test backend-specific index resolution helpers.""" - - @_production_backends - @_devices - def test_resolve_env_ids_handles_tensor_view_shape(self, backend, device): - obj, _ = get_rigid_object_collection(backend, num_instances=4, device=device) - - env_ids = torch.arange(4, dtype=torch.int32, device=device) - resolved_full = obj._resolve_env_ids(env_ids) - resolved_view = obj._resolve_env_ids(env_ids[:2]) - - assert resolved_full.shape[0] == 4 - assert resolved_view.shape[0] == 2 - cached = wp.to_torch(obj._ALL_ENV_INDICES) - for selection in (slice(None), slice(1, None, 2), slice(0, 0)): - resolved = wp.to_torch(obj._resolve_env_ids(selection)) - torch.testing.assert_close(resolved, cached[selection]) - assert resolved.data_ptr() == cached[selection].data_ptr() - assert resolved.stride() == cached[selection].stride() - - @_index_resolution_backends - def test_resolve_body_ids_handles_tensor_view_shape(self, backend): - obj, _ = get_rigid_object_collection(backend, num_bodies=4, device="cpu") - - body_ids = torch.arange(4, dtype=torch.int32, device="cpu") - resolved_full = obj._resolve_body_ids(body_ids) - resolved_view = obj._resolve_body_ids(body_ids[:2]) - - assert resolved_full.shape[0] == 4 - assert resolved_view.shape[0] == 2 - - -# --------------------------------------------------------------------------- -# Tests: View reshape helpers -# --------------------------------------------------------------------------- - - -class TestCollectionViewReshape: - """Test backend-specific view reshape helpers.""" - - @_reshape_3d_backends - @_devices - @pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) - def test_reshape_data_to_view_3d_accepts_torch_tensor(self, backend, device, dtype): - num_instances = 2 - num_bodies = 3 - data_dim = 4 - obj, _ = get_rigid_object_collection(backend, num_instances=num_instances, num_bodies=num_bodies, device=device) - data = torch.arange(num_instances * num_bodies * data_dim, dtype=dtype, device=device).reshape( - num_instances, num_bodies, data_dim - ) - - view = obj.reshape_data_to_view_3d(data, data_dim, device=device) - - assert isinstance(view, torch.Tensor) - assert view.dtype == dtype - assert view.device == data.device - assert view.is_contiguous() - torch.testing.assert_close(view, data.permute(1, 0, 2).reshape(num_bodies * num_instances, data_dim)) - - if dtype == torch.float32: - # Warp inputs keep returning a warp array with the same body-major layout. - warp_view = obj.reshape_data_to_view_3d(wp.from_torch(data, dtype=wp.float32), data_dim, device=device) - assert isinstance(warp_view, wp.array) - assert warp_view.shape == (num_bodies * num_instances, data_dim) - assert warp_view.dtype == wp.float32 - assert str(warp_view.device) == device - torch.testing.assert_close(wp.to_torch(warp_view), view) - - @_reshape_3d_backends - @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available") - def test_reshape_data_to_view_3d_moves_torch_tensor_to_requested_device(self, backend): - num_instances = 2 - num_bodies = 3 - data_dim = 4 - obj, _ = get_rigid_object_collection( - backend, num_instances=num_instances, num_bodies=num_bodies, device="cuda:0" - ) - data = torch.arange(num_instances * num_bodies * data_dim, dtype=torch.float32, device="cuda:0").reshape( - num_instances, num_bodies, data_dim - ) - - view = obj.reshape_data_to_view_3d(data, data_dim, device="cpu") - - assert view.device.type == "cpu" - torch.testing.assert_close(view, data.permute(1, 0, 2).reshape(num_bodies * num_instances, data_dim).cpu()) - - -# --------------------------------------------------------------------------- -# Tests: Collection properties and finders -# --------------------------------------------------------------------------- - - -class TestCollectionProperties: - """Test that collection properties return the correct types/values.""" - - @_backends - def test_collection_counts_and_names(self, backend): - from isaaclab.assets.rigid_object_collection.base_rigid_object_collection_data import ( - BaseRigidObjectCollectionData, - ) - - obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, "cpu") - - assert obj.num_instances == _NUM_INSTANCES - assert obj.num_bodies == _NUM_BODIES - names = obj.body_names - assert isinstance(names, list) - assert len(names) == _NUM_BODIES - assert all(isinstance(n, str) for n in names) - assert isinstance(obj.data, BaseRigidObjectCollectionData) - - -class TestCollectionFinderReturnModes: - """Test finder return modes on production collection backends.""" - - @_production_backends - def test_find_bodies_returns_legacy_tensor_or_cached_proxy(self, backend): - collection, _ = get_rigid_object_collection(backend, num_instances=2, num_bodies=3, device="cpu") - - indices, names = collection.find_bodies(".*") - proxy, proxy_names = collection.find_bodies(".*", as_proxy=True) - - assert isinstance(indices, torch.Tensor) - assert indices.dtype == torch.int32 - assert isinstance(names, list) - assert len(names) == 3 - assert indices.tolist() == proxy.torch.tolist() - assert names == proxy_names - assert proxy is collection.find_bodies(".*", as_proxy=True)[0] - assert proxy.dtype == wp.int32 - assert str(proxy.device) == collection.device - - first_body = collection.body_names[0] - _, first_names = collection.find_bodies(first_body) - assert first_names == [first_body] - - @_production_backends - def test_find_objects_forwards_return_mode_with_alias_warning(self, backend): - collection, _ = get_rigid_object_collection(backend, num_instances=2, num_bodies=3, device="cpu") - - with pytest.warns(DeprecationWarning): - proxy, _ = collection.find_objects(".*", as_proxy=True) - - assert proxy is collection.find_bodies(".*", as_proxy=True)[0] - - -# --------------------------------------------------------------------------- -# Tests: Collection data property contract -# --------------------------------------------------------------------------- - -# (property, shape kind, dtype). Shape kinds: "NB" = (num_instances, num_bodies), -# "NB9" = (num_instances, num_bodies, 9). -_COLLECTION_DATA_PROPERTIES = [ - # body state - ("body_link_pose_w", "NB", wp.transformf), - ("body_link_vel_w", "NB", wp.spatial_vectorf), - ("body_com_pose_w", "NB", wp.transformf), - ("body_com_vel_w", "NB", wp.spatial_vectorf), - ("body_com_acc_w", "NB", wp.spatial_vectorf), - ("body_com_pose_b", "NB", wp.transformf), - # sliced - ("body_link_pos_w", "NB", wp.vec3f), - ("body_link_quat_w", "NB", wp.quatf), - ("body_link_lin_vel_w", "NB", wp.vec3f), - ("body_link_ang_vel_w", "NB", wp.vec3f), - ("body_com_pos_w", "NB", wp.vec3f), - ("body_com_quat_w", "NB", wp.quatf), - ("body_com_lin_vel_w", "NB", wp.vec3f), - ("body_com_ang_vel_w", "NB", wp.vec3f), - ("body_com_lin_acc_w", "NB", wp.vec3f), - ("body_com_ang_acc_w", "NB", wp.vec3f), - ("body_com_pos_b", "NB", wp.vec3f), - ("body_com_quat_b", "NB", wp.quatf), - # derived - ("projected_gravity_b", "NB", wp.vec3f), - ("heading_w", "NB", wp.float32), - ("body_link_lin_vel_b", "NB", wp.vec3f), - ("body_link_ang_vel_b", "NB", wp.vec3f), - ("body_com_lin_vel_b", "NB", wp.vec3f), - ("body_com_ang_vel_b", "NB", wp.vec3f), - # mass properties - ("body_mass", "NB", wp.float32), - ("body_inertia", "NB9", wp.float32), - # defaults - ("default_body_pose", "NB", wp.transformf), - ("default_body_vel", "NB", wp.spatial_vectorf), -] - - -class TestCollectionDataProperties: - """Test that every data property is a ProxyArray with the advertised shape and dtype.""" - - @_backends - @_devices - def test_collection_data_property_contract(self, backend, device): - obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device) - obj.data.update(dt=0.01) - shapes = {"NB": (_NUM_INSTANCES, _NUM_BODIES), "NB9": (_NUM_INSTANCES, _NUM_BODIES, 9)} - for name, shape_kind, dtype in _COLLECTION_DATA_PROPERTIES: - _check_proxy_array( - getattr(obj.data, name), expected_shape=shapes[shape_kind], expected_dtype=dtype, name=name - ) - - -# --------------------------------------------------------------------------- -# Tests: Body pose/velocity writers -# --------------------------------------------------------------------------- - -# writer suffix -> data getter that reads the written quantity back -_BODY_POSE_METHODS = { - "body_pose": "body_link_pose_w", - "body_link_pose": "body_link_pose_w", - "body_com_pose": "body_com_pose_w", -} -_BODY_VEL_METHODS = { - "body_velocity": "body_com_vel_w", - "body_com_velocity": "body_com_vel_w", - "body_link_velocity": "body_link_vel_w", -} - - -class TestCollectionCacheInvalidation: - @_production_backends - def test_pose_write_invalidates_pose_dependent_caches(self, backend): - obj, _ = get_rigid_object_collection(backend, num_instances=2, num_bodies=3, device="cpu") - obj.data.update(dt=0.01) - buffers = _prime_timestamped_properties( - obj.data, - [ - ("body_link_vel_w", "_body_link_vel_w"), - ("projected_gravity_b", "_projected_gravity_b"), - ("heading_w", "_heading_w"), - ("body_link_lin_vel_b", "_body_link_lin_vel_b"), - ("body_link_ang_vel_b", "_body_link_ang_vel_b"), - ("body_com_lin_vel_b", "_body_com_lin_vel_b"), - ("body_com_ang_vel_b", "_body_com_ang_vel_b"), - ], - ) - body_pose = _make_data_warp((obj.num_instances, obj.num_bodies), "cpu", wp.transformf) - obj.write_body_link_pose_to_sim_index(body_poses=body_pose) - _assert_buffers_stale(obj.data, buffers) - - @_production_backends - def test_velocity_write_invalidates_body_frame_caches(self, backend): - obj, _ = get_rigid_object_collection(backend, num_instances=2, num_bodies=3, device="cpu") - obj.data.update(dt=0.01) - buffers = _prime_timestamped_properties( - obj.data, - [ - ("body_link_lin_vel_b", "_body_link_lin_vel_b"), - ("body_link_ang_vel_b", "_body_link_ang_vel_b"), - ("body_com_lin_vel_b", "_body_com_lin_vel_b"), - ("body_com_ang_vel_b", "_body_com_ang_vel_b"), - ], - ) - body_velocity = _make_data_warp((obj.num_instances, obj.num_bodies), "cpu", wp.spatial_vectorf) - obj.write_body_com_velocity_to_sim_index(body_velocities=body_velocity) - _assert_buffers_stale(obj.data, buffers) - - @_production_backends - @pytest.mark.parametrize("setter_kind", ["index", "mask"]) - def test_set_coms_invalidates_same_timestamp_dependents(self, backend, setter_kind): - obj, _ = get_rigid_object_collection(backend, num_instances=2, num_bodies=3, device="cpu") - obj.data.update(dt=0.01) - common_pairs = [ - ("body_com_pose_w", "_body_com_pose_w"), - ("body_link_vel_w", "_body_link_vel_w"), - ("body_link_lin_vel_b", "_body_link_lin_vel_b"), - ("body_link_ang_vel_b", "_body_link_ang_vel_b"), - ("body_com_lin_vel_b", "_body_com_lin_vel_b"), - ("body_com_ang_vel_b", "_body_com_ang_vel_b"), - ("body_state_w", "_body_state_w"), - ("body_link_state_w", "_body_link_state_w"), - ("body_com_state_w", "_body_com_state_w"), - ] - if backend != "newton": - common_pairs.append(("body_com_vel_w", "_body_com_vel_w")) - # Prime public properties before resolving private buffers so Newton allocates lazy caches. - buffers = _prime_timestamped_properties(obj.data, common_pairs) - if backend == "newton": - buffers += _prime_timestamped_properties(obj.data, [("body_com_pose_b", "_body_com_pose_b")]) - coms = _make_com_data(backend, (obj.num_instances, obj.num_bodies), "cpu") - - def set_coms() -> None: - if setter_kind == "index": - obj.set_coms_index(coms=coms) - else: - obj.set_coms_mask(coms=coms) - - if backend == "newton": - from isaaclab_newton.physics import NewtonManager - - with patch.object(NewtonManager, "add_model_change"): - set_coms() - else: - set_coms() - _assert_buffers_stale(obj.data, buffers) - - -class TestCollectionWritersPose: - """Test body pose/velocity writers with all input combinations.""" - - # -- index variants for pose -- - - @_backends - @_devices - @pytest.mark.parametrize("method_suffix", _BODY_POSE_METHODS) - def test_write_body_pose_to_sim_index(self, backend, device, method_suffix): - obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device) - num_instances, num_bodies = _NUM_INSTANCES, _NUM_BODIES - obj.data.update(dt=0.01) - method = getattr(obj, f"write_{method_suffix}_to_sim_index") - - # torch, all envs + all bodies - method(body_poses=_make_data_torch((num_instances, num_bodies), device, wp.transformf)) - # torch, subset envs - method(body_poses=_make_data_torch((1, num_bodies), device, wp.transformf), env_ids=_make_env_ids(device, True)) - # torch, subset bodies - method( - body_poses=_make_data_torch((num_instances, 1), device, wp.transformf), body_ids=_make_body_ids(device, [0]) - ) - # torch, subset both - method( - body_poses=_make_data_torch((1, 1), device, wp.transformf), - env_ids=_make_env_ids(device, True), - body_ids=_make_body_ids(device, [0]), - ) - # warp, all envs + all bodies: the matching getter reads the written poses back - method(body_poses=_make_payload_warp((num_instances, num_bodies), device, wp.transformf)) - getter = _BODY_POSE_METHODS[method_suffix] - expected = _make_payload_torch((num_instances, num_bodies), device, wp.transformf) - _assert_reads_back(getattr(obj.data, getter), expected, getter) - # warp, subset - method( - body_poses=_make_data_warp((1, 1), device, wp.transformf), - env_ids=_make_env_ids(device, True), - body_ids=_make_body_ids(device, [0]), - ) - # negative: bad torch shape - with pytest.raises((AssertionError, RuntimeError)): - method(body_poses=_make_bad_data_torch((num_instances, num_bodies), device, wp.transformf)) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(body_poses=_make_bad_data_warp((num_instances, num_bodies), device, wp.transformf)) - - # -- index variants for velocity -- - - @_backends - @_devices - @pytest.mark.parametrize("method_suffix", _BODY_VEL_METHODS) - def test_write_body_velocity_to_sim_index(self, backend, device, method_suffix): - obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device) - num_instances, num_bodies = _NUM_INSTANCES, _NUM_BODIES - obj.data.update(dt=0.01) - method = getattr(obj, f"write_{method_suffix}_to_sim_index") - - # torch, all envs + all bodies - method(body_velocities=_make_data_torch((num_instances, num_bodies), device, wp.spatial_vectorf)) - # torch, subset envs - method( - body_velocities=_make_data_torch((1, num_bodies), device, wp.spatial_vectorf), - env_ids=_make_env_ids(device, True), - ) - # torch, subset bodies - method( - body_velocities=_make_data_torch((num_instances, 1), device, wp.spatial_vectorf), - body_ids=_make_body_ids(device, [0]), - ) - # torch, subset both - method( - body_velocities=_make_data_torch((1, 1), device, wp.spatial_vectorf), - env_ids=_make_env_ids(device, True), - body_ids=_make_body_ids(device, [0]), - ) - # warp, all envs + all bodies: the matching getter reads the written velocities back - method(body_velocities=_make_payload_warp((num_instances, num_bodies), device, wp.spatial_vectorf)) - getter = _BODY_VEL_METHODS[method_suffix] - expected = _make_payload_torch((num_instances, num_bodies), device, wp.spatial_vectorf) - _assert_reads_back(getattr(obj.data, getter), expected, getter) - # warp, subset - method( - body_velocities=_make_data_warp((1, 1), device, wp.spatial_vectorf), - env_ids=_make_env_ids(device, True), - body_ids=_make_body_ids(device, [0]), - ) - # negative: bad torch shape - with pytest.raises((AssertionError, RuntimeError)): - method(body_velocities=_make_bad_data_torch((num_instances, num_bodies), device, wp.spatial_vectorf)) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(body_velocities=_make_bad_data_warp((num_instances, num_bodies), device, wp.spatial_vectorf)) - - # -- mask variants for pose -- - - @_backends - @_devices - @pytest.mark.parametrize("method_suffix", _BODY_POSE_METHODS) - def test_write_body_pose_to_sim_mask(self, backend, device, method_suffix): - obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device) - num_instances, num_bodies = _NUM_INSTANCES, _NUM_BODIES - obj.data.update(dt=0.01) - method = getattr(obj, f"write_{method_suffix}_to_sim_mask") - - # torch, no mask (all) - method(body_poses=_make_data_torch((num_instances, num_bodies), device, wp.transformf)) - # torch, partial env_mask - method( - body_poses=_make_data_torch((num_instances, num_bodies), device, wp.transformf), - env_mask=_make_env_mask(num_instances, device, True), - ) - # torch, partial body_mask - method( - body_poses=_make_data_torch((num_instances, num_bodies), device, wp.transformf), - body_mask=_make_item_mask(num_bodies, [0], device), - ) - # torch, both masks - method( - body_poses=_make_data_torch((num_instances, num_bodies), device, wp.transformf), - env_mask=_make_env_mask(num_instances, device, True), - body_mask=_make_item_mask(num_bodies, [0], device), - ) - # warp, no mask: the matching getter reads the written poses back - method(body_poses=_make_payload_warp((num_instances, num_bodies), device, wp.transformf)) - getter = _BODY_POSE_METHODS[method_suffix] - expected = _make_payload_torch((num_instances, num_bodies), device, wp.transformf) - _assert_reads_back(getattr(obj.data, getter), expected, getter) - # warp, partial env_mask - method( - body_poses=_make_data_warp((num_instances, num_bodies), device, wp.transformf), - env_mask=_make_env_mask(num_instances, device, True), - ) - # negative: bad torch shape - with pytest.raises((AssertionError, RuntimeError)): - method(body_poses=_make_bad_data_torch((num_instances, num_bodies), device, wp.transformf)) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(body_poses=_make_bad_data_warp((num_instances, num_bodies), device, wp.transformf)) - - # -- mask variants for velocity -- - - @_backends - @_devices - @pytest.mark.parametrize("method_suffix", _BODY_VEL_METHODS) - def test_write_body_velocity_to_sim_mask(self, backend, device, method_suffix): - obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device) - num_instances, num_bodies = _NUM_INSTANCES, _NUM_BODIES - obj.data.update(dt=0.01) - method = getattr(obj, f"write_{method_suffix}_to_sim_mask") - - # torch, no mask - method(body_velocities=_make_data_torch((num_instances, num_bodies), device, wp.spatial_vectorf)) - # torch, partial env_mask - method( - body_velocities=_make_data_torch((num_instances, num_bodies), device, wp.spatial_vectorf), - env_mask=_make_env_mask(num_instances, device, True), - ) - # torch, partial body_mask - method( - body_velocities=_make_data_torch((num_instances, num_bodies), device, wp.spatial_vectorf), - body_mask=_make_item_mask(num_bodies, [0], device), - ) - # warp, no mask: the matching getter reads the written velocities back - method(body_velocities=_make_payload_warp((num_instances, num_bodies), device, wp.spatial_vectorf)) - getter = _BODY_VEL_METHODS[method_suffix] - expected = _make_payload_torch((num_instances, num_bodies), device, wp.spatial_vectorf) - _assert_reads_back(getattr(obj.data, getter), expected, getter) - # warp, partial env_mask - method( - body_velocities=_make_data_warp((num_instances, num_bodies), device, wp.spatial_vectorf), - env_mask=_make_env_mask(num_instances, device, True), - ) - # negative: bad torch shape - with pytest.raises((AssertionError, RuntimeError)): - method(body_velocities=_make_bad_data_torch((num_instances, num_bodies), device, wp.spatial_vectorf)) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(body_velocities=_make_bad_data_warp((num_instances, num_bodies), device, wp.spatial_vectorf)) - - -# --------------------------------------------------------------------------- -# Tests: Body property setters -# --------------------------------------------------------------------------- - -_BODY_METHODS = [("set_masses", "masses"), ("set_coms", "coms"), ("set_inertias", "inertias")] - - -def _body_writer_layout(backend: str, method_base: str) -> tuple[type, int, str]: - """Return the warp dtype, torch trailing size, and read-back getter of a body property writer.""" - if method_base == "set_masses": - return wp.float32, 0, "body_mass" - if method_base == "set_inertias": - return wp.float32, 9, "body_inertia" - if backend == "newton": - # Newton stores the COM as a position only. - return wp.vec3f, 3, "body_com_pos_b" - return wp.transformf, 7, "body_com_pose_b" - - -def _make_body_torch(shape: tuple[int, int], device: str, wp_dtype: type, trailing: int, payload: bool = False): - """Create body-property torch data: valid ones by default, or per-element distinct values.""" - full_shape = (*shape, trailing) if trailing else shape - if payload: - data = torch.arange(1, math.prod(full_shape) + 1, dtype=torch.float32, device=device).reshape(full_shape) - else: - data = torch.ones(full_shape, device=device, dtype=torch.float32) - if wp_dtype == wp.transformf: - data[..., 3:6] = 0.0 - data[..., 6] = 1.0 - if not payload: - data[..., :3] = 0.0 - return data - - -def _make_body_warp(shape: tuple[int, int], device: str, wp_dtype: type, trailing: int, payload: bool = False): - """Create body-property data as a warp array (structured dtypes collapse the trailing dim).""" - t = _make_body_torch(shape, device, wp_dtype, trailing, payload).contiguous() - return wp.from_torch(t, dtype=wp.float32 if wp_dtype == wp.float32 else wp_dtype) - - -class TestCollectionWritersBody: - """Test body property writers/setters with all input combinations.""" - - @_backends - @_devices - @pytest.mark.parametrize("method_base, kwarg", _BODY_METHODS, ids=[m[0] for m in _BODY_METHODS]) - def test_body_writer_index(self, backend, device, method_base, kwarg): - obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device) - num_instances, num_bodies = _NUM_INSTANCES, _NUM_BODIES - wp_dtype, trailing, getter = _body_writer_layout(backend, method_base) - obj.data.update(dt=0.01) - method = getattr(obj, f"{method_base}_index") - sub_body_ids = [0] - - # torch, all envs + all bodies - method(**{kwarg: _make_body_torch((num_instances, num_bodies), device, wp_dtype, trailing)}) - # torch, subset - method( - **{ - kwarg: _make_body_torch((1, 1), device, wp_dtype, trailing), - "body_ids": sub_body_ids, - "env_ids": _make_env_ids(device, True), - } - ) - # warp, all envs + all bodies: the matching getter reads the written values back - method(**{kwarg: _make_body_warp((num_instances, num_bodies), device, wp_dtype, trailing, payload=True)}) - expected = _make_body_torch((num_instances, num_bodies), device, wp_dtype, trailing, payload=True) - _assert_reads_back(getattr(obj.data, getter), expected, getter) - # warp, subset - method( - **{ - kwarg: _make_body_warp((1, 1), device, wp_dtype, trailing), - "body_ids": sub_body_ids, - "env_ids": _make_env_ids(device, True), - } - ) - # negative: bad torch shape (extra env) - with pytest.raises((AssertionError, RuntimeError)): - method(**{kwarg: _make_body_torch((num_instances + 1, num_bodies), device, wp_dtype, trailing)}) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(**{kwarg: _make_body_warp((num_instances + 1, num_bodies), device, wp_dtype, trailing)}) - - @_backends - @_devices - @pytest.mark.parametrize("method_base, kwarg", _BODY_METHODS, ids=[m[0] for m in _BODY_METHODS]) - def test_body_writer_mask(self, backend, device, method_base, kwarg): - obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device) - num_instances, num_bodies = _NUM_INSTANCES, _NUM_BODIES - wp_dtype, trailing, getter = _body_writer_layout(backend, method_base) - obj.data.update(dt=0.01) - method = getattr(obj, f"{method_base}_mask") - sub_body_sel = [0] - - # torch, no mask - method(**{kwarg: _make_body_torch((num_instances, num_bodies), device, wp_dtype, trailing)}) - # torch, partial env_mask + body_mask - method( - **{ - kwarg: _make_body_torch((num_instances, num_bodies), device, wp_dtype, trailing), - "body_mask": _make_item_mask(num_bodies, sub_body_sel, device), - "env_mask": _make_env_mask(num_instances, device, True), - } - ) - # warp, no mask: the matching getter reads the written values back - method(**{kwarg: _make_body_warp((num_instances, num_bodies), device, wp_dtype, trailing, payload=True)}) - expected = _make_body_torch((num_instances, num_bodies), device, wp_dtype, trailing, payload=True) - _assert_reads_back(getattr(obj.data, getter), expected, getter) - # warp, partial env_mask + body_mask - method( - **{ - kwarg: _make_body_warp((num_instances, num_bodies), device, wp_dtype, trailing), - "body_mask": _make_item_mask(num_bodies, sub_body_sel, device), - "env_mask": _make_env_mask(num_instances, device, True), - } - ) - # negative: bad torch shape - with pytest.raises((AssertionError, RuntimeError)): - method(**{kwarg: _make_body_torch((num_instances + 1, num_bodies), device, wp_dtype, trailing)}) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(**{kwarg: _make_body_warp((num_instances + 1, num_bodies), device, wp_dtype, trailing)}) - - -# --------------------------------------------------------------------------- -# Tests: Alias/shorthand properties -# --------------------------------------------------------------------------- - -_COLLECTION_ALIASES = [ - ("body_pose_w", "body_link_pose_w"), - ("body_pos_w", "body_link_pos_w"), - ("body_quat_w", "body_link_quat_w"), - ("body_vel_w", "body_com_vel_w"), - ("body_lin_vel_w", "body_com_lin_vel_w"), - ("body_ang_vel_w", "body_com_ang_vel_w"), - ("body_acc_w", "body_com_acc_w"), - ("body_lin_acc_w", "body_com_lin_acc_w"), - ("body_ang_acc_w", "body_com_ang_acc_w"), - ("com_pos_b", "body_com_pos_b"), - ("com_quat_b", "body_com_quat_b"), -] - - -class TestCollectionDataAliases: - """Test that alias properties return the values of their canonical counterparts.""" - - @_backends - def test_aliases_match_canonical_values(self, backend): - # Random mock state makes link and COM quantities differ, so a retargeted alias fails. - obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, "cpu") - obj.data.update(dt=0.01) - d = obj.data - - for alias, canonical in _COLLECTION_ALIASES: - alias_value, canonical_value = getattr(d, alias), getattr(d, canonical) - assert alias_value.shape == canonical_value.shape, alias - assert alias_value.dtype == canonical_value.dtype, alias - assert torch.equal(alias_value.torch, canonical_value.torch), alias diff --git a/source/isaaclab/test/assets/test_rigid_object_iface.py b/source/isaaclab/test/assets/test_rigid_object_iface.py deleted file mode 100644 index 47d02fefe868..000000000000 --- a/source/isaaclab/test/assets/test_rigid_object_iface.py +++ /dev/null @@ -1,776 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -# ignore private usage of variables warning -# pyright: reportPrivateUsage=none - -""" -Checks that the rigid object interfaces are consistent across backends, and are providing the exact same data as what -the base rigid object class advertises. All rigid object interfaces need to comply with the same interface contract. - -The setup is a bit convoluted so that we can run these tests without requiring Isaac Sim or GPU simulation. -""" - -import math -from unittest.mock import MagicMock, patch - -import numpy as np -import pytest -import torch -import warp as wp -from _rigid_object_iface_test_utils import BACKENDS, get_rigid_object - -from isaaclab.test.utils import DeviceScope, test_devices - -pytestmark = pytest.mark.integration - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _check_proxy_array(arr, *, expected_shape: tuple, expected_dtype: type, name: str): - """Assert that `arr` is a ProxyArray with the expected shape and dtype.""" - from isaaclab.utils.warp import ProxyArray - - assert isinstance(arr, ProxyArray), f"{name}: expected ProxyArray, got {type(arr)}" - assert arr.shape == expected_shape, f"{name}: expected shape {expected_shape}, got {arr.shape}" - assert arr.dtype == expected_dtype, f"{name}: expected dtype {expected_dtype}, got {arr.dtype}" - - -# Common parametrize decorators. Pure bookkeeping (counts, names, finders, aliases) runs on CPU only; -# getters and writers keep every test device because PhysX stages through CPU-pinned buffers on CUDA. -_backends = pytest.mark.parametrize("backend", BACKENDS, indirect=False) -_devices = pytest.mark.parametrize("device", test_devices(DeviceScope.CPU_AND_DEFAULT_CUDA)) -_NUM_INSTANCES = 2 -_production_backends = pytest.mark.parametrize( - "backend", [backend for backend in ("physx", "newton", "ovphysx") if backend in BACKENDS], indirect=False -) - - -# --------------------------------------------------------------------------- -# Tests: Index resolution helpers -# --------------------------------------------------------------------------- - - -class TestRigidObjectIndexResolution: - """Test backend-specific index resolution helpers.""" - - @_production_backends - @_devices - def test_resolve_env_ids_handles_tensor_view_shape(self, backend, device): - obj, _ = get_rigid_object(backend, num_instances=4, device=device) - - env_ids = torch.arange(4, dtype=torch.int32, device=device) - resolved_full = obj._resolve_env_ids(env_ids) - resolved_view = obj._resolve_env_ids(env_ids[:2]) - - assert resolved_full.shape[0] == 4 - assert resolved_view.shape[0] == 2 - cached = wp.to_torch(obj._ALL_INDICES) - for selection in (slice(None), slice(1, None, 2), slice(0, 0)): - resolved = wp.to_torch(obj._resolve_env_ids(selection)) - torch.testing.assert_close(resolved, cached[selection]) - assert resolved.data_ptr() == cached[selection].data_ptr() - assert resolved.stride() == cached[selection].stride() - - -# --------------------------------------------------------------------------- -# Tests: RigidObject properties and finders -# --------------------------------------------------------------------------- - - -class TestRigidObjectProperties: - """Test that rigid object properties return the correct types/values.""" - - @_backends - def test_rigid_object_counts_and_names(self, backend): - from isaaclab.assets.rigid_object.base_rigid_object_data import BaseRigidObjectData - - obj, _ = get_rigid_object(backend, _NUM_INSTANCES, device="cpu") - - assert obj.num_instances == _NUM_INSTANCES - assert obj.num_bodies == 1 - names = obj.body_names - assert isinstance(names, list) - assert len(names) == 1 - assert all(isinstance(n, str) for n in names) - assert isinstance(obj.data, BaseRigidObjectData) - - -class TestRigidObjectFinderReturnModes: - """Test finder return modes on production rigid-object backends.""" - - @_production_backends - def test_find_bodies_returns_legacy_list_or_cached_proxy(self, backend): - obj, _ = get_rigid_object(backend, num_instances=2, device="cpu") - - indices, names = obj.find_bodies(".*") - proxy, proxy_names = obj.find_bodies(".*", as_proxy=True) - - assert isinstance(indices, list) and isinstance(names, list) - assert len(indices) == 1 - assert len(names) == 1 - assert all(isinstance(i, int) for i in indices) - assert all(isinstance(n, str) for n in names) - assert indices == proxy.torch.tolist() - assert names == proxy_names - assert proxy is obj.find_bodies(".*", as_proxy=True)[0] - assert proxy.dtype == wp.int32 - assert str(proxy.device) == obj.device - - first_body = obj.body_names[0] - assert obj.find_bodies(first_body) == ([0], [first_body]) - - -# --------------------------------------------------------------------------- -# Tests: RigidObjectData property contract -# --------------------------------------------------------------------------- - -# (property, shape kind, dtype). Shape kinds: "N" = (num_instances,), "N1" = (num_instances, 1), -# "N19" = (num_instances, 1, 9). -_RIGID_OBJECT_DATA_PROPERTIES = [ - # root state - ("root_link_pose_w", "N", wp.transformf), - ("root_link_vel_w", "N", wp.spatial_vectorf), - ("root_com_pose_w", "N", wp.transformf), - ("root_com_vel_w", "N", wp.spatial_vectorf), - ("root_link_pos_w", "N", wp.vec3f), - ("root_link_quat_w", "N", wp.quatf), - ("root_link_lin_vel_w", "N", wp.vec3f), - ("root_link_ang_vel_w", "N", wp.vec3f), - ("root_com_pos_w", "N", wp.vec3f), - ("root_com_quat_w", "N", wp.quatf), - ("root_com_lin_vel_w", "N", wp.vec3f), - ("root_com_ang_vel_w", "N", wp.vec3f), - # derived - ("projected_gravity_b", "N", wp.vec3f), - ("heading_w", "N", wp.float32), - ("root_link_lin_vel_b", "N", wp.vec3f), - ("root_link_ang_vel_b", "N", wp.vec3f), - ("root_com_lin_vel_b", "N", wp.vec3f), - ("root_com_ang_vel_b", "N", wp.vec3f), - # body state - ("body_link_pose_w", "N1", wp.transformf), - ("body_link_vel_w", "N1", wp.spatial_vectorf), - ("body_com_pose_w", "N1", wp.transformf), - ("body_com_vel_w", "N1", wp.spatial_vectorf), - ("body_com_acc_w", "N1", wp.spatial_vectorf), - ("body_com_pose_b", "N1", wp.transformf), - ("body_mass", "N1", wp.float32), - ("body_inertia", "N19", wp.float32), - ("body_link_pos_w", "N1", wp.vec3f), - ("body_link_quat_w", "N1", wp.quatf), - ("body_link_lin_vel_w", "N1", wp.vec3f), - ("body_link_ang_vel_w", "N1", wp.vec3f), - ("body_com_pos_w", "N1", wp.vec3f), - ("body_com_quat_w", "N1", wp.quatf), - ("body_com_pos_b", "N1", wp.vec3f), - ("body_com_quat_b", "N1", wp.quatf), - # defaults - ("default_root_pose", "N", wp.transformf), - ("default_root_vel", "N", wp.spatial_vectorf), -] - - -class TestRigidObjectDataProperties: - """Test that every data property is a ProxyArray with the advertised shape and dtype.""" - - @_backends - @_devices - def test_rigid_object_data_property_contract(self, backend, device): - obj, _ = get_rigid_object(backend, _NUM_INSTANCES, device) - obj.data.update(dt=0.01) - if backend == "newton": - for name in ("root_link_vel_w", "root_com_pose_w", "body_com_pose_b", "projected_gravity_b", "heading_w"): - assert getattr(obj.data, "_" + name).data is None - shapes = {"N": (_NUM_INSTANCES,), "N1": (_NUM_INSTANCES, 1), "N19": (_NUM_INSTANCES, 1, 9)} - for name, shape_kind, dtype in _RIGID_OBJECT_DATA_PROPERTIES: - _check_proxy_array( - getattr(obj.data, name), expected_shape=shapes[shape_kind], expected_dtype=dtype, name=name - ) - - for frame in ("root_link", "root_com", "body_link", "body_com"): - for quantity, components in (("pose", ("pos", "quat")), ("vel", ("lin_vel", "ang_vel"))): - packed = getattr(obj.data, f"{frame}_{quantity}_w").torch - for component, expected in zip(components, (packed[..., :3], packed[..., 3:]), strict=True): - view = getattr(obj.data, f"{frame}_{component}_w").torch - torch.testing.assert_close(view, expected) - assert view.data_ptr() == expected.data_ptr() - assert view.stride() == expected.stride() - - -# --------------------------------------------------------------------------- -# Tests: Alias/shorthand properties -# --------------------------------------------------------------------------- - -_RIGID_OBJECT_ALIASES = [ - ("root_pose_w", "root_link_pose_w"), - ("root_pos_w", "root_link_pos_w"), - ("root_quat_w", "root_link_quat_w"), - ("root_vel_w", "root_com_vel_w"), - ("root_lin_vel_w", "root_com_lin_vel_w"), - ("root_ang_vel_w", "root_com_ang_vel_w"), - ("body_pose_w", "body_link_pose_w"), - ("body_pos_w", "body_link_pos_w"), - ("body_quat_w", "body_link_quat_w"), - ("body_vel_w", "body_com_vel_w"), - ("body_lin_vel_w", "body_com_lin_vel_w"), - ("body_ang_vel_w", "body_com_ang_vel_w"), -] - - -class TestRigidObjectDataAliases: - """Test that alias properties return the values of their canonical counterparts.""" - - @_backends - def test_aliases_match_canonical_values(self, backend): - # Random mock state makes link and COM quantities differ, so a retargeted alias fails. - obj, _ = get_rigid_object(backend, _NUM_INSTANCES, device="cpu") - obj.data.update(dt=0.01) - d = obj.data - - for alias, canonical in _RIGID_OBJECT_ALIASES: - alias_value, canonical_value = getattr(d, alias), getattr(d, canonical) - assert alias_value.shape == canonical_value.shape, alias - assert alias_value.dtype == canonical_value.dtype, alias - assert torch.equal(alias_value.torch, canonical_value.torch), alias - - -# --------------------------------------------------------------------------- -# Writer/setter test helpers -# --------------------------------------------------------------------------- - -# Map warp structured dtypes to their torch trailing dimension size. -_WP_DTYPE_TO_TRAILING = { - wp.transformf: 7, - wp.spatial_vectorf: 6, - wp.vec2f: 2, - wp.float32: 0, # no trailing dimension -} - - -def _make_data_torch(shape: tuple, device: str, wp_dtype=wp.float32) -> torch.Tensor: - """Create valid torch test data for a given warp dtype.""" - trailing = _WP_DTYPE_TO_TRAILING[wp_dtype] - if trailing: - full_shape = (*shape, trailing) - else: - full_shape = shape - data = torch.zeros(full_shape, device=device, dtype=torch.float32) - if wp_dtype == wp.transformf: - data[..., 6] = 1.0 # identity quat w - elif wp_dtype == wp.vec2f: - data[..., 0] = -1.0 - data[..., 1] = 1.0 - elif wp_dtype == wp.float32: - data.fill_(1.0) - return data - - -def _make_data_warp(shape: tuple, device: str, wp_dtype=wp.float32) -> wp.array: - """Create valid warp test data for a given warp dtype.""" - t = _make_data_torch(shape, device, wp_dtype) - if wp_dtype == wp.float32: - return wp.from_torch(t, dtype=wp.float32) - return wp.from_torch(t.contiguous(), dtype=wp_dtype) - - -def _make_payload_torch(shape: tuple, device: str, wp_dtype=wp.float32) -> torch.Tensor: - """Create valid torch data whose entries differ per element, for writer read-back checks. - - Transforms get distinct positions and a fixed 90-degree rotation about Z. - """ - values = torch.arange(1, math.prod(shape) + 1, dtype=torch.float32, device=device).reshape(shape) + 0.25 - if wp_dtype == wp.spatial_vectorf: - return values.unsqueeze(-1) + torch.arange(6, dtype=torch.float32, device=device) / 10.0 - if wp_dtype == wp.transformf: - data = torch.zeros((*shape, 7), dtype=torch.float32, device=device) - data[..., :3] = values.unsqueeze(-1) + torch.arange(3, dtype=torch.float32, device=device) / 10.0 - data[..., 5] = data[..., 6] = 2.0**-0.5 - return data - raise ValueError(f"Unsupported payload dtype: {wp_dtype}") - - -def _make_payload_warp(shape: tuple, device: str, wp_dtype=wp.float32) -> wp.array: - """Create the per-element distinct payload of :func:`_make_payload_torch` as a warp array.""" - return wp.from_torch(_make_payload_torch(shape, device, wp_dtype).contiguous(), dtype=wp_dtype) - - -def _assert_reads_back(value, expected: torch.Tensor, name: str) -> None: - """Assert a data getter (ProxyArray) returns the values a writer just wrote.""" - torch.testing.assert_close(value.torch, expected, atol=1e-5, rtol=1e-5, msg=lambda msg: f"{name}: {msg}") - - -def _make_com_data(backend: str, shape: tuple[int, ...], device: str) -> wp.array: - """Create backend-compatible center-of-mass test data.""" - if backend == "newton": - return wp.zeros(shape, dtype=wp.vec3f, device=device) - return _make_data_warp(shape, device, wp.transformf) - - -def _prime_timestamped_properties(data, property_buffer_pairs: list[tuple[str, str]]): - """Prime public lazy properties and return their concrete timestamped buffers.""" - buffers = [] - for property_name, buffer_name in property_buffer_pairs: - getattr(data, property_name) - buffer = getattr(data, buffer_name) - assert buffer is not None, buffer_name - buffer.timestamp = data._sim_timestamp - buffers.append((buffer_name, buffer)) - return buffers - - -def _assert_buffers_stale(data, buffers) -> None: - for name, buffer in buffers: - assert buffer.timestamp < data._sim_timestamp, name - - -def _make_bad_data_torch(shape: tuple, device: str, wp_dtype=wp.float32) -> torch.Tensor: - """Create torch data with wrong leading shape for negative testing.""" - bad_shape = (shape[0] + 1,) + shape[1:] - return _make_data_torch(bad_shape, device, wp_dtype) - - -def _make_bad_data_warp(shape: tuple, device: str, wp_dtype=wp.float32) -> wp.array: - """Create warp data with wrong leading shape for negative testing.""" - bad_shape = (shape[0] + 1,) + shape[1:] - return _make_data_warp(bad_shape, device, wp_dtype) - - -def _make_env_mask(num_instances: int, device: str, partial: bool) -> wp.array | None: - """Create an env_mask: None for all envs, or a partial bool mask.""" - if not partial: - return None - mask_np = np.zeros(num_instances, dtype=bool) - mask_np[0] = True - return wp.array(mask_np, dtype=wp.bool, device=device) - - -def _make_env_ids(device: str, subset: bool) -> torch.Tensor | None: - """Create env_ids: None for all envs, or [0] for a subset.""" - if not subset: - return None - return torch.tensor([0], dtype=torch.int32, device=device) - - -def _make_item_mask(total: int, selected: list[int], device: str) -> wp.array: - """Create a bool warp mask with True at `selected` indices, False elsewhere.""" - mask_np = np.zeros(total, dtype=bool) - for i in selected: - mask_np[i] = True - return wp.array(mask_np, dtype=wp.bool, device=device) - - -# --------------------------------------------------------------------------- -# Tests: Root writers — torch/warp × index/mask × all/subset × negative -# --------------------------------------------------------------------------- - -# writer suffix -> data getter that reads the written quantity back -_ROOT_POSE_METHODS = { - "root_pose": "root_link_pose_w", - "root_link_pose": "root_link_pose_w", - "root_com_pose": "root_com_pose_w", -} -_ROOT_VEL_METHODS = { - "root_velocity": "root_com_vel_w", - "root_link_velocity": "root_link_vel_w", - "root_com_velocity": "root_com_vel_w", -} - - -class TestRigidObjectCacheInvalidation: - @_production_backends - def test_pose_write_invalidates_pose_dependent_caches(self, backend): - obj, _ = get_rigid_object(backend, num_instances=2, device="cpu") - obj.data.update(dt=0.01) - buffers = _prime_timestamped_properties( - obj.data, - [ - ("root_link_vel_w", "_root_link_vel_w"), - ("projected_gravity_b", "_projected_gravity_b"), - ("heading_w", "_heading_w"), - ("root_link_lin_vel_b", "_root_link_lin_vel_b"), - ("root_link_ang_vel_b", "_root_link_ang_vel_b"), - ("root_com_lin_vel_b", "_root_com_lin_vel_b"), - ("root_com_ang_vel_b", "_root_com_ang_vel_b"), - ], - ) - root_pose = _make_data_warp((obj.num_instances,), "cpu", wp.transformf) - obj.write_root_link_pose_to_sim_index(root_pose=root_pose) - _assert_buffers_stale(obj.data, buffers) - - @_production_backends - def test_velocity_write_invalidates_body_frame_caches(self, backend): - obj, _ = get_rigid_object(backend, num_instances=2, device="cpu") - obj.data.update(dt=0.01) - body_velocity = obj.data.body_link_vel_w - buffers = _prime_timestamped_properties( - obj.data, - [ - ("root_link_lin_vel_b", "_root_link_lin_vel_b"), - ("root_link_ang_vel_b", "_root_link_ang_vel_b"), - ("root_com_lin_vel_b", "_root_com_lin_vel_b"), - ("root_com_ang_vel_b", "_root_com_ang_vel_b"), - ], - ) - root_velocity = _make_data_warp((obj.num_instances,), "cpu", wp.spatial_vectorf) - obj.write_root_com_velocity_to_sim_index(root_velocity=root_velocity) - _assert_buffers_stale(obj.data, buffers) - # Read the body alias first; zero angular velocity makes link and COM velocities equal. - _assert_reads_back(obj.data.body_link_vel_w, wp.to_torch(root_velocity).unsqueeze(1), "body_link_vel_w") - assert obj.data.body_link_vel_w is body_velocity - - @_production_backends - @pytest.mark.parametrize("setter_kind", ["index", "mask"]) - def test_set_coms_invalidates_same_timestamp_dependents(self, backend, setter_kind): - obj, _ = get_rigid_object(backend, num_instances=2, device="cpu") - obj.data.update(dt=0.01) - common_pairs = [ - ("root_com_pose_w", "_root_com_pose_w"), - ("root_link_vel_w", "_root_link_vel_w"), - ("root_link_lin_vel_b", "_root_link_lin_vel_b"), - ("root_link_ang_vel_b", "_root_link_ang_vel_b"), - ("root_com_lin_vel_b", "_root_com_lin_vel_b"), - ("root_com_ang_vel_b", "_root_com_ang_vel_b"), - ("root_state_w", "_root_state_w"), - ("root_link_state_w", "_root_link_state_w"), - ("root_com_state_w", "_root_com_state_w"), - ] - if backend != "newton": - common_pairs.append(("root_com_vel_w", "_root_com_vel_w")) - # Prime public properties before resolving private buffers so Newton allocates lazy caches. - buffers = _prime_timestamped_properties(obj.data, common_pairs) - if backend == "newton": - buffers += _prime_timestamped_properties(obj.data, [("body_com_pose_b", "_body_com_pose_b")]) - coms = _make_com_data(backend, (obj.num_instances, obj.num_bodies), "cpu") - - def set_coms() -> None: - if setter_kind == "index": - obj.set_coms_index(coms=coms) - else: - obj.set_coms_mask(coms=coms) - - if backend == "newton": - from isaaclab_newton.physics import NewtonManager - - with patch.object(NewtonManager, "add_model_change"): - set_coms() - else: - set_coms() - _assert_buffers_stale(obj.data, buffers) - - -class TestRigidObjectWritersRoot: - """Test root pose/velocity writers with all input combinations.""" - - # -- index variants -- - - @_backends - @_devices - @pytest.mark.parametrize("method_suffix", _ROOT_POSE_METHODS) - def test_write_root_pose_to_sim_index(self, backend, device, method_suffix): - num_instances = _NUM_INSTANCES - obj, _ = get_rigid_object(backend, num_instances, device) - obj.data.update(dt=0.01) - method = getattr(obj, f"write_{method_suffix}_to_sim_index") - - # torch, all envs - method(root_pose=_make_data_torch((num_instances,), device, wp.transformf)) - # torch, subset - method(root_pose=_make_data_torch((1,), device, wp.transformf), env_ids=_make_env_ids(device, True)) - # warp, all envs: the matching getter reads the written poses back - method(root_pose=_make_payload_warp((num_instances,), device, wp.transformf)) - getter = _ROOT_POSE_METHODS[method_suffix] - _assert_reads_back( - getattr(obj.data, getter), _make_payload_torch((num_instances,), device, wp.transformf), getter - ) - # warp, subset - method(root_pose=_make_data_warp((1,), device, wp.transformf), env_ids=_make_env_ids(device, True)) - # negative: bad torch shape - with pytest.raises((AssertionError, RuntimeError)): - method(root_pose=_make_bad_data_torch((num_instances,), device, wp.transformf)) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(root_pose=_make_bad_data_warp((num_instances,), device, wp.transformf)) - - @_backends - @_devices - @pytest.mark.parametrize("method_suffix", _ROOT_VEL_METHODS) - def test_write_root_velocity_to_sim_index(self, backend, device, method_suffix): - num_instances = _NUM_INSTANCES - obj, _ = get_rigid_object(backend, num_instances, device) - obj.data.update(dt=0.01) - method = getattr(obj, f"write_{method_suffix}_to_sim_index") - - # torch, all envs - method(root_velocity=_make_data_torch((num_instances,), device, wp.spatial_vectorf)) - # torch, subset - method(root_velocity=_make_data_torch((1,), device, wp.spatial_vectorf), env_ids=_make_env_ids(device, True)) - # warp, all envs: the matching getter reads the written velocities back - method(root_velocity=_make_payload_warp((num_instances,), device, wp.spatial_vectorf)) - getter = _ROOT_VEL_METHODS[method_suffix] - _assert_reads_back( - getattr(obj.data, getter), _make_payload_torch((num_instances,), device, wp.spatial_vectorf), getter - ) - # warp, subset - method(root_velocity=_make_data_warp((1,), device, wp.spatial_vectorf), env_ids=_make_env_ids(device, True)) - # negative: bad torch shape - with pytest.raises((AssertionError, RuntimeError)): - method(root_velocity=_make_bad_data_torch((num_instances,), device, wp.spatial_vectorf)) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(root_velocity=_make_bad_data_warp((num_instances,), device, wp.spatial_vectorf)) - - # -- mask variants -- - - @_backends - @_devices - @pytest.mark.parametrize("method_suffix", _ROOT_POSE_METHODS) - def test_write_root_pose_to_sim_mask(self, backend, device, method_suffix): - num_instances = _NUM_INSTANCES - obj, _ = get_rigid_object(backend, num_instances, device) - obj.data.update(dt=0.01) - method = getattr(obj, f"write_{method_suffix}_to_sim_mask") - - # torch, no mask (all) - method(root_pose=_make_data_torch((num_instances,), device, wp.transformf)) - # torch, partial mask - method( - root_pose=_make_data_torch((num_instances,), device, wp.transformf), - env_mask=_make_env_mask(num_instances, device, True), - ) - # warp, no mask: the matching getter reads the written poses back - method(root_pose=_make_payload_warp((num_instances,), device, wp.transformf)) - getter = _ROOT_POSE_METHODS[method_suffix] - _assert_reads_back( - getattr(obj.data, getter), _make_payload_torch((num_instances,), device, wp.transformf), getter - ) - # warp, partial mask - method( - root_pose=_make_data_warp((num_instances,), device, wp.transformf), - env_mask=_make_env_mask(num_instances, device, True), - ) - # negative: bad torch shape - with pytest.raises((AssertionError, RuntimeError)): - method(root_pose=_make_bad_data_torch((num_instances,), device, wp.transformf)) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(root_pose=_make_bad_data_warp((num_instances,), device, wp.transformf)) - - @_backends - @_devices - @pytest.mark.parametrize("method_suffix", _ROOT_VEL_METHODS) - def test_write_root_velocity_to_sim_mask(self, backend, device, method_suffix): - num_instances = _NUM_INSTANCES - obj, _ = get_rigid_object(backend, num_instances, device) - obj.data.update(dt=0.01) - method = getattr(obj, f"write_{method_suffix}_to_sim_mask") - - # torch, no mask - method(root_velocity=_make_data_torch((num_instances,), device, wp.spatial_vectorf)) - # torch, partial mask - method( - root_velocity=_make_data_torch((num_instances,), device, wp.spatial_vectorf), - env_mask=_make_env_mask(num_instances, device, True), - ) - # warp, no mask: the matching getter reads the written velocities back - method(root_velocity=_make_payload_warp((num_instances,), device, wp.spatial_vectorf)) - getter = _ROOT_VEL_METHODS[method_suffix] - _assert_reads_back( - getattr(obj.data, getter), _make_payload_torch((num_instances,), device, wp.spatial_vectorf), getter - ) - # warp, partial mask - method( - root_velocity=_make_data_warp((num_instances,), device, wp.spatial_vectorf), - env_mask=_make_env_mask(num_instances, device, True), - ) - # negative: bad torch shape - with pytest.raises((AssertionError, RuntimeError)): - method(root_velocity=_make_bad_data_torch((num_instances,), device, wp.spatial_vectorf)) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(root_velocity=_make_bad_data_warp((num_instances,), device, wp.spatial_vectorf)) - - -# --------------------------------------------------------------------------- -# Tests: Body writers — torch/warp × index/mask × all/subset × negative -# --------------------------------------------------------------------------- - -_BODY_METHODS = [("set_masses", "masses"), ("set_coms", "coms"), ("set_inertias", "inertias")] - - -def _body_writer_layout(backend: str, method_base: str) -> tuple[type, int, str | None]: - """Return the warp dtype, torch trailing size, and read-back getter of a body property writer.""" - if method_base == "set_masses": - return wp.float32, 0, "body_mass" - if method_base == "set_inertias": - return wp.float32, 9, "body_inertia" - if backend == "newton": - # Newton stores the COM as a position only. - return wp.vec3f, 3, "body_com_pos_b" - if backend == "physx": - # PhysX re-reads the COM from its view after a write, and the mocked view drops writes. - return wp.transformf, 7, None - return wp.transformf, 7, "body_com_pose_b" - - -def _make_body_torch(shape: tuple[int, int], device: str, wp_dtype: type, trailing: int, payload: bool = False): - """Create body-property torch data: valid ones by default, or per-element distinct values.""" - full_shape = (*shape, trailing) if trailing else shape - if payload: - data = torch.arange(1, math.prod(full_shape) + 1, dtype=torch.float32, device=device).reshape(full_shape) - else: - data = torch.ones(full_shape, device=device, dtype=torch.float32) - if wp_dtype == wp.transformf: - data[..., 3:6] = 0.0 - data[..., 6] = 1.0 - if not payload: - data[..., :3] = 0.0 - return data - - -def _make_body_warp(shape: tuple[int, int], device: str, wp_dtype: type, trailing: int, payload: bool = False): - """Create body-property data as a warp array (structured dtypes collapse the trailing dim).""" - t = _make_body_torch(shape, device, wp_dtype, trailing, payload).contiguous() - return wp.from_torch(t, dtype=wp.float32 if wp_dtype == wp.float32 else wp_dtype) - - -class TestRigidObjectWritersBody: - """Test body property writers/setters with all input combinations.""" - - @_production_backends - def test_external_wrench_frames(self, backend): - """Forward local and world wrenches through the real writer in each backend's frame.""" - device = "cpu" # the composer and wrench-packing kernels are device-independent - obj, raw_backend = get_rigid_object(backend, num_instances=2, device=device) - # Seed a known 90-degree rotation about Z so the local-frame expectation is independent of production. - root_pose = torch.tensor( - [[1.0, 2.0, 3.0, 0.0, 0.0, 2.0**-0.5, 2.0**-0.5], [4.0, 5.0, 6.0, 0.0, 0.0, 2.0**-0.5, 2.0**-0.5]], - device=device, - ) - obj.write_root_link_pose_to_sim_index(root_pose=root_pose) - if backend == "newton": - # Seed the body pose the native FK would publish. - obj.data._sim_bind_body_link_pose_w.assign(wp.from_torch(root_pose, dtype=wp.transformf)) - composer = obj.permanent_wrench_composer - forces = torch.arange(1.0, 7.0, device=device).reshape(2, 1, 3) - torques = forces + 10.0 - if backend == "physx": - raw_backend.apply_forces_and_torques_at_position = MagicMock() - - for is_global in (False, True): - composer.reset() - composer.set_forces_and_torques_index(forces=forces, torques=torques, is_global=is_global) - with patch.object(composer, "compose_to_body_frame", wraps=composer.compose_to_body_frame) as compose: - obj.write_data_to_sim() - assert compose.call_count == int(is_global and backend == "newton") - - expected_force, expected_torque = forces, torques - if backend == "physx": - call = raw_backend.apply_forces_and_torques_at_position.call_args.kwargs - assert call["is_global"] is is_global - assert call["position_data"] is None - actual_force = call["force_data"].numpy().reshape(2, 1, 3) - actual_torque = call["torque_data"].numpy().reshape(2, 1, 3) - else: - if not is_global: - # A 90-degree Z rotation maps body (x, y, z) to world (-y, x, z). - expected_force = torch.stack((-forces[..., 1], forces[..., 0], forces[..., 2]), dim=-1) - expected_torque = torch.stack((-torques[..., 1], torques[..., 0], torques[..., 2]), dim=-1) - if backend == "newton": - packed = obj.data._sim_bind_body_external_wrench.numpy() - else: - from isaaclab_ov import tensor_types as TT - - packed = raw_backend.bindings[TT.RIGID_BODY_WRENCH]._data.reshape(2, 1, 9) - np.testing.assert_allclose(packed[..., 6:9], root_pose[:, None, :3].numpy()) - actual_force, actual_torque = packed[..., :3], packed[..., 3:6] - np.testing.assert_allclose(actual_force, expected_force.cpu().numpy(), atol=1e-5, rtol=1e-5) - np.testing.assert_allclose(actual_torque, expected_torque.cpu().numpy(), atol=1e-5, rtol=1e-5) - - @_backends - @_devices - @pytest.mark.parametrize("method_base, kwarg", _BODY_METHODS, ids=[m[0] for m in _BODY_METHODS]) - def test_body_writer_index(self, backend, device, method_base, kwarg): - num_instances, num_bodies = _NUM_INSTANCES, 1 - obj, _ = get_rigid_object(backend, num_instances, device) - wp_dtype, trailing, getter = _body_writer_layout(backend, method_base) - obj.data.update(dt=0.01) - method = getattr(obj, f"{method_base}_index") - sub_b = 1 # rigid object always has 1 body - sub_body_ids = [0] - - # torch, all envs + all bodies - method(**{kwarg: _make_body_torch((num_instances, num_bodies), device, wp_dtype, trailing)}) - # torch, subset - method( - **{ - kwarg: _make_body_torch((1, sub_b), device, wp_dtype, trailing), - "body_ids": sub_body_ids, - "env_ids": _make_env_ids(device, True), - } - ) - # warp, all envs + all bodies: the matching getter reads the written values back - method(**{kwarg: _make_body_warp((num_instances, num_bodies), device, wp_dtype, trailing, payload=True)}) - if getter is not None: - expected = _make_body_torch((num_instances, num_bodies), device, wp_dtype, trailing, payload=True) - _assert_reads_back(getattr(obj.data, getter), expected, getter) - # warp, subset - method( - **{ - kwarg: _make_body_warp((1, sub_b), device, wp_dtype, trailing), - "body_ids": sub_body_ids, - "env_ids": _make_env_ids(device, True), - } - ) - # negative: bad torch shape (extra env) - with pytest.raises((AssertionError, RuntimeError)): - method(**{kwarg: _make_body_torch((num_instances + 1, num_bodies), device, wp_dtype, trailing)}) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(**{kwarg: _make_body_warp((num_instances + 1, num_bodies), device, wp_dtype, trailing)}) - - @_backends - @_devices - @pytest.mark.parametrize("method_base, kwarg", _BODY_METHODS, ids=[m[0] for m in _BODY_METHODS]) - def test_body_writer_mask(self, backend, device, method_base, kwarg): - num_instances, num_bodies = _NUM_INSTANCES, 1 - obj, _ = get_rigid_object(backend, num_instances, device) - wp_dtype, trailing, getter = _body_writer_layout(backend, method_base) - obj.data.update(dt=0.01) - method = getattr(obj, f"{method_base}_mask") - sub_body_sel = [0] - - # torch, no mask - method(**{kwarg: _make_body_torch((num_instances, num_bodies), device, wp_dtype, trailing)}) - # torch, partial env_mask + body_mask - method( - **{ - kwarg: _make_body_torch((num_instances, num_bodies), device, wp_dtype, trailing), - "body_mask": _make_item_mask(num_bodies, sub_body_sel, device), - "env_mask": _make_env_mask(num_instances, device, True), - } - ) - # warp, no mask: the matching getter reads the written values back - method(**{kwarg: _make_body_warp((num_instances, num_bodies), device, wp_dtype, trailing, payload=True)}) - if getter is not None: - expected = _make_body_torch((num_instances, num_bodies), device, wp_dtype, trailing, payload=True) - _assert_reads_back(getattr(obj.data, getter), expected, getter) - # warp, partial env_mask + body_mask - method( - **{ - kwarg: _make_body_warp((num_instances, num_bodies), device, wp_dtype, trailing), - "body_mask": _make_item_mask(num_bodies, sub_body_sel, device), - "env_mask": _make_env_mask(num_instances, device, True), - } - ) - # negative: bad torch shape - with pytest.raises((AssertionError, RuntimeError)): - method(**{kwarg: _make_body_torch((num_instances + 1, num_bodies), device, wp_dtype, trailing)}) - # negative: bad warp shape - with pytest.raises((AssertionError, RuntimeError)): - method(**{kwarg: _make_body_warp((num_instances + 1, num_bodies), device, wp_dtype, trailing)}) diff --git a/source/isaaclab/test/envs/mdp/test_terminations.py b/source/isaaclab/test/envs/mdp/test_terminations.py new file mode 100644 index 000000000000..0c89b5d712e1 --- /dev/null +++ b/source/isaaclab/test/envs/mdp/test_terminations.py @@ -0,0 +1,31 @@ +# 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 + +"""Unit tests for MDP termination terms.""" + +from types import SimpleNamespace + +import torch + +from isaaclab.envs.mdp.terminations import joint_effort_out_of_limit +from isaaclab.managers import SceneEntityCfg + + +def test_joint_effort_limit_terminates_only_environments_with_clipped_selected_joints() -> None: + """Terminate exactly the environments whose applied effort differs from the computed effort on selected joints.""" + computed = torch.tensor([10.0, 20.0, 30.0]).repeat(4, 1) + # With joints 0 and 1 selected: environment 0 clips one selected joint, environment 1 clips nothing, + # environment 2 clips both selected joints, and environment 3 clips only the unselected joint. + applied = torch.tensor([[10.0, 19.0, 30.0], [10.0, 20.0, 30.0], [9.0, 17.0, 30.0], [10.0, 20.0, 29.0]]) + robot = SimpleNamespace( + actuators=SimpleNamespace( + computed_effort=SimpleNamespace(torch=computed), applied_effort=SimpleNamespace(torch=applied) + ) + ) + env = SimpleNamespace(scene={"robot": robot}) + + result = joint_effort_out_of_limit(env, SceneEntityCfg("robot", joint_ids=[0, 1])) + + torch.testing.assert_close(result, torch.tensor([True, False, True, False])) diff --git a/source/isaaclab/test/utils/test_wrench_composer.py b/source/isaaclab/test/utils/test_wrench_composer.py index a86d990193e8..10daf0d06612 100644 --- a/source/isaaclab/test/utils/test_wrench_composer.py +++ b/source/isaaclab/test/utils/test_wrench_composer.py @@ -3,10 +3,14 @@ # # SPDX-License-Identifier: BSD-3-Clause -import ast -import inspect +"""Literal unit tests for :class:`isaaclab.utils.wrench_composer.WrenchComposer`. + +Every case uses a two-environment, two-body composer on CPU with hand-computed expected values. A body rotated a +quarter turn about +Z maps a world vector ``(x, y, z)`` to the body vector ``(y, -x, z)``; a half turn maps it to +``(-x, -y, z)``. Real-solver delivery is covered by ``test_wrench_composer_integration.py``. +""" + from types import SimpleNamespace -from typing import get_type_hints from unittest.mock import patch import numpy as np @@ -14,669 +18,229 @@ import torch import warp as wp -from isaaclab.test.utils import test_devices from isaaclab.utils.warp import ProxyArray from isaaclab.utils.wrench_composer import WrenchComposer pytestmark = pytest.mark.unit +_IDENTITY = (0.0, 0.0, 0.0, 1.0) +_QUARTER_TURN_Z = (0.0, 0.0, 2.0**-0.5, 2.0**-0.5) +_HALF_TURN_Z = (0.0, 0.0, 1.0, 0.0) +_METHODS = [ + "add_forces_and_torques_index", + "add_forces_and_torques_mask", + "set_forces_and_torques_index", + "set_forces_and_torques_mask", +] +# ``(env, body)`` cell -> literal 3-vector or quaternion +_Cells = dict[tuple[int, int], tuple[float, ...]] -def test_wrench_composer_uses_asset_frame_conventions(): - """Keep frame selection as an is_global boolean, with no enum or content-classification layer.""" - tree = ast.parse(inspect.getsource(WrenchComposer)) - assert [node.name for node in ast.walk(tree) if isinstance(node, ast.ClassDef)] == ["WrenchComposer"] - assert not any( - isinstance(node, ast.Attribute) and node.attr in {"_content", "_classify"} for node in ast.walk(tree) - ) - module_tree = ast.parse(inspect.getsource(inspect.getmodule(WrenchComposer))) - assert not any(isinstance(node, ast.ImportFrom) and node.module == "enum" for node in ast.walk(module_tree)) - assert not hasattr(WrenchComposer, "resolve_submission") - assert get_type_hints(WrenchComposer.get_forces_and_torques)["return"] == tuple[wp.array, wp.array, bool] - - -def _get_wrench_without_pose_reads(composer: WrenchComposer) -> tuple[wp.array, wp.array, bool]: - """Read a fast-path wrench while rejecting the expensive body-pose queries.""" - with patch.object(composer._asset, "data", None): - return composer.get_forces_and_torques() - - -class _WrenchAssetDataFixture: - """Minimal asset data required by :class:`WrenchComposer`.""" - - def __init__(self, com_pos_w: torch.Tensor, link_quat_w: torch.Tensor, device: str) -> None: - self._device = device - self.body_com_pos_w = ProxyArray(wp.from_torch(com_pos_w.to(device), dtype=wp.vec3f)) - self.body_link_quat_w = ProxyArray(wp.from_torch(link_quat_w.to(device), dtype=wp.quatf)) - - def set_body_com_pose_w(self, pose_w: torch.Tensor) -> None: - self.body_com_pos_w = ProxyArray(wp.from_torch(pose_w[..., :3].to(self._device), dtype=wp.vec3f)) +# --------------------------------------------------------------------------- +# Literal fixture helpers +# --------------------------------------------------------------------------- -def create_mock_asset( - num_envs: int, - num_bodies: int, - device: str, - link_pos: torch.Tensor | None = None, - link_quat: torch.Tensor | None = None, -) -> SimpleNamespace: - """Create a minimal asset fixture with optional custom link poses. - Args: - num_envs: Number of environments. - num_bodies: Number of bodies. - device: Device to use. - link_pos: Optional link positions (num_envs, num_bodies, 3). Defaults to zeros. - link_quat: Optional link quaternions in (x, y, z, w) format (num_envs, num_bodies, 4). - Defaults to identity quaternion. +class _AssetData: + """Body poses that :class:`WrenchComposer` reads when it composes into the body frame.""" - Returns: - Asset fixture with the state required by WrenchComposer. - """ + def __init__(self, com_pos_w: torch.Tensor, link_quat_w: torch.Tensor) -> None: + self.set_pose(com_pos_w, link_quat_w) - # Build combined pose (N, B, 7) = pos(3) + quat_xyzw(4) matching wp.transformf layout - if link_pos is None: - pos = torch.zeros(num_envs, num_bodies, 3, dtype=torch.float32) - else: - pos = link_pos.float() - - if link_quat is None: - # Identity quaternion in (x, y, z, w) format = (0, 0, 0, 1) - quat = torch.zeros(num_envs, num_bodies, 4, dtype=torch.float32) - quat[..., 3] = 1.0 - else: - quat = link_quat.float() + def set_pose(self, com_pos_w: torch.Tensor, link_quat_w: torch.Tensor) -> None: + self.body_com_pos_w = ProxyArray(wp.from_torch(com_pos_w, dtype=wp.vec3f)) + self.body_link_quat_w = ProxyArray(wp.from_torch(link_quat_w, dtype=wp.quatf)) - return SimpleNamespace( - num_instances=num_envs, num_bodies=num_bodies, device=device, data=_WrenchAssetDataFixture(pos, quat, device) - ) - - -# --- Helper functions for quaternion math --- - - -def quat_rotate_inv_np(quat_xyzw: np.ndarray, vec: np.ndarray) -> np.ndarray: - """Rotate a vector by the inverse of a quaternion (numpy). - - Args: - quat_xyzw: Quaternion in (x, y, z, w) format. Shape: (..., 4) - vec: Vector to rotate. Shape: (..., 3) - - Returns: - Rotated vector. Shape: (..., 3) - """ - # Extract components - xyz = quat_xyzw[..., 0:3] - w = quat_xyzw[..., 3:4] - - # For inverse rotation, we conjugate the quaternion (negate xyz) - # q^-1 * v * q = q_conj * v * q_conj^-1 for unit quaternion - # Using the formula: v' = v + 2*w*(xyz x v) + 2*(xyz x (xyz x v)) - # But for inverse: use -xyz - - # Cross product: xyz x vec - t = 2.0 * np.cross(-xyz, vec, axis=-1) - # Result: vec + w*t + xyz x t - return vec + w * t + np.cross(-xyz, t, axis=-1) - - -def random_unit_quaternion_np(rng: np.random.Generator, shape: tuple) -> np.ndarray: - """Generate random unit quaternions in (x, y, z, w) format. - - Args: - rng: Random number generator. - shape: Output shape, e.g. (num_envs, num_bodies). - - Returns: - Random unit quaternions. Shape: (*shape, 4) - """ - # Generate random quaternion components - q = rng.standard_normal(shape + (4,)).astype(np.float32) - # Normalize to unit quaternion - q = q / np.linalg.norm(q, axis=-1, keepdims=True) - return q - - -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize( - "num_envs, num_bodies, use_forces, use_torques, use_positions", - [ - pytest.param(7, 3, True, False, False, id="forces"), - pytest.param(7, 3, False, True, False, id="torques"), - pytest.param(7, 3, True, False, True, id="forces_at_positions"), - pytest.param(7, 3, False, True, True, id="torques_ignore_positions"), - pytest.param(7, 3, True, True, True, id="forces_and_torques_at_positions"), - pytest.param(1, 1, True, True, True, id="forces_and_torques_at_positions_single_body"), - ], -) -def test_add_forces_and_torques_at_position( - device: str, num_envs: int, num_bodies: int, use_forces: bool, use_torques: bool, use_positions: bool -): - """Test accumulating local forces and torques, optionally at local positions, over random index subsets.""" - rng = np.random.default_rng(seed=4) - - mock_asset = create_mock_asset(num_envs, num_bodies, device) - wrench_composer = WrenchComposer(mock_asset) - # Initialize hand-calculated composed force and torque - hand_calculated_composed_force_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) - hand_calculated_composed_torque_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) - for _ in range(10): - # Get random number of envs and bodies and their indices - num_envs_np = rng.integers(1, num_envs, endpoint=True) - num_bodies_np = rng.integers(1, num_bodies, endpoint=True) - env_ids_np = rng.choice(num_envs, size=num_envs_np, replace=False) - body_ids_np = rng.choice(num_bodies, size=num_bodies_np, replace=False) - selected = (env_ids_np[:, None], body_ids_np[None, :]) - # Get random forces, torques, and positions - forces_np, torques_np, positions_np = rng.uniform(-100.0, 100.0, (3, num_envs_np, num_bodies_np, 3)).astype( - np.float32 - ) - wrench = {} - if use_forces: - wrench["forces"] = wp.from_numpy(forces_np, dtype=wp.vec3f, device=device) - hand_calculated_composed_force_np[selected] += forces_np - if use_torques: - wrench["torques"] = wp.from_numpy(torques_np, dtype=wp.vec3f, device=device) - hand_calculated_composed_torque_np[selected] += torques_np - if use_positions: - wrench["positions"] = wp.from_numpy(positions_np, dtype=wp.vec3f, device=device) - # torque = cross(position, force); positions are ignored for torque-only input - if use_forces: - hand_calculated_composed_torque_np[selected] += np.cross(positions_np, forces_np) - # Add the wrench to the wrench composer - wrench_composer.add_forces_and_torques_index( - **wrench, - body_ids=wp.from_numpy(body_ids_np, dtype=wp.int32, device=device), - env_ids=wp.from_numpy(env_ids_np, dtype=wp.int32, device=device), - ) - # Compose to body frame before checking output - wrench_composer.compose_to_body_frame() - composed_force_np = wrench_composer.out_force_b.warp.numpy() - composed_torque_np = wrench_composer.out_torque_b.warp.numpy() - np.testing.assert_allclose(composed_force_np, hand_calculated_composed_force_np, atol=1e-2, rtol=1e-5) - np.testing.assert_allclose(composed_torque_np, hand_calculated_composed_torque_np, atol=1e-2, rtol=1e-5) - - -# ============================================================================ -# Global Frame Tests -# ============================================================================ - - -@pytest.mark.parametrize("device", test_devices()) -def test_global_forces_with_rotation(device: str): - """Test that global forces are correctly rotated to the local frame.""" - num_envs, num_bodies = 7, 3 - rng = np.random.default_rng(seed=10) - - for method in ( - "add_forces_and_torques_index", - "add_forces_and_torques_mask", - "set_forces_and_torques_index", - "set_forces_and_torques_mask", - ): - # Create random link quaternions - link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) - link_quat_torch = torch.from_numpy(link_quat_np) - - # Create mock asset with custom quaternions - mock_asset = create_mock_asset(num_envs, num_bodies, device, link_quat=link_quat_torch) - wrench_composer = WrenchComposer(mock_asset, supports_world_at_com=True) - - # Generate random global forces for all envs and bodies - forces_global_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) - forces_global = wp.from_numpy(forces_global_np, dtype=wp.vec3f, device=device) - - # Apply global forces - getattr(wrench_composer, method)(forces=forces_global, is_global=True) - - # Compute expected local forces by rotating global forces by inverse quaternion - expected_forces_local = quat_rotate_inv_np(link_quat_np, forces_global_np) - - # Check raw global buffer has the global forces - force, torque, is_global = _get_wrench_without_pose_reads(wrench_composer) - assert is_global is True - np.testing.assert_array_equal(torque.numpy(), np.zeros_like(forces_global_np)) - global_force_np = force.numpy() - assert np.allclose(global_force_np, forces_global_np, atol=1e-4, rtol=1e-5), ( - f"Global force buffer mismatch.\nExpected:\n{forces_global_np}\nGot:\n{global_force_np}" - ) - - # Compose to body frame before checking output - wrench_composer.compose_to_body_frame() - - # Verify - composed_force_np = wrench_composer.out_force_b.warp.numpy() - assert np.allclose(composed_force_np, expected_forces_local, atol=1e-4, rtol=1e-5), ( - f"Global force rotation failed.\nExpected:\n{expected_forces_local}\nGot:\n{composed_force_np}" - ) +def _poses( + com: _Cells | None = None, quat: _Cells | None = None, shape: tuple[int, int] = (2, 2) +) -> tuple[torch.Tensor, torch.Tensor]: + """Return CoM positions and link quaternions: zero and identity except for the given ``(env, body)`` cells.""" + com_pos_w = torch.zeros((*shape, 3)) + link_quat_w = torch.tensor(_IDENTITY).repeat(*shape, 1) + for cell, value in (com or {}).items(): + com_pos_w[cell] = torch.tensor(value) + for cell, value in (quat or {}).items(): + link_quat_w[cell] = torch.tensor(value) + return com_pos_w, link_quat_w -@pytest.mark.parametrize("device", test_devices()) -def test_global_torques_with_rotation(device: str): - """Test that global torques are correctly rotated to the local frame.""" - num_envs, num_bodies = 7, 3 - rng = np.random.default_rng(seed=11) - for method in ( - "add_forces_and_torques_index", - "add_forces_and_torques_mask", - "set_forces_and_torques_index", - "set_forces_and_torques_mask", - ): - # Create random link quaternions - link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) - link_quat_torch = torch.from_numpy(link_quat_np) - - # Create mock asset with custom quaternions - mock_asset = create_mock_asset(num_envs, num_bodies, device, link_quat=link_quat_torch) - wrench_composer = WrenchComposer(mock_asset, supports_world_at_com=True) - - # Generate random global torques - torques_global_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) - torques_global = wp.from_numpy(torques_global_np, dtype=wp.vec3f, device=device) - - # Positions are unused when only torques are supplied. - getattr(wrench_composer, method)(torques=torques_global, positions=torques_global, is_global=True) - - # Compute expected local torques - expected_torques_local = quat_rotate_inv_np(link_quat_np, torques_global_np) - - # Check raw global buffer has the global torques - force, torque, is_global = _get_wrench_without_pose_reads(wrench_composer) - assert is_global is True - np.testing.assert_array_equal(force.numpy(), np.zeros_like(torques_global_np)) - global_torque_np = torque.numpy() - assert np.allclose(global_torque_np, torques_global_np, atol=1e-4, rtol=1e-5), ( - f"Global torque buffer mismatch.\nExpected:\n{torques_global_np}\nGot:\n{global_torque_np}" - ) +def _make_composer( + com: _Cells | None = None, quat: _Cells | None = None, *, supports_world_at_com: bool = False +) -> WrenchComposer: + """Create a two-environment, two-body CPU composer with literal body poses.""" + asset = SimpleNamespace(num_instances=2, num_bodies=2, device="cpu", data=_AssetData(*_poses(com, quat))) + return WrenchComposer(asset, supports_world_at_com=supports_world_at_com) - # Compose to body frame before checking output - wrench_composer.compose_to_body_frame() - # Verify - composed_torque_np = wrench_composer.out_torque_b.warp.numpy() - assert np.allclose(composed_torque_np, expected_torques_local, atol=1e-4, rtol=1e-5), ( - f"Global torque rotation failed.\nExpected:\n{expected_torques_local}\nGot:\n{composed_torque_np}" - ) +def _grid(cells: _Cells, shape: tuple[int, int] = (2, 2)) -> torch.Tensor: + """Return a ``(num_envs, num_bodies, 3)`` tensor that is zero except for the given ``(env, body)`` cells.""" + values = torch.zeros((*shape, 3)) + for cell, value in cells.items(): + values[cell] = torch.tensor(value) + return values -@pytest.mark.parametrize("device", test_devices()) -def test_global_forces_at_global_position(device: str): - """Test global forces at global positions with full coordinate transformation.""" - num_envs, num_bodies = 7, 3 - rng = np.random.default_rng(seed=12) +def _vectors(values: torch.Tensor) -> wp.array: + """Convert a contiguous ``(..., 3)`` tensor to a Warp vector array.""" + return wp.from_torch(values, dtype=wp.vec3f) - for method in ( - "add_forces_and_torques_index", - "add_forces_and_torques_mask", - "set_forces_and_torques_index", - "set_forces_and_torques_mask", - ): - # Create random link poses - link_pos_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) - link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) - link_pos_torch = torch.from_numpy(link_pos_np) - link_quat_torch = torch.from_numpy(link_quat_np) - - # Create mock asset - mock_asset = create_mock_asset(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) - wrench_composer = WrenchComposer(mock_asset, supports_world_at_com=True) - - # Generate random global forces and positions - forces_global_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) - positions_global_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) - forces_global = wp.from_numpy(forces_global_np, dtype=wp.vec3f, device=device) - positions_global = wp.from_numpy(positions_global_np, dtype=wp.vec3f, device=device) - - # Apply global forces at global positions - getattr(wrench_composer, method)(forces=forces_global, positions=positions_global, is_global=True) - - # Compute expected results: - # 1. Force in local frame = quat_rotate_inv(link_quat, global_force) - expected_forces_local = quat_rotate_inv_np(link_quat_np, forces_global_np) - - # 2. Torque about CoM in world frame = cross(P_global - link_pos, F_global) - # Then rotate to body frame - position_offset_global = positions_global_np - link_pos_np - expected_torques_local = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) - for i in range(num_envs): - for j in range(num_bodies): - torque_w = np.cross(position_offset_global[i, j], forces_global_np[i, j]) - expected_torques_local[i, j] = quat_rotate_inv_np( - link_quat_np[i : i + 1, j : j + 1], torque_w.reshape(1, 1, 3) - )[0, 0] - - # Check raw global force buffer has the global forces - global_force_np = wrench_composer.global_force_w.numpy() - assert np.allclose(global_force_np, forces_global_np, atol=1e-4, rtol=1e-5), ( - f"Global force buffer mismatch.\nExpected:\n{forces_global_np}\nGot:\n{global_force_np}" - ) - force, torque, is_global = wrench_composer.get_forces_and_torques() - assert is_global is False +def _mask(values: list[bool]) -> wp.array: + """Convert literal booleans to a CPU Warp mask array.""" + return wp.array(values, dtype=wp.bool, device="cpu") - # Verify forces - composed_force_np = force.numpy() - assert np.allclose(composed_force_np, expected_forces_local, atol=1e-3, rtol=1e-4), ( - f"Global force at position failed.\nExpected forces:\n{expected_forces_local}\nGot:\n{composed_force_np}" - ) - # Verify torques - composed_torque_np = torque.numpy() - assert np.allclose(composed_torque_np, expected_torques_local, atol=1e-3, rtol=1e-4), ( - f"Global force at position failed.\nExpected torques:\n{expected_torques_local}\nGot:\n{composed_torque_np}" - ) +def _apply_to_cell(composer: WrenchComposer, method: str, cell: tuple[int, int], is_global: bool, **vectors) -> None: + """Apply literal ``forces``/``torques``/``positions`` to one ``(env, body)`` cell through any entry point.""" + env, body = cell + if method.endswith("_index"): + inputs = {name: _vectors(torch.tensor(value).reshape(1, 1, 3)) for name, value in vectors.items()} + getattr(composer, method)(**inputs, env_ids=[env], body_ids=[body], is_global=is_global) + else: + # Fill the masked-out cells with a sentinel so a mask that leaks shows up in the result. + inputs = {} + for name, value in vectors.items(): + full = torch.full((2, 2, 3), 999.0) + full[cell] = torch.tensor(value) + inputs[name] = _vectors(full) + env_mask, body_mask = _mask([i == env for i in range(2)]), _mask([i == body for i in range(2)]) + getattr(composer, method)(**inputs, env_mask=env_mask, body_mask=body_mask, is_global=is_global) -@pytest.mark.parametrize("device", test_devices()) -def test_90_degree_rotation_global_force(device: str): - """Test global force with a known 90-degree rotation for easy verification.""" - num_envs, num_bodies = 1, 1 +def _get_wrench_without_pose_reads(composer: WrenchComposer) -> tuple[wp.array, wp.array, bool]: + """Read a fast-path wrench while rejecting the body-pose queries that composition needs.""" + with patch.object(composer._asset, "data", None): + return composer.get_forces_and_torques() - # 90-degree rotation around Z-axis: (x, y, z, w) = (0, 0, sin(45°), cos(45°)) - # This rotates X -> Y, Y -> -X - angle = np.pi / 2 - link_quat_np = np.array([[[[0, 0, np.sin(angle / 2), np.cos(angle / 2)]]]], dtype=np.float32).reshape(1, 1, 4) - link_quat_torch = torch.from_numpy(link_quat_np) - mock_asset = create_mock_asset(num_envs, num_bodies, device, link_quat=link_quat_torch) - wrench_composer = WrenchComposer(mock_asset) +def _assert_vectors(actual: wp.array | ProxyArray, expected: torch.Tensor, atol: float = 1.0e-6) -> None: + """Compare a Warp vector buffer (or output proxy) with an expected tensor.""" + actual = actual.torch if isinstance(actual, ProxyArray) else wp.to_torch(actual) + torch.testing.assert_close(actual, expected, atol=atol, rtol=0.0) - # Apply force in global +X direction - force_global = np.array([[[1.0, 0.0, 0.0]]], dtype=np.float32) - force_wp = wp.from_numpy(force_global, dtype=wp.vec3f, device=device) - wrench_composer.add_forces_and_torques_index(forces=force_wp, is_global=True) +# --------------------------------------------------------------------------- +# Frames and composition +# --------------------------------------------------------------------------- - # Expected: After inverse rotation (rotate by -90° around Z), X becomes -Y - # Actually, inverse rotation of +90° around Z applied to (1,0,0) gives (0,-1,0) - expected_force_local = np.array([[[0.0, -1.0, 0.0]]], dtype=np.float32) - force, _, is_global = wrench_composer.get_forces_and_torques() - assert is_global is False +@pytest.mark.parametrize("method", _METHODS) +def test_local_wrench_at_position_composes_without_pose_reads(method: str) -> None: + """Local forces at local offsets add ``cross(p, F)``; torque-only input ignores positions; no pose is read.""" + composer = _make_composer(quat={(1, 0): _QUARTER_TURN_Z}, supports_world_at_com=True) - composed_force_np = force.numpy() - assert np.allclose(composed_force_np, expected_force_local, atol=1e-5), ( - f"90-degree rotation test failed.\nExpected:\n{expected_force_local}\nGot:\n{composed_force_np}" + _apply_to_cell( + composer, method, (1, 0), False, forces=(2.0, 0.0, 0.0), torques=(0.0, 0.0, 5.0), positions=(0.0, 3.0, 0.0) ) + _apply_to_cell(composer, method, (0, 1), False, torques=(1.0, 0.0, 0.0), positions=(0.0, 3.0, 0.0)) - # A second read must use the live pose even though the wrench buffers have not changed. - quat_180_z = torch.tensor([[[0.0, 0.0, 1.0, 0.0]]], device=device) - mock_asset.data.body_link_quat_w = ProxyArray(wp.from_torch(quat_180_z, dtype=wp.quatf)) - force, _, is_global = wrench_composer.get_forces_and_torques() + force, torque, is_global = _get_wrench_without_pose_reads(composer) assert is_global is False - np.testing.assert_allclose(force.numpy(), [[[-1.0, 0.0, 0.0]]], atol=1e-6) - + _assert_vectors(force, _grid({(1, 0): (2.0, 0.0, 0.0)})) + # cross((0, 3, 0), (2, 0, 0)) = (0, 0, -6) is added to the local torque (0, 0, 5). + _assert_vectors(torque, _grid({(1, 0): (0.0, 0.0, -1.0), (0, 1): (1.0, 0.0, 0.0)})) -@pytest.mark.parametrize("device", test_devices()) -def test_composition_mixed_local_and_global(device: str): - """Test that local and global forces can be composed together correctly.""" - rng = np.random.default_rng(seed=14) - num_envs, num_bodies = 5, 3 - # Create random link quaternions - link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) - link_quat_torch = torch.from_numpy(link_quat_np) +def test_add_accumulates_local_wrenches() -> None: + composer = _make_composer() - mock_asset = create_mock_asset(num_envs, num_bodies, device, link_quat=link_quat_torch) - wrench_composer = WrenchComposer(mock_asset, supports_world_at_com=True) + for force, torque in (((1.0, 2.0, 3.0), (0.0, 1.0, 0.0)), ((4.0, -1.0, 0.0), (2.0, 0.0, -3.0))): + _apply_to_cell(composer, "add_forces_and_torques_index", (0, 1), False, forces=force, torques=torque) - # Generate random local and global forces - forces_local_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) - forces_global_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + _assert_vectors(composer.out_force_b, _grid({(0, 1): (5.0, 1.0, 3.0)})) + _assert_vectors(composer.out_torque_b, _grid({(0, 1): (2.0, 1.0, -3.0)})) - forces_local = wp.from_numpy(forces_local_np, dtype=wp.vec3f, device=device) - forces_global = wp.from_numpy(forces_global_np, dtype=wp.vec3f, device=device) - # Add local forces first - wrench_composer.add_forces_and_torques_index(forces=forces_local, is_global=False) +@pytest.mark.parametrize("method", _METHODS) +def test_global_wrench_at_com_uses_world_fast_path_or_rotates(method: str) -> None: + """A world-at-CoM consumer receives the world wrench as-is; a body-frame consumer receives it rotated.""" + composers = [_make_composer(quat={(0, 0): _QUARTER_TURN_Z}, supports_world_at_com=s) for s in (True, False)] + for composer in composers: + _apply_to_cell(composer, method, (0, 0), True, forces=(2.0, 0.0, 0.0)) + # Positions are unused when only torques are supplied, so the wrench stays a world-at-CoM wrench. + _apply_to_cell(composer, method, (1, 1), True, torques=(0.0, 3.0, 0.0), positions=(5.0, 5.0, 5.0)) + world_consumer, body_consumer = composers - # Add global forces - wrench_composer.add_forces_and_torques_index(forces=forces_global, is_global=True) + force, torque, is_global = _get_wrench_without_pose_reads(world_consumer) + assert is_global is True + _assert_vectors(force, _grid({(0, 0): (2.0, 0.0, 0.0)})) + _assert_vectors(torque, _grid({(1, 1): (0.0, 3.0, 0.0)})) - # Expected: local forces stay as-is, global forces get rotated, then sum - global_forces_in_local = quat_rotate_inv_np(link_quat_np, forces_global_np) - expected_total = forces_local_np + global_forces_in_local - - # Check raw buffer properties - local_force_np = wrench_composer.local_force_b.numpy() - assert np.allclose(local_force_np, forces_local_np, atol=1e-4, rtol=1e-5) - global_force_at_com_np = wrench_composer.global_force_at_com_w.numpy() - assert np.allclose(global_force_at_com_np, forces_global_np, atol=1e-4, rtol=1e-5) - - force, _, is_global = wrench_composer.get_forces_and_torques() + force, torque, is_global = body_consumer.get_forces_and_torques() assert is_global is False + _assert_vectors(force, _grid({(0, 0): (0.0, -2.0, 0.0)})) + _assert_vectors(torque, _grid({(1, 1): (0.0, 3.0, 0.0)})) + # The body-frame outputs of a world-at-CoM composer hold the same rotated wrench. + _assert_vectors(world_consumer.out_force_b, _grid({(0, 0): (0.0, -2.0, 0.0)})) - composed_force_np = force.numpy() - assert np.allclose(composed_force_np, expected_total, atol=1e-4, rtol=1e-5), ( - f"Mixed local/global composition failed.\nExpected:\n{expected_total}\nGot:\n{composed_force_np}" - ) - - -@pytest.mark.parametrize("device", test_devices()) -def test_local_forces_at_local_position(device: str): - """Test local forces at local positions (offset from link frame).""" - num_envs, num_bodies = 7, 3 - rng = np.random.default_rng(seed=15) - - for method in ( - "add_forces_and_torques_index", - "add_forces_and_torques_mask", - "set_forces_and_torques_index", - "set_forces_and_torques_mask", - ): - # Create random link poses (shouldn't affect local frame calculations) - link_pos_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) - link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) - link_pos_torch = torch.from_numpy(link_pos_np) - link_quat_torch = torch.from_numpy(link_quat_np) - - mock_asset = create_mock_asset(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) - wrench_composer = WrenchComposer(mock_asset, supports_world_at_com=True) - - # Generate random local forces and local positions (offsets) - forces_local_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) - positions_local_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) - forces_local = wp.from_numpy(forces_local_np, dtype=wp.vec3f, device=device) - positions_local = wp.from_numpy(positions_local_np, dtype=wp.vec3f, device=device) - - # Apply local forces at local positions - getattr(wrench_composer, method)(forces=forces_local, positions=positions_local, is_global=False) - - # Expected: forces stay as-is, torque = cross(position, force) - expected_forces = forces_local_np - expected_torques = np.cross(positions_local_np, forces_local_np) - - # Check raw local buffer - local_force_np = wrench_composer.local_force_b.numpy() - assert np.allclose(local_force_np, expected_forces, atol=1e-4, rtol=1e-5) - - force, torque, is_global = _get_wrench_without_pose_reads(wrench_composer) - assert is_global is False - - # Verify - composed_force_np = force.numpy() - composed_torque_np = torque.numpy() - assert np.allclose(composed_force_np, expected_forces, atol=1e-4, rtol=1e-5) - assert np.allclose(composed_torque_np, expected_torques, atol=1e-4, rtol=1e-5) - - -# ============================================================================ -# add_raw_buffers_from Tests -# ============================================================================ - - -@pytest.mark.parametrize("device", test_devices()) -def test_add_raw_buffers_from(device: str): - """Test that add_raw_buffers_from merges all five input buffers correctly.""" - num_envs, num_bodies = 7, 3 - rng = np.random.default_rng(seed=20) - - # Create two composers with random link poses - link_pos_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) - link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) - link_pos_torch = torch.from_numpy(link_pos_np) - link_quat_torch = torch.from_numpy(link_quat_np) - - mock_a = create_mock_asset(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) - mock_b = create_mock_asset(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) - - composer_a = WrenchComposer(mock_a, supports_world_at_com=True) - composer_b = WrenchComposer(mock_b) - - # Populate composer_a with local forces at positions - forces_local_a_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) - positions_local_a_np = rng.uniform(-5.0, 5.0, (num_envs, num_bodies, 3)).astype(np.float32) - composer_a.add_forces_and_torques_index( - forces=wp.from_numpy(forces_local_a_np, dtype=wp.vec3f, device=device), - positions=wp.from_numpy(positions_local_a_np, dtype=wp.vec3f, device=device), - is_global=False, +@pytest.mark.parametrize("method", _METHODS) +def test_global_force_at_position_uses_live_com_and_rotation(method: str) -> None: + """World positions add torque about the live CoM of each body, even far from the origin.""" + composer = _make_composer( + com={(0, 0): (1.0, 2.0, 3.0), (1, 1): (2000.0, 0.0, 1.0)}, + quat={(0, 0): _QUARTER_TURN_Z}, + supports_world_at_com=True, ) + _apply_to_cell(composer, method, (0, 0), True, forces=(2.0, 0.0, 0.0), positions=(1.0, 4.0, 3.0)) + _apply_to_cell(composer, method, (1, 1), True, forces=(0.0, 10.0, 0.0), positions=(2001.0, 0.0, 1.0)) - # Populate composer_b with global forces at global positions - forces_global_b_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) - positions_global_b_np = rng.uniform(-5.0, 5.0, (num_envs, num_bodies, 3)).astype(np.float32) - composer_b.add_forces_and_torques_index( - forces=wp.from_numpy(forces_global_b_np, dtype=wp.vec3f, device=device), - positions=wp.from_numpy(positions_global_b_np, dtype=wp.vec3f, device=device), - is_global=True, - ) - - # Merge b into a - composer_a.add_raw_buffers_from(composer_b) - - # Build a reference composer that receives both calls directly - mock_ref = create_mock_asset(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) - composer_ref = WrenchComposer(mock_ref) - composer_ref.add_forces_and_torques_index( - forces=wp.from_numpy(forces_local_a_np, dtype=wp.vec3f, device=device), - positions=wp.from_numpy(positions_local_a_np, dtype=wp.vec3f, device=device), - is_global=False, - ) - composer_ref.add_forces_and_torques_index( - forces=wp.from_numpy(forces_global_b_np, dtype=wp.vec3f, device=device), - positions=wp.from_numpy(positions_global_b_np, dtype=wp.vec3f, device=device), - is_global=True, - ) - - # Compose both and compare - force, torque, is_global = composer_a.get_forces_and_torques() + _assert_vectors(composer.global_force_w, _grid({(0, 0): (2.0, 0.0, 0.0), (1, 1): (0.0, 10.0, 0.0)})) + force, torque, is_global = composer.get_forces_and_torques() assert is_global is False - composer_ref.compose_to_body_frame() + _assert_vectors(force, _grid({(0, 0): (0.0, -2.0, 0.0), (1, 1): (0.0, 10.0, 0.0)})) + # Lever arms (0, 2, 0) and (1, 0, 0): world torques (0, 0, -4) and (0, 0, 10). + _assert_vectors(torque, _grid({(0, 0): (0.0, 0.0, -4.0), (1, 1): (0.0, 0.0, 10.0)}), atol=1.0e-3) - assert np.allclose(force.numpy(), composer_ref.out_force_b.warp.numpy(), atol=1e-4, rtol=1e-5), ( - "add_raw_buffers_from force mismatch vs direct accumulation" - ) - assert np.allclose(torque.numpy(), composer_ref.out_torque_b.warp.numpy(), atol=1e-4, rtol=1e-5), ( - "add_raw_buffers_from torque mismatch vs direct accumulation" + # Without new input, the next read uses the live pose: the CoM moved onto the force line and the body half-turned. + composer._asset.data.set_pose( + *_poses(com={(0, 0): (1.0, 4.0, 3.0), (1, 1): (2000.0, 0.0, 1.0)}, quat={(0, 0): _HALF_TURN_Z}) ) + force, torque, is_global = composer.get_forces_and_torques() + assert is_global is False + _assert_vectors(force, _grid({(0, 0): (-2.0, 0.0, 0.0), (1, 1): (0.0, 10.0, 0.0)})) + _assert_vectors(torque, _grid({(1, 1): (0.0, 0.0, 10.0)}), atol=1.0e-3) -@pytest.mark.parametrize("device", test_devices()) -def test_add_raw_buffers_from_inactive_is_noop(device: str): - """Test that add_raw_buffers_from is a no-op when the source composer is inactive.""" - num_envs, num_bodies = 4, 2 - rng = np.random.default_rng(seed=21) - - mock_a = create_mock_asset(num_envs, num_bodies, device) - mock_b = create_mock_asset(num_envs, num_bodies, device) - composer_a = WrenchComposer(mock_a) - composer_b = WrenchComposer(mock_b) - - # Populate composer_a with some forces - forces_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) - composer_a.add_forces_and_torques_index( - forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), - ) - - # composer_b is inactive (never written to) - assert not composer_b.active - - # Snapshot composer_a's local buffer before merge - local_force_before = composer_a.local_force_b.numpy().copy() - - # Merge inactive composer_b into composer_a -- should be a no-op - composer_a.add_raw_buffers_from(composer_b) - - assert np.allclose(composer_a.local_force_b.numpy(), local_force_before, atol=1e-7) - - # Merging an inactive source must not activate an inactive target - composer_c = WrenchComposer(create_mock_asset(num_envs, num_bodies, device)) - composer_c.add_raw_buffers_from(composer_b) - assert not composer_c.active - - -# ============================================================================ -# Mask-based API Tests -# ============================================================================ - - -@pytest.mark.parametrize("device", test_devices()) -def test_add_forces_mask(device: str): - """Test that add_forces_and_torques_mask produces the same result as the index variant.""" - num_envs, num_bodies = 7, 3 - rng = np.random.default_rng(seed=30) - - for _ in range(5): - # Random subset selection - env_select = rng.choice([True, False], size=num_envs, replace=True) - body_select = rng.choice([True, False], size=num_bodies, replace=True) - # Ensure at least one env and body are selected - env_select[0] = True - body_select[0] = True +def test_mixed_local_and_global_wrenches_compose_in_body_frame() -> None: + composer = _make_composer(quat={(0, 0): _QUARTER_TURN_Z}, supports_world_at_com=True) - env_ids_np = np.where(env_select)[0].astype(np.int32) - body_ids_np = np.where(body_select)[0].astype(np.int32) - env_mask_np = env_select.astype(np.bool_) - body_mask_np = body_select.astype(np.bool_) + _apply_to_cell(composer, "add_forces_and_torques_index", (0, 0), False, forces=(1.0, 0.0, 0.0)) + _apply_to_cell(composer, "add_forces_and_torques_index", (0, 0), True, forces=(2.0, 0.0, 0.0)) - # Random forces for the full grid (mask variant takes full-sized arrays) - forces_full_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + _assert_vectors(composer.local_force_b, _grid({(0, 0): (1.0, 0.0, 0.0)})) + _assert_vectors(composer.global_force_at_com_w, _grid({(0, 0): (2.0, 0.0, 0.0)})) + force, _, is_global = composer.get_forces_and_torques() + assert is_global is False + _assert_vectors(force, _grid({(0, 0): (1.0, -2.0, 0.0)})) - # Index-based composer - mock_idx = create_mock_asset(num_envs, num_bodies, device) - composer_idx = WrenchComposer(mock_idx) - # Extract the subset for index API - forces_subset_np = forces_full_np[env_ids_np[:, None], body_ids_np[None, :], :] - composer_idx.add_forces_and_torques_index( - forces=wp.from_numpy(forces_subset_np, dtype=wp.vec3f, device=device), - env_ids=wp.from_numpy(env_ids_np, dtype=wp.int32, device=device), - body_ids=wp.from_numpy(body_ids_np, dtype=wp.int32, device=device), - ) - # Mask-based composer - mock_mask = create_mock_asset(num_envs, num_bodies, device) - composer_mask = WrenchComposer(mock_mask) - composer_mask.add_forces_and_torques_mask( - forces=wp.from_numpy(forces_full_np, dtype=wp.vec3f, device=device), - env_mask=wp.from_numpy(env_mask_np, dtype=wp.bool, device=device), - body_mask=wp.from_numpy(body_mask_np, dtype=wp.bool, device=device), - ) +# --------------------------------------------------------------------------- +# Selection, set, reset, and merge +# --------------------------------------------------------------------------- - # Compose both - composer_idx.compose_to_body_frame() - composer_mask.compose_to_body_frame() - assert np.allclose( - composer_idx.out_force_b.warp.numpy(), composer_mask.out_force_b.warp.numpy(), atol=1e-4, rtol=1e-5 - ), f"Mask vs index force mismatch (envs={num_envs}, bodies={num_bodies})" - assert np.allclose( - composer_idx.out_torque_b.warp.numpy(), composer_mask.out_torque_b.warp.numpy(), atol=1e-4, rtol=1e-5 - ), f"Mask vs index torque mismatch (envs={num_envs}, bodies={num_bodies})" +def test_index_and_mask_selection_change_only_selected_cells() -> None: + composer = _make_composer() + composer.add_forces_and_torques_index( + forces=_vectors(torch.tensor([[[1.0, 0.0, 0.0]]])), body_ids=torch.tensor([1]), env_ids=torch.tensor([0]) + ) + composer.add_forces_and_torques_mask( + forces=_vectors(torch.tensor([[[10.0, 0.0, 0.0], [20.0, 0.0, 0.0]], [[30.0, 0.0, 0.0], [40.0, 0.0, 0.0]]])), + env_mask=_mask([False, True]), + body_mask=_mask([True, False]), + ) -# ============================================================================ -# set_forces_and_torques_index Tests -# ============================================================================ + _assert_vectors(composer.out_force_b, _grid({(0, 1): (1.0, 0.0, 0.0), (1, 0): (30.0, 0.0, 0.0)})) -@pytest.mark.parametrize("env_dtype", [torch.int32, torch.int64]) -@pytest.mark.parametrize("body_dtype", [torch.int32, torch.int64]) -def test_index_dtype_combinations_preserve_selected_wrench_cells( - env_dtype: torch.dtype, body_dtype: torch.dtype -) -> None: +@pytest.mark.parametrize("env_dtype, body_dtype", [(torch.int32, torch.int64), (torch.int64, torch.int32)]) +def test_index_dtypes_preserve_selected_wrench_cells(env_dtype: torch.dtype, body_dtype: torch.dtype) -> None: """Set, add, and reset selected cells with either index width.""" - composer = WrenchComposer(create_mock_asset(num_envs=3, num_bodies=3, device="cpu")) + asset = SimpleNamespace(num_instances=3, num_bodies=3, device="cpu", data=_AssetData(*_poses(shape=(3, 3)))) + composer = WrenchComposer(asset) env_ids = torch.tensor([2, 0], dtype=env_dtype) body_ids = torch.tensor([1, 2], dtype=body_dtype) reset_env_ids = env_ids[:1] @@ -712,580 +276,210 @@ def test_index_dtype_combinations_preserve_selected_wrench_cells( np.testing.assert_array_equal(composer.local_torque_b.numpy(), expected_torques) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("method", ["set_forces_and_torques_index", "set_forces_and_torques_mask"]) -def test_set_forces_overwrites_previous_add(device: str, method: str): - """Test that set_forces_and_torques_index/mask clears previously accumulated values.""" - num_envs, num_bodies = 4, 2 - rng = np.random.default_rng(seed=40) +def _fill_all_input_buffers(composer: WrenchComposer) -> None: + """Give every cell a distinct contribution in each of the five input buffers.""" - mock_asset = create_mock_asset(num_envs, num_bodies, device) - composer = WrenchComposer(mock_asset) + def every_cell(value: tuple[float, float, float]) -> wp.array: + return _vectors(torch.tensor(value).repeat(2, 2, 1)) - # First accumulate some forces via add - forces_a_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) composer.add_forces_and_torques_index( - forces=wp.from_numpy(forces_a_np, dtype=wp.vec3f, device=device), + forces=every_cell((0.0, 0.0, 1.0)), positions=every_cell((0.0, 1.0, 0.0)), is_global=True ) + composer.add_forces_and_torques_index(forces=every_cell((1.0, 0.0, 0.0)), is_global=True) + composer.add_forces_and_torques_index(forces=every_cell((0.0, 1.0, 0.0)), torques=every_cell((0.0, 0.0, 1.0))) - # Now set new forces -- should replace, not accumulate - forces_b_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) - getattr(composer, method)( - forces=wp.from_numpy(forces_b_np, dtype=wp.vec3f, device=device), - ) - force, _, is_global = _get_wrench_without_pose_reads(composer) - assert is_global is False +# buffer -> per-cell value written by ``_fill_all_input_buffers`` +_FILLED_BUFFERS = { + "global_force_w": (0.0, 0.0, 1.0), + "global_torque_w": (1.0, 0.0, 0.0), # cross((0, 1, 0), (0, 0, 1)) + "global_force_at_com_w": (1.0, 0.0, 0.0), + "local_force_b": (0.0, 1.0, 0.0), + "local_torque_b": (0.0, 0.0, 1.0), +} - # Output should match forces_b only (forces_a should be gone) - assert np.allclose(force.numpy(), forces_b_np, atol=1e-4, rtol=1e-5), f"{method} did not clear previous add" +@pytest.mark.parametrize("method", ["set_forces_and_torques_index", "set_forces_and_torques_mask"]) +def test_set_clears_every_buffer_of_targeted_environments_only(method: str) -> None: + composer = _make_composer() + _fill_all_input_buffers(composer) -@pytest.mark.parametrize("device", test_devices()) -def test_set_forces_clears_targeted_envs_only(device: str): - """Test that set_forces_and_torques_index clears only the targeted environments.""" - num_envs, num_bodies = 4, 3 - rng = np.random.default_rng(seed=41) + _apply_to_cell(composer, method, (0, 1), False, forces=(5.0, 0.0, 0.0)) - mock_asset = create_mock_asset(num_envs, num_bodies, device) - composer = WrenchComposer(mock_asset) + for name, value in _FILLED_BUFFERS.items(): + expected = _grid({(1, 0): value, (1, 1): value}) + if name == "local_force_b": + expected[0, 1] = torch.tensor([5.0, 0.0, 0.0]) + _assert_vectors(getattr(composer, name), expected) - # Add global forces at positions (populates global_force_w and global_torque_w) - forces_global_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) - positions_np = rng.uniform(-5.0, 5.0, (num_envs, num_bodies, 3)).astype(np.float32) - composer.add_forces_and_torques_index( - forces=wp.from_numpy(forces_global_np, dtype=wp.vec3f, device=device), - positions=wp.from_numpy(positions_np, dtype=wp.vec3f, device=device), - is_global=True, - ) - # Also add local torques (populates local_torque_b) - torques_local_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) +@pytest.mark.parametrize("selector, source_is_global", [("list", False), ("slice", True), ("mask", False)]) +def test_partial_reset_zeros_only_selected_environments(selector: str, source_is_global: bool) -> None: + """Partial resets keep the remaining merged wrench and its required body-frame composition.""" + composer = _make_composer(supports_world_at_com=True) composer.add_forces_and_torques_index( - torques=wp.from_numpy(torques_local_np, dtype=wp.vec3f, device=device), - is_global=False, - ) - - # Now set local forces for envs [0, 2] -- should clear only envs 0, 2 - env_ids_np = np.array([0, 2], dtype=np.int32) - kept_env_ids = np.array([1, 3], dtype=np.int32) - forces_new_np = rng.uniform(-50.0, 50.0, (2, num_bodies, 3)).astype(np.float32) - composer.set_forces_and_torques_index( - forces=wp.from_numpy(forces_new_np, dtype=wp.vec3f, device=device), - env_ids=wp.from_numpy(env_ids_np, dtype=wp.int32, device=device), - is_global=False, + forces=_vectors(_grid({(0, 0): (1.0, 0.0, 0.0), (1, 0): (1.0, 0.0, 0.0)})), is_global=True ) - - zeros = np.zeros((num_bodies, 3), dtype=np.float32) - - # Targeted envs [0, 2]: all buffers cleared, then local_force_b written - for eid in env_ids_np: - assert np.allclose(composer.global_force_w.numpy()[eid], zeros, atol=1e-7), ( - f"global_force_w not cleared for targeted env {eid}" - ) - assert np.allclose(composer.global_torque_w.numpy()[eid], zeros, atol=1e-7), ( - f"global_torque_w not cleared for targeted env {eid}" - ) - assert np.allclose(composer.local_torque_b.numpy()[eid], zeros, atol=1e-7), ( - f"local_torque_b not cleared for targeted env {eid}" - ) - - # Non-targeted envs [1, 3]: should retain original values - for eid in kept_env_ids: - assert np.allclose(composer.global_force_w.numpy()[eid], forces_global_np[eid], atol=1e-4, rtol=1e-5), ( - f"global_force_w changed for non-targeted env {eid}" - ) - assert np.allclose(composer.local_torque_b.numpy()[eid], torques_local_np[eid], atol=1e-4, rtol=1e-5), ( - f"local_torque_b changed for non-targeted env {eid}" - ) - - # local_force_b should have new values at env_ids [0, 2], zeros at [1, 3] - expected_local_force = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) - expected_local_force[env_ids_np] = forces_new_np - assert np.allclose(composer.local_force_b.numpy(), expected_local_force, atol=1e-4, rtol=1e-5), ( - "local_force_b has wrong values after set" - ) - - -# ============================================================================ -# Partial Reset Tests -# ============================================================================ - - -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("source_is_global", [False, True]) -@pytest.mark.parametrize("use_slice", [False, True]) -def test_partial_reset_zeros_only_specified_envs(device: str, source_is_global: bool, use_slice: bool): - """Partial resets preserve the remaining merged wrench and its required composition.""" - num_envs, num_bodies = 8, 3 - rng = np.random.default_rng(seed=50) - asset = create_mock_asset(num_envs, num_bodies, device) - composer = WrenchComposer(asset, supports_world_at_com=True) - source = WrenchComposer(asset) - forces = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) - source_forces = rng.uniform(-100.0, 100.0, forces.shape).astype(np.float32) - positions = rng.uniform(-5.0, 5.0, forces.shape).astype(np.float32) - composer.add_forces_and_torques_index(forces=wp.array(forces, dtype=wp.vec3f, device=device), is_global=True) assert _get_wrench_without_pose_reads(composer)[2] is True + source = _make_composer() source.add_forces_and_torques_index( - forces=wp.array(source_forces, dtype=wp.vec3f, device=device), - positions=wp.array(positions, dtype=wp.vec3f, device=device), + forces=_vectors(_grid({(0, 0): (0.0, 2.0, 0.0), (1, 0): (0.0, 2.0, 0.0)})), + positions=_vectors(_grid({(0, 0): (1.0, 0.0, 0.0), (1, 0): (1.0, 0.0, 0.0)})), is_global=source_is_global, ) composer.add_raw_buffers_from(source) - reset_env_ids = [1, 3, 5] - if use_slice: - composer.reset(env_ids=slice(1, 6, 2)) - elif source_is_global: - mask = np.zeros(num_envs, dtype=bool) - mask[reset_env_ids] = True - composer.reset(env_mask=wp.array(mask, dtype=wp.bool, device=device)) - else: - composer.reset(env_ids=reset_env_ids) + reset_selection = { + "list": {"env_ids": [1]}, + "slice": {"env_ids": slice(1, None)}, + "mask": {"env_mask": _mask([False, True])}, + } + composer.reset(**reset_selection[selector]) + assert composer.active assert composer._dirty - force, torque, is_global = composer.get_forces_and_torques() assert is_global is False - expected_force = forces + source_forces - expected_torque = np.cross(positions, source_forces) - expected_force[reset_env_ids] = 0.0 - expected_torque[reset_env_ids] = 0.0 - np.testing.assert_allclose(force.numpy(), expected_force, atol=1e-4, rtol=1e-5) - np.testing.assert_allclose(torque.numpy(), expected_torque, atol=1e-4, rtol=1e-5) + _assert_vectors(force, _grid({(0, 0): (1.0, 2.0, 0.0)})) + # cross((1, 0, 0), (0, 2, 0)) = (0, 0, 2) about a CoM at the origin, in either frame. + _assert_vectors(torque, _grid({(0, 0): (0.0, 0.0, 2.0)})) -@pytest.mark.parametrize("device", test_devices()) @pytest.mark.parametrize("env_ids", [None, slice(None)], ids=["none", "full_slice"]) -def test_full_reset_clears_active_flag(device: str, env_ids: slice | None): - """Test that either full-reset selector clears the _active flag.""" - num_envs, num_bodies = 4, 2 - - mock_asset = create_mock_asset(num_envs, num_bodies, device) - composer = WrenchComposer(mock_asset, supports_world_at_com=True) - - forces_np = np.ones((num_envs, num_bodies, 3), dtype=np.float32) - # Not parallel to the force, so the global torque buffer is non-zero before the reset. - positions_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) - positions_np[..., 0] = 1.0 - composer.add_forces_and_torques_index( - forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), - torques=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), - ) - composer.add_forces_and_torques_index( - forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), - torques=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), - positions=wp.from_numpy(positions_np, dtype=wp.vec3f, device=device), - is_global=True, - ) - composer.add_forces_and_torques_index( - forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), is_global=True - ) +def test_full_reset_clears_every_buffer_and_frame_flag(env_ids: slice | None) -> None: + composer = _make_composer(supports_world_at_com=True) + _fill_all_input_buffers(composer) composer.get_forces_and_torques() assert composer.active composer.reset(env_ids=env_ids) + assert not composer.active assert not composer._dirty - # Check all 7 buffers are zero (5 input + 2 output) - zeros = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) - np.testing.assert_array_equal(composer.global_force_w.numpy(), zeros) - np.testing.assert_array_equal(composer.global_torque_w.numpy(), zeros) - np.testing.assert_array_equal(composer.global_force_at_com_w.numpy(), zeros) - np.testing.assert_array_equal(composer.local_force_b.numpy(), zeros) - np.testing.assert_array_equal(composer.local_torque_b.numpy(), zeros) - np.testing.assert_array_equal(composer.out_force_b.warp.numpy(), zeros) - np.testing.assert_array_equal(composer.out_torque_b.warp.numpy(), zeros) - + for name in (*_FILLED_BUFFERS, "out_force_b", "out_torque_b"): + _assert_vectors(getattr(composer, name), torch.zeros((2, 2, 3))) force, torque, is_global = _get_wrench_without_pose_reads(composer) assert is_global is False - np.testing.assert_array_equal(force.numpy(), np.zeros_like(forces_np)) - np.testing.assert_array_equal(torque.numpy(), np.zeros_like(forces_np)) - composer.add_forces_and_torques_index( - forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), is_global=True - ) + _assert_vectors(force, torch.zeros((2, 2, 3))) + _assert_vectors(torque, torch.zeros((2, 2, 3))) + # The frame flags were cleared too: a world-at-CoM wrench takes the world fast path again. + _apply_to_cell(composer, "add_forces_and_torques_index", (0, 0), True, forces=(1.0, 0.0, 0.0)) assert _get_wrench_without_pose_reads(composer)[2] is True -# ============================================================================ -# Deprecated API Backward-Compatibility Tests -# ============================================================================ - - -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize( - "deprecated_name, name, input_name", - [("composed_force", "out_force_b", "forces"), ("composed_torque", "out_torque_b", "torques")], -) -def test_composed_wrench_emits_deprecation_warning(device: str, deprecated_name: str, name: str, input_name: str): - """Test that accessing composed_force/composed_torque emits a DeprecationWarning and aliases the output.""" - num_envs, num_bodies = 2, 1 - - mock_asset = create_mock_asset(num_envs, num_bodies, device) - composer = WrenchComposer(mock_asset) - - values_np = np.array([[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]]], dtype=np.float32) - composer.add_forces_and_torques_index(**{input_name: wp.from_numpy(values_np, dtype=wp.vec3f, device=device)}) - - with pytest.warns(DeprecationWarning, match=f"{deprecated_name}.*is deprecated"): - result = getattr(composer, deprecated_name) - - # Should return the same data as the new output property - assert np.allclose(result.warp.numpy(), getattr(composer, name).warp.numpy(), atol=1e-7) - - -@pytest.mark.parametrize("device", test_devices()) -def test_deprecated_add_forces_and_torques_emits_warning(device: str): - """Test that the deprecated add_forces_and_torques wrapper emits a warning and works.""" - num_envs, num_bodies = 4, 2 - rng = np.random.default_rng(seed=52) - - mock_asset = create_mock_asset(num_envs, num_bodies, device) - composer = WrenchComposer(mock_asset) - - forces_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) - - with pytest.warns(DeprecationWarning, match="add_forces_and_torques.*is deprecated"): - composer.add_forces_and_torques( - forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), - ) - - composer.compose_to_body_frame() - assert np.allclose(composer.out_force_b.warp.numpy(), forces_np, atol=1e-4, rtol=1e-5) - - -# ============================================================================ -# set_forces_and_torques_mask Tests -# ============================================================================ - - -@pytest.mark.parametrize("device", test_devices()) -def test_set_forces_mask_clears_targeted_envs_only(device: str): - """Test that set_forces_and_torques_mask clears only the masked environments.""" - num_envs, num_bodies = 4, 3 - rng = np.random.default_rng(seed=61) - - mock_asset = create_mock_asset(num_envs, num_bodies, device) - composer = WrenchComposer(mock_asset) - - # Populate global buffers for all envs - forces_global_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) - positions_np = rng.uniform(-5.0, 5.0, (num_envs, num_bodies, 3)).astype(np.float32) - composer.add_forces_and_torques_index( - forces=wp.from_numpy(forces_global_np, dtype=wp.vec3f, device=device), - positions=wp.from_numpy(positions_np, dtype=wp.vec3f, device=device), - is_global=True, - ) - - # Also add local torques for all envs - torques_local_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) - composer.add_forces_and_torques_index( - torques=wp.from_numpy(torques_local_np, dtype=wp.vec3f, device=device), - is_global=False, - ) - - # Set local forces via mask for envs [0, 2] -- should clear only masked envs - env_mask_np = np.array([True, False, True, False], dtype=np.bool_) - body_mask_np = np.array([True, True, False], dtype=np.bool_) - forces_new_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) - composer.set_forces_and_torques_mask( - forces=wp.from_numpy(forces_new_np, dtype=wp.vec3f, device=device), - env_mask=wp.from_numpy(env_mask_np, dtype=wp.bool, device=device), - body_mask=wp.from_numpy(body_mask_np, dtype=wp.bool, device=device), - is_global=False, - ) - - zeros = np.zeros((num_bodies, 3), dtype=np.float32) - - # Masked envs [0, 2]: all buffers cleared by reset, then local_force_b written where body_mask is True - for eid in [0, 2]: - assert np.allclose(composer.global_force_w.numpy()[eid], zeros, atol=1e-7), ( - f"global_force_w not cleared for masked env {eid}" - ) - assert np.allclose(composer.global_torque_w.numpy()[eid], zeros, atol=1e-7), ( - f"global_torque_w not cleared for masked env {eid}" - ) - assert np.allclose(composer.local_torque_b.numpy()[eid], zeros, atol=1e-7), ( - f"local_torque_b not cleared for masked env {eid}" - ) - - # Non-masked envs [1, 3]: should retain original values - for eid in [1, 3]: - assert np.allclose(composer.global_force_w.numpy()[eid], forces_global_np[eid], atol=1e-4, rtol=1e-5), ( - f"global_force_w changed for non-masked env {eid}" - ) - assert np.allclose(composer.local_torque_b.numpy()[eid], torques_local_np[eid], atol=1e-4, rtol=1e-5), ( - f"local_torque_b changed for non-masked env {eid}" - ) - - # local_force_b should have new values where both masks are True, zeros for masked envs otherwise - expected_local_force = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) - for e in range(num_envs): - for b in range(num_bodies): - if env_mask_np[e] and body_mask_np[b]: - expected_local_force[e, b] = forces_new_np[e, b] - assert np.allclose(composer.local_force_b.numpy(), expected_local_force, atol=1e-4, rtol=1e-5), ( - "local_force_b has wrong values after mask set" - ) - - -# ============================================================================ -# Lazy Composition (_ensure_composed) Tests -# ============================================================================ - - -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("input_name, output_name", [("forces", "out_force_b"), ("torques", "out_torque_b")]) -def test_output_property_triggers_lazy_composition(device: str, input_name: str, output_name: str): - """Test that reading out_force_b/out_torque_b without explicit compose_to_body_frame returns correct results.""" - num_envs, num_bodies = 4, 2 - rng = np.random.default_rng(seed=70) - - link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) - link_quat_torch = torch.from_numpy(link_quat_np) - - mock_asset = create_mock_asset(num_envs, num_bodies, device, link_quat=link_quat_torch) - composer = WrenchComposer(mock_asset) - - values_global_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) - composer.add_forces_and_torques_index( - **{input_name: wp.from_numpy(values_global_np, dtype=wp.vec3f, device=device)}, - is_global=True, - ) - - # Do NOT call compose_to_body_frame -- rely on lazy composition - expected_values_local = quat_rotate_inv_np(link_quat_np, values_global_np) - composed_np = getattr(composer, output_name).warp.numpy() - - assert np.allclose(composed_np, expected_values_local, atol=1e-4, rtol=1e-5), ( - f"Lazy composition via {output_name} failed" - ) - - -@pytest.mark.parametrize("device", test_devices()) -def test_lazy_composition_reflects_later_adds(device: str): - """Test that an add after a lazy composition is reflected by the next output read.""" - num_envs, num_bodies = 2, 1 - - mock_asset = create_mock_asset(num_envs, num_bodies, device) - composer = WrenchComposer(mock_asset) - - forces_np = np.ones((num_envs, num_bodies, 3), dtype=np.float32) - composer.add_forces_and_torques_index( - forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), - ) - # Reading out_force_b composes lazily - assert np.allclose(composer.out_force_b.warp.numpy(), forces_np, atol=1e-4, rtol=1e-5) - - # Another add must be reflected by the next read of either output - composer.add_forces_and_torques_index( - forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), - torques=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), - ) - assert np.allclose(composer.out_torque_b.warp.numpy(), forces_np, atol=1e-4, rtol=1e-5) - assert np.allclose(composer.out_force_b.warp.numpy(), 2.0 * forces_np, atol=1e-4, rtol=1e-5) - - -@pytest.mark.parametrize("device", test_devices()) -def test_compose_is_idempotent(device: str): - """Calling compose_to_body_frame twice without intervening writes produces the same result.""" - rng = np.random.default_rng(seed=456) - num_envs, num_bodies = 4, 3 - - # Non-trivial link pose so the rotation path is exercised - link_pos_np = rng.uniform(-2, 2, (num_envs, num_bodies, 3)).astype(np.float32) - link_quat_np = rng.standard_normal((num_envs, num_bodies, 4)).astype(np.float32) - link_quat_np /= np.linalg.norm(link_quat_np, axis=-1, keepdims=True) - - mock_asset = create_mock_asset( - num_envs, - num_bodies, - device, - link_pos=torch.from_numpy(link_pos_np), - link_quat=torch.from_numpy(link_quat_np), - ) - composer = WrenchComposer(mock_asset) - - # Add global forces with positions (exercises cross-product torque path) - forces_np = rng.uniform(-5, 5, (num_envs, num_bodies, 3)).astype(np.float32) - positions_np = rng.uniform(-1, 1, (num_envs, num_bodies, 3)).astype(np.float32) - torques_np = rng.uniform(-3, 3, (num_envs, num_bodies, 3)).astype(np.float32) - - composer.add_forces_and_torques_index( - forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), - torques=wp.from_numpy(torques_np, dtype=wp.vec3f, device=device), - positions=wp.from_numpy(positions_np, dtype=wp.vec3f, device=device), - is_global=True, +def test_raw_buffer_merge_accumulates_all_five_buffers_and_ignores_inactive_sources() -> None: + # Alone, the world-at-CoM destination would take the world fast path. + destination = _make_composer(supports_world_at_com=True) + source = _make_composer() + inactive_source = _make_composer() + _apply_to_cell(destination, "add_forces_and_torques_index", (0, 0), True, forces=(2.0, 0.0, 0.0)) + _apply_to_cell( + source, "add_forces_and_torques_index", (0, 0), False, forces=(1.0, 0.0, 0.0), torques=(0.0, 2.0, 0.0) ) + _apply_to_cell(source, "add_forces_and_torques_index", (0, 0), True, forces=(0.0, 0.0, 3.0)) + _apply_to_cell( + source, + "add_forces_and_torques_index", + (0, 0), + True, + forces=(0.0, 4.0, 0.0), + torques=(0.0, 0.0, 5.0), + positions=(0.0, 0.0, 0.0), + ) + # A raw write that never activated its composer must not be merged. + wp.copy(inactive_source.local_force_b, _vectors(_grid({(1, 1): (8.0, 0.0, 0.0)}))) + assert not inactive_source.active + + destination.add_raw_buffers_from(source) + destination.add_raw_buffers_from(inactive_source) + + _assert_vectors(destination.local_force_b, _grid({(0, 0): (1.0, 0.0, 0.0)})) + _assert_vectors(destination.local_torque_b, _grid({(0, 0): (0.0, 2.0, 0.0)})) + _assert_vectors(destination.global_force_w, _grid({(0, 0): (0.0, 4.0, 0.0)})) + _assert_vectors(destination.global_force_at_com_w, _grid({(0, 0): (2.0, 0.0, 3.0)})) + _assert_vectors(destination.global_torque_w, _grid({(0, 0): (0.0, 0.0, 5.0)})) + # The merged frame flags (local input and world positions) require body-frame composition. + force, torque, is_global = destination.get_forces_and_torques() + assert is_global is False + _assert_vectors(force, _grid({(0, 0): (3.0, 4.0, 3.0)})) + _assert_vectors(torque, _grid({(0, 0): (0.0, 2.0, 5.0)})) - # First compose - composer.compose_to_body_frame() - force_first = composer.out_force_b.warp.numpy().copy() - torque_first = composer.out_torque_b.warp.numpy().copy() - - # Second compose (no writes in between) - composer.compose_to_body_frame() - force_second = composer.out_force_b.warp.numpy() - torque_second = composer.out_torque_b.warp.numpy() - - np.testing.assert_array_equal(force_first, force_second) - np.testing.assert_array_equal(torque_first, torque_second) - - -# ============================================================================ -# CoM Offset from Link Origin Tests -# ============================================================================ - - -@pytest.mark.parametrize("device", test_devices()) -def test_global_force_with_com_offset(device: str): - """Test that torque correction uses CoM position, not link position, when they differ.""" - num_envs, num_bodies = 2, 1 - - # Link at origin, CoM offset by [1, 0, 0] - link_pos_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) - link_quat_np = np.zeros((num_envs, num_bodies, 4), dtype=np.float32) - link_quat_np[..., 3] = 1.0 # identity quaternion (xyzw) + # Merging an inactive source leaves an inactive destination inactive. + empty_destination = _make_composer() + empty_destination.add_raw_buffers_from(inactive_source) + assert not empty_destination.active - com_pos_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) - com_pos_np[..., 0] = 1.0 # CoM at [1, 0, 0] - mock_asset = create_mock_asset( - num_envs, - num_bodies, - device, - link_pos=torch.from_numpy(link_pos_np), - link_quat=torch.from_numpy(link_quat_np), - ) - # Set CoM pose separately (pos=[1,0,0], quat=identity) - com_pose = torch.cat( - ( - torch.from_numpy(com_pos_np), - torch.tensor([0.0, 0.0, 0.0, 1.0]).view(1, 1, 4).expand(num_envs, num_bodies, 4), - ), - dim=-1, - ) - mock_asset.data.set_body_com_pose_w(com_pose) +# --------------------------------------------------------------------------- +# Lazy composition +# --------------------------------------------------------------------------- - composer = WrenchComposer(mock_asset) - # Apply global force [0, 0, 10] at position [0, 0, 0] (world origin) - forces_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) - forces_np[..., 2] = 10.0 - positions_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) - - composer.add_forces_and_torques_index( - forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), - positions=wp.from_numpy(positions_np, dtype=wp.vec3f, device=device), - is_global=True, - ) +def test_output_properties_compose_lazily_and_reflect_later_adds() -> None: + """Reading an output composes pending input into the body frame; later adds are reflected by the next read.""" + composer = _make_composer(quat={(0, 1): _QUARTER_TURN_Z}) + _apply_to_cell(composer, "add_forces_and_torques_index", (0, 1), True, forces=(2.0, 0.0, 0.0)) + assert composer._dirty + _assert_vectors(composer.out_force_b, _grid({(0, 1): (0.0, -2.0, 0.0)})) + assert not composer._dirty + _apply_to_cell(composer, "add_forces_and_torques_index", (0, 1), True, torques=(2.0, 0.0, 0.0)) + _assert_vectors(composer.out_torque_b, _grid({(0, 1): (0.0, -2.0, 0.0)})) + _assert_vectors(composer.out_force_b, _grid({(0, 1): (0.0, -2.0, 0.0)})) + # Composing again without new input is idempotent. composer.compose_to_body_frame() - - # With identity quaternion: - # torque_w = cross(P, F) - cross(com, F) = cross([0,0,0], [0,0,10]) - cross([1,0,0], [0,0,10]) - # = [0,0,0] - [0*10-0*0, 0*0-1*10, 1*0-0*0] = [0,0,0] - [0, -10, 0] = [0, 10, 0] - # In body frame (identity rotation): [0, 10, 0] - expected_torque = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) - expected_torque[..., 1] = 10.0 - - actual_torque = composer.out_torque_b.warp.numpy() - assert np.allclose(actual_torque, expected_torque, atol=1e-4, rtol=1e-5), ( - f"CoM offset torque correction failed.\nExpected:\n{expected_torque}\nGot:\n{actual_torque}" - ) - - # Force should be unchanged (identity rotation) - assert np.allclose(composer.out_force_b.warp.numpy(), forces_np, atol=1e-4, rtol=1e-5) + _assert_vectors(composer.out_force_b, _grid({(0, 1): (0.0, -2.0, 0.0)})) + _assert_vectors(composer.out_torque_b, _grid({(0, 1): (0.0, -2.0, 0.0)})) -@pytest.mark.parametrize("device", test_devices()) -def test_com_offset_with_rotation(device: str): - """Test torque correction with both CoM offset and non-identity rotation.""" - num_envs, num_bodies = 1, 1 - rng = np.random.default_rng(seed=73) +# --------------------------------------------------------------------------- +# Deprecated API and input validation +# --------------------------------------------------------------------------- - # Random rotation - link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) - link_pos_np = rng.uniform(-5.0, 5.0, (num_envs, num_bodies, 3)).astype(np.float32) - # CoM offset from link - com_offset_np = rng.uniform(0.5, 2.0, (num_envs, num_bodies, 3)).astype(np.float32) - com_pos_np = link_pos_np + com_offset_np # simple world-frame offset for test clarity - - mock_asset = create_mock_asset( - num_envs, - num_bodies, - device, - link_pos=torch.from_numpy(link_pos_np), - link_quat=torch.from_numpy(link_quat_np), - ) - com_pose = torch.cat( - ( - torch.from_numpy(com_pos_np), - torch.tensor([0.0, 0.0, 0.0, 1.0]).view(1, 1, 4).expand(num_envs, num_bodies, 4), - ), - dim=-1, +def test_deprecated_composed_wrench_properties_alias_the_outputs() -> None: + composer = _make_composer() + _apply_to_cell( + composer, "add_forces_and_torques_index", (1, 0), False, forces=(1.0, 2.0, 3.0), torques=(4.0, 5.0, 6.0) ) - mock_asset.data.set_body_com_pose_w(com_pose) - - composer = WrenchComposer(mock_asset) - - # Apply global force at a random world position - forces_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) - positions_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) - - composer.add_forces_and_torques_index( - forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), - positions=wp.from_numpy(positions_np, dtype=wp.vec3f, device=device), - is_global=True, - ) - - composer.compose_to_body_frame() - # Expected: torque_w = cross(P, F) - cross(com, F) = cross(P - com, F) - lever_arm = positions_np - com_pos_np - torque_w = np.cross(lever_arm, forces_np) - expected_torque_b = quat_rotate_inv_np(link_quat_np, torque_w) - expected_force_b = quat_rotate_inv_np(link_quat_np, forces_np) + for deprecated_name, name, expected in ( + ("composed_force", "out_force_b", (1.0, 2.0, 3.0)), + ("composed_torque", "out_torque_b", (4.0, 5.0, 6.0)), + ): + with pytest.warns(DeprecationWarning, match=f"{deprecated_name}.*is deprecated"): + result = getattr(composer, deprecated_name) + assert result is getattr(composer, name) + _assert_vectors(result, _grid({(1, 0): expected})) - assert np.allclose(composer.out_force_b.warp.numpy(), expected_force_b, atol=1e-3, rtol=1e-4), ( - "Force mismatch with CoM offset + rotation" - ) - assert np.allclose(composer.out_torque_b.warp.numpy(), expected_torque_b, atol=1e-3, rtol=1e-4), ( - f"Torque mismatch with CoM offset + rotation.\n" - f"Expected:\n{expected_torque_b}\nGot:\n{composer.out_torque_b.warp.numpy()}" - ) +def test_deprecated_writers_warn_and_keep_add_and_set_behavior() -> None: + composer = _make_composer() -# ============================================================================ -# Deprecated set_forces_and_torques Tests -# ============================================================================ + with pytest.warns(DeprecationWarning, match="add_forces_and_torques.*is deprecated"): + composer.add_forces_and_torques( + forces=_vectors(torch.tensor([[[3.0, 0.0, 0.0]]])), + torques=_vectors(torch.tensor([[[0.0, 0.0, 2.0]]])), + body_ids=[1], + env_ids=[1], + ) + _assert_vectors(composer.out_force_b, _grid({(1, 1): (3.0, 0.0, 0.0)})) + _assert_vectors(composer.out_torque_b, _grid({(1, 1): (0.0, 0.0, 2.0)})) + with pytest.warns(DeprecationWarning, match="set_forces_and_torques.*is deprecated"): + composer.set_forces_and_torques(forces=_vectors(torch.tensor([[[4.0, 0.0, 0.0]]])), body_ids=[1], env_ids=[1]) -@pytest.mark.parametrize("device", test_devices()) -def test_deprecated_set_forces_and_torques_clears_previous(device: str): - """Test that deprecated set_forces_and_torques actually replaces previous values.""" - num_envs, num_bodies = 4, 2 - rng = np.random.default_rng(seed=81) + # The deprecated set replaces the earlier wrench of the targeted environment, torque included. + _assert_vectors(composer.out_force_b, _grid({(1, 1): (4.0, 0.0, 0.0)})) + _assert_vectors(composer.out_torque_b, torch.zeros((2, 2, 3))) - mock_asset = create_mock_asset(num_envs, num_bodies, device) - composer = WrenchComposer(mock_asset) - # First add some forces - forces_a_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) - composer.add_forces_and_torques_index( - forces=wp.from_numpy(forces_a_np, dtype=wp.vec3f, device=device), - ) - - # Then set via deprecated method -- should replace - forces_b_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) - with pytest.warns(DeprecationWarning, match="set_forces_and_torques.*is deprecated"): - composer.set_forces_and_torques( - forces=wp.from_numpy(forces_b_np, dtype=wp.vec3f, device=device), - ) +def test_invalid_selection_and_empty_wrench_input_are_reported() -> None: + composer = _make_composer() - composer.compose_to_body_frame() - assert np.allclose(composer.out_force_b.warp.numpy(), forces_b_np, atol=1e-4, rtol=1e-5), ( - "Deprecated set_forces_and_torques did not replace previous values" - ) + with pytest.raises(TypeError, match="env_ids must be"): + composer.add_forces_and_torques_index(forces=_vectors(torch.ones((1, 2, 3))), env_ids=(0,)) + for method in ("add_forces_and_torques_index", "set_forces_and_torques_mask"): + with pytest.warns(UserWarning, match="No forces or torques"): + getattr(composer, method)() + assert not composer.active diff --git a/source/isaaclab/test/utils/test_wrench_composer_integration.py b/source/isaaclab/test/utils/test_wrench_composer_integration.py index 39ac9b57a4fc..4095291fbbee 100644 --- a/source/isaaclab/test/utils/test_wrench_composer_integration.py +++ b/source/isaaclab/test/utils/test_wrench_composer_integration.py @@ -3,646 +3,117 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Integration tests for wrench composer with rigid objects. +"""Real PhysX delivery of :class:`~isaaclab.utils.wrench_composer.WrenchComposer` wrenches. -These tests validate that global forces/torques remain invariant under body rotation +The composer arithmetic is covered by literal units in ``test_wrench_composer.py``. This module checks the two ways a +rigid object hands the composed wrench to PhysX: body-frame vectors for a world force at a world position, and +world-frame vectors applied at the center of mass. """ from isaaclab.test.utils import launch_test_simulation launch_test_simulation() +import math + import pytest import torch import warp as wp +from isaaclab_physx.sim.schemas import PhysxRigidBodyCfg import isaaclab.sim as sim_utils from isaaclab.assets import RigidObject, RigidObjectCfg from isaaclab.sim import build_simulation_context -from isaaclab.test.utils import DeviceScope, test_devices -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.test.utils import test_devices pytestmark = pytest.mark.integration - -def generate_cubes_scene( - num_cubes: int = 1, - height: float = 1.0, - device: str = "cuda:0", -) -> tuple[RigidObject, torch.Tensor]: - """Generate a scene with the provided number of cubes.""" - origins = torch.tensor([(i * 1.0, 0, height) for i in range(num_cubes)]).to(device) - for i, origin in enumerate(origins): - sim_utils.create_prim(f"/World/Table_{i}", "Xform", translation=origin) - - spawn_cfg = sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", - rigid_props=sim_utils.UsdPhysicsRigidBodyCfg(), +_ROTATION_45_Z = (0.0, 0.0, math.sin(math.pi / 8), math.cos(math.pi / 8)) +# A 120-degree turn about (1, 1, 1) maps the body axes x, y, z onto the world axes y, z, x. +_ROTATION_120_XYZ = (0.5, 0.5, 0.5, 0.5) +_NUM_STEPS = 4 + + +def _make_cube(name: str, position: tuple[float, float, float], rotation: tuple[float, ...]) -> RigidObject: + """Spawn a free, undamped unit-mass cube with a rotated initial pose.""" + spawn = sim_utils.CuboidCfg( + size=(0.2, 0.2, 0.2), + rigid_props=[ + sim_utils.UsdPhysicsRigidBodyCfg(rigid_body_enabled=True), + PhysxRigidBodyCfg(disable_gravity=True, linear_damping=0.0, angular_damping=0.0), + ], + mass_props=sim_utils.MassCfg(mass=1.0), + collision_props=sim_utils.UsdPhysicsCollisionCfg(collision_enabled=True), ) - - cube_object_cfg = RigidObjectCfg( - prim_path="/World/Table_[^/]*/Object", - spawn=spawn_cfg, - init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, height)), + cfg = RigidObjectCfg( + prim_path=f"/World/{name}", spawn=spawn, init_state=RigidObjectCfg.InitialStateCfg(pos=position, rot=rotation) ) - cube_object = RigidObject(cfg=cube_object_cfg) - return cube_object, origins - - -N_STEPS = 100 -FORCE_MAGNITUDE = 10.0 -TORQUE_MAGNITUDE = 1.0 - - -@pytest.mark.parametrize("device", test_devices()) -def test_global_force_invariant_under_rotation(device): - """Test that a permanent global force produces the same acceleration before and after body rotation. - - A global +X force is applied. After 100 steps the body is rotated 180deg about Z. - The acceleration (delta_v per phase) should be the same in both phases because the - force is in the global frame and should not rotate with the body. - """ - with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_object, _ = generate_cubes_scene(num_cubes=1, device=device) - - sim.reset() - - body_ids, _ = cube_object.find_bodies(".*") - mass = float(wp.to_torch(cube_object.root_view.get_masses())[0]) - com = cube_object.data.body_com_pos_w.torch.clone() - - # Apply permanent global force along +X at CoM - forces = torch.zeros(1, len(body_ids), 3, device=device) - forces[..., 0] = FORCE_MAGNITUDE - torques = torch.zeros(1, len(body_ids), 3, device=device) - - cube_object.permanent_wrench_composer.set_forces_and_torques_index( - forces=forces, - torques=torques, - positions=com, - body_ids=body_ids, - is_global=True, - ) - - # Phase 1: run N_STEPS - for _ in range(N_STEPS): - cube_object.write_data_to_sim() - sim.step() - cube_object.update(sim.cfg.dt) - - vel_after_phase1 = cube_object.data.root_lin_vel_w.torch[0].clone() - - # Rotate body 180deg about Z (quat wxyz = [0, 0, 0, 1]) while keeping velocity - root_pose = cube_object.data.root_pose_w.torch[0].clone().unsqueeze(0) - root_pose[0, 3:7] = torch.tensor([0.0, 0.0, 1.0, 0.0], device=device) # 180deg about Z (xyzw) - cube_object.write_root_pose_to_sim_index(root_pose=root_pose) - - # Phase 2: run N_STEPS more - for _ in range(N_STEPS): - cube_object.write_data_to_sim() - sim.step() - cube_object.update(sim.cfg.dt) - - vel_after_phase2 = cube_object.data.root_lin_vel_w.torch[0].clone() - - # Acceleration should be same in both phases: delta_v_phase2 ≈ delta_v_phase1 - delta_v_phase1 = vel_after_phase1[0].item() # vx after phase 1 - delta_v_phase2 = vel_after_phase2[0].item() - vel_after_phase1[0].item() # vx gained in phase 2 - - expected_dv = FORCE_MAGNITUDE / mass * sim.cfg.dt * N_STEPS - - torch.testing.assert_close( - torch.tensor(delta_v_phase1), - torch.tensor(expected_dv), - rtol=0.001, - atol=0.0001, - ) - torch.testing.assert_close( - torch.tensor(delta_v_phase2), - torch.tensor(expected_dv), - rtol=0.001, - atol=0.0001, - ) - - # Y and Z velocity should remain ~0 - assert abs(vel_after_phase2[1].item()) < 0.5, f"Unexpected Y velocity: {vel_after_phase2[1].item()}" - assert abs(vel_after_phase2[2].item()) < 0.5, f"Unexpected Z velocity: {vel_after_phase2[2].item()}" - - -@pytest.mark.parametrize("device", test_devices()) -def test_local_force_follows_rotation(device): - """Test that a permanent local force rotates with the body. - - A local +X force is applied. After 100 steps the body is rotated 180deg about Z. - Since local +X is now world -X, the force should decelerate the body back towards zero velocity. - """ - with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_object, _ = generate_cubes_scene(num_cubes=1, device=device) - - sim.reset() - - body_ids, _ = cube_object.find_bodies(".*") - - # Apply permanent local force along body +X - forces = torch.zeros(1, len(body_ids), 3, device=device) - forces[..., 0] = FORCE_MAGNITUDE - torques = torch.zeros(1, len(body_ids), 3, device=device) - - cube_object.permanent_wrench_composer.set_forces_and_torques_index( - forces=forces, - torques=torques, - body_ids=body_ids, - is_global=False, - ) - - # Phase 1: run N_STEPS — object accelerates along world +X - for _ in range(N_STEPS): - cube_object.write_data_to_sim() - sim.step() - cube_object.update(sim.cfg.dt) - - vel_after_phase1 = cube_object.data.root_lin_vel_w.torch[0].clone() - assert vel_after_phase1[0].item() > 1.0, "Object should be moving in +X" - - # Rotate body 180deg about Z while keeping velocity - root_pose = cube_object.data.root_pose_w.torch[0].clone().unsqueeze(0) - root_pose[0, 3:7] = torch.tensor([0.0, 0.0, 1.0, 0.0], device=device) # 180deg about Z (xyzw) - cube_object.write_root_pose_to_sim_index(root_pose=root_pose) - - # Phase 2: run N_STEPS — local +X is now world -X, so force decelerates - for _ in range(N_STEPS): - cube_object.write_data_to_sim() - sim.step() - cube_object.update(sim.cfg.dt) - - vel_after_phase2 = cube_object.data.root_lin_vel_w.torch[0].clone() - - # Velocity should be approximately zero: decelerated by the same amount as it accelerated - torch.testing.assert_close( - vel_after_phase2[0], - torch.tensor(0.0, device=device), - atol=0.0001, - rtol=0.001, - ) - - -@pytest.mark.parametrize("device", test_devices()) -def test_global_torque_invariant_under_rotation(device): - """Test that a permanent global torque produces the same angular acceleration before and after rotation. - - A global +Z torque is applied. After 100 steps the body is rotated 90deg about X. - The angular acceleration (delta_omega per phase) about Z should be the same in both phases - because the torque is in the global frame. - """ - with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_object, _ = generate_cubes_scene(num_cubes=1, device=device) - - sim.reset() - - body_ids, _ = cube_object.find_bodies(".*") - - # Apply permanent global torque about +Z - forces = torch.zeros(1, len(body_ids), 3, device=device) - torques = torch.zeros(1, len(body_ids), 3, device=device) - torques[..., 2] = TORQUE_MAGNITUDE - - cube_object.permanent_wrench_composer.set_forces_and_torques_index( - forces=forces, - torques=torques, - body_ids=body_ids, - is_global=True, - ) - - # Phase 1: run N_STEPS - for _ in range(N_STEPS): - cube_object.write_data_to_sim() - sim.step() - cube_object.update(sim.cfg.dt) - - omega_z_after_phase1 = cube_object.data.root_ang_vel_w.torch[0, 2].clone().item() - - # Rotate body 90deg about X and zero out velocities so phase 2 starts from rest - # (avoids gyroscopic cross-coupling at high omega) - root_pose = cube_object.data.root_pose_w.torch[0].clone().unsqueeze(0) - root_pose[0, 3:7] = torch.tensor([0.7071, 0.0, 0.0, 0.7071], device=device) # 90deg about X (xyzw) - cube_object.write_root_pose_to_sim_index(root_pose=root_pose) - - root_vel = cube_object.data.root_vel_w.torch.clone() - root_vel[0, :] = 0.0 - cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) - - # Phase 2: run N_STEPS from rest with different body orientation - for _ in range(N_STEPS): - cube_object.write_data_to_sim() - sim.step() - cube_object.update(sim.cfg.dt) - - omega_z_after_phase2 = cube_object.data.root_ang_vel_w.torch[0, 2].clone().item() - - # Both phases start from rest — angular acceleration about Z should be the same - torch.testing.assert_close( - torch.tensor(omega_z_after_phase1), - torch.tensor(omega_z_after_phase2), - rtol=0.001, - atol=0.0001, - ) + return RigidObject(cfg=cfg) @pytest.mark.parametrize("device", test_devices()) -def test_global_force_torque_after_translation(device): - """Test that global force torque updates dynamically when the body translates. - - Phase 1: Cube at (1,0,0). Global force F=(0,10,0) applied at explicit position (1,0,0). - stored_torque = cross((1,0,0), (0,10,0)) = (0,0,10) - correction = -cross((1,0,0), (0,10,0)) = (0,0,-10) - net torque = 0 → no rotation, only linear acceleration in +Y. - - Phase 2: Teleport cube to origin (0,0,0), zero velocity, don't re-apply force. - stored_torque = (0,0,10) (unchanged in buffer) - correction = -cross((0,0,0), (0,10,0)) = (0,0,0) - net torque = (0,0,10) → rotation about +Z. - """ - with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_object, _ = generate_cubes_scene(num_cubes=1, height=1.0, device=device) - - sim.reset() - - body_ids, _ = cube_object.find_bodies(".*") - - # Phase 1 setup: Move cube to (1, 0, 1) and apply force at (1, 0, 1) - root_pose = cube_object.data.root_pose_w.torch.clone() - root_pose[0, 0] = 1.0 # x = 1 - root_pose[0, 1] = 0.0 # y = 0 - root_pose[0, 2] = 1.0 # z = 1 - root_pose[0, 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) # identity quat (xyzw) - cube_object.write_root_pose_to_sim_index(root_pose=root_pose) - - root_vel = cube_object.data.root_vel_w.torch.clone() - root_vel[0, :] = 0.0 # zero velocity - cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) - - # Step once to let the state settle - sim.step() - cube_object.update(sim.cfg.dt) - - # Get current CoM position for the force application point - com_pos = cube_object.data.body_com_pos_w.torch[:, body_ids, :3].clone() - - forces = torch.zeros(1, len(body_ids), 3, device=device) - forces[..., 1] = FORCE_MAGNITUDE # +Y force - torques = torch.zeros(1, len(body_ids), 3, device=device) - - cube_object.permanent_wrench_composer.set_forces_and_torques_index( - forces=forces, - torques=torques, - positions=com_pos, - body_ids=body_ids, - is_global=True, - ) - - # Phase 1: run 50 steps — force at CoM, expect no rotation - for _ in range(50): - cube_object.write_data_to_sim() - sim.step() - cube_object.update(sim.cfg.dt) - - ang_vel_phase1 = cube_object.data.root_ang_vel_w.torch[0].clone() - lin_vel_phase1 = cube_object.data.root_lin_vel_w.torch[0].clone() - - # Should have linear velocity in +Y - assert lin_vel_phase1[1].item() > 0.1, f"Expected positive Y velocity, got {lin_vel_phase1[1].item()}" - - # Angular velocity should be ~0 (force applied at CoM → no torque) - assert abs(ang_vel_phase1[2].item()) < 0.1, ( - f"Expected ~0 Z angular velocity in phase 1, got {ang_vel_phase1[2].item()}" - ) - - # Phase 2: Teleport cube to origin, zero velocity, don't re-apply force - root_pose2 = cube_object.data.root_pose_w.torch.clone() - root_pose2[0, 0] = 0.0 # x = 0 - root_pose2[0, 1] = 0.0 - root_pose2[0, 2] = 1.0 # z = 1 - root_pose2[0, 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) # identity (xyzw) - cube_object.write_root_pose_to_sim_index(root_pose=root_pose2) - - root_vel2 = cube_object.data.root_vel_w.torch.clone() - root_vel2[0, :] = 0.0 # zero velocity - cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel2) - - # Step once to let state settle - sim.step() - cube_object.update(sim.cfg.dt) - - # Phase 2: run 50 steps — body at origin but stored torque = cross((1,0,1), (0,10,0)) = (-10,0,10) - # correction = -cross((0,0,1), (0,10,0)) = -(0,0,0 - but wait, z=1) - # Actually: stored = cross((com_x,com_y,com_z), (0,10,0)) - # After teleport: correction = -cross(new_pos, F), net torque ≠ 0 since positions differ - for _ in range(50): - cube_object.write_data_to_sim() - sim.step() - cube_object.update(sim.cfg.dt) - - ang_vel_phase2 = cube_object.data.root_ang_vel_w.torch[0].clone() - - # The X component of position changed from ~1 to ~0, so torque about Z changes. - # stored_torque_z = com_x * Fy = ~1 * 10 = ~10 - # After teleport, correction_z = -new_x * Fy = ~0 * 10 = ~0 - # net torque_z ≈ 10 → positive Z angular velocity - assert ang_vel_phase2[2].item() > 0.5, ( - f"Expected positive Z angular velocity in phase 2, got {ang_vel_phase2[2].item()}" - ) - - -@pytest.mark.parametrize("device", test_devices()) -def test_global_force_no_position_no_torque(device): - """Test that global force without positions produces no torque (applied at CoM). - - A body at (2, 0, 1) with global F=(0, 10, 0) and no positions should experience - only linear acceleration, no rotation. The force is applied at the body's CoM. - """ - with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_object, _ = generate_cubes_scene(num_cubes=1, height=1.0, device=device) - - sim.reset() - - body_ids, _ = cube_object.find_bodies(".*") - - # Move cube to (2, 0, 1) - root_pose = cube_object.data.root_pose_w.torch.clone() - root_pose[0, 0] = 2.0 - root_pose[0, 1] = 0.0 - root_pose[0, 2] = 1.0 - root_pose[0, 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) # identity (xyzw) - cube_object.write_root_pose_to_sim_index(root_pose=root_pose) - - root_vel = cube_object.data.root_vel_w.torch.clone() - root_vel[0, :] = 0.0 - cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) - sim.step() - cube_object.update(sim.cfg.dt) - - # Apply global F=(0, 10, 0) WITHOUT positions → force at CoM, no torque - forces = torch.zeros(1, len(body_ids), 3, device=device) - forces[..., 1] = FORCE_MAGNITUDE - torques = torch.zeros(1, len(body_ids), 3, device=device) - - cube_object.permanent_wrench_composer.set_forces_and_torques_index( - forces=forces, - torques=torques, - body_ids=body_ids, - is_global=True, - ) - - # Run 50 steps - for _ in range(50): - cube_object.write_data_to_sim() - sim.step() - cube_object.update(sim.cfg.dt) - - omega_z = cube_object.data.root_ang_vel_w.torch[0, 2].item() - # No positions → force at CoM → zero torque → zero angular velocity - assert abs(omega_z) < 0.01, f"Expected ~zero omega_z for force at CoM, got {omega_z}" - - # Should still have linear acceleration in +Y - lin_vel_y = cube_object.data.root_lin_vel_w.torch[0, 1].item() - assert lin_vel_y > 0.1, f"Expected positive Y velocity from applied force, got {lin_vel_y}" - - -@pytest.mark.parametrize("device", test_devices()) -def test_multi_cube_different_torques_from_same_force(device): - """Test kernel indexing across multiple envs with different CoM positions. - - 2 cubes: Cube 0 at (-1, 0, 1), Cube 1 at (+1, 0, 1). - Same global F=(0, 10, 0) at same world point P=(0, 0, 1) to both cubes. - Cube 0: torque_z = cross((1,0,0), (0,10,0))_z = +10 → omega_z > 0 - Cube 1: torque_z = cross((-1,0,0), (0,10,0))_z = -10 → omega_z < 0 - Both have same linear acceleration in +Y. - """ - with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_object, _ = generate_cubes_scene(num_cubes=2, height=1.0, device=device) - - sim.reset() - - body_ids, _ = cube_object.find_bodies(".*") - - # Position cubes: Cube 0 at (-1, 0, 1), Cube 1 at (+1, 0, 1) - root_pose = cube_object.data.root_pose_w.torch.clone() - root_pose[0, 0] = -1.0 - root_pose[0, 1] = 0.0 - root_pose[0, 2] = 1.0 - root_pose[0, 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) # identity (xyzw) - - root_pose[1, 0] = 1.0 - root_pose[1, 1] = 0.0 - root_pose[1, 2] = 1.0 - root_pose[1, 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) # identity (xyzw) - cube_object.write_root_pose_to_sim_index(root_pose=root_pose) - - root_vel = cube_object.data.root_vel_w.torch.clone() - root_vel[:, :] = 0.0 - cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) - sim.step() - cube_object.update(sim.cfg.dt) - - # Apply same global F=(0, 10, 0) at P=(0, 0, 1) to both cubes - forces = torch.zeros(2, len(body_ids), 3, device=device) - forces[..., 1] = FORCE_MAGNITUDE - torques = torch.zeros(2, len(body_ids), 3, device=device) - positions = torch.zeros(2, len(body_ids), 3, device=device) - positions[..., 2] = 1.0 # P = (0, 0, 1) - - cube_object.permanent_wrench_composer.set_forces_and_torques_index( - forces=forces, - torques=torques, - positions=positions, - body_ids=body_ids, - is_global=True, - ) - - # Run 50 steps - for _ in range(50): - cube_object.write_data_to_sim() - sim.step() - cube_object.update(sim.cfg.dt) - - # Cube 0: omega_z > 0 (force point is to the right of CoM) - omega_z_0 = cube_object.data.root_ang_vel_w.torch[0, 2].item() - assert omega_z_0 > 0.1, f"Cube 0: expected positive omega_z, got {omega_z_0}" - - # Cube 1: omega_z < 0 (force point is to the left of CoM) - omega_z_1 = cube_object.data.root_ang_vel_w.torch[1, 2].item() - assert omega_z_1 < -0.1, f"Cube 1: expected negative omega_z, got {omega_z_1}" - - # Both cubes should have same linear velocity in +Y (same force magnitude) - lin_vel_y_0 = cube_object.data.root_lin_vel_w.torch[0, 1].item() - lin_vel_y_1 = cube_object.data.root_lin_vel_w.torch[1, 1].item() - assert abs(lin_vel_y_0 - lin_vel_y_1) < 0.5, ( - f"Both cubes should have similar Y velocity, got {lin_vel_y_0} and {lin_vel_y_1}" - ) - - -@pytest.mark.parametrize("device", test_devices()) -def test_global_force_torque_far_from_origin(device): - """Test that global force torque correction produces correct physics at large world coordinates. - - Two cubes with identical relative geometry (force offset = (1, 0, 0) from CoM): - Cube 0 at (0, 0, 1) — near origin (reference) - Cube 1 at (2000, 0, 1) — far from origin - - Both get global F=(0, 10, 0) at offset (1, 0, 0) from their respective CoMs. - Expected torque: cross((1,0,0), (0,10,0)) = (0, 0, 10) for both. - - The compose kernel computes cross(P, F) - cross(link_pos, F): - Cube 0: cross((1,0,1), F) - cross((0,0,1), F) — small values, no cancellation - Cube 1: cross((2001,0,1), F) - cross((2000,0,1), F) — large values nearly cancel - - Both cubes should produce the same angular and linear velocities. - """ - with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_object, _ = generate_cubes_scene(num_cubes=2, height=1.0, device=device) - - sim.reset() - - body_ids, _ = cube_object.find_bodies(".*") - - # Position cubes: Cube 0 near origin, Cube 1 far from origin - root_pose = cube_object.data.root_pose_w.torch.clone() - # Cube 0 at (0, 0, 1) - root_pose[0, 0] = 0.0 - root_pose[0, 1] = 0.0 - root_pose[0, 2] = 1.0 - root_pose[0, 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) # identity (xyzw) - # Cube 1 at (2000, 0, 1) - root_pose[1, 0] = 2000.0 - root_pose[1, 1] = 0.0 - root_pose[1, 2] = 1.0 - root_pose[1, 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) # identity (xyzw) - cube_object.write_root_pose_to_sim_index(root_pose=root_pose) - - root_vel = cube_object.data.root_vel_w.torch.clone() - root_vel[:, :] = 0.0 - cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) - sim.step() - cube_object.update(sim.cfg.dt) - - # Apply F=(0, 10, 0) at +1m X offset from each cube's CoM - forces = torch.zeros(2, len(body_ids), 3, device=device) - forces[..., 1] = FORCE_MAGNITUDE # +Y force - torques = torch.zeros(2, len(body_ids), 3, device=device) - - # Positions: each cube's CoM + (1, 0, 0) - com_pos = cube_object.data.body_com_pos_w.torch[:, body_ids, :3].clone() - positions = com_pos.clone() - positions[..., 0] += 1.0 # +1m X offset from CoM - - cube_object.permanent_wrench_composer.set_forces_and_torques_index( - forces=forces, - torques=torques, - positions=positions, - body_ids=body_ids, - is_global=True, - ) - - # Run 50 steps - for _ in range(50): - cube_object.write_data_to_sim() - sim.step() - cube_object.update(sim.cfg.dt) - - # Both cubes should have positive omega_z (cross((1,0,0), (0,10,0)) = (0,0,10)) - omega_z_0 = cube_object.data.root_ang_vel_w.torch[0, 2].item() - omega_z_1 = cube_object.data.root_ang_vel_w.torch[1, 2].item() - assert omega_z_0 > 0.1, f"Cube 0: expected positive omega_z, got {omega_z_0}" - assert omega_z_1 > 0.1, f"Cube 1: expected positive omega_z, got {omega_z_1}" - - # omega_z values should match within 1% (same relative geometry) - torch.testing.assert_close( - torch.tensor(omega_z_0), - torch.tensor(omega_z_1), - rtol=0.01, - atol=0.0, - msg=lambda msg: ( - f"Angular velocity mismatch between near-origin and far-from-origin cubes:\n" - f" Cube 0 (near): omega_z = {omega_z_0:.6f}\n" - f" Cube 1 (far): omega_z = {omega_z_1:.6f}\n{msg}" - ), - ) - - # Linear velocity in +Y should also match - lin_vel_y_0 = cube_object.data.root_lin_vel_w.torch[0, 1].item() - lin_vel_y_1 = cube_object.data.root_lin_vel_w.torch[1, 1].item() - torch.testing.assert_close( - torch.tensor(lin_vel_y_0), - torch.tensor(lin_vel_y_1), - rtol=0.01, - atol=0.0, - msg=lambda msg: ( - f"Linear velocity mismatch between near-origin and far-from-origin cubes:\n" - f" Cube 0 (near): lin_vel_y = {lin_vel_y_0:.6f}\n" - f" Cube 1 (far): lin_vel_y = {lin_vel_y_1:.6f}\n{msg}" - ), - ) - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -def test_global_force_no_position_no_rotation_large_offset(device): - """Test that a global force without positions produces no rotation at large offsets. - - A cube is placed at (2000, 0, 1) and a global force F=(0, 10, 0) is applied - without positions. The cube should accelerate linearly but not rotate. - Before the fix, this would produce torque proportional to 2000 and cause rotation. - """ +def test_composed_wrenches_reach_physx_in_the_frame_they_are_delivered(device: str) -> None: + """Match raw PhysX for a world force at an offset, and world-frame physics for a world wrench at the CoM.""" with build_simulation_context( - device=device, add_ground_plane=False, auto_add_lighting=True, gravity_enabled=False + device=device, gravity_enabled=False, add_ground_plane=False, auto_add_lighting=True ) as sim: sim._app_control_on_stop_handle = None - cube_object, _ = generate_cubes_scene(num_cubes=1, height=1.0, device=device) - + # Off the world z axis, so the CoM correction cross(com, F) of the composed torque is not zero. + composed = _make_cube("Composed", (1.0, 0.0, 1.0), _ROTATION_45_Z) + raw = _make_cube("Raw", (1.0, 3.0, 1.0), _ROTATION_45_Z) + # Far from the origin, a world force at the CoM must not pick up a torque from the body position. + world = _make_cube("World", (2000.0, 0.0, 1.0), _ROTATION_120_XYZ) sim.reset() - body_ids, _ = cube_object.find_bodies(".*") - - # Place cube at large X offset - root_pose = cube_object.data.default_root_pose.torch.clone() - root_pose[0, 0] = 2000.0 # large X position - root_pose[0, 1] = 0.0 - root_pose[0, 2] = 1.0 - cube_object.write_root_pose_to_sim_index(root_pose=root_pose) - - root_vel = cube_object.data.default_root_vel.torch.clone() - cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) - cube_object.reset() - - # Apply global force without positions (should go to CoM, no torque) - forces = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=device) - forces[0, :, 1] = 10.0 # F_y = 10 N - - cube_object.permanent_wrench_composer.set_forces_and_torques_index( - forces=forces, - body_ids=body_ids, + # A world force at a world offset from the CoM is composed into body-frame force and induced torque. + force = torch.tensor([[[0.0, 0.0, 10.0]]], device=device) + offset = torch.tensor([[[0.0, 1.0, 0.0]]], device=device) + composed_position = composed.data.body_com_pos_w.torch + offset + raw_position = raw.data.body_com_pos_w.torch + offset + composed.permanent_wrench_composer.set_forces_and_torques_index( + forces=force, positions=composed_position, is_global=True + ) + assert composed.permanent_wrench_composer.get_forces_and_torques()[2] is False + + # A world force and torque at the CoM take the world-frame fast path without composition. + world_force, world_torque = 10.0, 0.01 + world.permanent_wrench_composer.set_forces_and_torques_index( + forces=torch.tensor([[[world_force, 0.0, 0.0]]], device=device), + torques=torch.tensor([[[0.0, 0.0, world_torque]]], device=device), is_global=True, ) - - # Step simulation - for _ in range(50): - cube_object.write_data_to_sim() + assert world.permanent_wrench_composer.get_forces_and_torques()[2] is True + + raw_force = wp.from_torch(force.view(-1, 3), dtype=wp.float32) + raw_torque = wp.zeros((1, 3), dtype=wp.float32, device=device) + raw_position_data = wp.from_torch(raw_position.view(-1, 3), dtype=wp.float32) + for _ in range(_NUM_STEPS): + composed.write_data_to_sim() + world.write_data_to_sim() + raw.root_view.apply_forces_and_torques_at_position( + force_data=raw_force, + torque_data=raw_torque, + position_data=raw_position_data, + indices=raw._ALL_INDICES, + is_global=True, + ) sim.step() - cube_object.update(sim.cfg.dt) + for asset in (composed, raw, world): + asset.update(sim.cfg.dt) - # Check: angular velocity should be near zero (no rotation) - ang_vel = cube_object.data.root_ang_vel_w.torch[0] - assert torch.allclose(ang_vel, torch.zeros(3, device=device), atol=0.01), ( - f"Expected near-zero angular velocity, got {ang_vel}. " - "Global force without positions should not produce torque." + torch.testing.assert_close( + composed.data.root_lin_vel_w.torch, raw.data.root_lin_vel_w.torch, rtol=1e-4, atol=1e-4 ) - - # Check: linear velocity in Y should be positive (force is in +Y) - lin_vel = cube_object.data.root_lin_vel_w.torch[0] - assert lin_vel[1] > 0.1, f"Expected positive Y velocity from applied force, got {lin_vel[1]}" + torch.testing.assert_close( + composed.data.root_ang_vel_w.torch, raw.data.root_ang_vel_w.torch, rtol=1e-4, atol=1e-4 + ) + assert torch.abs(composed.data.root_ang_vel_w.torch[0, :2]).max().item() > 0.01 + + # World-frame delivery: the velocities follow the world axes, not the rotated body axes. + duration = _NUM_STEPS * sim.cfg.dt + mass = world.data.body_mass.torch[0, 0].item() + inertia_zz = world.data.body_inertia.torch[0, 0, 8].item() + expected_lin_vel = torch.tensor([[world_force / mass * duration, 0.0, 0.0]], device=device) + expected_ang_vel = torch.tensor([[0.0, 0.0, world_torque / inertia_zz * duration]], device=device) + torch.testing.assert_close(world.data.root_lin_vel_w.torch, expected_lin_vel, rtol=1e-3, atol=1e-4) + torch.testing.assert_close(world.data.root_ang_vel_w.torch, expected_ang_vel, rtol=1e-3, atol=1e-4) diff --git a/source/isaaclab/test/utils/test_wrench_composer_vs_physx.py b/source/isaaclab/test/utils/test_wrench_composer_vs_physx.py deleted file mode 100644 index d3b4051f3507..000000000000 --- a/source/isaaclab/test/utils/test_wrench_composer_vs_physx.py +++ /dev/null @@ -1,701 +0,0 @@ -# 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 - -"""Integration tests comparing WrenchComposer output vs raw PhysX apply_forces_and_torques_at_position. - -Two identical rigid objects are placed in the same scene. One uses the WrenchComposer path -(set_forces_and_torques → write_data_to_sim → compose → PhysX apply with is_global=False), -the other uses the raw PhysX API directly (apply_forces_and_torques_at_position with matching -is_global flag). After N steps, both objects should have identical velocities. -""" - -from isaaclab.test.utils import launch_test_simulation - -launch_test_simulation() - -import math - -import pytest -import torch -import warp as wp - -import isaaclab.sim as sim_utils -from isaaclab.assets import RigidObject, RigidObjectCfg -from isaaclab.sim import build_simulation_context -from isaaclab.test.utils import test_devices -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR - -pytestmark = pytest.mark.integration - - -def generate_dual_cube_scene( - num_cubes: int = 1, - height: float = 1.0, - device: str = "cuda:0", - initial_rot: tuple[float, ...] | None = None, - spacing: float = 2.0, -) -> tuple[RigidObject, RigidObject]: - """Generate a scene with two sets of cubes: one for the composer path, one for raw PhysX. - - Both sets share the same spawn config and initial state (except a Y offset to avoid overlap). - - Args: - num_cubes: Number of cubes per group (environments). - height: Spawn height. - device: Simulation device. - initial_rot: Initial quaternion (x, y, z, w). Defaults to identity. - spacing: Distance between env origins in X. Defaults to 2.0. - - Returns: - Tuple of (cube_composer, cube_raw) RigidObject instances. - """ - if initial_rot is None: - initial_rot = (0.0, 0.0, 0.0, 1.0) # identity in (x,y,z,w) - - y_offset = max(spacing, 3.0) - - # Create Xform prims for both groups - for i in range(num_cubes): - origin_composer = (i * spacing, 0.0, height) - origin_raw = (i * spacing, y_offset, height) # Y offset to avoid overlap - sim_utils.create_prim(f"/World/Composer_{i}", "Xform", translation=origin_composer) - sim_utils.create_prim(f"/World/Raw_{i}", "Xform", translation=origin_raw) - - spawn_cfg = sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", - rigid_props=sim_utils.UsdPhysicsRigidBodyCfg(), - ) - - cube_composer_cfg = RigidObjectCfg( - prim_path="/World/Composer_[^/]*/Object", - spawn=spawn_cfg, - init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, height), rot=initial_rot), - ) - cube_composer = RigidObject(cfg=cube_composer_cfg) - - cube_raw_cfg = RigidObjectCfg( - prim_path="/World/Raw_[^/]*/Object", - spawn=spawn_cfg, - init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, y_offset, height), rot=initial_rot), - ) - cube_raw = RigidObject(cfg=cube_raw_cfg) - - return cube_composer, cube_raw - - -N_STEPS = 50 -FORCE_MAGNITUDE = 10.0 -TORQUE_MAGNITUDE = 1.0 -# 45 degrees about Z: (cos(22.5°), 0, 0, sin(22.5°)) -ROT_45_Z = (0.0, 0.0, math.sin(math.pi / 8), math.cos(math.pi / 8)) # 45deg about Z in (x,y,z,w) - - -@pytest.mark.parametrize("device", test_devices()) -def test_composer_vs_physx_local_force(device): - """Baseline: local force at identity orientation. Composer and raw PhysX should match exactly.""" - with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device) - - sim.reset() - - body_ids, _ = cube_composer.find_bodies(".*") - - # Composer path: local force +X - forces = torch.zeros(1, len(body_ids), 3, device=device) - forces[..., 0] = FORCE_MAGNITUDE - torques = torch.zeros(1, len(body_ids), 3, device=device) - - cube_composer.permanent_wrench_composer.set_forces_and_torques_index( - forces=forces, - torques=torques, - body_ids=body_ids, - is_global=False, - ) - - # Raw PhysX data (flattened for PhysX view API) - raw_forces = torch.zeros(1, 3, device=device) - raw_forces[:, 0] = FORCE_MAGNITUDE - raw_torques = torch.zeros(1, 3, device=device) - raw_indices = cube_raw._ALL_INDICES - - for _ in range(N_STEPS): - cube_composer.write_data_to_sim() - cube_raw.write_data_to_sim() # no-op (composer inactive) - cube_raw.root_view.apply_forces_and_torques_at_position( - force_data=wp.from_torch(raw_forces.contiguous(), dtype=wp.float32), - torque_data=wp.from_torch(raw_torques.contiguous(), dtype=wp.float32), - position_data=None, - indices=raw_indices, - is_global=False, - ) - sim.step() - cube_composer.update(sim.cfg.dt) - cube_raw.update(sim.cfg.dt) - - # Compare velocities - torch.testing.assert_close( - cube_composer.data.root_lin_vel_w.torch, - cube_raw.data.root_lin_vel_w.torch, - rtol=1e-4, - atol=1e-4, - ) - # Both should have ~zero angular velocity (force at CoM, no torque) - torch.testing.assert_close( - cube_composer.data.root_ang_vel_w.torch, - torch.zeros(1, 3, device=device), - rtol=0.0, - atol=1e-4, - ) - torch.testing.assert_close( - cube_raw.data.root_ang_vel_w.torch, - torch.zeros(1, 3, device=device), - rtol=0.0, - atol=1e-4, - ) - - -@pytest.mark.parametrize("device", test_devices()) -def test_composer_vs_physx_local_force_at_position(device): - """Local force at a local offset. Both paths should produce identical cross-product torque.""" - with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device) - - sim.reset() - - body_ids, _ = cube_composer.find_bodies(".*") - - # Local force +X at local offset +0.5m Y - forces = torch.zeros(1, len(body_ids), 3, device=device) - forces[..., 0] = FORCE_MAGNITUDE - torques = torch.zeros(1, len(body_ids), 3, device=device) - positions = torch.zeros(1, len(body_ids), 3, device=device) - positions[..., 1] = 0.5 # +0.5m Y offset in local frame - - cube_composer.permanent_wrench_composer.set_forces_and_torques_index( - forces=forces, - torques=torques, - positions=positions, - body_ids=body_ids, - is_global=False, - ) - - # Raw PhysX data (local force at local position) - raw_forces = torch.zeros(1, 3, device=device) - raw_forces[:, 0] = FORCE_MAGNITUDE - raw_torques = torch.zeros(1, 3, device=device) - raw_positions = torch.zeros(1, 3, device=device) - raw_positions[:, 1] = 0.5 - raw_indices = cube_raw._ALL_INDICES - - for _ in range(N_STEPS): - cube_composer.write_data_to_sim() - cube_raw.write_data_to_sim() - cube_raw.root_view.apply_forces_and_torques_at_position( - force_data=wp.from_torch(raw_forces.contiguous(), dtype=wp.float32), - torque_data=wp.from_torch(raw_torques.contiguous(), dtype=wp.float32), - position_data=wp.from_torch(raw_positions.contiguous(), dtype=wp.float32), - indices=raw_indices, - is_global=False, - ) - sim.step() - cube_composer.update(sim.cfg.dt) - cube_raw.update(sim.cfg.dt) - - # Both linear and angular velocities should match - torch.testing.assert_close( - cube_composer.data.root_lin_vel_w.torch, - cube_raw.data.root_lin_vel_w.torch, - rtol=1e-4, - atol=1e-4, - ) - torch.testing.assert_close( - cube_composer.data.root_ang_vel_w.torch, - cube_raw.data.root_ang_vel_w.torch, - rtol=1e-4, - atol=1e-4, - ) - - # Sanity: angular velocity should be nonzero (cross-product torque) - assert torch.abs(cube_composer.data.root_ang_vel_w.torch[0, 2]).item() > 0.1, ( - "Expected nonzero Z angular velocity from cross-product torque" - ) - - -@pytest.mark.parametrize("device", test_devices()) -def test_composer_vs_physx_global_force_at_position(device): - """Global force at world position with non-identity rotation. Both rotation AND position correction matter.""" - with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device, initial_rot=ROT_45_Z) - - sim.reset() - - body_ids, _ = cube_composer.find_bodies(".*") - - # Global force +X - forces = torch.zeros(1, len(body_ids), 3, device=device) - forces[..., 0] = FORCE_MAGNITUDE - torques = torch.zeros(1, len(body_ids), 3, device=device) - - # Position = each cube's link_pos + offset (same offset for both) - offset = torch.zeros(1, len(body_ids), 3, device=device) - offset[..., 1] = 1.0 # +1m Y offset in world frame - - pos_composer = cube_composer.data.body_com_pos_w.torch[:, body_ids, :3].clone() + offset - pos_raw = cube_raw.data.body_com_pos_w.torch[:, body_ids, :3].clone() + offset - - cube_composer.permanent_wrench_composer.set_forces_and_torques_index( - forces=forces, - torques=torques, - positions=pos_composer, - body_ids=body_ids, - is_global=True, - ) - - # Raw PhysX data - raw_forces = torch.zeros(1, 3, device=device) - raw_forces[:, 0] = FORCE_MAGNITUDE - raw_torques = torch.zeros(1, 3, device=device) - raw_positions = pos_raw.view(-1, 3) - raw_indices = cube_raw._ALL_INDICES - - for _ in range(N_STEPS): - cube_composer.write_data_to_sim() - cube_raw.write_data_to_sim() - cube_raw.root_view.apply_forces_and_torques_at_position( - force_data=wp.from_torch(raw_forces.contiguous(), dtype=wp.float32), - torque_data=wp.from_torch(raw_torques.contiguous(), dtype=wp.float32), - position_data=wp.from_torch(raw_positions.contiguous(), dtype=wp.float32), - indices=raw_indices, - is_global=True, - ) - sim.step() - cube_composer.update(sim.cfg.dt) - cube_raw.update(sim.cfg.dt) - - # Both linear and angular velocities should match - torch.testing.assert_close( - cube_composer.data.root_lin_vel_w.torch, - cube_raw.data.root_lin_vel_w.torch, - rtol=1e-4, - atol=1e-4, - ) - torch.testing.assert_close( - cube_composer.data.root_ang_vel_w.torch, - cube_raw.data.root_ang_vel_w.torch, - rtol=1e-4, - atol=1e-4, - ) - - # Sanity: angular velocity should be nonzero (cross-product torque) - assert torch.abs(cube_composer.data.root_ang_vel_w.torch[0, 2]).item() > 0.1, ( - "Expected nonzero Z angular velocity from positional torque" - ) - - -@pytest.mark.parametrize("device", test_devices()) -def test_composer_vs_physx_local_torque(device): - """Local torque at identity orientation. Should produce matching angular velocity.""" - with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device) - - sim.reset() - - body_ids, _ = cube_composer.find_bodies(".*") - - # Composer path: local torque about +Z - forces = torch.zeros(1, len(body_ids), 3, device=device) - torques = torch.zeros(1, len(body_ids), 3, device=device) - torques[..., 2] = TORQUE_MAGNITUDE - - cube_composer.permanent_wrench_composer.set_forces_and_torques_index( - forces=forces, - torques=torques, - body_ids=body_ids, - is_global=False, - ) - - # Raw PhysX data - raw_forces = torch.zeros(1, 3, device=device) - raw_torques = torch.zeros(1, 3, device=device) - raw_torques[:, 2] = TORQUE_MAGNITUDE - raw_indices = cube_raw._ALL_INDICES - - for _ in range(N_STEPS): - cube_composer.write_data_to_sim() - cube_raw.write_data_to_sim() - cube_raw.root_view.apply_forces_and_torques_at_position( - force_data=wp.from_torch(raw_forces.contiguous(), dtype=wp.float32), - torque_data=wp.from_torch(raw_torques.contiguous(), dtype=wp.float32), - position_data=None, - indices=raw_indices, - is_global=False, - ) - sim.step() - cube_composer.update(sim.cfg.dt) - cube_raw.update(sim.cfg.dt) - - # Angular velocities should match - torch.testing.assert_close( - cube_composer.data.root_ang_vel_w.torch, - cube_raw.data.root_ang_vel_w.torch, - rtol=1e-4, - atol=1e-4, - ) - # Linear velocity should be ~zero for both (no force) - torch.testing.assert_close( - cube_composer.data.root_lin_vel_w.torch, - torch.zeros(1, 3, device=device), - rtol=0.0, - atol=1e-4, - ) - torch.testing.assert_close( - cube_raw.data.root_lin_vel_w.torch, - torch.zeros(1, 3, device=device), - rtol=0.0, - atol=1e-4, - ) - - -@pytest.mark.parametrize("device", test_devices()) -def test_composer_vs_physx_global_torque(device): - """Global torque with non-identity rotation (45 deg Z). Composer rotates to body frame internally.""" - with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device, initial_rot=ROT_45_Z) - - sim.reset() - - body_ids, _ = cube_composer.find_bodies(".*") - - # Composer path: global torque about +Z - forces = torch.zeros(1, len(body_ids), 3, device=device) - torques = torch.zeros(1, len(body_ids), 3, device=device) - torques[..., 2] = TORQUE_MAGNITUDE - - cube_composer.permanent_wrench_composer.set_forces_and_torques_index( - forces=forces, - torques=torques, - body_ids=body_ids, - is_global=True, - ) - - # Raw PhysX data - raw_forces = torch.zeros(1, 3, device=device) - raw_torques = torch.zeros(1, 3, device=device) - raw_torques[:, 2] = TORQUE_MAGNITUDE - raw_indices = cube_raw._ALL_INDICES - - for _ in range(N_STEPS): - cube_composer.write_data_to_sim() - cube_raw.write_data_to_sim() - cube_raw.root_view.apply_forces_and_torques_at_position( - force_data=wp.from_torch(raw_forces.contiguous(), dtype=wp.float32), - torque_data=wp.from_torch(raw_torques.contiguous(), dtype=wp.float32), - position_data=None, - indices=raw_indices, - is_global=True, - ) - sim.step() - cube_composer.update(sim.cfg.dt) - cube_raw.update(sim.cfg.dt) - - # Angular velocities should match - torch.testing.assert_close( - cube_composer.data.root_ang_vel_w.torch, - cube_raw.data.root_ang_vel_w.torch, - rtol=1e-4, - atol=1e-4, - ) - - -NUM_CUBES_MULTI = 4 - - -@pytest.mark.parametrize("device", test_devices()) -def test_composer_vs_physx_global_force_with_reset(device): - """Global force (no position) with a mid-simulation reset of half the envs. - - Regression: after reset the permanent wrench is cleared. Re-setting it should - produce correct behavior even though the object state was just reset. - """ - with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_composer, cube_raw = generate_dual_cube_scene( - num_cubes=NUM_CUBES_MULTI, device=device, initial_rot=ROT_45_Z, spacing=20.0 - ) - - sim.reset() - - # Capture initial world-frame state (includes env origin offsets) - cube_composer.update(sim.cfg.dt) - cube_raw.update(sim.cfg.dt) - initial_state_composer = torch.cat( - [ - cube_composer.data.root_link_pos_w.torch, - cube_composer.data.root_link_quat_w.torch, - cube_composer.data.root_com_vel_w.torch, - ], - dim=-1, - ).clone() - initial_state_raw = torch.cat( - [ - cube_raw.data.root_link_pos_w.torch, - cube_raw.data.root_link_quat_w.torch, - cube_raw.data.root_com_vel_w.torch, - ], - dim=-1, - ).clone() - - body_ids, _ = cube_composer.find_bodies(".*") - - def apply_global_force(): - """Set the same global +X force on the composer cube.""" - forces = torch.zeros(NUM_CUBES_MULTI, len(body_ids), 3, device=device) - forces[..., 0] = FORCE_MAGNITUDE - torques = torch.zeros(NUM_CUBES_MULTI, len(body_ids), 3, device=device) - cube_composer.permanent_wrench_composer.set_forces_and_torques_index( - forces=forces, - torques=torques, - body_ids=body_ids, - is_global=True, - ) - - apply_global_force() - - # Raw PhysX data - raw_forces = torch.zeros(NUM_CUBES_MULTI, 3, device=device) - raw_forces[:, 0] = FORCE_MAGNITUDE - raw_torques = torch.zeros(NUM_CUBES_MULTI, 3, device=device) - raw_indices = cube_raw._ALL_INDICES - - # Phase 1: run N_STEPS / 2 - half = N_STEPS // 2 - for _ in range(half): - cube_composer.write_data_to_sim() - cube_raw.write_data_to_sim() - cube_raw.root_view.apply_forces_and_torques_at_position( - force_data=wp.from_torch(raw_forces.contiguous(), dtype=wp.float32), - torque_data=wp.from_torch(raw_torques.contiguous(), dtype=wp.float32), - position_data=None, - indices=raw_indices, - is_global=True, - ) - sim.step() - cube_composer.update(sim.cfg.dt) - cube_raw.update(sim.cfg.dt) - - # Reset first half of envs on both cubes - reset_ids = list(range(NUM_CUBES_MULTI // 2)) - reset_ids_torch = torch.tensor(reset_ids, dtype=torch.long, device=device) - - # Reset root state using captured world-frame initial state (includes env origins) - cube_composer.write_root_link_pose_to_sim_index( - root_pose=initial_state_composer[reset_ids_torch, :7], env_ids=reset_ids_torch - ) - cube_composer.write_root_com_velocity_to_sim_index( - root_velocity=initial_state_composer[reset_ids_torch, 7:], env_ids=reset_ids_torch - ) - cube_raw.write_root_link_pose_to_sim_index( - root_pose=initial_state_raw[reset_ids_torch, :7], env_ids=reset_ids_torch - ) - cube_raw.write_root_com_velocity_to_sim_index( - root_velocity=initial_state_raw[reset_ids_torch, 7:], env_ids=reset_ids_torch - ) - - cube_composer.reset(reset_ids) - cube_raw.reset(reset_ids) - - # Re-apply the force (reset cleared the permanent wrench) - apply_global_force() - - # Phase 2: run N_STEPS / 2 more - for _ in range(half): - cube_composer.write_data_to_sim() - cube_raw.write_data_to_sim() - cube_raw.root_view.apply_forces_and_torques_at_position( - force_data=wp.from_torch(raw_forces.contiguous(), dtype=wp.float32), - torque_data=wp.from_torch(raw_torques.contiguous(), dtype=wp.float32), - position_data=None, - indices=raw_indices, - is_global=True, - ) - sim.step() - cube_composer.update(sim.cfg.dt) - cube_raw.update(sim.cfg.dt) - - # All envs: composer vs raw should match - torch.testing.assert_close( - cube_composer.data.root_lin_vel_w.torch, - cube_raw.data.root_lin_vel_w.torch, - rtol=1e-4, - atol=1e-4, - ) - torch.testing.assert_close( - cube_composer.data.root_ang_vel_w.torch, - cube_raw.data.root_ang_vel_w.torch, - rtol=1e-4, - atol=1e-4, - ) - # All envs should have ~zero angular velocity - torch.testing.assert_close( - cube_composer.data.root_ang_vel_w.torch, - torch.zeros(NUM_CUBES_MULTI, 3, device=device), - rtol=0.0, - atol=1e-4, - ) - - -@pytest.mark.parametrize("device", test_devices()) -def test_composer_vs_physx_payload_scenario(device): - """Mirrors the apply_payload MDP: permanent global downward force at CoM with gravity. - - A constant world-frame downward force (payload weight) is applied via the composer - path vs raw PhysX. The body falls under gravity + payload, contacts the ground, and - orientation changes. The composer does a world->body->world round-trip each step; - this test catches any precision drift from that. - """ - with build_simulation_context(device=device, gravity_enabled=True, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_composer, cube_raw = generate_dual_cube_scene( - num_cubes=1, height=0.5, device=device, initial_rot=ROT_45_Z, spacing=20.0 - ) - - sim.reset() - cube_composer.update(sim.cfg.dt) - cube_raw.update(sim.cfg.dt) - - # Record initial positions to compare displacements (cubes spawn at different Y) - init_pos_composer = cube_composer.data.root_pos_w.torch.clone() - init_pos_raw = cube_raw.data.root_pos_w.torch.clone() - - body_ids, _ = cube_composer.find_bodies(".*") - - payload_force = 2.0 * 9.81 - forces = torch.zeros(1, len(body_ids), 3, device=device) - forces[..., 2] = -payload_force - torques = torch.zeros(1, len(body_ids), 3, device=device) - - cube_composer.permanent_wrench_composer.set_forces_and_torques_index( - forces=forces, - torques=torques, - body_ids=body_ids, - is_global=True, - ) - - raw_forces = torch.zeros(1, 3, device=device) - raw_forces[:, 2] = -payload_force - raw_torques = torch.zeros(1, 3, device=device) - raw_indices = cube_raw._ALL_INDICES - - for _ in range(N_STEPS): - cube_composer.write_data_to_sim() - cube_raw.write_data_to_sim() - cube_raw.root_view.apply_forces_and_torques_at_position( - force_data=wp.from_torch(raw_forces.contiguous(), dtype=wp.float32), - torque_data=wp.from_torch(raw_torques.contiguous(), dtype=wp.float32), - position_data=None, - indices=raw_indices, - is_global=True, - ) - sim.step() - cube_composer.update(sim.cfg.dt) - cube_raw.update(sim.cfg.dt) - - # Compare displacements (not absolute positions — cubes have different spawn Y) - disp_composer = cube_composer.data.root_pos_w.torch - init_pos_composer - disp_raw = cube_raw.data.root_pos_w.torch - init_pos_raw - - torch.testing.assert_close(disp_composer, disp_raw, rtol=1e-4, atol=1e-4) - torch.testing.assert_close( - cube_composer.data.root_lin_vel_w.torch, - cube_raw.data.root_lin_vel_w.torch, - rtol=1e-4, - atol=1e-4, - ) - torch.testing.assert_close( - cube_composer.data.root_ang_vel_w.torch, - cube_raw.data.root_ang_vel_w.torch, - rtol=1e-4, - atol=1e-4, - ) - - -@pytest.mark.parametrize("device", test_devices()) -def test_composer_vs_physx_permanent_global_force_at_position_long_run(device): - """Permanent global force at a world-frame offset, run long enough for significant body motion. - - This test catches temporal drift bugs where the stored positional torque diverges from - what PhysX computes each step as the body moves. The force is large enough that the body - translates and rotates significantly over 100 steps, but not so large that it causes - numerical instability. - """ - with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device, initial_rot=ROT_45_Z) - - sim.reset() - - body_ids, _ = cube_composer.find_bodies(".*") - - # Global force +Z at +1m Y offset from CoM — produces torque around X - forces = torch.zeros(1, len(body_ids), 3, device=device) - forces[..., 2] = FORCE_MAGNITUDE - torques = torch.zeros(1, len(body_ids), 3, device=device) - - offset = torch.zeros(1, len(body_ids), 3, device=device) - offset[..., 1] = 1.0 - - pos_composer = cube_composer.data.body_com_pos_w.torch[:, body_ids, :3].clone() + offset - pos_raw = cube_raw.data.body_com_pos_w.torch[:, body_ids, :3].clone() + offset - - cube_composer.permanent_wrench_composer.set_forces_and_torques_index( - forces=forces, - torques=torques, - positions=pos_composer, - body_ids=body_ids, - is_global=True, - ) - - raw_forces = torch.zeros(1, 3, device=device) - raw_forces[:, 2] = FORCE_MAGNITUDE - raw_torques = torch.zeros(1, 3, device=device) - raw_positions = pos_raw.view(-1, 3) - raw_indices = cube_raw._ALL_INDICES - - for _ in range(100): - cube_composer.write_data_to_sim() - cube_raw.write_data_to_sim() - cube_raw.root_view.apply_forces_and_torques_at_position( - force_data=wp.from_torch(raw_forces.contiguous(), dtype=wp.float32), - torque_data=wp.from_torch(raw_torques.contiguous(), dtype=wp.float32), - position_data=wp.from_torch(raw_positions.contiguous(), dtype=wp.float32), - indices=raw_indices, - is_global=True, - ) - sim.step() - cube_composer.update(sim.cfg.dt) - cube_raw.update(sim.cfg.dt) - - torch.testing.assert_close( - cube_composer.data.root_lin_vel_w.torch, - cube_raw.data.root_lin_vel_w.torch, - rtol=1e-3, - atol=1e-3, - ) - torch.testing.assert_close( - cube_composer.data.root_ang_vel_w.torch, - cube_raw.data.root_ang_vel_w.torch, - rtol=1e-3, - atol=1e-3, - ) - - # Sanity: angular velocity should be nonzero - assert torch.abs(cube_composer.data.root_ang_vel_w.torch).max().item() > 0.1, ( - "Expected nonzero angular velocity from positional torque over 100 steps" - ) diff --git a/source/isaaclab_newton/changelog.d/antoiner-asset-tests-redesign.skip b/source/isaaclab_newton/changelog.d/antoiner-asset-tests-redesign.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_newton/test/assets/data/fixed_cartpole.usda b/source/isaaclab_newton/test/assets/data/fixed_cartpole.usda new file mode 100644 index 000000000000..ca71b53ab2e7 --- /dev/null +++ b/source/isaaclab_newton/test/assets/data/fixed_cartpole.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0192a110bc566113bebf4f91260f4e01614877d391bbe216faccdafc6e5adad8 +size 1374 diff --git a/source/isaaclab_newton/test/assets/data/fixed_spatial_chain.usda b/source/isaaclab_newton/test/assets/data/fixed_spatial_chain.usda new file mode 100644 index 000000000000..7bb00bdc88e5 --- /dev/null +++ b/source/isaaclab_newton/test/assets/data/fixed_spatial_chain.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fba7dc964ecf2cbd2a8991eddb30459694b8a31f3a652d10ee3f3e68f3245d41 +size 5270 diff --git a/source/isaaclab_newton/test/assets/data/fixed_tendon_hand.usda b/source/isaaclab_newton/test/assets/data/fixed_tendon_hand.usda new file mode 100644 index 000000000000..598edc4b8bea --- /dev/null +++ b/source/isaaclab_newton/test/assets/data/fixed_tendon_hand.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5b9cf48d11088e1288469e9121ce10e1fa79b3036a0808dad1c3f4ad8c91e061 +size 4509 diff --git a/source/isaaclab_newton/test/assets/data/floating_two_leg.usda b/source/isaaclab_newton/test/assets/data/floating_two_leg.usda new file mode 100644 index 000000000000..372a984d188d --- /dev/null +++ b/source/isaaclab_newton/test/assets/data/floating_two_leg.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:993f2b5c29cd9082308c395199876ab418ad5886e7332b1ce52243998791bdb7 +size 4910 diff --git a/source/isaaclab_newton/test/assets/data/floating_two_link.usda b/source/isaaclab_newton/test/assets/data/floating_two_link.usda new file mode 100644 index 000000000000..d1b7fcc35fc2 --- /dev/null +++ b/source/isaaclab_newton/test/assets/data/floating_two_link.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4b5d751ea914f448fbebbaf8a6bd16d3847ef0f03e7b27adde5b61749c67951b +size 1429 diff --git a/source/isaaclab_newton/test/assets/data/revolute_pendulum.usda b/source/isaaclab_newton/test/assets/data/revolute_pendulum.usda new file mode 100644 index 000000000000..67657cfe0e84 --- /dev/null +++ b/source/isaaclab_newton/test/assets/data/revolute_pendulum.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5e011baeb4d866706a8d19e9c4684d3f422dc8c1bb358899ef82d78eeef1603e +size 1087 diff --git a/source/isaaclab_newton/test/assets/data/rigid_body_with_articulation_root.usda b/source/isaaclab_newton/test/assets/data/rigid_body_with_articulation_root.usda new file mode 100644 index 000000000000..2f3ca98a09fe --- /dev/null +++ b/source/isaaclab_newton/test/assets/data/rigid_body_with_articulation_root.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:74d8aa1f1e23ae04ecb0bab5cd92cb6edcffcfd20980c113f47ca05eb19d1a53 +size 797 diff --git a/source/isaaclab_newton/test/assets/newton_test_utils.py b/source/isaaclab_newton/test/assets/newton_test_utils.py new file mode 100644 index 000000000000..b98d6b67e715 --- /dev/null +++ b/source/isaaclab_newton/test/assets/newton_test_utils.py @@ -0,0 +1,157 @@ +# 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 + +"""Locally authored scenes shared by the kitless Newton asset tests. + +The small USD scenes live in ``test/assets/data`` and need neither Kit nor Nucleus. +Tests load them directly, so their bodies, joints, mass properties, and drives are readable without running +a Python scene generator. +""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from pathlib import Path +from typing import Any + +import torch +import warp as wp +from isaaclab_newton.assets import Articulation, RigidObject, RigidObjectCollection +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +from isaaclab_newton.physics import NewtonManager as SimulationManager +from newton import ModelFlags + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationCfg, RigidObjectCfg, RigidObjectCollectionCfg +from isaaclab.cloner import CloneCfg, clone_plan_from_env_0, replicate +from isaaclab.sim import SimulationCfg, SimulationContext + +NUM_ENVS = 2 +"""Environments per scene: enough to prove a partial write leaves the other environment untouched.""" + +WRIST_USD_STIFFNESS = (0.5, 0.4, 0.3) +"""Angular drive stiffness authored on the fixed chain's wrist joints [N·m/deg]; the damping is a tenth of it.""" + +Vec3 = tuple[float, float, float] + +_DATA_DIR = Path(__file__).resolve().parent / "data" +"""Directory holding the checked-in USD fixtures.""" + +_ENV_SPACING = 40.0 +"""Distance between environment origins along x [m].""" + + +def local_usd(name: str, **kwargs: Any) -> sim_utils.UsdFileCfg: + """Return a spawner for one of the local USD fixtures. + + Args: + name: File name in ``test/assets/data``. + **kwargs: Additional :class:`~isaaclab.sim.UsdFileCfg` fields. + + Returns: + The spawner configuration. + """ + return sim_utils.UsdFileCfg(usd_path=str(_DATA_DIR / name), **kwargs) + + +def newton_sim_cfg( + device: str = "cpu", + *, + dt: float = 1.0 / 120.0, + gravity: Vec3 = (0.0, 0.0, 0.0), + use_newton_actuators: bool = True, + newton_contacts: bool = False, + solver_cfg: MJWarpSolverCfg | None = None, + **newton_kwargs: Any, +) -> SimulationCfg: + """Return an MJWarp simulation configuration, without CUDA-graph capture unless requested. + + Args: + device: Simulation device. + dt: Physics time step [s]. + gravity: Scene gravity [m/s^2]. Defaults to none, so free bodies only move when a test drives them. + use_newton_actuators: Whether explicit actuator groups execute as Newton-native actuators. + newton_contacts: Whether to use Newton's collision pipeline instead of MuJoCo contacts. + solver_cfg: Solver configuration. Defaults to :class:`MJWarpSolverCfg` with the selected contacts. + **newton_kwargs: Additional :class:`NewtonCfg` fields. + + Returns: + The simulation configuration. + """ + if solver_cfg is None: + solver_cfg = MJWarpSolverCfg(use_mujoco_contacts=not newton_contacts) + newton_kwargs.setdefault("use_cuda_graph", False) + return SimulationCfg( + device=device, + dt=dt, + gravity=gravity, + use_newton_actuators=use_newton_actuators, + physics=NewtonCfg(solver_cfg=solver_cfg, **newton_kwargs), + ) + + +def env_origins(device: str, num_envs: int = NUM_ENVS) -> torch.Tensor: + """Return the world origins :func:`spawn_assets` places the environments at [m].""" + origins = torch.zeros((num_envs, 3), device=device) + origins[:, 0] = _ENV_SPACING * torch.arange(num_envs, device=device) + return origins + + +def spawn_assets( + asset_cfgs: Mapping[str, ArticulationCfg | RigidObjectCfg | RigidObjectCollectionCfg], num_envs: int = NUM_ENVS +) -> dict[str, Any]: + """Plan, construct, and replicate assets into ``/World/Env_*`` of the active simulation. + + Call this inside a simulation context and before :meth:`~isaaclab.sim.SimulationContext.reset`. + Each configuration's ``prim_path`` must address one asset per environment. + + Args: + asset_cfgs: Asset configurations keyed by a name the caller uses to look the asset up. + num_envs: Number of environments. + + Returns: + The Newton assets keyed like ``asset_cfgs``. + """ + positions = env_origins("cpu", num_envs).numpy() + # a collection is planned as its member rigid objects + planned_cfgs = [] + for cfg in asset_cfgs.values(): + planned_cfgs.extend(cfg.rigid_objects.values() if isinstance(cfg, RigidObjectCollectionCfg) else (cfg,)) + sim_utils.create_prim("/World/Env_0", "Xform") + clone_plan_from_env_0( + CloneCfg(clone_template="/World/Env_{}"), planned_cfgs, num_envs, _ENV_SPACING, positions=positions + ) + assets = {} + for name, cfg in asset_cfgs.items(): + if isinstance(cfg, ArticulationCfg): + assets[name] = Articulation(cfg) + elif isinstance(cfg, RigidObjectCollectionCfg): + assets[name] = RigidObjectCollection(cfg) + elif isinstance(cfg, RigidObjectCfg): + assets[name] = RigidObject(cfg) + else: + raise TypeError(f"Unsupported asset configuration for '{name}': {type(cfg).__name__}.") + replicate(SimulationContext.instance().get_clone_plan()) + return assets + + +@contextmanager +def world_gravity(gravity: Vec3) -> Iterator[None]: + """Apply the same gravity to every world of the live Newton model, then restore the previous values. + + Args: + gravity: Gravity applied to every world [m/s^2]. + """ + model = SimulationManager.get_model() + world_gravity_view = wp.to_torch(model.gravity[: model.world_count]) + previous = world_gravity_view.clone() + world_gravity_view.copy_(torch.tensor(gravity, device=previous.device)) + SimulationManager.add_model_change(ModelFlags.MODEL_PROPERTIES) + try: + yield + finally: + world_gravity_view.copy_(previous) + SimulationManager.add_model_change(ModelFlags.MODEL_PROPERTIES) diff --git a/source/isaaclab_newton/test/assets/test_articulation.py b/source/isaaclab_newton/test/assets/test_articulation.py index 511de3ef2426..89efea0f137a 100644 --- a/source/isaaclab_newton/test/assets/test_articulation.py +++ b/source/isaaclab_newton/test/assets/test_articulation.py @@ -6,386 +6,85 @@ # ignore private usage of variables warning # pyright: reportPrivateUsage=none -from isaaclab.test.utils import DeviceScope, launch_test_simulation, test_devices -from isaaclab.test.utils.articulation_ordering import ( - ANYMAL_C_PHYSX_JOINT_NAMES, - BRANCHING_MJWARP_BODY_NAMES, - BRANCHING_MJWARP_JOINT_NAMES, - BRANCHING_PHYSX_BODY_NAMES, - BRANCHING_PHYSX_JOINT_NAMES, - PANDA_JOINT_NAMES, - PANDA_ROOT_PRESERVING_REVERSED_BODY_NAMES, -) -from isaaclab.utils import clone, replace +"""Kitless real-Newton articulation coverage on one persistent composite scene. + +Every articulation is a small locally authored fixture (see ``articulation_test_utils``). One module-scoped +scene holds isolated two-environment islands, one per fixture configuration, and is reset once. Each test drives +the islands it names and targets partial writes at environment 1 so it can prove environment 0 keeps its real +backend state. The islands of one environment share a solver world, so every test starts with all islands at +rest and restores the properties it changes. Scene gravity is off; a test that needs gravity applies it to every +world for its own duration. -HEADLESS = True +Configurations that fail initialization, the rebind scenario, which swaps the live Newton state, and the +CUDA-graph check build their own small scenes. Only one simulation context can be alive, so these tests come +first and fail if selected after a composite-scene test. -launch_test_simulation() +The composite scene runs on the CPU without CUDA-graph capture. On CUDA, one own-scene test covers the +ordered-state republish recorded into a captured graph. +""" + +from isaaclab_newton.physics import NewtonCfg + +from isaaclab.sim import SimulationCfg +from isaaclab.test.utils import DeviceScope, launch_test_simulation, test_devices + +launch_test_simulation(SimulationCfg(physics=NewtonCfg())) import logging import sys -from copy import copy, deepcopy -from pathlib import Path +from collections.abc import Callable, Iterator +from copy import copy +from dataclasses import dataclass from types import SimpleNamespace from unittest.mock import patch +import isaaclab_newton.physics.newton_manager as newton_manager_module import newton import numpy as np import pytest import torch import warp as wp from isaaclab_newton.assets import Articulation -from isaaclab_newton.assets.articulation.actuator_control import NewtonActuatorControl +from isaaclab_newton.assets.articulation import kernels as articulation_kernels from isaaclab_newton.assets.articulation.articulation_data import ArticulationData -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonBuilderCfg, NewtonCfg from isaaclab_newton.physics import NewtonManager as SimulationManager from isaaclab_physx.sim.schemas import PhysxJointCfg -from newton import JointTargetMode, JointType, ModelBuilder, ModelFlags, ShapeFlags +from newton import JointType, Model, ModelBuilder, ModelFlags, ShapeFlags, State from newton.selection import ArticulationView -from newton.solvers import SolverMuJoCo +from newton_test_utils import ( + NUM_ENVS, + WRIST_USD_STIFFNESS, + env_origins, + local_usd, + newton_sim_cfg, + spawn_assets, + world_gravity, +) from pxr import UsdPhysics import isaaclab.assets.articulation.ordering_resolvers as ordering_resolvers import isaaclab.sim as sim_utils import isaaclab.utils.math as math_utils -import isaaclab.utils.string as string_utils -from isaaclab.actuators import ( - IdealPDActuatorCfg, - ImplicitActuator, - ImplicitActuatorCfg, -) -from isaaclab.assets import ArticulationCfg, AssetBaseCfg +from isaaclab.actuators import ActuatorBaseCfg, IdealPDActuatorCfg, ImplicitActuatorCfg +from isaaclab.assets import ArticulationCfg from isaaclab.assets.articulation.ordering_resolvers import get_articulation_name_ordering -from isaaclab.cloner import CloneCfg, clone_plan_from_env_0, replicate from isaaclab.controllers import OperationalSpaceController, OperationalSpaceControllerCfg +from isaaclab.envs.mdp import randomize_physics_scene_gravity from isaaclab.envs.mdp.events import randomize_rigid_body_collider_offsets, randomize_rigid_body_material from isaaclab.managers import EventTermCfg, SceneEntityCfg -from isaaclab.sim import SimulationCfg, build_simulation_context -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR -from isaaclab.utils.math import compute_pose_error, matrix_from_quat, quat_inv, subtract_frame_transforms - -## -# Pre-defined configs -## -from isaaclab_assets import ANYMAL_C_CFG, FRANKA_PANDA_CFG, FRANKA_PANDA_HIGH_PD_CFG # isort:skip -from isaaclab_assets.robots.shadow_hand import SHADOW_HAND_NEWTON_CFG - -_FRANKA_PANDA_NEWTON_CFG = clone(FRANKA_PANDA_CFG) -_FRANKA_PANDA_NEWTON_CFG.spawn.variants = {"Physics": "mujoco", "Colliders": "gripper_only"} -_FRANKA_PANDA_HIGH_PD_NEWTON_CFG = clone(FRANKA_PANDA_HIGH_PD_CFG) -_FRANKA_PANDA_HIGH_PD_NEWTON_CFG.spawn.variants = {"Physics": "mujoco", "Colliders": "gripper_only"} - -SIM_CFGs = { - "humanoid": SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=80, - nconmax=25, - ls_iterations=20, - cone="pyramidal", - update_data_interval=2, - integrator="implicitfast", - impratio=1, - ), - num_substeps=2, - debug_mode=False, - ), - ), - "anymal": SimulationCfg( - dt=1 / 200, - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=70, - nconmax=70, - ls_iterations=40, - cone="elliptic", - impratio=100, - integrator="implicitfast", - ), - num_substeps=2, - debug_mode=True, - ), - ), - "panda": SimulationCfg( - dt=1 / 120, - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=20, - nconmax=20, - ls_iterations=20, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ), - ), - # "panda" with 4 solver substeps: at a single 1/120 substep, MJWarp implicitfast - # carries a ~0.4-1.0 mm integration limit cycle under a gravity load that never - # settles; 4 substeps integrate the same frozen per-control-step torques to a - # dead-still ~1 um hold. Used by the gravity-compensation precision test. - "panda_fine": SimulationCfg( - dt=1 / 120, - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=20, - nconmax=20, - ls_iterations=20, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=4, - debug_mode=False, - ), - ), - "single_joint_implicit": SimulationCfg( - dt=1 / 120, - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=20, - nconmax=20, - ls_iterations=20, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ), - ), - "single_joint_explicit": SimulationCfg( - dt=1 / 120, - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=20, - nconmax=20, - ls_iterations=20, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ), - ), - "shadow_hand": SimulationCfg( - dt=1 / 120, - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=70, - nconmax=70, - ls_iterations=40, - cone="elliptic", - impratio=100, - integrator="implicitfast", - ), - num_substeps=2, - debug_mode=True, - ), - ), -} - - -class CustomDrive(ImplicitActuator): - """Implicit actuator with a class name that does not encode its execution type.""" - - -def generate_articulation_cfg( - articulation_type: str, - stiffness: float | None = 10.0, - damping: float | None = 2.0, - actuator_velocity_limit: float | None = None, - actuator_effort_limit: float | None = None, - joint_velocity_limit: float | None = None, - joint_effort_limit: 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" or "spatial_tendon_test_asset". - 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. - actuator_velocity_limit: Velocity limit for the actuators. Only currently used for "single_joint_implicit" - and "single_joint_explicit". - actuator_effort_limit: Effort limit for explicit actuators. Only currently used for - "single_joint_explicit". - joint_velocity_limit: Velocity limit for the actuators (set into the simulation). - Only currently used for "single_joint_implicit" and "single_joint_explicit". - joint_effort_limit: 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_NEWTON_CFG - elif articulation_type == "anymal": - articulation_cfg = ANYMAL_C_CFG - elif articulation_type == "shadow_hand": - # The MuJoCo variant, not the PhysX one: only that variant authors the hand's tendons. - articulation_cfg = SHADOW_HAND_NEWTON_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.UsdPhysicsDriveCfg(max_force=80.0), PhysxJointCfg(max_joint_velocity=5.0)], - ), - actuators={ - "joint": ImplicitActuatorCfg( - joint_names_expr=[".*"], - joint_effort_limit=joint_effort_limit, - joint_velocity_limit=joint_velocity_limit, - actuator_velocity_limit=actuator_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.UsdPhysicsDriveCfg(max_force=80.0), PhysxJointCfg(max_joint_velocity=5.0)], - ), - actuators={ - "joint": IdealPDActuatorCfg( - joint_names_expr=[".*"], - joint_effort_limit=joint_effort_limit, - joint_velocity_limit=joint_velocity_limit, - actuator_effort_limit=actuator_effort_limit, - actuator_velocity_limit=actuator_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 fix_reversed_joints(stage): - """Fix reversed joints on the USD stage. - - Some USD assets have joints where physics:body0 is the child and physics:body1 is the parent, - which is the opposite of what Newton expects. This function detects reversed joints by building - a graph of body connections and identifying the root body (the one attached to world via a joint - with a missing body target). Any joint where body0 is closer to the root than body1 is swapped. - """ - from pxr import UsdPhysics - - # First pass: find root bodies (bodies with a joint that has only one target, i.e. attached to world) - root_bodies: set[str] = set() - joints_to_check = [] - for prim in stage.Traverse(): - if not prim.IsA(UsdPhysics.Joint): - continue - body0_targets = prim.GetRelationship("physics:body0").GetTargets() - body1_targets = prim.GetRelationship("physics:body1").GetTargets() - if body0_targets and not body1_targets: - root_bodies.add(str(body0_targets[0])) - elif body1_targets and not body0_targets: - root_bodies.add(str(body1_targets[0])) - elif body0_targets and body1_targets: - joints_to_check.append(prim) - - if not root_bodies: - return - - # Second pass: for each joint with two bodies, ensure body0 is the parent (closer to root) - for prim in joints_to_check: - body0_rel = prim.GetRelationship("physics:body0") - body1_rel = prim.GetRelationship("physics:body1") - body0_path = str(body0_rel.GetTargets()[0]) - body1_path = str(body1_rel.GetTargets()[0]) - - # Determine if we need to swap: body1 is root or ancestor → need to swap - body1_is_parent = body1_path in root_bodies or body0_path.startswith(body1_path + "/") - body0_is_parent = body0_path in root_bodies or body1_path.startswith(body0_path + "/") - - if body0_is_parent or not body1_is_parent: - continue # already correct or ambiguous - - # Swap body0 and body1 - body0_rel.SetTargets(body1_rel.GetTargets()) - body1_rel.SetTargets([body0_path]) - - # Swap local transforms - for attr_suffix in ("localPos", "localRot"): - attr0 = prim.GetAttribute(f"physics:{attr_suffix}0") - attr1 = prim.GetAttribute(f"physics:{attr_suffix}1") - val0, val1 = attr0.Get(), attr1.Get() - if val0 is not None and val1 is not None: - attr0.Set(val1) - attr1.Set(val0) - - -_REVERSED_JOINT_USD_FILES = {"revolute_articulation.usd"} -"""USD filenames with known reversed joint body0/body1 ordering.""" - - -_ANYMAL_C_BODY_NAMES = ( - "base", - "LF_HIP", - "LF_THIGH", - "LF_SHANK", - "LF_FOOT", - "LH_HIP", - "LH_THIGH", - "LH_SHANK", - "LH_FOOT", - "RF_HIP", - "RF_THIGH", - "RF_SHANK", - "RF_FOOT", - "RH_HIP", - "RH_THIGH", - "RH_SHANK", - "RH_FOOT", +from isaaclab.sim import SimulationContext, build_simulation_context +from isaaclab.test.utils.articulation_ordering import ( + BRANCHING_MJWARP_BODY_NAMES, + BRANCHING_MJWARP_JOINT_NAMES, + BRANCHING_PHYSX_BODY_NAMES, + BRANCHING_PHYSX_JOINT_NAMES, ) -_ANYMAL_C_ROOT_PRESERVING_REVERSED_BODY_NAMES = (_ANYMAL_C_BODY_NAMES[0], *reversed(_ANYMAL_C_BODY_NAMES[1:])) +from isaaclab.utils import replace +pytestmark = [pytest.mark.integration, pytest.mark.kitless] + +_GRAVITY = (0.0, 0.0, -9.81) _NEWTON_USER_ORDER_STATE_CACHES = ( "_joint_pos_user", @@ -394,484 +93,322 @@ def fix_reversed_joints(stage): "_body_com_vel_w_user", ) +_USD_LIMIT_PROPS = [sim_utils.UsdPhysicsDriveCfg(max_force=80.0), PhysxJointCfg(max_joint_velocity=5.0)] +"""Solver clamps authored on the limit islands' joints [N·m, rad/s].""" -def generate_articulation( - articulation_cfg: ArticulationCfg, num_articulations: int, device: str, add_ground_plane: bool = False -) -> 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. - add_ground_plane: Whether the simulation context authored a shared ground plane. - - Returns: - The articulation and environment translations. +_LEG_PUBLIC_JOINT_NAMES = ("LF_KFE", "RH_HAA", "LF_HAA", "RH_KFE", "LF_HFE", "RH_HFE") +_LEG_PUBLIC_BODY_NAMES = ("base", "RH_THIGH", "LF_SHANK", "RH_HIP", "LF_HIP", "RH_SHANK", "LF_THIGH") +"""Public orders that permute ``floating_two_leg.usda``'s backend joints and non-root bodies. - """ - # Generate translations of 2.5 m in x for each articulation - translations = np.zeros((num_articulations, 3), dtype=np.float32) - translations[:, 0] = np.arange(num_articulations) * 2.5 - - sim_utils.create_prim("/World/Env_0", "Xform", translation=translations[0]) - articulation_cfg = replace(articulation_cfg, prim_path="/World/Env_[^/]*/Robot") - cfgs = [articulation_cfg] - if add_ground_plane: - cfgs.append(AssetBaseCfg(prim_path="/World/defaultGroundPlane")) - clone_plan_from_env_0( - CloneCfg(clone_template="/World/Env_{}"), cfgs, num_articulations, 2.5, positions=translations - ) - articulation = Articulation(articulation_cfg) - - # Fix reversed joints for known-broken USD assets (body0/body1 swapped) - usd_path = getattr(articulation_cfg.spawn, "usd_path", "") - if any(name in usd_path for name in _REVERSED_JOINT_USD_FILES): - import omni.usd - - fix_reversed_joints(omni.usd.get_context().get_stage()) - - return articulation, torch.as_tensor(translations, device=device) - - -# --------------------------------------------------------------------------- -# Franka task-space tracking helpers (shared between the OSC tests). -# --------------------------------------------------------------------------- - - -def _setup_franka_at_home_pose(sim, *, zero_actuator_pd: bool = False, disable_gravity: bool = True): - """Build a Franka articulation at its configured home pose. - - Constructs :data:`_FRANKA_PANDA_HIGH_PD_NEWTON_CFG`, optionally zeroes the - arm-actuator PD gains, resets the simulator, and teleports the - arm joints to :attr:`default_joint_pos` (the env reset path that - normally does this is not invoked for standalone tests, so the - robot would otherwise sit at the URDF-neutral pose where the - Franka wrist is near-singular). - - Args: - sim: The simulation context to use. - zero_actuator_pd: If True, sets the ``panda_arm`` actuator stiffness - and damping to zero. Used by the OSC test - so OSC's joint-effort output is not opposed by the - implicit-PD's residual ``kp·(target − q)``. - disable_gravity: Per-body gravity flag written to the spawn config. - :data:`_FRANKA_PANDA_HIGH_PD_NEWTON_CFG` ships with gravity disabled; - pass False for tests where the arm must feel scene gravity. - - Returns: - Tuple of ``(robot, ee_frame_idx, ee_jacobi_idx, arm_joint_ids)``. - """ - cfg = replace(_FRANKA_PANDA_HIGH_PD_NEWTON_CFG, prim_path="/World/Env_[^/]*/Robot") - if zero_actuator_pd: - cfg.actuators["panda_arm"].stiffness = 0.0 - cfg.actuators["panda_arm"].damping = 0.0 - cfg.spawn.rigid_props.disable_gravity = disable_gravity - sim_utils.create_prim("/World/Env_0", "Xform", translation=(0.0, 0.0, 0.0)) - clone_plan_from_env_0(CloneCfg(clone_template="/World/Env_{}"), (cfg,), 1, 0.0) - robot = Articulation(cfg) - replicate(sim.get_clone_plan()) - sim.reset() - assert robot.is_initialized - - ee_frame_idx = robot.find_bodies("panda_hand")[0][0] - ee_jacobi_idx = ee_frame_idx - 1 - arm_joint_ids = robot.find_joints(["panda_joint.*"])[0] - - robot.write_joint_position_to_sim_index(position=robot.data.default_joint_pos.torch[:, :].clone()) - robot.write_joint_velocity_to_sim_index(velocity=robot.data.default_joint_vel.torch[:, :].clone()) - return robot, ee_frame_idx, ee_jacobi_idx, arm_joint_ids - - -def _compute_ee_pose_root(robot, ee_frame_idx): - """Return ``(ee_pos_b, ee_quat_b, root_pose_w)`` in the root frame.""" - ee_pose_w = robot.data.body_pose_w.torch[:, ee_frame_idx] - root_pose_w = robot.data.root_pose_w.torch - ee_pos_b, ee_quat_b = subtract_frame_transforms( - root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_pose_w[:, 0:3], ee_pose_w[:, 3:7] - ) - return ee_pos_b, ee_quat_b, root_pose_w - - -def _compute_jacobian_root_frame(robot, ee_jacobi_idx, arm_joint_ids): - """Return the EE Jacobian sliced to ``arm_joint_ids`` and rotated to the root frame.""" - jacobian = robot.data.body_link_jacobian_w.torch[:, ee_jacobi_idx, :, arm_joint_ids] - base_rot_matrix = matrix_from_quat(quat_inv(robot.data.root_pose_w.torch[:, 3:7])) - jacobian[:, :3, :] = torch.bmm(base_rot_matrix, jacobian[:, :3, :]) - jacobian[:, 3:, :] = torch.bmm(base_rot_matrix, jacobian[:, 3:, :]) - return jacobian +Neither permutation is its own inverse, so a map applied in the wrong direction shows. +""" -def _compute_ee_vel_root(jacobian_b, joint_vel): - """Return the EE 6D velocity in the root frame as ``J · q_dot``. +def _limit_actuators( + actuator_type: type[ImplicitActuatorCfg] | type[IdealPDActuatorCfg], +) -> dict[str, ActuatorBaseCfg]: + """Return one limit-resolution group per branching-fixture joint. - Required to make OSC's ``kd * ee_vel_b`` damping term meaningful. - Passing zero EE velocity (the convenient hack) leaves the impedance - undamped and the EE oscillates around the target. We use ``J · q_dot`` - rather than reading ``data.body_vel_w`` because Newton's lazy - velocity buffers can return stale/zero values until forced - materialization, while ``joint_vel`` and ``J`` are already pulled - by the loop. ``J`` correctness is pinned independently by - ``test_get_gravity_compensation_forces_matches_jacobian_gravity``. + Mirrors the single-joint implicit and explicit rows of the solver-clamp resolution contract: joint limits + only, actuator limits only, both, and neither. Implicit actuators take no separate rated effort limit. """ - return torch.bmm(jacobian_b, joint_vel.unsqueeze(-1)).squeeze(-1) - - -def _build_relative_pose_target(robot, ee_frame_idx, delta_xyz, device): - """Build a target pose = (current EE pose) + ``delta_xyz``, preserving orientation.""" - initial_ee_pos_b, initial_ee_quat_b, _ = _compute_ee_pose_root(robot, ee_frame_idx) - target_pos_b = initial_ee_pos_b + torch.tensor([list(delta_xyz)], device=device, dtype=initial_ee_pos_b.dtype) - return torch.cat([target_pos_b, initial_ee_quat_b], dim=-1) - - -def _tail_mean(history, tail: int = 200): - """Return the mean over the last ``tail`` samples.""" - tail_slice = history[-tail:] - return sum(tail_slice) / len(tail_slice) - - -@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") + if actuator_type is ImplicitActuatorCfg: + gains = {"stiffness": 10.0, "damping": 1.0} + actuator_limits = {"actuator_velocity_limit": 1e2} 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 - articulation_type = request.getfixturevalue("articulation_type") - sim_cfg = deepcopy(SIM_CFGs[articulation_type]) - sim_cfg.device = device - if "use_newton_actuators" in request.fixturenames: - sim_cfg.use_newton_actuators = request.getfixturevalue("use_newton_actuators") - # ``gravity_enabled`` is silently ignored by ``build_simulation_context`` - # when an explicit ``sim_cfg`` is also passed; apply it here so the - # fixture honors what its parameter advertises. - if not gravity_enabled: - sim_cfg.gravity = (0.0, 0.0, 0.0) - with build_simulation_context( - device=device, - auto_add_lighting=True, - gravity_enabled=gravity_enabled, - add_ground_plane=add_ground_plane, - sim_cfg=sim_cfg, - ) as sim: - sim._app_control_on_stop_handle = None - yield sim - - -def _make_target_mode_builder( - monkeypatch, joint_names: list[str], modes: list[JointTargetMode], stiffness: list[float], damping: list[float] -) -> ModelBuilder: - """Build a zero-gain articulated model builder for target-mode tests.""" - sim = object.__new__(sim_utils.SimulationContext) - sim.cfg = SimpleNamespace(physics=NewtonCfg()) - sim._backend_registry = [] - monkeypatch.setattr(sim_utils.SimulationContext, "instance", lambda: sim) - builder = sim.get_or_create_backend(NewtonBuilderCfg(physics_cfg=sim.cfg.physics)) - inertia = wp.mat33(1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0) - parent = -1 - joint_ids = [] - for joint_name in joint_names: - link = builder.add_link(mass=1.0, inertia=inertia, label=f"/World/Env_0/Robot/{joint_name}_link") - joint_ids.append( - builder.add_joint_revolute( - parent, - link, - target_ke=0.0, - target_kd=0.0, - label=f"/World/Env_0/Robot/{joint_name}", - ) - ) - parent = link - builder.add_articulation(joint_ids, label="/World/Env_0/Robot") - builder.articulation_label = ["/World/Env_0/Robot"] - builder.joint_target_mode = [int(mode) for mode in modes] - builder.joint_target_ke = stiffness - builder.joint_target_kd = damping - return builder - - -@pytest.mark.parametrize( - ("actuator_cfg", "expected_native_groups"), - [ - (ImplicitActuatorCfg(joint_names_expr=["joint"], stiffness=10.0, damping=1.0), set()), - (IdealPDActuatorCfg(joint_names_expr=["joint"], stiffness=None, damping=None), {"explicit"}), - ], - ids=["implicit", "explicit"], -) -def test_prepare_native_actuators_activates_only_explicit_groups(monkeypatch, actuator_cfg, expected_native_groups): - """Keep implicit-only articulations on the solver-drive path and leave solver gains untouched. - - Explicit groups activate the Newton-actuator path without writing gains; collection construction resolves - the actuator defaults later. - """ - activation_calls = [] - gain_writes = [] - articulation = SimpleNamespace( - _sim_cfg=SimpleNamespace(use_newton_actuators=True), - device="cpu", - find_joints=lambda _: ([0], ["joint"]), - write_joint_stiffness_to_sim_index=lambda **_: gain_writes.append("stiffness"), - write_joint_damping_to_sim_index=lambda **_: gain_writes.append("damping"), - ) - monkeypatch.setattr(SimulationManager, "activate_newton_actuator_path", lambda: activation_calls.append(True)) - - control = NewtonActuatorControl(articulation) - group_name = "explicit" if expected_native_groups else "implicit" - native_groups = control.prepare_native_actuators(collection=None, actuator_cfgs={group_name: actuator_cfg}) + gains = {"stiffness": 0.0, "damping": 0.1} + actuator_limits = {"actuator_velocity_limit": 1e2, "actuator_effort_limit": 1e2} + joint_limits = {"joint_velocity_limit": 1e5, "joint_effort_limit": 1e5} + return { + "joint_limit": actuator_type(joint_names_expr=["left_shoulder"], **joint_limits, **gains), + "actuator_limit": actuator_type(joint_names_expr=["left_elbow"], **actuator_limits, **gains), + "both_limits": actuator_type(joint_names_expr=["right_shoulder"], **joint_limits, **actuator_limits, **gains), + "no_limit": actuator_type(joint_names_expr=["right_elbow"], **gains), + } - assert native_groups == expected_native_groups - assert gain_writes == [] - if expected_native_groups: - assert activation_calls == [True] - else: - assert not control.native_actuator_path_active - assert not articulation._has_newton_actuators - assert activation_calls == [] +def _island_cfgs() -> dict[str, ArticulationCfg]: + """Return the composite scene's islands, each at its own offset inside every environment.""" + + def island( + name: str, + spawn: sim_utils.UsdFileCfg, + y: float, + actuators: dict[str, ActuatorBaseCfg], + init_state: ArticulationCfg.InitialStateCfg | None = None, + **kwargs, + ) -> ArticulationCfg: + """Return the island ``name`` at lateral offset ``y`` [m] inside every environment.""" + init_state = replace(init_state or ArticulationCfg.InitialStateCfg(), pos=(0.0, y, 1.0)) + return ArticulationCfg( + prim_path=f"/World/Env_[^/]*/{name}", spawn=spawn, init_state=init_state, actuators=actuators, **kwargs + ) -@pytest.mark.parametrize( - ("actuator_cfg", "expected_mode", "expected_actuator_indices"), - [ - ( - ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=10.0, damping=0.0), - JointTargetMode.POSITION, - [0, 1], + return { + # The floating island comes first: its root owns the first body and joint coordinates of each world. + "floating": island( + "Floating", + local_usd("floating_two_link.usda"), + 0.0, + {"joint": ImplicitActuatorCfg(joint_names_expr=["Joint"], stiffness=20.0, damping=2.0)}, + articulation_root_prim_path="/Root", ), - ( - ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=0.0, damping=2.0), - JointTargetMode.VELOCITY, - [-2, -3], + "fixed": island( + "Fixed", + local_usd("fixed_spatial_chain.usda"), + 2.0, + { + "slides": ImplicitActuatorCfg( + joint_names_expr=["Joint_[0-2]"], stiffness=2000.0, damping=100.0, viscous_friction=0.25 + ), + "wrist": ImplicitActuatorCfg(joint_names_expr=["Joint_[3-5]"], stiffness=None, damping=None), + }, ), - ( - ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=10.0, damping=2.0), - JointTargetMode.POSITION_VELOCITY, - [0, -2, 1, -3], + "pendulum": island( + "Pendulum", + local_usd("revolute_pendulum.usda"), + 6.0, + {"joint": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=0.0, damping=0.0)}, ), - ( - ImplicitActuatorCfg( - class_type=f"{__name__}:CustomDrive", joint_names_expr=[".*"], stiffness=10.0, damping=2.0 - ), - JointTargetMode.POSITION_VELOCITY, - [0, -2, 1, -3], + "ordered": island( + "Ordered", + local_usd("floating_two_leg.usda"), + 8.0, + {"legs": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=10.0, damping=1.0)}, + init_state=ArticulationCfg.InitialStateCfg(joint_pos={".*HAA": 0.2, ".*HFE": 0.3, ".*KFE": -0.6}), + joint_ordering=_LEG_PUBLIC_JOINT_NAMES, + body_ordering=_LEG_PUBLIC_BODY_NAMES, ), - ( - ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=0.0, damping=0.0), - JointTargetMode.EFFORT, - None, + "branching": island( + "Branching", + local_usd("articulation_ordering_branching.usda"), + 20.0, + {"arms": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=10.0, damping=1.0)}, + joint_ordering="physx", + body_ordering="physx", ), - ( - IdealPDActuatorCfg(joint_names_expr=[".*"], stiffness=10.0, damping=2.0), - JointTargetMode.EFFORT, - None, + "tendon": island( + "Hand", + local_usd("fixed_tendon_hand.usda"), + 10.0, + { + # the tendons drive these joints: gains stay as authored + "tendon_joints": ImplicitActuatorCfg( + joint_names_expr=[".*J(1|2)"], stiffness=None, damping=None, armature=1e-3 + ) + }, ), - ], -) -def test_actuator_cfg_sets_newton_target_mode_before_solver_init( - monkeypatch, actuator_cfg, expected_mode, expected_actuator_indices -): - """Resolve configured modes before finalization constructs MuJoCo actuators.""" - articulation_cfg = ArticulationCfg( - prim_path="/World/Env_[^/]*/Robot", - articulation_root_prim_path="", - actuators={"joint": actuator_cfg}, - ) - builder = _make_target_mode_builder( - monkeypatch, ["left_joint", "right_joint"], [JointTargetMode.NONE, JointTargetMode.NONE], [0.0, 0.0], [0.0, 0.0] - ) - Articulation._configure_joint_target_modes(SimpleNamespace(cfg=articulation_cfg), None) - model = builder.finalize(device="cpu") - solver = SolverMuJoCo(model, use_mujoco_cpu=True) - assert model.joint_target_mode.numpy().tolist() == [int(expected_mode), int(expected_mode)] - assert ( - solver.mjc_actuator_to_newton_idx.numpy().tolist() if solver.mjc_actuator_to_newton_idx is not None else None - ) == expected_actuator_indices + "root_fixed": island( + "RootFixed", + local_usd("floating_two_link.usda", fix_root_link=True), + 12.0, + {"joint": ImplicitActuatorCfg(joint_names_expr=["Joint"], stiffness=20.0, damping=2.0)}, + ), + "root_freed": island( + "RootFreed", + local_usd("fixed_spatial_chain.usda", fix_root_link=False), + 14.0, + {"all": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=1.0, damping=0.1)}, + ), + "limits_implicit": island( + "LimitsImplicit", + local_usd("articulation_ordering_branching.usda", joint_drive_props=_USD_LIMIT_PROPS), + 16.0, + _limit_actuators(ImplicitActuatorCfg), + ), + "limits_explicit": island( + "LimitsExplicit", + local_usd("articulation_ordering_branching.usda", joint_drive_props=_USD_LIMIT_PROPS), + 18.0, + _limit_actuators(IdealPDActuatorCfg), + ), + "osc": island( + "Osc", + local_usd("fixed_spatial_chain.usda"), + 22.0, + # unpowered, so the operational-space controller's efforts are the only actuation + {"arm": ImplicitActuatorCfg(joint_names_expr=["Joint_.*"], stiffness=0.0, damping=0.0)}, + init_state=ArticulationCfg.InitialStateCfg(joint_pos={"Joint_3": 0.3, "Joint_4": -0.2, "Joint_5": 0.4}), + ), + } -def test_actuator_cfg_matches_explicit_descendant_articulation_root(monkeypatch): - """Match target modes against an explicitly configured descendant articulation root.""" - articulation_cfg = ArticulationCfg( - prim_path="/World/Env_[^/]*/Robot", - articulation_root_prim_path="/base", - actuators={"joint": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=10.0, damping=0.0)}, - ) - builder = _make_target_mode_builder(monkeypatch, ["joint"], [JointTargetMode.NONE], [0.0], [0.0]) - builder.articulation_label = ["/World/Env_0/Robot/base"] - Articulation._configure_joint_target_modes(SimpleNamespace(cfg=articulation_cfg), None) - assert builder.joint_target_mode == [int(JointTargetMode.POSITION)] +@dataclass +class _Scene: + """Articulations that share one real Newton model and solver lifecycle.""" + sim: SimulationContext + articulations: dict[str, Articulation] + cfgs: dict[str, ArticulationCfg] + origins: torch.Tensor + device: str -def test_actuator_cfg_matches_clone_plan_root_expr(monkeypatch): - """Match builder labels against the clone slot spelling clone-plan root resolution returns.""" - articulation_cfg = ArticulationCfg( - prim_path="{ENV_REGEX_NS}/Robot", - actuators={"joint": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=10.0, damping=0.0)}, - ) - monkeypatch.setattr( - "isaaclab_newton.assets.articulation.articulation.resolve_matching_prims_from_source", - lambda *_args, **_kwargs: [(None, "/World/envs/env_[^/]+/Robot/base")], - ) - builder = _make_target_mode_builder(monkeypatch, ["joint"], [JointTargetMode.NONE], [0.0], [0.0]) - builder.articulation_label = ["/World/envs/env_0/Robot/base"] - Articulation._configure_joint_target_modes(SimpleNamespace(cfg=articulation_cfg), None) - assert builder.joint_target_mode == [int(JointTargetMode.POSITION)] + def step(self, *names: str, num_steps: int = 1) -> None: + """Write the named islands' commands, step every island, and update the named islands.""" + for _ in range(num_steps): + for name in names: + self.articulations[name].write_data_to_sim() + self.sim.step() + for name in names: + self.articulations[name].update(self.sim.cfg.dt) + + def default_root_pose_w(self, articulation: Articulation) -> torch.Tensor: + """Return the configured root pose placed at each environment origin.""" + root_pose = articulation.data.default_root_pose.torch.clone() + root_pose[:, :3] += self.origins + return root_pose + def rest(self, articulation: Articulation) -> None: + """Put an island at rest at its default configuration and hold it there. -@pytest.mark.parametrize("joint_type", [JointType.FREE, JointType.FIXED]) -def test_actuator_cfg_leaves_excluded_joint_types_imported(monkeypatch, joint_type): - """Leave target modes for free and fixed joints unchanged.""" - articulation_cfg = ArticulationCfg( - prim_path="/World/Env_[^/]*/Robot", - articulation_root_prim_path="", - actuators={"joint": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=10.0, damping=0.0)}, - ) - builder = _make_target_mode_builder(monkeypatch, ["joint"], [JointTargetMode.NONE], [0.0], [0.0]) - builder.joint_type[0] = joint_type - Articulation._configure_joint_target_modes(SimpleNamespace(cfg=articulation_cfg), None) - assert builder.joint_target_mode == [int(JointTargetMode.NONE)] + The targets reach the solver immediately, so the island stays at rest while other tests step the scene. + """ + if not articulation.is_fixed_base: + articulation.write_root_pose_to_sim_index(root_pose=self.default_root_pose_w(articulation)) + articulation.write_root_velocity_to_sim_index( + root_velocity=torch.zeros_like(articulation.data.default_root_vel.torch) + ) + default_joint_pos = articulation.data.default_joint_pos.torch.clone() + articulation.write_joint_state_to_sim_index( + position=default_joint_pos, velocity=torch.zeros_like(default_joint_pos) + ) + commands = articulation.actuators.target_command + commands.set_position_index(value=default_joint_pos) + commands.set_velocity_index(value=torch.zeros_like(default_joint_pos)) + commands.set_effort_index(value=torch.zeros_like(default_joint_pos)) + articulation.reset() + articulation.write_data_to_sim() -def test_actuator_cfg_uses_imported_gain_for_none_stiffness(monkeypatch): - """Retain the imported stiffness when an implicit actuator config leaves it unset.""" - articulation_cfg = ArticulationCfg( - prim_path="/World/Env_[^/]*/Robot", - articulation_root_prim_path="", - actuators={"joint": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=None, damping=0.0)}, - ) - builder = _make_target_mode_builder(monkeypatch, ["joint"], [JointTargetMode.EFFORT], [10.0], [0.0]) - Articulation._configure_joint_target_modes(SimpleNamespace(cfg=articulation_cfg), None) - assert builder.joint_target_mode == [int(JointTargetMode.POSITION)] +## +# Kernels and model-level accessors. These build no simulation context. +## -def test_actuator_cfg_leaves_unconfigured_newton_target_modes_imported(monkeypatch): - """Leave target modes for DOFs outside an actuator group unchanged.""" - subset_cfg = ArticulationCfg( - prim_path="/World/Env_[^/]*/Robot", - articulation_root_prim_path="", - actuators={ - "shoulder": ImplicitActuatorCfg(joint_names_expr=["left_shoulder"], stiffness=10.0, damping=0.0), - }, - ) - builder = _make_target_mode_builder( - monkeypatch, - ["left_shoulder", "right_shoulder"], - [JointTargetMode.NONE, JointTargetMode.VELOCITY], - [0.0, 0.0], - [0.0, 2.0], +def _selector(values: list[int], dtype: type) -> wp.array: + """Create a CPU Warp selector with the requested integer width.""" + return wp.array(values, dtype=dtype, device="cpu") + + +@pytest.mark.parametrize(("env_dtype", "joint_dtype"), [(wp.int32, wp.int32), (wp.int64, wp.int64)]) +def test_write_joint_limit_data_to_user_and_backend_index_accepts_index_dtypes( + env_dtype: type, joint_dtype: type +) -> None: + """Write partial user-order joint limits into user and backend-order buffers.""" + limits_np = np.asarray([[[1.0, 3.0], [2.0, 5.0]], [[-1.0, 1.0], [4.0, 8.0]]], dtype=np.float32) + limits = wp.array(limits_np, dtype=wp.vec2f, device="cpu") + env_ids = _selector([0, 1], env_dtype) + user_ids = _selector([2, 0], joint_dtype) + user_to_backend = wp.array(np.asarray([1, 2, 0], dtype=np.int32), dtype=wp.int32, device="cpu") + user_lower = wp.zeros((2, 3), dtype=wp.float32, device="cpu") + user_upper = wp.zeros((2, 3), dtype=wp.float32, device="cpu") + user_limits = wp.zeros((2, 3), dtype=wp.vec2f, device="cpu") + backend_lower = wp.zeros((2, 3), dtype=wp.float32, device="cpu") + backend_upper = wp.zeros((2, 3), dtype=wp.float32, device="cpu") + soft_limits = wp.zeros((2, 3), dtype=wp.vec2f, device="cpu") + default_pos = wp.array( + np.asarray([[0.0, 0.0, 4.0], [0.0, 0.0, 0.0]], dtype=np.float32), dtype=wp.float32, device="cpu" + ) + clamped_defaults = wp.zeros(1, dtype=wp.int32, device="cpu") + kernel = articulation_kernels.write_joint_limit_data_to_user_and_backend_index + if env_dtype != wp.int32 or joint_dtype != wp.int32: + kernel = articulation_kernels.write_joint_limit_data_to_user_and_backend_index_kernel(env_ids, user_ids) + + wp.launch( + kernel, + dim=limits.shape, + inputs=[limits, 1.0, env_ids, user_ids, user_to_backend, True], + outputs=[ + user_lower, + user_upper, + user_limits, + backend_lower, + backend_upper, + soft_limits, + default_pos, + clamped_defaults, + ], + device="cpu", ) - Articulation._configure_joint_target_modes(SimpleNamespace(cfg=subset_cfg), None) - assert builder.joint_target_mode == [int(JointTargetMode.POSITION), int(JointTargetMode.VELOCITY)] - -@pytest.mark.parametrize( - ("stiffness", "damping", "expected_modes"), - [ - ({"left_joint": 10.0}, 0.0, [JointTargetMode.POSITION, JointTargetMode.EFFORT]), - ({"left_joint": 10.0}, {"right_joint": 2.0}, [JointTargetMode.POSITION, JointTargetMode.VELOCITY]), - ], -) -def test_actuator_cfg_aligns_partial_dictionary_gains_by_joint_name(monkeypatch, stiffness, damping, expected_modes): - """Resolve sparse stiffness and damping dictionaries independently by joint name.""" - articulation_cfg = ArticulationCfg( - prim_path="/World/Env_[^/]*/Robot", - articulation_root_prim_path="", - actuators={"joint": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=stiffness, damping=damping)}, - ) - builder = _make_target_mode_builder( - monkeypatch, ["left_joint", "right_joint"], [JointTargetMode.NONE, JointTargetMode.NONE], [0.0, 0.0], [0.0, 0.0] + np.testing.assert_allclose(user_lower.numpy(), np.asarray([[2.0, 0.0, 1.0], [4.0, 0.0, -1.0]], dtype=np.float32)) + np.testing.assert_allclose(user_upper.numpy(), np.asarray([[5.0, 0.0, 3.0], [8.0, 0.0, 1.0]], dtype=np.float32)) + np.testing.assert_allclose(backend_lower.numpy(), np.asarray([[1.0, 2.0, 0.0], [-1.0, 4.0, 0.0]], dtype=np.float32)) + np.testing.assert_allclose(backend_upper.numpy(), np.asarray([[3.0, 5.0, 0.0], [1.0, 8.0, 0.0]], dtype=np.float32)) + np.testing.assert_allclose( + user_limits.numpy(), + np.asarray([[[2.0, 5.0], [0.0, 0.0], [1.0, 3.0]], [[4.0, 8.0], [0.0, 0.0], [-1.0, 1.0]]], dtype=np.float32), + ) + np.testing.assert_allclose(default_pos.numpy(), np.asarray([[2.0, 0.0, 3.0], [4.0, 0.0, 0.0]], dtype=np.float32)) + assert clamped_defaults.numpy()[0] == 3 + + +def test_write_joint_limit_data_to_user_and_backend_mask_reorders_backend_buffers() -> None: + """Write masked user-order joint limits into user and backend-order buffers.""" + limits_np = np.asarray( + [[[1.0, 2.0], [3.0, 6.0], [4.0, 9.0]], [[-2.0, 2.0], [5.0, 7.0], [8.0, 10.0]]], dtype=np.float32 + ) + limits = wp.array(limits_np, dtype=wp.vec2f, device="cpu") + env_mask = wp.array(np.asarray([True, False], dtype=bool), dtype=wp.bool, device="cpu") + user_mask = wp.array(np.asarray([False, True, True], dtype=bool), dtype=wp.bool, device="cpu") + user_to_backend = wp.array(np.asarray([1, 2, 0], dtype=np.int32), dtype=wp.int32, device="cpu") + user_lower = wp.zeros((2, 3), dtype=wp.float32, device="cpu") + user_upper = wp.zeros((2, 3), dtype=wp.float32, device="cpu") + user_limits = wp.zeros((2, 3), dtype=wp.vec2f, device="cpu") + backend_lower = wp.zeros((2, 3), dtype=wp.float32, device="cpu") + backend_upper = wp.zeros((2, 3), dtype=wp.float32, device="cpu") + soft_limits = wp.zeros((2, 3), dtype=wp.vec2f, device="cpu") + default_pos = wp.zeros((2, 3), dtype=wp.float32, device="cpu") + clamped_defaults = wp.zeros(1, dtype=wp.int32, device="cpu") + + wp.launch( + articulation_kernels.write_joint_limit_data_to_user_and_backend_mask, + dim=limits.shape, + inputs=[limits, 1.0, env_mask, user_mask, user_to_backend, True], + outputs=[ + user_lower, + user_upper, + user_limits, + backend_lower, + backend_upper, + soft_limits, + default_pos, + clamped_defaults, + ], + device="cpu", ) - Articulation._configure_joint_target_modes(SimpleNamespace(cfg=articulation_cfg), None) - assert builder.joint_target_mode == [int(mode) for mode in expected_modes] - - -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("gravity_enabled", [False]) -@pytest.mark.parametrize("articulation_type", ["single_joint_explicit"]) -def test_branching_fixture_physx_ordering_reorders_newton_to_bfs(sim, device, gravity_enabled, articulation_type): - """Resolve the documented Newton ``joint_ordering="physx"`` sim-to-sim workflow on a branching asset. - - Mirrors :func:`isaaclab_physx.test.assets.test_articulation.test_branching_fixture_resolves_distinct_conventions` - with the backend roles swapped: here the live backend is Newton (depth-first, so its native view is - the MJWarp order), and the request is ``physx``/``body_ordering="physx"``. Cross-backend discovery must - resolve the breadth-first PhysX order and reorder the public joint/body axes to it. This is the headline - workflow documented in - ``docs/source/overview/core-concepts/physical-backends/joint_and_body_ordering.rst``. - - The branching fixture is shared between both backends; a copy lives in this package's - test data directory so the two backends assert against the same ground-truth asset. - The same articulation also checks the MJWarp-order emulation used for cross-backend discovery and that - selected inertial-property writes keep Newton's inverse arrays current under the body ordering. - """ - fixture_path = Path(__file__).parent / "data" / "articulation_ordering_branching.usda" - sim_utils.create_prim("/World/Env_0", "Xform") - articulation_cfg = ArticulationCfg( - prim_path="/World/Env_0/Robot", - spawn=sim_utils.UsdFileCfg(usd_path=str(fixture_path)), - actuators={}, - joint_ordering="physx", - body_ordering="physx", + np.testing.assert_allclose(user_lower.numpy(), np.asarray([[0.0, 3.0, 4.0], [0.0, 0.0, 0.0]], dtype=np.float32)) + np.testing.assert_allclose(user_upper.numpy(), np.asarray([[0.0, 6.0, 9.0], [0.0, 0.0, 0.0]], dtype=np.float32)) + np.testing.assert_allclose(backend_lower.numpy(), np.asarray([[4.0, 0.0, 3.0], [0.0, 0.0, 0.0]], dtype=np.float32)) + np.testing.assert_allclose(backend_upper.numpy(), np.asarray([[9.0, 0.0, 6.0], [0.0, 0.0, 0.0]], dtype=np.float32)) + np.testing.assert_allclose( + user_limits.numpy(), + np.asarray([[[0.0, 0.0], [3.0, 6.0], [4.0, 9.0]], [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]]], dtype=np.float32), ) - clone_plan_from_env_0(CloneCfg(clone_template="/World/Env_{}"), (articulation_cfg,), 1, 0.0) - articulation = Articulation(articulation_cfg) - - replicate(sim.get_clone_plan()) - sim.reset() - assert articulation.is_initialized + np.testing.assert_allclose(default_pos.numpy(), np.asarray([[0.0, 3.0, 4.0], [0.0, 0.0, 0.0]], dtype=np.float32)) + assert clamped_defaults.numpy()[0] == 2 - # Newton's native traversal is depth-first, so the live backend view already reflects MJWarp order. - assert tuple(articulation.backend_joint_names) == BRANCHING_MJWARP_JOINT_NAMES - assert tuple(articulation.backend_body_names) == BRANCHING_MJWARP_BODY_NAMES - - # The same-backend "mjwarp" request takes an identity fast path, so also force the cross-backend - # emulation (the temporary Newton USD builder a PhysX-backed articulation would use) and compare its - # independently rebuilt view against the live backend view. BFS and DFS differ on this branching fixture. - emulated_names = ordering_resolvers._get_mjwarp_names_from_newton_usd_builder(articulation) - assert emulated_names is not None - assert emulated_names["joint"] == tuple(articulation.backend_joint_names) - assert emulated_names["body"] == tuple(articulation.backend_body_names) - assert get_articulation_name_ordering(articulation, "mjwarp", kind="joint") == tuple( - articulation.backend_joint_names - ) - assert get_articulation_name_ordering(articulation, "mjwarp", kind="body") == tuple(articulation.backend_body_names) - # Cross-backend discovery (bypassing the same-backend fast path) resolves the breadth-first PhysX order. - assert get_articulation_name_ordering(articulation, "physx", kind="joint") == BRANCHING_PHYSX_JOINT_NAMES - assert get_articulation_name_ordering(articulation, "physx", kind="body") == BRANCHING_PHYSX_BODY_NAMES +@pytest.mark.parametrize("index_dtype", [wp.int32, wp.int64]) +def test_scatter_reset_masks_from_ids_accepts_index_dtype(index_dtype: type) -> None: + """Set exact world and articulation reset masks from nonidentity environment IDs.""" - # The requested PhysX ordering reorders the public joint/body axes to the BFS convention. - assert tuple(articulation.joint_names) == BRANCHING_PHYSX_JOINT_NAMES - assert tuple(articulation.body_names) == BRANCHING_PHYSX_BODY_NAMES - assert articulation.joint_ordering is not None - assert articulation.body_ordering is not None + env_ids = _selector([2, 0], index_dtype) + articulation_ids = wp.array(np.asarray([[0, 1], [2, 3], [4, 5]], dtype=np.int32), dtype=int, device="cpu") + world_mask = wp.zeros(3, dtype=wp.bool, device="cpu") + fk_mask = wp.zeros(6, dtype=wp.bool, device="cpu") - # Selected mass and inertia writes with int64 selectors update Newton's inverse arrays in backend order. - env_ids = torch.tensor([0], dtype=torch.int64, device=device) - body_ids = torch.tensor([2, articulation.num_bodies - 1], dtype=torch.int64, device=device) - backend_body_ids = torch.tensor( - [articulation.data.body_ordering.user_to_backend_indices[index] for index in body_ids.tolist()], - dtype=torch.int64, - device=device, + wp.launch( + newton_manager_module._scatter_reset_masks_from_ids, + dim=(env_ids.shape[0], articulation_ids.shape[1]), + inputs=[env_ids, articulation_ids], + outputs=[world_mask, fk_mask], + device="cpu", ) - assert backend_body_ids[0] != body_ids[0] - model = SimulationManager.get_model() - - masses = articulation.data.body_mass.torch[env_ids][:, body_ids].clone() + torch.tensor([[1.0, 2.0]], device=device) - articulation.set_masses_index(masses=masses, env_ids=env_ids, body_ids=body_ids) - model_inv_mass = wp.to_torch(articulation.root_view.get_attribute("body_inv_mass", model)[:, 0]) - torch.testing.assert_close(model_inv_mass[env_ids][:, backend_body_ids], masses.reciprocal()) - inertia_matrices = torch.diag_embed(torch.tensor([[[2.0, 3.0, 4.0], [5.0, 6.0, 7.0]]], device=device)) - articulation.set_inertias_index(inertias=inertia_matrices.reshape(1, 2, 9), env_ids=env_ids, body_ids=body_ids) - model_inv_inertia = wp.to_torch(articulation.root_view.get_attribute("body_inv_inertia", model)[:, 0]) - torch.testing.assert_close(model_inv_inertia[env_ids][:, backend_body_ids], torch.linalg.inv(inertia_matrices)) + np.testing.assert_array_equal(world_mask.numpy(), np.asarray([True, False, True])) + np.testing.assert_array_equal(fk_mask.numpy(), np.asarray([True, True, False, False, True, True])) def test_num_shapes_per_body_follows_public_body_order() -> None: @@ -894,138 +431,235 @@ class _ShapeCountSurface: @pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("articulation_type", ["anymal"]) # consumed by the sim fixture -@pytest.mark.parametrize("use_newton_actuators", [True]) # consumed by the sim fixture -def test_newton_native_actuator_gain_write_maps_public_joint_subset_to_backend( - sim, articulation_type, use_newton_actuators, device -): - """Map selected public joint IDs to Newton-controller columns.""" - articulation_cfg = replace( - generate_articulation_cfg("anymal"), - actuators={ - "legs": IdealPDActuatorCfg( - joint_names_expr=[".*HAA", ".*HFE", ".*KFE"], stiffness=40.0, damping=5.0, actuator_effort_limit=80.0 - ) - }, - joint_ordering=tuple(reversed(ANYMAL_C_PHYSX_JOINT_NAMES)), - ) - articulation, _ = generate_articulation(articulation_cfg, 2, device=sim.device) - replicate(sim.get_clone_plan()) - sim.reset() - assert articulation.joint_ordering is not None - assert articulation.newton_actuator_adapter is not None - - def gather_stiffness() -> torch.Tensor: - stiffness = torch.zeros( - (articulation.num_instances, articulation.num_joints), - device=articulation.device, - ) - for actuator in articulation.newton_actuator_adapter.actuators: - if hasattr(actuator.controller, "kp"): - stiffness += wp.to_torch( - articulation.root_view.get_actuator_parameter(actuator, actuator.controller, "kp") - ) - return stiffness - - stiffness_before = gather_stiffness() - env_ids = torch.tensor([1], device=articulation.device, dtype=torch.long) - joint_ids = torch.tensor([1, 6, 10], device=articulation.device, dtype=torch.long) - stiffness = torch.tensor([[101.0, 106.0, 110.0]], device=articulation.device) - - with pytest.warns(DeprecationWarning, match="write_actuator_stiffness_to_sim"): - articulation.write_actuator_stiffness_to_sim( - stiffness=stiffness, - env_ids=env_ids, - joint_ids=joint_ids, - ) - - backend_joint_ids = torch.tensor( - articulation.joint_ordering.user_to_backend_indices, - device=articulation.device, - dtype=torch.long, - )[joint_ids] - expected_stiffness = stiffness_before.clone() - expected_stiffness[env_ids.unsqueeze(1), backend_joint_ids.unsqueeze(0)] = stiffness - torch.testing.assert_close(gather_stiffness(), expected_stiffness) - - -@pytest.mark.parametrize("num_articulations", [1]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("gravity_enabled", [False]) -@pytest.mark.parametrize("articulation_type", ["anymal"]) -def test_newton_ordered_body_state_cache_invalidates_on_same_timestamp_root_write( - sim, num_articulations, device, gravity_enabled, articulation_type -): - """Refresh ordered body pose and velocity after root writes at the current simulation timestamp.""" - articulation_cfg = replace( - generate_articulation_cfg(articulation_type=articulation_type), - body_ordering=_ANYMAL_C_ROOT_PRESERVING_REVERSED_BODY_NAMES, +@pytest.mark.parametrize( + "first_property", + ["body_com_jacobian_w", "body_link_jacobian_w", "mass_matrix", "gravity_compensation_forces"], +) +def test_task_space_allocation_and_capture(monkeypatch, device, first_property): + """Allocate only requested outputs and retain them across changed-state graph replay.""" + builder = ModelBuilder() + builder.begin_world() + base = builder.add_link(mass=4.0, inertia=wp.mat33(np.eye(3)), label="Robot/base") + arm = builder.add_link(mass=2.0, inertia=wp.mat33(np.eye(3)), com=(0.5, 0.0, 0.0), label="Robot/arm") + slider = builder.add_link(mass=3.0, inertia=wp.mat33(np.eye(3)), com=(0.25, 0.0, 0.0), label="Robot/slider") + builder.add_articulation( + [ + builder.add_joint_fixed(-1, base), + builder.add_joint_revolute(base, arm, axis=(0.0, 1.0, 0.0), label="hinge"), + builder.add_joint_prismatic(arm, slider, axis=(1.0, 0.0, 0.0), label="slide"), + ], + label="Robot", ) - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) - - replicate(sim.get_clone_plan()) - sim.reset() - sim.step() - articulation.update(sim.cfg.dt) + builder.end_world() + model = builder.finalize(device=device) + state, control = model.state(), model.control() + monkeypatch.setattr(SimulationManager, "get_model", lambda: model) + monkeypatch.setattr(SimulationManager, "get_state_0", lambda: state) + monkeypatch.setattr(SimulationManager, "get_control", lambda: control) + view = ArticulationView(model, "Robot", exclude_joint_types=[JointType.FIXED]) + eager = ArticulationData(view, device) + eager._apply_ordering_maps_after_resolve() + newton.eval_fk(model, state.joint_q, state.joint_qd, state) + getattr(eager, first_property) # Compile kernels before first-use capture on a fresh container. + data = ArticulationData(view, device) + data._apply_ordering_maps_after_resolve() + properties = ("body_com_jacobian_w", "body_link_jacobian_w", "mass_matrix", "gravity_compensation_forces") + assert all(getattr(data, f"_{name}_ta") is None for name in properties) + assert data._jacobian_buf_flat is data._mass_matrix_full_buf is data._gravity_force_full_buf is None + with wp.ScopedCapture(device=device) as capture: + output = getattr(data, first_property) + required = {first_property} + if first_property == "body_link_jacobian_w": + required.add("body_com_jacobian_w") + for name in properties: + assert (getattr(data, f"_{name}_ta") is not None) == (name in required) + assert (data._jacobian_buf_flat is not None) == (first_property != "gravity_compensation_forces") + assert (data._mass_matrix_full_buf is not None) == (first_property == "mass_matrix") + assert (data._gravity_force_full_buf is not None) == (first_property == "gravity_compensation_forces") - data = articulation.data - assert data.body_ordering is not None - root_body_idx = articulation.find_bodies("base")[0][0] - sim_timestamp = data._sim_timestamp + for angle, displacement in ((0.0, 0.0), (0.7, 0.4), (-0.3, 0.2)): + state.joint_q.assign(np.asarray([angle, displacement], dtype=np.float32)) + newton.eval_fk(model, state.joint_q, state.joint_qd, state) + expected = getattr(eager, first_property).warp.numpy() + wp.capture_launch(capture.graph) + np.testing.assert_allclose(output.warp.numpy(), expected, atol=1e-5) + data._create_simulation_bindings() + data._apply_ordering_maps_after_resolve() + assert getattr(data, first_property) is output - cached_body_pose = data.body_link_pose_w.torch[:, root_body_idx].clone() - written_root_pose = data.root_link_pose_w.torch.clone() - written_root_pose[:, 0] += 0.25 - articulation.write_root_link_pose_to_sim_index(root_pose=written_root_pose) - assert data._sim_timestamp == sim_timestamp - torch.testing.assert_close(data.root_link_pose_w.torch, written_root_pose) - SimulationManager.forward() # Another consumer may resolve shared FK before this view reads. - refreshed_body_pose = data.body_link_pose_w.torch[:, root_body_idx] - torch.testing.assert_close(refreshed_body_pose, written_root_pose) - assert not torch.equal(refreshed_body_pose, cached_body_pose) - # Populate the velocity cache after the pose write so only the velocity write can invalidate it. - cached_body_vel = data.body_com_vel_w.torch[:, root_body_idx].clone() - written_root_vel = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]], device=device, dtype=cached_body_vel.dtype) - articulation.write_root_com_velocity_to_sim_index(root_velocity=written_root_vel) - assert data._sim_timestamp == sim_timestamp - torch.testing.assert_close(data.root_com_vel_w.torch, written_root_vel) - refreshed_body_vel = data.body_com_vel_w.torch[:, root_body_idx] - torch.testing.assert_close(refreshed_body_vel, written_root_vel) - assert not torch.equal(refreshed_body_vel, cached_body_vel) +## +# Own scenes. These tests run before the composite scene exists: only one simulation context can be alive. +## -@pytest.mark.parametrize("num_articulations", [1]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("gravity_enabled", [False]) -@pytest.mark.parametrize("articulation_type", ["panda"]) -@pytest.mark.parametrize("ordering_mode", ["none", "reversed"]) -@pytest.mark.parametrize("use_newton_actuators", [False]) -def test_newton_ordered_state_caches_invalidate_on_rebind( - sim, num_articulations, device, gravity_enabled, articulation_type, ordering_mode, use_newton_actuators -): - """Rebind public state to recreated Newton arrays and invalidate ordered caches.""" - articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) - if ordering_mode == "reversed": - articulation_cfg = replace( - articulation_cfg, - joint_ordering=tuple(reversed(PANDA_JOINT_NAMES)), - body_ordering=PANDA_ROOT_PRESERVING_REVERSED_BODY_NAMES, - ) - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) +@pytest.mark.parametrize("invalid_field", ["joint_pos", "joint_vel", "articulation_root_prim_path"]) +def test_invalid_articulation_cfg_fails_initialization(device: str, invalid_field: str) -> None: + """Initialization fails when a default joint state exceeds the joint limits or the explicit root does not exist.""" + articulation_cfg = ArticulationCfg( + prim_path="/World/Env_[^/]*/Robot", + spawn=local_usd("fixed_spatial_chain.usda", joint_drive_props=[PhysxJointCfg(max_joint_velocity=5.0)]), + actuators={"all": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=100.0, damping=10.0)}, + ) + if invalid_field == "joint_pos": + articulation_cfg.init_state.joint_pos = {"Joint_0": 10.0, "Joint_[35]": -20.0} + error, match = ValueError, "default positions out of the limits" + elif invalid_field == "joint_vel": + articulation_cfg.init_state.joint_vel = {"Joint_0": 100.0, "Joint_[35]": -60.0} + error, match = ValueError, "default velocities out of the limits" + else: + articulation_cfg.articulation_root_prim_path = "/non_existing_prim_path" + error, match = KeyError, "No articulations matching pattern" + with build_simulation_context(sim_cfg=newton_sim_cfg(device)) as sim: + articulation = spawn_assets({"robot": articulation_cfg}, num_envs=1)["robot"] + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + with pytest.raises(error, match=match): + sim.reset() + + +@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) +def test_newton_rebind_refreshes_ordered_state_and_preserves_lab_owned_actuator_gains(device: str) -> None: + """Rebind public state to recreated Newton arrays, invalidate ordered caches, and keep Lab-owned gains. + + A full sim reset recreates the solver's state and model arrays; the scenario shallow-copies the live state and + model, swaps in sentinel-filled arrays, and rebinds every island's data. - replicate(sim.get_clone_plan()) - sim.reset() - assert articulation.is_initialized - has_ordering = ordering_mode == "reversed" - assert (articulation.data.joint_ordering is not None) is has_ordering - assert (articulation.data.body_ordering is not None) is has_ordering + * Implicit islands, floating with and without non-identity ordering and fixed with it: every binding, public + proxy, component view, and actuator kernel input follows the new arrays, and the acceleration caches restart + from the new velocities. + * Explicit islands: named actuator groups own their actuator kp/kd while the solver's sim gains are zeroed for + explicit DOFs. Rebind must NOT resync the actuator-owned values from the freshly rebuilt (sentinel) solver + gains, while the sim-owned mirrors track them. The identity-ordering island is the control that passes with + or without the fix. - sim.step() - articulation.update(sim.cfg.dt) + Finally ``_clear_callbacks`` deregisters exactly the articulation's post-step hook, without touching other + callbacks. + """ + leg_cfg = local_usd("floating_two_leg.usda") + explicit_legs = IdealPDActuatorCfg(joint_names_expr=[".*"], stiffness=40.0, damping=5.0, actuator_effort_limit=80.0) + implicit_legs = ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=40.0, damping=5.0) + orderings = {"joint_ordering": _LEG_PUBLIC_JOINT_NAMES, "body_ordering": _LEG_PUBLIC_BODY_NAMES} + cfgs = { + f"{actuation}_{ordering}": ArticulationCfg( + prim_path=f"/World/Env_[^/]*/Robot_{actuation}_{ordering}", + spawn=leg_cfg, + init_state=ArticulationCfg.InitialStateCfg(pos=(0.0, 2.0 * index, 1.0)), + actuators={"legs": implicit_legs if actuation == "implicit" else explicit_legs}, + **(orderings if ordering == "reordered" else {}), + ) + for index, (actuation, ordering) in enumerate( + (actuation, ordering) for actuation in ("implicit", "explicit") for ordering in ("none", "reordered") + ) + } + # a fixed base publishes no root velocity binding of its own + cfgs["implicit_fixed_reordered"] = ArticulationCfg( + prim_path="/World/Env_[^/]*/Robot_implicit_fixed_reordered", + spawn=local_usd("fixed_spatial_chain.usda"), + init_state=ArticulationCfg.InitialStateCfg(pos=(0.0, 8.0, 1.0)), + actuators={"all": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=40.0, damping=5.0)}, + joint_ordering=("Joint_2", "Joint_5", "Joint_0", "Joint_4", "Joint_1", "Joint_3"), + body_ordering=("Root", "Link_3", "Link_0", "Link_5", "Link_1", "Link_4", "Link_2"), + ) + with build_simulation_context(sim_cfg=newton_sim_cfg(device, use_newton_actuators=False)) as sim: + articulations = spawn_assets(cfgs) + sim.reset() + sim.step() + for articulation in articulations.values(): + assert articulation.is_initialized + articulation.update(sim.cfg.dt) + implicit_checks = [ + _prime_ordered_state_rebind(articulations[name], sim.cfg.dt) + for name in ("implicit_none", "implicit_reordered", "implicit_fixed_reordered") + ] + for name in ("explicit_none", "explicit_reordered"): + # Prime: explicit (IdealPD) actuators keep their PD in actuator-owned records, + # while the solver's sim gains are zeroed so it applies no PD on these DOFs. + articulation = articulations[name] + data = articulation.data + assert (data.joint_ordering is not None) is name.endswith("reordered") + np.testing.assert_allclose(articulation.actuators["legs"].stiffness.cpu().numpy(), 40.0) + np.testing.assert_allclose(articulation.actuators["legs"].damping.cpu().numpy(), 5.0) + np.testing.assert_allclose(data._sim_bind_joint_stiffness_sim.numpy(), 0.0) + np.testing.assert_allclose(data._sim_bind_joint_damping_sim.numpy(), 0.0) + + # Simulate a full sim reset: shallow-copy the state and the model, swap their arrays for sentinel-filled + # arrays (standing in for whatever the solver rebuilds), and rebind the data-side sim bindings. + old_state = SimulationManager.get_state_0() + old_model = SimulationManager.get_model() + new_state = copy(old_state) + new_model = copy(old_model) + + joint_q_values = np.arange(len(old_state.joint_q), dtype=np.float32) + 1000.0 + joint_qd_values = np.arange(len(old_state.joint_qd), dtype=np.float32) + 2000.0 + body_indices = np.arange(len(old_state.body_q), dtype=np.float32)[:, None] + body_q_values = np.zeros((len(old_state.body_q), 7), dtype=np.float32) + body_q_values[:, :3] = 3000.0 + 10.0 * body_indices + np.arange(3, dtype=np.float32) + body_q_values[:, 6] = 1.0 + body_qd_values = 4000.0 + 10.0 * body_indices + np.arange(6, dtype=np.float32) + limit_indices = np.arange(len(old_model.joint_limit_lower), dtype=np.float32) + sentinel_ke = 12345.0 + sentinel_kd = 678.0 + device = old_state.joint_q.device + new_state.joint_q = wp.array(joint_q_values, dtype=wp.float32, device=device) + new_state.joint_qd = wp.array(joint_qd_values, dtype=wp.float32, device=device) + new_state.body_q = wp.array(body_q_values, dtype=wp.transformf, device=device) + new_state.body_qd = wp.array(body_qd_values, dtype=wp.spatial_vectorf, device=device) + new_model.joint_limit_lower = wp.array(-5000.0 - limit_indices, dtype=wp.float32, device=device) + new_model.joint_limit_upper = wp.array(5000.0 + limit_indices, dtype=wp.float32, device=device) + num_dofs = len(old_model.joint_target_ke) + new_model.joint_target_ke = wp.array(np.full(num_dofs, sentinel_ke, dtype=np.float32), device=device) + new_model.joint_target_kd = wp.array(np.full(num_dofs, sentinel_kd, dtype=np.float32), device=device) + SimulationManager.backend.state_0 = new_state + SimulationManager.backend.model = new_model + + for check_rebind in implicit_checks: + check_rebind(new_state, new_model) + + for name in ("explicit_none", "explicit_reordered"): + articulation = articulations[name] + data = articulation.data + data._create_simulation_bindings() + # The actuator-owned gains must survive the rebind unchanged... + np.testing.assert_allclose(articulation.actuators["legs"].stiffness.cpu().numpy(), 40.0) + np.testing.assert_allclose(articulation.actuators["legs"].damping.cpu().numpy(), 5.0) + # ...while the sim-owned mirrors track the solver's freshly seeded (sentinel) gains. + if data.joint_ordering is not None: + np.testing.assert_allclose(data._joint_stiffness_user.numpy(), sentinel_ke) + np.testing.assert_allclose(data._joint_damping_user.numpy(), sentinel_kd) + else: + np.testing.assert_allclose(data._sim_bind_joint_stiffness_sim.numpy(), sentinel_ke) + np.testing.assert_allclose(data._sim_bind_joint_damping_sim.numpy(), sentinel_kd) + + # ``_clear_callbacks`` must deregister exactly this hook so it does not leak on the class-level list, + # and leave an unrelated callback (standing in for another articulation's hook) untouched. + articulation = articulations["implicit_reordered"] + registered_callback = articulation._post_step_callback + assert registered_callback is not None + assert registered_callback in SimulationManager._post_step_callbacks + + def _other_callback() -> None: + return None + + SimulationManager.register_post_step_callback(_other_callback) + articulation._clear_callbacks() + + assert articulation._post_step_callback is None + assert registered_callback not in SimulationManager._post_step_callbacks + assert _other_callback in SimulationManager._post_step_callbacks + SimulationManager.unregister_post_step_callback(_other_callback) + + +def _prime_ordered_state_rebind(articulation: Articulation, dt: float) -> Callable[[State, Model], None]: + """Prime an implicit island's ordered state caches and return the check to run after the arrays are swapped.""" data = articulation.data previous_body_com_vel = data._previous_body_com_vel.numpy().copy() + has_ordering = data.joint_ordering is not None + assert (data.body_ordering is not None) is has_ordering primed_joint_vel_values = ( np.arange(np.prod(data._sim_bind_joint_vel.shape), dtype=np.float32).reshape(data._sim_bind_joint_vel.shape) + 100.0 @@ -1046,22 +680,23 @@ def test_newton_ordered_state_caches_invalidate_on_rebind( # ``joint_vel`` shadow) observes the primed backend state. No-op under identity # ordering, where the getters alias the sim-bound arrays directly. data._refresh_user_order_state() - # One update spanning two physics steps, as when Newton owns the decimation loop. - update_dt = 2 * sim.cfg.dt + # Newton decimation may collect several physics steps into one data update. + update_dt = 2 * dt data.update(update_dt) primed_joint_acc = data.joint_acc.warp.numpy().copy() primed_body_com_acc_w = data.body_com_acc_w.warp.numpy().copy() assert data._joint_acc.timestamp == data._sim_timestamp assert data._body_com_acc_w.timestamp == data._sim_timestamp assert np.any(primed_joint_acc != 0.0) - primed_body_user_to_backend = ( + assert np.any(primed_body_com_acc_w != 0.0) + body_user_to_backend = ( np.asarray(articulation.body_ordering.user_to_backend_indices) if articulation.body_ordering is not None else np.arange(articulation.num_bodies) ) np.testing.assert_allclose( primed_body_com_acc_w, - ((primed_body_com_vel_values - previous_body_com_vel) / update_dt)[:, primed_body_user_to_backend], + ((primed_body_com_vel_values - previous_body_com_vel) / update_dt)[:, body_user_to_backend], rtol=1e-5, ) @@ -1093,467 +728,486 @@ def test_newton_ordered_state_caches_invalidate_on_rebind( data.joint_pos_limits.torch.clone() # The Tier-1 state shadows are plain wp.arrays (no timestamp): they are # allocated for non-identity ordering and stay ``None`` for identity ordering. - if has_ordering: - for cache_name in _NEWTON_USER_ORDER_STATE_CACHES: - assert getattr(data, cache_name) is not None - else: - for cache_name in _NEWTON_USER_ORDER_STATE_CACHES: - assert getattr(data, cache_name) is None - - old_state = SimulationManager.get_state_0() - old_model = SimulationManager.get_model() - new_state = copy(old_state) - new_model = copy(old_model) - - joint_q_values = np.arange(len(old_state.joint_q), dtype=np.float32) + 1000.0 - joint_qd_values = np.arange(len(old_state.joint_qd), dtype=np.float32) + 2000.0 - body_indices = np.arange(len(old_state.body_q), dtype=np.float32)[:, None] - body_q_values = np.zeros((len(old_state.body_q), 7), dtype=np.float32) - body_q_values[:, :3] = 3000.0 + 10.0 * body_indices + np.arange(3, dtype=np.float32) - body_q_values[:, 6] = 1.0 - body_qd_values = 4000.0 + 10.0 * body_indices + np.arange(6, dtype=np.float32) - limit_indices = np.arange(len(old_model.joint_limit_lower), dtype=np.float32) - limit_lower_values = -5000.0 - limit_indices - limit_upper_values = 5000.0 + limit_indices - - new_state.joint_q = wp.array(joint_q_values, dtype=wp.float32, device=old_state.joint_q.device) - new_state.joint_qd = wp.array(joint_qd_values, dtype=wp.float32, device=old_state.joint_qd.device) - new_state.body_q = wp.array(body_q_values, dtype=wp.transformf, device=old_state.body_q.device) - new_state.body_qd = wp.array(body_qd_values, dtype=wp.spatial_vectorf, device=old_state.body_qd.device) - new_model.joint_limit_lower = wp.array( - limit_lower_values, dtype=wp.float32, device=old_model.joint_limit_lower.device - ) - new_model.joint_limit_upper = wp.array( - limit_upper_values, dtype=wp.float32, device=old_model.joint_limit_upper.device - ) - SimulationManager.backend.state_0 = new_state - SimulationManager.backend.model = new_model - - body_velocities = articulation.root_view.get_link_velocities(new_state) - assert body_velocities is not None - new_source_bindings = { - "_sim_bind_joint_pos": articulation.root_view.get_dof_positions(new_state)[:, 0], - "_sim_bind_joint_vel": articulation.root_view.get_dof_velocities(new_state)[:, 0], - "_sim_bind_body_link_pose_w": articulation.root_view.get_link_transforms(new_state)[:, 0], - "_sim_bind_body_com_vel_w": body_velocities[:, 0], - "_sim_bind_joint_pos_limits_lower": articulation.root_view.get_attribute("joint_limit_lower", new_model)[:, 0], - "_sim_bind_joint_pos_limits_upper": articulation.root_view.get_attribute("joint_limit_upper", new_model)[:, 0], - } - for binding_name, new_source in new_source_bindings.items(): - if binding_name in old_binding_ptrs: - assert int(new_source.ptr) != old_binding_ptrs[binding_name] - - data._create_simulation_bindings() - - for binding_name, old_binding in old_bindings.items(): - rebound = getattr(data, binding_name) - assert rebound is not old_binding - assert int(rebound.ptr) != old_binding_ptrs[binding_name] - assert int(rebound.ptr) == int(new_source_bindings[binding_name].ptr) - - for inputs in actuator_state_inputs: - assert inputs[3].ptr == data.joint_pos.warp.ptr - assert inputs[4].ptr == data.joint_vel.warp.ptr - - assert data._joint_acc.timestamp == -1.0 - assert data._body_com_acc_w.timestamp == -1.0 + for cache_name in _NEWTON_USER_ORDER_STATE_CACHES: + assert (getattr(data, cache_name) is not None) is has_ordering + + def check_rebind(new_state: State, new_model: Model) -> None: + body_velocities = articulation.root_view.get_link_velocities(new_state) + assert body_velocities is not None + new_source_bindings = { + "_sim_bind_joint_pos": articulation.root_view.get_dof_positions(new_state)[:, 0], + "_sim_bind_joint_vel": articulation.root_view.get_dof_velocities(new_state)[:, 0], + "_sim_bind_body_link_pose_w": articulation.root_view.get_link_transforms(new_state)[:, 0], + "_sim_bind_body_com_vel_w": body_velocities[:, 0], + "_sim_bind_joint_pos_limits_lower": articulation.root_view.get_attribute("joint_limit_lower", new_model)[ + :, 0 + ], + "_sim_bind_joint_pos_limits_upper": articulation.root_view.get_attribute("joint_limit_upper", new_model)[ + :, 0 + ], + } + for binding_name, new_source in new_source_bindings.items(): + if binding_name in old_binding_ptrs: + assert int(new_source.ptr) != old_binding_ptrs[binding_name] + + data._create_simulation_bindings() + + for binding_name, old_binding in old_bindings.items(): + rebound = getattr(data, binding_name) + assert rebound is not old_binding + assert int(rebound.ptr) != old_binding_ptrs[binding_name] + assert int(rebound.ptr) == int(new_source_bindings[binding_name].ptr) + + for inputs in actuator_state_inputs: + assert inputs[3].ptr == data.joint_pos.warp.ptr + assert inputs[4].ptr == data.joint_vel.warp.ptr + + assert data._joint_acc.timestamp == -1.0 + assert data._body_com_acc_w.timestamp == -1.0 + + joint_user_to_backend = ( + np.asarray(articulation.joint_ordering.user_to_backend_indices) + if articulation.joint_ordering is not None + else np.arange(articulation.num_joints) + ) + body_user_to_backend = ( + np.asarray(articulation.body_ordering.user_to_backend_indices) + if articulation.body_ordering is not None + else np.arange(articulation.num_bodies) + ) + expected_previous_joint_vel = new_source_bindings["_sim_bind_joint_vel"].numpy()[:, joint_user_to_backend] + expected_previous_body_com_vel = new_source_bindings["_sim_bind_body_com_vel_w"].numpy() + np.testing.assert_array_equal(data._previous_joint_vel.numpy(), expected_previous_joint_vel) + np.testing.assert_array_equal(data._previous_body_com_vel.numpy(), expected_previous_body_com_vel) + + joint_acc = data.joint_acc.warp.numpy() + body_com_acc_w = data.body_com_acc_w.warp.numpy() + np.testing.assert_array_equal(joint_acc, np.zeros_like(expected_previous_joint_vel)) + np.testing.assert_array_equal( + body_com_acc_w, np.zeros_like(expected_previous_body_com_vel[:, body_user_to_backend]) + ) + assert data._joint_acc.timestamp == data._sim_timestamp + assert data._body_com_acc_w.timestamp == data._sim_timestamp + assert not np.array_equal(joint_acc, primed_joint_acc) + assert not np.array_equal(body_com_acc_w, primed_body_com_acc_w) + + expected_public = { + "joint_pos": new_source_bindings["_sim_bind_joint_pos"].numpy()[:, joint_user_to_backend], + "joint_vel": new_source_bindings["_sim_bind_joint_vel"].numpy()[:, joint_user_to_backend], + "body_link_pose_w": new_source_bindings["_sim_bind_body_link_pose_w"].numpy()[:, body_user_to_backend], + "body_com_vel_w": new_source_bindings["_sim_bind_body_com_vel_w"].numpy()[:, body_user_to_backend], + } + for property_name, expected in expected_public.items(): + proxy = getattr(data, property_name) + assert proxy is not old_public_proxies[property_name] + binding_name = public_to_binding[property_name] + assert int(proxy.warp.ptr) != old_binding_ptrs[binding_name] + if has_ordering: + shadow = getattr(data, public_to_shadow[property_name]) + assert int(proxy.warp.ptr) == int(shadow.ptr) + else: + assert int(proxy.warp.ptr) == int(getattr(data, binding_name).ptr) + np.testing.assert_array_equal(proxy.warp.numpy(), expected) + + for name, (parent, selection) in component_sources.items(): + component = getattr(data, name) + expected = getattr(data, parent).torch[..., selection] + assert component is not old_components[name] + assert component.torch.data_ptr() == expected.data_ptr() + assert component.torch.stride() == expected.stride() + torch.testing.assert_close(component.torch, expected) + + expected_limits = np.stack( + ( + new_source_bindings["_sim_bind_joint_pos_limits_lower"].numpy()[:, joint_user_to_backend], + new_source_bindings["_sim_bind_joint_pos_limits_upper"].numpy()[:, joint_user_to_backend], + ), + axis=-1, + ) + np.testing.assert_array_equal(data.joint_pos_limits.warp.numpy(), expected_limits) - joint_user_to_backend = ( - np.asarray(articulation.joint_ordering.user_to_backend_indices) - if articulation.joint_ordering is not None - else np.arange(articulation.num_joints) - ) - body_user_to_backend = ( - np.asarray(articulation.body_ordering.user_to_backend_indices) - if articulation.body_ordering is not None - else np.arange(articulation.num_bodies) - ) - expected_previous_joint_vel = new_source_bindings["_sim_bind_joint_vel"].numpy()[:, joint_user_to_backend] - expected_previous_body_com_vel = new_source_bindings["_sim_bind_body_com_vel_w"].numpy() - np.testing.assert_array_equal(data._previous_joint_vel.numpy(), expected_previous_joint_vel) - np.testing.assert_array_equal(data._previous_body_com_vel.numpy(), expected_previous_body_com_vel) - - joint_acc = data.joint_acc.warp.numpy() - body_com_acc_w = data.body_com_acc_w.warp.numpy() - np.testing.assert_array_equal(joint_acc, np.zeros_like(expected_previous_joint_vel)) - np.testing.assert_array_equal( - body_com_acc_w, np.zeros_like(expected_previous_body_com_vel[:, body_user_to_backend]) - ) - assert data._joint_acc.timestamp == data._sim_timestamp - assert data._body_com_acc_w.timestamp == data._sim_timestamp - assert not np.array_equal(joint_acc, primed_joint_acc) - assert not np.array_equal(body_com_acc_w, primed_body_com_acc_w) - - expected_public = { - "joint_pos": new_source_bindings["_sim_bind_joint_pos"].numpy()[:, joint_user_to_backend], - "joint_vel": new_source_bindings["_sim_bind_joint_vel"].numpy()[:, joint_user_to_backend], - "body_link_pose_w": new_source_bindings["_sim_bind_body_link_pose_w"].numpy()[:, body_user_to_backend], - "body_com_vel_w": new_source_bindings["_sim_bind_body_com_vel_w"].numpy()[:, body_user_to_backend], - } - for property_name, expected in expected_public.items(): - proxy = getattr(data, property_name) - assert proxy is not old_public_proxies[property_name] - binding_name = public_to_binding[property_name] - assert int(proxy.warp.ptr) != old_binding_ptrs[binding_name] - if has_ordering: - shadow = getattr(data, public_to_shadow[property_name]) - assert int(proxy.warp.ptr) == int(shadow.ptr) - else: - assert int(proxy.warp.ptr) == int(getattr(data, binding_name).ptr) - np.testing.assert_array_equal(proxy.warp.numpy(), expected) - - for name, (parent, selection) in component_sources.items(): - component = getattr(data, name) - expected = getattr(data, parent).torch[..., selection] - assert component is not old_components[name] - assert component.torch.data_ptr() == expected.data_ptr() - assert component.torch.stride() == expected.stride() - torch.testing.assert_close(component.torch, expected) - - expected_limits = np.stack( - ( - new_source_bindings["_sim_bind_joint_pos_limits_lower"].numpy()[:, joint_user_to_backend], - new_source_bindings["_sim_bind_joint_pos_limits_upper"].numpy()[:, joint_user_to_backend], - ), - axis=-1, - ) - np.testing.assert_array_equal(data.joint_pos_limits.warp.numpy(), expected_limits) + return check_rebind -@pytest.mark.parametrize("num_articulations", [1]) @pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("gravity_enabled", [True]) -@pytest.mark.parametrize("articulation_type", ["anymal"]) -@pytest.mark.parametrize("ordering_mode", ["none", "reversed"]) -@pytest.mark.parametrize("use_newton_actuators", [False]) -def test_newton_rebind_preserves_lab_owned_actuator_gains( - sim, num_articulations, device, gravity_enabled, articulation_type, ordering_mode, use_newton_actuators -): - """Keep Lab-owned actuator gains across a rebind that re-seeds the solver's sim gains. - - Part 2 (D3) regression: named actuator groups own their actuator kp/kd; the solver's - sim gains are deliberately zeroed for explicit DOFs. A full sim reset recreates the solver arrays. - Rebind must NOT resync the actuator-owned values from freshly rebuilt (here: sentinel) solver gains. - ``none`` is the identity-ordering control that must pass with or without the fix. - """ - articulation_cfg = replace( - generate_articulation_cfg(articulation_type=articulation_type), - actuators={ - "legs": IdealPDActuatorCfg( - joint_names_expr=[".*HAA", ".*HFE", ".*KFE"], stiffness=40.0, damping=5.0, actuator_effort_limit=80.0 - ) - }, - ) - if ordering_mode == "reversed": - articulation_cfg = replace(articulation_cfg, joint_ordering=tuple(reversed(ANYMAL_C_PHYSX_JOINT_NAMES))) - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) - replicate(sim.get_clone_plan()) - sim.reset() - assert articulation.is_initialized - - has_ordering = ordering_mode == "reversed" - data = articulation.data - assert (data.joint_ordering is not None) is has_ordering - - # Prime: explicit (IdealPD) actuators keep their PD in actuator-owned records, - # while the solver's sim gains are zeroed so it applies no PD on these DOFs. - np.testing.assert_allclose(articulation.actuators["legs"].stiffness.cpu().numpy(), 40.0) - np.testing.assert_allclose(articulation.actuators["legs"].damping.cpu().numpy(), 5.0) - np.testing.assert_allclose(data._sim_bind_joint_stiffness_sim.numpy(), 0.0) - np.testing.assert_allclose(data._sim_bind_joint_damping_sim.numpy(), 0.0) - - # Simulate a full sim reset: shallow-copy the model, swap the joint gain arrays - # for sentinel-filled arrays (standing in for whatever the solver rebuilds), and - # rebind the data-side sim bindings. - old_model = SimulationManager.get_model() - new_model = copy(old_model) - sentinel_ke = 12345.0 - sentinel_kd = 678.0 - new_model.joint_target_ke = wp.array( - np.full(len(old_model.joint_target_ke), sentinel_ke, dtype=np.float32), - dtype=wp.float32, - device=old_model.joint_target_ke.device, - ) - new_model.joint_target_kd = wp.array( - np.full(len(old_model.joint_target_kd), sentinel_kd, dtype=np.float32), - dtype=wp.float32, - device=old_model.joint_target_kd.device, - ) - SimulationManager.backend.model = new_model - data._create_simulation_bindings() - - # The actuator-owned gains must survive the rebind unchanged... - np.testing.assert_allclose(articulation.actuators["legs"].stiffness.cpu().numpy(), 40.0) - np.testing.assert_allclose(articulation.actuators["legs"].damping.cpu().numpy(), 5.0) - # ...while the sim-owned mirrors track the solver's freshly seeded (sentinel) gains. - if has_ordering: - np.testing.assert_allclose(data._joint_stiffness_user.numpy(), sentinel_ke) - np.testing.assert_allclose(data._joint_damping_user.numpy(), sentinel_kd) - else: - np.testing.assert_allclose(data._sim_bind_joint_stiffness_sim.numpy(), sentinel_ke) - np.testing.assert_allclose(data._sim_bind_joint_damping_sim.numpy(), sentinel_kd) - +def test_newton_ordered_state_publishes_inside_captured_cuda_graph(device: str) -> None: + """Republish the user-order state shadows from inside a captured CUDA graph. -@pytest.mark.parametrize("num_articulations", [1]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("gravity_enabled", [True]) -@pytest.mark.parametrize("articulation_type", ["anymal"]) -def test_newton_post_step_hook_publishes_ordered_state_and_deregisters( - sim, num_articulations, device, gravity_enabled, articulation_type -): - """Republish the user-order Tier-1 shadows inside the sim step, without any read. - - Part 1 (D1) regression: with non-identity ordering the passthrough state getters no - longer reorder on read; the post-step callback republishes the shadows from live - backend state inside the stepped region. We deliberately clobber the four shadows, - step the simulation WITHOUT reading any state property, and assert the shadows again - equal the reordered backend state -- which can only hold if the hook ran inside the - step. Under the lazy design (no hook), only a property read would refresh them, so - the clobbered shadows would stay stale and the assertions would fail. - - Ships the eager-mode invariant variant: CUDA-graph capture is not reliably reachable - from this CPU test harness, and this invariant directly proves the in-step republish. - - Finally checks that ``_clear_callbacks`` deregisters the hook without touching other callbacks. + With CUDA-graph capture a step replays the recorded launches, so the post-step reorder of the Tier-1 shadows + must be recorded into the graph rather than run from Python. The shadows are clobbered after the graph exists + and must equal the reordered backend state after one replayed step, without a state property read in between. """ - articulation_cfg = replace( - generate_articulation_cfg(articulation_type=articulation_type), + articulation_cfg = ArticulationCfg( + prim_path="/World/Env_[^/]*/Robot", + spawn=local_usd("floating_two_leg.usda"), + init_state=ArticulationCfg.InitialStateCfg(pos=(0.0, 0.0, 1.0)), actuators={"legs": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=40.0, damping=5.0)}, - joint_ordering=tuple(reversed(ANYMAL_C_PHYSX_JOINT_NAMES)), - body_ordering=_ANYMAL_C_ROOT_PRESERVING_REVERSED_BODY_NAMES, + joint_ordering=_LEG_PUBLIC_JOINT_NAMES, + body_ordering=_LEG_PUBLIC_BODY_NAMES, ) - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) - replicate(sim.get_clone_plan()) - sim.reset() - assert articulation.is_initialized - - data = articulation.data - assert data.joint_ordering is not None - assert data.body_ordering is not None - joint_u2b = np.asarray(articulation.joint_ordering.user_to_backend_indices) - body_u2b = np.asarray(articulation.body_ordering.user_to_backend_indices) - - # Clobber every Tier-1 shadow with a large sentinel so a stale (unrepublished) shadow - # is unmistakably detectable -- the true backend joint/velocity state sits near zero, - # so a plain zero-fill would coincide with it. Then step WITHOUT touching joint_pos / - # joint_vel / body_link_pose_w / body_com_vel_w. - data._joint_pos_user.fill_(1000.0) - data._joint_vel_user.fill_(1000.0) - data._body_link_pose_w_user.fill_(wp.transformf(1000.0, 1000.0, 1000.0, 0.0, 0.0, 0.0, 1.0)) - data._body_com_vel_w_user.fill_(wp.spatial_vectorf(1000.0, 1000.0, 1000.0, 1000.0, 1000.0, 1000.0)) - sim.step() + sim_cfg = newton_sim_cfg(device, use_newton_actuators=False, use_cuda_graph=True) + with build_simulation_context(sim_cfg=sim_cfg) as sim: + articulation = spawn_assets({"robot": articulation_cfg})["robot"] + sim.reset() + # the first step captures the graph that later steps replay + sim.step() + assert SimulationManager._graph is not None + data = articulation.data + joint_u2b = np.asarray(articulation.joint_ordering.user_to_backend_indices) + body_u2b = np.asarray(articulation.body_ordering.user_to_backend_indices) + + # a sentinel far from the resting backend state, so a stale shadow cannot pass for a republished one + data._joint_pos_user.fill_(1000.0) + data._joint_vel_user.fill_(1000.0) + data._body_link_pose_w_user.fill_(wp.transformf(1000.0, 1000.0, 1000.0, 0.0, 0.0, 0.0, 1.0)) + data._body_com_vel_w_user.fill_(wp.spatial_vectorf(1000.0, 1000.0, 1000.0, 1000.0, 1000.0, 1000.0)) + sim.step() - np.testing.assert_allclose(data._joint_pos_user.numpy(), data._sim_bind_joint_pos.numpy()[:, joint_u2b]) - np.testing.assert_allclose(data._joint_vel_user.numpy(), data._sim_bind_joint_vel.numpy()[:, joint_u2b]) - np.testing.assert_allclose( - data._body_link_pose_w_user.numpy(), data._sim_bind_body_link_pose_w.numpy()[:, body_u2b] - ) - np.testing.assert_allclose(data._body_com_vel_w_user.numpy(), data._sim_bind_body_com_vel_w.numpy()[:, body_u2b]) + np.testing.assert_allclose(data._joint_pos_user.numpy(), data._sim_bind_joint_pos.numpy()[:, joint_u2b]) + np.testing.assert_allclose(data._joint_vel_user.numpy(), data._sim_bind_joint_vel.numpy()[:, joint_u2b]) + np.testing.assert_allclose( + data._body_link_pose_w_user.numpy(), data._sim_bind_body_link_pose_w.numpy()[:, body_u2b] + ) + np.testing.assert_allclose( + data._body_com_vel_w_user.numpy(), data._sim_bind_body_com_vel_w.numpy()[:, body_u2b] + ) - # ``_clear_callbacks`` must deregister exactly this hook so it does not leak on the class-level list, - # and leave an unrelated callback (standing in for another articulation's hook) untouched. - registered_callback = articulation._post_step_callback - assert registered_callback is not None - assert registered_callback in SimulationManager._post_step_callbacks - def _other_callback() -> None: - return None +## +# Composite scene. +## - SimulationManager.register_post_step_callback(_other_callback) - articulation._clear_callbacks() - assert articulation._post_step_callback is None - assert registered_callback not in SimulationManager._post_step_callbacks - assert _other_callback in SimulationManager._post_step_callbacks +@pytest.fixture(scope="module", params=test_devices(DeviceScope.CPU)) +def composite_scene(request: pytest.FixtureRequest) -> Iterator[_Scene]: + """Initialize every articulation island once for this module.""" + device = request.param + sim_cfg = newton_sim_cfg(device, use_newton_actuators=False, newton_contacts=True) + with build_simulation_context(sim_cfg=sim_cfg) as sim: + cfgs = _island_cfgs() + articulations = spawn_assets(cfgs) + sim.reset() + yield _Scene(sim=sim, articulations=articulations, cfgs=cfgs, origins=env_origins(device), device=device) -@pytest.mark.parametrize("num_articulations", [1]) -@pytest.mark.parametrize("device", ["cpu"]) -@pytest.mark.parametrize("gravity_enabled", [False]) -@pytest.mark.parametrize("articulation_type", ["anymal"]) -@pytest.mark.parametrize("use_newton_actuators", [False, True]) -@pytest.mark.parametrize("ordering_mode", ["none", "reversed"]) -def test_write_data_to_sim_writes_joint_targets_in_backend_order( - sim, num_articulations, device, gravity_enabled, articulation_type, use_newton_actuators, ordering_mode -): - """Write the published joint position targets into the backend buffer in backend joint order. +@pytest.fixture +def scene(composite_scene: _Scene) -> _Scene: + """Hand each test the composite scene with every island at rest. - Explicit actuators on the Newton-actuator path publish their raw targets; implicit actuators on the Lab - path publish the processed targets. Both must reach the solver-bound buffer permuted to backend order. + Islands of one environment share a solver world, so a state one test leaves behind must not reach the next. """ - actuator_cfg = ( - IdealPDActuatorCfg(joint_names_expr=[".*"], stiffness=40.0, damping=5.0) - if use_newton_actuators - else ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=40.0, damping=5.0) - ) - articulation_cfg = replace( - generate_articulation_cfg(articulation_type=articulation_type), actuators={"legs": actuator_cfg} - ) - if ordering_mode == "reversed": - articulation_cfg = replace(articulation_cfg, joint_ordering=tuple(reversed(ANYMAL_C_PHYSX_JOINT_NAMES))) - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) - replicate(sim.get_clone_plan()) - sim.reset() - assert articulation.is_initialized - assert (articulation.data.joint_ordering is not None) is (ordering_mode == "reversed") - assert articulation._has_newton_actuators is use_newton_actuators + for articulation in composite_scene.articulations.values(): + composite_scene.rest(articulation) + return composite_scene - # Distinct per-joint targets away from the defaults, so a skipped or unpermuted write is visible. - target = articulation.data.default_joint_pos.torch.clone() - target += 0.01 * torch.arange(1, articulation.num_joints + 1, device=device) - articulation.set_joint_position_target_index(target=target) - articulation.write_data_to_sim() - - source = ( - articulation.actuators.target_command.position.torch - if use_newton_actuators - else articulation.actuators.output_command.position.torch - ) - torch.testing.assert_close(source, target) - user_to_backend = ( - list(articulation.joint_ordering.user_to_backend_indices) - if articulation.joint_ordering is not None - else list(range(articulation.num_joints)) - ) - expected_backend_target = torch.empty_like(source) - expected_backend_target[:, user_to_backend] = source - torch.testing.assert_close(wp.to_torch(articulation.data._sim_bind_joint_position_target), expected_backend_target) +def _model_attribute(articulation: Articulation, name: str) -> torch.Tensor: + """Read one live Newton model attribute through the articulation's view, in backend order.""" + return wp.to_torch(articulation.root_view.get_attribute(name, SimulationManager.get_model()))[:, 0] -@pytest.mark.isaacsim_ci -@pytest.mark.parametrize("num_articulations", [3]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("add_ground_plane", [True]) -@pytest.mark.parametrize("articulation_type", ["humanoid"]) -def test_gravity_vec_w_tracks_model_gravity(sim, num_articulations, device, add_ground_plane, articulation_type): - """Per-env mutations to Newton's ``model.gravity`` reach ``GRAVITY_VEC_W`` and ``projected_gravity_b``. - Regression for the pre-fix snapshot: ``GRAVITY_VEC_W`` used to be env 0's - gravity broadcast to every env, hiding per-env gravity randomization (e.g. - :class:`~isaaclab.envs.mdp.randomize_physics_scene_gravity`). +def test_articulation_initialization_and_partial_state(scene: _Scene) -> None: + """Initialize every island and prove partial joint-state writes against live backend state. - The humanoid's articulation root sits on a rigid body, so this also checks floating-base initialization. + Covers fixed and floating roots, a root fixed and a root freed by spawner fragments, the per-island buffer + shapes and actuator models, a partial write with unsorted int64 selectors on a reordered articulation, body + state refresh after joint writes without a step, and joint position limit writes that keep or clamp the + default joint positions. """ - articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type, stiffness=0.0, damping=0.0) - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device, add_ground_plane=True) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(articulation) < 10 - - replicate(sim.get_clone_plan()) - sim.reset() + device = scene.device + expected_fixed_base = { + "floating": False, + "fixed": True, + "pendulum": True, + "ordered": False, + "tendon": True, + "root_fixed": True, + "root_freed": False, + "limits_implicit": False, + "limits_explicit": False, + "branching": False, + "osc": True, + } + for name, articulation in scene.articulations.items(): + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10, name + assert articulation.is_initialized, name + assert articulation.is_fixed_base == expected_fixed_base[name], name + assert articulation.num_instances == 2 + assert articulation.data.root_pos_w.torch.shape == (2, 3) + assert articulation.data.root_quat_w.torch.shape == (2, 4) + assert articulation.data.joint_pos.torch.shape == (2, articulation.num_joints) + assert articulation.data.body_mass.torch.shape == (2, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (2, articulation.num_bodies, 9) + for actuator_name, actuator in articulation.actuators.items(): + is_implicit_model_cfg = isinstance(scene.cfgs[name].actuators[actuator_name], ImplicitActuatorCfg) + assert actuator.is_implicit_model == is_implicit_model_cfg + + floating = scene.articulations["floating"] + assert floating.num_bodies == 2 + assert floating.num_joints == 1 + assert floating.joint_names == ["Joint"] + assert floating.data.body_com_pos_b.shape == (2, 2) + assert scene.articulations["tendon"].num_fixed_tendons > 0 + + # Newton consumes the base manager's world joint without relocating the root API. + root = sim_utils.get_first_matching_child_prim( + "/World/Env_0/RootFixed", lambda prim: prim.HasAPI(UsdPhysics.ArticulationRootAPI), stage=scene.sim.stage + ) + assert root is not None and root.HasAPI(UsdPhysics.RigidBodyAPI) + assert sim_utils.find_global_fixed_joint_prim("/World/Env_0/RootFixed", stage=scene.sim.stage) is not None + + # A partial write to environment 1 leaves environment 0 untouched. + scene.rest(floating) + env_ids = torch.tensor([1], dtype=torch.int32, device=device) + joint_ids = torch.tensor([0], dtype=torch.int32, device=device) + initial_joint_pos = floating.data.joint_pos.torch.clone() + initial_joint_vel = floating.data.joint_vel.torch.clone() + target_joint_pos = torch.tensor([[0.25]], device=device) + target_joint_vel = torch.tensor([[-0.5]], device=device) + floating.write_joint_state_to_sim_index( + position=target_joint_pos, velocity=target_joint_vel, env_ids=env_ids, joint_ids=joint_ids + ) + torch.testing.assert_close(floating.data.joint_pos.torch[env_ids], target_joint_pos) + torch.testing.assert_close(floating.data.joint_vel.torch[env_ids], target_joint_vel) + torch.testing.assert_close(floating.data.joint_pos.torch[:1], initial_joint_pos[:1]) + torch.testing.assert_close(floating.data.joint_vel.torch[:1], initial_joint_vel[:1]) + + # The same contract on a reordered articulation, with unsorted int64 selectors. + articulation = scene.articulations["ordered"] + scene.rest(articulation) + num_articulations = articulation.num_instances + default_joint_pos = articulation.data.default_joint_pos.torch.clone() + authored_limits = articulation.data.joint_pos_limits.torch.clone() - # 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, 21) + limits = torch.stack( + ( + -5.0 - torch.rand(num_articulations, articulation.num_joints, device=device), + 5.0 + torch.rand(num_articulations, articulation.num_joints, device=device), + ), + dim=-1, + ) + articulation.write_joint_position_limit_to_sim_index(limits=limits) - # -- actuator type - for actuator_name, actuator in articulation.actuators.items(): - is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) - assert getattr(actuator, "is_implicit_model", False) == is_implicit_model_cfg + # Check new limits are in place and the in-limit defaults are preserved + torch.testing.assert_close(articulation.data.joint_pos_limits.torch, limits) + torch.testing.assert_close(articulation.data.default_joint_pos.torch, default_joint_pos) - # GRAVITY_VEC_W must share storage with Newton's per-env gravity array. - model = SimulationManager.get_model() - model_gravity_arr = model.gravity[: model.world_count] - global_gravity = wp.to_torch(model.gravity)[-1].clone() - assert articulation.data.GRAVITY_VEC_W.warp.ptr == model_gravity_arr.ptr - assert articulation.data.GRAVITY_VEC_W.shape == (num_articulations,) + # Write joint state with unsorted int64 selectors + env_ids = torch.tensor([1, 0], dtype=torch.int64, device=device) + joint_ids = torch.tensor([articulation.num_joints - 1, 0], dtype=torch.int64, device=device) + position = torch.tensor([[0.21, 0.11], [0.22, 0.12]], device=device) + velocity = torch.tensor([[1.21, 1.11], [1.22, 1.12]], device=device) + expected_position = articulation.data.joint_pos.torch.clone() + expected_velocity = articulation.data.joint_vel.torch.clone() + articulation.write_joint_state_to_sim_index( + position=position, velocity=velocity, env_ids=env_ids, joint_ids=joint_ids + ) + expected_position[env_ids[:, None], joint_ids[None, :]] = position + expected_velocity[env_ids[:, None], joint_ids[None, :]] = velocity + torch.testing.assert_close(articulation.data.joint_pos.torch, expected_position) + torch.testing.assert_close(articulation.data.joint_vel.torch, expected_velocity) - # Mutate model.gravity per-env in place, as randomize_physics_scene_gravity does. - new_gravity = torch.tensor( - [[0.1 * (i + 1), 0.2 * (i + 1), -3.0 - float(i)] for i in range(num_articulations)], - device=device, - dtype=torch.float32, + # A joint state write moves the bodies and refreshes the derived body state without a step. + joint_pos_limits = articulation.data.joint_pos_limits.torch + joint_vel_limits = articulation.data.joint_vel_limits.torch.clamp(max=10.0) + pos_dist = torch.distributions.Uniform(joint_pos_limits[..., 0], joint_pos_limits[..., 1]) + vel_dist = torch.distributions.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() + articulation.write_joint_position_to_sim_index(position=pos_dist.sample()) + articulation.write_joint_velocity_to_sim_index(velocity=vel_dist.sample()) + 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 + ) + # 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 + body_link_vel_w = articulation.data.body_link_vel_w.torch + torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) + # validate link - com consistency + 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), + ) + torch.testing.assert_close(expected_com_pos.view(num_articulations, -1, 3), articulation.data.body_com_pos_w.torch) + torch.testing.assert_close( + expected_com_quat.view(num_articulations, -1, 4), articulation.data.body_com_quat_w.torch ) - wp.to_torch(model_gravity_arr).copy_(new_gravity) - SimulationManager.add_model_change(ModelFlags.MODEL_PROPERTIES) - # Live view: new per-env values are visible immediately, no invalidation step. - torch.testing.assert_close(articulation.data.GRAVITY_VEC_W.torch, new_gravity) - torch.testing.assert_close(wp.to_torch(model.gravity)[-1], global_gravity) + # Set new joint limits with indexing that invalidate the selected default joint positions + env_ids = torch.arange(1, device=device, dtype=torch.int32) + joint_ids = torch.nonzero(default_joint_pos[0].abs() > 0.1).squeeze(-1)[:2].to(torch.int32) + assert len(joint_ids) == 2 + limits = torch.stack( + ( + -0.1 * torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device), + 0.1 * torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device), + ), + dim=-1, + ) + articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) - # Recompute the lazily-cached projected_gravity_b without sim.step (which would - # drift root orientation from the reset state). Project against the same quat - # buffer the kernel reads so the expectation holds for any default orientation. - articulation.update(sim.cfg.dt) - root_quat = articulation.data.root_link_quat_w.torch - expected = math_utils.quat_apply_inverse(root_quat, torch.nn.functional.normalize(new_gravity, dim=-1)) - torch.testing.assert_close(articulation.data.projected_gravity_b.torch, expected, atol=1e-5, rtol=1e-5) + # Check new limits are in place and the defaults are clamped into them + torch.testing.assert_close(articulation.data.joint_pos_limits.torch[env_ids][:, joint_ids], limits) + 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) + # Islands share their environment's solver world, so leave the sampled states and narrow limits behind. + articulation.write_joint_position_limit_to_sim_index(limits=authored_limits) + scene.rest(articulation) + scene.rest(floating) -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("articulation_type", ["panda"]) -def test_fixed_base_reports_body_velocities(sim, num_articulations, device, articulation_type): - """Test that fixed-base articulations report live body velocities while their joints move. - Regression test: the fixed-base fallback in ``_create_buffers`` zeroed the body - center-of-mass velocity binding together with the (genuinely unavailable) root velocity, - so :attr:`body_lin_vel_w` and :attr:`body_ang_vel_w` read zeros for every fixed-base - robot regardless of motion. +def test_articulation_joint_and_body_properties_round_trip(scene: _Scene, monkeypatch: pytest.MonkeyPatch) -> None: + """Prove joint and body property writes reach the live Newton model and notify the solver. - This test verifies that: - 1. The articulation is initialized as fixed base with correctly shaped buffers - 2. Commanding a joint-space motion moves the bodies (finite difference of positions) while the root holds - its default state - 3. The reported body velocities track the finite-difference ground truth - - Args: - sim: The simulation fixture - num_articulations: Number of articulations to test - device: The device to run the simulation on + Covers partial friction and inertial writes with their exact model-change notifications, per-environment + static friction rows in the model, passive viscous friction kept apart from the actuator's derivative gain + through both writer selectors, and actuator gains taken from the configuration or, when unset, from the + USD drives. """ - articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) - 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 - replicate(sim.get_clone_plan()) - 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) - - # -- actuator type - for actuator_name, actuator in articulation.actuators.items(): - is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) - assert getattr(actuator, "is_implicit_model", False) == is_implicit_model_cfg - - # the root holds its default state as it is fixed base - default_root_pose = articulation.data.default_root_pose.torch.clone() - default_root_pose[:, :3] = default_root_pose[:, :3] + translations - default_root_vel = articulation.data.default_root_vel.torch.clone() - - # command a step away from the default pose so the distal bodies move - joint_pos_target = articulation.data.default_joint_pos.torch.clone() - joint_pos_target[:, 1] += 0.5 - prev_body_pos = articulation.data.body_link_pos_w.torch.clone() - reported_max = [] - fin_diff_max = [] + device = scene.device + articulation = scene.articulations["floating"] + env_ids = torch.tensor([1], dtype=torch.int32, device=device) + body_ids = torch.tensor([1], dtype=torch.int32, device=device) + joint_ids = torch.tensor([0], dtype=torch.int32, device=device) + notifications = [] + add_model_change = SimulationManager.add_model_change + + def record_model_change(change: ModelFlags) -> None: + notifications.append(change) + add_model_change(change) + + monkeypatch.setattr(SimulationManager, "add_model_change", staticmethod(record_model_change)) + initial_friction = articulation.data.joint_friction_coeff.torch.clone() + friction = torch.tensor([[0.3]], device=device) + articulation.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff=friction, env_ids=env_ids, joint_ids=joint_ids + ) + torch.testing.assert_close(articulation.data.joint_friction_coeff.torch[env_ids][:, joint_ids], friction) + torch.testing.assert_close(articulation.data.joint_friction_coeff.torch[:1], initial_friction[:1]) + torch.testing.assert_close(_model_attribute(articulation, "joint_friction")[env_ids][:, joint_ids], friction) + torch.testing.assert_close(_model_attribute(articulation, "joint_friction")[:1], initial_friction[:1]) + assert notifications == [ModelFlags.JOINT_DOF_PROPERTIES] + + notifications.clear() + initial_mass = articulation.data.body_mass.torch.clone() + masses = torch.tensor([[3.0]], device=device) + articulation.set_masses_index(masses=masses, env_ids=env_ids, body_ids=body_ids) + torch.testing.assert_close(articulation.data.body_mass.torch[env_ids][:, body_ids], masses) + torch.testing.assert_close(articulation.data.body_mass.torch[:1], initial_mass[:1]) + torch.testing.assert_close(_model_attribute(articulation, "body_mass")[env_ids][:, body_ids], masses) + assert notifications == [ModelFlags.BODY_INERTIAL_PROPERTIES] + + notifications.clear() + initial_com = articulation.data.body_com_pos_b.torch.clone() + coms = torch.tensor([[[0.05, -0.02, 0.01]]], device=device) + articulation.set_coms_index(coms=coms, env_ids=env_ids, body_ids=body_ids) + torch.testing.assert_close(articulation.data.body_com_pos_b.torch[env_ids][:, body_ids], coms) + torch.testing.assert_close(articulation.data.body_com_pos_b.torch[:1], initial_com[:1]) + assert notifications == [ModelFlags.BODY_INERTIAL_PROPERTIES] + + notifications.clear() + initial_inertia = articulation.data.body_inertia.torch.clone() + inertias = torch.diag_embed(torch.tensor([[[2.0, 3.0, 4.0]]], device=device)).reshape(1, 1, 9) + articulation.set_inertias_index(inertias=inertias, env_ids=env_ids, body_ids=body_ids) + torch.testing.assert_close(articulation.data.body_inertia.torch[env_ids][:, body_ids], inertias) + torch.testing.assert_close(articulation.data.body_inertia.torch[:1], initial_inertia[:1]) + assert notifications == [ModelFlags.BODY_INERTIAL_PROPERTIES] + monkeypatch.undo() + + # restore the configured friction and inertial properties for the tests that drive this island + articulation.write_joint_friction_coefficient_to_sim_index(joint_friction_coeff=initial_friction) + articulation.set_masses_index(masses=initial_mass) + articulation.set_coms_index(coms=initial_com) + articulation.set_inertias_index(inertias=initial_inertia) + + # Passive viscous joint damping is distinct from actuator derivative gains. + articulation = scene.articulations["fixed"] + slide_joint_ids = articulation.actuators["slides"].joint_indices + expected_viscous_friction = torch.full((articulation.num_instances, 3), 0.25, device=device) + torch.testing.assert_close( + articulation.data.joint_viscous_friction_coeff.torch[:, slide_joint_ids], expected_viscous_friction + ) + torch.testing.assert_close( + _model_attribute(articulation, "joint_damping")[:, slide_joint_ids], expected_viscous_friction + ) + expected_pd_damping = torch.full_like(expected_viscous_friction, 100.0) + torch.testing.assert_close(articulation.data.joint_damping.torch[:, slide_joint_ids], expected_pd_damping) + torch.testing.assert_close( + _model_attribute(articulation, "joint_target_kd")[:, slide_joint_ids], expected_pd_damping + ) + for writer, value in ( + (articulation.write_joint_viscous_friction_coefficient_to_sim_index, 0.25), + (articulation.write_joint_viscous_friction_coefficient_to_sim_mask, 0.5), + ): + values = torch.full((articulation.num_instances, articulation.num_joints), value, device=device) + writer(joint_viscous_friction_coeff=values) + torch.testing.assert_close(articulation.data.joint_viscous_friction_coeff.torch, values) + torch.testing.assert_close(_model_attribute(articulation, "joint_damping"), values) + + # Distinct per-env rows catch writers that ignore the env index + friction = torch.rand(articulation.num_instances, articulation.num_joints, device=device) + assert not torch.allclose(friction[0], friction[1]) + articulation.write_joint_friction_coefficient_to_sim_index(joint_friction_coeff=friction) + torch.testing.assert_close(_model_attribute(articulation, "joint_friction"), friction) + # restore the fixture's passive joint properties + articulation.write_joint_friction_coefficient_to_sim_index(joint_friction_coeff=torch.zeros_like(friction)) + viscous_friction = torch.zeros_like(friction) + viscous_friction[:, slide_joint_ids] = 0.25 + articulation.write_joint_viscous_friction_coefficient_to_sim_index(joint_viscous_friction_coeff=viscous_friction) + + # Configured gains reach the actuator; unset gains fall back to the USD drives, which author angular gains + # per degree. + slides, wrist = articulation.actuators["slides"], articulation.actuators["wrist"] + torch.testing.assert_close(slides.stiffness, torch.full((2, 3), 2000.0, device=device)) + torch.testing.assert_close(slides.damping, torch.full((2, 3), 100.0, device=device)) + usd_stiffness_per_radian = (torch.tensor(WRIST_USD_STIFFNESS, device=device) * 180.0 / torch.pi).expand(2, -1) + torch.testing.assert_close(wrist.stiffness, usd_stiffness_per_radian, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(wrist.damping, 0.1 * usd_stiffness_per_radian, rtol=1e-5, atol=1e-5) + + +def test_articulation_drive_and_dynamics(scene: _Scene) -> None: + """Prove fixed-root actuation, live body velocities, and a partial drive command. + + Fixed roots, including one fixed by a spawner fragment, hold their default state while the joints move. + The reported body velocities track the finite-difference motion of the driven links, which stay zero if + a fixed-base fallback zeroes the body velocity binding. + """ + articulation = scene.articulations["fixed"] + root_fixed = scene.articulations["root_fixed"] + scene.rest(articulation) + scene.rest(root_fixed) + fixed_roots = {name: scene.default_root_pose_w(scene.articulations[name]) for name in ("fixed", "root_fixed")} + default_root_vel = articulation.data.default_root_vel.torch.clone() + initial_joint_pos = articulation.data.joint_pos.torch.clone() + + # command a step away from the default pose on environment 1 only + joint_pos_target = articulation.data.default_joint_pos.torch.clone() + joint_pos_target[1, 0] += 0.1 + articulation.actuators.target_command.set_position_index(value=joint_pos_target) + prev_body_pos = articulation.data.body_link_pos_w.torch.clone() + reported_max = [] + fin_diff_max = [] for _ in range(20): - articulation.set_joint_position_target(joint_pos_target) - articulation.write_data_to_sim() - sim.step() - articulation.update(sim.cfg.dt) - 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) + scene.step("fixed", "root_fixed") + for name, default_root_pose in fixed_roots.items(): + torch.testing.assert_close(scene.articulations[name].data.root_link_pose_w.torch, default_root_pose) + torch.testing.assert_close(scene.articulations[name].data.root_com_vel_w.torch, default_root_vel) body_pos = articulation.data.body_link_pos_w.torch - reported_max.append(articulation.data.body_lin_vel_w.torch.norm(dim=-1).amax()) - fin_diff_max.append(((body_pos - prev_body_pos) / sim.cfg.dt).norm(dim=-1).amax()) + reported_max.append(articulation.data.body_lin_vel_w.torch[1].norm(dim=-1).amax()) + fin_diff_max.append(((body_pos[1] - prev_body_pos[1]) / scene.sim.cfg.dt).norm(dim=-1).amax()) prev_body_pos = body_pos.clone() reported_max = torch.stack(reported_max).amax() fin_diff_max = torch.stack(fin_diff_max).amax() @@ -1562,524 +1216,211 @@ def test_fixed_base_reports_body_velocities(sim, num_articulations, device, arti assert fin_diff_max > 0.1 # and the reported body velocities track it (identically zero under the regression) assert reported_max > 0.5 * fin_diff_max + assert articulation.data.joint_pos.torch[1, 0] > initial_joint_pos[1, 0] + 0.05 + # the environment that was not commanded stays at rest + torch.testing.assert_close(articulation.data.joint_pos.torch[0], initial_joint_pos[0], atol=1e-6, rtol=0) -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("articulation_type", ["shadow_hand"]) -def test_hand_with_tendons_initializes_and_targets_only_given_envs(sim, num_articulations, device, articulation_type): - """Initialize a fixed-base hand with tendons; a tendon command for one environment must leave the others alone. - - ``set_fixed_tendon_position_target_index`` is declared backend-neutral and documented to accept - partial data. Newton took ``env_ids`` and never forwarded it, so a partial command was sized - against every instance and raised rather than commanding the environment asked for. - """ - articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(articulation) < 10 - - replicate(sim.get_clone_plan()) - sim.reset() - assert articulation.is_initialized - assert articulation.is_fixed_base - assert articulation.num_fixed_tendons > 0 - 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) - for actuator_name, actuator in articulation.actuators.items(): - is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) - assert getattr(actuator, "is_implicit_model", False) == is_implicit_model_cfg - - target = torch.full((1, articulation.num_fixed_tendons), 1.0, dtype=torch.float32, device=device) - articulation.set_fixed_tendon_position_target_index(target=target, env_ids=[0]) - - for _ in range(30): - articulation.write_data_to_sim() - sim.step() - articulation.update(sim.cfg.dt) - - # Both environments start from the same pose under the same gravity, so any divergence comes - # from the command -- and identical poses would mean it reached both. - commanded, untouched = articulation.data.joint_pos.torch[0], articulation.data.joint_pos.torch[1] - assert not torch.allclose(commanded, untouched) - # Each Shadow Hand tendon ``rh_XFJ0`` is the sum of joints ``rh_XFJ1`` and ``rh_XFJ2``, so the - # commanded environment's tendon lengths must have moved toward the 1.0 target and away from - # the uncommanded environment, which the actuator holds at its 0.0 control. - for tendon_name in articulation.fixed_tendon_names: - joint_ids, _ = articulation.find_joints([tendon_name[:-1] + "1", tendon_name[:-1] + "2"]) - assert commanded[joint_ids].sum() > untouched[joint_ids].sum() - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", ["cuda:0"]) -@pytest.mark.parametrize("articulation_type", ["shadow_hand"]) -def test_fixed_tendon_properties_reach_solver(sim, num_articulations, device, articulation_type): - """Written fixed tendon stiffness, damping, and position limits reach the MuJoCo solver. - - Covers both the index and the mask setters and writers. - """ - articulation, _ = generate_articulation( - generate_articulation_cfg(articulation_type=articulation_type), num_articulations, device - ) - replicate(sim.get_clone_plan()) - sim.reset() - shape = (num_articulations, articulation.num_fixed_tendons) - limits = torch.tensor([-0.1, 0.2], device=device).expand(*shape, 2) - - articulation.set_fixed_tendon_stiffness_mask(stiffness=torch.full(shape, 12.0, device=device)) - articulation.set_fixed_tendon_damping_index(damping=torch.full(shape, 3.0, device=device)) - articulation.set_fixed_tendon_position_limit_index(limit=limits) - articulation.write_fixed_tendon_properties_to_sim_mask() - sim.step() - - solver_model = SimulationManager._solver.mjw_model - np.testing.assert_allclose(solver_model.tendon_stiffness.numpy(), 12.0) - np.testing.assert_allclose(solver_model.tendon_damping.numpy(), 3.0) - np.testing.assert_allclose(solver_model.tendon_range.numpy(), limits.cpu().numpy(), rtol=1e-6) - torch.testing.assert_close(articulation.data.fixed_tendon_pos_limits.torch, limits) - - articulation.set_fixed_tendon_stiffness_index(stiffness=42.0, env_ids=[1], fixed_tendon_ids=[1]) - articulation.set_fixed_tendon_damping_index(damping=6.0, env_ids=[1], fixed_tendon_ids=[1]) - articulation.write_fixed_tendon_properties_to_sim_index(env_ids=[1], fixed_tendon_ids=[1]) - sim.step() - expected_stiffness = np.full(shape, 12.0) - expected_damping = np.full(shape, 3.0) - expected_stiffness[1, 1] = 42.0 - expected_damping[1, 1] = 6.0 - np.testing.assert_allclose(solver_model.tendon_stiffness.numpy(), expected_stiffness) - np.testing.assert_allclose(solver_model.tendon_damping.numpy(), expected_damping) - - -@pytest.mark.parametrize("device", ["cpu"]) -@pytest.mark.parametrize("add_ground_plane", [True]) -@pytest.mark.parametrize("articulation_type", ["anymal"]) -def test_fragment_fix_root_link_uses_base_manager(sim, device, add_ground_plane, articulation_type): - """Newton consumes the base manager's world joint without relocating the root API. - - The floating-base ANYmal made fixed-base must then hold its root at the default state. - """ - articulation_cfg = deepcopy(generate_articulation_cfg(articulation_type=articulation_type)) - articulation_cfg.spawn.articulation_props = [] - articulation_cfg.spawn.fix_root_link = True - articulation, translations = generate_articulation( - articulation_cfg, num_articulations=1, device=device, add_ground_plane=True - ) - - root = sim_utils.get_first_matching_child_prim( - "/World/Env_0/Robot", - lambda prim: prim.HasAPI(UsdPhysics.ArticulationRootAPI), - stage=sim.stage, - ) - assert root is not None and root.HasAPI(UsdPhysics.RigidBodyAPI) - assert sim_utils.find_global_fixed_joint_prim("/World/Env_0/Robot", stage=sim.stage) is not None - - replicate(sim.get_clone_plan()) - sim.reset() - assert articulation.is_initialized - assert articulation.is_fixed_base - assert articulation.data.joint_pos.torch.shape == (1, 12) - - for _ in range(10): - sim.step() - 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", [2]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("add_ground_plane", [True]) -@pytest.mark.parametrize("articulation_type", ["panda"]) -def test_initialization_fixed_base_made_floating_base( - sim, num_articulations, device, add_ground_plane, articulation_type -): - """Test initialization for fixed base made floating-base using schema properties. +def test_floating_articulation_root_and_wrench_response(scene: _Scene, monkeypatch: pytest.MonkeyPatch) -> None: + """Prove floating-root state writes, center-of-mass updates, and wrench delivery. - This test verifies that: - 1. The articulation is properly initialized - 2. The articulation is floating base after modification - 3. All buffers have correct shapes - - Args: - sim: The simulation fixture - num_articulations: Number of articulations to test + A root pose or velocity write reads back in the written frame and refreshes the derived frame without a + step; a partial write changes only the selected environment. Forces on several bodies and forces at a + position, in the local and the global frame, reach the solver as the matching net force and moment, and a + reset clears the wrenches of the selected environments only. """ - articulation_cfg = clone(generate_articulation_cfg(articulation_type=articulation_type)) - # Unfix root link by making it non-kinematic - articulation_cfg.spawn.fix_root_link = False - articulation, _ = generate_articulation( - articulation_cfg, num_articulations, device=sim.device, add_ground_plane=True - ) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(articulation) < 10 - - # Play sim - replicate(sim.get_clone_plan()) - 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) - - -@pytest.mark.parametrize("device", ["cpu"]) -@pytest.mark.parametrize("articulation_type", ["panda"]) -@pytest.mark.parametrize("state_field", ["joint_pos", "joint_vel"]) -def test_out_of_range_default_joint_state(sim, device, articulation_type, state_field): - """Initialization fails when the configured default joint position or velocity exceeds the joint limits.""" - articulation_cfg = clone(generate_articulation_cfg(articulation_type=articulation_type)) - if state_field == "joint_pos": - articulation_cfg.init_state.joint_pos = {"panda_joint1": 10.0, "panda_joint[2, 4]": -20.0} - else: - articulation_cfg.init_state.joint_vel = {"panda_joint1": 100.0, "panda_joint[2, 4]": -60.0} - articulation, _ = generate_articulation(articulation_cfg, 1, device=device) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(articulation) < 10 - - quantity = "positions" if state_field == "joint_pos" else "velocities" - replicate(sim.get_clone_plan()) - with pytest.raises(ValueError, match=f"default {quantity} out of the limits"): - sim.reset() + device = scene.device + articulation = scene.articulations["floating"] + scene.rest(articulation) + num_articulations = articulation.num_instances + env_pos = scene.origins + assert articulation.data.root_link_pose_w.torch.shape == (2, 7) + assert articulation.data.root_com_pose_w.torch.shape == (2, 7) + assert articulation.data.body_link_pose_w.torch.shape == (2, articulation.num_bodies, 7) + assert articulation.data.body_com_pose_w.torch.shape == (2, articulation.num_bodies, 7) + # Resolve root body index by name (ordering may differ across physics backends) + root_idx = articulation.find_bodies("Root")[0][0] -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("articulation_type", ["anymal"]) -def test_external_force_on_multiple_bodies(sim, num_articulations, device, articulation_type): - """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=articulation_type) - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + # change center of mass offset from link frame on the root body + original_com = articulation.data.body_com_pos_b.torch.clone() + # Populate the derived world-frame cache so a missing invalidation would surface as a stale read. + original_com_w = articulation.data.body_com_pos_w.torch.clone() + new_com = original_com.clone() + new_com[:, root_idx] = torch.tensor([1.0, 0.0, 0.0], device=device) + env_ids = torch.arange(num_articulations, device=device, dtype=torch.int32) + # Full poses are accepted too; Newton uses the position and ignores the orientation. + com_poses = articulation.data.body_com_pose_b.torch.clone() + com_poses[..., :3] = new_com + articulation.set_coms_index(coms=com_poses, env_ids=env_ids) + torch.testing.assert_close(articulation.data.body_com_pos_b.torch, new_com, atol=1e-5, rtol=1e-5) + articulation.set_coms_index(coms=original_com, env_ids=env_ids) + torch.testing.assert_close(articulation.data.body_com_pos_b.torch, original_com, atol=1e-5, rtol=1e-5) + _ = articulation.data.body_com_pos_w.torch + articulation.set_coms_index(coms=new_com, env_ids=env_ids) - # Play the simulator - replicate(sim.get_clone_plan()) - sim.reset() + torch.testing.assert_close(articulation.data.body_com_pos_b.torch, new_com, atol=1e-5, rtol=1e-5) + # Without a sim step the links stay put, so the world-frame CoM is the link pose applied to the new offset. + link_pos_w = articulation.data.body_link_pos_w.torch + link_quat_w = articulation.data.body_link_quat_w.torch + expected_com_w = link_pos_w + math_utils.quat_apply(link_quat_w, new_com) + updated_com_w = articulation.data.body_com_pos_w.torch + torch.testing.assert_close(updated_com_w, expected_com_w, atol=1e-5, rtol=1e-5) + assert not torch.allclose(updated_com_w, original_com_w) - # 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] = 200.0 + def random_root_state(num_envs: int) -> torch.Tensor: + state = torch.rand(num_envs, 13, device=device) + state[..., :3] += env_pos[:num_envs] + # make quaternion a unit vector + state[..., 3:7] = torch.nn.functional.normalize(state[..., 3:7], dim=-1) + return state - # 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 + for state_location in ("com", "link", "root"): + + def write_root_state(state: torch.Tensor, env_ids: torch.Tensor | None = None) -> None: + if state_location == "com": + articulation.write_root_com_pose_to_sim_index(root_pose=state[..., :7], env_ids=env_ids) + articulation.write_root_com_velocity_to_sim_index(root_velocity=state[..., 7:], env_ids=env_ids) + elif state_location == "link": + articulation.write_root_link_pose_to_sim_index(root_pose=state[..., :7], env_ids=env_ids) + articulation.write_root_link_velocity_to_sim_index(root_velocity=state[..., 7:], env_ids=env_ids) + elif state_location == "root": + articulation.write_root_pose_to_sim_index(root_pose=state[..., :7], env_ids=env_ids) + articulation.write_root_velocity_to_sim_index(root_velocity=state[..., 7:], env_ids=env_ids) + + def read_written_root_state() -> torch.Tensor: + # the root pose is the link pose and the root velocity is the center-of-mass velocity + pose_w = ( + articulation.data.root_com_pose_w if state_location == "com" else articulation.data.root_link_pose_w + ) + vel_w = articulation.data.root_link_vel_w if state_location == "link" else articulation.data.root_com_vel_w + return torch.cat((pose_w.torch, vel_w.torch), dim=-1) + rand_state = random_root_state(num_articulations) + scene.step("floating") -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("articulation_type", ["anymal"]) -def test_external_force_on_multiple_bodies_at_position(sim, num_articulations, device, articulation_type): - """Test application of external force on the legs of the articulation at a given position. + # Prime the lazily-derived caches at the current sim timestamp. Without this they would + # recompute on first access after the write regardless of invalidation; priming them makes a + # missing reset_pose/reset_velocity observable as a stale read in the assertions below. + _ = articulation.data.root_link_pose_w.torch + _ = articulation.data.root_com_pose_w.torch + _ = articulation.data.root_link_vel_w.torch + _ = articulation.data.root_com_vel_w.torch - 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=articulation_type) - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) - - # Play the simulator - replicate(sim.get_clone_plan()) - 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] = 50.0 - external_wrench_positions_b = torch.zeros(articulation.num_instances, len(body_ids), 3, device=sim.device) - external_wrench_positions_b[..., 1] = 1.0 - - # Now we are ready! - for i in range(2): - # 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() + write_root_state(rand_state) - 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 + body_com_pose_b = articulation.data.body_com_pose_b.torch + if state_location == "com": + # the com pose/vel was written, so the derived link pose/vel must be refreshed + root_com_pose_w = articulation.data.root_com_pose_w.torch + root_com_vel_w = articulation.data.root_com_vel_w.torch + expected_root_link_pos, expected_root_link_quat = math_utils.combine_frame_transforms( + root_com_pose_w[:, :3], + root_com_pose_w[:, 3:], + math_utils.quat_rotate( + math_utils.quat_inv(body_com_pose_b[:, root_idx, 3:7]), -body_com_pose_b[:, root_idx, :3] + ), + math_utils.quat_inv(body_com_pose_b[:, root_idx, 3:7]), + ) + expected_root_link_pose = torch.cat((expected_root_link_pos, expected_root_link_quat), dim=1) + root_link_pose_w = articulation.data.root_link_pose_w.torch + root_link_vel_w = articulation.data.root_link_vel_w.torch + torch.testing.assert_close(expected_root_link_pose, root_link_pose_w) + # skip lin_vel because it differs from the link frame; angular velocity is frame-independent + # and only matches when the derived velocity was actually refreshed after the write + torch.testing.assert_close(root_com_vel_w[:, 3:], root_link_vel_w[:, 3:]) 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): - # the response axis depends on the link frames, so check that the articulation rotates - assert torch.linalg.vector_norm(articulation.data.root_ang_vel_w.torch[i]).item() > 0.1 - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("articulation_type", ["humanoid"]) -def test_loading_gains_from_usd(sim, num_articulations, device, articulation_type): - """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=articulation_type, stiffness=None, damping=None) - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) - - # Play sim - replicate(sim.get_clone_plan()) - 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", [2]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize( - ("joint_limit", "actuator_limit"), - [(1e5, None), (None, 1e2), (1e5, 1e2)], - ids=["joint_limit", "actuator_limit", "both_limits"], -) -@pytest.mark.parametrize("articulation_type", ["single_joint_implicit", "single_joint_explicit"]) -@pytest.mark.parametrize("use_newton_actuators", [False]) # consumed by the sim fixture -def test_setting_joint_limits_from_cfg( - sim, articulation_type, num_articulations, device, joint_limit, actuator_limit, use_newton_actuators -): - """Test the velocity and effort limit resolution for implicit and explicit actuators. - - This test verifies that: - 1. The solver clamps ``joint_velocity_limit`` and ``joint_effort_limit`` are applied to the simulation; - when unset, the USD-authored values are kept - 2. The actuator limits keep their configured values and are never pushed to the solver - 3. When unset, the actuator velocity limit falls back to the solver clamp, implicit actuators track the - solver effort clamp, and an explicit actuator's effort limit falls back to the USD-authored value - - Args: - sim: The simulation fixture - articulation_type: The implicit or explicit single-joint articulation - num_articulations: Number of articulations to test - device: The device to run the simulation on - joint_limit: The velocity and effort limits to set in simulation - actuator_limit: The velocity and effort limits to set in the actuator - """ - articulation_cfg = generate_articulation_cfg( - articulation_type=articulation_type, - joint_velocity_limit=joint_limit, - joint_effort_limit=joint_limit, - actuator_velocity_limit=actuator_limit, - actuator_effort_limit=actuator_limit, - ) - articulation, _ = generate_articulation( - articulation_cfg=articulation_cfg, - num_articulations=num_articulations, - device=device, - ) - # Play sim - replicate(sim.get_clone_plan()) - sim.reset() - - # read the values set into the simulation - model = SimulationManager.get_model() - newton_vel_limit = wp.to_torch(articulation.root_view.get_attribute("joint_velocity_limit", model)).to(device)[ - :, 0, : - ] - newton_effort_limit = wp.to_torch(articulation.root_view.get_attribute("joint_effort_limit", model)).to(device)[ - :, 0, : - ] - usd_vel_limit = next( - p.max_joint_velocity for p in articulation_cfg.spawn.joint_drive_props if isinstance(p, PhysxJointCfg) - ) - usd_effort_limit = next( - p.max_force for p in articulation_cfg.spawn.joint_drive_props if isinstance(p, sim_utils.UsdPhysicsDriveCfg) - ) - actuator = articulation.actuators["joint"] - - # check data buffers - torch.testing.assert_close(articulation.data.joint_vel_limits.torch, newton_vel_limit) - torch.testing.assert_close(articulation.data.joint_effort_limits.torch, newton_effort_limit) - # the solver clamps come from the joint limits when set, otherwise the USD-authored values - expected_vel_limit = usd_vel_limit if joint_limit is None else joint_limit - expected_effort_limit = usd_effort_limit if joint_limit is None else joint_limit - torch.testing.assert_close(newton_vel_limit, torch.full_like(newton_vel_limit, expected_vel_limit)) - torch.testing.assert_close(newton_effort_limit, torch.full_like(newton_effort_limit, expected_effort_limit)) - - # the actuator velocity limit keeps its configured value and is not pushed to the solver; - # when unset it falls back to the solver clamp - if actuator_limit is not None: - torch.testing.assert_close(actuator.actuator_velocity_limit, torch.full_like(newton_vel_limit, actuator_limit)) - assert not torch.allclose(actuator.actuator_velocity_limit, newton_vel_limit) - else: - torch.testing.assert_close(actuator.actuator_velocity_limit, newton_vel_limit) - - if articulation_type == "single_joint_implicit": - # without a separately configured rated limit, the implicit actuator limits track the solver clamp - torch.testing.assert_close(actuator.joint_effort_limit, newton_effort_limit) - torch.testing.assert_close(actuator.actuator_effort_limit, newton_effort_limit) - elif actuator_limit is not None: - torch.testing.assert_close(actuator.actuator_effort_limit, torch.full_like(newton_effort_limit, actuator_limit)) - else: - # an unset explicit actuator effort limit falls back to the USD-authored value, not the solver clamp + # the link pose/vel was written, so the derived com pose/vel must be refreshed + root_link_pose_w = articulation.data.root_link_pose_w.torch + root_link_vel_w = articulation.data.root_link_vel_w.torch + expected_com_pos, expected_com_quat = math_utils.combine_frame_transforms( + root_link_pose_w[:, :3], + root_link_pose_w[:, 3:], + body_com_pose_b[:, root_idx, :3], + body_com_pose_b[:, root_idx, 3:7], + ) + expected_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1) + root_com_pose_w = articulation.data.root_com_pose_w.torch + root_com_vel_w = articulation.data.root_com_vel_w.torch + torch.testing.assert_close(expected_com_pose, root_com_pose_w) + torch.testing.assert_close(root_link_vel_w[:, 3:], root_com_vel_w[:, 3:]) + + # the written state reads back in the written frame + torch.testing.assert_close(read_written_root_state(), rand_state) + + # a partial write only changes the selected environment + partial_state = random_root_state(1) + write_root_state(partial_state, env_ids=torch.tensor([0], device=device, dtype=torch.int32)) + expected_state = rand_state.clone() + expected_state[0] = partial_state[0] + torch.testing.assert_close(read_written_root_state(), expected_state) + + articulation.set_coms_index(coms=original_com) + + # Wrenches reach the solver as the matching net force and moment. Environment 0 receives none and stays at + # rest. The bodies lie on the x axis and the joint turns about z, so the island turns about x as one rigid + # body whose inertia is the sum of the body inertias. + total_mass = articulation.data.body_mass.torch[0].sum() + rigid_inertia_xx = articulation.data.body_inertia.torch[0, :, 0].sum() + dt = scene.sim.cfg.dt + env_1 = torch.tensor([1], dtype=torch.int32, device=device) + zeros = torch.zeros((1, articulation.num_bodies, 3), device=device) + + def one_step_response() -> tuple[torch.Tensor, torch.Tensor]: + """Step once; return environment 1's center-of-mass linear velocity and root angular velocity.""" + scene.step("floating") torch.testing.assert_close( - actuator.actuator_effort_limit, torch.full_like(newton_effort_limit, usd_effort_limit) + articulation.data.body_com_vel_w.torch[0], torch.zeros(2, 6, device=device), atol=1e-6, rtol=0.0 ) - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("articulation_type", ["humanoid"]) -def test_reset(sim, num_articulations, device, articulation_type, monkeypatch): - """Test that the actuator gains come from the configuration and that reset method works properly.""" - 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 - replicate(sim.get_clone_plan()) - sim.reset() - - # Check that gains are loaded from the configuration - expected_stiffness = torch.full( - (articulation.num_instances, articulation.num_joints), 10.0, device=articulation.device - ) - expected_damping = torch.full_like(expected_stiffness, 2.0) - torch.testing.assert_close(articulation.actuators["body"].stiffness, expected_stiffness) - torch.testing.assert_close(articulation.actuators["body"].damping, expected_damping) - - # Now we are ready! - # reset articulation - actuator = next(iter(articulation.actuators.values())) + body_masses = articulation.data.body_mass.torch[1].unsqueeze(-1) + com_lin_vel = (body_masses * articulation.data.body_com_lin_vel_w.torch[1]).sum(dim=0) / total_mass + return com_lin_vel, articulation.data.root_com_ang_vel_w.torch[1] + + # Opposite forces on the two bodies form a pure moment about -z that turns both links. + scene.rest(articulation) + forces = zeros.clone() + forces[0, :, 1] = torch.tensor([10.0, -10.0], device=device) + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=forces, torques=zeros, env_ids=env_1, body_ids=torch.arange(2, device=device, dtype=torch.int32) + ) + _, ang_vel = one_step_response() + assert ang_vel[2] < -0.1 + torch.testing.assert_close(ang_vel[:2], torch.zeros(2, device=device), atol=1e-5, rtol=0.0) + + # A force at a position offset along y, set then added in the local frame and set then added in the global + # frame, applies twice the force and twice its moment about x. + force = 1.0 + for is_global in (False, True): + scene.rest(articulation) + positions = torch.tensor([[[0.0, 0.1, 0.0]]], device=device) + if is_global: + positions += articulation.data.body_com_pos_w.torch[1:, :1] + forces = torch.tensor([[[0.0, 0.0, force]]], device=device) + for write in ( + articulation.permanent_wrench_composer.set_forces_and_torques_index, + articulation.permanent_wrench_composer.add_forces_and_torques_index, + ): + write( + forces=forces, + torques=zeros[:, :1], + positions=positions, + env_ids=env_1, + body_ids=torch.tensor([root_idx], dtype=torch.int32, device=device), + is_global=is_global, + ) + lin_vel, ang_vel = one_step_response() + torch.testing.assert_close(lin_vel[2], 2.0 * force * dt / total_mass, rtol=1e-3, atol=0.0) + torch.testing.assert_close(ang_vel[0], 2.0 * 0.1 * force * dt / rigid_inertia_xx, rtol=1e-3, atol=0.0) + + # Check that the gains come from the configuration and that reset works per environment + expected_stiffness = torch.full((articulation.num_instances, articulation.num_joints), 20.0, device=device) + torch.testing.assert_close(articulation.actuators["joint"].stiffness, expected_stiffness) + torch.testing.assert_close(articulation.actuators["joint"].damping, torch.full_like(expected_stiffness, 2.0)) + actuator = articulation.actuators["joint"] actuator_reset = actuator.reset reset_env_ids = [] - def record_actuator_reset(env_ids=None): + def record_actuator_reset(env_ids: torch.Tensor | None = None) -> None: reset_env_ids.append(env_ids) actuator_reset(env_ids) @@ -2115,633 +1456,308 @@ def record_actuator_reset(env_ids=None): assert torch.count_nonzero(instantaneous_composer.out_torque_b.torch) == num_bodies * 3 assert torch.count_nonzero(permanent_composer.out_force_b.torch) == num_bodies * 3 assert torch.count_nonzero(permanent_composer.out_torque_b.torch) == num_bodies * 3 + articulation.reset() -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("articulation_type", ["single_joint_implicit"]) -def test_body_root_state(sim, num_articulations, device, articulation_type): - """Test for reading the `body_state_w` property. - - This test verifies that: - 1. The single-joint articulation initializes as fixed base with correctly shaped buffers - 2. Body link and center-of-mass states match an analytic pendulum with a center-of-mass offset - 3. The fixed root holds its default state - - 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=articulation_type) - articulation, env_pos = generate_articulation(articulation_cfg, num_articulations, device) - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(articulation) < 10, "Possible reference leak for articulation" - # Play sim - replicate(sim.get_clone_plan()) - 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" - # 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) - - # -- actuator type - for actuator_name, actuator in articulation.actuators.items(): - is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) - assert getattr(actuator, "is_implicit_model", False) == is_implicit_model_cfg +def test_branching_fixture_physx_ordering_reorders_newton_to_bfs(scene: _Scene) -> None: + """Resolve the documented Newton ``joint_ordering="physx"`` sim-to-sim workflow on a branching asset. - # 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") + Mirrors :func:`isaaclab_physx.test.assets.test_articulation.test_branching_fixture_resolves_distinct_conventions` + with the backend roles swapped: here the live backend is Newton (depth-first, so its native view is + the MJWarp order), and the request is ``physx``/``body_ordering="physx"``. Cross-backend discovery must + resolve the breadth-first PhysX order and reorder the public joint/body axes to it. This is the headline + workflow documented in + ``docs/source/overview/core-concepts/physical-backends/joint_and_body_ordering.rst``. - # change center of mass offset from link frame - offset = [0.5, 0.0, 0.0] + The same articulation also checks the MJWarp-order emulation used for cross-backend discovery and that + selected inertial-property writes keep Newton's inverse arrays current under the body ordering. + """ + articulation = scene.articulations["branching"] + device = scene.device - # create com offsets — apply offset to the Arm body - num_bodies = articulation.num_bodies - com = wp.to_torch(articulation.root_view.get_attribute("body_com", SimulationManager.get_model())) - 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[:, 0, arm_idx, :] = new_com.squeeze(-2) - articulation.root_view.set_attribute("body_com", SimulationManager.get_model(), wp.from_torch(com, dtype=wp.vec3f)) - with wp.ScopedDevice(device): - SimulationManager._solver.notify_model_changed(ModelFlags.BODY_INERTIAL_PROPERTIES) + # Newton's native traversal is depth-first, so the live backend view already reflects MJWarp order. + assert tuple(articulation.backend_joint_names) == BRANCHING_MJWARP_JOINT_NAMES + assert tuple(articulation.backend_body_names) == BRANCHING_MJWARP_BODY_NAMES - # check they are set - torch.testing.assert_close( - wp.to_torch(articulation.root_view.get_attribute("body_com", SimulationManager.get_model())), com + # The same-backend "mjwarp" request takes an identity fast path, so also force the cross-backend + # emulation (the temporary Newton USD builder a PhysX-backed articulation would use) and compare its + # independently rebuilt view against the live backend view. BFS and DFS differ on this branching fixture. + emulated_names = ordering_resolvers._get_mjwarp_names_from_newton_usd_builder(articulation) + assert emulated_names is not None + assert emulated_names["joint"] == tuple(articulation.backend_joint_names) + assert emulated_names["body"] == tuple(articulation.backend_body_names) + assert get_articulation_name_ordering(articulation, "mjwarp", kind="joint") == tuple( + articulation.backend_joint_names ) + assert get_articulation_name_ordering(articulation, "mjwarp", kind="body") == tuple(articulation.backend_body_names) - for i in range(50): - # perform step - sim.step() - # update buffers - 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] + env_pos - 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) - - # get state properties - 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 - - # 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) - - # 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:]) - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("state_location", ["com", "link", "root"]) -@pytest.mark.parametrize("gravity_enabled", [False]) -@pytest.mark.parametrize("articulation_type", ["anymal"]) -def test_write_root_state_functions_data_consistency( - sim, num_articulations, device, state_location, gravity_enabled, articulation_type -): - """A root pose/velocity write must read back in the written frame and refresh the derived cross-frame caches. - - Regression coverage for the velocity invalidation cleanup: writing the root center-of-mass - (or link) velocity must invalidate the derived root link (or com) velocity so the next read - re-derives it. Linear velocity differs between the two frames, so - as in the rigid object - test - we compare angular velocity, which is frame-independent and therefore only matches when - the derived velocity was actually refreshed. - - The center-of-mass offset is set through ``set_coms_index``, which must update the body-frame CoM and - the world-frame CoM derived from it without a sim step. - """ - sim._app_control_on_stop_handle = None - articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) - articulation, env_pos = generate_articulation(articulation_cfg, num_articulations, device) + # Cross-backend discovery (bypassing the same-backend fast path) resolves the breadth-first PhysX order. + assert get_articulation_name_ordering(articulation, "physx", kind="joint") == BRANCHING_PHYSX_JOINT_NAMES + assert get_articulation_name_ordering(articulation, "physx", kind="body") == BRANCHING_PHYSX_BODY_NAMES - # Play sim - replicate(sim.get_clone_plan()) - sim.reset() + # The requested PhysX ordering reorders the public joint/body axes to the BFS convention. + assert tuple(articulation.joint_names) == BRANCHING_PHYSX_JOINT_NAMES + assert tuple(articulation.body_names) == BRANCHING_PHYSX_BODY_NAMES + assert articulation.joint_ordering is not None + assert articulation.body_ordering is not None - # Resolve root body index by name (ordering may differ across physics backends) - root_idx = articulation.find_bodies("base")[0][0] + # Selected mass and inertia writes with int64 selectors update Newton's inverse arrays in backend order. + env_ids = torch.tensor([0], dtype=torch.int64, device=device) + body_ids = torch.tensor([2, articulation.num_bodies - 1], dtype=torch.int64, device=device) + backend_body_ids = torch.tensor( + [articulation.data.body_ordering.user_to_backend_indices[index] for index in body_ids.tolist()], + dtype=torch.int64, + device=device, + ) + assert backend_body_ids[0] != body_ids[0] + model = SimulationManager.get_model() + initial_masses = articulation.data.body_mass.torch.clone() + initial_inertias = articulation.data.body_inertia.torch.clone() - # change center of mass offset from link frame on the root body - original_com = articulation.data.body_com_pos_b.torch.clone() - # Populate the derived world-frame cache so a missing invalidation would surface as a stale read. - original_com_w = articulation.data.body_com_pos_w.torch.clone() - new_com = original_com.clone() - new_com[:, root_idx] = torch.tensor([1.0, 0.0, 0.0], device=device) - env_ids = torch.arange(num_articulations, device=device, dtype=torch.int32) - # Full poses are accepted too; Newton uses the position and ignores the orientation. - com_poses = articulation.data.body_com_pose_b.torch.clone() - com_poses[..., :3] = new_com - articulation.set_coms_index(coms=com_poses, env_ids=env_ids) - torch.testing.assert_close(articulation.data.body_com_pos_b.torch, new_com, atol=1e-5, rtol=1e-5) - articulation.set_coms_index(coms=original_com, env_ids=env_ids) - torch.testing.assert_close(articulation.data.body_com_pos_b.torch, original_com, atol=1e-5, rtol=1e-5) - _ = articulation.data.body_com_pos_w.torch - articulation.set_coms_index(coms=new_com, env_ids=env_ids) + masses = articulation.data.body_mass.torch[env_ids][:, body_ids].clone() + torch.tensor([[1.0, 2.0]], device=device) + articulation.set_masses_index(masses=masses, env_ids=env_ids, body_ids=body_ids) + model_inv_mass = wp.to_torch(articulation.root_view.get_attribute("body_inv_mass", model)[:, 0]) + torch.testing.assert_close(model_inv_mass[env_ids][:, backend_body_ids], masses.reciprocal()) - torch.testing.assert_close(articulation.data.body_com_pos_b.torch, new_com, atol=1e-5, rtol=1e-5) - # Without a sim step the links stay put, so the world-frame CoM is the link pose applied to the new offset. - link_pos_w = articulation.data.body_link_pos_w.torch - link_quat_w = articulation.data.body_link_quat_w.torch - expected_com_w = link_pos_w + math_utils.quat_apply(link_quat_w, new_com) - updated_com_w = articulation.data.body_com_pos_w.torch - torch.testing.assert_close(updated_com_w, expected_com_w, atol=1e-5, rtol=1e-5) - assert not torch.allclose(updated_com_w, original_com_w) + inertia_matrices = torch.diag_embed(torch.tensor([[[2.0, 3.0, 4.0], [5.0, 6.0, 7.0]]], device=device)) + articulation.set_inertias_index(inertias=inertia_matrices.reshape(1, 2, 9), env_ids=env_ids, body_ids=body_ids) + model_inv_inertia = wp.to_torch(articulation.root_view.get_attribute("body_inv_inertia", model)[:, 0]) + torch.testing.assert_close(model_inv_inertia[env_ids][:, backend_body_ids], torch.linalg.inv(inertia_matrices)) - def random_root_state(num_envs: int) -> torch.Tensor: - state = torch.rand(num_envs, 13, device=device) - state[..., :3] += env_pos[:num_envs] - # make quaternion a unit vector - state[..., 3:7] = torch.nn.functional.normalize(state[..., 3:7], dim=-1) - return state + articulation.set_masses_index(masses=initial_masses) + articulation.set_inertias_index(inertias=initial_inertias) - def write_root_state(state: torch.Tensor, env_ids: torch.Tensor | None = None): - if state_location == "com": - articulation.write_root_com_pose_to_sim_index(root_pose=state[..., :7], env_ids=env_ids) - articulation.write_root_com_velocity_to_sim_index(root_velocity=state[..., 7:], env_ids=env_ids) - elif state_location == "link": - articulation.write_root_link_pose_to_sim_index(root_pose=state[..., :7], env_ids=env_ids) - articulation.write_root_link_velocity_to_sim_index(root_velocity=state[..., 7:], env_ids=env_ids) - elif state_location == "root": - articulation.write_root_pose_to_sim_index(root_pose=state[..., :7], env_ids=env_ids) - articulation.write_root_velocity_to_sim_index(root_velocity=state[..., 7:], env_ids=env_ids) - - def read_written_root_state() -> torch.Tensor: - # the root pose is the link pose and the root velocity is the center-of-mass velocity - pose_w = articulation.data.root_com_pose_w if state_location == "com" else articulation.data.root_link_pose_w - vel_w = articulation.data.root_link_vel_w if state_location == "link" else articulation.data.root_com_vel_w - return torch.cat((pose_w.torch, vel_w.torch), dim=-1) - - rand_state = random_root_state(num_articulations) - - # perform a step then update the buffers - sim.step() - articulation.update(sim.cfg.dt) - - # Prime the lazily-derived caches at the current sim timestamp. Without this they would - # recompute on first access after the write regardless of invalidation; priming them makes a - # missing reset_pose/reset_velocity observable as a stale read in the assertions below. - _ = articulation.data.root_link_pose_w.torch - _ = articulation.data.root_com_pose_w.torch - _ = articulation.data.root_link_vel_w.torch - _ = articulation.data.root_com_vel_w.torch - - write_root_state(rand_state) - - body_com_pose_b = articulation.data.body_com_pose_b.torch - if state_location == "com": - # the com pose/vel was written, so the derived link pose/vel must be refreshed - root_com_pose_w = articulation.data.root_com_pose_w.torch - root_com_vel_w = articulation.data.root_com_vel_w.torch - expected_root_link_pos, expected_root_link_quat = math_utils.combine_frame_transforms( - root_com_pose_w[:, :3], - root_com_pose_w[:, 3:], - math_utils.quat_rotate( - math_utils.quat_inv(body_com_pose_b[:, root_idx, 3:7]), -body_com_pose_b[:, root_idx, :3] - ), - math_utils.quat_inv(body_com_pose_b[:, root_idx, 3:7]), - ) - expected_root_link_pose = torch.cat((expected_root_link_pos, expected_root_link_quat), dim=1) - root_link_pose_w = articulation.data.root_link_pose_w.torch - root_link_vel_w = articulation.data.root_link_vel_w.torch - torch.testing.assert_close(expected_root_link_pose, root_link_pose_w) - # skip lin_vel because it differs from the link frame; angular velocity is frame-independent - # and only matches when the derived velocity was actually refreshed after the write - torch.testing.assert_close(root_com_vel_w[:, 3:], root_link_vel_w[:, 3:]) - else: - # the link pose/vel was written, so the derived com pose/vel must be refreshed - root_link_pose_w = articulation.data.root_link_pose_w.torch - root_link_vel_w = articulation.data.root_link_vel_w.torch - expected_com_pos, expected_com_quat = math_utils.combine_frame_transforms( - root_link_pose_w[:, :3], - root_link_pose_w[:, 3:], - body_com_pose_b[:, root_idx, :3], - body_com_pose_b[:, root_idx, 3:7], - ) - expected_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1) - root_com_pose_w = articulation.data.root_com_pose_w.torch - root_com_vel_w = articulation.data.root_com_vel_w.torch - torch.testing.assert_close(expected_com_pose, root_com_pose_w) - torch.testing.assert_close(root_link_vel_w[:, 3:], root_com_vel_w[:, 3:]) - - # the written state reads back in the written frame - torch.testing.assert_close(read_written_root_state(), rand_state) - - # a partial write only changes the selected environment - partial_state = random_root_state(1) - write_root_state(partial_state, env_ids=torch.tensor([0], device=device, dtype=torch.int32)) - expected_state = rand_state.clone() - expected_state[0] = partial_state[0] - torch.testing.assert_close(read_written_root_state(), expected_state) - - -@pytest.mark.parametrize("device", ["cpu"]) -@pytest.mark.parametrize("articulation_type", ["humanoid"]) -@pytest.mark.parametrize("root_prim_path", ["/torso", "/non_existing_prim_path"]) -def test_setting_articulation_root_prim_path(sim, device, articulation_type, root_prim_path): - """An explicit articulation root prim path initializes when it exists and fails clearly otherwise.""" - sim._app_control_on_stop_handle = None - articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) - articulation_cfg.articulation_root_prim_path = root_prim_path - articulation, _ = generate_articulation(articulation_cfg, 1, device) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(articulation) < 10 - - if root_prim_path == "/torso": - replicate(sim.get_clone_plan()) - sim.reset() - assert articulation.is_initialized - else: - replicate(sim.get_clone_plan()) - with pytest.raises(KeyError, match="No articulations matching pattern"): - sim.reset() +def test_newton_ordered_state_publishes_in_step_and_refreshes_same_timestamp_writes(scene: _Scene) -> None: + """Republish the user-order Tier-1 shadows inside the sim step and refresh them after same-timestamp writes. -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("gravity_enabled", [False]) -@pytest.mark.parametrize("articulation_type", ["anymal"]) -def test_write_joint_state_data_consistency(sim, num_articulations, device, gravity_enabled, articulation_type, caplog): - """Joint limit and joint state writes update the joint buffers and refresh the body state without a step. + With non-identity ordering the passthrough state getters no longer reorder on read; the post-step + callback republishes the shadows from live backend state inside the stepped region. We deliberately + clobber the four shadows, step the simulation WITHOUT reading any state property, and assert the shadows + again equal the reordered backend state -- which can only hold if the hook ran inside the step. Under the + lazy design (no hook), only a property read would refresh them, so the clobbered shadows would stay stale. - This test verifies that: - 1. Joint position limits are written and keep in-limit default joint positions - 2. A partial joint state write with unsorted int64 selectors updates only the selected entries - 3. A joint state write moves the bodies and refreshes the derived body poses and velocities - 4. Joint limits that exclude a default joint position clamp it into the new limits and report the - clamping at the requested log level - - Args: - sim: The simulation fixture - num_articulations: Number of articulations to test - device: The device to run the simulation on + Root pose and velocity writes at the current simulation timestamp must then refresh the ordered body + pose and velocity, including after another consumer has resolved the shared forward kinematics. """ - sim._app_control_on_stop_handle = None - articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) - articulation, env_pos = generate_articulation(articulation_cfg, num_articulations, device) - env_idx = torch.tensor([x for x in range(num_articulations)]) - - # Play sim - replicate(sim.get_clone_plan()) - sim.reset() - - # Get current default joint pos - default_joint_pos = articulation.data.default_joint_pos.torch.clone() - - 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) + device = scene.device + articulation = scene.articulations["ordered"] + scene.rest(articulation) + data = articulation.data + assert data.joint_ordering is not None + assert data.body_ordering is not None + joint_u2b = np.asarray(articulation.joint_ordering.user_to_backend_indices) + body_u2b = np.asarray(articulation.body_ordering.user_to_backend_indices) - # Check new limits are in place and the in-limit defaults are preserved - torch.testing.assert_close(articulation.data.joint_pos_limits.torch, limits) - torch.testing.assert_close(articulation.data.default_joint_pos.torch, default_joint_pos) + # Clobber every Tier-1 shadow with a large sentinel so a stale (unrepublished) shadow + # is unmistakably detectable -- the true backend joint/velocity state sits near zero, + # so a plain zero-fill would coincide with it. Then step WITHOUT touching joint_pos / + # joint_vel / body_link_pose_w / body_com_vel_w. + data._joint_pos_user.fill_(1000.0) + data._joint_vel_user.fill_(1000.0) + data._body_link_pose_w_user.fill_(wp.transformf(1000.0, 1000.0, 1000.0, 0.0, 0.0, 0.0, 1.0)) + data._body_com_vel_w_user.fill_(wp.spatial_vectorf(1000.0, 1000.0, 1000.0, 1000.0, 1000.0, 1000.0)) + scene.sim.step() - # Write joint state with unsorted int64 selectors - env_ids = torch.tensor([1, 0], dtype=torch.int64, device=device) - joint_ids = torch.tensor([articulation.num_joints - 1, 0], dtype=torch.int64, device=device) - position = torch.tensor([[0.21, 0.11], [0.22, 0.12]], device=device) - velocity = torch.tensor([[1.21, 1.11], [1.22, 1.12]], device=device) - expected_position = articulation.data.joint_pos.torch.clone() - expected_velocity = articulation.data.joint_vel.torch.clone() - articulation.write_joint_state_to_sim_index( - position=position, velocity=velocity, env_ids=env_ids, joint_ids=joint_ids + np.testing.assert_allclose(data._joint_pos_user.numpy(), data._sim_bind_joint_pos.numpy()[:, joint_u2b]) + np.testing.assert_allclose(data._joint_vel_user.numpy(), data._sim_bind_joint_vel.numpy()[:, joint_u2b]) + np.testing.assert_allclose( + data._body_link_pose_w_user.numpy(), data._sim_bind_body_link_pose_w.numpy()[:, body_u2b] ) - expected_position[env_ids[:, None], joint_ids[None, :]] = position - expected_velocity[env_ids[:, None], joint_ids[None, :]] = velocity - torch.testing.assert_close(articulation.data.joint_pos.torch, expected_position) - torch.testing.assert_close(articulation.data.joint_vel.torch, expected_velocity) - - 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() + np.testing.assert_allclose(data._body_com_vel_w_user.numpy(), data._sim_bind_body_com_vel_w.numpy()[:, body_u2b]) - rand_joint_pos = pos_dist.sample() - rand_joint_vel = vel_dist.sample() + articulation.update(scene.sim.cfg.dt) + root_body_idx = articulation.find_bodies("base")[0][0] + sim_timestamp = data._sim_timestamp - 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 - ) - # 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 - body_link_vel_w = articulation.data.body_link_vel_w.torch - torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) + cached_body_pose = data.body_link_pose_w.torch[:, root_body_idx].clone() + written_root_pose = data.root_link_pose_w.torch.clone() + written_root_pose[:, 0] += 0.25 + articulation.write_root_link_pose_to_sim_index(root_pose=written_root_pose) + assert data._sim_timestamp == sim_timestamp + torch.testing.assert_close(data.root_link_pose_w.torch, written_root_pose) + SimulationManager.forward() # Another consumer may resolve shared FK before this view reads. + refreshed_body_pose = data.body_link_pose_w.torch[:, root_body_idx] + torch.testing.assert_close(refreshed_body_pose, written_root_pose) + assert not torch.equal(refreshed_body_pose, cached_body_pose) - # 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), - ) - torch.testing.assert_close(expected_com_pos.view(len(env_idx), -1, 3), articulation.data.body_com_pos_w.torch) - torch.testing.assert_close(expected_com_quat.view(len(env_idx), -1, 4), articulation.data.body_com_quat_w.torch) + # Populate the velocity cache after the pose write so only the velocity write can invalidate it. + cached_body_vel = data.body_com_vel_w.torch[:, root_body_idx].clone() + written_root_vel = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]], device=device).repeat(2, 1) + articulation.write_root_com_velocity_to_sim_index(root_velocity=written_root_vel) + assert data._sim_timestamp == sim_timestamp + torch.testing.assert_close(data.root_com_vel_w.torch, written_root_vel) + refreshed_body_vel = data.body_com_vel_w.torch[:, root_body_idx] + torch.testing.assert_close(refreshed_body_vel, written_root_vel) + assert not torch.equal(refreshed_body_vel, cached_body_vel) + scene.rest(articulation) - # Set new joint limits with indexing that invalidate the selected default joint positions - env_ids = torch.arange(1, device=device, dtype=torch.int32) - joint_ids = torch.nonzero(default_joint_pos[0].abs() > 0.1).squeeze(-1)[:2].to(torch.int32) - assert len(joint_ids) == 2 - 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_logger = Articulation.__module__ - with caplog.at_level(logging.WARNING, logger=articulation_logger): - articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) - assert any(record.levelno == logging.WARNING for record in caplog.records) - # Check new limits are in place and the defaults are clamped into them - torch.testing.assert_close(articulation.data.joint_pos_limits.torch[env_ids][:, joint_ids], limits) - 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("island", ["fixed", "ordered"], ids=["identity", "reordered"]) +def test_write_data_to_sim_writes_joint_targets_in_backend_order(scene: _Scene, island: str) -> None: + """Write the published joint position targets into the backend buffer in backend joint order. - # Info-level clamping reports are logged only when that level is enabled. - articulation.data.default_joint_pos.torch.fill_(1.0) - full_limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) - full_limits[..., 1] = 0.5 - caplog.clear() - with caplog.at_level(logging.INFO, logger=articulation_logger): - articulation.write_joint_position_limit_to_sim_mask(limits=full_limits, warn_limit_violation=False) - assert [record.levelno for record in caplog.records] == [logging.INFO] - articulation.data.default_joint_pos.torch.fill_(1.0) - caplog.clear() - with caplog.at_level(logging.WARNING, logger=articulation_logger): - articulation.write_joint_position_limit_to_sim_mask(limits=full_limits, warn_limit_violation=False) - assert not caplog.records - torch.testing.assert_close(articulation.data.default_joint_pos.torch, torch.full_like(full_limits[..., 1], 0.5)) - - -@pytest.mark.parametrize("selector_kind", ["index", "mask"]) -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("articulation_type", ["panda"]) -@pytest.mark.parametrize("device", test_devices()) -def test_write_joint_viscous_friction_to_sim(sim, num_articulations, device, articulation_type, selector_kind): - """Test passive viscous joint damping is distinct from actuator derivative gains. - - Static joint friction writes also propagate directly to the Newton model. + Implicit actuators on the Lab path publish the processed targets, which must reach the solver-bound + buffer permuted to backend order. The Newton-actuator path is covered with the native actuators. """ - articulation_cfg = clone(generate_articulation_cfg(articulation_type)) - articulation_cfg.actuators["panda_arm"].viscous_friction = 0.25 - articulation_cfg.actuators["panda_arm"].damping = 2.0 - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device) - replicate(sim.get_clone_plan()) - sim.reset() - - arm_joint_ids = articulation.actuators["panda_arm"].joint_indices - expected_viscous_friction = torch.full((articulation.num_instances, len(arm_joint_ids)), 0.25, device=device) - torch.testing.assert_close( - articulation.data.joint_viscous_friction_coeff.torch[:, arm_joint_ids], expected_viscous_friction - ) - torch.testing.assert_close( - wp.to_torch(articulation.root_view.get_attribute("joint_damping", SimulationManager.get_model()))[ - :, 0, arm_joint_ids - ], - expected_viscous_friction, - ) + device = scene.device + articulation = scene.articulations[island] + scene.rest(articulation) + assert (articulation.data.joint_ordering is not None) is (island == "ordered") + assert articulation._has_newton_actuators is False - expected_pd_damping = torch.full_like(expected_viscous_friction, 2.0) - torch.testing.assert_close(articulation.data.joint_damping.torch[:, arm_joint_ids], expected_pd_damping) - torch.testing.assert_close( - wp.to_torch(articulation.root_view.get_attribute("joint_target_kd", SimulationManager.get_model()))[ - :, 0, arm_joint_ids - ], - expected_pd_damping, - ) - - values = torch.full((articulation.num_instances, articulation.num_joints), 0.25, device=device) - if selector_kind == "index": - articulation.write_joint_viscous_friction_coefficient_to_sim_index( - joint_viscous_friction_coeff=values, - ) - else: - articulation.write_joint_viscous_friction_coefficient_to_sim_mask( - joint_viscous_friction_coeff=values, - ) + # Distinct per-joint targets away from the defaults, so a skipped or unpermuted write is visible. + target = articulation.data.default_joint_pos.torch.clone() + target += 0.01 * torch.arange(1, articulation.num_joints + 1, device=device) + articulation.set_joint_position_target_index(target=target) + articulation.write_data_to_sim() - torch.testing.assert_close(articulation.data.joint_viscous_friction_coeff.torch, values) - torch.testing.assert_close( - wp.to_torch(articulation.root_view.get_attribute("joint_damping", SimulationManager.get_model()))[:, 0], - values, + source = articulation.actuators.output_command.position.torch + torch.testing.assert_close(source, target) + user_to_backend = ( + list(articulation.joint_ordering.user_to_backend_indices) + if articulation.joint_ordering is not None + else list(range(articulation.num_joints)) ) - - # Distinct per-env rows catch writers that ignore the env index - friction = torch.rand(articulation.num_instances, articulation.num_joints, device=device) - assert not torch.allclose(friction[0], friction[1]) - articulation.write_joint_friction_coefficient_to_sim_index(joint_friction_coeff=friction) - joint_friction_coeff_sim = wp.to_torch( - articulation.root_view.get_attribute("joint_friction", SimulationManager.get_model()) - )[:, 0, :] - torch.testing.assert_close(joint_friction_coeff_sim, friction) + expected_backend_target = torch.empty_like(source) + expected_backend_target[:, user_to_backend] = source + torch.testing.assert_close(wp.to_torch(articulation.data._sim_bind_joint_position_target), expected_backend_target) + scene.rest(articulation) -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("articulation_type", ["anymal"]) -def test_body_q_consistent_after_root_write(num_articulations, device, articulation_type): - """Test that body_q is fresh when collide() runs after a root pose write. +@pytest.mark.parametrize("island", ["limits_implicit", "limits_explicit"], ids=["implicit", "explicit"]) +def test_setting_joint_limits_from_cfg(scene: _Scene, island: str) -> None: + """Test the velocity and effort limit resolution for implicit and explicit actuators. - Regression test for a NaN bug where collide() used stale body_q after env - reset because eval_fk was not called between write_root_pose and collide. + This test verifies that: + 1. The solver clamps ``joint_velocity_limit`` and ``joint_effort_limit`` are applied to the simulation; + when unset, the USD-authored values are kept + 2. The actuator limits keep their configured values and are never pushed to the solver + 3. When unset, the actuator velocity limit falls back to the solver clamp, implicit actuators track the + solver effort clamp, and an explicit actuator's effort limit falls back to the USD-authored value - Uses ``use_mujoco_contacts=False`` so the Newton collision pipeline is - active, then patches ``_simulate_physics_only`` to capture body_q at - the moment collide() is called and asserts it matches joint_q. + Each joint of the island belongs to one group: joint limits only, actuator limits only, both, or neither. """ - - sim_cfg = SimulationCfg( - dt=1 / 200, - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=70, - nconmax=70, - integrator="implicitfast", - use_mujoco_contacts=False, - ), - num_substeps=1, - use_cuda_graph=False, - ), + articulation = scene.articulations[island] + articulation_cfg = scene.cfgs[island] + newton_vel_limit = _model_attribute(articulation, "joint_velocity_limit") + newton_effort_limit = _model_attribute(articulation, "joint_effort_limit") + usd_vel_limit = next( + p.max_joint_velocity for p in articulation_cfg.spawn.joint_drive_props if isinstance(p, PhysxJointCfg) + ) + usd_effort_limit = next( + p.max_force for p in articulation_cfg.spawn.joint_drive_props if isinstance(p, sim_utils.UsdPhysicsDriveCfg) ) - with build_simulation_context(sim_cfg=sim_cfg, device=device) as sim: - sim._app_control_on_stop_handle = None - articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) - articulation, env_pos = generate_articulation(articulation_cfg, num_articulations, device) - - replicate(sim.get_clone_plan()) - sim.reset() - - model = SimulationManager.get_model() - jc_starts = model.joint_coord_world_start.numpy() - body_starts = model.body_world_start.numpy() - - for _ in range(5): - sim.step() - articulation.update(sim.cfg.dt) - # Teleport env 0 by 10m (simulating a reset) - new_pose = articulation.data.default_root_pose.torch.clone() - new_pose[0, 0] += 10.0 - new_pose[0, 1] += 5.0 - articulation.write_root_pose_to_sim_index( - root_pose=new_pose[0:1], - env_ids=torch.tensor([0], device=device, dtype=torch.int32), - ) + # check data buffers + torch.testing.assert_close(articulation.data.joint_vel_limits.torch, newton_vel_limit) + torch.testing.assert_close(articulation.data.joint_effort_limits.torch, newton_effort_limit) - # Patch _simulate_physics_only to capture body_q before collide runs - captured = {} - original_simulate = SimulationManager._simulate_physics_only.__func__ - - @classmethod # type: ignore[misc] - def _patched_simulate(cls): - if cls._needs_collision_pipeline: - bq = wp.to_torch(cls.backend.state_0.body_q) - jq = wp.to_torch(cls.backend.state_0.joint_q) - b0 = int(body_starts[0]) - jc0 = int(jc_starts[0]) - captured["bq_root"] = bq[b0, :3].clone() - captured["jq_root"] = jq[jc0 : jc0 + 3].clone() - original_simulate(cls) - - with patch.object(SimulationManager, "_simulate_physics_only", _patched_simulate): - sim.step() - articulation.update(sim.cfg.dt) - - assert captured, "collision pipeline did not run — _needs_collision_pipeline is False" - - bq_root = captured["bq_root"] - jq_root = captured["jq_root"] - diff = (jq_root - bq_root).abs().max().item() - assert diff < 0.01, ( - f"body_q was stale when collide() ran: diff={diff:.4f}m, jq={jq_root.tolist()}, bq={bq_root.tolist()}" - ) + for group_name, actuator in articulation.actuators.items(): + group_cfg = articulation_cfg.actuators[group_name] + joint_limit = group_cfg.joint_velocity_limit + actuator_limit = group_cfg.actuator_velocity_limit + joint_ids = actuator.joint_indices + group_vel_limit = newton_vel_limit[:, joint_ids] + group_effort_limit = newton_effort_limit[:, joint_ids] + + # the solver clamps come from the joint limits when set, otherwise the USD-authored values + expected_vel_limit = usd_vel_limit if joint_limit is None else joint_limit + expected_effort_limit = usd_effort_limit if joint_limit is None else joint_limit + torch.testing.assert_close(group_vel_limit, torch.full_like(group_vel_limit, expected_vel_limit)) + torch.testing.assert_close(group_effort_limit, torch.full_like(group_effort_limit, expected_effort_limit)) + + # the actuator velocity limit keeps its configured value and is not pushed to the solver; + # when unset it falls back to the solver clamp + if actuator_limit is not None: + torch.testing.assert_close( + actuator.actuator_velocity_limit, torch.full_like(group_vel_limit, actuator_limit) + ) + assert not torch.allclose(actuator.actuator_velocity_limit, group_vel_limit) + else: + torch.testing.assert_close(actuator.actuator_velocity_limit, group_vel_limit) + + if island == "limits_implicit": + # without a separately configured rated limit, the implicit actuator limits track the solver clamp + torch.testing.assert_close(actuator.joint_effort_limit, group_effort_limit) + torch.testing.assert_close(actuator.actuator_effort_limit, group_effort_limit) + elif actuator_limit is not None: + torch.testing.assert_close( + actuator.actuator_effort_limit, torch.full_like(group_effort_limit, actuator_limit) + ) + else: + # an unset explicit actuator effort limit falls back to the USD-authored value, not the solver clamp + torch.testing.assert_close( + actuator.actuator_effort_limit, torch.full_like(group_effort_limit, usd_effort_limit) + ) -@pytest.mark.parametrize("add_ground_plane", [True]) -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("articulation_type", ["panda"]) -@pytest.mark.parametrize("body_subset", [False, True]) -def test_set_material_properties(sim, num_articulations, device, add_ground_plane, articulation_type, body_subset): - """Material and collider-offset randomization write through to the robot's shapes in the Newton model. +def test_set_material_properties(scene: _Scene) -> None: + """Material and collider-offset randomization write through to the articulation's shapes in the Newton model. The event terms write the asset's view-level shape bindings; the assertions read the flat Newton model - arrays at the collision shapes of the selected bodies and environments. Visual shapes are ignored because - their materials and offsets have no physical effect. + arrays at the collision shapes of the selected bodies and environments. Material randomization runs once on + a body subset and once on every body. """ - articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) - articulation, _ = generate_articulation( - articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=device, add_ground_plane=True - ) + articulation = scene.articulations["floating"] + num_articulations = articulation.num_instances + device = scene.device - # Play the simulator - replicate(sim.get_clone_plan()) - sim.reset() - - # Resolve the robot's shapes per environment and body from the flat Newton model. + # Resolve this articulation's shapes per environment and body from the flat Newton model. model = SimulationManager.get_model() body_world = model.body_world.numpy() - body_names = [label.rsplit("/", 1)[-1] for label in model.body_label] + body_labels = list(model.body_label) shape_body = model.shape_body.numpy() is_collision_shape = (model.shape_flags.numpy() & int(ShapeFlags.COLLIDE_SHAPES)) != 0 def robot_shapes(env_index: int, selected_body_names: list[str] | None = None) -> np.ndarray: + prefix = f"/World/Env_{env_index}/Floating/" bodies = [ body for body in np.flatnonzero(body_world == env_index) - if selected_body_names is None or body_names[body] in selected_body_names + if body_labels[body].startswith(prefix) + and (selected_body_names is None or body_labels[body].rsplit("/", 1)[-1] in selected_body_names) ] return np.flatnonzero(np.isin(shape_body, bodies) & is_collision_shape) - env = SimpleNamespace(scene={"robot": articulation}, sim=sim, device=device, num_envs=num_articulations) + env = SimpleNamespace(scene={"robot": articulation}, sim=scene.sim, device=device, num_envs=num_articulations) env_ids = torch.tensor([num_articulations - 1], device=device) + other_env_shapes = robot_shapes(0) + configured_mu = model.shape_material_mu.numpy().copy() + configured_restitution = model.shape_material_restitution.numpy().copy() # Randomize the materials in the last environment, with degenerate ranges. - asset_cfg = SceneEntityCfg("robot") - selected_body_names = None - if body_subset: - asset_cfg.body_ids, selected_body_names = articulation.find_bodies(["panda_link3", "panda_hand"]) - selected_shapes = robot_shapes(num_articulations - 1, selected_body_names) - unselected_shapes = np.setdiff1d(robot_shapes(num_articulations - 1), selected_shapes) - other_env_shapes = robot_shapes(0) - assert len(selected_shapes) > 0 and len(other_env_shapes) > 0 - original_mu = model.shape_material_mu.numpy().copy() - original_restitution = model.shape_material_restitution.numpy().copy() - params = { - "static_friction_range": (0.55, 0.55), - "dynamic_friction_range": (0.55, 0.55), - "restitution_range": (0.15, 0.15), - "num_buckets": 1, - "asset_cfg": asset_cfg, - } - material_term = randomize_rigid_body_material( - EventTermCfg(func=randomize_rigid_body_material, mode="startup", params=params), env - ) - material_term(env, env_ids, **params) + for body_subset, friction, restitution_value in ((True, 0.55, 0.15), (False, 0.65, 0.25)): + asset_cfg = SceneEntityCfg("robot") + selected_body_names = None + if body_subset: + asset_cfg.body_ids, selected_body_names = articulation.find_bodies(["Child"]) + selected_shapes = robot_shapes(num_articulations - 1, selected_body_names) + unselected_shapes = np.setdiff1d(robot_shapes(num_articulations - 1), selected_shapes) + assert len(selected_shapes) > 0 and len(other_env_shapes) > 0 + assert body_subset == (len(unselected_shapes) > 0) + original_mu = model.shape_material_mu.numpy().copy() + original_restitution = model.shape_material_restitution.numpy().copy() + params = { + "static_friction_range": (friction, friction), + "dynamic_friction_range": (friction, friction), + "restitution_range": (restitution_value, restitution_value), + "num_buckets": 1, + "asset_cfg": asset_cfg, + } + material_term = randomize_rigid_body_material( + EventTermCfg(func=randomize_rigid_body_material, mode="startup", params=params), env + ) + material_term(env, env_ids, **params) + scene.step("floating") + + mu = model.shape_material_mu.numpy() + restitution = model.shape_material_restitution.numpy() + np.testing.assert_allclose(mu[selected_shapes], friction) + np.testing.assert_allclose(restitution[selected_shapes], restitution_value) + for shapes in (unselected_shapes, other_env_shapes): + np.testing.assert_array_equal(mu[shapes], original_mu[shapes]) + np.testing.assert_array_equal(restitution[shapes], original_restitution[shapes]) # Randomize the collider offsets of every shape in the last environment. original_margin = model.shape_margin.numpy().copy() @@ -2755,18 +1771,7 @@ def robot_shapes(env_index: int, selected_body_names: list[str] | None = None) - EventTermCfg(func=randomize_rigid_body_collider_offsets, mode="startup", params=offset_params), env ) offset_term(env, env_ids, **offset_params) - - # Simulate physics - sim.step() - articulation.update(sim.cfg.dt) - - mu = model.shape_material_mu.numpy() - restitution = model.shape_material_restitution.numpy() - np.testing.assert_allclose(mu[selected_shapes], 0.55) - np.testing.assert_allclose(restitution[selected_shapes], 0.15) - for shapes in (unselected_shapes, other_env_shapes): - np.testing.assert_array_equal(mu[shapes], original_mu[shapes]) - np.testing.assert_array_equal(restitution[shapes], original_restitution[shapes]) + scene.step("floating") # Newton maps the rest offset to the shape margin and the contact offset to margin + gap. randomized_shapes = robot_shapes(num_articulations - 1) @@ -2775,191 +1780,85 @@ def robot_shapes(env_index: int, selected_body_names: list[str] | None = None) - np.testing.assert_array_equal(model.shape_margin.numpy()[other_env_shapes], original_margin[other_env_shapes]) np.testing.assert_array_equal(model.shape_gap.numpy()[other_env_shapes], original_gap[other_env_shapes]) - -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize( - "first_property", - ["body_com_jacobian_w", "body_link_jacobian_w", "mass_matrix", "gravity_compensation_forces"], -) -def test_task_space_allocation_and_capture(monkeypatch, device, first_property): - """Allocate only requested outputs and retain them across changed-state graph replay.""" - builder = ModelBuilder() - builder.begin_world() - base = builder.add_link(mass=4.0, inertia=wp.mat33(np.eye(3)), label="Robot/base") - arm = builder.add_link(mass=2.0, inertia=wp.mat33(np.eye(3)), com=(0.5, 0.0, 0.0), label="Robot/arm") - slider = builder.add_link(mass=3.0, inertia=wp.mat33(np.eye(3)), com=(0.25, 0.0, 0.0), label="Robot/slider") - builder.add_articulation( - [ - builder.add_joint_fixed(-1, base), - builder.add_joint_revolute(base, arm, axis=(0.0, 1.0, 0.0), label="hinge"), - builder.add_joint_prismatic(arm, slider, axis=(1.0, 0.0, 0.0), label="slide"), - ], - label="Robot", - ) - builder.end_world() - model = builder.finalize(device=device) - state, control = model.state(), model.control() - monkeypatch.setattr(SimulationManager, "get_model", lambda: model) - monkeypatch.setattr(SimulationManager, "get_state_0", lambda: state) - monkeypatch.setattr(SimulationManager, "get_control", lambda: control) - view = ArticulationView(model, "Robot", exclude_joint_types=[JointType.FIXED]) - eager = ArticulationData(view, device) - eager._apply_ordering_maps_after_resolve() - newton.eval_fk(model, state.joint_q, state.joint_qd, state) - getattr(eager, first_property) # Compile kernels before first-use capture on a fresh container. - data = ArticulationData(view, device) - data._apply_ordering_maps_after_resolve() - properties = ("body_com_jacobian_w", "body_link_jacobian_w", "mass_matrix", "gravity_compensation_forces") - assert all(getattr(data, f"_{name}_ta") is None for name in properties) - assert data._jacobian_buf_flat is data._mass_matrix_full_buf is data._gravity_force_full_buf is None - if wp.get_device(device).is_cuda: - with wp.ScopedCapture(device=device) as capture: - output = getattr(data, first_property) - else: - output = getattr(data, first_property) - required = {first_property} - if first_property == "body_link_jacobian_w": - required.add("body_com_jacobian_w") - for name in properties: - assert (getattr(data, f"_{name}_ta") is not None) == (name in required) - assert (data._jacobian_buf_flat is not None) == (first_property != "gravity_compensation_forces") - assert (data._mass_matrix_full_buf is not None) == (first_property == "mass_matrix") - assert (data._gravity_force_full_buf is not None) == (first_property == "gravity_compensation_forces") - - for angle, displacement in ((0.0, 0.0), (0.7, 0.4), (-0.3, 0.2)): - state.joint_q.assign(np.asarray([angle, displacement], dtype=np.float32)) - newton.eval_fk(model, state.joint_q, state.joint_qd, state) - expected = getattr(eager, first_property).warp.numpy() - if wp.get_device(device).is_cuda: - wp.capture_launch(capture.graph) - else: - assert getattr(data, first_property) is output - np.testing.assert_allclose(output.warp.numpy(), expected, atol=1e-5) - data._create_simulation_bindings() - data._apply_ordering_maps_after_resolve() - assert getattr(data, first_property) is output + # restore the configured shape properties, which the other tests' contacts rely on + model.shape_material_mu.assign(configured_mu) + model.shape_material_restitution.assign(configured_restitution) + model.shape_margin.assign(original_margin) + model.shape_gap.assign(original_gap) + SimulationManager.add_model_change(ModelFlags.SHAPE_PROPERTIES) -## -# Shape-contract regression tests for the new BaseArticulation accessors. -# These pin the public shape contract so future regressions (e.g., reverting -# to model-wide max sizing or to the wrong fixed-base row offset) fail fast. -## +def test_hand_with_tendons_targets_only_given_envs_and_properties_reach_solver(scene: _Scene) -> None: + """Command fixed tendons for one environment and write fixed-tendon properties through to the MuJoCo solver. + ``set_fixed_tendon_position_target_index`` is declared backend-neutral and documented to accept + partial data. Newton took ``env_ids`` and never forwarded it, so a partial command was sized + against every instance and raised rather than commanding the environment asked for. -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("add_ground_plane", [True]) -@pytest.mark.parametrize("articulation_type", ["anymal"]) -@pytest.mark.isaacsim_ci -def test_heterogeneous_scene_per_view_shapes(sim, device, add_ground_plane, articulation_type): - """Mixed-articulation scene: each view returns ITS OWN asset's shape. - - Direct regression test for the Codex round-2 finding. With Franka - (9 DoFs) and Anymal-C (18 DoFs) co-resident in the model, - ``model.max_dofs_per_articulation == 18`` and - ``model.max_joints_per_articulation == anymal.num_bodies``. The Franka - view's ``body_link_jacobian_w`` / ``mass_matrix`` outputs must use - Franka's per-asset counts, NOT the model-wide maxima — otherwise - Franka's mass matrix would carry zero-padded rows/cols and be - singular. - - Uses the ``anymal`` ``SIM_CFGs`` entry (more capable solver settings) - for the host sim; the ``articulation_type`` parametrize is only there - so the ``sim`` fixture picks a config — the test itself constructs - both Anymal and Franka articulations directly. + Written fixed tendon stiffness, damping, and position limits reach the MuJoCo solver through both the + index and the mask setters and writers. """ - # ``num_per_type=1`` keeps the actuator-default replication path off — - # Newton's USD default loader hits a (1, num_joints) vs (num_envs, - # num_joints) shape mismatch with multi-instance multi-type scenes; one - # of each is the minimum heterogeneous setup that still exercises the - # per-articulation shape gate without that pre-existing quirk. - num_per_type = 1 - - franka_cfg = replace(_FRANKA_PANDA_NEWTON_CFG, prim_path="/World/Env_[^/]*/Franka") - anymal_cfg = replace(ANYMAL_C_CFG, prim_path="/World/Env_[^/]*/Anymal") - anymal_cfg.init_state.pos = (0.0, 5.0, anymal_cfg.init_state.pos[2]) - - sim_utils.create_prim("/World/Env_0", "Xform") - clone_plan_from_env_0( - CloneCfg(clone_template="/World/Env_{}"), - (franka_cfg, anymal_cfg, AssetBaseCfg(prim_path="/World/defaultGroundPlane")), - num_per_type, - 2.5, - ) - - franka = Articulation(franka_cfg) - anymal = Articulation(anymal_cfg) - replicate(sim.get_clone_plan()) - sim.reset() - assert franka.is_initialized and anymal.is_initialized - assert franka.is_fixed_base and not anymal.is_fixed_base - - # Sanity: the model-wide maxima are larger than at least one view's - # per-asset count, so a regression to model-wide sizing would manifest - # as wrong shapes here. Assert that precondition explicitly so the test - # fails clearly if the fixture stops being heterogeneous. - model = SimulationManager.get_model() - assert model.max_dofs_per_articulation > min(franka.num_joints, anymal.num_joints), ( - "scene is no longer heterogeneous; this test relies on model.max_dofs > one view's num_joints" - ) + device = scene.device + articulation = scene.articulations["tendon"] + scene.rest(articulation) + num_articulations = articulation.num_instances + assert articulation.num_fixed_tendons == 2 - franka_J = franka.data.body_link_jacobian_w.torch - anymal_J = anymal.data.body_link_jacobian_w.torch + target = torch.full((1, articulation.num_fixed_tendons), 1.0, dtype=torch.float32, device=device) + articulation.set_fixed_tendon_position_target_index(target=target, env_ids=[0]) + scene.step("tendon", num_steps=30) - # Each view's output uses its OWN per-asset count, not the model-wide max. - # Floating-base assets prepend ``num_base_dofs`` floating-base columns; fixed-base - # assets have ``num_base_dofs == 0``. - franka_dofs = franka.num_joints + franka.num_base_dofs - anymal_dofs = anymal.num_joints + anymal.num_base_dofs - assert franka_J.shape == torch.Size((num_per_type, franka.num_bodies - 1, 6, franka_dofs)), ( - f"Franka jacobian leaked model-wide shape: got {tuple(franka_J.shape)}" - ) - assert anymal_J.shape == torch.Size((num_per_type, anymal.num_bodies, 6, anymal_dofs)), ( - f"Anymal jacobian leaked model-wide shape: got {tuple(anymal_J.shape)}" - ) + # Both environments start from the same pose, so any divergence comes from the command -- and identical + # poses would mean it reached both. + commanded, untouched = articulation.data.joint_pos.torch[0], articulation.data.joint_pos.torch[1] + assert not torch.allclose(commanded, untouched) + # Each tendon ``XJ0`` is the sum of joints ``XJ1`` and ``XJ2``, so the commanded environment's tendon lengths + # must have moved toward the 1.0 target and away from the uncommanded environment, which the actuator holds + # at its 0.0 control. + for tendon_name in articulation.fixed_tendon_names: + joint_ids, _ = articulation.find_joints([tendon_name[:-1] + "1", tendon_name[:-1] + "2"]) + assert commanded[joint_ids].sum() > untouched[joint_ids].sum() + articulation.set_fixed_tendon_position_target_index(target=torch.zeros_like(target), env_ids=[0]) - sim.step() - franka.update(sim.cfg.dt) - anymal.update(sim.cfg.dt) + shape = (num_articulations, articulation.num_fixed_tendons) + limits = torch.tensor([-0.1, 0.2], device=device).expand(*shape, 2) + authored_stiffness = articulation.data.fixed_tendon_stiffness.torch.clone() + authored_damping = articulation.data.fixed_tendon_damping.torch.clone() + authored_limits = articulation.data.fixed_tendon_pos_limits.torch.clone() - franka_M = franka.data.mass_matrix.torch - anymal_M = anymal.data.mass_matrix.torch + articulation.set_fixed_tendon_stiffness_mask(stiffness=torch.full(shape, 12.0, device=device)) + articulation.set_fixed_tendon_damping_index(damping=torch.full(shape, 3.0, device=device)) + articulation.set_fixed_tendon_position_limit_index(limit=limits) + articulation.write_fixed_tendon_properties_to_sim_mask() + scene.step("tendon") - assert franka_M.shape == torch.Size((num_per_type, franka_dofs, franka_dofs)) - assert anymal_M.shape == torch.Size((num_per_type, anymal_dofs, anymal_dofs)) + solver_model = SimulationManager._solver.mjw_model + np.testing.assert_allclose(solver_model.tendon_stiffness.numpy(), 12.0) + np.testing.assert_allclose(solver_model.tendon_damping.numpy(), 3.0) + np.testing.assert_allclose(solver_model.tendon_range.numpy(), limits.cpu().numpy(), rtol=1e-6) + torch.testing.assert_close(articulation.data.fixed_tendon_pos_limits.torch, limits) - # Each view's mass matrix must have positive diagonals — padded zero - # rows/cols (the round-2 bug) would surface as zero diagonals on the - # smaller-DoF view. Using a per-diagonal check here instead of det() - # because det of a real Franka mass matrix is naturally ~1e-13. - assert (franka_M.diagonal(dim1=-2, dim2=-1) > 1e-6).all(), ( - "Franka mass matrix has non-positive diagonal under heterogeneous scene" - ) - assert (anymal_M.diagonal(dim1=-2, dim2=-1) > 1e-6).all(), ( - "Anymal mass matrix has non-positive diagonal under heterogeneous scene" - ) + articulation.set_fixed_tendon_stiffness_index(stiffness=42.0, env_ids=[1], fixed_tendon_ids=[1]) + articulation.set_fixed_tendon_damping_index(damping=6.0, env_ids=[1], fixed_tendon_ids=[1]) + articulation.write_fixed_tendon_properties_to_sim_index(env_ids=[1], fixed_tendon_ids=[1]) + scene.step("tendon") + expected_stiffness = np.full(shape, 12.0) + expected_damping = np.full(shape, 3.0) + expected_stiffness[1, 1] = 42.0 + expected_damping[1, 1] = 6.0 + np.testing.assert_allclose(solver_model.tendon_stiffness.numpy(), expected_stiffness) + np.testing.assert_allclose(solver_model.tendon_damping.numpy(), expected_damping) - # Gravity compensation gathers a FLAT model-wide DoF buffer, so a padded-layout - # regression (indexing by ``art_id * max_dofs`` instead of - # ``joint_qd_start[articulation_start[art_id]]``) is numerically invisible in - # homogeneous scenes — this mixed scene is the only place it can surface. - franka_g = franka.data.gravity_compensation_forces.torch - anymal_g = anymal.data.gravity_compensation_forces.torch - assert franka_g.shape == torch.Size((num_per_type, franka_dofs)) - assert anymal_g.shape == torch.Size((num_per_type, anymal_dofs)) - assert franka_g.abs().max() > 1e-3, "Franka gravity compensation is all-zero under heterogeneous scene" - assert anymal_g.abs().max() > 1e-3, "Anymal gravity compensation is all-zero under heterogeneous scene" + # restore the authored tendon properties, which the other tests' dynamics rely on + articulation.set_fixed_tendon_stiffness_index(stiffness=authored_stiffness) + articulation.set_fixed_tendon_damping_index(damping=authored_damping) + articulation.set_fixed_tendon_position_limit_index(limit=authored_limits) + articulation.write_fixed_tendon_properties_to_sim_index() -@pytest.mark.parametrize("num_articulations", [4]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("add_ground_plane", [True]) -@pytest.mark.parametrize("articulation_type", ["panda", "anymal"]) -@pytest.mark.parametrize("ordering_mode", ["none", "reversed"]) @pytest.mark.isaacsim_ci -def test_get_gravity_compensation_forces_matches_jacobian_gravity( - sim, num_articulations, device, add_ground_plane, articulation_type, ordering_mode -): +@pytest.mark.parametrize("island", ["fixed", "floating", "ordered"], ids=["fixed", "floating", "floating_reordered"]) +def test_get_gravity_compensation_forces_matches_jacobian_gravity(scene: _Scene, island: str) -> None: """``g(q)`` must equal ``-sum_b J_com_b^T (m_b * g_w)`` at the current configuration. Newton computes the gravity compensation force through an RNEA pass @@ -2971,139 +1870,123 @@ def test_get_gravity_compensation_forces_matches_jacobian_gravity( fixed upstream in newton#2625 (wrong gravity compensation under non-identity root pose). - With ``ordering_mode="reversed"`` a nonidentity joint ordering is active and - both sides of the identity must be expressed in user joint order: the - Jacobian gather applies the user->backend permutation, so a - ``gather_dof_force_rows`` that skips it returns backend-ordered forces and - breaks the identity row-wise. + On the reordered island a nonidentity joint ordering is active and both sides of the identity must be + expressed in user joint order: the Jacobian gather applies the user->backend permutation, so a + ``gather_dof_force_rows`` that skips it returns backend-ordered forces and breaks the identity row-wise. - The same fixture also pins: + The same islands also pin: * the per-articulation shapes of the Jacobian, mass matrix and gravity compensation - accessors, and a symmetric, positive-definite mass matrix. Fixed-base (panda): - ``body_link_jacobian_w`` drops the fixed-root row, so its shape is - ``(N, num_bodies - 1, 6, num_joints)``. Floating-base (anymal): every body row is kept and + accessors in a scene whose articulations have different DoF counts, and a symmetric, + positive-definite mass matrix. Fixed-base: ``body_link_jacobian_w`` drops the fixed-root row, so its + shape is ``(N, num_bodies - 1, 6, num_joints)``. Floating-base: every body row is kept and ``num_base_dofs`` floating-base columns/entries are prepended on the DoF axis, matching the cross-library convention (Pinocchio, Drake, MuJoCo, RBDL, OCS2, iDynTree). A zero-padded - model-wide sizing would surface as non-positive mass-matrix diagonals; + model-wide sizing would surface as wrong shapes or non-positive mass-matrix diagonals; * that every dynamics accessor reflects a manual joint write without a sim step (the FK trigger before ``eval_jacobian``, ``eval_mass_matrix`` and the RNEA pass); * the link-origin Jacobian contract: ``J · q_dot`` must encode the link-origin twist ``v_origin = v_com - omega x (R · body_com_pos_b)``, which the IsaacLab task-space controllers (IK / OSC / RMPFlow) rely on. Newton's ``eval_jacobian`` natively produces COM-referenced rows, so the ground truth reads Newton's per-body twist directly from the - ArticulationView state. + ArticulationView state. The fixed island's links carry center-of-mass offsets. """ - articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) - if ordering_mode == "reversed": - joint_names = PANDA_JOINT_NAMES if articulation_type == "panda" else ANYMAL_C_PHYSX_JOINT_NAMES - articulation_cfg = replace(articulation_cfg, joint_ordering=tuple(reversed(joint_names))) - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device, add_ground_plane=True) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(articulation) < 10 - - replicate(sim.get_clone_plan()) - sim.reset() - assert articulation.is_initialized - assert articulation.is_fixed_base == (articulation_type == "panda") - # 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 if articulation_type == "panda" else 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) - - # -- actuator type - for actuator_name, actuator in articulation.actuators.items(): - is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) - assert getattr(actuator, "is_implicit_model", False) == is_implicit_model_cfg + articulation = scene.articulations[island] + num_articulations = articulation.num_instances + device = scene.device + scene.rest(articulation) + + # Sanity: the islands have different DoF counts, so a regression to model-wide sizing would manifest as wrong + # shapes on the smaller views. + model = SimulationManager.get_model() + island_dofs = [other.num_joints + other.num_base_dofs for other in scene.articulations.values()] + assert model.max_dofs_per_articulation == max(island_dofs) > min(island_dofs) num_dofs = articulation.num_joints + articulation.num_base_dofs num_jacobian_bodies = articulation.num_bodies - 1 if articulation.is_fixed_base else articulation.num_bodies - J = articulation.data.body_link_jacobian_w.torch - assert J.shape == torch.Size((num_articulations, num_jacobian_bodies, 6, num_dofs)), tuple(J.shape) - assert J.dtype == torch.float32 - - g = articulation.data.gravity_compensation_forces.torch - assert g.shape == torch.Size((num_articulations, num_dofs)), tuple(g.shape) - assert g.dtype == torch.float32 - - sim.step() - articulation.update(sim.cfg.dt) - - M = articulation.data.mass_matrix.torch - assert M.shape == torch.Size((num_articulations, num_dofs, num_dofs)), tuple(M.shape) - assert M.dtype == torch.float32 - diag = M.diagonal(dim1=-2, dim2=-1) - assert (diag > 1e-6).all(), f"mass matrix has non-positive diagonal entries: min={diag.min()}" - - # The joint-space inertia is symmetric by construction; asymmetry means a wrong-axis gather or a - # half-populated buffer. OSC inverts ``J M^-1 J^T`` every step, so ``M`` must also be positive-definite. - asym = (M - M.transpose(-1, -2)).abs().max().item() - assert asym < 1e-4, f"|M - M^T|_max = {asym:.3e} — mass matrix is not symmetric" - # A tiny jitter tolerates the float32 eigenvalue floor without masking real non-PD bugs. - eye = torch.eye(M.shape[-1], device=M.device, dtype=M.dtype).expand_as(M) - torch.linalg.cholesky(M + 1e-6 * eye) - - # Read every accessor at the stepped joint state. - J_link_0 = articulation.data.body_link_jacobian_w.torch.clone() - J_com_0 = articulation.data.body_com_jacobian_w.torch.clone() - M_0 = articulation.data.mass_matrix.torch.clone() - g_0 = articulation.data.gravity_compensation_forces.torch.clone() - - # Non-trivial configuration via manual writes (no sim step, so the assert - # compares both quantities at exactly this state): random joint offsets, - # and for floating-base a non-identity root pose. - torch.manual_seed(0) - q = articulation.data.default_joint_pos.torch + 0.3 * torch.randn( - num_articulations, articulation.num_joints, device=device - ) - articulation.write_joint_position_to_sim_index(position=q) - if not articulation.is_fixed_base: - root_pose = articulation.data.default_root_pose.torch.clone() - root_pose[:, 2] += 1.0 - # (x, y, z, w) quaternion — 30 deg roll about x. - root_pose[:, 3:] = torch.tensor([0.2588, 0.0, 0.0, 0.9659], device=device) - articulation.write_root_pose_to_sim_index(root_pose=root_pose) - # Guard against a vacuous identity: if root-pose FK invalidation ever - # regressed, both sides would be evaluated at the stale identity pose and - # agree trivially, voiding the newton#2625 rotated-root coverage. - torch.testing.assert_close(articulation.data.root_link_pose_w.torch, root_pose, atol=1e-5, rtol=0.0) - - # With the FK trigger, forward() refreshes body_q to the written state before each accessor evaluates. - # Without it, body_q stays at the previous state and every accessor returns its previous value. - assert not torch.allclose(J_link_0, articulation.data.body_link_jacobian_w.torch, atol=1e-3), ( - "body_link_jacobian_w did not change after manual joint write; FK trigger likely missing" - ) - assert not torch.allclose(J_com_0, articulation.data.body_com_jacobian_w.torch, atol=1e-3), ( - "body_com_jacobian_w did not change after manual joint write; FK trigger likely missing" - ) - assert not torch.allclose(M_0, articulation.data.mass_matrix.torch, atol=1e-3), ( - "mass_matrix did not change after manual joint write; FK trigger likely missing" - ) - assert not torch.allclose(g_0, articulation.data.gravity_compensation_forces.torch, atol=1e-3), ( - "gravity_compensation_forces did not change after manual joint write; FK trigger likely missing" - ) + with world_gravity(_GRAVITY): + J = articulation.data.body_link_jacobian_w.torch + assert J.shape == torch.Size((num_articulations, num_jacobian_bodies, 6, num_dofs)), tuple(J.shape) + assert J.dtype == torch.float32 + + g = articulation.data.gravity_compensation_forces.torch + assert g.shape == torch.Size((num_articulations, num_dofs)), tuple(g.shape) + assert g.dtype == torch.float32 + assert g.abs().max() > 1e-3, "gravity compensation is all-zero" + + scene.step(island) + + M = articulation.data.mass_matrix.torch + assert M.shape == torch.Size((num_articulations, num_dofs, num_dofs)), tuple(M.shape) + assert M.dtype == torch.float32 + diag = M.diagonal(dim1=-2, dim2=-1) + assert (diag > 1e-6).all(), f"mass matrix has non-positive diagonal entries: min={diag.min()}" + + # The joint-space inertia is symmetric by construction; asymmetry means a wrong-axis gather or a + # half-populated buffer. OSC inverts ``J M^-1 J^T`` every step, so ``M`` must also be positive-definite. + asym = (M - M.transpose(-1, -2)).abs().max().item() + assert asym < 1e-4, f"|M - M^T|_max = {asym:.3e} — mass matrix is not symmetric" + # A tiny jitter tolerates the float32 eigenvalue floor without masking real non-PD bugs. + eye = torch.eye(M.shape[-1], device=M.device, dtype=M.dtype).expand_as(M) + torch.linalg.cholesky(M + 1e-6 * eye) + + # Read every accessor at the stepped joint state. + J_link_0 = articulation.data.body_link_jacobian_w.torch.clone() + J_com_0 = articulation.data.body_com_jacobian_w.torch.clone() + M_0 = articulation.data.mass_matrix.torch.clone() + g_0 = articulation.data.gravity_compensation_forces.torch.clone() + + # Non-trivial configuration via manual writes (no sim step, so the assert + # compares both quantities at exactly this state): random joint offsets, + # and for floating-base a non-identity root pose. + torch.manual_seed(0) + q = articulation.data.default_joint_pos.torch + 0.3 * torch.randn( + num_articulations, articulation.num_joints, device=device + ) + articulation.write_joint_position_to_sim_index(position=q) + if not articulation.is_fixed_base: + root_pose = scene.default_root_pose_w(articulation) + root_pose[:, 2] += 1.0 + # (x, y, z, w) quaternion — 30 deg roll about x. + root_pose[:, 3:] = torch.tensor([0.2588, 0.0, 0.0, 0.9659], device=device) + articulation.write_root_pose_to_sim_index(root_pose=root_pose) + # Guard against a vacuous identity: if root-pose FK invalidation ever + # regressed, both sides would be evaluated at the stale identity pose and + # agree trivially, voiding the newton#2625 rotated-root coverage. + torch.testing.assert_close(articulation.data.root_link_pose_w.torch, root_pose, atol=1e-5, rtol=0.0) + + # With the FK trigger, forward() refreshes body_q to the written state before each accessor evaluates. + # Without it, body_q stays at the previous state and every accessor returns its previous value. + assert not torch.allclose(J_link_0, articulation.data.body_link_jacobian_w.torch, atol=1e-3), ( + "body_link_jacobian_w did not change after manual joint write; FK trigger likely missing" + ) + assert not torch.allclose(J_com_0, articulation.data.body_com_jacobian_w.torch, atol=1e-3), ( + "body_com_jacobian_w did not change after manual joint write; FK trigger likely missing" + ) + assert not torch.allclose(M_0, articulation.data.mass_matrix.torch, atol=1e-3), ( + "mass_matrix did not change after manual joint write; FK trigger likely missing" + ) + assert not torch.allclose(g_0, articulation.data.gravity_compensation_forces.torch, atol=1e-3), ( + "gravity_compensation_forces did not change after manual joint write; FK trigger likely missing" + ) - g_meas = articulation.data.gravity_compensation_forces.torch + g_meas = articulation.data.gravity_compensation_forces.torch - # Independent derivation: generalized gravity load tau_g = sum_b J_lin_b^T (m_b g_w); - # the compensation force is its negation. The COM-referenced Jacobian is exactly the - # right lever arm for a point gravity force acting at each body's COM. - J_com = articulation.data.body_com_jacobian_w.torch # (N, B_jac, 6, D) - masses = articulation.data.body_mass.torch # (N, num_bodies) - if articulation.is_fixed_base: - # jacobi_body_idx == body_idx - 1 for fixed-base (fixed-root row excluded). - masses = masses[:, 1:] - model = SimulationManager.get_model() - gravity_w = wp.to_torch(model.gravity[: model.world_count]) # (num_worlds, 3) - assert gravity_w.shape[0] == num_articulations, "fixture must place one articulation per world" - f_gravity = masses.unsqueeze(-1) * gravity_w.unsqueeze(1) # (N, B_jac, 3) - g_expected = -torch.einsum("nbij,nbi->nj", J_com[:, :, 0:3, :], f_gravity) + # Independent derivation: generalized gravity load tau_g = sum_b J_lin_b^T (m_b g_w); + # the compensation force is its negation. The COM-referenced Jacobian is exactly the + # right lever arm for a point gravity force acting at each body's COM. + J_com = articulation.data.body_com_jacobian_w.torch # (N, B_jac, 6, D) + masses = articulation.data.body_mass.torch # (N, num_bodies) + if articulation.is_fixed_base: + # jacobi_body_idx == body_idx - 1 for fixed-base (fixed-root row excluded). + masses = masses[:, 1:] + gravity_w = wp.to_torch(model.gravity[: model.world_count]) # (num_worlds, 3) + assert gravity_w.shape[0] == num_articulations, "fixture must place one articulation per world" + f_gravity = masses.unsqueeze(-1) * gravity_w.unsqueeze(1) # (N, B_jac, 3) + g_expected = -torch.einsum("nbij,nbi->nj", J_com[:, :, 0:3, :], f_gravity) - torch.testing.assert_close(g_meas, g_expected, atol=1e-2, rtol=1e-3) + torch.testing.assert_close(g_meas, g_expected, atol=1e-2, rtol=1e-3) # Link-origin Jacobian contract. Reproducible non-trivial q_dot — large enough to drive omega # well above the floor where COM offset effects would round into noise. The root is at rest so @@ -3114,8 +1997,8 @@ def test_get_gravity_compensation_forces_matches_jacobian_gravity( articulation.write_root_velocity_to_sim_index(root_velocity=torch.zeros(num_articulations, 6, device=device)) # Refresh kinematics without integrating: for floating bases, an actuator step can # introduce root motion that is intentionally absent from the actuated-only J slice. - sim.forward() - articulation.update(sim.cfg.dt) + scene.sim.forward() + articulation.update(scene.sim.cfg.dt) # body_link_jacobian_w prepends ``num_base_dofs`` floating-base columns; slice past # them so the joint axis aligns with joint_vel (actuated-only). @@ -3143,117 +2026,264 @@ def test_get_gravity_compensation_forces_matches_jacobian_gravity( body_v_com = body_v_com[:, 1:] body_omega = body_omega[:, 1:] c_world = c_world[:, 1:] + elif articulation.body_ordering is not None: + # the view reports backend body order; the public Jacobian rows follow the public body order + body_u2b = list(articulation.body_ordering.user_to_backend_indices) + body_v_com = body_v_com[:, body_u2b] + body_omega = body_omega[:, body_u2b] # Expected v_origin = v_com - omega x c_world. v_origin_expected = body_v_com - torch.cross(body_omega, c_world, dim=-1) - # Tolerance: 5 mm absolute. The COM-offset bug produces a ~3 cm bias - # on the panda hand under the 0.5-rad/s injected qdot, well above - # this floor; numerical noise from kernel ordering stays under 1 mm. + # Tolerance: 5 mm absolute; a missing COM-offset correction is centimeters off at these velocities. torch.testing.assert_close(v_pred_ang, body_omega, atol=5e-3, rtol=1e-2) torch.testing.assert_close(v_pred_lin, v_origin_expected, atol=5e-3, rtol=1e-2) + scene.rest(articulation) + + +def test_body_root_state(scene: _Scene) -> None: + """Body link and center-of-mass states must follow a revolute pendulum with a center-of-mass offset. + + The pendulum turns freely about z; the fixed pivot holds its default state, the angular velocity is the same + in the link and center-of-mass frames, and the link and center-of-mass positions and velocities match the + rigid-rotation ground truth. The center-of-mass offset is set through ``set_coms_index``. + """ + articulation = scene.articulations["pendulum"] + scene.rest(articulation) + num_articulations = articulation.num_instances + num_bodies = articulation.num_bodies + device = scene.device + + # 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") + link_offset = [1.0, 0.0, 0.0] # the offset from CenterPivot to Arm frames + offset = [0.5, 0.0, 0.0] + initial_com = articulation.data.body_com_pos_b.torch.clone() + coms = initial_com.clone() + coms[:, arm_idx] = torch.tensor(offset, device=device) + articulation.set_coms_index(coms=coms) + torch.testing.assert_close(articulation.data.body_com_pos_b.torch, coms) + + articulation.write_joint_velocity_to_sim_index(velocity=torch.tensor([[2.0], [-1.5]], device=device)) + pivot_pos_w = scene.default_root_pose_w(articulation)[:, :3] + default_root_pose = scene.default_root_pose_w(articulation) + default_root_vel = articulation.data.default_root_vel.torch.clone() + for _ in range(50): + scene.step("pendulum") + + # the fixed pivot holds its default state + 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) + + 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 + joint_pos = articulation.data.joint_pos.torch.unsqueeze(-1) + joint_vel = articulation.data.joint_vel.torch.unsqueeze(-1) + + # the angular velocity is the same in the link and center-of-mass 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:]) + + # the arm's link origin circles the pivot at the link offset; the pivot's link does not move + lin_vel_gt = torch.zeros(num_articulations, num_bodies, 3, device=device) + vx = -(link_offset[0]) * joint_vel * torch.sin(joint_pos) + vy = (link_offset[0]) * joint_vel * torch.cos(joint_pos) + lin_vel_gt[:, arm_idx, :] = torch.cat([vx, vy, torch.zeros_like(vx)], dim=-1).squeeze(-2) + torch.testing.assert_close(lin_vel_gt[:, root_idx, :], root_link_vel_w[..., :3], atol=1e-3, rtol=1e-1) + torch.testing.assert_close(lin_vel_gt, body_link_vel_w[..., :3], atol=1e-3, rtol=1e-1) + + # the arm's center of mass circles the pivot at the link offset plus the center-of-mass offset + pos_gt = torch.zeros(num_articulations, num_bodies, 3, device=device) + px = (link_offset[0] + offset[0]) * torch.cos(joint_pos) + py = (link_offset[0] + offset[0]) * torch.sin(joint_pos) + pos_gt[:, arm_idx, :] = torch.cat([px, py, torch.zeros_like(px)], dim=-1).squeeze(-2) + pos_gt += pivot_pos_w.unsqueeze(-2) + 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) + + # the center-of-mass orientation is the link orientation composed with the body-frame offset rotation + 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:]) + + # the pendulum turned far enough for the trigonometric ground truth to discriminate + assert (articulation.data.joint_pos.torch.abs() > 0.5).all() + articulation.set_coms_index(coms=initial_com) + scene.rest(articulation) -@pytest.mark.parametrize("num_articulations", [1]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("articulation_type", ["panda"]) @pytest.mark.isaacsim_ci -def test_get_gravity_compensation_forces_static_equilibrium(sim, num_articulations, device, articulation_type): - """Newton accuracy: ``τ_gc`` must hold the manipulator in static equilibrium. - - Newton-side variant of the PhysX test of the same name (backend parity). - The contract is the EOM identity ``M(q) q̈ + C(q,q̇) q̇ + g(q) = τ_input``. - Setting ``τ_input = g(q)`` at ``q̇ = 0`` gives ``q̈ = 0`` — the arm should - not move. This pins - :attr:`~isaaclab.assets.BaseArticulationData.gravity_compensation_forces` - in isolation: sign errors, frame errors, and DoF-ordering errors all - surface as joint drift, while a controller-level test would have those - bugs averaged out by PD damping. +def test_gravity_vec_w_tracks_model_gravity(scene: _Scene) -> None: + """Per-env mutations to Newton's ``model.gravity`` reach ``GRAVITY_VEC_W`` and ``projected_gravity_b``. + + Regression for the pre-fix snapshot: ``GRAVITY_VEC_W`` used to be env 0's + gravity broadcast to every env, hiding per-env gravity randomization (e.g. + :class:`~isaaclab.envs.mdp.randomize_physics_scene_gravity`). """ - base_cfg = generate_articulation_cfg(articulation_type=articulation_type) - # Replace default Franka actuators with a passthrough implicit actuator - # (stiffness = 0, damping = 0). With both gains zero the effort target - # we set IS the joint torque applied — no PD spring-damper masks the - # gravity-comp signal. Default Franka cfg has stiffness=80 / damping=4 - # which would absorb gravity through PD bias and hide accessor bugs. - cfg = replace(base_cfg, actuators={"all": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=0.0, damping=0.0)}) - # FRANKA_PANDA_CFG has rigid_props.disable_gravity=False already, but be - # defensive — gravity must be ON for τ_gc to have anything to cancel. - cfg = replace(cfg, spawn=replace(cfg.spawn, rigid_props=replace(cfg.spawn.rigid_props, disable_gravity=False))) - - articulation, _ = generate_articulation(cfg, num_articulations, device=device) - replicate(sim.get_clone_plan()) - sim.reset() - assert articulation.is_initialized - - # Force a clean static state: default joint positions, zero velocities. - # ``sim.reset`` may leave residual ``q_dot`` from solver settling under - # gravity, so we pin it explicitly here. - default_q = articulation.data.default_joint_pos.torch.clone() - default_qd = torch.zeros_like(default_q) - articulation.write_joint_position_to_sim_index(position=default_q) - articulation.write_joint_velocity_to_sim_index(velocity=default_qd) - articulation.update(sim.cfg.dt) - - # Default joint pose from FRANKA_PANDA_CFG bends the elbow - # (joint2=-0.569, joint4=-2.81, joint6=3.04) so several links carry a - # gravity load — τ_gc is non-trivial in this configuration. A natural- - # hang pose (all zeros) would produce near-zero τ_gc and make this - # test uninformative. - init_q = articulation.data.joint_pos.torch.clone() - - # Step 100 times applying only τ_gc as joint efforts. - for _ in range(100): - # ``gravity_compensation_forces`` shape is ``(N, num_joints + num_base_dofs)`` - # — leading ``num_base_dofs`` floating-base entries (0 on fixed-base) followed - # by the actuated-joint entries. Slice past the floating-base entries so the - # remaining tensor aligns with ``set_joint_effort_target_index`` (actuated only). - tau_gc = articulation.data.gravity_compensation_forces.torch[:, articulation.num_base_dofs :] - articulation.set_joint_effort_target_index(target=tau_gc) - articulation.write_data_to_sim() - sim.step() - articulation.update(sim.cfg.dt) - - final_q = articulation.data.joint_pos.torch - drift = (final_q - init_q).abs().max() - # Tight bound: 5e-3 rad ≈ 0.3°. Numerical integration over 100 steps will - # accumulate some floor (sub-millirad on Franka), but a sign or frame bug - # in τ_gc produces drift of at least a degree per step on bent-elbow - # poses. This bound separates "correct" from "broken" cleanly. - assert drift < 5e-3, ( - f"max joint drift {drift:.5f} rad after 100 gravity-comp-only steps —" - " τ_gc did not hold static equilibrium. Check sign, DoF ordering, and" - " whether gravity_compensation_forces returns g(q) (positive) or" - " its negation." + articulation = scene.articulations["floating"] + scene.rest(articulation) + num_articulations = articulation.num_instances + device = scene.device + + # GRAVITY_VEC_W must share storage with Newton's per-env gravity array. + model = SimulationManager.get_model() + model_gravity_arr = model.gravity[: model.world_count] + global_gravity = wp.to_torch(model.gravity)[-1].clone() + assert articulation.data.GRAVITY_VEC_W.warp.ptr == model_gravity_arr.ptr + assert articulation.data.GRAVITY_VEC_W.shape == (num_articulations,) + + # Randomize the per-env gravity through the public event term. + new_gravity = torch.tensor( + [[0.1 * (i + 1), 0.2 * (i + 1), -3.0 - float(i)] for i in range(num_articulations)], + device=device, + dtype=torch.float32, ) + with world_gravity((0.0, 0.0, 0.0)): + env = SimpleNamespace(sim=scene.sim, device=device, num_envs=num_articulations) + params = {"gravity_distribution_params": ((0.0, 0.0, 0.0), (0.0, 0.0, 0.0)), "operation": "abs"} + event = randomize_physics_scene_gravity(EventTermCfg(func=randomize_physics_scene_gravity, params=params), env) + for row, values in enumerate(new_gravity.tolist()): + event(env, torch.tensor([row], device=device), (values, values), operation="abs") + # Live view: new per-env values are visible immediately, no invalidation step. + torch.testing.assert_close(articulation.data.GRAVITY_VEC_W.torch, new_gravity) + torch.testing.assert_close(wp.to_torch(model.gravity)[-1], global_gravity) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("articulation_type", ["panda"]) -@pytest.mark.parametrize("gravity_enabled", [False]) -@pytest.mark.isaacsim_ci -def test_franka_osc_tracking_accuracy(sim, device, articulation_type, gravity_enabled): - """Newton-side OSC pose tracking sentinel. - - Mirror of the existing PhysX-side OSC tests in - :mod:`isaaclab.test.controllers.test_operational_space`, scoped to - Franka pose-abs tracking on Newton. This test exercises the full controller-bridge pipeline - (:attr:`~isaaclab.assets.BaseArticulationData.body_link_jacobian_w` + - :attr:`~isaaclab.assets.BaseArticulationData.mass_matrix`) end-to-end - and asserts a loose regression bound rather than a tight correctness - oracle. - - OSC runs with ``gravity_compensation=False`` and scene gravity disabled - so the sentinel isolates the J/M bridge; the gravity-compensation path is - covered by :func:`test_franka_osc_gravity_compensation_precision`. - ``inertial_dynamics_decoupling=True`` - exercises ``mass_matrix`` and the Newton COM-referenced J → - M_b → J product. The actuator PD is zeroed at cfg time so OSC's - joint-effort output is not opposed by ``kp·(target − q)``. + # Recompute the lazily-cached projected_gravity_b without sim.step. Project against the same quat + # buffer the kernel reads so the expectation holds for any orientation. + root_pose = scene.default_root_pose_w(articulation) + # (x, y, z, w) quaternion — 30 deg roll about x. + half_angle = torch.tensor(torch.pi / 12.0, device=device) + root_pose[:, 3:] = torch.stack( + (half_angle.sin(), torch.zeros_like(half_angle), torch.zeros_like(half_angle), half_angle.cos()) + ) + articulation.write_root_pose_to_sim_index(root_pose=root_pose) + articulation.update(scene.sim.cfg.dt) + root_quat = articulation.data.root_link_quat_w.torch + expected = math_utils.quat_apply_inverse(root_quat, torch.nn.functional.normalize(new_gravity, dim=-1)) + torch.testing.assert_close(articulation.data.projected_gravity_b.torch, expected, atol=1e-5, rtol=1e-5) + scene.rest(articulation) + + +def test_body_q_consistent_after_root_write(scene: _Scene) -> None: + """Test that body_q is fresh when collide() runs after a root pose write. + + Regression test for a NaN bug where collide() used stale body_q after env + reset because eval_fk was not called between write_root_pose and collide. + + The scene uses Newton's collision pipeline, and the test patches ``_simulate_physics_only`` to capture + body_q at the moment collide() is called and asserts it matches joint_q. """ - robot, ee_frame_idx, ee_jacobi_idx, arm_joint_ids = _setup_franka_at_home_pose(sim, zero_actuator_pd=True) + device = scene.device + articulation = scene.articulations["floating"] + scene.rest(articulation) + model = SimulationManager.get_model() + jc_starts = model.joint_coord_world_start.numpy() + body_starts = model.body_world_start.numpy() + # The floating island is authored first, so its root owns the first body and joint coordinates of each world. + assert model.body_label[int(body_starts[0])] == "/World/Env_0/Floating/Root" + + scene.step("floating", num_steps=5) + + # Teleport env 0 by 10m (simulating a reset) + new_pose = scene.default_root_pose_w(articulation) + new_pose[0, 0] += 10.0 + new_pose[0, 1] += 5.0 + articulation.write_root_pose_to_sim_index( + root_pose=new_pose[0:1], + env_ids=torch.tensor([0], device=device, dtype=torch.int32), + ) + + # Patch _simulate_physics_only to capture body_q before collide runs + captured = {} + original_simulate = SimulationManager._simulate_physics_only.__func__ + + @classmethod # type: ignore[misc] + def _patched_simulate(cls) -> None: + if cls._needs_collision_pipeline: + bq = wp.to_torch(cls.backend.state_0.body_q) + jq = wp.to_torch(cls.backend.state_0.joint_q) + b0 = int(body_starts[0]) + jc0 = int(jc_starts[0]) + captured["bq_root"] = bq[b0, :3].clone() + captured["jq_root"] = jq[jc0 : jc0 + 3].clone() + original_simulate(cls) + + with patch.object(SimulationManager, "_simulate_physics_only", _patched_simulate): + scene.sim.step() + articulation.update(scene.sim.cfg.dt) + + assert captured, "collision pipeline did not run — _needs_collision_pipeline is False" + + bq_root = captured["bq_root"] + jq_root = captured["jq_root"] + diff = (jq_root - bq_root).abs().max().item() + assert diff < 0.01, ( + f"body_q was stale when collide() ran: diff={diff:.4f}m, jq={jq_root.tolist()}, bq={bq_root.tolist()}" + ) + scene.rest(articulation) + + +def test_joint_position_limit_clamping_respects_logging(scene: _Scene, caplog) -> None: + """Logging level controls reporting without changing clamping or reusing a stale violation count.""" + articulation = scene.articulations["ordered"] + logger = type(articulation).__module__ + original_limits = articulation.data.joint_pos_limits.torch.clone() + original_defaults = articulation.data.default_joint_pos.torch.clone() + limits = torch.zeros_like(original_limits) + limits[..., 1] = 0.5 + try: + for level in (logging.WARNING, logging.INFO): + articulation.data.default_joint_pos.torch.fill_(1.0) + caplog.clear() + with caplog.at_level(level, logger=logger): + articulation.write_joint_position_limit_to_sim_index(limits=limits, warn_limit_violation=False) + assert [record.levelno for record in caplog.records] == ([logging.INFO] if level == logging.INFO else []) + torch.testing.assert_close(articulation.data.default_joint_pos.torch, torch.full_like(limits[..., 1], 0.5)) + caplog.clear() + with caplog.at_level(logging.INFO, logger=logger): + articulation.write_joint_position_limit_to_sim_mask(limits=limits, warn_limit_violation=False) + assert not caplog.records + finally: + articulation.write_joint_position_limit_to_sim_index(limits=original_limits) + articulation.data.default_joint_pos.torch.copy_(original_defaults) + + +## +# Operational-space control. The OSC island's six-DOF chain carries center-of-mass offsets and its joints are +# unpowered, so a wrong Jacobian, mass matrix, gravity force, or DoF ordering pushes the controller's steady-state +# error well past the bounds. +## + + +def _compute_ee_pose_root(robot: Articulation, ee_frame_idx: int) -> tuple[torch.Tensor, torch.Tensor]: + """Return the end-effector position [m] and quaternion ``(x, y, z, w)`` in the root frame.""" + ee_pose_w = robot.data.body_pose_w.torch[:, ee_frame_idx] + root_pose_w = robot.data.root_pose_w.torch + return math_utils.subtract_frame_transforms( + root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_pose_w[:, 0:3], ee_pose_w[:, 3:7] + ) + + +def _compute_jacobian_root_frame(robot: Articulation, ee_jacobi_idx: int, arm_joint_ids: list[int]) -> torch.Tensor: + """Return the end-effector Jacobian sliced to ``arm_joint_ids`` and rotated to the root frame, shape [N, 6, D].""" + jacobian = robot.data.body_link_jacobian_w.torch[:, ee_jacobi_idx, :, arm_joint_ids] + base_rot_matrix = math_utils.matrix_from_quat(math_utils.quat_inv(robot.data.root_pose_w.torch[:, 3:7])) + jacobian[:, :3, :] = torch.bmm(base_rot_matrix, jacobian[:, :3, :]) + jacobian[:, 3:, :] = torch.bmm(base_rot_matrix, jacobian[:, 3:, :]) + return jacobian + - osc = OperationalSpaceController( +def _make_osc(device: str) -> OperationalSpaceController: + """Return a fixed-impedance absolute-pose controller with inertial decoupling and no gravity compensation.""" + return OperationalSpaceController( OperationalSpaceControllerCfg( target_types=["pose_abs"], impedance_mode="fixed", @@ -3263,23 +2293,37 @@ def test_franka_osc_tracking_accuracy(sim, device, articulation_type, gravity_en motion_stiffness_task=500.0, motion_damping_ratio_task=1.0, ), - num_envs=1, + num_envs=NUM_ENVS, device=device, ) - sim.step() - robot.update(sim.cfg.dt) - target_pose_b = _build_relative_pose_target(robot, ee_frame_idx, (0.05, 0.0, 0.0), device) +def _osc_chain(scene: _Scene) -> tuple[Articulation, int, int, list[int]]: + """Step the OSC island once and return it with its end-effector body and Jacobian indices and arm joints.""" + robot = scene.articulations["osc"] + scene.step("osc") + ee_frame_idx = robot.find_bodies("Link_5")[0][0] + # the fixed root has no Jacobian row + return robot, ee_frame_idx, ee_frame_idx - 1, robot.find_joints(["Joint_.*"])[0] + + +def _run_osc( + scene: _Scene, osc: OperationalSpaceController, target_pose_b: torch.Tensor, num_steps: int, *, gravity: bool +) -> tuple[list[float], list[float]]: + """Close the OSC loop for ``num_steps`` steps; return the per-step max position and rotation errors.""" + robot, ee_frame_idx, ee_jacobi_idx, arm_joint_ids = _osc_chain(scene) pos_history: list[float] = [] rot_history: list[float] = [] - for _ in range(800): + for _ in range(num_steps): jacobian_b = _compute_jacobian_root_frame(robot, ee_jacobi_idx, arm_joint_ids) mass_matrix = robot.data.mass_matrix.torch[:, arm_joint_ids, :][:, :, arm_joint_ids] - ee_pos_b, ee_quat_b, _ = _compute_ee_pose_root(robot, ee_frame_idx) + gravity_forces = robot.data.gravity_compensation_forces.torch[:, arm_joint_ids] if gravity else None + ee_pos_b, ee_quat_b = _compute_ee_pose_root(robot, ee_frame_idx) ee_pose_b = torch.cat([ee_pos_b, ee_quat_b], dim=-1) + # OSC's damping term ``kd * ee_vel_b`` needs the end-effector velocity ``J · q_dot``; a zero velocity + # leaves the impedance undamped. joint_vel = robot.data.joint_vel.torch[:, arm_joint_ids] - ee_vel_b = _compute_ee_vel_root(jacobian_b, joint_vel) + ee_vel_b = torch.bmm(jacobian_b, joint_vel.unsqueeze(-1)).squeeze(-1) osc.set_command(target_pose_b, current_ee_pose_b=ee_pose_b) joint_efforts = osc.compute( @@ -3287,148 +2331,87 @@ def test_franka_osc_tracking_accuracy(sim, device, articulation_type, gravity_en current_ee_pose_b=ee_pose_b, current_ee_vel_b=ee_vel_b, mass_matrix=mass_matrix, - gravity=None, + gravity=gravity_forces, ) + robot.actuators.target_command.set_effort_index(value=joint_efforts, joint_ids=arm_joint_ids) + scene.step("osc") - robot.set_joint_effort_target(joint_efforts, joint_ids=arm_joint_ids) - robot.write_data_to_sim() - sim.step() - robot.update(sim.cfg.dt) - - pos_error, rot_error = compute_pose_error(ee_pos_b, ee_quat_b, target_pose_b[:, 0:3], target_pose_b[:, 3:7]) + pos_error, rot_error = math_utils.compute_pose_error( + ee_pos_b, ee_quat_b, target_pose_b[:, 0:3], target_pose_b[:, 3:7] + ) pos_history.append(pos_error.norm(dim=-1).max().item()) rot_history.append(rot_error.norm(dim=-1).max().item()) + return pos_history, rot_history + + +@pytest.mark.isaacsim_ci +def test_osc_tracking_accuracy(scene: _Scene) -> None: + """OSC pose tracking sentinel for the Jacobian and mass-matrix bridge. - pos_mean = _tail_mean(pos_history) - rot_mean = _tail_mean(rot_history) - - # Regression sentinel: assert on tail mean rather than min. With - # ``current_ee_vel_b = J · q_dot`` providing OSC's damping term and - # the actuator PD zeroed, the impedance settles to machine - # precision. The 5 mm bound is a - # bridge regression sentinel: a wrong J, wrong mass matrix, or - # DoF mis-ordering pushes the steady-state error well past it - # because OSC consumes both ``body_link_jacobian_w`` and - # ``mass_matrix`` per step. + OSC runs with ``gravity_compensation=False`` and scene gravity disabled so the sentinel isolates the J/M + bridge; the gravity-compensation path is covered by :func:`test_osc_gravity_compensation_precision`. + ``inertial_dynamics_decoupling=True`` exercises ``mass_matrix`` and the COM-referenced J → M_b → J product. + """ + robot, ee_frame_idx, _, _ = _osc_chain(scene) + ee_pos_b, ee_quat_b = _compute_ee_pose_root(robot, ee_frame_idx) + target_pose_b = torch.cat([ee_pos_b + ee_pos_b.new_tensor((0.05, 0.0, 0.0)), ee_quat_b], dim=-1) + pos_history, rot_history = _run_osc(scene, _make_osc(scene.device), target_pose_b, 150, gravity=False) + + pos_mean = sum(pos_history[-100:]) / 100 + rot_mean = sum(rot_history[-100:]) / 100 + + # Regression sentinel: assert on tail mean rather than min. With ``current_ee_vel_b = J · q_dot`` providing + # OSC's damping term and no joint PD, the impedance settles to machine precision. A wrong J, wrong mass + # matrix, or DoF mis-ordering pushes the steady-state error well past the 5 mm bound because OSC consumes + # both ``body_link_jacobian_w`` and ``mass_matrix`` per step. assert pos_mean < 5e-3, f"OSC pos_mean {pos_mean:.5f} > 5 mm — bridge regression?" assert rot_mean < 5e-2, f"OSC rot_mean {rot_mean:.5f} > 0.05 rad — bridge regression?" -@pytest.mark.parametrize("device", ["cuda:0"]) -@pytest.mark.parametrize("articulation_type", ["panda_fine"]) -@pytest.mark.parametrize("gravity_enabled", [True]) @pytest.mark.isaacsim_ci -def test_franka_osc_gravity_compensation_precision(sim, device, articulation_type, gravity_enabled): +def test_osc_gravity_compensation_precision(scene: _Scene) -> None: """Two-phase EE hold: gravity sag without compensation, tight hold with it. - Same OSC pose-hold loop as :func:`test_franka_osc_tracking_accuracy`, but - with scene and per-body gravity ON and the target pinned to the initial EE - pose, so any steady-state error is pure gravity sag. Phase 1 runs with - ``gravity_compensation=False`` and must sag past a floor; phase 2 flips - ``osc.cfg.gravity_compensation`` — read per :meth:`compute` call, so the - flag is the only variable across phases (the gravity tensor is fetched and - passed in both) — and must recover the hold to under 0.1 mm. - - The floor assertion keeps the test discriminating: if the task stiffness - is ever raised high enough to mask gravity, phase 1 stops clearing the - floor and the test fails loudly instead of silently passing on a - non-discriminating setup. The gravity feed-forward consumes - :attr:`~isaaclab.assets.BaseArticulationData.gravity_compensation_forces` - (Newton RNEA via ``eval_inverse_dynamics_passive``) live in the loop, covering the - FK-staleness refresh on every step of phase 2. - - The task stiffness (500) deliberately matches the PhysX-side OSC - gravity-compensation test for a cross-backend-comparable setup, and the - ``panda_fine`` sim config (4 solver substeps) is load-bearing: at a single - 1/120 substep MJWarp implicitfast mis-integrates the gravity-loaded hold - into a 0.4-1.0 mm limit cycle that never settles (dt-linear, worsened by - higher damping gains, gone in zero gravity — solver integration error, - not a compensation error). With 4 substeps the same per-control-step - torques hold dead-still at ~1 um, quieter than PhysX (OVPhysX) at ~8 um. - The uncompensated sag agrees with PhysX to ~1% (23.7 vs 23.9 mm), - independently validating the gravity forces. Both phases reach a true - steady state here, enforced by tail-half stationarity guards. + Same OSC pose-hold loop as :func:`test_osc_tracking_accuracy`, but with gravity on and the target pinned + to the initial EE pose, so any steady-state error is pure gravity sag. Phase 1 runs with + ``gravity_compensation=False`` and must sag past a floor; phase 2 flips ``osc.cfg.gravity_compensation`` + — read per :meth:`compute` call, so the flag is the only variable across phases (the gravity tensor is + fetched and passed in both) — and must recover the hold to under 0.1 mm. + + The floor assertion keeps the test discriminating: if the task stiffness is ever raised high enough to + mask gravity, phase 1 stops clearing the floor and the test fails loudly instead of silently passing on a + non-discriminating setup. The gravity feed-forward consumes ``gravity_compensation_forces`` (Newton RNEA via + ``eval_inverse_dynamics_passive``) live in the loop, covering the FK-staleness refresh on every step of + phase 2. Both phases reach a true steady state, enforced by tail-half stationarity guards. """ - robot, ee_frame_idx, ee_jacobi_idx, arm_joint_ids = _setup_franka_at_home_pose( - sim, zero_actuator_pd=True, disable_gravity=False - ) - - osc = OperationalSpaceController( - OperationalSpaceControllerCfg( - target_types=["pose_abs"], - impedance_mode="fixed", - inertial_dynamics_decoupling=True, - partial_inertial_dynamics_decoupling=False, - gravity_compensation=False, - motion_stiffness_task=500.0, - motion_damping_ratio_task=1.0, - ), - num_envs=1, - device=device, - ) - - sim.step() - robot.update(sim.cfg.dt) + robot, ee_frame_idx, _, _ = _osc_chain(scene) + osc = _make_osc(scene.device) # Hold the initial EE pose: phase-1 steady-state error is pure gravity sag. - target_pose_b = _build_relative_pose_target(robot, ee_frame_idx, (0.0, 0.0, 0.0), device) + ee_pos_b, ee_quat_b = _compute_ee_pose_root(robot, ee_frame_idx) + target_pose_b = torch.cat([ee_pos_b, ee_quat_b], dim=-1) - def run_phase(num_steps: int) -> list[float]: - pos_history: list[float] = [] - for _ in range(num_steps): - jacobian_b = _compute_jacobian_root_frame(robot, ee_jacobi_idx, arm_joint_ids) - mass_matrix = robot.data.mass_matrix.torch[:, arm_joint_ids, :][:, :, arm_joint_ids] - gravity = robot.data.gravity_compensation_forces.torch[:, arm_joint_ids] - ee_pos_b, ee_quat_b, _ = _compute_ee_pose_root(robot, ee_frame_idx) - ee_pose_b = torch.cat([ee_pos_b, ee_quat_b], dim=-1) - joint_vel = robot.data.joint_vel.torch[:, arm_joint_ids] - ee_vel_b = _compute_ee_vel_root(jacobian_b, joint_vel) - - osc.set_command(target_pose_b, current_ee_pose_b=ee_pose_b) - joint_efforts = osc.compute( - jacobian_b=jacobian_b, - current_ee_pose_b=ee_pose_b, - current_ee_vel_b=ee_vel_b, - mass_matrix=mass_matrix, - gravity=gravity, - ) - robot.set_joint_effort_target(joint_efforts, joint_ids=arm_joint_ids) - robot.write_data_to_sim() - sim.step() - robot.update(sim.cfg.dt) - - pos_error, _ = compute_pose_error(ee_pos_b, ee_quat_b, target_pose_b[:, 0:3], target_pose_b[:, 3:7]) - pos_history.append(pos_error.norm(dim=-1).max().item()) - return pos_history - - def _stationary_tail_mean(history, label): - """Mean of the last 200 samples, asserting the two tail halves agree within 25%. - - The relative check carries a 10 µm absolute floor: at the ~1 µm solver noise - floor of the compensated hold, tail jitter is far below the 0.1 mm verdict - threshold and cannot flip the outcome, so demanding 25% relative agreement - of micrometer-scale means would only add GPU-dependent flakiness. + def _stationary_tail_mean(history: list[float], label: str) -> float: + """Mean of the last 100 samples, asserting the two tail halves agree within 25%. + + The relative check carries a 10 µm absolute floor: at the solver noise floor of the compensated hold, + tail jitter is far below the 0.1 mm verdict threshold and cannot flip the outcome. """ - a = sum(history[-200:-100]) / 100 - b = sum(history[-100:]) / 100 + a = sum(history[-100:-50]) / 50 + b = sum(history[-50:]) / 50 mean = (a + b) / 2.0 assert abs(a - b) < 0.25 * max(mean, 1e-5), ( f"{label} not stationary: tail halves {a:.6f} vs {b:.6f} — extend the phase" ) return mean - hist_off = run_phase(400) - osc.cfg.gravity_compensation = True - hist_on = run_phase(600) + with world_gravity((0.0, 0.0, -9.81)): + hist_off, _ = _run_osc(scene, osc, target_pose_b, 200, gravity=True) + osc.cfg.gravity_compensation = True + hist_on, _ = _run_osc(scene, osc, target_pose_b, 200, gravity=True) pos_off = _stationary_tail_mean(hist_off, "phase-1 sag") pos_on = _stationary_tail_mean(hist_on, "phase-2 hold") - # Re-validated on newton 81cdcfc2 / mujoco-warp 3.10.0.2 with 4 substeps: - # pos_off ~= 0.024, pos_on ~= 1e-6. assert pos_off > 1.2e-2, f"uncompensated sag {pos_off:.5f} < 1.2 cm — setup no longer discriminates gravity" assert pos_on < 1e-4, f"compensated hold {pos_on:.6f} > 0.1 mm — gravity compensation inaccurate" assert pos_on < pos_off / 10.0, f"compensation only improved sag {pos_off:.5f} -> {pos_on:.6f} (<10x)" - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "--maxfail=1"]) diff --git a/source/isaaclab_newton/test/assets/test_articulation_ordering_kernels.py b/source/isaaclab_newton/test/assets/test_articulation_ordering_kernels.py deleted file mode 100644 index fcb4e7759dee..000000000000 --- a/source/isaaclab_newton/test/assets/test_articulation_ordering_kernels.py +++ /dev/null @@ -1,149 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -import warnings - -import numpy as np -import pytest -import warp as wp -from isaaclab_newton.assets.articulation import kernels as articulation_kernels - - -def _selector(values: list[int], dtype: type) -> wp.array: - """Create a CPU Warp selector with the requested integer width.""" - return wp.array(values, dtype=dtype, device="cpu") - - -@pytest.mark.parametrize(("env_dtype", "joint_dtype"), [(wp.int32, wp.int32), (wp.int64, wp.int64)]) -def test_write_joint_limit_data_to_user_and_backend_index_accepts_index_dtypes( - env_dtype: type, joint_dtype: type -) -> None: - """Write partial user-order joint limits into user and backend-order buffers.""" - limits_np = np.asarray([[[1.0, 3.0], [2.0, 5.0]], [[-1.0, 1.0], [4.0, 8.0]]], dtype=np.float32) - limits = wp.array(limits_np, dtype=wp.vec2f, device="cpu") - env_ids = _selector([0, 1], env_dtype) - user_ids = _selector([2, 0], joint_dtype) - user_to_backend = wp.array(np.asarray([1, 2, 0], dtype=np.int32), dtype=wp.int32, device="cpu") - user_lower = wp.zeros((2, 3), dtype=wp.float32, device="cpu") - user_upper = wp.zeros((2, 3), dtype=wp.float32, device="cpu") - user_limits = wp.zeros((2, 3), dtype=wp.vec2f, device="cpu") - backend_lower = wp.zeros((2, 3), dtype=wp.float32, device="cpu") - backend_upper = wp.zeros((2, 3), dtype=wp.float32, device="cpu") - soft_limits = wp.zeros((2, 3), dtype=wp.vec2f, device="cpu") - default_pos = wp.array( - np.asarray([[0.0, 0.0, 4.0], [0.0, 0.0, 0.0]], dtype=np.float32), dtype=wp.float32, device="cpu" - ) - clamped_defaults = wp.zeros(1, dtype=wp.int32, device="cpu") - kernel = articulation_kernels.write_joint_limit_data_to_user_and_backend_index - if env_dtype != wp.int32 or joint_dtype != wp.int32: - kernel = articulation_kernels.write_joint_limit_data_to_user_and_backend_index_kernel(env_ids, user_ids) - - wp.launch( - kernel, - dim=limits.shape, - inputs=[limits, 1.0, env_ids, user_ids, user_to_backend, True], - outputs=[ - user_lower, - user_upper, - user_limits, - backend_lower, - backend_upper, - soft_limits, - default_pos, - clamped_defaults, - ], - device="cpu", - ) - - np.testing.assert_allclose(user_lower.numpy(), np.asarray([[2.0, 0.0, 1.0], [4.0, 0.0, -1.0]], dtype=np.float32)) - np.testing.assert_allclose(user_upper.numpy(), np.asarray([[5.0, 0.0, 3.0], [8.0, 0.0, 1.0]], dtype=np.float32)) - np.testing.assert_allclose(backend_lower.numpy(), np.asarray([[1.0, 2.0, 0.0], [-1.0, 4.0, 0.0]], dtype=np.float32)) - np.testing.assert_allclose(backend_upper.numpy(), np.asarray([[3.0, 5.0, 0.0], [1.0, 8.0, 0.0]], dtype=np.float32)) - np.testing.assert_allclose( - user_limits.numpy(), - np.asarray([[[2.0, 5.0], [0.0, 0.0], [1.0, 3.0]], [[4.0, 8.0], [0.0, 0.0], [-1.0, 1.0]]], dtype=np.float32), - ) - np.testing.assert_allclose(default_pos.numpy(), np.asarray([[2.0, 0.0, 3.0], [4.0, 0.0, 0.0]], dtype=np.float32)) - assert clamped_defaults.numpy()[0] == 3 - - -def test_write_joint_limit_data_to_user_and_backend_mask_reorders_backend_buffers() -> None: - """Write masked user-order joint limits into user and backend-order buffers.""" - limits_np = np.asarray( - [[[1.0, 2.0], [3.0, 6.0], [4.0, 9.0]], [[-2.0, 2.0], [5.0, 7.0], [8.0, 10.0]]], dtype=np.float32 - ) - limits = wp.array(limits_np, dtype=wp.vec2f, device="cpu") - env_mask = wp.array(np.asarray([True, False], dtype=bool), dtype=wp.bool, device="cpu") - user_mask = wp.array(np.asarray([False, True, True], dtype=bool), dtype=wp.bool, device="cpu") - user_to_backend = wp.array(np.asarray([1, 2, 0], dtype=np.int32), dtype=wp.int32, device="cpu") - user_lower = wp.zeros((2, 3), dtype=wp.float32, device="cpu") - user_upper = wp.zeros((2, 3), dtype=wp.float32, device="cpu") - user_limits = wp.zeros((2, 3), dtype=wp.vec2f, device="cpu") - backend_lower = wp.zeros((2, 3), dtype=wp.float32, device="cpu") - backend_upper = wp.zeros((2, 3), dtype=wp.float32, device="cpu") - soft_limits = wp.zeros((2, 3), dtype=wp.vec2f, device="cpu") - default_pos = wp.zeros((2, 3), dtype=wp.float32, device="cpu") - clamped_defaults = wp.zeros(1, dtype=wp.int32, device="cpu") - - wp.launch( - articulation_kernels.write_joint_limit_data_to_user_and_backend_mask, - dim=limits.shape, - inputs=[limits, 1.0, env_mask, user_mask, user_to_backend, True], - outputs=[ - user_lower, - user_upper, - user_limits, - backend_lower, - backend_upper, - soft_limits, - default_pos, - clamped_defaults, - ], - device="cpu", - ) - - np.testing.assert_allclose(user_lower.numpy(), np.asarray([[0.0, 3.0, 4.0], [0.0, 0.0, 0.0]], dtype=np.float32)) - np.testing.assert_allclose(user_upper.numpy(), np.asarray([[0.0, 6.0, 9.0], [0.0, 0.0, 0.0]], dtype=np.float32)) - np.testing.assert_allclose(backend_lower.numpy(), np.asarray([[4.0, 0.0, 3.0], [0.0, 0.0, 0.0]], dtype=np.float32)) - np.testing.assert_allclose(backend_upper.numpy(), np.asarray([[9.0, 0.0, 6.0], [0.0, 0.0, 0.0]], dtype=np.float32)) - np.testing.assert_allclose( - user_limits.numpy(), - np.asarray([[[0.0, 0.0], [3.0, 6.0], [4.0, 9.0]], [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]]], dtype=np.float32), - ) - np.testing.assert_allclose(default_pos.numpy(), np.asarray([[0.0, 3.0, 4.0], [0.0, 0.0, 0.0]], dtype=np.float32)) - assert clamped_defaults.numpy()[0] == 2 - - -@pytest.mark.parametrize("index_dtype", [wp.int32, wp.int64]) -def test_scatter_reset_masks_from_ids_accepts_index_dtype(index_dtype: type) -> None: - """Set exact world and articulation reset masks from nonidentity environment IDs.""" - with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", - message=( - "^'RigidBodyMaterialCfg' is deprecated and will be removed in 5\\.0\\. Use " - "'isaaclab_physx\\.sim\\.spawners\\.materials\\.PhysxRigidBodyMaterialCfg' for PhysX properties, " - "or 'isaaclab\\.sim\\.spawners\\.materials\\.RigidBodyMaterialBaseCfg' for solver-common properties " - "only\\.$" - ), - category=DeprecationWarning, - ) - from isaaclab_newton.physics import newton_manager - - env_ids = _selector([2, 0], index_dtype) - articulation_ids = wp.array(np.asarray([[0, 1], [2, 3], [4, 5]], dtype=np.int32), dtype=int, device="cpu") - world_mask = wp.zeros(3, dtype=wp.bool, device="cpu") - fk_mask = wp.zeros(6, dtype=wp.bool, device="cpu") - - wp.launch( - newton_manager._scatter_reset_masks_from_ids, - dim=(env_ids.shape[0], articulation_ids.shape[1]), - inputs=[env_ids, articulation_ids], - outputs=[world_mask, fk_mask], - device="cpu", - ) - - np.testing.assert_array_equal(world_mask.numpy(), np.asarray([True, False, True])) - np.testing.assert_array_equal(fk_mask.numpy(), np.asarray([True, True, False, False, True, True])) diff --git a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py index 95e183e5163d..025371088a75 100644 --- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py +++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py @@ -3,51 +3,67 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""PD actuator equivalence tests on ANYmal-C (floating-base quadruped). +# ignore private usage of variables warning +# pyright: reportPrivateUsage=none -Compares IsaacLab-native actuators against Newton-native actuators (created -from the same Lab configs via USD authoring) on the Newton physics backend. -Both paths must produce identical joint trajectories within tolerance. +"""Lab-versus-Newton actuator equivalence on locally authored articulations. -Using ANYmal-C — a 12-DOF quadruped on a floating base — exercises the -coordinate-vs-DOF index separation that is critical when free joints shift -the mapping between ``joint_q`` (coordinate layout) and ``joint_qd`` -(DOF layout). +Isaac Lab actuators and Newton-native actuators (created from the same Lab configs via USD authoring) must +produce identical joint trajectories and torque telemetry. Each actuator configuration drives its own +two-environment island in one composite scene per execution path, so the Lab and the Newton paths each build +one scene. The legged island is floating based, which exercises the coordinate-vs-DOF index separation that +free joints introduce between ``joint_q`` and ``joint_qd``, and the cartpole island shares the Newton path's +model-wide actuator adapter with articulations of a different DOF count and base type. -Each test class overrides ANYmal's default actuators with a specific Lab -config (IdealPD, DCMotor, or mixed) and verifies Lab vs Newton equivalence. +The Newton scene stays alive for the tests that drive native actuators directly: gain writes and domain +randomization, target submission, and per-environment state reset. """ +from isaaclab_newton.physics import NewtonCfg + +from isaaclab.sim import SimulationCfg from isaaclab.test.utils import launch_test_simulation -from isaaclab.utils import replace -launch_test_simulation() +launch_test_simulation(SimulationCfg(physics=NewtonCfg())) -import functools import os -import unittest +from collections.abc import Iterator +from dataclasses import dataclass +from types import SimpleNamespace import numpy as np +import pytest import torch import warp as wp from isaaclab_newton.assets import Articulation -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +from isaaclab_newton.assets.articulation.actuator_control import NewtonActuatorControl +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonBuilderCfg from isaaclab_newton.physics import NewtonManager as SimulationManager +from newton import JointTargetMode, JointType, ModelBuilder +from newton.solvers import SolverMuJoCo +from newton_test_utils import NUM_ENVS, local_usd, newton_sim_cfg, spawn_assets import isaaclab.sim as sim_utils -from isaaclab.actuators import DCMotorCfg, DelayedPDActuatorCfg, IdealPDActuatorCfg +from isaaclab.actuators import ( + ActuatorBaseCfg, + DCMotorCfg, + DelayedPDActuatorCfg, + IdealPDActuatorCfg, + ImplicitActuator, + ImplicitActuatorCfg, +) +from isaaclab.actuators.actuator_net_cfg import ActuatorNetLSTMCfg, ActuatorNetMLPCfg +from isaaclab.actuators.actuator_pd_cfg import RemotizedPDActuatorCfg from isaaclab.actuators.newton import read_group_parameter from isaaclab.actuators.newton.kernels import sync_torque_telemetry -from isaaclab.assets import AssetBaseCfg -from isaaclab.cloner import CloneCfg, clone_plan_from_env_0, replicate -from isaaclab.sim import SimulationCfg, build_simulation_context +from isaaclab.assets import ArticulationCfg +from isaaclab.sim import SimulationContext, build_simulation_context +from isaaclab.test.utils import DeviceScope, test_devices from isaaclab.test.utils.actuator_equivalence import ( CARTPOLE_EXPLICIT_ACTUATORS, - DELAYED_PD_ACTUATORS, IDEAL_PD_ACTUATORS, IMPLICIT_ONLY_ACTUATORS, MIXED_WITH_IMPLICIT_ACTUATORS, - ActuatorStateResetBase, EquivalenceAssertionsMixin, MockEnv, build_dr_term, @@ -55,32 +71,38 @@ make_dummy_mlp_checkpoint, ) from isaaclab.test.utils.articulation_ordering import assert_articulation_ordering_trace_matches +from isaaclab.utils import replace -from isaaclab_assets import ANYMAL_C_CFG from isaaclab_assets.robots.spot import joint_parameter_lookup as SPOT_KNEE_LOOKUP +pytestmark = [pytest.mark.integration, pytest.mark.kitless] + # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- -NUM_ENVS = 2 NUM_STEPS = 10 DT = 1.0 / 120.0 TARGET_OFFSET = 0.1 # [rad] added to initial joint positions -NEWTON_CFG = NewtonCfg( - solver_cfg=MJWarpSolverCfg( + +def _solver_cfg() -> MJWarpSolverCfg: + """Return the MJWarp solver configuration shared by the actuator scenes.""" + return MJWarpSolverCfg( njmax=500, nconmax=500, ls_iterations=20, cone="pyramidal", impratio=1, integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - use_cuda_graph=False, -) + ) + + +_LEG_REORDERED_JOINT_NAMES = ("LF_KFE", "RH_HAA", "LF_HAA", "RH_KFE", "LF_HFE", "RH_HFE") +"""A public order that permutes ``floating_two_leg.usda``'s backend joints without being its own inverse.""" + +_LEG_INIT_STATE = ArticulationCfg.InitialStateCfg(joint_pos={".*HAA": 0.0, ".*HFE": 0.4, ".*KFE": -0.8}) +"""A bent-knee pose inside the Spot knee lookup table.""" # The PD demand (~kp * TARGET_OFFSET = 4 N·m) exceeds the 2 N·m saturation effort, so the DC-motor # torque-speed clamp binds and Newton must author it for the paths to match. @@ -108,749 +130,861 @@ ), } +# RemotizedPD (Spot knee lookup) on KFE with IdealPD on HAA/HFE. +REMOTIZED_PD_ACTUATORS = { + "hips": IdealPDActuatorCfg( + joint_names_expr=[".*HAA", ".*HFE"], + stiffness=40.0, + damping=5.0, + actuator_effort_limit=80.0, + ), + "knees": RemotizedPDActuatorCfg( + joint_names_expr=[".*KFE"], + stiffness=60.0, + damping=1.5, + actuator_effort_limit=80.0, + min_delay=3, + max_delay=3, + joint_parameter_lookup=SPOT_KNEE_LOOKUP, + ), +} + + +def _neural_actuators(mlp_path: str, lstm_path: str) -> dict[str, ActuatorBaseCfg]: + """MLP on HAA, LSTM on HFE, and IdealPD on KFE.""" + return { + "mlp_legs": ActuatorNetMLPCfg( + joint_names_expr=[".*HAA"], + network_file=mlp_path, + saturation_effort=120.0, + actuator_effort_limit=80.0, + actuator_velocity_limit=7.5, + pos_scale=-1.0, + vel_scale=1.0, + torque_scale=1.0, + input_order="pos_vel", + input_idx=[0, 1, 2], + ), + "lstm_legs": ActuatorNetLSTMCfg( + joint_names_expr=[".*HFE"], + network_file=lstm_path, + saturation_effort=120.0, + actuator_effort_limit=80.0, + actuator_velocity_limit=7.5, + ), + "pd_legs": IdealPDActuatorCfg( + joint_names_expr=[".*KFE"], + stiffness=40.0, + damping=5.0, + actuator_effort_limit=80.0, + ), + } + + +@dataclass +class _Island: + """One actuator configuration under test and the commands its rollout sends.""" + + actuators: dict[str, ActuatorBaseCfg] + usd: str = "floating_two_leg.usda" + joint_ordering: tuple[str, ...] | None = None + feedforward: float | None = None + """Constant per-DOF feedforward effort target, when not None.""" + ramp_targets: bool = False + """Whether to ramp the position target over the rollout so command delays are observable.""" + permutation_sensitive_commands: bool = False + """Whether to command distinct position, velocity, and effort values by physical joint name.""" + newton_only: bool = False + """Whether only the Newton path runs the island.""" + torque_atol: float = EquivalenceAssertionsMixin.torque_atol + """Absolute tolerance of the torque-telemetry oracles [N·m].""" + + +def _islands(mlp_path: str | None = None, lstm_path: str | None = None) -> dict[str, _Island]: + """Return the actuator islands, with the neural island when the network checkpoints are given.""" + islands = { + "ideal": _Island(IDEAL_PD_ACTUATORS), + "ideal_reordered": _Island(IDEAL_PD_ACTUATORS, joint_ordering=_LEG_REORDERED_JOINT_NAMES), + "cartpole": _Island(CARTPOLE_EXPLICIT_ACTUATORS, usd="fixed_cartpole.usda"), + "dc_motor": _Island(SATURATING_DC_MOTOR_ACTUATORS), + "mixed": _Island(MIXED_WITH_IMPLICIT_ACTUATORS), + "implicit_feedforward": _Island(IMPLICIT_ONLY_ACTUATORS, feedforward=2.0, torque_atol=0.5), + "delayed": _Island(FIXED_DELAYED_PD_ACTUATORS, ramp_targets=True), + "remotized": _Island(REMOTIZED_PD_ACTUATORS, ramp_targets=True), + # The implicit hip group reads its torque telemetry from backend-order effort buffers, so the + # permutation-sensitive feedforward effort also checks that gather. + "mixed_permuted": _Island(MIXED_WITH_IMPLICIT_ACTUATORS, permutation_sensitive_commands=True, newton_only=True), + "mixed_permuted_reordered": _Island( + MIXED_WITH_IMPLICIT_ACTUATORS, + joint_ordering=_LEG_REORDERED_JOINT_NAMES, + permutation_sensitive_commands=True, + newton_only=True, + ), + } + if mlp_path is not None: + islands["neural"] = _Island(_neural_actuators(mlp_path, lstm_path), newton_only=True) + return islands + + # --------------------------------------------------------------------------- # Simulation runner # --------------------------------------------------------------------------- -def _run_simulation( - actuators: dict, +@dataclass +class _Run: + """Articulations of one execution path and their recorded rollouts.""" + + sim: SimulationContext + articulations: dict[str, Articulation] + results: dict[str, dict] + + +def _island_cfg(name: str, island: _Island, index: int) -> ArticulationCfg: + """Return the articulation of one island, offset along y by its index.""" + init_state = ArticulationCfg.InitialStateCfg() if island.usd != "floating_two_leg.usda" else _LEG_INIT_STATE + init_state = replace(init_state, pos=(0.0, 3.0 * index, 1.0)) + return ArticulationCfg( + prim_path=f"/World/Env_[^/]*/{name}", + spawn=local_usd(island.usd), + init_state=init_state, + actuators=island.actuators, + joint_ordering=island.joint_ordering, + ) + + +def _run( + sim: SimulationContext, + islands: dict[str, _Island], use_newton_actuators: bool, *, dt: float = DT, - newton_cfg: NewtonCfg = NEWTON_CFG, num_steps: int = NUM_STEPS, decimation: int = 1, - feedforward: float | None = None, - joint_ordering: tuple[str, ...] | None = None, - permutation_sensitive_commands: bool = False, - ramp_targets: bool = False, -) -> dict: - """Run ANYmal-C and return recorded trajectories + telemetry. - - Always records ``joint_pos``, ``joint_vel``, ``computed_effort``, and - ``applied_effort`` so callers don't need a separate "with telemetry" - runner. Optionally applies a constant per-DOF feedforward effort target. +) -> _Run: + """Spawn one island per actuator configuration, roll every island out, and record trajectories and telemetry. - Args: - actuators: Actuator config dict overriding ANYmal's defaults. - use_newton_actuators: Use Newton-native actuators when ``True``. - dt: Physics timestep [s]. - newton_cfg: Newton physics configuration. - num_steps: Number of policy-level steps. - decimation: Actuator steps per policy step (Newton's CUDA-graph - d-loop is used when all-graphable; otherwise an explicit Python - inner loop). - feedforward: When not ``None``, set a constant per-DOF feedforward - effort target. Used by the implicit-FF equivalence test. - joint_ordering: Optional explicit public joint-name order. - permutation_sensitive_commands: Whether to command distinct position, velocity, and effort values by - physical joint name. - ramp_targets: Whether to ramp the position target over the rollout so command delays are observable. - - Returns: - Recorded joint-name metadata, commands, public trajectories and torque telemetry, backend-order - adapter effort traces, and the Newton actuators built from the model. + Records ``joint_pos``, ``joint_vel``, ``computed_effort``, and ``applied_effort`` per step, the + backend-order adapter efforts on the Newton path, joint-name metadata, the commands, and the Newton actuators + that execute each group. """ - sim_cfg = SimulationCfg(dt=dt, physics=newton_cfg, use_newton_actuators=use_newton_actuators) - with build_simulation_context( - device="cuda:0", - gravity_enabled=True, - add_ground_plane=True, - sim_cfg=sim_cfg, - ) as sim: - sim._app_control_on_stop_handle = None - sim_utils.create_prim("/World/Env_0", "Xform") - art_cfg = replace( - ANYMAL_C_CFG, actuators=actuators, prim_path="/World/Env_[^/]*/Robot", joint_ordering=joint_ordering - ) - clone_plan_from_env_0( - CloneCfg(clone_template="/World/Env_{}"), - (art_cfg, AssetBaseCfg(prim_path="/World/defaultGroundPlane")), - NUM_ENVS, - 3.0, - positions=np.asarray([(i * 3.0, 0.0, 0.0) for i in range(NUM_ENVS)]), - ) - articulation = Articulation(art_cfg) - replicate(sim.get_clone_plan()) - sim.reset() + articulations = spawn_assets( + {name: _island_cfg(name, island, index) for index, (name, island) in enumerate(islands.items())} + ) + sim.reset() + for articulation in articulations.values(): assert articulation.is_initialized + # Start from the configured joint state, as an environment reset would. + articulation.write_joint_state_to_sim_index( + position=articulation.data.default_joint_pos.torch.clone(), + velocity=articulation.data.default_joint_vel.torch.clone(), + ) # Reset samples the Lab-path actuator delay; without it the Lab lag stays at zero. articulation.reset() - if use_newton_actuators and decimation > 1: - SimulationManager.set_decimation(decimation) - - handles_dec = ( - use_newton_actuators - and decimation > 1 - and SimulationManager._is_all_graphable() - and SimulationManager._decimation > 1 - ) + if use_newton_actuators and decimation > 1: + SimulationManager.set_decimation(decimation) + handles_dec = ( + use_newton_actuators + and decimation > 1 + and SimulationManager._is_all_graphable() + and SimulationManager._decimation > 1 + ) + results = {} + for name, island in islands.items(): + articulation = articulations[name] joint_names = tuple(articulation.joint_names) backend_joint_names = tuple(articulation.backend_joint_names) installed_ordering = articulation.joint_ordering - joint_ordering_state = ( - None - if installed_ordering is None - else { - "user_names": joint_names, - "backend_names": backend_joint_names, - "user_to_backend_indices": installed_ordering.user_to_backend_indices, - "backend_to_user_indices": installed_ordering.backend_to_user_indices, - "is_identity": False, - } - ) - init_pos = wp.to_torch(articulation.data.joint_pos).clone() - if permutation_sensitive_commands: - scale_by_name = {name: index + 1 for index, name in enumerate(backend_joint_names)} + init_pos = articulation.data.joint_pos.torch.clone() + if island.permutation_sensitive_commands: + scale_by_name = {joint_name: index + 1 for index, joint_name in enumerate(backend_joint_names)} joint_scale = torch.tensor( - [scale_by_name[name] for name in joint_names], - device=articulation.device, - dtype=init_pos.dtype, - ).unsqueeze(0) - joint_scale = joint_scale.expand_as(init_pos) + [scale_by_name[joint_name] for joint_name in joint_names], device=sim.device, dtype=init_pos.dtype + ).expand_as(init_pos) target_pos = init_pos + 0.01 * joint_scale target_vel = 0.001 * joint_scale effort_target = 0.1 * joint_scale else: target_pos = init_pos + TARGET_OFFSET target_vel = torch.zeros_like(init_pos) - effort_target = None if feedforward is None else torch.full_like(init_pos, feedforward) - - articulation.set_joint_position_target_index(target=target_pos) - articulation.set_joint_velocity_target_index(target=target_vel) + effort_target = None if island.feedforward is None else torch.full_like(init_pos, island.feedforward) + commands = articulation.actuators.target_command + commands.set_position_index(value=target_pos) + commands.set_velocity_index(value=target_vel) if effort_target is not None: - articulation.set_joint_effort_target_index(target=effort_target) - - recorded_pos, recorded_vel = [], [] - recorded_computed_effort, recorded_applied_effort = [], [] - recorded_adapter_applied = [] - for step in range(num_steps): - if ramp_targets: - target_pos = init_pos + TARGET_OFFSET * (step + 1) / num_steps - articulation.set_joint_position_target_index(target=target_pos) - if handles_dec: + commands.set_effort_index(value=effort_target) + results[name] = { + "joint_names": joint_names, + "backend_joint_names": backend_joint_names, + "joint_ordering": ( + None + if installed_ordering is None + else { + "user_names": joint_names, + "backend_names": backend_joint_names, + "user_to_backend_indices": installed_ordering.user_to_backend_indices, + "backend_to_user_indices": installed_ordering.backend_to_user_indices, + "is_identity": False, + } + ), + "adapter_joint_names": backend_joint_names, + "init_pos": init_pos, + "target_pos": target_pos.clone(), + "target_vel": target_vel.clone(), + "effort_target": None if effort_target is None else effort_target.clone(), + "joint_pos": [], + "joint_vel": [], + "computed_effort": [], + "applied_effort": [], + "adapter_applied_effort": [], + } + + for step in range(num_steps): + for name, island in islands.items(): + if island.ramp_targets: + target_pos = results[name]["init_pos"] + TARGET_OFFSET * (step + 1) / num_steps + articulations[name].actuators.target_command.set_position_index(value=target_pos) + results[name]["target_pos"] = target_pos.clone() + for _ in range(1 if handles_dec else decimation): + for articulation in articulations.values(): articulation.write_data_to_sim() - sim.step() - articulation.update(dt * decimation) - else: - for _ in range(decimation): - articulation.write_data_to_sim() - sim.step() - articulation.update(dt) - recorded_pos.append(wp.to_torch(articulation.data.joint_pos).clone()) - recorded_vel.append(wp.to_torch(articulation.data.joint_vel).clone()) - recorded_computed_effort.append(articulation.actuators.computed_effort.torch.clone()) - recorded_applied_effort.append(articulation.actuators.applied_effort.torch.clone()) + sim.step() + for articulation in articulations.values(): + articulation.update(dt * decimation if handles_dec else dt) + for name, articulation in articulations.items(): + result = results[name] + result["joint_pos"].append(articulation.data.joint_pos.torch.clone()) + result["joint_vel"].append(articulation.data.joint_vel.torch.clone()) + result["computed_effort"].append(articulation.actuators.computed_effort.torch.clone()) + result["applied_effort"].append(articulation.actuators.applied_effort.torch.clone()) if use_newton_actuators: - recorded_adapter_applied.append(wp.to_torch(articulation.data._sim_bind_joint_effort).clone()) + result["adapter_applied_effort"].append(wp.to_torch(articulation.data._sim_bind_joint_effort).clone()) + for name, articulation in articulations.items(): actuator_info = [] if use_newton_actuators: - for act in SimulationManager.get_model().actuators: - actuator_info.append( - { - "controller_type": type(act.controller).__name__, - "clamping_types": sorted(type(c).__name__ for c in (act.clamping or [])), - "has_delay": act.delay is not None, - } - ) - - return { - "joint_names": joint_names, - "backend_joint_names": backend_joint_names, - "joint_ordering": joint_ordering_state, - "adapter_joint_names": backend_joint_names, - "joint_pos": recorded_pos, - "joint_vel": recorded_vel, - "computed_effort": recorded_computed_effort, - "applied_effort": recorded_applied_effort, - "adapter_applied_effort": recorded_adapter_applied, - "target_pos": target_pos.clone(), - "target_vel": target_vel.clone(), - "effort_target": None if effort_target is None else effort_target.clone(), - "actuator_info": actuator_info, - } - - -def test_newton_actuator_rollout_matches_reversed_joint_ordering() -> None: - """Match Newton-backend actuator traces under reversed public joint ordering. - - The implicit hip group reads its torque telemetry from backend-order effort buffers, so the - permutation-sensitive feedforward effort also checks that gather. + for group_name in islands[name].actuators: + if group_name not in articulation.actuators._native_group_names: + continue + group_actuators = articulation.actuators[group_name] + if not isinstance(group_actuators, tuple): + group_actuators = (group_actuators,) + for act in group_actuators: + actuator_info.append( + { + "group": group_name, + "controller_type": type(act.controller).__name__, + "clamping_types": sorted(type(c).__name__ for c in (act.clamping or [])), + "has_delay": act.delay is not None, + } + ) + results[name]["actuator_info"] = actuator_info + return _Run(sim=sim, articulations=articulations, results=results) + + +def _record_lab_state_reset(run: _Run) -> dict[str, tuple[torch.Tensor, torch.Tensor]]: + """Reset one Lab-path environment of the delayed island and record each group's effort and PD demand. + + Reset environments must accept fresh commands while the remaining environments retain their delay history. """ - identity_result = _run_simulation( - MIXED_WITH_IMPLICIT_ACTUATORS, - use_newton_actuators=True, - permutation_sensitive_commands=True, - ) - requested_joint_names = tuple(reversed(identity_result["joint_names"])) - reversed_result = _run_simulation( - MIXED_WITH_IMPLICIT_ACTUATORS, - use_newton_actuators=True, - joint_ordering=requested_joint_names, - permutation_sensitive_commands=True, - ) - - installed_ordering = reversed_result["joint_ordering"] - assert installed_ordering is not None - assert not installed_ordering["is_identity"] - assert_articulation_ordering_trace_matches(identity_result, reversed_result, requested_joint_names) - - -# --------------------------------------------------------------------------- -# Base test class -# --------------------------------------------------------------------------- - - -class _EquivalenceTestBase(EquivalenceAssertionsMixin, unittest.TestCase): - """Base for Lab-vs-Newton equivalence tests. - - Subclasses set ``actuators`` to the config under test. ``setUpClass`` - runs the simulation with both ``use_newton_actuators=False`` (Lab path) - and ``True`` (Newton path) and stores the results. The ``test_*_match`` - oracles come from :class:`EquivalenceAssertionsMixin`. - """ - - __test__ = False - actuators: dict = {} - feedforward: float | None = None - dt: float = DT - newton_cfg: NewtonCfg = NEWTON_CFG - num_steps: int = NUM_STEPS - decimation: int = 1 - ramp_targets: bool = False - - @classmethod - def setUpClass(cls): - kwargs = dict( - feedforward=cls.feedforward, - dt=cls.dt, - newton_cfg=cls.newton_cfg, - num_steps=cls.num_steps, - decimation=cls.decimation, - ramp_targets=cls.ramp_targets, + articulation = run.articulations["delayed"] + commands = articulation.actuators.target_command + old_target = commands.position.torch.clone() + # Episode reset samples the configured lag; construction alone leaves actuator lag at zero. + articulation.reset() + commands.set_position_index(value=old_target) + articulation.write_data_to_sim() + new_target = old_target + 0.02 + articulation.reset(env_ids=torch.tensor([0], device=articulation.device)) + commands.set_position_index(value=new_target) + articulation.write_data_to_sim() + expected_target = old_target.clone() + expected_target[0] = new_target[0] + recorded = {} + for name, actuator in articulation.actuators.items(): + joints = actuator.joint_indices + demand = ( + actuator.stiffness * (expected_target[:, joints] - articulation.data.joint_pos.torch[:, joints]) + - actuator.damping * articulation.data.joint_vel.torch[:, joints] ) - cls.lab_result = _run_simulation(cls.actuators, use_newton_actuators=False, **kwargs) - cls.newton_result = _run_simulation(cls.actuators, use_newton_actuators=True, **kwargs) - - -# --------------------------------------------------------------------------- -# Equivalence tests with different actuator types -# --------------------------------------------------------------------------- - - -class TestDCMotorEquivalence(_EquivalenceTestBase): - """Saturating DCMotor actuator on all 12 joints: Lab vs Newton.""" - - __test__ = True - actuators = SATURATING_DC_MOTOR_ACTUATORS - - def test_dc_motor_clamp_binds(self): - for result in (self.lab_result, self.newton_result): - assert any( - not torch.allclose(applied, computed) - for applied, computed in zip(result["applied_effort"], result["computed_effort"]) - ), "the DC-motor clamp never bound, so the equivalence cannot detect missing clamping" + recorded[name] = (actuator.computed_effort.clone(), demand) + return recorded + + +@pytest.fixture(scope="module", params=test_devices(DeviceScope.CUDA)) +def device(request: pytest.FixtureRequest) -> str: + """Simulation device of the actuator scenes.""" + return request.param + + +@pytest.fixture(scope="module") +def lab_run(device: str) -> dict: + """Roll out every shared island on the Isaac Lab actuator path.""" + islands = {name: island for name, island in _islands().items() if not island.newton_only} + sim_cfg = newton_sim_cfg(device, dt=DT, use_newton_actuators=False, solver_cfg=_solver_cfg()) + with build_simulation_context(sim_cfg=sim_cfg) as sim: + run = _run(sim, islands, use_newton_actuators=False) + return {"results": run.results, "state_reset": _record_lab_state_reset(run)} + + +@pytest.fixture(scope="module") +def decimated_runs(device: str) -> dict[str, dict]: + """Roll out the RemotizedPD island with decimation 2 and CUDA-graph capture on both execution paths.""" + islands = {"remotized": _Island(REMOTIZED_PD_ACTUATORS, ramp_targets=True)} + runs = {} + for use_newton_actuators in (False, True): + sim_cfg = newton_sim_cfg( + device, + dt=1.0 / 100.0, + use_newton_actuators=use_newton_actuators, + solver_cfg=_solver_cfg(), + num_substeps=2, + use_cuda_graph=True, + ) + with build_simulation_context(sim_cfg=sim_cfg) as sim: + run = _run(sim, islands, use_newton_actuators, dt=1.0 / 100.0, num_steps=5, decimation=2) + runs["newton" if use_newton_actuators else "lab"] = run.results["remotized"] + return runs -class TestMixedWithImplicitEquivalence(_EquivalenceTestBase): - """Implicit HAA + IdealPD HFE + DCMotor KFE: Lab vs Newton. +@pytest.fixture(scope="module") +def newton_run(device: str, lab_run: dict, decimated_runs: dict) -> Iterator[_Run]: + """Roll out every island on the Newton actuator path and keep the scene alive. - Verifies that implicit actuators (handled by the physics engine's - built-in joint drives) coexist correctly with explicit Newton actuators. + Depends on the Lab-path and decimated runs so their scenes are built and closed first. """ + mlp_path = make_dummy_mlp_checkpoint() + lstm_path = make_dummy_lstm_checkpoint() + sim_cfg = newton_sim_cfg(device, dt=DT, use_newton_actuators=True, solver_cfg=_solver_cfg()) + try: + with build_simulation_context(sim_cfg=sim_cfg) as sim: + yield _run(sim, _islands(mlp_path, lstm_path), use_newton_actuators=True) + finally: + os.unlink(mlp_path) + os.unlink(lstm_path) + + +def _assert_equivalent( + lab_result: dict, newton_result: dict, torque_atol: float = EquivalenceAssertionsMixin.torque_atol +) -> None: + """Run the shared trajectory and telemetry oracles on one island's Lab and Newton rollouts. - __test__ = True - actuators = MIXED_WITH_IMPLICIT_ACTUATORS - - -# --------------------------------------------------------------------------- -# Implicit + non-zero feedforward effort target -# --------------------------------------------------------------------------- - - -class TestImplicitWithFeedforwardEquivalence(_EquivalenceTestBase): - """Implicit-only actuators with a non-zero feedforward effort target. - - Verifies that the user's FF effort lands additively on top of the - simulator's joint-drive PD identically on both Lab and Newton paths. + Args: + lab_result: Rollout recorded on the Isaac Lab actuator path. + newton_result: Rollout recorded on the Newton actuator path. + torque_atol: Absolute tolerance of the torque-telemetry oracles [N·m]. """ - - __test__ = True - actuators = IMPLICIT_ONLY_ACTUATORS - feedforward = 2.0 - torque_atol = 0.5 + oracle = EquivalenceAssertionsMixin() + oracle.torque_atol = torque_atol + oracle.lab_result = lab_result + oracle.newton_result = newton_result + oracle.test_joint_positions_match() + oracle.test_joint_velocities_match() + oracle.test_applied_effort_match() + oracle.test_computed_effort_match() # --------------------------------------------------------------------------- -# Multi-articulation Newton scene (regression test for class-attr clobber) +# Adaptation on a model builder and on Warp arrays; these step no simulation and run before the scenes exist # --------------------------------------------------------------------------- -def _run_anymal_and_cartpole(use_newton_actuators: bool, *, num_steps: int = NUM_STEPS) -> dict: - """Spawn ANYmal-C + Cartpole per env (different DOF counts, different base types).""" - from isaaclab_assets import CARTPOLE_CFG # noqa: PLC0415 +class CustomDrive(ImplicitActuator): + """Implicit actuator with a class name that does not encode its execution type.""" + + +def _make_target_mode_builder( + monkeypatch, joint_names: list[str], modes: list[JointTargetMode], stiffness: list[float], damping: list[float] +) -> ModelBuilder: + """Build a zero-gain articulated model builder for target-mode tests.""" + sim = object.__new__(sim_utils.SimulationContext) + sim.cfg = SimpleNamespace(physics=NewtonCfg()) + sim._backend_registry = [] + monkeypatch.setattr(sim_utils.SimulationContext, "instance", lambda: sim) + builder = sim.get_or_create_backend(NewtonBuilderCfg(physics_cfg=sim.cfg.physics)) + inertia = wp.mat33(1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0) + parent = -1 + joint_ids = [] + for joint_name in joint_names: + link = builder.add_link(mass=1.0, inertia=inertia, label=f"/World/Env_0/Robot/{joint_name}_link") + joint_ids.append( + builder.add_joint_revolute( + parent, + link, + target_ke=0.0, + target_kd=0.0, + label=f"/World/Env_0/Robot/{joint_name}", + ) + ) + parent = link + builder.add_articulation(joint_ids, label="/World/Env_0/Robot") + builder.articulation_label = ["/World/Env_0/Robot"] + builder.joint_target_mode = [int(mode) for mode in modes] + builder.joint_target_ke = stiffness + builder.joint_target_kd = damping + return builder + + +@pytest.mark.parametrize( + ("actuator_cfg", "expected_native_groups"), + [ + (ImplicitActuatorCfg(joint_names_expr=["joint"], stiffness=10.0, damping=1.0), set()), + (IdealPDActuatorCfg(joint_names_expr=["joint"], stiffness=None, damping=None), {"explicit"}), + ], + ids=["implicit", "explicit"], +) +def test_prepare_native_actuators_activates_only_explicit_groups(monkeypatch, actuator_cfg, expected_native_groups): + """Keep implicit-only articulations on the solver-drive path and leave solver gains untouched. - sim_cfg = SimulationCfg(dt=DT, physics=NEWTON_CFG, use_newton_actuators=use_newton_actuators) - with build_simulation_context( - device="cuda:0", - gravity_enabled=True, - add_ground_plane=True, - sim_cfg=sim_cfg, - ) as sim: - sim._app_control_on_stop_handle = None + Explicit groups activate the Newton-actuator path without writing gains; collection construction resolves + the actuator defaults later. + """ + activation_calls = [] + gain_writes = [] + articulation = SimpleNamespace( + _sim_cfg=SimpleNamespace(use_newton_actuators=True), + device="cpu", + find_joints=lambda _: ([0], ["joint"]), + write_joint_stiffness_to_sim_index=lambda **_: gain_writes.append("stiffness"), + write_joint_damping_to_sim_index=lambda **_: gain_writes.append("damping"), + ) + monkeypatch.setattr(SimulationManager, "activate_newton_actuator_path", lambda: activation_calls.append(True)) + + control = NewtonActuatorControl(articulation) + group_name = "explicit" if expected_native_groups else "implicit" + native_groups = control.prepare_native_actuators(collection=None, actuator_cfgs={group_name: actuator_cfg}) + + assert native_groups == expected_native_groups + assert gain_writes == [] + if expected_native_groups: + assert activation_calls == [True] + else: + assert not control.native_actuator_path_active + assert not articulation._has_newton_actuators + assert activation_calls == [] + + +@pytest.mark.parametrize( + ("actuator_cfg", "expected_mode", "expected_actuator_indices"), + [ + ( + ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=10.0, damping=0.0), + JointTargetMode.POSITION, + [0, 1], + ), + ( + ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=0.0, damping=2.0), + JointTargetMode.VELOCITY, + [-2, -3], + ), + ( + ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=10.0, damping=2.0), + JointTargetMode.POSITION_VELOCITY, + [0, -2, 1, -3], + ), + ( + ImplicitActuatorCfg( + class_type=f"{__name__}:CustomDrive", joint_names_expr=[".*"], stiffness=10.0, damping=2.0 + ), + JointTargetMode.POSITION_VELOCITY, + [0, -2, 1, -3], + ), + ( + ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=0.0, damping=0.0), + JointTargetMode.EFFORT, + None, + ), + ( + IdealPDActuatorCfg(joint_names_expr=[".*"], stiffness=10.0, damping=2.0), + JointTargetMode.EFFORT, + None, + ), + ], +) +def test_actuator_cfg_sets_newton_target_mode_before_solver_init( + monkeypatch, actuator_cfg, expected_mode, expected_actuator_indices +): + """Resolve configured modes before finalization constructs MuJoCo actuators.""" + articulation_cfg = ArticulationCfg( + prim_path="/World/Env_[^/]*/Robot", + articulation_root_prim_path="", + actuators={"joint": actuator_cfg}, + ) + builder = _make_target_mode_builder( + monkeypatch, ["left_joint", "right_joint"], [JointTargetMode.NONE, JointTargetMode.NONE], [0.0, 0.0], [0.0, 0.0] + ) + Articulation._configure_joint_target_modes(SimpleNamespace(cfg=articulation_cfg), None) + model = builder.finalize(device="cpu") + solver = SolverMuJoCo(model, use_mujoco_cpu=True) + assert model.joint_target_mode.numpy().tolist() == [int(expected_mode), int(expected_mode)] + assert ( + solver.mjc_actuator_to_newton_idx.numpy().tolist() if solver.mjc_actuator_to_newton_idx is not None else None + ) == expected_actuator_indices + + +def test_actuator_cfg_matches_explicit_descendant_articulation_root(monkeypatch): + """Match target modes against an explicitly configured descendant articulation root.""" + articulation_cfg = ArticulationCfg( + prim_path="/World/Env_[^/]*/Robot", + articulation_root_prim_path="/base", + actuators={"joint": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=10.0, damping=0.0)}, + ) + builder = _make_target_mode_builder(monkeypatch, ["joint"], [JointTargetMode.NONE], [0.0], [0.0]) + builder.articulation_label = ["/World/Env_0/Robot/base"] + Articulation._configure_joint_target_modes(SimpleNamespace(cfg=articulation_cfg), None) + assert builder.joint_target_mode == [int(JointTargetMode.POSITION)] - sim_utils.create_prim("/World/Env_0", "Xform") - anymal_cfg = replace(ANYMAL_C_CFG, actuators=IDEAL_PD_ACTUATORS, prim_path="/World/Env_[^/]*/Anymal") - cartpole_cfg = replace( - CARTPOLE_CFG, actuators=CARTPOLE_EXPLICIT_ACTUATORS, prim_path="/World/Env_[^/]*/Cartpole" - ) - # Stand the cartpole well clear of the anymal. - cartpole_cfg.init_state = replace(cartpole_cfg.init_state, pos=(0.0, 3.0, 2.0)) - - clone_plan_from_env_0( - CloneCfg(clone_template="/World/Env_{}"), - (anymal_cfg, cartpole_cfg, AssetBaseCfg(prim_path="/World/defaultGroundPlane")), - NUM_ENVS, - 6.0, - positions=np.asarray([(i * 6.0, 0.0, 0.0) for i in range(NUM_ENVS)]), - ) - anymal = Articulation(anymal_cfg) - cartpole = Articulation(cartpole_cfg) - replicate(sim.get_clone_plan()) - sim.reset() - assert anymal.is_initialized and cartpole.is_initialized - - init_anymal = wp.to_torch(anymal.data.joint_pos).clone() - init_cartpole = wp.to_torch(cartpole.data.joint_pos).clone() - anymal.set_joint_position_target_index(target=init_anymal + TARGET_OFFSET) - anymal.set_joint_velocity_target_index(target=torch.zeros_like(init_anymal)) - cartpole.set_joint_position_target_index(target=init_cartpole + TARGET_OFFSET) - cartpole.set_joint_velocity_target_index(target=torch.zeros_like(init_cartpole)) - - pos_anymal, pos_cartpole = [], [] - for _ in range(num_steps): - anymal.write_data_to_sim() - cartpole.write_data_to_sim() - sim.step() - anymal.update(DT) - cartpole.update(DT) - pos_anymal.append(wp.to_torch(anymal.data.joint_pos).clone()) - pos_cartpole.append(wp.to_torch(cartpole.data.joint_pos).clone()) +def test_actuator_cfg_matches_clone_plan_root_expr(monkeypatch): + """Match builder labels against the clone slot spelling clone-plan root resolution returns.""" + articulation_cfg = ArticulationCfg( + prim_path="{ENV_REGEX_NS}/Robot", + actuators={"joint": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=10.0, damping=0.0)}, + ) + monkeypatch.setattr( + "isaaclab_newton.assets.articulation.articulation.resolve_matching_prims_from_source", + lambda *_args, **_kwargs: [(None, "/World/envs/env_[^/]+/Robot/base")], + ) + builder = _make_target_mode_builder(monkeypatch, ["joint"], [JointTargetMode.NONE], [0.0], [0.0]) + builder.articulation_label = ["/World/envs/env_0/Robot/base"] + Articulation._configure_joint_target_modes(SimpleNamespace(cfg=articulation_cfg), None) + assert builder.joint_target_mode == [int(JointTargetMode.POSITION)] + + +@pytest.mark.parametrize("joint_type", [JointType.FREE, JointType.FIXED]) +def test_actuator_cfg_leaves_excluded_joint_types_imported(monkeypatch, joint_type): + """Leave target modes for free and fixed joints unchanged.""" + articulation_cfg = ArticulationCfg( + prim_path="/World/Env_[^/]*/Robot", + articulation_root_prim_path="", + actuators={"joint": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=10.0, damping=0.0)}, + ) + builder = _make_target_mode_builder(monkeypatch, ["joint"], [JointTargetMode.NONE], [0.0], [0.0]) + builder.joint_type[0] = joint_type + Articulation._configure_joint_target_modes(SimpleNamespace(cfg=articulation_cfg), None) + assert builder.joint_target_mode == [int(JointTargetMode.NONE)] + + +def test_actuator_cfg_uses_imported_gain_for_none_stiffness(monkeypatch): + """Retain the imported stiffness when an implicit actuator config leaves it unset.""" + articulation_cfg = ArticulationCfg( + prim_path="/World/Env_[^/]*/Robot", + articulation_root_prim_path="", + actuators={"joint": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=None, damping=0.0)}, + ) + builder = _make_target_mode_builder(monkeypatch, ["joint"], [JointTargetMode.EFFORT], [10.0], [0.0]) + Articulation._configure_joint_target_modes(SimpleNamespace(cfg=articulation_cfg), None) + assert builder.joint_target_mode == [int(JointTargetMode.POSITION)] + + +def test_actuator_cfg_leaves_unconfigured_newton_target_modes_imported(monkeypatch): + """Leave target modes for DOFs outside an actuator group unchanged.""" + subset_cfg = ArticulationCfg( + prim_path="/World/Env_[^/]*/Robot", + articulation_root_prim_path="", + actuators={ + "shoulder": ImplicitActuatorCfg(joint_names_expr=["left_shoulder"], stiffness=10.0, damping=0.0), + }, + ) + builder = _make_target_mode_builder( + monkeypatch, + ["left_shoulder", "right_shoulder"], + [JointTargetMode.NONE, JointTargetMode.VELOCITY], + [0.0, 0.0], + [0.0, 2.0], + ) + Articulation._configure_joint_target_modes(SimpleNamespace(cfg=subset_cfg), None) + assert builder.joint_target_mode == [int(JointTargetMode.POSITION), int(JointTargetMode.VELOCITY)] - return {"joint_pos_anymal": pos_anymal, "joint_pos_cartpole": pos_cartpole} +@pytest.mark.parametrize( + ("stiffness", "damping", "expected_modes"), + [ + ({"left_joint": 10.0}, 0.0, [JointTargetMode.POSITION, JointTargetMode.EFFORT]), + ({"left_joint": 10.0}, {"right_joint": 2.0}, [JointTargetMode.POSITION, JointTargetMode.VELOCITY]), + ], +) +def test_actuator_cfg_aligns_partial_dictionary_gains_by_joint_name(monkeypatch, stiffness, damping, expected_modes): + """Resolve sparse stiffness and damping dictionaries independently by joint name.""" + articulation_cfg = ArticulationCfg( + prim_path="/World/Env_[^/]*/Robot", + articulation_root_prim_path="", + actuators={"joint": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=stiffness, damping=damping)}, + ) + builder = _make_target_mode_builder( + monkeypatch, ["left_joint", "right_joint"], [JointTargetMode.NONE, JointTargetMode.NONE], [0.0, 0.0], [0.0, 0.0] + ) + Articulation._configure_joint_target_modes(SimpleNamespace(cfg=articulation_cfg), None) + assert builder.joint_target_mode == [int(mode) for mode in expected_modes] -class TestHeterogeneousMultiArticulationNewton(unittest.TestCase): - """Two structurally-different articulations (ANYmal floating + Cartpole fixed) on Newton. - Regression for the singleton-clobber bug in ``NewtonManager._adapter`` - / ``_post_actuator_callback`` — fixed by the global-adapter refactor - + callback-list multiplexing. Heterogeneous DOF counts (12 vs 2) and - base types (floating vs fixed) stress the global adapter's handling - of varied actuator index patterns. Equivalence against the Lab - actuator path is the meaningful end-to-end check: divergence on - either robot would indicate broken stepping. - """ +def test_sync_torque_telemetry_keeps_user_order_effort_buffers_unmapped() -> None: + """Report torque telemetry directly from user-order actuator buffers.""" + joint_pos = wp.zeros((1, 3), dtype=wp.float32, device="cpu") + joint_modes = wp.array(np.asarray([0, 1, 0], dtype=np.int32), dtype=wp.int32, device="cpu") + user_to_backend = wp.array(np.asarray([2, 0, 1], dtype=np.int32), dtype=wp.int32, device="cpu") + user_effort = wp.array(np.asarray([[100.0, 200.0, 300.0]], dtype=np.float32), dtype=wp.float32, device="cpu") + user_computed_effort = wp.array(np.asarray([[10.0, 20.0, 30.0]], dtype=np.float32), dtype=wp.float32, device="cpu") + computed = wp.zeros_like(joint_pos) + applied = wp.zeros_like(joint_pos) - @classmethod - def setUpClass(cls): - cls.lab_result = _run_anymal_and_cartpole(use_newton_actuators=False) - cls.newton_result = _run_anymal_and_cartpole(use_newton_actuators=True) - - def test_anymal_matches_lab(self): - for step_i, (lab, newton) in enumerate( - zip(self.lab_result["joint_pos_anymal"], self.newton_result["joint_pos_anymal"]) - ): - torch.testing.assert_close( - newton, - lab, - atol=2e-3, - rtol=1e-3, - msg=f"ANYmal joint_pos diverged from Lab path at step {step_i}", - ) + wp.launch( + sync_torque_telemetry, + dim=joint_pos.shape, + inputs=[ + joint_pos, + wp.zeros_like(joint_pos), + wp.zeros_like(joint_pos), + wp.zeros_like(joint_pos), + wp.zeros_like(joint_pos), + wp.zeros_like(joint_pos), + wp.full((1, 3), 1000.0, dtype=wp.float32, device="cpu"), + joint_modes, + user_effort, + user_computed_effort, + user_to_backend, + False, + ], + outputs=[computed, applied], + device="cpu", + ) - def test_cartpole_matches_lab(self): - for step_i, (lab, newton) in enumerate( - zip(self.lab_result["joint_pos_cartpole"], self.newton_result["joint_pos_cartpole"]) - ): - torch.testing.assert_close( - newton, - lab, - atol=2e-3, - rtol=1e-3, - msg=f"Cartpole joint_pos diverged from Lab path at step {step_i}", - ) + np.testing.assert_allclose(computed.numpy(), np.asarray([[10.0, 200.0, 30.0]], dtype=np.float32)) + np.testing.assert_allclose(applied.numpy(), np.asarray([[100.0, 200.0, 300.0]], dtype=np.float32)) # --------------------------------------------------------------------------- -# Domain randomization via events.py — Newton backend +# Equivalence tests with different actuator types # --------------------------------------------------------------------------- -class TestRandomizeActuatorGainsViaEventsNewton(unittest.TestCase): - """End-to-end DR test for the Newton backend. - - Drives ``randomize_actuator_gains`` and verifies that kp/kd values reach - the controllers of the articulation's Newton actuators through - ``write_group_parameter``; the assertions read the - controllers back via the public ``read_group_parameter``. +@pytest.mark.parametrize( + "island", + ["ideal", "cartpole", "dc_motor", "mixed", "implicit_feedforward", "delayed", "remotized"], +) +def test_newton_actuators_match_lab_actuators(lab_run: dict, newton_run: _Run, island: str) -> None: + """Newton-native actuators reproduce the Isaac Lab actuator trajectories and torque telemetry. - With ``operation="abs"`` and ``distribution="uniform"`` over a - degenerate range ``(K, K)``, every randomized cell is set to exactly - ``K`` — so the assertions are deterministic. + The islands cover IdealPD on the floating leg and on the fixed cartpole, which share the model-wide adapter, + a saturating DC motor, implicit hips beside explicit Newton actuators, an implicit feedforward effort added on + top of the solver's joint drive, and command delays with and without position-based clamping. """ + _assert_equivalent( + lab_run["results"][island], newton_run.results[island], torque_atol=_islands()[island].torque_atol + ) - def test_two_articulations(self): - from isaaclab_assets import CARTPOLE_CFG # noqa: PLC0415 - - sim_cfg = SimulationCfg(dt=DT, physics=NEWTON_CFG, use_newton_actuators=True) - with build_simulation_context( - device="cuda:0", - gravity_enabled=True, - add_ground_plane=True, - sim_cfg=sim_cfg, - ) as sim: - sim._app_control_on_stop_handle = None - sim_utils.create_prim("/World/Env_0", "Xform") - - anymal_cfg = replace(ANYMAL_C_CFG, actuators=IDEAL_PD_ACTUATORS, prim_path="/World/Env_[^/]*/Anymal") - cartpole_cfg = replace( - CARTPOLE_CFG, actuators=CARTPOLE_EXPLICIT_ACTUATORS, prim_path="/World/Env_[^/]*/Cartpole" - ) - cartpole_cfg.init_state = replace(cartpole_cfg.init_state, pos=(0.0, 3.0, 2.0)) - clone_plan_from_env_0( - CloneCfg(clone_template="/World/Env_{}"), - (anymal_cfg, cartpole_cfg, AssetBaseCfg(prim_path="/World/defaultGroundPlane")), - NUM_ENVS, - 6.0, - positions=np.asarray([(i * 6.0, 0.0, 0.0) for i in range(NUM_ENVS)]), - ) - anymal = Articulation(anymal_cfg) - cartpole = Articulation(cartpole_cfg) - replicate(sim.get_clone_plan()) - sim.reset() - - self.assertIsNotNone(SimulationManager._adapter) - - anymal_read = functools.partial(read_group_parameter, anymal.actuators) - cartpole_read = functools.partial(read_group_parameter, cartpole.actuators) - anymal_stiffness_before = anymal_read("legs", "controller", "kp").clone() - anymal_damping_before = anymal_read("legs", "controller", "kd").clone() - # Before DR, native gain reads must return the configured values for *every* env. This is - # also the regression check for the env-major DOF stride decoding on floating-base - # articulations (ANYmal-C has 6 free-root DOFs + 12 joints): a wrong stride corrupts every - # env past the first. - n = anymal.num_joints - torch.testing.assert_close(anymal_stiffness_before, torch.full((NUM_ENVS, n), 40.0, device=anymal.device)) - torch.testing.assert_close(anymal_damping_before, torch.full((NUM_ENVS, n), 5.0, device=anymal.device)) - cartpole_stiffness_before = cartpole_read("all_joints", "controller", "kp").clone() - cartpole_damping_before = cartpole_read("all_joints", "controller", "kd").clone() - - env = MockEnv({"anymal": anymal, "cartpole": cartpole}, NUM_ENVS, anymal.device) - term, asset_cfg = build_dr_term(env, "cartpole") - env_ids = torch.tensor([0], device=anymal.device, dtype=torch.long) - - term( - env, - env_ids=env_ids, - asset_cfg=asset_cfg, - stiffness_distribution_params=(100.0, 100.0), - damping_distribution_params=(5.0, 5.0), - operation="abs", - distribution="uniform", - ) - n_cp = cartpole.num_joints - torch.testing.assert_close( - cartpole_read("all_joints", "controller", "kp")[0], torch.full((n_cp,), 100.0, device=anymal.device) - ) - torch.testing.assert_close( - cartpole_read("all_joints", "controller", "kd")[0], torch.full((n_cp,), 5.0, device=anymal.device) - ) +def test_decimated_remotized_pd_matches_lab_actuators(decimated_runs: dict) -> None: + """RemotizedPD with decimation 2 and CUDA-graph capture: Lab vs Newton.""" + _assert_equivalent(decimated_runs["lab"], decimated_runs["newton"]) - # ANYmal is untouched (DR was scoped to cartpole). - torch.testing.assert_close(anymal_read("legs", "controller", "kp"), anymal_stiffness_before) - torch.testing.assert_close(anymal_read("legs", "controller", "kd"), anymal_damping_before) - # Cartpole's other envs are also untouched (env_ids=[0] only). - for env_idx in range(1, NUM_ENVS): - torch.testing.assert_close( - cartpole_read("all_joints", "controller", "kp")[env_idx], cartpole_stiffness_before[env_idx] - ) - torch.testing.assert_close( - cartpole_read("all_joints", "controller", "kd")[env_idx], cartpole_damping_before[env_idx] - ) +def test_dc_motor_clamp_binds(lab_run: dict, newton_run: _Run) -> None: + """The saturating DC motor clamps on both paths, so the equivalence can detect missing clamping.""" + for result in (lab_run["results"]["dc_motor"], newton_run.results["dc_motor"]): + assert any( + not torch.allclose(applied, computed) + for applied, computed in zip(result["applied_effort"], result["computed_effort"]) + ), "the DC-motor clamp never bound, so the equivalence cannot detect missing clamping" - # DR scoped to the floating-base ANYmal updates only its selected env. - anymal_term, anymal_asset_cfg = build_dr_term(env, "anymal") - anymal_term( - env, - env_ids=env_ids, - asset_cfg=anymal_asset_cfg, - stiffness_distribution_params=(100.0, 100.0), - damping_distribution_params=(5.0, 5.0), - operation="abs", - distribution="uniform", - ) - torch.testing.assert_close( - anymal_read("legs", "controller", "kp")[0], torch.full((n,), 100.0, device=anymal.device) - ) - torch.testing.assert_close( - anymal_read("legs", "controller", "kd")[0], torch.full((n,), 5.0, device=anymal.device) - ) - for env_idx in range(1, NUM_ENVS): - torch.testing.assert_close( - anymal_read("legs", "controller", "kp")[env_idx], anymal_stiffness_before[env_idx] - ) - torch.testing.assert_close( - anymal_read("legs", "controller", "kd")[env_idx], anymal_damping_before[env_idx] - ) + +@pytest.mark.parametrize( + "island, group_names, controller_type, clamping, has_delay", + [ + ("delayed", ("legs",), "DrivePD", None, True), + ("remotized", ("knees",), "DrivePD", "ClampingPositionBased", True), + ("neural", ("mlp_legs",), "DriveNeuralMLP", "ClampingDCMotor", False), + ("neural", ("lstm_legs",), "DriveNeuralLSTM", "ClampingDCMotor", False), + ], + ids=["delayed_pd", "remotized_pd", "mlp", "lstm"], +) +def test_newton_actuator_authoring( + newton_run: _Run, island: str, group_names: tuple[str, ...], controller_type: str, clamping: str | None, has_delay +) -> None: + """Lab actuator configurations are authored as the matching Newton controller, clamping, and delay.""" + info = [entry for entry in newton_run.results[island]["actuator_info"] if entry["group"] in group_names] + assert info, f"no Newton actuators were created for {group_names}" + for entry in info: + assert entry["controller_type"] == controller_type + assert clamping is None or clamping in entry["clamping_types"] + assert entry["has_delay"] is has_delay + if island == "neural": + # the neural controllers run without producing non-finite joint positions + assert all(torch.isfinite(pos).all() for pos in newton_run.results[island]["joint_pos"]) # --------------------------------------------------------------------------- -# DelayedPD equivalence: PD with actuator command delay +# Joint ordering # --------------------------------------------------------------------------- -class TestDelayedPDEquivalence(_EquivalenceTestBase): - """DelayedPDActuator on all 12 joints: Lab vs Newton. +def test_newton_actuator_rollout_matches_reordered_joints(newton_run: _Run) -> None: + """Match Newton-backend actuator traces under a permuted public joint ordering. - Verifies that actuator command delays are correctly authored as - ``NewtonActuatorDelayAPI`` and produce matching trajectories. The ramped - position target makes a missing delay change the trajectory. + The implicit hip group reads its torque telemetry from backend-order effort buffers, so the + permutation-sensitive feedforward effort also checks that gather. """ + identity_result = newton_run.results["mixed_permuted"] + reordered_result = newton_run.results["mixed_permuted_reordered"] - __test__ = True - actuators = FIXED_DELAYED_PD_ACTUATORS - ramp_targets = True - - def test_newton_actuators_are_delayed_pd(self): - self.assertTrue(self.newton_result["actuator_info"], "No Newton actuators were created") - for a in self.newton_result["actuator_info"]: - self.assertTrue(a["has_delay"], "Delay not found on delayed PD actuator") - self.assertEqual(a["controller_type"], "DrivePD") - - -# --------------------------------------------------------------------------- -# Decimation tests: re-run equivalence with decimation > 1 + CUDA graph capture -# --------------------------------------------------------------------------- - -NEWTON_CFG_DEC = NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=500, - nconmax=500, - ls_iterations=20, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=2, - debug_mode=False, - use_cuda_graph=True, -) + installed_ordering = reordered_result["joint_ordering"] + assert installed_ordering is not None + assert not installed_ordering["is_identity"] + assert_articulation_ordering_trace_matches(identity_result, reordered_result, _LEG_REORDERED_JOINT_NAMES) + + +@pytest.mark.parametrize("island", ["ideal", "ideal_reordered"], ids=["identity", "reordered"]) +def test_write_data_to_sim_writes_joint_targets_in_backend_order(newton_run: _Run, island: str) -> None: + """Explicit actuators on the Newton-actuator path publish their raw targets to the backend in backend order.""" + articulation = newton_run.articulations[island] + assert (articulation.data.joint_ordering is not None) is (island == "ideal_reordered") + assert articulation._has_newton_actuators is True + + # Distinct per-joint targets away from the defaults, so a skipped or unpermuted write is visible. + target = articulation.data.default_joint_pos.torch.clone() + target += 0.01 * torch.arange(1, articulation.num_joints + 1, device=target.device) + articulation.set_joint_position_target_index(target=target) + articulation.write_data_to_sim() + + source = articulation.actuators.target_command.position.torch + torch.testing.assert_close(source, target) + user_to_backend = ( + list(articulation.joint_ordering.user_to_backend_indices) + if articulation.joint_ordering is not None + else list(range(articulation.num_joints)) + ) + expected_backend_target = torch.empty_like(source) + expected_backend_target[:, user_to_backend] = source + torch.testing.assert_close(wp.to_torch(articulation.data._sim_bind_joint_position_target), expected_backend_target) + + +def test_newton_native_actuator_gain_write_maps_public_joint_subset_to_backend(newton_run: _Run) -> None: + """Map selected public joint IDs to Newton-controller columns.""" + articulation = newton_run.articulations["ideal_reordered"] + assert articulation.joint_ordering is not None + assert articulation.newton_actuator_adapter is not None + + def gather_stiffness() -> torch.Tensor: + stiffness = torch.zeros((articulation.num_instances, articulation.num_joints), device=articulation.device) + for actuator in articulation.newton_actuator_adapter.actuators: + if hasattr(actuator.controller, "kp"): + stiffness += wp.to_torch( + articulation.root_view.get_actuator_parameter(actuator, actuator.controller, "kp") + ) + return stiffness + stiffness_before = gather_stiffness() + env_ids = torch.tensor([1], device=articulation.device, dtype=torch.long) + joint_ids = torch.tensor([1, 3, 4], device=articulation.device, dtype=torch.long) + stiffness = torch.tensor([[101.0, 103.0, 104.0]], device=articulation.device) -class _DecimationMixin: - """Common knobs for decimation/CUDA-graph variants of equivalence classes.""" + with pytest.warns(DeprecationWarning, match="write_actuator_stiffness_to_sim"): + articulation.write_actuator_stiffness_to_sim(stiffness=stiffness, env_ids=env_ids, joint_ids=joint_ids) - __test__ = True - dt = 1.0 / 100.0 - newton_cfg = NEWTON_CFG_DEC - num_steps = 5 - decimation = 2 + backend_joint_ids = torch.tensor( + articulation.joint_ordering.user_to_backend_indices, device=articulation.device, dtype=torch.long + )[joint_ids] + expected_stiffness = stiffness_before.clone() + expected_stiffness[env_ids.unsqueeze(1), backend_joint_ids.unsqueeze(0)] = stiffness + torch.testing.assert_close(gather_stiffness(), expected_stiffness) # --------------------------------------------------------------------------- -# Per-env reset: actuator state isolation +# Domain randomization via events.py — Newton backend # --------------------------------------------------------------------------- -class TestActuatorStateReset(ActuatorStateResetBase, unittest.TestCase): - """Per-env actuator state reset isolation on the Newton backend. +def test_randomize_actuator_gains_reaches_newton_controllers(newton_run: _Run) -> None: + """``randomize_actuator_gains`` writes kp/kd into the selected environment of one articulation's controllers. - The scenario and assertions live in :class:`ActuatorStateResetBase`; - this subclass provides the Newton sim config and the model-wide adapter. + The event writes through ``write_group_parameter`` and the assertions read back through the public + ``read_group_parameter``. A degenerate ``(K, K)`` range with ``operation="abs"`` sets each randomized cell to + exactly ``K``. """ - - def _make_sim_cfg(self, use_newton_actuators: bool) -> SimulationCfg: - return SimulationCfg(dt=DT, physics=NEWTON_CFG, use_newton_actuators=use_newton_actuators) - - def _make_articulation(self) -> Articulation: - cfg = replace(ANYMAL_C_CFG, actuators=DELAYED_PD_ACTUATORS, prim_path="/World/Env_.*/Robot") - plan = clone_plan_from_env_0( - CloneCfg(clone_template="/World/Env_{}"), - (cfg, AssetBaseCfg(prim_path="/World/defaultGroundPlane")), - self.NUM_ENVS, - 3.0, - positions=np.asarray([(i * 3.0, 0.0, 0.0) for i in range(self.NUM_ENVS)]), + groups = { + "legs": (newton_run.articulations["ideal"], "legs"), + "cartpole": (newton_run.articulations["cartpole"], "all_joints"), + } + legs = groups["legs"][0] + assert SimulationManager._adapter is not None + + def gains(name: str) -> torch.Tensor: + """Return the ``(kp, kd)`` gains of one articulation's actuator group, shape ``(2, num_envs, num_joints)``.""" + articulation, group = groups[name] + return torch.stack([read_group_parameter(articulation.actuators, group, "controller", p) for p in ("kp", "kd")]) + + # Every environment reads the configured gains. On the floating leg this pins the env-major DOF stride + # decoding (6 free-root DOFs + leg joints): a wrong stride corrupts every environment past the first. + configured = torch.tensor([40.0, 5.0], device=legs.device).view(2, 1, 1).expand(2, NUM_ENVS, legs.num_joints) + torch.testing.assert_close(gains("legs"), configured) + + env = MockEnv({name: articulation for name, (articulation, _) in groups.items()}, NUM_ENVS, legs.device) + for name in ("cartpole", "legs"): + expected = {other: gains(other).clone() for other in groups} + expected[name][:, 0] = torch.tensor([100.0, 7.0], device=legs.device).view(2, 1) + term, asset_cfg = build_dr_term(env, name) + term( + env, + env_ids=torch.tensor([0], device=legs.device, dtype=torch.long), + asset_cfg=asset_cfg, + stiffness_distribution_params=(100.0, 100.0), + damping_distribution_params=(7.0, 7.0), + operation="abs", + distribution="uniform", ) - articulation = Articulation(cfg) - replicate(plan) - return articulation - - def _get_adapter(self, articulation): - return SimulationManager._adapter + # only environment 0 of the selected articulation changes + for other in groups: + torch.testing.assert_close(gains(other), expected[other]) # --------------------------------------------------------------------------- -# RemotizedPD actuator: PD + delay + position-based clamping lookup table +# Per-env reset: actuator state isolation # --------------------------------------------------------------------------- -def _remotized_pd_actuators() -> dict: - """RemotizedPD (Spot knee lookup) on KFE with IdealPD on HAA/HFE.""" - from isaaclab.actuators.actuator_pd_cfg import RemotizedPDActuatorCfg # noqa: PLC0415 - - return { - "hips": IdealPDActuatorCfg( - joint_names_expr=[".*HAA", ".*HFE"], - stiffness=40.0, - damping=5.0, - actuator_effort_limit=80.0, - ), - "knees": RemotizedPDActuatorCfg( - joint_names_expr=[".*KFE"], - stiffness=60.0, - damping=1.5, - actuator_effort_limit=80.0, - min_delay=3, - max_delay=3, - joint_parameter_lookup=SPOT_KNEE_LOOKUP, - ), - } - - -class TestRemotizedPDEquivalence(_EquivalenceTestBase): - """RemotizedPD (PD + delay + position-based clamping): Lab vs Newton.""" - - __test__ = True - ramp_targets = True +def test_newton_state_reset_isolated_to_reset_env(newton_run: _Run) -> None: + """Newton: ``num_pushes`` zeroes for env 0's DOFs only after reset of [0]. - @classmethod - def setUpClass(cls): - cls.actuators = _remotized_pd_actuators() - super().setUpClass() - - def test_newton_knee_actuators_are_delayed_position_clamped_pd(self): - kfe_acts = [a for a in self.newton_result["actuator_info"] if "ClampingPositionBased" in a["clamping_types"]] - self.assertTrue(len(kfe_acts) > 0, "No actuator with position-based clamping found") - for a in kfe_acts: - self.assertEqual(a["controller_type"], "DrivePD") - self.assertTrue(a["has_delay"], "Delay not found on remotized KFE actuator") - - -class TestDecimationRemotizedPD(_DecimationMixin, TestRemotizedPDEquivalence): - """RemotizedPD — decimation=2 + CUDA graph.""" - - -# --------------------------------------------------------------------------- -# Neural network actuator authoring: MLP and LSTM -# --------------------------------------------------------------------------- - - -class TestNeuralActuatorAuthoring(unittest.TestCase): - """Verify ActuatorNetMLPCfg and ActuatorNetLSTMCfg are authored as Newton neural - controllers with DC motor clamping and run on the Newton backend. + Only the reset articulation's own delayed actuators are checked. The adapter is model-wide, so the other + islands' actuators share its state buffers; their state is not part of this articulation's contract. """ - - @classmethod - def setUpClass(cls): - from isaaclab.actuators.actuator_net_cfg import ActuatorNetLSTMCfg, ActuatorNetMLPCfg # noqa: PLC0415 - - cls.mlp_path = make_dummy_mlp_checkpoint() - cls.lstm_path = make_dummy_lstm_checkpoint() - cls.result = _run_simulation( - { - "mlp_legs": ActuatorNetMLPCfg( - joint_names_expr=[".*HAA"], - network_file=cls.mlp_path, - saturation_effort=120.0, - actuator_effort_limit=80.0, - actuator_velocity_limit=7.5, - pos_scale=-1.0, - vel_scale=1.0, - torque_scale=1.0, - input_order="pos_vel", - input_idx=[0, 1, 2], - ), - "lstm_legs": ActuatorNetLSTMCfg( - joint_names_expr=[".*HFE"], - network_file=cls.lstm_path, - saturation_effort=120.0, - actuator_effort_limit=80.0, - actuator_velocity_limit=7.5, - ), - "pd_legs": IdealPDActuatorCfg( - joint_names_expr=[".*KFE"], - stiffness=40.0, - damping=5.0, - actuator_effort_limit=80.0, - ), - }, - use_newton_actuators=True, - ) - - @classmethod - def tearDownClass(cls): - os.unlink(cls.mlp_path) - os.unlink(cls.lstm_path) - - def test_mlp_has_dc_motor_clamping(self): - mlp_acts = [a for a in self.result["actuator_info"] if a["controller_type"] == "DriveNeuralMLP"] - self.assertTrue(len(mlp_acts) > 0, "No NeuralMLP controller found") - for a in mlp_acts: - self.assertIn("ClampingDCMotor", a["clamping_types"]) - - def test_lstm_has_dc_motor_clamping(self): - lstm_acts = [a for a in self.result["actuator_info"] if a["controller_type"] == "DriveNeuralLSTM"] - self.assertTrue(len(lstm_acts) > 0, "No NeuralLSTM controller found") - for a in lstm_acts: - self.assertIn("ClampingDCMotor", a["clamping_types"]) - - def test_positions_finite(self): - for step_i, pos in enumerate(self.result["joint_pos"]): - self.assertTrue(torch.isfinite(pos).all(), f"Non-finite positions at step {step_i}") - - -def test_sync_torque_telemetry_keeps_user_order_effort_buffers_unmapped() -> None: - """Report torque telemetry directly from user-order actuator buffers.""" - joint_pos = wp.zeros((1, 3), dtype=wp.float32, device="cpu") - joint_modes = wp.array(np.asarray([0, 1, 0], dtype=np.int32), dtype=wp.int32, device="cpu") - user_to_backend = wp.array(np.asarray([2, 0, 1], dtype=np.int32), dtype=wp.int32, device="cpu") - user_effort = wp.array(np.asarray([[100.0, 200.0, 300.0]], dtype=np.float32), dtype=wp.float32, device="cpu") - user_computed_effort = wp.array(np.asarray([[10.0, 20.0, 30.0]], dtype=np.float32), dtype=wp.float32, device="cpu") - computed = wp.zeros_like(joint_pos) - applied = wp.zeros_like(joint_pos) - - wp.launch( - sync_torque_telemetry, - dim=joint_pos.shape, - inputs=[ - joint_pos, - wp.zeros_like(joint_pos), - wp.zeros_like(joint_pos), - wp.zeros_like(joint_pos), - wp.zeros_like(joint_pos), - wp.zeros_like(joint_pos), - wp.full((1, 3), 1000.0, dtype=wp.float32, device="cpu"), - joint_modes, - user_effort, - user_computed_effort, - user_to_backend, - False, - ], - outputs=[computed, applied], - device="cpu", - ) - - np.testing.assert_allclose(computed.numpy(), np.asarray([[10.0, 200.0, 30.0]], dtype=np.float32)) - np.testing.assert_allclose(applied.numpy(), np.asarray([[100.0, 200.0, 300.0]], dtype=np.float32)) + articulation = newton_run.articulations["delayed"] + adapter = SimulationManager._adapter + assert adapter is not None + own_actuators = [] + for group_name in articulation.actuators._native_group_names: + group_actuators = articulation.actuators[group_name] + own_actuators.extend(group_actuators if isinstance(group_actuators, tuple) else (group_actuators,)) + stateful_pairs = [ + (act, st) + for act, st in zip(adapter.actuators, adapter._states_a) + if any(act is own for own in own_actuators) and st is not None and st.delay_state is not None + ] + assert len(stateful_pairs) > 0, "expected at least one DelayedPD actuator with delay_state" + + for act, state in stateful_pairs: + pushes_before = state.delay_state.num_pushes.numpy() + assert (pushes_before > 0).all(), "expected non-zero num_pushes for all DOFs after warmup" + + articulation.reset(env_ids=torch.tensor([0], device=articulation.device, dtype=torch.long)) + + # Map each entry of ``act.indices`` to its env via the adapter's per-env DOF count. The adapter is + # model-wide (includes free-joint DOFs on floating-base articulations), so ``adapter.num_joints`` is the + # stride. + for act, state in stateful_pairs: + pushes_after = state.delay_state.num_pushes.numpy() + indices_np = act.indices.numpy() + for i, global_dof in enumerate(indices_np): + env = int(global_dof) // adapter.num_joints + if env == 0: + assert int(pushes_after[i]) == 0, f"DOF {i} (env {env}) should be reset to 0, got {pushes_after[i]}" + else: + assert int(pushes_after[i]) > 0, f"DOF {i} (env {env}) was NOT in reset env_ids but num_pushes is 0" -if __name__ == "__main__": - unittest.main() +def test_lab_state_reset_isolated_to_reset_env(lab_run: dict) -> None: + """Reset environments accept fresh commands while the remaining environments retain their history.""" + for computed_effort, demand in lab_run["state_reset"].values(): + torch.testing.assert_close(computed_effort, demand) diff --git a/source/isaaclab_newton/test/assets/test_rigid_object.py b/source/isaaclab_newton/test/assets/test_rigid_object.py index a49149b30455..fac10ee7bdfa 100644 --- a/source/isaaclab_newton/test/assets/test_rigid_object.py +++ b/source/isaaclab_newton/test/assets/test_rigid_object.py @@ -6,33 +6,43 @@ # ignore private usage of variables warning # pyright: reportPrivateUsage=none -from isaaclab.test.utils import DeviceScope, launch_test_simulation, test_devices +"""Kitless real-Newton rigid-object and rigid-object-collection coverage on one persistent two-environment CPU scene. -launch_test_simulation() +Every environment holds a rigid-object cube and a three-cube collection, all locally spawned shapes. The cube +also stands beside the collection as an unrelated rigid body, so the collection view must select exactly its +configured bodies. Scene gravity is off; a test that needs gravity applies it to every world for its own duration. +Configurations that fail initialization and the single-instance cases build their own scenes. Only one simulation +context can be alive, so those tests come first and fail if selected after a shared-scene test. +""" + +from isaaclab_newton.physics import NewtonCfg + +from isaaclab.sim import SimulationCfg +from isaaclab.test.utils import launch_test_simulation + +launch_test_simulation(SimulationCfg(physics=NewtonCfg())) import sys +from collections.abc import Iterator +from dataclasses import dataclass from types import SimpleNamespace -from typing import Literal import numpy as np import pytest import torch import warp as wp -from flaky import flaky -from isaaclab_newton.assets import RigidObject -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +from isaaclab_newton.assets import RigidObject, RigidObjectCollection from isaaclab_newton.physics import NewtonManager as SimulationManager -from isaaclab_physx.sim.schemas import PhysxRigidBodyCfg from newton import ModelFlags +from newton_test_utils import env_origins, local_usd, newton_sim_cfg, spawn_assets, world_gravity import isaaclab.sim as sim_utils -from isaaclab.assets import AssetBaseCfg, RigidObjectCfg -from isaaclab.cloner import CloneCfg, clone_plan_from_env_0, replicate +from isaaclab.assets import RigidObjectCfg, RigidObjectCollectionCfg from isaaclab.envs.mdp import randomize_physics_scene_gravity from isaaclab.envs.mdp.events import randomize_rigid_body_material from isaaclab.managers import EventTermCfg, SceneEntityCfg -from isaaclab.sim import SimulationCfg, build_simulation_context -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR +from isaaclab.sim import SimulationContext, build_simulation_context +from isaaclab.test.utils import DeviceScope, test_devices from isaaclab.utils.math import ( combine_frame_transforms, quat_apply_inverse, @@ -40,410 +50,396 @@ quat_mul, quat_rotate, random_orientation, + subtract_frame_transforms, ) -NEWTON_SIM_CFG = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg(), - ), -) +pytestmark = [pytest.mark.integration, pytest.mark.kitless] +_GRAVITY = (0.0, 0.0, -9.81) -def _newton_sim_context(device, gravity_enabled=True, dt=None, **kwargs): - """Helper to create a Newton simulation context with the correct device. - When sim_cfg is provided to build_simulation_context, the device, gravity_enabled, and dt - kwargs are ignored. This helper applies them to the shared NEWTON_SIM_CFG before calling. - """ - NEWTON_SIM_CFG.device = device - NEWTON_SIM_CFG.gravity = (0.0, 0.0, -9.81) if gravity_enabled else (0.0, 0.0, 0.0) - if dt is not None: - NEWTON_SIM_CFG.dt = dt - return build_simulation_context(device=device, sim_cfg=NEWTON_SIM_CFG, **kwargs) - - -def generate_cubes_scene( - num_cubes: int = 1, - height=1.0, - api: Literal["none", "rigid_body", "articulation_root"] = "rigid_body", - kinematic_enabled: bool = False, - device: str = "cuda:0", - add_ground_plane: bool = False, -) -> tuple[RigidObject, torch.Tensor]: - """Generate a scene with the provided number of cubes. - - Args: - num_cubes: Number of cubes to generate. - height: Height of the cubes. - api: The type of API that the cubes should have. - kinematic_enabled: Whether the cubes are kinematic. - device: Device to use for the simulation. - add_ground_plane: Whether the simulation context authored a shared ground plane. - - Returns: - A tuple containing the rigid object representing the cubes and the origins of the cubes. +def _cube_cfg(height: float = 1.0) -> RigidObjectCfg: + """Return a dynamic unit-mass cube with a collider.""" + return RigidObjectCfg( + prim_path="/World/Env_[^/]*/Object", + spawn=sim_utils.CuboidCfg( + size=(0.2, 0.2, 0.2), + rigid_props=sim_utils.UsdPhysicsRigidBodyCfg(), + mass_props=sim_utils.MassCfg(mass=1.0), + collision_props=sim_utils.UsdPhysicsCollisionCfg(), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, -3.0, height)), + ) - """ - origins = np.asarray([(i * 1.0, 0, height) for i in range(num_cubes)], dtype=np.float32) - sim_utils.create_prim("/World/Env_0", "Xform", translation=origins[0]) - # Resolve spawn configuration - if api == "none": +_NUM_CUBES = 3 + + +def _cube_spawn_cfg(rigid: bool = True) -> sim_utils.CuboidCfg: + """Return a unit-mass cube, or a static collider without a rigid body.""" + if not rigid: # since no rigid body properties defined, this is just a static collider - spawn_cfg = sim_utils.CuboidCfg( - size=(0.1, 0.1, 0.1), - collision_props=sim_utils.UsdPhysicsCollisionCfg(), - ) - elif api == "rigid_body": - spawn_cfg = sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", - rigid_props=sim_utils.UsdPhysicsRigidBodyCfg(kinematic_enabled=kinematic_enabled), - ) - elif api == "articulation_root": - spawn_cfg = sim_utils.UsdFileCfg( - usd_path=f"{ISAACLAB_NUCLEUS_DIR}/Tests/RigidObject/Cube/dex_cube_instanceable_with_articulation_root.usd", - # Only tune existing bodies; do not create one on this invalid articulation fixture. - rigid_props={"(/.*)?": [sim_utils.UsdPhysicsRigidBodyCfg(kinematic_enabled=kinematic_enabled)]}, - ) - else: - raise ValueError(f"Unknown api: {api}") + return sim_utils.CuboidCfg(size=(0.1, 0.1, 0.1), collision_props=sim_utils.UsdPhysicsCollisionCfg()) + return sim_utils.CuboidCfg( + size=(0.2, 0.2, 0.2), + rigid_props=sim_utils.UsdPhysicsRigidBodyCfg(), + mass_props=sim_utils.MassCfg(mass=1.0), + collision_props=sim_utils.UsdPhysicsCollisionCfg(), + ) - # Create rigid object - cube_object_cfg = RigidObjectCfg( - prim_path="/World/Env_[^/]*/Object", - spawn=spawn_cfg, - init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, height)), + +def _collection_cfg( + num_cubes: int, height: float = 1.0, rigid: bool = True, *, object_paths: tuple[str, ...] | None = None +) -> RigidObjectCollectionCfg: + """Return a collection of cubes spaced 3 m apart along y.""" + if object_paths is None: + object_paths = tuple(f"Object_{i}" for i in range(num_cubes)) + return RigidObjectCollectionCfg( + rigid_objects={ + f"cube_{i}": RigidObjectCfg( + prim_path=f"/World/Env_[^/]*/{path}", + spawn=_cube_spawn_cfg(rigid), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 3.0 * i, height)), + ) + for i, path in enumerate(object_paths) + } ) - cfgs = [cube_object_cfg] - if add_ground_plane: - cfgs.append(AssetBaseCfg(prim_path="/World/defaultGroundPlane")) - clone_plan_from_env_0(CloneCfg(clone_template="/World/Env_{}"), cfgs, num_cubes, 1.0, positions=origins) - cube_object = RigidObject(cfg=cube_object_cfg) - return cube_object, torch.as_tensor(origins, device=device) + +## +# Own scenes. These tests run before the shared scene exists: only one simulation context can be alive. +## @pytest.mark.isaacsim_ci -@pytest.mark.parametrize("api", ["none", "articulation_root"]) -def test_initialization_rejects_non_rigid_body_prims(api): - """Initialization fails when the prim path has no rigid body API or carries an articulation root.""" - with _newton_sim_context("cpu", auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_object, _ = generate_cubes_scene(num_cubes=1, api=api, device="cpu") +@pytest.mark.parametrize("api", ["none", "articulation"]) +def test_initialization_rejects_non_rigid_body_prims(api: str) -> None: + """Initialization fails unless the prim path holds exactly one rigid body. + + A static collider has no rigid body, and an articulation rooted on a rigid body has several. + """ + if api == "none": + # since no rigid body properties defined, this is just a static collider + spawn_cfg = sim_utils.CuboidCfg(size=(0.1, 0.1, 0.1), collision_props=sim_utils.UsdPhysicsCollisionCfg()) + else: + spawn_cfg = local_usd("rigid_body_with_articulation_root.usda") + cube_cfg = RigidObjectCfg(prim_path="/World/Env_[^/]*/Object", spawn=spawn_cfg) + with build_simulation_context(sim_cfg=newton_sim_cfg("cpu")) as sim: + cube_object = spawn_assets({"cube": cube_cfg}, num_envs=1)["cube"] # Check that the framework doesn't hold excessive strong references. assert sys.getrefcount(cube_object) < 10 - replicate(sim.get_clone_plan()) with pytest.raises(RuntimeError, match="Expected 1 prims at"): sim.reset() -@pytest.mark.isaacsim_ci -@pytest.mark.parametrize("num_cubes", [4]) -@pytest.mark.parametrize("device", test_devices()) -def test_external_force_on_single_body(num_cubes, device): - """Test initialization and external forces on the base of the object. - - In the first phase, we apply a force equal to the weight of an object on the base of - one of the objects. We check that the object does not move. For the other object, - we do not apply any force and check that it falls down. In the second phase, the force - is applied at 1m in the Y direction and the object must rotate around its X axis. - - We validate that this works when we apply the force in the global frame and in the local frame, - and that resetting the object clears the wrench applied in the previous iteration. - """ - # Generate cubes scene - with _newton_sim_context(device, add_ground_plane=True, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_object, origins = generate_cubes_scene(num_cubes=num_cubes, device=device, add_ground_plane=True) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(cube_object) < 10 - - # Play the simulator - replicate(sim.get_clone_plan()) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) +def test_single_instance_initialization_and_root_state_write(device: str) -> None: + """A single rigid object initializes with singleton buffers and reads a root state write back without a step.""" + with build_simulation_context(sim_cfg=newton_sim_cfg(device)) as sim: + cube_object = spawn_assets({"cube": _cube_cfg()}, num_envs=1)["cube"] sim.reset() - - # Check if object is initialized assert cube_object.is_initialized - assert len(cube_object.body_names) == 1 - - # Check buffers that exists and have correct shapes - assert cube_object.data.root_pos_w.torch.shape == (num_cubes, 3) - assert cube_object.data.root_quat_w.torch.shape == (num_cubes, 4) - assert cube_object.data.body_mass.torch.shape == (num_cubes, 1) - assert cube_object.data.body_inertia.torch.shape == (num_cubes, 1, 9) - - # Find bodies to apply the force - body_ids, body_names = cube_object.find_bodies(".*") - - def reset_cubes(): - # reset root state; shift the cubes to their origins so they are not on top of each other - root_pose = cube_object.data.default_root_pose.torch.clone() - root_pose[:, :3] = origins - cube_object.write_root_pose_to_sim_index(root_pose=root_pose) - cube_object.write_root_velocity_to_sim_index(root_velocity=cube_object.data.default_root_vel.torch.clone()) - cube_object.reset() - - # Reset should zero external forces and torques - assert not cube_object.instantaneous_wrench_composer.active - assert not cube_object.permanent_wrench_composer.active - assert torch.count_nonzero(cube_object.instantaneous_wrench_composer.out_force_b.torch) == 0 - assert torch.count_nonzero(cube_object.instantaneous_wrench_composer.out_torque_b.torch) == 0 - assert torch.count_nonzero(cube_object.permanent_wrench_composer.out_force_b.torch) == 0 - assert torch.count_nonzero(cube_object.permanent_wrench_composer.out_torque_b.torch) == 0 - - def simulate(): - for _ in range(5): - cube_object.write_data_to_sim() - sim.step() - cube_object.update(sim.cfg.dt) - - # Sample a force equal to the weight of the object - external_wrench_b = torch.zeros(cube_object.num_instances, len(body_ids), 6, device=sim.device) - # Every 2nd cube should have a force applied to it - external_wrench_b[0::2, :, 2] = 9.81 * cube_object.data.body_mass.torch[0] + assert cube_object.num_instances == 1 + assert cube_object.data.root_pos_w.torch.shape == (1, 3) + assert cube_object.data.body_mass.torch.shape == (1, 1) + assert cube_object.data.body_inertia.torch.shape == (1, 1, 9) - for i in range(2): - reset_cubes() + target_pose = torch.cat( + [torch.tensor([[0.3, -0.2, 1.5]], device=device), random_orientation(1, device)], dim=-1 + ) + target_vel = torch.randn(1, 6, device=device) + cube_object.write_root_pose_to_sim_index(root_pose=target_pose) + cube_object.write_root_velocity_to_sim_index(root_velocity=target_vel) + torch.testing.assert_close(cube_object.data.root_link_pose_w.torch, target_pose) + torch.testing.assert_close(cube_object.data.root_com_vel_w.torch, target_vel) + torch.testing.assert_close(cube_object.data.body_link_pose_w.torch.squeeze(1), target_pose) - is_global = False - if i % 2 == 0: - is_global = True - positions = cube_object.data.body_com_pos_w.torch[:, body_ids, :3] - else: - positions = None - - # apply force - cube_object.permanent_wrench_composer.set_forces_and_torques_index( - forces=external_wrench_b[..., :3], - torques=external_wrench_b[..., 3:], - positions=positions, - body_ids=body_ids, - is_global=is_global, - ) - simulate() - # First object should still be at the same Z position (1.0) - torch.testing.assert_close( - cube_object.data.root_pos_w.torch[0::2, 2], torch.ones(num_cubes // 2, device=sim.device) - ) - # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(cube_object.data.root_pos_w.torch[1::2, 2] < 1.0) +def test_collection_initialization_with_no_rigid_body() -> None: + """Test that initialization fails when no rigid body is found at the provided prim path.""" + with build_simulation_context(sim_cfg=newton_sim_cfg("cpu")) as sim: + object_collection = spawn_assets({"collection": _collection_cfg(2, rigid=False)}, num_envs=1)["collection"] - # Apply a force at 1m in the Y direction on every 2nd cube - external_wrench_b = torch.zeros(cube_object.num_instances, len(body_ids), 6, device=sim.device) - external_wrench_positions_b = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=sim.device) - external_wrench_b[0::2, :, 2] = 50.0 - external_wrench_positions_b[0::2, :, 1] = 1.0 + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(object_collection) < 10 - for i in range(2): - reset_cubes() - - is_global = False - if i % 2 == 0: - is_global = True - body_com_pos_w = cube_object.data.body_com_pos_w.torch[:, body_ids, :3] - external_wrench_positions_b[..., 0] = 0.0 - external_wrench_positions_b[..., 1] = 1.0 - external_wrench_positions_b[..., 2] = 0.0 - external_wrench_positions_b += body_com_pos_w - else: - external_wrench_positions_b[..., 0] = 0.0 - external_wrench_positions_b[..., 1] = 1.0 - external_wrench_positions_b[..., 2] = 0.0 - - # apply force - cube_object.permanent_wrench_composer.set_forces_and_torques_index( - forces=external_wrench_b[..., :3], - torques=external_wrench_b[..., 3:], - positions=external_wrench_positions_b, - body_ids=body_ids, - is_global=is_global, - ) - cube_object.permanent_wrench_composer.add_forces_and_torques_index( - forces=external_wrench_b[..., :3], - torques=external_wrench_b[..., 3:], - positions=external_wrench_positions_b, - body_ids=body_ids, - is_global=is_global, - ) - simulate() + with pytest.raises(RuntimeError, match="Expected 1 prims at"): + sim.reset() - # The first object should be rotating around it's X axis - assert torch.all(torch.abs(cube_object.data.root_ang_vel_b.torch[0::2, 0]) > 0.1) - # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(cube_object.data.root_pos_w.torch[1::2, 2] < 1.0) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) +def test_collection_single_instance_initialization(device: str) -> None: + """A single-environment, single-body collection initializes with singleton buffers.""" + with build_simulation_context(sim_cfg=newton_sim_cfg(device)) as sim: + object_collection = spawn_assets({"collection": _collection_cfg(1)}, num_envs=1)["collection"] + sim.reset() + assert object_collection.is_initialized + assert object_collection.num_instances == 1 + assert object_collection.root_view.count == 1 + assert len(object_collection.body_names) == 1 + assert object_collection.data.default_body_pose.torch.shape == (1, 1, 7) + assert object_collection.data.body_link_pos_w.torch.shape == (1, 1, 3) + assert object_collection.data.body_link_quat_w.torch.shape == (1, 1, 4) + assert object_collection.data.body_mass.torch.shape == (1, 1) + assert object_collection.data.body_inertia.torch.shape == (1, 1, 9) + + +## +# Shared scene. +## + + +@dataclass +class _Scene: + """Rigid-object cubes and cube collections that share one real Newton model and solver lifecycle.""" + + sim: SimulationContext + cube: RigidObject + collection: RigidObjectCollection + origins: torch.Tensor + device: str + inertial: dict[str, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] + """Configured masses [kg], center-of-mass positions [m], and inertias [kg·m^2] per asset, restored before + each test.""" + + def step(self, num_steps: int = 1) -> None: + """Step ``num_steps`` times, writing the assets' commands before and updating them after each step.""" + for _ in range(num_steps): + self.cube.write_data_to_sim() + self.collection.write_data_to_sim() + self.sim.step() + self.cube.update(self.sim.cfg.dt) + self.collection.update(self.sim.cfg.dt) + + def rest(self) -> None: + """Put every cube at rest at its configured pose.""" + root_pose = self.cube.data.default_root_pose.torch.clone() + root_pose[:, :3] += self.origins + self.cube.write_root_pose_to_sim_index(root_pose=root_pose) + self.cube.write_root_velocity_to_sim_index( + root_velocity=torch.zeros_like(self.cube.data.default_root_vel.torch) + ) + self.cube.reset() + body_pose = self.collection.data.default_body_pose.torch.clone() + body_pose[..., :3] += self.origins.unsqueeze(1) + self.collection.write_body_link_pose_to_sim_index(body_poses=body_pose) + self.collection.write_body_com_velocity_to_sim_index( + body_velocities=torch.zeros_like(self.collection.data.default_body_vel.torch) + ) + self.collection.reset() + + +@pytest.fixture(scope="module", params=test_devices(DeviceScope.CPU)) +def shared_scene(request: pytest.FixtureRequest) -> Iterator[_Scene]: + """Initialize the cubes and the collections once for this module.""" + device = request.param + with build_simulation_context(sim_cfg=newton_sim_cfg(device)) as sim: + assets = spawn_assets( + { + "cube": _cube_cfg(), + "collection": _collection_cfg( + _NUM_CUBES, object_paths=("Object_0", "RobotLeft/Body", "RobotRight/Body") + ), + } + ) + sim.reset() + yield _Scene( + sim=sim, + cube=assets["cube"], + collection=assets["collection"], + origins=env_origins(device), + device=device, + inertial={ + name: ( + asset.data.body_mass.torch.clone(), + asset.data.body_com_pos_b.torch.clone(), + asset.data.body_inertia.torch.clone(), + ) + for name, asset in assets.items() + }, + ) -@pytest.mark.isaacsim_ci -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_rigid_body_set_material_properties(num_cubes, device): - """Material randomization writes friction and restitution into the Newton model shapes of the selected envs.""" - with _newton_sim_context(device, gravity_enabled=True, add_ground_plane=True, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - # Generate cubes scene - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, device=device, add_ground_plane=True) - # Play sim - replicate(sim.get_clone_plan()) - sim.reset() +@pytest.fixture +def scene(shared_scene: _Scene) -> _Scene: + """Hand each test the shared scene at rest with the configured inertial properties.""" + for asset, (masses, coms, inertias) in ( + (shared_scene.cube, shared_scene.inertial["cube"]), + (shared_scene.collection, shared_scene.inertial["collection"]), + ): + asset.set_masses_index(masses=masses) + asset.set_coms_index(coms=coms) + asset.set_inertias_index(inertias=inertias) + shared_scene.rest() + return shared_scene - # Resolve each cube's shapes from the flat Newton model, independent of the asset's view binding. - model = SimulationManager.get_model() - body_world = model.body_world.numpy() - shape_body = model.shape_body.numpy() - cube_shapes = [ - np.flatnonzero(np.isin(shape_body, np.flatnonzero(body_world == index))) for index in range(num_cubes) - ] - assert all(len(shapes) > 0 for shapes in cube_shapes) - original_mu = model.shape_material_mu.numpy().copy() - original_restitution = model.shape_material_restitution.numpy().copy() - - # Randomize the materials of the last cube through the event term, with degenerate ranges. - env = SimpleNamespace(scene={"cube": cube_object}, sim=sim, device=device, num_envs=num_cubes) - params = { - "static_friction_range": (0.55, 0.55), - "dynamic_friction_range": (0.55, 0.55), - "restitution_range": (0.15, 0.15), - "num_buckets": 1, - "asset_cfg": SceneEntityCfg("cube"), - } - term = randomize_rigid_body_material( - EventTermCfg(func=randomize_rigid_body_material, mode="startup", params=params), env - ) - term(env, torch.tensor([num_cubes - 1], device=device), **params) - # Simulate physics - sim.step() - cube_object.update(sim.cfg.dt) +@pytest.mark.isaacsim_ci +def test_external_force_on_single_body(scene: _Scene) -> None: + """A permanent wrench reaches the solver in the local and the global frame, and a reset clears it. - mu = model.shape_material_mu.numpy() - restitution = model.shape_material_restitution.numpy() - np.testing.assert_allclose(mu[cube_shapes[-1]], 0.55) - np.testing.assert_allclose(restitution[cube_shapes[-1]], 0.15) - # Shapes of the other cubes are untouched. - for shapes in cube_shapes[:-1]: - np.testing.assert_array_equal(mu[shapes], original_mu[shapes]) - np.testing.assert_array_equal(restitution[shapes], original_restitution[shapes]) + Cube 0 receives a force equal to its weight and hovers while cube 1 falls. Cube 0 then receives a force set + and added 1 m off its center along y and must turn about x while cube 1 keeps falling. + """ + cube = scene.cube + composer = cube.permanent_wrench_composer + zeros = torch.zeros(cube.num_instances, 1, 3, device=scene.device) + weight, lift, lever = zeros.clone(), zeros.clone(), zeros.clone() + weight[0, :, 2] = 9.81 * cube.data.body_mass.torch[0] + lift[0, :, 2] = 50.0 + lever[..., 1] = 1.0 + + with world_gravity(_GRAVITY): + for is_global in (True, False): + for forces, offset, writes in ((weight, zeros, 1), (lift, lever, 2)): + scene.rest() + for wrench_composer in (composer, cube.instantaneous_wrench_composer): + assert not wrench_composer.active + assert torch.count_nonzero(wrench_composer.out_force_b.torch) == 0 + assert torch.count_nonzero(wrench_composer.out_torque_b.torch) == 0 + positions = offset + cube.data.body_com_pos_w.torch if is_global else offset + for write in (composer.set_forces_and_torques_index, composer.add_forces_and_torques_index)[:writes]: + write(forces=forces, torques=zeros, positions=positions, is_global=is_global) + scene.step(5) + + if forces is weight: + torch.testing.assert_close(cube.data.root_pos_w.torch[0, 2], torch.tensor(1.0, device=scene.device)) + else: + assert cube.data.root_ang_vel_b.torch[0, 0].abs() > 0.1 + assert cube.data.root_pos_w.torch[1, 2] < 1.0 @pytest.mark.isaacsim_ci -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_rigid_body_set_mass(num_cubes, device): - """Test that selected mass writes update inverse mass and inertia across static transitions.""" - with _newton_sim_context(device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - origins = np.asarray([(float(index), 0.0, 1.0) for index in range(num_cubes)], dtype=np.float32) - sim_utils.create_prim("/World/Env_0", "Xform", translation=origins[0]) - cfg = RigidObjectCfg( - prim_path="/World/Env_[^/]*/Object", - spawn=sim_utils.CuboidCfg( - size=(0.2, 0.2, 0.2), - rigid_props=PhysxRigidBodyCfg(disable_gravity=True), - mass_props=sim_utils.MassCfg(mass=1.0), - collision_props=sim_utils.UsdPhysicsCollisionCfg(), - ), - ) - clone_plan_from_env_0(CloneCfg(clone_template="/World/Env_{}"), (cfg,), num_cubes, 1.0, positions=origins) - cube_object = RigidObject(cfg) +def test_rigid_body_set_material_properties(scene: _Scene) -> None: + """Material randomization writes friction and restitution into the Newton model shapes of the selected envs.""" + cube_object = scene.cube + num_cubes = cube_object.num_instances + device = scene.device + + # Resolve each cube's shapes from the flat Newton model, independent of the asset's view binding. + model = SimulationManager.get_model() + body_world = model.body_world.numpy() + shape_body = model.shape_body.numpy() + is_cube = np.asarray([label.endswith("/Object") for label in model.body_label]) + cube_shapes = [ + np.flatnonzero(np.isin(shape_body, np.flatnonzero(is_cube & (body_world == index)))) + for index in range(num_cubes) + ] + assert all(len(shapes) > 0 for shapes in cube_shapes) + original_mu = model.shape_material_mu.numpy().copy() + original_restitution = model.shape_material_restitution.numpy().copy() + + # Randomize the materials of the last cube through the event term, with degenerate ranges. + env = SimpleNamespace(scene={"cube": cube_object}, sim=scene.sim, device=device, num_envs=num_cubes) + params = { + "static_friction_range": (0.55, 0.55), + "dynamic_friction_range": (0.55, 0.55), + "restitution_range": (0.15, 0.15), + "num_buckets": 1, + "asset_cfg": SceneEntityCfg("cube"), + } + term = randomize_rigid_body_material( + EventTermCfg(func=randomize_rigid_body_material, mode="startup", params=params), env + ) + term(env, torch.tensor([num_cubes - 1], device=device), **params) - # Play sim - replicate(sim.get_clone_plan()) - sim.reset() + # Simulate physics + scene.step() - # Get masses before updating one environment. - original_masses = cube_object.data.body_mass.torch.clone() - raw_model_inv_mass = cube_object.root_view.get_attribute("body_inv_mass", SimulationManager.get_model())[:, 0] - raw_model_inv_inertia = cube_object.root_view.get_attribute("body_inv_inertia", SimulationManager.get_model())[ - :, 0 - ] - assert cube_object.data._sim_bind_body_inv_mass.ptr == raw_model_inv_mass.ptr - assert cube_object.data._sim_bind_body_inv_inertia.ptr == raw_model_inv_inertia.ptr - model_inv_mass = cube_object.data._sim_bind_body_inv_mass - model_inv_inertia = cube_object.data._sim_bind_body_inv_inertia - original_inv_mass = wp.to_torch(model_inv_mass).clone() - original_inv_inertia = wp.to_torch(model_inv_inertia).clone() - - assert original_masses.shape == (num_cubes, 1) - - env_ids = torch.tensor([1], dtype=torch.int32, device=device) - body_ids = torch.tensor([0], dtype=torch.int32, device=device) - - # A positive-to-zero transition makes the selected body static. - zero_mass = torch.zeros(1, 1, device=device) - cube_object.set_masses_index(masses=zero_mass, env_ids=env_ids, body_ids=body_ids) - torch.testing.assert_close(wp.to_torch(model_inv_mass)[1], torch.zeros_like(original_inv_mass[1])) - torch.testing.assert_close(wp.to_torch(model_inv_inertia)[1], torch.zeros_like(original_inv_inertia[1])) - torch.testing.assert_close(wp.to_torch(model_inv_mass)[0], original_inv_mass[0]) - torch.testing.assert_close(wp.to_torch(model_inv_inertia)[0], original_inv_inertia[0]) - - # Inertia writes keep inverse mass unchanged and respect the body's current static state. - wp.to_torch(model_inv_inertia)[1].copy_(torch.eye(3, device=device).reshape(1, 3, 3)) - inertia_matrix = torch.diag(torch.tensor([2.0, 3.0, 4.0], device=device)) - inertias = inertia_matrix.reshape(1, 1, 9) - cube_object.set_inertias_index(inertias=inertias, env_ids=env_ids, body_ids=body_ids) - torch.testing.assert_close(wp.to_torch(model_inv_mass)[1], torch.zeros_like(original_inv_mass[1])) - torch.testing.assert_close(wp.to_torch(model_inv_inertia)[1], torch.zeros_like(original_inv_inertia[1])) - - # A zero-to-positive transition restores both inverse arrays from current primary data. - masses = original_masses[env_ids][:, body_ids] + 4.0 - cube_object.set_masses_index(masses=masses, env_ids=env_ids, body_ids=body_ids) - torch.testing.assert_close(cube_object.data.body_mass.torch[env_ids][:, body_ids], masses) - torch.testing.assert_close(wp.to_torch(model_inv_mass)[env_ids][:, body_ids], masses.reciprocal()) - torch.testing.assert_close( - wp.to_torch(model_inv_inertia)[env_ids][:, body_ids], - torch.linalg.inv(inertia_matrix).reshape(1, 1, 3, 3), - ) + mu = model.shape_material_mu.numpy() + restitution = model.shape_material_restitution.numpy() + np.testing.assert_allclose(mu[cube_shapes[-1]], 0.55) + np.testing.assert_allclose(restitution[cube_shapes[-1]], 0.15) + # Shapes of the other cubes are untouched. + for shapes in cube_shapes[:-1]: + np.testing.assert_array_equal(mu[shapes], original_mu[shapes]) + np.testing.assert_array_equal(restitution[shapes], original_restitution[shapes]) - # Simulate physics - # perform rendering - sim.step() - # update object - cube_object.update(sim.cfg.dt) - masses_to_check = cube_object.data.body_mass.torch[env_ids][:, body_ids] +@pytest.mark.isaacsim_ci +def test_rigid_body_set_mass(scene: _Scene) -> None: + """Test that selected mass writes update inverse mass and inertia across static transitions.""" + cube_object = scene.cube + num_cubes = cube_object.num_instances + device = scene.device + + # Get masses before updating one environment. + original_masses = cube_object.data.body_mass.torch.clone() + raw_model_inv_mass = cube_object.root_view.get_attribute("body_inv_mass", SimulationManager.get_model())[:, 0] + raw_model_inv_inertia = cube_object.root_view.get_attribute("body_inv_inertia", SimulationManager.get_model())[:, 0] + assert cube_object.data._sim_bind_body_inv_mass.ptr == raw_model_inv_mass.ptr + assert cube_object.data._sim_bind_body_inv_inertia.ptr == raw_model_inv_inertia.ptr + model_inv_mass = cube_object.data._sim_bind_body_inv_mass + model_inv_inertia = cube_object.data._sim_bind_body_inv_inertia + original_inv_mass = wp.to_torch(model_inv_mass).clone() + original_inv_inertia = wp.to_torch(model_inv_inertia).clone() + + assert original_masses.shape == (num_cubes, 1) + + env_ids = torch.tensor([1], dtype=torch.int32, device=device) + body_ids = torch.tensor([0], dtype=torch.int32, device=device) + + # A positive-to-zero transition makes the selected body static. + zero_mass = torch.zeros(1, 1, device=device) + cube_object.set_masses_index(masses=zero_mass, env_ids=env_ids, body_ids=body_ids) + torch.testing.assert_close(wp.to_torch(model_inv_mass)[1], torch.zeros_like(original_inv_mass[1])) + torch.testing.assert_close(wp.to_torch(model_inv_inertia)[1], torch.zeros_like(original_inv_inertia[1])) + torch.testing.assert_close(wp.to_torch(model_inv_mass)[0], original_inv_mass[0]) + torch.testing.assert_close(wp.to_torch(model_inv_inertia)[0], original_inv_inertia[0]) + + # Inertia writes keep inverse mass unchanged and respect the body's current static state. + wp.to_torch(model_inv_inertia)[1].copy_(torch.eye(3, device=device).reshape(1, 3, 3)) + inertia_matrix = torch.diag(torch.tensor([2.0, 3.0, 4.0], device=device)) + inertias = inertia_matrix.reshape(1, 1, 9) + cube_object.set_inertias_index(inertias=inertias, env_ids=env_ids, body_ids=body_ids) + torch.testing.assert_close(wp.to_torch(model_inv_mass)[1], torch.zeros_like(original_inv_mass[1])) + torch.testing.assert_close(wp.to_torch(model_inv_inertia)[1], torch.zeros_like(original_inv_inertia[1])) + + # A zero-to-positive transition restores both inverse arrays from current primary data. + masses = original_masses[env_ids][:, body_ids] + 4.0 + cube_object.set_masses_index(masses=masses, env_ids=env_ids, body_ids=body_ids) + torch.testing.assert_close(cube_object.data.body_mass.torch[env_ids][:, body_ids], masses) + torch.testing.assert_close(wp.to_torch(model_inv_mass)[env_ids][:, body_ids], masses.reciprocal()) + torch.testing.assert_close( + wp.to_torch(model_inv_inertia)[env_ids][:, body_ids], + torch.linalg.inv(inertia_matrix).reshape(1, 1, 3, 3), + ) + + # Simulate physics + scene.step() - # Check if mass is set correctly - torch.testing.assert_close(masses, masses_to_check) + # Check if mass is set correctly + torch.testing.assert_close(masses, cube_object.data.body_mass.torch[env_ids][:, body_ids]) @pytest.mark.isaacsim_ci -@pytest.mark.parametrize("num_cubes", [3]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -def test_gravity_vec_w_tracks_model_gravity(num_cubes, device): +def test_gravity_vec_w_tracks_model_gravity(scene: _Scene) -> None: """Per-env mutations to Newton's ``model.gravity`` reach ``GRAVITY_VEC_W`` and ``projected_gravity_b``. Regression for the pre-fix snapshot: ``GRAVITY_VEC_W`` used to be env 0's gravity broadcast to every env, hiding per-env gravity randomization (e.g. :class:`~isaaclab.envs.mdp.randomize_physics_scene_gravity`). """ - with _newton_sim_context(device, gravity_enabled=True) as sim: - sim._app_control_on_stop_handle = None - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, device=device) - replicate(sim.get_clone_plan()) - sim.reset() + cube_object = scene.cube + num_cubes = cube_object.num_instances + device = scene.device + with world_gravity(_GRAVITY): # Check that gravity is set correctly - torch.testing.assert_close( - cube_object.data.GRAVITY_VEC_W.torch[0], torch.tensor((0.0, 0.0, -9.81), device=device) - ) + torch.testing.assert_close(cube_object.data.GRAVITY_VEC_W.torch[0], torch.tensor(_GRAVITY, device=device)) # The free-falling cubes accelerate with gravity and keep their identity orientation. - for _ in range(2): - sim.step() - cube_object.update(sim.cfg.dt) + scene.step(2) gravity = torch.zeros(num_cubes, 1, 6, device=device) gravity[:, :, 2] = -9.81 torch.testing.assert_close(cube_object.data.body_acc_w.torch, gravity) # One update may span several physics steps, e.g. once per env step when Newton owns decimation. - sim.step() - sim.step() - cube_object.update(2 * sim.cfg.dt) + scene.sim.step() + scene.sim.step() + cube_object.update(2 * scene.sim.cfg.dt) torch.testing.assert_close(cube_object.data.body_acc_w.torch, gravity) # GRAVITY_VEC_W must share storage with Newton's per-env gravity array. @@ -459,7 +455,7 @@ def test_gravity_vec_w_tracks_model_gravity(num_cubes, device): device=device, dtype=torch.float32, ) - env = SimpleNamespace(sim=sim, device=device, num_envs=num_cubes) + env = SimpleNamespace(sim=scene.sim, device=device, num_envs=num_cubes) params = {"gravity_distribution_params": ((0.0, 0.0, 0.0), (0.0, 0.0, 0.0)), "operation": "abs"} event = randomize_physics_scene_gravity(EventTermCfg(func=randomize_physics_scene_gravity, params=params), env) for row, values in enumerate(new_gravity.tolist()): @@ -471,149 +467,115 @@ def test_gravity_vec_w_tracks_model_gravity(num_cubes, device): # Recompute the lazily-cached projected_gravity_b without sim.step, so cube # orientation stays at identity and the projection equals unit-direction gravity. - cube_object.update(sim.cfg.dt) + scene.rest() + cube_object.update(scene.sim.cfg.dt) expected = torch.nn.functional.normalize(new_gravity, dim=-1) torch.testing.assert_close(cube_object.data.projected_gravity_b.torch, expected, atol=1e-5, rtol=1e-5) @pytest.mark.isaacsim_ci -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@flaky(max_runs=3, min_passes=1) -def test_body_root_state_properties(num_cubes, device): +def test_body_root_state_properties(scene: _Scene) -> None: """Test the root_com_state_w, root_link_state_w, body_com_state_w, and body_link_state_w properties.""" - with _newton_sim_context(device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - # Create a scene with random cubes - cube_object, env_pos = generate_cubes_scene(num_cubes=num_cubes, height=0.0, device=device) - - # Play sim - replicate(sim.get_clone_plan()) - sim.reset() - - # Check if cube_object is initialized - assert cube_object.is_initialized - - # change center of mass offset from link frame - offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_cubes, 1) - - # Set center of mass offset via Newton API (position only, no quaternion) - com_pos = offset.unsqueeze(1) # (N, 1, 3) - cube_object.set_coms_index(coms=wp.from_torch(com_pos, dtype=wp.vec3f)) - with wp.ScopedDevice(device): - SimulationManager._solver.notify_model_changed(ModelFlags.BODY_INERTIAL_PROPERTIES) - - # check center of mass has been set - torch.testing.assert_close(cube_object.data.body_com_pos_b.torch.squeeze(1), offset) - - # random z spin velocity (bounded to keep numerical drift within the position tolerance below) - spin_twist = torch.zeros(6, device=device) - spin_twist[5] = 0.5 * torch.randn(1, device=device).clamp(-1.0, 1.0) - - # Simulate physics - for _ in range(100): - # spin the object around Z axis (com) - cube_object.write_root_velocity_to_sim_index(root_velocity=spin_twist.repeat(num_cubes, 1)) - # perform rendering - sim.step() - # update object - cube_object.update(sim.cfg.dt) - - # get state properties - root_link_pose_w = cube_object.data.root_link_pose_w.torch - root_link_vel_w = cube_object.data.root_link_vel_w.torch - root_com_pose_w = cube_object.data.root_com_pose_w.torch - root_com_vel_w = cube_object.data.root_com_vel_w.torch - body_link_pose_w = cube_object.data.body_link_pose_w.torch - body_link_vel_w = cube_object.data.body_link_vel_w.torch - body_com_pose_w = cube_object.data.body_com_pose_w.torch - body_com_vel_w = cube_object.data.body_com_vel_w.torch - - # cubes are spinning around center of mass - # position will not match - # center of mass position will be constant (i.e. spinning around com) - _tol = dict(atol=2e-3, rtol=2e-3) - torch.testing.assert_close(env_pos + offset, root_com_pose_w[..., :3], **_tol) - torch.testing.assert_close(env_pos + offset, body_com_pose_w[..., :3].squeeze(-2), **_tol) - # link position will be moving but should stay constant away from center of mass - root_link_state_pos_rel_com = quat_apply_inverse( - root_link_pose_w[..., 3:], - root_link_pose_w[..., :3] - root_com_pose_w[..., :3], - ) - torch.testing.assert_close(-offset, root_link_state_pos_rel_com, **_tol) - body_link_state_pos_rel_com = quat_apply_inverse( - body_link_pose_w[..., 3:], - body_link_pose_w[..., :3] - body_com_pose_w[..., :3], - ) - torch.testing.assert_close(-offset, body_link_state_pos_rel_com.squeeze(-2), **_tol) - - # orientation of com will be a constant rotation from link orientation - com_quat_b = cube_object.data.body_com_quat_b.torch - com_quat_w = quat_mul(body_link_pose_w[..., 3:], com_quat_b) - torch.testing.assert_close(com_quat_w, body_com_pose_w[..., 3:], **_tol) - torch.testing.assert_close(com_quat_w.squeeze(-2), root_com_pose_w[..., 3:], **_tol) - - # root and body link orientations describe the same rigid body - torch.testing.assert_close(root_link_pose_w[..., 3:], body_link_pose_w[..., 3:].squeeze(-2), **_tol) - - # lin_vel will not match - # center of mass vel will be constant (i.e. spinning around com) - torch.testing.assert_close(torch.zeros_like(root_com_vel_w[..., :3]), root_com_vel_w[..., :3], **_tol) - torch.testing.assert_close(torch.zeros_like(body_com_vel_w[..., :3]), body_com_vel_w[..., :3], **_tol) - # link frame will be moving, and should be equal to input angular velocity cross offset - lin_vel_rel_root_gt = quat_apply_inverse(root_link_pose_w[..., 3:], root_link_vel_w[..., :3]) - lin_vel_rel_body_gt = quat_apply_inverse(body_link_pose_w[..., 3:], body_link_vel_w[..., :3]) - lin_vel_rel_gt = torch.linalg.cross(spin_twist.repeat(num_cubes, 1)[..., 3:], -offset) - torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_root_gt, **_tol) - torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_body_gt.squeeze(-2), **_tol) - - # ang_vel will always match - torch.testing.assert_close(root_com_vel_w[..., 3:], root_link_vel_w[..., 3:]) - torch.testing.assert_close(root_com_vel_w[..., 3:], body_com_vel_w[..., 3:].squeeze(-2)) - torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) + cube_object = scene.cube + num_cubes = cube_object.num_instances + device = scene.device + env_pos = scene.origins + cube_object.data.default_root_pose.torch[:, :3] + + # change center of mass offset from link frame + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_cubes, 1) + cube_object.set_coms_index(coms=wp.from_torch(offset.unsqueeze(1), dtype=wp.vec3f)) + + # check center of mass has been set + torch.testing.assert_close(cube_object.data.body_com_pos_b.torch.squeeze(1), offset) + + # spin about z, fast enough for the link-frame velocity to discriminate + spin_twist = torch.zeros(6, device=device) + spin_twist[5] = 2.0 + + # Simulate physics + for _ in range(20): + # spin the object around Z axis (com) + cube_object.write_root_velocity_to_sim_index(root_velocity=spin_twist.repeat(num_cubes, 1)) + scene.step() + + # get state properties + root_link_pose_w = cube_object.data.root_link_pose_w.torch + root_link_vel_w = cube_object.data.root_link_vel_w.torch + root_com_pose_w = cube_object.data.root_com_pose_w.torch + root_com_vel_w = cube_object.data.root_com_vel_w.torch + body_link_pose_w = cube_object.data.body_link_pose_w.torch + body_link_vel_w = cube_object.data.body_link_vel_w.torch + body_com_pose_w = cube_object.data.body_com_pose_w.torch + body_com_vel_w = cube_object.data.body_com_vel_w.torch + + # cubes are spinning around center of mass + # position will not match + # center of mass position will be constant (i.e. spinning around com) + _tol = dict(atol=2e-3, rtol=2e-3) + torch.testing.assert_close(env_pos + offset, root_com_pose_w[..., :3], **_tol) + torch.testing.assert_close(env_pos + offset, body_com_pose_w[..., :3].squeeze(-2), **_tol) + # link position will be moving but should stay constant away from center of mass + root_link_state_pos_rel_com = quat_apply_inverse( + root_link_pose_w[..., 3:], + root_link_pose_w[..., :3] - root_com_pose_w[..., :3], + ) + torch.testing.assert_close(-offset, root_link_state_pos_rel_com, **_tol) + body_link_state_pos_rel_com = quat_apply_inverse( + body_link_pose_w[..., 3:], + body_link_pose_w[..., :3] - body_com_pose_w[..., :3], + ) + torch.testing.assert_close(-offset, body_link_state_pos_rel_com.squeeze(-2), **_tol) + + # orientation of com will be a constant rotation from link orientation + com_quat_b = cube_object.data.body_com_quat_b.torch + com_quat_w = quat_mul(body_link_pose_w[..., 3:], com_quat_b) + torch.testing.assert_close(com_quat_w, body_com_pose_w[..., 3:], **_tol) + torch.testing.assert_close(com_quat_w.squeeze(-2), root_com_pose_w[..., 3:], **_tol) + + # root and body link orientations describe the same rigid body + torch.testing.assert_close(root_link_pose_w[..., 3:], body_link_pose_w[..., 3:].squeeze(-2), **_tol) + + # lin_vel will not match + # center of mass vel will be constant (i.e. spinning around com) + torch.testing.assert_close(torch.zeros_like(root_com_vel_w[..., :3]), root_com_vel_w[..., :3], **_tol) + torch.testing.assert_close(torch.zeros_like(body_com_vel_w[..., :3]), body_com_vel_w[..., :3], **_tol) + # link frame will be moving, and should be equal to input angular velocity cross offset + lin_vel_rel_root_gt = quat_apply_inverse(root_link_pose_w[..., 3:], root_link_vel_w[..., :3]) + lin_vel_rel_body_gt = quat_apply_inverse(body_link_pose_w[..., 3:], body_link_vel_w[..., :3]) + lin_vel_rel_gt = torch.linalg.cross(spin_twist.repeat(num_cubes, 1)[..., 3:], -offset) + torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_root_gt, **_tol) + torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_body_gt.squeeze(-2), **_tol) + + # ang_vel will always match + torch.testing.assert_close(root_com_vel_w[..., 3:], root_link_vel_w[..., 3:]) + torch.testing.assert_close(root_com_vel_w[..., 3:], body_com_vel_w[..., 3:].squeeze(-2)) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) @pytest.mark.isaacsim_ci -@pytest.mark.parametrize( - ("num_cubes", "state_location"), [(2, "com"), (2, "link"), (2, "root"), (1, "root")] -) # num_cubes=1 covers single-instance initialization -@pytest.mark.parametrize("device", test_devices()) -def test_write_root_state(num_cubes, device, state_location): +def test_write_root_state(scene: _Scene) -> None: """Test the root state setters in the center-of-mass frame, the link frame, and the default root frames. A write must be readable in the written frame and refresh the derived frame and the body-frame caches without a sim step, and the written state must persist into the solver across a step. """ - with _newton_sim_context(device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - # Create a scene with random cubes - cube_object, env_pos = generate_cubes_scene(num_cubes=num_cubes, height=0.0, device=device) - env_idx = torch.tensor([x for x in range(num_cubes)], dtype=torch.int32, device=device) - - # Play sim - replicate(sim.get_clone_plan()) - sim.reset() + cube_object = scene.cube + num_cubes = cube_object.num_instances + device = scene.device + env_pos = scene.origins + env_idx = torch.tensor([x for x in range(num_cubes)], dtype=torch.int32, device=device) - # Check if cube_object is initialized - assert cube_object.is_initialized - - # change center of mass offset from link frame - offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_cubes, 1) - - # Set center of mass offset via Newton API (position only) - com_pos = offset.unsqueeze(1) # (N, 1, 3) - cube_object.set_coms_index(coms=wp.from_torch(com_pos, dtype=wp.vec3f)) - with wp.ScopedDevice(device): - SimulationManager._solver.notify_model_changed(ModelFlags.BODY_INERTIAL_PROPERTIES) + # change center of mass offset from link frame + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_cubes, 1) + cube_object.set_coms_index(coms=wp.from_torch(offset.unsqueeze(1), dtype=wp.vec3f)) - # check center of mass has been set - torch.testing.assert_close(cube_object.data.body_com_pos_b.torch.squeeze(1), offset) + # check center of mass has been set + torch.testing.assert_close(cube_object.data.body_com_pos_b.torch.squeeze(1), offset) + for state_location in ("com", "link", "root"): for i in range(2): # perform step - sim.step() - # update buffers - cube_object.update(sim.cfg.dt) + scene.step() # A target state distinct from the current one in position, orientation, and velocity. target_pose = torch.cat( @@ -644,10 +606,11 @@ def test_write_root_state(num_cubes, device, state_location): # Snapshot the body-frame caches *before* reading the root-frame caches: touching a # root cache lazily recomputes the shared buffer and would mask a stale body cache. - # The body-frame caches must already reflect the write on their own (regression: - # body_com_pose_w returned the pre-write buffer after a link-frame pose write). + # The body-frame caches must reflect the write on their own, including ``body_com_pose_w`` + # after a link-frame pose write and ``body_link_vel_w`` after a center-of-mass velocity write. body_link_pose_w = cube_object.data.body_link_pose_w.torch.squeeze(1).clone() body_com_pose_w = cube_object.data.body_com_pose_w.torch.squeeze(1).clone() + body_link_vel_w = cube_object.data.body_link_vel_w.torch.squeeze(1).clone() body_com_vel_w = cube_object.data.body_com_vel_w.torch.squeeze(1).clone() root_link_pose_w = cube_object.data.root_link_pose_w.torch @@ -689,23 +652,21 @@ def test_write_root_state(num_cubes, device, state_location): # caches reshaped, so they must stay consistent after a write without a sim step. torch.testing.assert_close(root_link_pose_w, body_link_pose_w) torch.testing.assert_close(root_com_pose_w, body_com_pose_w) - torch.testing.assert_close(root_link_vel_w, cube_object.data.body_link_vel_w.torch.squeeze(1)) + torch.testing.assert_close(root_link_vel_w, body_link_vel_w) torch.testing.assert_close(root_com_vel_w, body_com_vel_w) # The written state persists into the solver: with gravity off, one step only integrates the # written velocity. written_pose_w = (root_com_pose_w if state_location == "com" else root_link_pose_w).clone() written_com_vel_w = root_com_vel_w.clone() - sim.step() - cube_object.update(sim.cfg.dt) + scene.step() pose_w = cube_object.data.root_com_pose_w if state_location == "com" else cube_object.data.root_link_pose_w torch.testing.assert_close(pose_w.torch, written_pose_w, rtol=1e-1, atol=1e-1) torch.testing.assert_close(cube_object.data.root_com_vel_w.torch, written_com_vel_w, rtol=1e-1, atol=1e-1) -@pytest.mark.parametrize("device", test_devices()) @pytest.mark.isaacsim_ci -def test_body_link_pose_w_fresh_after_root_pose_write(device): +def test_body_link_pose_w_fresh_after_root_pose_write(scene: _Scene) -> None: """Regression: ``body_link_pose_w`` must reflect a freshly written root pose without an intervening sim step. After ``write_root_{link,com}_pose_to_sim_{index,mask}``, the cached ``_sim_bind_body_link_pose_w`` @@ -720,59 +681,408 @@ def _fk_reset_mask_dirty() -> bool: assert SimulationManager._fk_reset_mask is not None return bool(wp.to_torch(SimulationManager._fk_reset_mask).any().item()) - num_cubes = 2 - with _newton_sim_context(device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=0.5, device=device) + cube_object = scene.cube + num_cubes = cube_object.num_instances + + # Step once so that _sim_timestamp > 0 and caches are primed. + scene.step() + + for writer in ("link_pose_to_sim_index", "link_pose_to_sim_mask", "com_pose_to_sim_index", "com_pose_to_sim_mask"): + # Prime the body_link_pose_w cache with the current pose. + pre_write_pose = cube_object.data.body_link_pose_w.torch.clone().view(num_cubes, 7) + + # Clear the dirty flag so we can observe that the write sets it. + SimulationManager.forward() + assert not _fk_reset_mask_dirty() + + # A target pose distinct from the current one in translation and in orientation (90 deg about z). + target_pose = cube_object.data.root_link_pose_w.torch.clone() + target_pose[:, :3] += torch.tensor([10.0, 5.0, 2.0], device=target_pose.device) + target_pose[:, 3:] = torch.tensor([0.0, 0.0, 0.5**0.5, 0.5**0.5], device=target_pose.device) + getattr(cube_object, f"write_root_{writer}")(root_pose=target_pose) + + # The simulator-side dirty flag must be set before any property read clears it via forward(). + assert _fk_reset_mask_dirty(), f"{writer} pose write must call SimulationManager.invalidate_fk()" + + # Read without stepping: getter must trigger forward kinematics and return the fresh pose. + body_link = cube_object.data.body_link_pose_w.torch.view(num_cubes, 7) + assert not torch.allclose(body_link[..., :3], pre_write_pose[..., :3], rtol=1e-4, atol=1e-4), ( + f"body_link_pose_w returned the pre-write cached pose after {writer}; forward() was not invoked" + ) + torch.testing.assert_close(body_link[..., :3], target_pose[..., :3], rtol=1e-4, atol=1e-4) + # Orientation: compare via |q1 · q2| ≈ 1 to account for the q ≡ -q double cover. + quat_dot = torch.abs((body_link[..., 3:7] * target_pose[..., 3:7]).sum(dim=-1)) + torch.testing.assert_close(quat_dot, torch.ones_like(quat_dot), rtol=1e-4, atol=1e-4) + + +## +# Rigid-object collection. +## + + +def test_collection_initialization(scene: _Scene) -> None: + """Test initialization for prims with rigid body API at the provided prim paths. + + With the rigid-object cube next to the collection in each environment, the collection view must still + select only its configured rigid objects. + """ + object_collection = scene.collection + num_envs = object_collection.num_instances + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(object_collection) < 10 + + # Check if object is initialized + assert object_collection.is_initialized + assert scene.cube.is_initialized + assert object_collection.num_instances == 2 + assert object_collection.root_view.count == num_envs * _NUM_CUBES + assert object_collection.body_names == [f"cube_{i}" for i in range(_NUM_CUBES)] + + # Check buffers that exist and have correct shapes + assert object_collection.data.default_body_pose.torch.shape == (num_envs, _NUM_CUBES, 7) + assert object_collection.data.body_link_pos_w.torch.shape == (num_envs, _NUM_CUBES, 3) + assert object_collection.data.body_link_quat_w.torch.shape == (num_envs, _NUM_CUBES, 4) + assert object_collection.data.body_mass.torch.shape == (num_envs, _NUM_CUBES) + assert object_collection.data.body_inertia.torch.shape == (num_envs, _NUM_CUBES, 9) + + # The selected bodies are the configured cubes, at their configured places, and not the rigid-object cube. + expected_pos = object_collection.data.default_body_pose.torch[..., :3] + scene.origins.unsqueeze(1) + torch.testing.assert_close(object_collection.data.body_link_pos_w.torch, expected_pos) + + +def test_set_body_inertial_properties_updates_inverses(scene: _Scene) -> None: + """Masked inertial-property writes update only selected Newton inverse entries.""" + object_collection = scene.collection + device = scene.device + env_mask = wp.array([True, False], dtype=wp.bool, device=device) + body_mask = wp.array([False, True, False], dtype=wp.bool, device=device) + selected = (0, 1) + + raw_model_inv_mass = object_collection.root_view.get_attribute("body_inv_mass", SimulationManager.get_model())[ + :, :, 0 + ] + assert object_collection.data._sim_bind_body_inv_mass.ptr == raw_model_inv_mass.ptr + model_inv_mass = object_collection.data._sim_bind_body_inv_mass + original_inv_mass = wp.to_torch(model_inv_mass).clone() + masses = object_collection.data.body_mass.torch.clone() + masses[selected] = 4.0 + object_collection.set_masses_mask(masses=masses, env_mask=env_mask, body_mask=body_mask) + + updated_inv_mass = wp.to_torch(model_inv_mass).clone() + torch.testing.assert_close(updated_inv_mass[selected], masses[selected].reciprocal()) + unselected = torch.ones_like(updated_inv_mass, dtype=torch.bool) + unselected[selected] = False + torch.testing.assert_close(updated_inv_mass[unselected], original_inv_mass[unselected]) + + raw_model_inv_inertia = object_collection.root_view.get_attribute( + "body_inv_inertia", SimulationManager.get_model() + )[:, :, 0] + assert object_collection.data._sim_bind_body_inv_inertia.ptr == raw_model_inv_inertia.ptr + model_inv_inertia = object_collection.data._sim_bind_body_inv_inertia + original_inv_inertia = wp.to_torch(model_inv_inertia).clone() + inertias = object_collection.data.body_inertia.torch.clone() + inertia_matrix = torch.diag(torch.tensor([2.0, 3.0, 5.0], device=device)) + inertias[selected] = inertia_matrix.reshape(9) + object_collection.set_inertias_mask(inertias=inertias, env_mask=env_mask, body_mask=body_mask) + + updated_inv_inertia = wp.to_torch(model_inv_inertia) + torch.testing.assert_close(updated_inv_inertia[selected], torch.linalg.inv(inertia_matrix)) + torch.testing.assert_close(updated_inv_inertia[unselected], original_inv_inertia[unselected]) + torch.testing.assert_close(wp.to_torch(model_inv_mass), updated_inv_mass) + + +def test_collection_external_force_on_single_body(scene: _Scene) -> None: + """A permanent wrench reaches the solver in the local and the global frame, and a reset clears it. + + Every other cube receives a force equal to its weight and hovers while the rest fall. Those cubes then + receive a force set and added 1 m off their centers along y and must turn about x. + """ + collection = scene.collection + composer = collection.permanent_wrench_composer + zeros = torch.zeros(collection.num_instances, collection.num_bodies, 3, device=scene.device) + weight, lift, lever = zeros.clone(), zeros.clone(), zeros.clone() + weight[:, 0::2, 2] = 9.81 * collection.data.body_mass.torch[:, 0::2] + lift[:, 0::2, 2] = 50.0 + lever[..., 1] = 1.0 + + with world_gravity(_GRAVITY): + for is_global in (True, False): + for forces, offset, writes in ((weight, zeros, 1), (lift, lever, 2)): + scene.rest() + for wrench_composer in (composer, collection.instantaneous_wrench_composer): + assert torch.count_nonzero(wrench_composer.out_force_b.torch) == 0 + assert torch.count_nonzero(wrench_composer.out_torque_b.torch) == 0 + positions = offset + collection.data.body_link_pos_w.torch if is_global else offset + for write in (composer.set_forces_and_torques_index, composer.add_forces_and_torques_index)[:writes]: + write(forces=forces, torques=zeros, positions=positions, is_global=is_global) + scene.step(10) + + height = collection.data.body_link_pos_w.torch[..., 2] + if forces is weight: + torch.testing.assert_close(height[:, 0::2], torch.ones_like(height[:, 0::2])) + else: + assert torch.all(collection.data.body_com_ang_vel_b.torch[:, 0::2, 0] > 0.1) + assert torch.all(height[:, 1::2] < 1.0) + + +@pytest.mark.isaacsim_ci +def test_collection_gravity_vec_w_tracks_model_gravity(scene: _Scene) -> None: + """Per-env mutations to Newton's ``model.gravity`` reach ``GRAVITY_VEC_W`` and ``projected_gravity_b``. + + Regression for the pre-fix snapshot: ``GRAVITY_VEC_W`` used to be env 0's + gravity broadcast to every env and body, hiding per-env gravity + randomization (e.g. :class:`~isaaclab.envs.mdp.randomize_physics_scene_gravity`). + """ + object_collection = scene.collection + num_envs = object_collection.num_instances + num_cubes = object_collection.num_bodies + device = scene.device + + with world_gravity(_GRAVITY): + # Check if gravity vector is set correctly + torch.testing.assert_close(object_collection.data.GRAVITY_VEC_W.torch[0], torch.tensor(_GRAVITY, device=device)) + + # The free-falling cubes accelerate with gravity and keep their identity orientation. + scene.step(2) + gravity = torch.zeros(num_envs, num_cubes, 6, device=device) + gravity[..., 2] = -9.81 + torch.testing.assert_close(object_collection.data.body_com_acc_w.torch, gravity) + + # An environment update may span multiple physics steps under Newton decimation. + scene.sim.step() + scene.sim.step() + object_collection.update(2 * scene.sim.cfg.dt) + torch.testing.assert_close(object_collection.data.body_com_acc_w.torch, gravity) + + # GRAVITY_VEC_W must share storage with Newton's per-env gravity array. + model = SimulationManager.get_model() + model_gravity_arr = model.gravity[: model.world_count] + global_gravity = wp.to_torch(model.gravity)[-1].clone() + assert object_collection.data.GRAVITY_VEC_W.warp.ptr == model_gravity_arr.ptr + assert object_collection.data.GRAVITY_VEC_W.shape == (num_envs,) + + # Mutate model.gravity per-env in place, as randomize_physics_scene_gravity does. + new_gravity = torch.tensor( + [[0.1 * (i + 1), 0.2 * (i + 1), -3.0 - float(i)] for i in range(num_envs)], + device=device, + dtype=torch.float32, + ) + wp.to_torch(model_gravity_arr).copy_(new_gravity) + SimulationManager.add_model_change(ModelFlags.MODEL_PROPERTIES) + torch.testing.assert_close(wp.to_torch(model.gravity)[-1], global_gravity) + + # Recompute the lazily-cached projected_gravity_b without sim.step: bodies stay + # at identity orientation, so each env's unit gravity broadcasts across its bodies. + object_collection.update(scene.sim.cfg.dt) + expected_per_env = torch.nn.functional.normalize(new_gravity, dim=-1) + expected = expected_per_env.unsqueeze(1).expand(-1, num_cubes, -1).contiguous() + torch.testing.assert_close(object_collection.data.projected_gravity_b.torch, expected, atol=1e-5, rtol=1e-5) + + +def test_object_state_properties(scene: _Scene) -> None: + """Test the object_com_state_w and object_link_state_w properties.""" + cube_object = scene.collection + num_envs = cube_object.num_instances + num_cubes = cube_object.num_bodies + device = scene.device + + # change center of mass offset from link frame + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) + + # Set center of mass offset via Newton API (position only, shape (E, B, 3)) + cube_object.set_coms_index(coms=wp.from_torch(offset, dtype=wp.vec3f)) + + # check center of mass has been set + torch.testing.assert_close(cube_object.data.body_com_pos_b.torch, offset) + + # spin about z, fast enough for the link-frame velocity to discriminate + spin_twist = torch.zeros(6, device=device) + spin_twist[5] = 2.0 + + # initial spawn point + init_com = cube_object.data.body_com_pose_w.torch[..., :3] + + for _ in range(10): + # spin the object around Z axis (com) + cube_object.write_body_com_velocity_to_sim_index(body_velocities=spin_twist.repeat(num_envs, num_cubes, 1)) + scene.step() + + # get state properties + object_link_pose_w = cube_object.data.body_link_pose_w.torch + object_link_vel_w = cube_object.data.body_link_vel_w.torch + object_com_pose_w = cube_object.data.body_com_pose_w.torch + object_com_vel_w = cube_object.data.body_com_vel_w.torch + + _tol = dict(atol=2e-3, rtol=2e-3) + # cubes are spinning around center of mass + # position will not match + # center of mass position will be constant (i.e. spinning around com) + torch.testing.assert_close(init_com, object_com_pose_w[..., :3], **_tol) + + # link position will be moving but should stay constant away from center of mass + object_link_state_pos_rel_com = quat_apply_inverse( + object_link_pose_w[..., 3:], + object_link_pose_w[..., :3] - object_com_pose_w[..., :3], + ) + + torch.testing.assert_close(-offset, object_link_state_pos_rel_com, **_tol) + + # orientation of com will be a constant rotation from link orientation + com_quat_b = cube_object.data.body_com_quat_b.torch + com_quat_w = quat_mul(object_link_pose_w[..., 3:], com_quat_b) + torch.testing.assert_close(com_quat_w, object_com_pose_w[..., 3:], **_tol) + + # lin_vel will not match + # center of mass vel will be constant (i.e. spinning around com) + torch.testing.assert_close(torch.zeros_like(object_com_vel_w[..., :3]), object_com_vel_w[..., :3], **_tol) + + # link frame will be moving, and should be equal to input angular velocity cross offset + lin_vel_rel_object_gt = quat_apply_inverse(object_link_pose_w[..., 3:], object_link_vel_w[..., :3]) + lin_vel_rel_gt = torch.linalg.cross(spin_twist.repeat(num_envs, num_cubes, 1)[..., 3:], -offset) + torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_object_gt, **_tol) + + # ang_vel will always match + torch.testing.assert_close(object_com_vel_w[..., 3:], object_link_vel_w[..., 3:]) - replicate(sim.get_clone_plan()) - sim.reset() - assert cube_object.is_initialized - # Step once so that _sim_timestamp > 0 and caches are primed. - sim.step() - cube_object.update(sim.cfg.dt) - - for writer in ("link_index", "link_mask", "com_index", "com_mask"): - # Prime the body_link_pose_w cache with the current pose. - pre_write_pose = wp.to_torch(cube_object.data.body_link_pose_w).clone().view(num_cubes, 7) - - # Clear the dirty flag so we can observe that the write sets it. - SimulationManager.forward() - assert not _fk_reset_mask_dirty() - - # Build a target pose clearly distinct from the current one in both translation and orientation. - # Quaternion in (x, y, z, w) for 90° about z: [0, 0, sin(pi/4), cos(pi/4)] = [0, 0, sqrt(0.5), sqrt(0.5)]. - target_pose = wp.to_torch(cube_object.data.root_link_pose_w).clone() - target_pose[..., 0] += 10.0 - target_pose[..., 1] += 5.0 - target_pose[..., 2] += 2.0 - sqrt_half = 0.7071067811865476 - target_pose[..., 3] = 0.0 - target_pose[..., 4] = 0.0 - target_pose[..., 5] = sqrt_half - target_pose[..., 6] = sqrt_half - - if writer == "link_index": - cube_object.write_root_link_pose_to_sim_index(root_pose=target_pose) - elif writer == "link_mask": - cube_object.write_root_link_pose_to_sim_mask(root_pose=target_pose) - elif writer == "com_index": - cube_object.write_root_com_pose_to_sim_index(root_pose=target_pose) - elif writer == "com_mask": - cube_object.write_root_com_pose_to_sim_mask(root_pose=target_pose) - - # The simulator-side dirty flag must be set before any property read clears it via forward(). - assert _fk_reset_mask_dirty(), f"{writer} pose write must call SimulationManager.invalidate_fk()" - - # Read without stepping: getter must trigger forward kinematics and return the fresh pose. - body_link = wp.to_torch(cube_object.data.body_link_pose_w).view(num_cubes, 7) - # Defeat alias accidents: the property must not still return the pre-write value. - assert not torch.allclose(body_link[..., :3], pre_write_pose[..., :3], rtol=1e-4, atol=1e-4), ( - f"body_link_pose_w returned the pre-write cached pose after {writer}; forward() was not invoked" +def test_write_object_state(scene: _Scene) -> None: + """Test the object state setters in the center-of-mass frame, the link frame, and the default root frames. + + A write must be readable in the written frame and refresh the derived frame without a sim step, and the + written state must persist into the solver across a step. + """ + cube_object = scene.collection + num_envs = cube_object.num_instances + num_cubes = cube_object.num_bodies + device = scene.device + env_ids = torch.tensor([x for x in range(num_envs)], dtype=torch.int32, device=device) + object_ids = torch.tensor([x for x in range(num_cubes)], dtype=torch.int32, device=device) + + # change center of mass offset from link frame + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) + + # Set center of mass offset via Newton API (position only, shape (E, B, 3)) + cube_object.set_coms_index(coms=wp.from_torch(offset, dtype=wp.vec3f)) + + # check center of mass has been set + torch.testing.assert_close(cube_object.data.body_com_pos_b.torch, offset) + + for state_location in ("com", "link", "root"): + for i in range(2): + scene.step() + + body_link_pose_w = cube_object.data.body_link_pose_w.torch + body_com_pose_w = cube_object.data.body_com_pose_w.torch + object_link_to_com_pos, object_link_to_com_quat = subtract_frame_transforms( + body_link_pose_w[..., :3].reshape(-1, 3), + body_link_pose_w[..., 3:7].reshape(-1, 4), + body_com_pose_w[..., :3].reshape(-1, 3), + body_com_pose_w[..., 3:7].reshape(-1, 4), + ) + + # A target state distinct from the current one in position, orientation, and velocity. + target_pose = torch.cat( + [ + body_link_pose_w[..., :3] + 0.3 * torch.rand(num_envs, num_cubes, 3, device=device), + random_orientation(num_envs * num_cubes, device).view(num_envs, num_cubes, 4), + ], + dim=-1, ) - # Translation must match the write. - torch.testing.assert_close(body_link[..., :3], target_pose[..., :3], rtol=1e-4, atol=1e-4) - # Orientation: compare via |q1 · q2| ≈ 1 to account for the q ≡ -q double cover. - quat_dot = torch.abs((body_link[..., 3:7] * target_pose[..., 3:7]).sum(dim=-1)) - torch.testing.assert_close(quat_dot, torch.ones_like(quat_dot), rtol=1e-4, atol=1e-4) + target_vel = torch.randn(num_envs, num_cubes, 6, device=device) + + # Alternate between the default and explicit environment and body selectors. + selectors = {} if i == 0 else {"env_ids": env_ids, "body_ids": object_ids} + if state_location == "com": + cube_object.write_body_com_pose_to_sim_index(body_poses=target_pose, **selectors) + cube_object.write_body_com_velocity_to_sim_index(body_velocities=target_vel, **selectors) + elif state_location == "link": + cube_object.write_body_link_pose_to_sim_index(body_poses=target_pose, **selectors) + cube_object.write_body_link_velocity_to_sim_index(body_velocities=target_vel, **selectors) + elif state_location == "root": + cube_object.write_body_link_pose_to_sim_index(body_poses=target_pose, **selectors) + cube_object.write_body_com_velocity_to_sim_index(body_velocities=target_vel, **selectors) + + link_pose_w = cube_object.data.body_link_pose_w.torch + link_vel_w = cube_object.data.body_link_vel_w.torch + com_pose_w = cube_object.data.body_com_pose_w.torch + com_vel_w = cube_object.data.body_com_vel_w.torch + if state_location == "com": + torch.testing.assert_close(target_pose, com_pose_w) + torch.testing.assert_close(target_vel, com_vel_w) + # the com pose was written, so the derived link pose must be refreshed + expected_link_pos, expected_link_quat = combine_frame_transforms( + com_pose_w[..., :3].reshape(-1, 3), + com_pose_w[..., 3:].reshape(-1, 4), + quat_rotate(quat_inv(object_link_to_com_quat), -object_link_to_com_pos), + quat_inv(object_link_to_com_quat), + ) + expected_link_pose = torch.cat((expected_link_pos, expected_link_quat), dim=1).view(num_envs, -1, 7) + torch.testing.assert_close(expected_link_pose, link_pose_w) + else: + torch.testing.assert_close(target_pose, link_pose_w) + written_vel_w = link_vel_w if state_location == "link" else com_vel_w + torch.testing.assert_close(target_vel, written_vel_w) + # the link pose was written, so the derived com pose must be refreshed + expected_com_pos, expected_com_quat = combine_frame_transforms( + link_pose_w[..., :3].reshape(-1, 3), + link_pose_w[..., 3:].reshape(-1, 4), + object_link_to_com_pos, + object_link_to_com_quat, + ) + expected_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1).view(num_envs, -1, 7) + torch.testing.assert_close(expected_com_pose, com_pose_w) + # skip lin_vel because it differs between the frames; angular velocity is frame-independent + # and only matches when the derived velocity was actually refreshed after the write + torch.testing.assert_close(com_vel_w[..., 3:], link_vel_w[..., 3:]) + + # The written state persists into the solver: with gravity off, one step only integrates the + # written velocity. + written_pose_w = (com_pose_w if state_location == "com" else link_pose_w).clone() + written_com_vel_w = com_vel_w.clone() + scene.step() + pose_w = cube_object.data.body_com_pose_w if state_location == "com" else cube_object.data.body_link_pose_w + torch.testing.assert_close(pose_w.torch, written_pose_w, rtol=1e-1, atol=1e-1) + torch.testing.assert_close(cube_object.data.body_com_vel_w.torch, written_com_vel_w, rtol=1e-1, atol=1e-1) + + +@pytest.mark.isaacsim_ci +def test_body_pose_write_marks_fk_reset_mask(scene: _Scene) -> None: + """Regression: ``write_body_{link,com}_pose_to_sim_{index,mask}`` must mark FK dirty. + + For a collection, ``_sim_bind_body_link_pose_w`` is bound directly to the simulator's root-transforms + buffer, so the property read is not what becomes stale — the simulator's internal ``body_q`` used by + collision detection is. The write methods must therefore call :meth:`SimulationManager.invalidate_fk` + so downstream consumers re-run forward kinematics before the next step. Without the fix, + ``_fk_reset_mask`` remains unset after an explicit pose write. The buffer-aliasing invariant is + also pinned: a refactor that decouples ``_sim_bind_body_link_pose_w`` from the write target would + silently make the property stale, so we check the post-write pose matches the written value. + """ + + def _fk_reset_mask_dirty() -> bool: + assert SimulationManager._fk_reset_mask is not None + return bool(wp.to_torch(SimulationManager._fk_reset_mask).any().item()) + + cube_object = scene.collection + scene.step() + + for writer in ("link_pose_to_sim_index", "link_pose_to_sim_mask", "com_pose_to_sim_index", "com_pose_to_sim_mask"): + # Clear the dirty flag so we can observe that the write sets it. + SimulationManager.forward() + assert not _fk_reset_mask_dirty() + + pre_write_pose = cube_object.data.body_link_pose_w.torch.clone() + target_pose = pre_write_pose.clone() + target_pose[..., :3] += torch.tensor([10.0, 5.0, 2.0], device=target_pose.device) + getattr(cube_object, f"write_body_{writer}")(body_poses=target_pose) + + assert _fk_reset_mask_dirty(), f"{writer} pose write must call SimulationManager.invalidate_fk()" + + # body_link_pose_w must reflect the write immediately — its underlying buffer is the write + # target. A regression that moves this property to a separate cached buffer (mirroring the + # single-object case) would silently break this invariant. + body_link = cube_object.data.body_link_pose_w.torch + assert not torch.allclose(body_link[..., :3], pre_write_pose[..., :3], rtol=1e-4, atol=1e-4), ( + f"body_link_pose_w still aliases the pre-write pose after {writer}; the buffer was not written" + ) + torch.testing.assert_close(body_link[..., :3], target_pose[..., :3], rtol=1e-4, atol=1e-4) diff --git a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py deleted file mode 100644 index 3612013271fe..000000000000 --- a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py +++ /dev/null @@ -1,702 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -# ignore private usage of variables warning -# pyright: reportPrivateUsage=none - -from isaaclab.test.utils import DeviceScope, launch_test_simulation, test_devices -from isaaclab.utils import clone - -launch_test_simulation() - -import sys - -import numpy as np -import pytest -import torch -import warp as wp -from isaaclab_newton.assets import RigidObjectCollection -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg -from isaaclab_newton.physics import NewtonManager as SimulationManager -from isaaclab_physx.sim.schemas import PhysxRigidBodyCfg -from newton import ModelFlags - -import isaaclab.sim as sim_utils -from isaaclab.assets import AssetBaseCfg, RigidObjectCfg, RigidObjectCollectionCfg -from isaaclab.cloner import CloneCfg, clone_plan_from_env_0, replicate -from isaaclab.sim import SimulationCfg, build_simulation_context -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR -from isaaclab.utils.math import ( - combine_frame_transforms, - quat_apply_inverse, - quat_inv, - quat_mul, - quat_rotate, - random_orientation, - subtract_frame_transforms, -) - -NEWTON_SIM_CFG = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg(), - ), -) - - -def _newton_sim_context(device, gravity_enabled=True, dt=None, **kwargs): - """Helper to create a Newton simulation context with the correct device. - - When sim_cfg is provided to build_simulation_context, the device, gravity_enabled, and dt - kwargs are ignored. This helper applies them to the shared NEWTON_SIM_CFG before calling. - """ - NEWTON_SIM_CFG.device = device - NEWTON_SIM_CFG.gravity = (0.0, 0.0, -9.81) if gravity_enabled else (0.0, 0.0, 0.0) - if dt is not None: - NEWTON_SIM_CFG.dt = dt - return build_simulation_context(device=device, sim_cfg=NEWTON_SIM_CFG, **kwargs) - - -def generate_cubes_scene( - num_envs: int = 1, - num_cubes: int = 1, - height=1.0, - has_api: bool = True, - kinematic_enabled: bool = False, - device: str = "cuda:0", - spawn_unrelated_sibling: bool = False, - object_paths: tuple[str, ...] | None = None, -) -> tuple[RigidObjectCollection, torch.Tensor]: - """Generate a scene with the provided number of cubes. - - Args: - num_envs: Number of envs to generate. - num_cubes: Number of cubes to generate when object_paths is not specified. - height: Height of the cubes. - has_api: Whether the cubes have a rigid body API on them. - kinematic_enabled: Whether the cubes are kinematic. - device: Device to use for the simulation. - spawn_unrelated_sibling: Whether to spawn a rigid body outside the collection in each environment. - object_paths: Member paths relative to each environment. Defaults to Object_0, Object_1, etc. - - Returns: - A tuple containing the rigid object collection representing the cubes and the origins of the cubes. - - """ - origins = np.asarray([(i * 3.0, 0, height) for i in range(num_envs)], dtype=np.float32) - sim_utils.create_prim("/World/Env_0", "Xform", translation=origins[0]) - - # Resolve spawn configuration - if has_api: - spawn_cfg = sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", - rigid_props=sim_utils.UsdPhysicsRigidBodyCfg(kinematic_enabled=kinematic_enabled), - ) - else: - # since no rigid body properties defined, this is just a static collider - spawn_cfg = sim_utils.CuboidCfg( - size=(0.1, 0.1, 0.1), - collision_props=sim_utils.UsdPhysicsCollisionCfg(), - ) - - # create the rigid object configs - cube_config_dict = {} - if object_paths is None: - object_paths = tuple(f"Object_{i}" for i in range(num_cubes)) - for i, path in enumerate(object_paths): - cube_object_cfg = RigidObjectCfg( - prim_path=f"/World/Env_[^/]*/{path}", - spawn=clone(spawn_cfg), - init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 3 * i, height)), - ) - cube_config_dict[f"cube_{i}"] = cube_object_cfg - cfgs = list(cube_config_dict.values()) - if spawn_unrelated_sibling: - cfgs.append(AssetBaseCfg(prim_path="/World/Env_[^/]*/Object_Target")) - clone_plan_from_env_0(CloneCfg(clone_template="/World/Env_{}"), cfgs, num_envs, 3.0, positions=origins) - if spawn_unrelated_sibling: - spawn_cfg.func("/World/Env_0/Object_Target", spawn_cfg, translation=(0.0, -3.0, height)) - # create the rigid object collection - cube_object_collection_cfg = RigidObjectCollectionCfg(rigid_objects=cube_config_dict) - cube_object_collection = RigidObjectCollection(cfg=cube_object_collection_cfg) - - return cube_object_collection, torch.as_tensor(origins, device=device) - - -@pytest.mark.parametrize( - ("num_envs", "object_paths", "spawn_unrelated_sibling"), - [ - (1, ("Object_A",), False), - (2, ("Object_A", "Object_B", "Object_C"), True), - (2, ("Object_A", "Shelf/Object_B"), False), - ], - ids=["single_member", "sibling_sharing_prefix", "nested_member"], -) -@pytest.mark.parametrize("device", test_devices()) -def test_initialization(num_envs, object_paths, spawn_unrelated_sibling, device): - """Bind only configured members, including nested paths, and preserve each body's name-to-pose mapping.""" - num_cubes = len(object_paths) - with _newton_sim_context(device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - object_collection, origins = generate_cubes_scene( - num_envs, device=device, spawn_unrelated_sibling=spawn_unrelated_sibling, object_paths=object_paths - ) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(object_collection) < 10 - - # Play sim - replicate(sim.get_clone_plan()) - sim.reset() - - # Check if object is initialized - assert object_collection.is_initialized - assert object_collection.num_instances == num_envs - assert object_collection.root_view.count == num_envs * num_cubes - assert len(object_collection.body_names) == num_cubes - - # Check buffers that exist and have correct shapes - assert object_collection.data.default_body_pose.torch.shape == (num_envs, num_cubes, 7) - assert object_collection.data.body_link_pos_w.torch.shape == (num_envs, num_cubes, 3) - assert object_collection.data.body_link_quat_w.torch.shape == (num_envs, num_cubes, 4) - assert object_collection.data.body_mass.torch.shape == (num_envs, num_cubes) - assert object_collection.data.body_inertia.torch.shape == (num_envs, num_cubes, 9) - - for name, member_cfg in object_collection.cfg.rigid_objects.items(): - body_ids, _ = object_collection.find_bodies(name) - body_pos = object_collection.data.body_link_pos_w.torch[:, body_ids[0]] - origins - expected_pos = torch.tensor(member_cfg.init_state.pos, device=device).expand(num_envs, 3) - torch.testing.assert_close(body_pos, expected_pos, atol=1e-5, rtol=0.0) - - -@pytest.mark.parametrize("device", test_devices()) -def test_set_body_inertial_properties_updates_inverses(device): - """Masked inertial-property writes update only selected Newton inverse entries.""" - num_envs = 2 - num_cubes = 3 - with _newton_sim_context(device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - origins = np.asarray([(float(env_index), 0.0, 1.0) for env_index in range(num_envs)], dtype=np.float32) - sim_utils.create_prim("/World/Env_0", "Xform", translation=origins[0]) - spawn_cfg = sim_utils.CuboidCfg( - size=(0.2, 0.2, 0.2), - rigid_props=PhysxRigidBodyCfg(disable_gravity=True), - mass_props=sim_utils.MassCfg(mass=1.0), - collision_props=sim_utils.UsdPhysicsCollisionCfg(), - ) - cfg = RigidObjectCollectionCfg( - rigid_objects={ - f"cube_{body_index}": RigidObjectCfg( - prim_path=f"/World/Env_[^/]*/Object_{body_index}", - spawn=clone(spawn_cfg), - init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, float(body_index), 0.0)), - ) - for body_index in range(num_cubes) - } - ) - clone_plan_from_env_0( - CloneCfg(clone_template="/World/Env_{}"), - cfg.rigid_objects.values(), - num_envs, - 1.0, - positions=origins, - ) - object_collection = RigidObjectCollection(cfg) - replicate(sim.get_clone_plan()) - sim.reset() - - env_mask = wp.array([True, False], dtype=wp.bool, device=device) - body_mask = wp.array([False, True, False], dtype=wp.bool, device=device) - selected = (0, 1) - - raw_model_inv_mass = object_collection.root_view.get_attribute("body_inv_mass", SimulationManager.get_model())[ - :, :, 0 - ] - assert object_collection.data._sim_bind_body_inv_mass.ptr == raw_model_inv_mass.ptr - model_inv_mass = object_collection.data._sim_bind_body_inv_mass - original_inv_mass = wp.to_torch(model_inv_mass).clone() - masses = object_collection.data.body_mass.torch.clone() - masses[selected] = 4.0 - object_collection.set_masses_mask(masses=masses, env_mask=env_mask, body_mask=body_mask) - - updated_inv_mass = wp.to_torch(model_inv_mass).clone() - torch.testing.assert_close(updated_inv_mass[selected], masses[selected].reciprocal()) - unselected = torch.ones_like(updated_inv_mass, dtype=torch.bool) - unselected[selected] = False - torch.testing.assert_close(updated_inv_mass[unselected], original_inv_mass[unselected]) - - raw_model_inv_inertia = object_collection.root_view.get_attribute( - "body_inv_inertia", SimulationManager.get_model() - )[:, :, 0] - assert object_collection.data._sim_bind_body_inv_inertia.ptr == raw_model_inv_inertia.ptr - model_inv_inertia = object_collection.data._sim_bind_body_inv_inertia - original_inv_inertia = wp.to_torch(model_inv_inertia).clone() - inertias = object_collection.data.body_inertia.torch.clone() - inertia_matrix = torch.diag(torch.tensor([2.0, 3.0, 5.0], device=device)) - inertias[selected] = inertia_matrix.reshape(9) - object_collection.set_inertias_mask(inertias=inertias, env_mask=env_mask, body_mask=body_mask) - - updated_inv_inertia = wp.to_torch(model_inv_inertia) - torch.testing.assert_close(updated_inv_inertia[selected], torch.linalg.inv(inertia_matrix)) - torch.testing.assert_close(updated_inv_inertia[unselected], original_inv_inertia[unselected]) - torch.testing.assert_close(wp.to_torch(model_inv_mass), updated_inv_mass) - - -def test_initialization_with_no_rigid_body(): - """Test that initialization fails when no rigid body is found at the provided prim path.""" - with _newton_sim_context("cpu", auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - object_collection, _ = generate_cubes_scene(num_cubes=2, has_api=False, device="cpu") - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(object_collection) < 10 - - # Play sim - replicate(sim.get_clone_plan()) - with pytest.raises(RuntimeError, match="Expected 1 prims at"): - sim.reset() - - -@pytest.mark.parametrize("num_envs", [2]) -@pytest.mark.parametrize("num_cubes", [4]) -@pytest.mark.parametrize("device", test_devices()) -def test_external_force_on_single_body(num_envs, num_cubes, device): - """Test application of external force on the base of the object. - - In the first phase, a force equal to the weight of every 2nd object keeps it in place while the others - fall. In the second phase, the force is applied at 1m in the Y direction and the object must rotate - around its X axis. Both phases run in the global and the local frame, and resetting the collection - must clear the wrench applied in the previous iteration. - """ - with _newton_sim_context(device, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - object_collection, origins = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) - replicate(sim.get_clone_plan()) - sim.reset() - - # find objects to apply the force - object_ids, object_names = object_collection.find_bodies(".*") - - def reset_objects(): - # reset object state; shift the cubes to their origins so they are not on top of each other - body_pose = object_collection.data.default_body_pose.torch.clone() - body_vel = object_collection.data.default_body_vel.torch.clone() - body_pose[..., :2] += origins.unsqueeze(1)[..., :2] - object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) - object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) - object_collection.reset() - - # Reset should zero external forces and torques - assert torch.count_nonzero(object_collection.instantaneous_wrench_composer.out_force_b.torch) == 0 - assert torch.count_nonzero(object_collection.instantaneous_wrench_composer.out_torque_b.torch) == 0 - assert torch.count_nonzero(object_collection.permanent_wrench_composer.out_force_b.torch) == 0 - assert torch.count_nonzero(object_collection.permanent_wrench_composer.out_torque_b.torch) == 0 - - def simulate(): - for _ in range(10): - object_collection.write_data_to_sim() - sim.step() - object_collection.update(sim.cfg.dt) - - # Sample a force equal to the weight of the object - external_wrench_b = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=sim.device) - # Every 2nd cube should have a force applied to it - external_wrench_b[:, 0::2, 2] = 9.81 * object_collection.data.body_mass.torch[:, 0::2] - - for i in range(2): - reset_objects() - - is_global = False - if i % 2 == 0: - positions = object_collection.data.body_link_pos_w.torch[:, object_ids, :3] - is_global = True - else: - positions = None - - # apply force - object_collection.permanent_wrench_composer.set_forces_and_torques_index( - forces=external_wrench_b[..., :3], - torques=external_wrench_b[..., 3:], - positions=positions, - body_ids=object_ids, - env_ids=None, - is_global=is_global, - ) - simulate() - - # First object should still be at the same Z position (1.0) - torch.testing.assert_close( - object_collection.data.body_link_pos_w.torch[:, 0::2, 2], - torch.ones_like(object_collection.data.body_link_pos_w.torch[:, 0::2, 2]), - ) - # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(object_collection.data.body_link_pos_w.torch[:, 1::2, 2] < 1.0) - - # Apply a force at 1m in the Y direction on every 2nd cube - external_wrench_b = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=sim.device) - external_wrench_positions_b = torch.zeros( - object_collection.num_instances, len(object_ids), 3, device=sim.device - ) - external_wrench_b[:, 0::2, 2] = 50.0 - external_wrench_positions_b[:, 0::2, 1] = 1.0 - - for i in range(2): - reset_objects() - - is_global = False - if i % 2 == 0: - body_com_pos_w = object_collection.data.body_link_pos_w.torch[:, object_ids, :3] - external_wrench_positions_b[..., 0] = 0.0 - external_wrench_positions_b[..., 1] = 1.0 - external_wrench_positions_b[..., 2] = 0.0 - external_wrench_positions_b += body_com_pos_w - is_global = True - else: - external_wrench_positions_b[..., 0] = 0.0 - external_wrench_positions_b[..., 1] = 1.0 - external_wrench_positions_b[..., 2] = 0.0 - - # apply force - object_collection.permanent_wrench_composer.set_forces_and_torques_index( - forces=external_wrench_b[..., :3], - torques=external_wrench_b[..., 3:], - positions=external_wrench_positions_b, - body_ids=object_ids, - env_ids=None, - is_global=is_global, - ) - object_collection.permanent_wrench_composer.add_forces_and_torques_index( - forces=external_wrench_b[..., :3], - torques=external_wrench_b[..., 3:], - positions=external_wrench_positions_b, - body_ids=object_ids, - is_global=is_global, - ) - simulate() - - # First object should be rotating around it's X axis - assert torch.all(object_collection.data.body_com_ang_vel_b.torch[:, 0::2, 0] > 0.1) - # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(object_collection.data.body_link_pos_w.torch[:, 1::2, 2] < 1.0) - - -@pytest.mark.isaacsim_ci -@pytest.mark.parametrize("num_envs", [3]) -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -def test_gravity_vec_w_tracks_model_gravity(num_envs, num_cubes, device): - """Per-env mutations to Newton's ``model.gravity`` reach ``GRAVITY_VEC_W`` and ``projected_gravity_b``. - - Regression for the pre-fix snapshot: ``GRAVITY_VEC_W`` used to be env 0's - gravity broadcast to every env and body, hiding per-env gravity - randomization (e.g. :class:`~isaaclab.envs.mdp.randomize_physics_scene_gravity`). - """ - with _newton_sim_context(device, gravity_enabled=True, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) - replicate(sim.get_clone_plan()) - sim.reset() - - # Check if gravity vector is set correctly - torch.testing.assert_close( - object_collection.data.GRAVITY_VEC_W.torch[0], torch.tensor((0.0, 0.0, -9.81), device=device) - ) - - # The free-falling cubes accelerate with gravity and keep their identity orientation. - for _ in range(2): - sim.step() - object_collection.update(sim.cfg.dt) - gravity = torch.zeros(num_envs, num_cubes, 6, device=device) - gravity[..., 2] = -9.81 - torch.testing.assert_close(object_collection.data.body_com_acc_w.torch, gravity) - - # One update may span several physics steps, e.g. once per env step when Newton owns decimation. - sim.step() - sim.step() - object_collection.update(2 * sim.cfg.dt) - torch.testing.assert_close(object_collection.data.body_com_acc_w.torch, gravity) - - # GRAVITY_VEC_W must share storage with Newton's per-env gravity array. - model = SimulationManager.get_model() - model_gravity_arr = model.gravity[: model.world_count] - global_gravity = wp.to_torch(model.gravity)[-1].clone() - assert object_collection.data.GRAVITY_VEC_W.warp.ptr == model_gravity_arr.ptr - assert object_collection.data.GRAVITY_VEC_W.shape == (num_envs,) - - # Mutate model.gravity per-env in place, as randomize_physics_scene_gravity does. - new_gravity = torch.tensor( - [[0.1 * (i + 1), 0.2 * (i + 1), -3.0 - float(i)] for i in range(num_envs)], - device=device, - dtype=torch.float32, - ) - wp.to_torch(model_gravity_arr).copy_(new_gravity) - SimulationManager.add_model_change(ModelFlags.MODEL_PROPERTIES) - torch.testing.assert_close(wp.to_torch(model.gravity)[-1], global_gravity) - - # Recompute the lazily-cached projected_gravity_b without sim.step: bodies stay - # at identity orientation, so each env's unit gravity broadcasts across its bodies. - object_collection.update(sim.cfg.dt) - expected_per_env = torch.nn.functional.normalize(new_gravity, dim=-1) - expected = expected_per_env.unsqueeze(1).expand(-1, num_cubes, -1).contiguous() - torch.testing.assert_close(object_collection.data.projected_gravity_b.torch, expected, atol=1e-5, rtol=1e-5) - - -@pytest.mark.parametrize("num_envs", [4]) -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_object_state_properties(num_envs, num_cubes, device): - """Test the object_com_state_w and object_link_state_w properties.""" - with _newton_sim_context(device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_object, env_pos = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, height=0.0, device=device) - - replicate(sim.get_clone_plan()) - sim.reset() - - # check if cube_object is initialized - assert cube_object.is_initialized - - # change center of mass offset from link frame - offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) - - # Set center of mass offset via Newton API (position only, shape (E, B, 3)) - cube_object.set_coms_index(coms=wp.from_torch(offset, dtype=wp.vec3f)) - # Flush the model change immediately so it takes effect before the next step - with wp.ScopedDevice(device): - SimulationManager._solver.notify_model_changed(ModelFlags.BODY_INERTIAL_PROPERTIES) - - # check center of mass has been set - torch.testing.assert_close(cube_object.data.body_com_pos_b.torch, offset) - - # random z spin velocity - spin_twist = torch.zeros(6, device=device) - spin_twist[5] = torch.randn(1, device=device) - - # initial spawn point - init_com = cube_object.data.body_com_pose_w.torch[..., :3] - - for i in range(10): - # spin the object around Z axis (com) - cube_object.write_body_com_velocity_to_sim_index(body_velocities=spin_twist.repeat(num_envs, num_cubes, 1)) - sim.step() - cube_object.update(sim.cfg.dt) - - # get state properties - object_link_pose_w = cube_object.data.body_link_pose_w.torch - object_link_vel_w = cube_object.data.body_link_vel_w.torch - object_com_pose_w = cube_object.data.body_com_pose_w.torch - object_com_vel_w = cube_object.data.body_com_vel_w.torch - - _tol = dict(atol=2e-3, rtol=2e-3) - # cubes are spinning around center of mass - # position will not match - # center of mass position will be constant (i.e. spinning around com) - torch.testing.assert_close(init_com, object_com_pose_w[..., :3], **_tol) - - # link position will be moving but should stay constant away from center of mass - object_link_state_pos_rel_com = quat_apply_inverse( - object_link_pose_w[..., 3:], - object_link_pose_w[..., :3] - object_com_pose_w[..., :3], - ) - - torch.testing.assert_close(-offset, object_link_state_pos_rel_com, **_tol) - - # orientation of com will be a constant rotation from link orientation - com_quat_b = cube_object.data.body_com_quat_b.torch - com_quat_w = quat_mul(object_link_pose_w[..., 3:], com_quat_b) - torch.testing.assert_close(com_quat_w, object_com_pose_w[..., 3:], **_tol) - - # lin_vel will not match - # center of mass vel will be constant (i.e. spinning around com) - torch.testing.assert_close( - torch.zeros_like(object_com_vel_w[..., :3]), - object_com_vel_w[..., :3], - **_tol, - ) - - # link frame will be moving, and should be equal to input angular velocity cross offset - lin_vel_rel_object_gt = quat_apply_inverse(object_link_pose_w[..., 3:], object_link_vel_w[..., :3]) - lin_vel_rel_gt = torch.linalg.cross(spin_twist.repeat(num_envs, num_cubes, 1)[..., 3:], -offset) - torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_object_gt, **_tol) - - # ang_vel will always match - torch.testing.assert_close(object_com_vel_w[..., 3:], object_link_vel_w[..., 3:]) - - -@pytest.mark.parametrize("num_envs", [3]) -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("state_location", ["com", "link", "root"]) -def test_write_object_state(num_envs, num_cubes, device, state_location): - """Test the object state setters in the center-of-mass frame, the link frame, and the default root frames. - - A write must be readable in the written frame and refresh the derived frame without a sim step, and the - written state must persist into the solver across a step. - """ - with _newton_sim_context(device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - # Create a scene with random cubes - cube_object, env_pos = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, height=0.0, device=device) - env_ids = torch.tensor([x for x in range(num_envs)], dtype=torch.int32, device=device) - object_ids = torch.tensor([x for x in range(num_cubes)], dtype=torch.int32, device=device) - - replicate(sim.get_clone_plan()) - sim.reset() - - # Check if cube_object is initialized - assert cube_object.is_initialized - - # change center of mass offset from link frame - offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) - - # Set center of mass offset via Newton API (position only, shape (E, B, 3)) - cube_object.set_coms_index(coms=wp.from_torch(offset, dtype=wp.vec3f)) - # Flush the model change immediately so it takes effect before the next step - with wp.ScopedDevice(device): - SimulationManager._solver.notify_model_changed(ModelFlags.BODY_INERTIAL_PROPERTIES) - - # check center of mass has been set - torch.testing.assert_close(cube_object.data.body_com_pos_b.torch, offset) - - for i in range(2): - sim.step() - cube_object.update(sim.cfg.dt) - - body_link_pose_w = cube_object.data.body_link_pose_w.torch - body_com_pose_w = cube_object.data.body_com_pose_w.torch - object_link_to_com_pos, object_link_to_com_quat = subtract_frame_transforms( - body_link_pose_w[..., :3].view(-1, 3), - body_link_pose_w[..., 3:7].view(-1, 4), - body_com_pose_w[..., :3].view(-1, 3), - body_com_pose_w[..., 3:7].view(-1, 4), - ) - - # A target state distinct from the current one in position, orientation, and velocity. - target_pose = torch.cat( - [ - body_link_pose_w[..., :3] + 0.3 * torch.rand(num_envs, num_cubes, 3, device=device), - random_orientation(num_envs * num_cubes, device).view(num_envs, num_cubes, 4), - ], - dim=-1, - ) - target_vel = torch.randn(num_envs, num_cubes, 6, device=device) - - # Alternate between the default and explicit environment and body selectors. - selectors = {} if i == 0 else {"env_ids": env_ids, "body_ids": object_ids} - if state_location == "com": - cube_object.write_body_com_pose_to_sim_index(body_poses=target_pose, **selectors) - cube_object.write_body_com_velocity_to_sim_index(body_velocities=target_vel, **selectors) - elif state_location == "link": - cube_object.write_body_link_pose_to_sim_index(body_poses=target_pose, **selectors) - cube_object.write_body_link_velocity_to_sim_index(body_velocities=target_vel, **selectors) - elif state_location == "root": - cube_object.write_body_link_pose_to_sim_index(body_poses=target_pose, **selectors) - cube_object.write_body_com_velocity_to_sim_index(body_velocities=target_vel, **selectors) - - link_pose_w = cube_object.data.body_link_pose_w.torch - link_vel_w = cube_object.data.body_link_vel_w.torch - com_pose_w = cube_object.data.body_com_pose_w.torch - com_vel_w = cube_object.data.body_com_vel_w.torch - if state_location == "com": - torch.testing.assert_close(target_pose, com_pose_w) - torch.testing.assert_close(target_vel, com_vel_w) - # the com pose was written, so the derived link pose must be refreshed - expected_link_pos, expected_link_quat = combine_frame_transforms( - com_pose_w[..., :3].view(-1, 3), - com_pose_w[..., 3:].view(-1, 4), - quat_rotate(quat_inv(object_link_to_com_quat), -object_link_to_com_pos), - quat_inv(object_link_to_com_quat), - ) - expected_link_pose = torch.cat((expected_link_pos, expected_link_quat), dim=1).view(num_envs, -1, 7) - torch.testing.assert_close(expected_link_pose, link_pose_w) - else: - torch.testing.assert_close(target_pose, link_pose_w) - written_vel_w = link_vel_w if state_location == "link" else com_vel_w - torch.testing.assert_close(target_vel, written_vel_w) - # the link pose was written, so the derived com pose must be refreshed - expected_com_pos, expected_com_quat = combine_frame_transforms( - link_pose_w[..., :3].view(-1, 3), - link_pose_w[..., 3:].view(-1, 4), - object_link_to_com_pos, - object_link_to_com_quat, - ) - expected_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1).view(num_envs, -1, 7) - torch.testing.assert_close(expected_com_pose, com_pose_w) - # skip lin_vel because it differs between the frames; angular velocity is frame-independent - # and only matches when the derived velocity was actually refreshed after the write - torch.testing.assert_close(com_vel_w[..., 3:], link_vel_w[..., 3:]) - - # The written state persists into the solver: with gravity off, one step only integrates the - # written velocity. - written_pose_w = (com_pose_w if state_location == "com" else link_pose_w).clone() - written_com_vel_w = com_vel_w.clone() - sim.step() - cube_object.update(sim.cfg.dt) - pose_w = cube_object.data.body_com_pose_w if state_location == "com" else cube_object.data.body_link_pose_w - torch.testing.assert_close(pose_w.torch, written_pose_w, rtol=1e-1, atol=1e-1) - torch.testing.assert_close(cube_object.data.body_com_vel_w.torch, written_com_vel_w, rtol=1e-1, atol=1e-1) - - -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.isaacsim_ci -def test_body_pose_write_marks_fk_reset_mask(device): - """Regression: ``write_body_{link,com}_pose_to_sim_{index,mask}`` must mark FK dirty. - - For a collection, ``_sim_bind_body_link_pose_w`` is bound directly to the simulator's root-transforms - buffer, so the property read is not what becomes stale — the simulator's internal ``body_q`` used by - collision detection is. The write methods must therefore call :meth:`SimulationManager.invalidate_fk` - so downstream consumers re-run forward kinematics before the next step. Without the fix, - ``_fk_reset_mask`` remains unset after an explicit pose write. The buffer-aliasing invariant is - also pinned: a refactor that decouples ``_sim_bind_body_link_pose_w`` from the write target would - silently make the property stale, so we check the post-write pose matches the written value. - """ - - def _fk_reset_mask_dirty() -> bool: - assert SimulationManager._fk_reset_mask is not None - return bool(wp.to_torch(SimulationManager._fk_reset_mask).any().item()) - - num_envs = 2 - num_cubes = 2 - with _newton_sim_context(device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_object, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, height=0.5, device=device) - - replicate(sim.get_clone_plan()) - sim.reset() - assert cube_object.is_initialized - - sim.step() - cube_object.update(sim.cfg.dt) - - for writer in ("link_index", "link_mask", "com_index", "com_mask"): - # Clear the dirty flag so we can observe that the write sets it. - SimulationManager.forward() - assert not _fk_reset_mask_dirty() - - pre_write_pose = wp.to_torch(cube_object.data.body_link_pose_w).clone() - - target_pose = wp.to_torch(cube_object.data.body_link_pose_w).clone() - target_pose[..., 0] += 10.0 - target_pose[..., 1] += 5.0 - target_pose[..., 2] += 2.0 - - if writer == "link_index": - cube_object.write_body_link_pose_to_sim_index(body_poses=target_pose) - elif writer == "link_mask": - cube_object.write_body_link_pose_to_sim_mask(body_poses=target_pose) - elif writer == "com_index": - cube_object.write_body_com_pose_to_sim_index(body_poses=target_pose) - elif writer == "com_mask": - cube_object.write_body_com_pose_to_sim_mask(body_poses=target_pose) - - assert _fk_reset_mask_dirty(), f"{writer} pose write must call SimulationManager.invalidate_fk()" - - # body_link_pose_w must reflect the write immediately — its underlying buffer is the write - # target. A regression that moves this property to a separate cached buffer (mirroring the - # single-object case) would silently break this invariant. - body_link = wp.to_torch(cube_object.data.body_link_pose_w) - assert not torch.allclose(body_link[..., :3], pre_write_pose[..., :3], rtol=1e-4, atol=1e-4), ( - f"body_link_pose_w still aliases the pre-write pose after {writer}; the buffer was not written" - ) - torch.testing.assert_close(body_link[..., :3], target_pose[..., :3], rtol=1e-4, atol=1e-4) diff --git a/source/isaaclab_ov/changelog.d/antoiner-asset-tests-redesign.changed.rst b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-redesign.changed.rst new file mode 100644 index 000000000000..bb391540052f --- /dev/null +++ b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-redesign.changed.rst @@ -0,0 +1,3 @@ +* Changed CPU simulations to select CPU dynamics through the attributes of every physics scene in the stage instead + of the OVPhysX process-wide CPU-only mode, which cannot be reverted. To keep OVPhysX off CUDA entirely, set + ``OVPHYSX_DISABLE_GPU``. diff --git a/source/isaaclab_ov/changelog.d/antoiner-asset-tests-redesign.fixed.rst b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-redesign.fixed.rst new file mode 100644 index 000000000000..c1c5dc293ba9 --- /dev/null +++ b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-redesign.fixed.rst @@ -0,0 +1,2 @@ +* Fixed :class:`~isaaclab_ov.physics.OvPhysxManager` refusing to start a CPU simulation after a CUDA simulation, or + the reverse, in the same process. diff --git a/source/isaaclab_ov/isaaclab_ov/physics/ovphysx_manager.py b/source/isaaclab_ov/isaaclab_ov/physics/ovphysx_manager.py index ae5ce203e9e4..9756cb67e106 100644 --- a/source/isaaclab_ov/isaaclab_ov/physics/ovphysx_manager.py +++ b/source/isaaclab_ov/isaaclab_ov/physics/ovphysx_manager.py @@ -25,7 +25,7 @@ import numpy as np import warp as wp -from pxr import Sdf, UsdPhysics +from pxr import Sdf, Usd, UsdPhysics from isaaclab.physics import PhysicsEvent, PhysicsManager from isaaclab.scene_data import SceneDataBackend, SceneDataFormat @@ -86,6 +86,30 @@ def _prepare_default_cache_dir(cache_dir: str) -> str: return cache_dir +def _set_scene_dynamics_device(scene_prim: Usd.Prim, device: str) -> None: + """Select CPU or GPU dynamics, with the matching broadphase, on a physics scene prim. + + The PhysxSchema USD plugin may not be loaded in standalone ovphysx mode, so the ``PhysxSceneAPI`` + entry and the scene attributes are written through raw metadata and attributes. + + Args: + scene_prim: A ``UsdPhysics.Scene`` prim. + device: Resolved physics device, either ``"cpu"`` or ``"gpu"``. + """ + schemas = Sdf.TokenListOp() + current = scene_prim.GetMetadata("apiSchemas") or Sdf.TokenListOp() + items = list(current.prependedItems) if current.prependedItems else [] + if "PhysxSceneAPI" not in items: + items.append("PhysxSceneAPI") + schemas.prependedItems = items + scene_prim.SetMetadata("apiSchemas", schemas) + + # The schema default enables GPU dynamics, so CPU scenes must author it explicitly. + use_gpu = device == "gpu" + scene_prim.CreateAttribute("physxScene:enableGPUDynamics", Sdf.ValueTypeNames.Bool).Set(use_gpu) + scene_prim.CreateAttribute("physxScene:broadphaseType", Sdf.ValueTypeNames.String).Set("GPU" if use_gpu else "MBP") + + logger = logging.getLogger(__name__) @@ -295,7 +319,8 @@ def __init__(self, cfg: OvPhysxBackendCfg): } if is_gpu: carbonite_overrides.update({"/physics/suppressReadback": True, "/physics/suppressFabricUpdate": True}) - ovphysx.PhysX.set_cpu_mode(not is_gpu) + # CPU scenes select CPU dynamics through their scene attributes (see _set_scene_dynamics_device). + # OVPhysX's process-wide CPU-only mode is irreversible and would block later CUDA scenes. self.physx = ovphysx.PhysX( config=ovphysx.PhysXConfig( num_threads=8, cooked_collider_cache_dir=cache_dir, carbonite_overrides=carbonite_overrides @@ -373,8 +398,6 @@ class OvPhysxManager(PhysicsManager): _warmup_done: ClassVar[bool] = False _next_control_ordinal: ClassVar[int] = 2 _requires_full_stage: ClassVar[bool] = False - # Device mode is process-wide; later contexts must reuse the first selected device. - _locked_device: ClassVar[str | None] = None # Retain construction inputs for hard reset; serialization and replay never consume them. _clone_recipes: ClassVar[list[CloneRecipe]] = [] _atexit_registered: ClassVar[bool] = False @@ -479,7 +502,6 @@ def initialize(cls, sim_context: SimulationContext) -> None: happens lazily in :meth:`reset`. The simulation registry retains its native resource across reinitialization. - ``cls._locked_device`` carries the process-wide first-device policy. """ super().initialize(sim_context) cls._ensure_physx_schemas_registered() @@ -728,20 +750,15 @@ def get_scene_data_backend(cls) -> SceneDataBackend: def _warmup_and_load(cls) -> None: """Serialize the USD stage and attach it to the ovphysx runtime. - When no runtime is active, constructs a new :class:`ovphysx.PhysX` - instance. The first construction also records IsaacLab's process device - choice and registers process-exit cleanup. On a forced re-warm before - :meth:`close`, it reuses the active instance, attaches the new USD through - OVStage, rebuilds active clone recipes through full-stage materialization - or runtime replay, and re-runs the supported warmup entry point + When no runtime is active, constructs a new :class:`ovphysx.PhysX` instance for the + simulation device and registers process-exit cleanup. CPU and CUDA simulations may follow + each other in one process. On a forced re-warm before :meth:`close`, it reuses the active + instance, attaches the new USD through OVStage, rebuilds active clone recipes through + full-stage materialization or runtime replay, and re-runs the supported warmup entry point so the new stage's bodies are resident. Raises: - RuntimeError: If ``SimulationContext`` is not set, or if a device - different from IsaacLab's first device choice is requested. - OVPhysX CPU-only mode is process-wide and cannot be reversed; - IsaacLab applies the same conservative policy in both - directions for a predictable lifecycle. + RuntimeError: If ``SimulationContext`` is not set. """ sim = PhysicsManager._sim if sim is None: @@ -754,18 +771,10 @@ def _warmup_and_load(cls) -> None: ovphysx_device = "gpu" if "cuda" in PhysicsManager._device else "cpu" - if cls._locked_device is not None and ovphysx_device != cls._locked_device: - raise RuntimeError( - f"OvPhysxManager is locked to device {cls._locked_device!r} for the lifetime of this process; " - f"cannot switch to {ovphysx_device!r}. IsaacLab pins the first OVPhysX device choice because " - "CPU-only mode cannot be reversed; restart the process to use a different device." - ) - scene_prim = sim.stage.GetPrimAtPath(sim.cfg.physics_prim_path) - if scene_prim.IsValid(): - if cls._clone_recipes: - scene_prim.CreateAttribute("physxScene:envIdInBoundsBitCount", Sdf.ValueTypeNames.Int).Set(4) - cls._configure_physx_scene_prim(scene_prim, PhysicsManager._cfg, ovphysx_device) + if scene_prim.IsValid() and cls._clone_recipes: + scene_prim.CreateAttribute("physxScene:envIdInBoundsBitCount", Sdf.ValueTypeNames.Int).Set(4) + cls._configure_physics_scenes(sim.stage, sim.cfg.physics_prim_path, PhysicsManager._cfg, ovphysx_device) full_stage = cls._requires_full_stage or ovphysx_device == "cpu" stage_usda, native_clones = _serialize_stage(sim.stage, cls._clone_recipes, full_stage, plan) @@ -778,7 +787,6 @@ def _warmup_and_load(cls) -> None: cooked_collider_cache_dir=sim.cfg.physics.cooked_collider_cache_dir, ) ) - cls._locked_device = ovphysx_device if not cls._atexit_registered: atexit.register(cls._close_at_exit) cls._atexit_registered = True @@ -825,19 +833,36 @@ def _close_at_exit(cls) -> None: except Exception: logger.exception("Failed to close OVPhysX during process exit.") + @classmethod + def _configure_physics_scenes( + cls, stage: Usd.Stage, physics_prim_path: str, cfg: OvPhysxCfg | None, device: str + ) -> None: + """Configure every physics scene in a stage for the simulation device. + + The simulation's physics scene receives the full OVPhysX scene configuration. OVPhysX also + simulates every other ``UsdPhysics.Scene`` in the stage, such as one authored by an asset, so + each of those selects the simulation device's dynamics and broadphase. + + Args: + stage: The USD stage to configure. + physics_prim_path: Path of the simulation's physics scene prim. + cfg: The configuration of the simulation's physics scene. + device: Resolved physics device, either ``"cpu"`` or ``"gpu"``. + """ + scene_prim = stage.GetPrimAtPath(physics_prim_path) + if scene_prim.IsValid(): + cls._configure_physx_scene_prim(scene_prim, cfg, device) + for prim in stage.Traverse(): + if prim.IsA(UsdPhysics.Scene) and prim != scene_prim: + _set_scene_dynamics_device(prim, device) + @staticmethod def _configure_physx_scene_prim(scene_prim, cfg, device: str) -> None: - """Apply PhysxSceneAPI schema and device-specific scene attributes to the - scene prim. - - The PhysxSchema USD plugin may not be loaded in standalone ovphysx mode, - so we write the apiSchemas list entry and scene attributes directly via - raw Sdf metadata manipulation instead of using the high-level USD API. + """Apply the OVPhysX scene configuration to the simulation's physics scene prim. - The schema, scene-query-support, and solver-determinism/accuracy attributes are applied - regardless of device. The GPU-specific dynamics/broadphase/capacity attributes are - applied only when ``device == "gpu"`` — without them PhysX defaults to - CPU broadphase even when OVPhysX is configured for GPU execution. + The schema, dynamics device and broadphase (see :func:`_set_scene_dynamics_device`), + scene-query-support, and solver-determinism/accuracy attributes are applied regardless of + device. The GPU buffer-capacity attributes are applied only to GPU scenes. Args: scene_prim: The /World/PhysicsScene prim to configure. @@ -845,13 +870,7 @@ def _configure_physx_scene_prim(scene_prim, cfg, device: str) -> None: values. The GPU buffer-capacity values are only consulted when ``device == "gpu"``. device: Resolved physics device — one of ``"cpu"`` or ``"gpu"``. """ - schemas = Sdf.TokenListOp() - current = scene_prim.GetMetadata("apiSchemas") or Sdf.TokenListOp() - items = list(current.prependedItems) if current.prependedItems else [] - if "PhysxSceneAPI" not in items: - items.append("PhysxSceneAPI") - schemas.prependedItems = items - scene_prim.SetMetadata("apiSchemas", schemas) + _set_scene_dynamics_device(scene_prim, device) # PhysX uses the declared timestep to derive automatic collision contact offsets. sim_cfg = PhysicsManager._sim.cfg @@ -870,17 +889,13 @@ def _configure_physx_scene_prim(scene_prim, cfg, device: str) -> None: cfg.enable_external_forces_every_iteration ) - if device == "gpu": - scene_prim.CreateAttribute("physxScene:enableGPUDynamics", Sdf.ValueTypeNames.Bool).Set(True) - scene_prim.CreateAttribute("physxScene:broadphaseType", Sdf.ValueTypeNames.String).Set("GPU") - - if cfg is not None: - for attr, val in [ - ("gpuMaxRigidContactCount", cfg.gpu_max_rigid_contact_count), - ("gpuMaxRigidPatchCount", cfg.gpu_max_rigid_patch_count), - ("gpuFoundLostPairsCapacity", cfg.gpu_found_lost_pairs_capacity), - ("gpuFoundLostAggregatePairsCapacity", cfg.gpu_found_lost_aggregate_pairs_capacity), - ("gpuTotalAggregatePairsCapacity", cfg.gpu_total_aggregate_pairs_capacity), - ("gpuCollisionStackSize", cfg.gpu_collision_stack_size), - ]: - scene_prim.CreateAttribute(f"physxScene:{attr}", Sdf.ValueTypeNames.UInt).Set(val) + if device == "gpu" and cfg is not None: + for attr, val in [ + ("gpuMaxRigidContactCount", cfg.gpu_max_rigid_contact_count), + ("gpuMaxRigidPatchCount", cfg.gpu_max_rigid_patch_count), + ("gpuFoundLostPairsCapacity", cfg.gpu_found_lost_pairs_capacity), + ("gpuFoundLostAggregatePairsCapacity", cfg.gpu_found_lost_aggregate_pairs_capacity), + ("gpuTotalAggregatePairsCapacity", cfg.gpu_total_aggregate_pairs_capacity), + ("gpuCollisionStackSize", cfg.gpu_collision_stack_size), + ]: + scene_prim.CreateAttribute(f"physxScene:{attr}", Sdf.ValueTypeNames.UInt).Set(val) diff --git a/source/isaaclab_ov/isaaclab_ov/physics/ovphysx_manager_cfg.py b/source/isaaclab_ov/isaaclab_ov/physics/ovphysx_manager_cfg.py index 99c080981368..6c843397b58d 100644 --- a/source/isaaclab_ov/isaaclab_ov/physics/ovphysx_manager_cfg.py +++ b/source/isaaclab_ov/isaaclab_ov/physics/ovphysx_manager_cfg.py @@ -33,7 +33,7 @@ class OvPhysxBackendCfg(BackendCfg): class_type: type[OvPhysxBackend] | str = "{DIR}.ovphysx_manager:OvPhysxBackend" device: str = MISSING - """Normalized device: ``cpu`` or ``cuda:``. CPU-only mode is process-wide.""" + """Normalized device: ``cpu`` or ``cuda:``.""" cooked_collider_cache_dir: str | None = DEFAULT_COOKED_COLLIDER_CACHE_DIR """Cooked-collider cache location, fixed by the first runtime created in the process.""" diff --git a/source/isaaclab_ov/test/assets/test_articulation.py b/source/isaaclab_ov/test/assets/test_articulation.py index f116a632706e..aa1afb4b8737 100644 --- a/source/isaaclab_ov/test/assets/test_articulation.py +++ b/source/isaaclab_ov/test/assets/test_articulation.py @@ -6,131 +6,375 @@ # ignore private usage of variables warning # pyright: reportPrivateUsage=none +"""Real OVPhysX articulation coverage on one module-scoped composite scene per device. -"""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 -the Kit launcher 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 -:attr:`~isaaclab_ov.assets.Articulation.root_view`, an -:class:`~isaaclab_ov.sim.views.OvPhysxView` over the per-tensor-type bindings -(``root_view.get_attribute(tensor_type)`` / -:meth:`~isaaclab_ov.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 --------------------------- - -The OVPhysX runtime fixes device mode (CPU vs GPU) when the process creates -its first ``ovphysx.PhysX`` instance and cannot switch it without a process -restart. :class:`~isaaclab_ov.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'``. Until the wheel exposes a way to reset Carbonite -device state, this is the supported pattern. +Each device builds one scene of isolated articulation islands from the local branching fixture, +resets it once, and keeps it alive for every test that uses it. Each island belongs to one test +unless noted, so those tests do not depend on each other's order. Partial writes select a subset of +environments, joints, or bodies and prove the unselected entries keep their real backend state. + +The CPU scene covers the host-resident property path, where CPU-only OVPhysX bindings are read and +written directly, and the host actuator path. The CUDA scene covers the rest, including pinned-host +staging of CPU-only bindings and the native Newton actuator path. + +Host-side units (data caches, tendon scoping, kernels, actuator control) need no simulation and run first. +Initialization failures need their own scenes. A second simulation context cannot start while a shared +scene is alive, so these tests must run before any shared-scene test in the same session; they are +defined first and pytest runs them before it creates the module-scoped scenes. """ from __future__ import annotations import importlib import logging -import sys +import math +from collections.abc import Iterator +from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace from unittest.mock import Mock +import numpy as np import pytest import torch import warp as wp from pxr import Gf, Sdf, Usd, UsdGeom, UsdPhysics -from isaaclab.test.utils import test_devices -from isaaclab.test.utils.articulation_ordering import ( - ANYMAL_C_PHYSX_JOINT_NAMES, - BRANCHING_MJWARP_BODY_NAMES, - BRANCHING_MJWARP_JOINT_NAMES, - BRANCHING_PHYSX_BODY_NAMES, - BRANCHING_PHYSX_JOINT_NAMES, - PANDA_ROOT_PRESERVING_REVERSED_BODY_NAMES, -) -from isaaclab.utils import clone, replace - -# The OVPhysX runtime wheel is optional. Skip gracefully when it is not installed; -# CI jobs that need OVPhysX coverage install it explicitly. pytest.importorskip("ovphysx.types", reason="ovphysx wheel not installed") from isaaclab_ov import tensor_types as TT # noqa: E402 -from isaaclab_ov.assets import Articulation # noqa: E402 +from isaaclab_ov.assets import Articulation, kernels # noqa: E402 +from isaaclab_ov.assets.articulation import actuator_control # noqa: E402 from isaaclab_ov.assets.articulation.actuator_control import OvPhysxActuatorControl # noqa: E402 from isaaclab_ov.assets.articulation.articulation_data import ArticulationData # noqa: E402 from isaaclab_ov.physics import OvPhysxCfg, OvPhysxManager # noqa: E402 +from isaaclab_ov.test.fixtures.views import MockOvPhysxBindingSet # noqa: E402 from isaaclab_physx.sim.schemas import PhysxJointCfg # noqa: E402 import isaaclab.sim as sim_utils # noqa: E402 import isaaclab.utils.math as math_utils # noqa: E402 -import isaaclab.utils.string as string_utils # noqa: E402 from isaaclab.actuators import DelayedPDActuatorCfg, IdealPDActuatorCfg, ImplicitActuatorCfg # noqa: E402 from isaaclab.assets import ArticulationCfg, get_articulation_name_ordering # noqa: E402 +from isaaclab.envs.mdp import randomize_actuator_gains, randomize_rigid_body_material # noqa: E402 from isaaclab.managers import EventTermCfg, SceneEntityCfg # noqa: E402 from isaaclab.scene import InteractiveScene, InteractiveSceneCfg # noqa: E402 -from isaaclab.sim import SimulationCfg, build_simulation_context # noqa: E402 -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR # noqa: E402 +from isaaclab.sim import SimulationCfg, SimulationContext, build_simulation_context # noqa: E402 +from isaaclab.test.utils import DeviceScope, test_devices # noqa: E402 +from isaaclab.test.utils.articulation_ordering import ( # noqa: E402 + BRANCHING_MJWARP_BODY_NAMES, + BRANCHING_MJWARP_JOINT_NAMES, + BRANCHING_PHYSX_BODY_NAMES, + BRANCHING_PHYSX_JOINT_NAMES, +) +from isaaclab.utils import replace # noqa: E402 from isaaclab.utils.warp.launch_cache import _WarpLaunchCache # noqa: E402 -## -# Pre-defined configs -## -from isaaclab_assets import ANYMAL_C_CFG, CARTPOLE_CFG, FRANKA_PANDA_CFG, SHADOW_HAND_PHYSX_CFG # isort:skip - -wp.init() +from isaaclab_assets import FRANKA_PANDA_CFG # noqa: E402 + +pytestmark = pytest.mark.integration + +_FIXTURE = Path(__file__).parent / "data" / "articulation_ordering_branching.usda" +_STIFFNESS, _DAMPING, _MAX_FORCE = 5.0, 0.5, 100.0 +_ALL_DEVICES = test_devices() +_CUDA_DEVICES = test_devices(DeviceScope.CUDA) +# CUDA-only scene tests keep the CPU row as a skip so that pytest groups every test of one scene device together +# and builds each module-scoped scene once. +_CUDA_SCENES = [ + device if device in _CUDA_DEVICES else pytest.param(device, marks=pytest.mark.skip(reason="CUDA-only path")) + for device in _ALL_DEVICES +] +_ROTATED_JOINT_NAMES = (*BRANCHING_PHYSX_JOINT_NAMES[1:], BRANCHING_PHYSX_JOINT_NAMES[0]) +# The USD default max force (1e10) is effectively unlimited, so the limit islands author a finite one. +_LIMIT_DRIVE_PROPS = (sim_utils.UsdPhysicsDriveCfg(max_force=80.0), PhysxJointCfg(max_joint_velocity=5.0)) +_ROOT_PRESERVING_REVERSED_BODY_NAMES = (BRANCHING_PHYSX_BODY_NAMES[0], *reversed(BRANCHING_PHYSX_BODY_NAMES[1:])) + + +def _implicit(**kwargs) -> dict[str, ImplicitActuatorCfg]: + return {"joints": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=_STIFFNESS, damping=_DAMPING, **kwargs)} + + +@dataclass(frozen=True) +class _IslandCfg: + """Authoring recipe for one isolated articulation island.""" + + actuators: dict + num_envs: int = 2 + fixed_base: bool = True + joint_ordering: tuple[str, ...] | str | None = None + body_ordering: tuple[str, ...] | str | None = None + drives: bool = True + joint_drive_props: tuple | None = None + y_axis_joints: bool = False + com_offset: float | None = None + distinct_coms: bool = False + reversed_elbow: bool = False + collision_shapes: bool = False + tendons: bool = False + articulation_root_prim_path: str | None = None + devices: DeviceScope = DeviceScope.ALL + + +_ISLANDS: dict[str, _IslandCfg] = { + # Nonidentity MJWarp order through cross-backend discovery; shared by the state, property, + # material, and FK-refresh tests, which each check deltas against their own starting state. + "ordered": _IslandCfg(_implicit(), joint_ordering="mjwarp", body_ordering="mjwarp", collision_shapes=True), + "drive": _IslandCfg(_implicit(), joint_ordering=_ROTATED_JOINT_NAMES, devices=DeviceScope.CUDA), + "dynamics": _IslandCfg( + {}, + joint_ordering="mjwarp", + body_ordering="mjwarp", + drives=False, + y_axis_joints=True, + com_offset=0.2, + devices=DeviceScope.CUDA, + ), + "reversed": _IslandCfg( + {}, num_envs=1, drives=False, y_axis_joints=True, com_offset=0.2, reversed_elbow=True, devices=DeviceScope.CUDA + ), + "reversed_ordered": _IslandCfg( + {}, + num_envs=1, + joint_ordering=BRANCHING_MJWARP_JOINT_NAMES, + body_ordering=BRANCHING_MJWARP_BODY_NAMES, + drives=False, + y_axis_joints=True, + com_offset=0.2, + reversed_elbow=True, + devices=DeviceScope.CUDA, + ), + "floating": _IslandCfg( + _implicit(), + fixed_base=False, + joint_ordering=tuple(reversed(BRANCHING_PHYSX_JOINT_NAMES)), + devices=DeviceScope.CUDA, + ), + "tendon": _IslandCfg(_implicit(), tendons=True, devices=DeviceScope.CUDA), + "native": _IslandCfg( + { + "joints": IdealPDActuatorCfg( + joint_names_expr=[".*"], stiffness=_STIFFNESS, damping=_DAMPING, actuator_effort_limit=_MAX_FORCE + ) + }, + joint_ordering=_ROTATED_JOINT_NAMES, + devices=DeviceScope.CUDA, + ), + "delayed": _IslandCfg( + { + "joint": DelayedPDActuatorCfg( + joint_names_expr=[".*"], + stiffness=20.0, + damping=1.0, + actuator_effort_limit=80.0, + min_delay=1, + max_delay=1, + ) + }, + devices=DeviceScope.CUDA, + ), + # Implicit shoulders beside explicit elbows; the CPU scene keeps the explicit group on the host path. + "mixed": _IslandCfg( + { + "cart": ImplicitActuatorCfg( + joint_names_expr=[".*_shoulder"], joint_effort_limit=400.0, stiffness=20.0, damping=0.0 + ), + "pole": IdealPDActuatorCfg( + joint_names_expr=[".*_elbow"], stiffness=20.0, damping=0.0, actuator_effort_limit=400.0 + ), + }, + joint_ordering=_ROTATED_JOINT_NAMES, + ), + "staging": _IslandCfg({}, num_envs=4, drives=False, devices=DeviceScope.CUDA), + # Unset gains, so the actuators adopt the gains of the authored drives. + "usd_gains": _IslandCfg( + {"joints": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=None, damping=None)}, + devices=DeviceScope.CUDA, + ), + # A floating articulation whose root API sits on a rigid body below the asset prim. + "non_root": _IslandCfg( + _implicit(), fixed_base=False, articulation_root_prim_path="/base", devices=DeviceScope.CUDA + ), + "com_order": _IslandCfg( + _implicit(), num_envs=1, body_ordering=_ROOT_PRESERVING_REVERSED_BODY_NAMES, distinct_coms=True + ), + # The authored limits are not the defaults, so solver values equal to them prove the USD path. + "limits_usd": _IslandCfg(_implicit(), num_envs=1, drives=False, joint_drive_props=_LIMIT_DRIVE_PROPS), + "limits_cfg": _IslandCfg( + _implicit(joint_velocity_limit=1e5, joint_effort_limit=1e5), + num_envs=1, + drives=False, + joint_drive_props=_LIMIT_DRIVE_PROPS, + ), +} + + +@dataclass +class _ArticulationScene: + """Articulation islands that share one real OVPhysX lifecycle.""" + + sim: SimulationContext + device: str + native_actuators: bool + islands: dict[str, Articulation] + + def step(self, articulation: Articulation, count: int = 1) -> None: + """Step the shared scene while writing and refreshing one island.""" + for _ in range(count): + articulation.write_data_to_sim() + self.sim.step() + articulation.update(self.sim.cfg.dt) + + +def _apply_api_schema(prim, schema_name: str) -> None: + """Author an applied-schema token without loading the Kit PhysX schema module.""" + schemas = list(prim.GetAppliedSchemas()) + schemas.append(schema_name) + api_schemas = Sdf.TokenListOp() + api_schemas.explicitItems = schemas + prim.SetMetadata("apiSchemas", api_schemas) + + +def _author_tendons(stage, prim_path: str) -> None: + """Author one spatial tendon and one fixed tendon per arm.""" + root_prim = stage.GetPrimAtPath(f"{prim_path}/base") + _apply_api_schema(root_prim, "PhysxTendonAttachmentRootAPI:root") + root_prim.CreateAttribute("physxTendon:root:localPos", Sdf.ValueTypeNames.Point3f).Set(Gf.Vec3f(0.0)) + root_prim.CreateAttribute("physxTendon:root:stiffness", Sdf.ValueTypeNames.Float).Set(5.0) + root_prim.CreateAttribute("physxTendon:root:damping", Sdf.ValueTypeNames.Float).Set(0.5) + root_prim.CreateAttribute("physxTendon:root:limitStiffness", Sdf.ValueTypeNames.Float).Set(1.0) + root_prim.CreateAttribute("physxTendon:root:offset", Sdf.ValueTypeNames.Float).Set(0.0) + leaf_prim = stage.GetPrimAtPath(f"{prim_path}/left_tip") + _apply_api_schema(leaf_prim, "PhysxTendonAttachmentLeafAPI:leaf") + leaf_prim.CreateAttribute("physxTendon:leaf:localPos", Sdf.ValueTypeNames.Point3f).Set(Gf.Vec3f(0.0)) + leaf_prim.CreateAttribute("physxTendon:leaf:parentAttachment", Sdf.ValueTypeNames.Token).Set("root") + leaf_prim.CreateRelationship("physxTendon:leaf:parentLink").SetTargets([root_prim.GetPath()]) + leaf_prim.CreateAttribute("physxTendon:leaf:restLength", Sdf.ValueTypeNames.Float).Set(0.5) + leaf_prim.CreateAttribute("physxTendon:leaf:lowerLimit", Sdf.ValueTypeNames.Float).Set(0.0) + leaf_prim.CreateAttribute("physxTendon:leaf:upperLimit", Sdf.ValueTypeNames.Float).Set(2.0) + + for index, side in enumerate(("left", "right")): + instance = f"t{index}" + shoulder = stage.GetPrimAtPath(f"{prim_path}/{side}_shoulder") + _apply_api_schema(shoulder, f"PhysxTendonAxisRootAPI:{instance}") + for name, value in ( + ("stiffness", 10.0), + ("damping", 1.0), + ("limitStiffness", 0.0), + ("offset", 0.0), + ("restLength", 0.2 + 0.1 * index), + ("lowerLimit", -1.0), + ("upperLimit", 1.0), + ): + shoulder.CreateAttribute(f"physxTendon:{instance}:{name}", Sdf.ValueTypeNames.Float).Set(value) + shoulder.CreateAttribute(f"physxTendon:{instance}:gearing", Sdf.ValueTypeNames.FloatArray).Set([1.0]) + elbow = stage.GetPrimAtPath(f"{prim_path}/{side}_elbow") + _apply_api_schema(elbow, f"PhysxTendonAxisAPI:{instance}") + elbow.CreateAttribute(f"physxTendon:{instance}:gearing", Sdf.ValueTypeNames.FloatArray).Set([1.0]) + + +def _spawn_island(name: str, cfg: _IslandCfg, y_offset: float) -> Articulation: + """Spawn one articulation island and author its USD variations.""" + island_path = f"/World/{name}" + sim_utils.create_prim(island_path, "Xform", translation=(0.0, y_offset, 0.0)) + for env_index in range(cfg.num_envs): + sim_utils.create_prim(f"{island_path}/Env_{env_index}", "Xform", translation=(3.0 * env_index, 0.0, 0.0)) + articulation = Articulation( + ArticulationCfg( + prim_path=f"{island_path}/Env_[^/]*/Robot", + spawn=sim_utils.UsdFileCfg( + usd_path=str(_FIXTURE), + joint_drive_props=None if cfg.joint_drive_props is None else list(cfg.joint_drive_props), + ), + actuators=cfg.actuators, + joint_ordering=cfg.joint_ordering, + body_ordering=cfg.body_ordering, + articulation_root_prim_path=cfg.articulation_root_prim_path, + ) + ) + stage = sim_utils.get_current_stage() + for env_index in range(cfg.num_envs): + prim_path = f"{island_path}/Env_{env_index}/Robot" + if cfg.articulation_root_prim_path is not None: + stage.GetPrimAtPath(prim_path).RemoveAPI(UsdPhysics.ArticulationRootAPI) + UsdPhysics.ArticulationRootAPI.Apply(stage.GetPrimAtPath(prim_path + cfg.articulation_root_prim_path)) + for joint_name in BRANCHING_PHYSX_JOINT_NAMES: + joint_prim = stage.GetPrimAtPath(f"{prim_path}/{joint_name}") + if cfg.drives: + drive = UsdPhysics.DriveAPI.Apply(joint_prim, "angular") + drive.CreateStiffnessAttr(_STIFFNESS) + drive.CreateDampingAttr(_DAMPING) + drive.CreateMaxForceAttr(_MAX_FORCE) + if cfg.y_axis_joints: + UsdPhysics.RevoluteJoint(joint_prim).GetAxisAttr().Set("Y") + for body_index, body_name in enumerate(BRANCHING_PHYSX_BODY_NAMES): + mass = UsdPhysics.MassAPI(stage.GetPrimAtPath(f"{prim_path}/{body_name}")) + if cfg.com_offset is not None: + mass.CreateCenterOfMassAttr(Gf.Vec3f(cfg.com_offset, 0.0, 0.0)) + # Isotropic inertia makes the energy check independent of body rotation. + mass.CreateDiagonalInertiaAttr(Gf.Vec3f(0.1)) + if cfg.distinct_coms: + # Non-identity principal axes make the native setter's float32 normalization observable. + mass.CreateCenterOfMassAttr(Gf.Vec3f(0.01 * body_index, -0.02, 0.03)) + mass.CreatePrincipalAxesAttr(Gf.Quatf(0.9, 0.1 * body_index, 0.2, -0.3).GetNormalized()) + if cfg.reversed_elbow: + joint = UsdPhysics.RevoluteJoint.Get(stage, f"{prim_path}/left_elbow") + body0, body1 = joint.GetBody0Rel().GetTargets(), joint.GetBody1Rel().GetTargets() + joint.GetBody0Rel().SetTargets(body1) + joint.GetBody1Rel().SetTargets(body0) + if cfg.fixed_base: + fixed_joint = UsdPhysics.FixedJoint.Define(stage, f"{prim_path}/fixed_root") + fixed_joint.GetBody1Rel().SetTargets([f"{prim_path}/base"]) + if cfg.collision_shapes: + # Offset along the joint axes so the shapes never meet while the joints rotate. + for body_name, height in (("left_tip", 0.3), ("right_tip", -0.3)): + shape = UsdGeom.Cube.Define(stage, f"{prim_path}/{body_name}/shape") + shape.CreateSizeAttr(0.05) + shape.AddTranslateOp().Set(Gf.Vec3d(0.0, 0.0, height)) + UsdPhysics.CollisionAPI.Apply(shape.GetPrim()) + if cfg.tendons: + _author_tendons(stage, prim_path) + return articulation + + +@pytest.fixture(scope="module") +def scene(request: pytest.FixtureRequest) -> Iterator[_ArticulationScene]: + """Initialize every articulation island for one device once for this module.""" + device = request.param + device_scope = DeviceScope.CUDA if device.startswith("cuda") else DeviceScope.CPU + # The native Newton actuator runtime requires CUDA; the CPU scene keeps explicit actuators on the host. + native_actuators = device.startswith("cuda") + sim_cfg = SimulationCfg(physics=OvPhysxCfg(), device=device, use_newton_actuators=native_actuators) + with build_simulation_context(device=device, sim_cfg=sim_cfg, auto_add_lighting=False) as sim: + islands = { + name: _spawn_island(name, cfg, y_offset=3.0 * index) + for index, (name, cfg) in enumerate(_ISLANDS.items()) + if cfg.devices & device_scope + } + sim.reset() + yield _ArticulationScene(sim=sim, device=device, native_actuators=native_actuators, islands=islands) -pytestmark = pytest.mark.device_split +def _read_binding_to_torch(articulation: Articulation, tensor_type: int, device: str | torch.device) -> torch.Tensor: + """Read an OVPhysX attribute on its native device and move it to *device*.""" + return wp.to_torch(articulation.root_view.get_attribute(tensor_type)).to(device) -def test_prepare_native_actuators_leaves_implicit_only_articulation_on_standard_path(monkeypatch): - """Keep implicit-only articulations on the unchanged solver-drive path.""" - from isaaclab_ov.assets.articulation import actuator_control - runtime_prepare_calls = [] - runtime = SimpleNamespace( - prepare=lambda *args, **kwargs: runtime_prepare_calls.append(True), wrapper=None, adapter=None - ) - articulation = SimpleNamespace( - _sim_cfg=SimpleNamespace(use_newton_actuators=True), - cfg=SimpleNamespace(prim_path="/World/Robot"), +def _spawn_own_articulation(**cfg) -> Articulation: + """Spawn one branching articulation for a test that needs its own scene.""" + return Articulation( + ArticulationCfg( + prim_path="/World/Robot", + spawn=sim_utils.UsdFileCfg(usd_path=str(_FIXTURE), joint_drive_props=list(_LIMIT_DRIVE_PROPS)), + actuators=_implicit(), + **cfg, + ) ) - monkeypatch.setattr(actuator_control, "PhysxActuatorRuntime", lambda *args, **kwargs: runtime) - monkeypatch.setattr(actuator_control, "find_first_matching_prim", lambda _: None) - control = OvPhysxActuatorControl(articulation) - native_groups = control.prepare_native_actuators( - collection=None, - actuator_cfgs={"implicit": ImplicitActuatorCfg(joint_names_expr=["joint"], stiffness=10.0, damping=1.0)}, - ) - assert native_groups == set() - assert not control.native_actuator_path_active - assert not articulation._has_newton_actuators - assert runtime_prepare_calls == [] +## +# Host-side units: data caches, joint directions, tendon scoping, kernels, and actuator control. +# They need no simulation, so they run before the scenes below. +## def test_cached_read_launches_reset_on_ordering_and_invalidation(): @@ -228,82 +472,245 @@ def __init__(self, shape): torch.testing.assert_close(wp.to_torch(user_buffer.data), torch.tensor([[2.0, 1.0]])) -def _read_binding_to_torch(articulation: Articulation, tensor_type: int, device: str | torch.device) -> torch.Tensor: - """Read an OVPhysX attribute 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 we go through the equivalent - :meth:`~isaaclab_ov.sim.views.OvPhysxView.get_attribute`, which returns a - freshly allocated ``float32`` array on the attribute's native device (CPU for - CPU-only property types), then move the result to *device*. - """ - arr = articulation.root_view.get_attribute(tensor_type) - return wp.to_torch(arr).to(device) +def test_joint_dof_sign_resolution_traverses_instance_proxies(): + """Resolve reversed joints inside an instanceable articulation.""" + source_stage = Usd.Stage.CreateInMemory() + UsdGeom.Xform.Define(source_stage, "/Robot") + UsdGeom.Xform.Define(source_stage, "/Robot/base") + UsdGeom.Xform.Define(source_stage, "/Robot/link") + joint = UsdPhysics.RevoluteJoint.Define(source_stage, "/Robot/joint") + joint.GetBody0Rel().SetTargets(["/Robot/link"]) + joint.GetBody1Rel().SetTargets(["/Robot/base"]) + stage = Usd.Stage.CreateInMemory() + instance = UsdGeom.Xform.Define(stage, "/World/Robot").GetPrim() + instance.GetReferences().AddReference(source_stage.GetRootLayer().identifier, "/Robot") + instance.SetInstanceable(True) + + articulation = Mock( + cfg=Mock(prim_path="/World/Robot"), + _joint_names=["joint"], + _body_names=["base", "link"], + ) + assert Articulation._resolve_joint_dof_signs(articulation, stage) == (-1,) -# 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] +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) -@pytest.fixture(autouse=True) -def _ovphysx_skip_other_device(request): - """Skip tests whose ``device`` parameter mismatches the session-locked device. - The OVPhysX runtime locks process-global device mode when the process - creates its first ``ovphysx.PhysX`` instance, so any test parametrized to a - different device after the first ``sim.reset()`` would hit the manager's - :exc:`RuntimeError`. 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 _make_articulation_root_stage_usda() -> str: + """Serialize one relevant articulation subtree and unrelated joints in memory.""" + 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", + ) + return stage.Flatten().ExportToString() -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) - use_newton_actuators = kwargs.pop("use_newton_actuators", False) - 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, - use_newton_actuators=use_newton_actuators, +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, + ) + # The migrated Articulation reads tendon counts off its OvPhysxView; inject the mock + # view over these bindings so the metadata passthrough resolves without a real view. + object.__setattr__(articulation, "_root_view", bindings.view) + 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(): + """Tendon discovery should ignore joints that live outside the current articulation subtree.""" + articulation = _make_articulation_shell() + stage_usda = _make_articulation_root_stage_usda() + old_stage_usda = OvPhysxManager._stage_usda + OvPhysxManager._stage_usda = stage_usda + try: + articulation._process_tendons() + finally: + OvPhysxManager._stage_usda = old_stage_usda + + # the tendon is reported by its schema INSTANCE name, matching PhysX; scope leakage would + # add the identically-named instance from /World/unrelated, giving two entries + assert articulation.fixed_tendon_names == ["inst0"] + assert articulation.spatial_tendon_names == ["spatial_joint"] + + +def _selector(values: list[int], dtype: type) -> wp.array: + return wp.array(values, dtype=dtype, device="cpu") + + +@pytest.mark.parametrize("env_dtype", [wp.int32, wp.int64]) +def test_root_worker_accepts_selector_widths(env_dtype: type) -> None: + env_ids = _selector([1, 0], env_dtype) + data = wp.array( + [[11.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0], [21.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]], + dtype=wp.transformf, + device="cpu", ) - return build_simulation_context(device=device, sim_cfg=sim_cfg, **kwargs) + output = wp.zeros(2, dtype=wp.transformf, device="cpu") + sim_env_ids = wp.empty(2, dtype=wp.int32, device="cpu") + kernel = kernels.set_root_link_pose_to_sim_index + if env_dtype == wp.int64: + kernel = kernels.set_root_link_pose_to_sim_index_kernel(env_ids) + + wp.launch(kernel, dim=2, inputs=[data, env_ids], outputs=[output, sim_env_ids], device="cpu") + + np.testing.assert_array_equal(output.numpy(), data.numpy()[[1, 0]]) + np.testing.assert_array_equal(sim_env_ids.numpy(), [1, 0]) + + +@pytest.mark.parametrize(("env_dtype", "item_dtype"), [(wp.int32, wp.int32), (wp.int64, wp.int64)]) +def test_item_worker_accepts_selector_widths(env_dtype: type, item_dtype: type) -> None: + env_ids = _selector([1, 0], env_dtype) + item_ids = _selector([2, 0], item_dtype) + data = wp.array([[11.0, 12.0], [21.0, 22.0]], dtype=wp.float32, device="cpu") + output = wp.full((2, 3), value=-1.0, dtype=wp.float32, device="cpu") + kernel = kernels.write_2d_data_to_buffer_with_indices + if env_dtype != wp.int32 or item_dtype != wp.int32: + kernel = kernels.write_2d_data_to_buffer_with_indices_kernel(env_ids, item_ids) + + wp.launch(kernel, dim=(2, 2), inputs=[data, env_ids, item_ids], outputs=[output], device="cpu") + + np.testing.assert_array_equal(output.numpy(), [[22.0, -1.0, 21.0], [12.0, -1.0, 11.0]]) + + +def test_prepare_native_actuators_leaves_implicit_only_articulation_on_standard_path(monkeypatch): + """Keep implicit-only articulations on the unchanged solver-drive path.""" + runtime_prepare_calls = [] + runtime = SimpleNamespace( + prepare=lambda *args, **kwargs: runtime_prepare_calls.append(True), wrapper=None, adapter=None + ) + articulation = SimpleNamespace( + _sim_cfg=SimpleNamespace(use_newton_actuators=True), + cfg=SimpleNamespace(prim_path="/World/Robot"), + ) + monkeypatch.setattr(actuator_control, "PhysxActuatorRuntime", lambda *args, **kwargs: runtime) + monkeypatch.setattr(actuator_control, "find_first_matching_prim", lambda _: None) + + control = OvPhysxActuatorControl(articulation) + native_groups = control.prepare_native_actuators( + collection=None, + actuator_cfgs={"implicit": ImplicitActuatorCfg(joint_names_expr=["joint"], stiffness=10.0, damping=1.0)}, + ) + + assert native_groups == set() + assert not control.native_actuator_path_active + assert not articulation._has_newton_actuators + assert runtime_prepare_calls == [] + + +@pytest.mark.parametrize( + "module_name", + [ + "isaaclab_physx.assets.articulation.actuator_control", + "isaaclab_ov.assets.articulation.actuator_control", + ], +) +def test_host_actuator_control_import_does_not_probe_optional_newton_runtime(monkeypatch, module_name): + """Import host controls without probing an unrequested Newton optional dependency.""" + original_find_spec = importlib.util.find_spec + + def reject_newton_probe(name, *args, **kwargs): + if name.startswith("isaaclab_newton"): + raise AssertionError("host actuator-control import eagerly probed Newton") + return original_find_spec(name, *args, **kwargs) + + monkeypatch.setattr(importlib.util, "find_spec", reject_newton_probe) + importlib.reload(importlib.import_module(module_name)) + + +## +# Tests that own their simulation context. Keep them above the composite scene. +## + + +@pytest.mark.parametrize("device", _CUDA_DEVICES) +def test_mimic_finger_follows_commanded_finger(device: str) -> None: + """The passive Franka finger follows its driven leader after a position command. + + Lab submits the same host-side targets on both devices; the mimic constraint itself is solved by PhysX. + """ + cfg = SimulationCfg(physics=OvPhysxCfg(), device=device, gravity=(0.0, 0.0, 0.0), use_newton_actuators=False) + with build_simulation_context(device=device, sim_cfg=cfg) as sim: + articulation = Articulation(replace(FRANKA_PANDA_CFG, prim_path="/World/Franka")) + sim.reset() + leader_id = articulation.find_joints("panda_finger_joint1")[0][0] + follower_id = articulation.find_joints("panda_finger_joint2")[0][0] + initial_leader = articulation.data.joint_pos.torch[:, leader_id].clone() + target = articulation.data.joint_pos.torch.clone() + target[:, [leader_id, follower_id]] = 0.01 + articulation.actuators.target_command.set_position_index(value=target) + for _ in range(50): + articulation.write_data_to_sim() + sim.step() + articulation.update(sim.cfg.dt) + leader = articulation.data.joint_pos.torch[:, leader_id] + follower = articulation.data.joint_pos.torch[:, follower_id] + assert torch.all(torch.abs(leader - initial_leader) > 0.005) + torch.testing.assert_close(follower, leader, rtol=0.0, atol=5.0e-4) -@pytest.mark.parametrize("device", test_devices()) @pytest.mark.parametrize( - "variant", ["geometry", "joint_type", "d6_rotation", "d6_translation", "disabled_joint", "fixed_tendon"] + ("device", "variant"), + # CUDA replays native clones for every compatible variant; CPU serializes the full stage, one variant. + [ + (device, variant) + for device in _CUDA_DEVICES + for variant in ("geometry", "joint_type", "d6_translation", "disabled_joint") + ] + + [(device, "geometry") for device in test_devices(DeviceScope.CPU)] + # The cloner rejects incompatible layouts from USD before any device-specific physics. + + [(_ALL_DEVICES[0], variant) for variant in ("d6_rotation", "fixed_tendon")], ) def test_heterogeneous_articulation_clone_indexed_state(device, variant, tmp_path): """Compatible variants preserve indexed state; incompatible action/tendon layouts are rejected.""" @@ -355,7 +762,9 @@ def test_heterogeneous_articulation_clone_indexed_state(device, variant, tmp_pat stage.Export(path) variants.append(sim_utils.UsdFileCfg(usd_path=path)) - with _ovphysx_sim_context(device=device, gravity_enabled=False) as sim: + with build_simulation_context( + device=device, sim_cfg=SimulationCfg(physics=OvPhysxCfg(), device=device, gravity=(0.0, 0.0, 0.0)) + ) as sim: cfg = InteractiveSceneCfg(num_envs=6, env_spacing=2.0) cfg.robot = ArticulationCfg( prim_path="{ENV_REGEX_NS}/Robot", @@ -397,462 +806,406 @@ def test_heterogeneous_articulation_clone_indexed_state(device, variant, tmp_pat binding.destroy() -def generate_articulation_cfg( - articulation_type: str, - stiffness: float | None = 10.0, - damping: float | None = 2.0, - actuator_velocity_limit: float | None = None, - actuator_effort_limit: float | None = None, - joint_velocity_limit: float | None = None, - joint_effort_limit: 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. - actuator_velocity_limit: Velocity limit for the actuators. Only currently used for "single_joint_implicit" - and "single_joint_explicit". - actuator_effort_limit: Effort limit for explicit actuators. Only currently used for - "single_joint_explicit". - joint_velocity_limit: Velocity limit for the actuators (set into the simulation). - Only currently used for "single_joint_implicit" and "single_joint_explicit". - joint_effort_limit: 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_PHYSX_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.UsdPhysicsDriveCfg(max_force=80.0), - PhysxJointCfg(max_joint_velocity=5.0), - ], - ), - actuators={ - "joint": ImplicitActuatorCfg( - joint_names_expr=[".*"], - joint_effort_limit=joint_effort_limit, - joint_velocity_limit=joint_velocity_limit, - actuator_velocity_limit=actuator_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.UsdPhysicsDriveCfg(max_force=80.0), - PhysxJointCfg(max_joint_velocity=5.0), - ], - ), - actuators={ - "joint": IdealPDActuatorCfg( - joint_names_expr=[".*"], - joint_effort_limit=joint_effort_limit, - joint_velocity_limit=joint_velocity_limit, - actuator_effort_limit=actuator_effort_limit, - actuator_velocity_limit=actuator_velocity_limit, - stiffness=0.0, - damping=10.0, - ), - }, - ) - else: - raise ValueError( - f"Invalid articulation type: {articulation_type}, valid options are 'humanoid', 'panda', 'anymal'," - " 'shadow_hand', 'single_joint_implicit' or 'single_joint_explicit'." - ) - - return articulation_cfg - - -def generate_articulation( - articulation_cfg: ArticulationCfg, num_articulations: int, device: str -) -> tuple[Articulation, torch.tensor]: - """Generate an articulation from a configuration. +@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) +@pytest.mark.parametrize( + ("init_state", "match"), + [ + (ArticulationCfg.InitialStateCfg(joint_pos={"left_shoulder": 10.0}), "default positions out of the limits"), + (ArticulationCfg.InitialStateCfg(joint_vel={"left_shoulder": 100.0}), "default velocities out of the limits"), + ], + ids=["position", "velocity"], +) +def test_out_of_range_default_joint_state(device: str, init_state: ArticulationCfg.InitialStateCfg, match: str) -> None: + """Reject default joint positions outside the joint limits and default velocities above the velocity limits.""" + sim_cfg = SimulationCfg(physics=OvPhysxCfg(), device=device) + with build_simulation_context(device=device, sim_cfg=sim_cfg, auto_add_lighting=False) as sim: + articulation = _spawn_own_articulation(init_state=init_state) + joint = UsdPhysics.RevoluteJoint.Get(sim.stage, "/World/Robot/left_shoulder") + joint.CreateLowerLimitAttr(-90.0) + joint.CreateUpperLimitAttr(90.0) + with pytest.raises(ValueError, match=match): + sim.reset() + assert not articulation.is_initialized + + +@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) +def test_setting_invalid_articulation_root_prim_path(device: str) -> None: + """Reject an explicit articulation root path that matches no prim.""" + sim_cfg = SimulationCfg(physics=OvPhysxCfg(), device=device) + with build_simulation_context(device=device, sim_cfg=sim_cfg, auto_add_lighting=False) as sim: + articulation = _spawn_own_articulation(articulation_root_prim_path="/non_existing_prim_path") + with pytest.raises(RuntimeError, match="Failed to find articulation root prim"): + sim.reset() + assert not articulation.is_initialized + + +@pytest.mark.parametrize("scene", _ALL_DEVICES, indirect=True) +def test_articulation_initialization_and_partial_state(scene: _ArticulationScene) -> None: + """Prove ordering and indexed state writes against the real OVPhysX view.""" + articulation = scene.islands["ordered"] + device = scene.device + assert articulation.is_initialized + assert articulation.is_fixed_base + assert tuple(articulation.joint_names) == BRANCHING_MJWARP_JOINT_NAMES + assert tuple(articulation.body_names) == BRANCHING_MJWARP_BODY_NAMES + assert articulation.joint_ordering is not None + assert articulation.body_ordering is not None + assert articulation.num_instances == 2 + + env_ids = torch.tensor([1], dtype=torch.int32, device=device) + joint_ids = torch.tensor([articulation.num_joints - 1, 0], dtype=torch.int32, device=device) + target_position = torch.tensor([[0.21, -0.13]], device=device) + target_velocity = torch.tensor([[0.41, -0.23]], device=device) + expected_position = articulation.data.joint_pos.torch.clone() + expected_velocity = articulation.data.joint_vel.torch.clone() + expected_position[env_ids[:, None], joint_ids] = target_position + expected_velocity[env_ids[:, None], joint_ids] = target_velocity + articulation.write_joint_state_to_sim_index( + position=target_position, + velocity=target_velocity, + env_ids=env_ids, + joint_ids=joint_ids, + ) + torch.testing.assert_close(articulation.data.joint_pos.torch, expected_position) + torch.testing.assert_close(articulation.data.joint_vel.torch, expected_velocity) + backend_to_user = list(articulation.joint_ordering.backend_to_user_indices) + torch.testing.assert_close( + wp.to_torch(articulation.root_view.get_attribute(TT.DOF_POSITION)), + expected_position[:, backend_to_user], + ) - Handles the creation of the articulation, the environment prims and the articulation's environment - translations + # A position-only partial write (a distinct kernel) preserves every unselected backend joint. + backend_before = wp.to_torch(articulation.root_view.get_attribute(TT.DOF_POSITION)).clone() + backend_joint_id = articulation.joint_ordering.user_to_backend_indices[1] + selected_value = backend_before[0, backend_joint_id] + 0.001 + articulation.write_joint_position_to_sim_index( + position=selected_value.reshape(1, 1), + env_ids=wp.array([0], dtype=wp.int32, device=device), + joint_ids=wp.array([1], dtype=wp.int32, device=device), + ) + expected_backend = backend_before.clone() + expected_backend[0, backend_joint_id] = selected_value + torch.testing.assert_close( + wp.to_torch(articulation.root_view.get_attribute(TT.DOF_POSITION)), expected_backend, rtol=0.0, atol=0.0 + ) - 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. +@pytest.mark.parametrize("scene", _ALL_DEVICES, indirect=True) +def test_articulation_joint_and_body_properties_round_trip(scene: _ArticulationScene) -> None: + """Prove selected joint and body properties reach real OVPhysX bindings and public getters.""" + articulation = scene.islands["ordered"] + device = scene.device + env_ids = torch.tensor([1], dtype=torch.int32, device=device) + joint_ids = torch.tensor([articulation.num_joints - 1, 0], dtype=torch.int32, device=device) + backend_friction_before = wp.to_torch(articulation.root_view.get_attribute(TT.DOF_FRICTION_PROPERTIES)).clone() + public_friction_before = [ + articulation.data.joint_friction_coeff.torch.clone(), + articulation.data.joint_dynamic_friction_coeff.torch.clone(), + articulation.data.joint_viscous_friction_coeff.torch.clone(), + ] + static_friction = torch.tensor([[0.9, 0.7]], device=device) + dynamic_friction = torch.tensor([[0.4, 0.3]], device=device) + viscous_friction = torch.tensor([[0.11, 0.22]], device=device) + articulation.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff=static_friction, + joint_dynamic_friction_coeff=dynamic_friction, + joint_viscous_friction_coeff=viscous_friction, + env_ids=env_ids, + joint_ids=joint_ids, + ) + backend_joint_ids = torch.as_tensor(articulation.joint_ordering.user_to_backend_indices)[joint_ids.cpu()] + expected_backend_friction = backend_friction_before.clone() + expected_backend_friction[1, backend_joint_ids, 0] = static_friction.cpu() + expected_backend_friction[1, backend_joint_ids, 1] = dynamic_friction.cpu() + expected_backend_friction[1, backend_joint_ids, 2] = viscous_friction.cpu() + torch.testing.assert_close( + wp.to_torch(articulation.root_view.get_attribute(TT.DOF_FRICTION_PROPERTIES)), + expected_backend_friction, + ) + # The public getters expose the same coefficients. + for getter, before, written in zip( + ( + articulation.data.joint_friction_coeff, + articulation.data.joint_dynamic_friction_coeff, + articulation.data.joint_viscous_friction_coeff, + ), + public_friction_before, + (static_friction, dynamic_friction, viscous_friction), + ): + expected = before.clone() + expected[env_ids[:, None], joint_ids] = written + torch.testing.assert_close(getter.torch, expected) + + body_ids = torch.tensor([articulation.num_bodies - 1, 1], dtype=torch.int32, device=device) + backend_body_ids = torch.as_tensor(articulation.body_ordering.user_to_backend_indices)[body_ids.cpu()] + raw_mass_before = wp.to_torch(articulation.root_view.get_attribute(TT.BODY_MASS)).clone() + masses = torch.tensor([[2.5, 3.5]], device=device) + articulation.set_masses_index(masses=masses, env_ids=env_ids, body_ids=body_ids) + expected_raw_mass = raw_mass_before.clone() + expected_raw_mass[1, backend_body_ids] = masses.cpu() + torch.testing.assert_close(wp.to_torch(articulation.root_view.get_attribute(TT.BODY_MASS)), expected_raw_mass) + + raw_com_before = wp.to_torch(articulation.root_view.get_attribute(TT.BODY_COM_POSE)).clone() + coms = articulation.data.body_com_pose_b.torch[env_ids][:, body_ids].clone() + coms[0, 0, :3] = torch.tensor([0.02, -0.01, 0.03], device=device) + coms[0, 1, :3] = torch.tensor([-0.03, 0.01, 0.02], device=device) + articulation.set_coms_index(coms=wp.from_torch(coms, dtype=wp.transformf), env_ids=env_ids, body_ids=body_ids) + expected_raw_com = raw_com_before.clone() + expected_raw_com[1, backend_body_ids] = coms.cpu() + torch.testing.assert_close(wp.to_torch(articulation.root_view.get_attribute(TT.BODY_COM_POSE)), expected_raw_com) + + raw_inertia_before = wp.to_torch(articulation.root_view.get_attribute(TT.BODY_INERTIA)).clone() + inertias = articulation.data.body_inertia.torch[env_ids][:, body_ids].clone() + inertias[0, 0, 0] *= 1.2 + inertias[0, 1, 4] *= 1.3 + articulation.set_inertias_index(inertias=inertias, env_ids=env_ids, body_ids=body_ids) + expected_raw_inertia = raw_inertia_before.clone() + expected_raw_inertia[1, backend_body_ids] = inertias.cpu() + torch.testing.assert_close(wp.to_torch(articulation.root_view.get_attribute(TT.BODY_INERTIA)), expected_raw_inertia) + torch.testing.assert_close(articulation.data.body_mass.torch[env_ids][:, body_ids], masses) + torch.testing.assert_close(articulation.data.body_com_pose_b.torch[env_ids][:, body_ids], coms) + torch.testing.assert_close(articulation.data.body_inertia.torch[env_ids][:, body_ids], inertias) + + +@pytest.mark.parametrize("scene", _ALL_DEVICES, indirect=True) +def test_set_material_properties(scene: _ArticulationScene) -> None: + """Randomize per-shape articulation materials through the OVPhysX material event. + OVPhysX exposes per-collision-shape material as the + ``articulation_shape_friction_and_restitution`` tensor binding (shape ``[N, S, 3]`` = + static friction, dynamic friction, restitution), addressed through the + :class:`~isaaclab_ov.sim.views.OvPhysxView`. The binding is CPU-native, so the + buffer lives in host memory. """ - # 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 + articulation = scene.islands["ordered"] + num_articulations = articulation.num_instances + device = scene.device + sim = scene.sim - # 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(replace(articulation_cfg, prim_path="/World/Env_[^/]*/Robot")) - - return articulation, translations + # The ranges exclude the asset's default material, so values inside them prove the write happened. + static_range, dynamic_range, restitution_range = (1.5, 2.0), (1.2, 1.4), (0.6, 0.8) + view = articulation.root_view + materials_before = wp.to_torch(view.get_attribute(TT.SHAPE_FRICTION_AND_RESTITUTION)).clone() + assert materials_before.shape[1] > 0 + assert (materials_before[..., 0] < static_range[0]).all(), f"default material overlaps: {materials_before}" + params = { + "static_friction_range": static_range, + "dynamic_friction_range": dynamic_range, + "restitution_range": restitution_range, + "num_buckets": 16, + "asset_cfg": SceneEntityCfg("robot"), + } + env = SimpleNamespace(scene={"robot": articulation}, sim=sim) + cfg = EventTermCfg(func=randomize_rigid_body_material, mode="reset", params=params) + randomize = randomize_rigid_body_material(cfg, env) + # Randomize only the last environment; the others keep their materials. + randomize(env, torch.tensor([num_articulations - 1], device=device), **cfg.params) + scene.step(articulation) -@pytest.mark.parametrize("device", ["cuda:0"]) -def test_newton_native_explicit_actuator_submits_ovphysx_effort(device): - """Run a Newton-native explicit actuator through the current OVPhysX state and effort binding.""" - stiffness, damping, actuator_effort_limit = 20.0, 1.0, 80.0 - with _ovphysx_sim_context(device=device, gravity_enabled=False, use_newton_actuators=True) as sim: - sim._app_control_on_stop_handle = None - articulation_cfg = replace( - generate_articulation_cfg("single_joint_explicit"), - actuators={ - "joint": IdealPDActuatorCfg( - joint_names_expr=[".*"], - stiffness=stiffness, - damping=damping, - actuator_effort_limit=actuator_effort_limit, - ) - }, - ) - articulation, _ = generate_articulation(articulation_cfg, 1, device) - sim.reset() + materials = wp.to_torch(view.get_attribute(TT.SHAPE_FRICTION_AND_RESTITUTION)) + torch.testing.assert_close(materials[:-1], materials_before[:-1]) + eps = 1e-5 + for component, (lo, hi) in enumerate((static_range, dynamic_range, restitution_range)): + values = materials[-1, :, component] + assert ((values >= lo - eps) & (values <= hi + eps)).all(), values - initial_pos = articulation.data.joint_pos.torch.clone() - target = initial_pos + 0.5 - articulation.actuators.target_command.set_position_index(value=target) - articulation.write_data_to_sim() - - assert articulation._actuator_control.native_actuator_path_active - assert articulation.newton_actuator_adapter is not None - assert torch.any(articulation.actuators.computed_effort.torch != 0.0) - assert torch.any(articulation.actuators.applied_effort.torch != 0.0) - torch.testing.assert_close( - _read_binding_to_torch(articulation, TT.DOF_ACTUATION_FORCE, device), - articulation.actuators.applied_effort.torch, - ) - sim.step() - articulation.update(sim.cfg.dt) - # Use raw OV bindings so the observation cannot refresh the public state shadow. - current_pos = _read_binding_to_torch(articulation, TT.DOF_POSITION, device) - current_vel = _read_binding_to_torch(articulation, TT.DOF_VELOCITY, device) - assert not torch.allclose(current_pos, initial_pos) - - articulation.write_data_to_sim() - expected_effort = torch.clamp( - stiffness * (target - current_pos) - damping * current_vel, - -actuator_effort_limit, - actuator_effort_limit, - ) - torch.testing.assert_close(articulation.actuators.applied_effort.torch, expected_effort) +@pytest.mark.parametrize("scene", _ALL_DEVICES, indirect=True) +@pytest.mark.parametrize("island", ["limits_cfg", "limits_usd"]) +def test_setting_velocity_and_effort_limits_write_to_solver(scene: _ArticulationScene, island: str) -> None: + """Test that the resolved joint velocity and effort limits reach the OVPhysX solver. + The full limit-resolution matrix (config override vs. USD default, implicit and explicit + actuators, actuator-limit soft fallback) is covered on the Newton backend and at unit + level. This test only verifies the OVPhysX write path: the configured limits (or the + USD-authored defaults when unset) land in the native solver buffers and match + ``data.joint_vel_limits`` and ``data.joint_effort_limits``. + """ + articulation = scene.islands[island] + device = scene.device + actuator_cfg = _ISLANDS[island].actuators["joints"] + joint_drive_props = _ISLANDS[island].joint_drive_props -@pytest.mark.parametrize( - "module_name", - [ - "isaaclab_physx.assets.articulation.actuator_control", - "isaaclab_ov.assets.articulation.actuator_control", - ], -) -def test_host_actuator_control_import_does_not_probe_optional_newton_runtime(monkeypatch, module_name): - """Import host controls without probing an unrequested Newton optional dependency.""" - original_find_spec = importlib.util.find_spec + physx_vel_limit = _read_binding_to_torch(articulation, TT.DOF_MAX_VELOCITY, device) + torch.testing.assert_close(articulation.data.joint_vel_limits.torch, physx_vel_limit) + # the solver clamp comes from joint_velocity_limit when set, otherwise the USD-authored value + if actuator_cfg.joint_velocity_limit is None: + limit = next(p.max_joint_velocity for p in joint_drive_props if isinstance(p, PhysxJointCfg)) + else: + limit = actuator_cfg.joint_velocity_limit + expected_velocity_limit = torch.full_like(physx_vel_limit, limit) + torch.testing.assert_close(physx_vel_limit, expected_velocity_limit) - def reject_newton_probe(name, *args, **kwargs): - if name.startswith("isaaclab_newton"): - raise AssertionError("host actuator-control import eagerly probed Newton") - return original_find_spec(name, *args, **kwargs) + physx_effort_limit = _read_binding_to_torch(articulation, TT.DOF_MAX_FORCE, device) + torch.testing.assert_close(articulation.data.joint_effort_limits.torch, physx_effort_limit) + # the solver keeps the USD-authored limit unless the user overrides it explicitly + if actuator_cfg.joint_effort_limit is None: + limit = next(p.max_force for p in joint_drive_props if isinstance(p, sim_utils.UsdPhysicsDriveCfg)) + else: + limit = actuator_cfg.joint_effort_limit + expected_effort_limit = torch.full_like(physx_effort_limit, limit) + torch.testing.assert_close(physx_effort_limit, expected_effort_limit) - monkeypatch.setattr(importlib.util, "find_spec", reject_newton_probe) - importlib.reload(importlib.import_module(module_name)) +@pytest.mark.parametrize("scene", _ALL_DEVICES, indirect=True) +def test_effort_binding_excludes_implicit_pd(scene: _ArticulationScene) -> None: + """Submit implicit feedforward and explicit PD effort without submitting implicit PD telemetry. -@pytest.mark.parametrize("device", ["cuda:0"]) -@pytest.mark.parametrize("use_newton_actuators", [False, True]) -def test_ovphysx_effort_binding_excludes_implicit_pd(device, use_newton_actuators): - """Submit implicit feedforward and explicit PD effort without submitting implicit PD telemetry.""" + The CUDA scene runs the explicit group on the native Newton actuator path; the CPU scene keeps it + on the host actuator path. + """ + articulation = scene.islands["mixed"] stiffness, effort_limit = 20.0, 400.0 - with _ovphysx_sim_context(device=device, gravity_enabled=False, use_newton_actuators=use_newton_actuators) as sim: - sim._app_control_on_stop_handle = None - articulation_cfg = replace( - CARTPOLE_CFG, - actuators={ - "cart": ImplicitActuatorCfg( - joint_names_expr=["slider_to_cart"], - joint_effort_limit=effort_limit, - stiffness=stiffness, - damping=0.0, - ), - "pole": IdealPDActuatorCfg( - joint_names_expr=["cart_to_pole"], - stiffness=stiffness, - damping=0.0, - actuator_effort_limit=effort_limit, - ), - }, - ) - articulation, _ = generate_articulation(articulation_cfg, 1, device) - sim.reset() - - joint_names = articulation.joint_names - cart_id = joint_names.index("slider_to_cart") - pole_id = joint_names.index("cart_to_pole") - initial_pos = articulation.data.joint_pos.torch.clone() - position_target = initial_pos.clone() - position_target[:, cart_id] += 0.25 - position_target[:, pole_id] += 0.5 - feedforward = torch.zeros_like(initial_pos) - feedforward[:, cart_id] = 1.5 - feedforward[:, pole_id] = -0.75 - articulation.actuators.target_command.set_position_index(value=position_target) - articulation.actuators.target_command.set_effort_index(value=feedforward) - articulation.write_data_to_sim() - - expected_pd_effort = torch.clamp( - stiffness * (position_target - initial_pos) + feedforward, -effort_limit, effort_limit - ) - applied_effort = articulation.actuators.applied_effort.torch - torch.testing.assert_close(applied_effort, expected_pd_effort) - assert torch.all(applied_effort[:, cart_id] != feedforward[:, cart_id]) - expected_force = expected_pd_effort.clone() - expected_force[:, cart_id] = feedforward[:, cart_id] - backend_to_user = [joint_names.index(name) for name in articulation.root_view.dof_names] - backend_force = _read_binding_to_torch(articulation, TT.DOF_ACTUATION_FORCE, device) - torch.testing.assert_close(backend_force, expected_force[:, backend_to_user]) - - expected_stiffness = torch.zeros_like(initial_pos) - expected_stiffness[:, cart_id] = stiffness - torch.testing.assert_close( - _read_binding_to_torch(articulation, TT.DOF_STIFFNESS, device), - expected_stiffness[:, backend_to_user], - ) - assert articulation._actuator_control.native_actuator_path_active == use_newton_actuators - - -@pytest.mark.parametrize("device", ["cuda:0"]) -def test_newton_native_actuator_reset_and_gain_event_are_environment_selective(device): - """Reset and randomize only the selected OVPhysX native-controller environment.""" - from isaaclab.envs.mdp.events import randomize_actuator_gains # noqa: PLC0415 - from isaaclab.managers import EventTermCfg # noqa: PLC0415 - - class Env: - def __init__(self, asset): - self.scene = self - self.num_envs = asset.num_instances - self.device = asset.device - self._asset = asset - - def __getitem__(self, name): - assert name == "robot" - return self._asset - - with _ovphysx_sim_context(device=device, use_newton_actuators=True) as sim: - sim._app_control_on_stop_handle = None - articulation_cfg = replace( - generate_articulation_cfg("single_joint_explicit"), - actuators={ - "joint": DelayedPDActuatorCfg( - joint_names_expr=[".*"], - stiffness=20.0, - damping=1.0, - actuator_effort_limit=80.0, - min_delay=1, - max_delay=1, - ) - }, - ) - articulation, _ = generate_articulation(articulation_cfg, 2, device) - sim.reset() - for _ in range(3): - articulation.write_data_to_sim() - sim.step() - articulation.update(sim.cfg.dt) - - adapter = articulation.newton_actuator_adapter - stateful_pairs = [ - state - for actuator, state in zip(adapter.actuators, adapter._states_a) - if state is not None and getattr(state, "delay_state", None) is not None - ] - assert len(stateful_pairs) == 1 - articulation.reset(env_ids=torch.tensor([0], device=device, dtype=torch.long)) - assert stateful_pairs[0].delay_state.num_pushes.numpy().tolist() == [0, 1] - - env = Env(articulation) - asset_cfg = SceneEntityCfg("robot") - event_params = { - "asset_cfg": asset_cfg, - "stiffness_distribution_params": (101.0, 101.0), - "damping_distribution_params": (3.0, 3.0), - "operation": "abs", - "distribution": "uniform", - } - event = randomize_actuator_gains(EventTermCfg(func=randomize_actuator_gains, params=event_params), env) - event(env, env_ids=torch.tensor([0], device=device), **event_params) - - from isaaclab.actuators.newton import read_group_parameter - - stiffness = read_group_parameter(articulation.actuators, "joint", "controller", "kp") - damping = read_group_parameter(articulation.actuators, "joint", "controller", "kd") - torch.testing.assert_close(stiffness, torch.tensor([[101.0], [20.0]], device=device)) - torch.testing.assert_close(damping, torch.tensor([[3.0], [1.0]], device=device)) - - -@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 - dt = request.getfixturevalue("dt") if "dt" in request.fixturenames else 1.0 / 60.0 - with _ovphysx_sim_context( - device=device, auto_add_lighting=True, gravity_enabled=gravity_enabled, add_ground_plane=add_ground_plane, dt=dt - ) as sim: - sim._app_control_on_stop_handle = None - yield sim - + joint_names = articulation.joint_names + shoulder_ids = [joint_names.index(name) for name in ("left_shoulder", "right_shoulder")] + elbow_ids = [joint_names.index(name) for name in ("left_elbow", "right_elbow")] + initial_pos = articulation.data.joint_pos.torch.clone() + position_target = initial_pos.clone() + position_target[:, shoulder_ids] += 0.25 + position_target[:, elbow_ids] += 0.5 + feedforward = torch.zeros_like(initial_pos) + feedforward[:, shoulder_ids] = 1.5 + feedforward[:, elbow_ids] = -0.75 + articulation.actuators.target_command.set_position_index(value=position_target) + articulation.actuators.target_command.set_effort_index(value=feedforward) + articulation.write_data_to_sim() -@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) -@pytest.mark.parametrize("gravity_enabled", [True]) -@pytest.mark.parametrize("user_ordering", [False, True]) -# Keep integration-time pose/velocity differences below the dynamics tolerance. -@pytest.mark.parametrize("dt", [1e-4]) -def test_reversed_joint_dynamics_use_public_joint_basis(sim, device, gravity_enabled, user_ordering, dt): - """Check velocity, kinetic energy and gravity in the public joint basis.""" - articulation = Articulation( - ArticulationCfg( - prim_path="/World/Robot", - spawn=sim_utils.UsdFileCfg( - usd_path=str(Path(__file__).parent / "data" / "articulation_ordering_branching.usda") - ), - actuators={}, - joint_ordering=BRANCHING_MJWARP_JOINT_NAMES if user_ordering else None, - body_ordering=BRANCHING_MJWARP_BODY_NAMES if user_ordering else None, - ) + expected_pd_effort = torch.clamp( + stiffness * (position_target - initial_pos) + feedforward, -effort_limit, effort_limit ) - UsdPhysics.FixedJoint.Define(sim.stage, "/World/Robot/fixed_root").GetBody1Rel().SetTargets(["/World/Robot/base"]) - joint = UsdPhysics.RevoluteJoint.Get(sim.stage, "/World/Robot/left_elbow") - body0, body1 = joint.GetBody0Rel().GetTargets(), joint.GetBody1Rel().GetTargets() - joint.GetBody0Rel().SetTargets(body1) - joint.GetBody1Rel().SetTargets(body0) - for prim in Usd.PrimRange(sim.stage.GetPrimAtPath("/World/Robot")): - if prim.IsA(UsdPhysics.RevoluteJoint): - UsdPhysics.RevoluteJoint(prim).GetAxisAttr().Set("Y") - if prim.HasAPI(UsdPhysics.MassAPI): - mass = UsdPhysics.MassAPI(prim) - mass.CreateCenterOfMassAttr(Gf.Vec3f(0.2, 0.0, 0.0)) - # Isotropic inertia makes the energy check independent of body rotation. - mass.CreateDiagonalInertiaAttr(Gf.Vec3f(0.1)) - sim.reset() + applied_effort = articulation.actuators.applied_effort.torch + torch.testing.assert_close(applied_effort, expected_pd_effort) + assert torch.all(applied_effort[:, shoulder_ids] != feedforward[:, shoulder_ids]) + expected_force = expected_pd_effort.clone() + expected_force[:, shoulder_ids] = feedforward[:, shoulder_ids] + backend_to_user = [joint_names.index(name) for name in articulation.root_view.dof_names] + backend_force = _read_binding_to_torch(articulation, TT.DOF_ACTUATION_FORCE, scene.device) + torch.testing.assert_close(backend_force, expected_force[:, backend_to_user]) + + expected_stiffness = torch.zeros_like(initial_pos) + expected_stiffness[:, shoulder_ids] = stiffness + torch.testing.assert_close( + _read_binding_to_torch(articulation, TT.DOF_STIFFNESS, scene.device), + expected_stiffness[:, backend_to_user], + ) + assert articulation._actuator_control.native_actuator_path_active == scene.native_actuators - velocity = torch.zeros((1, articulation.num_joints), device=device) - velocity[:, articulation.find_joints("left_shoulder")[0][0]] = 0.4 - velocity[:, articulation.find_joints("left_elbow")[0][0]] = 0.7 - articulation.write_joint_velocity_to_sim_index(velocity=velocity) - sim.step() - articulation.update(dt) - joint_velocity = articulation.data.joint_vel.torch - predicted_velocity = torch.einsum("nbij,nj->nbi", articulation.data.body_com_jacobian_w.torch, joint_velocity) - torch.testing.assert_close(predicted_velocity, articulation.data.body_com_vel_w.torch[:, 1:], atol=1e-5, rtol=1e-5) +@pytest.mark.parametrize("scene", _ALL_DEVICES, indirect=True) +def test_com_orientation_write_invalidates_static_inertia_cache_with_body_ordering(scene: _ArticulationScene) -> None: + """Partial COM writes and COM rotations follow non-identity body order. - generalized_energy = 0.5 * torch.einsum( - "ni,nij,nj->n", joint_velocity, articulation.data.mass_matrix.torch, joint_velocity - ) - body_velocity = articulation.data.body_com_vel_w.torch - body_inertia = articulation.data.body_inertia.torch.reshape(1, articulation.num_bodies, 3, 3) - body_energy = 0.5 * ( - (articulation.data.body_mass.torch.unsqueeze(-1) * body_velocity[..., :3].square()).sum((-1, -2)) - + torch.einsum("nbi,nbij,nbj->n", body_velocity[..., 3:], body_inertia, body_velocity[..., 3:]) - ) - torch.testing.assert_close(generalized_energy, body_energy, atol=1e-5, rtol=1e-5) + COM pose is a CPU-resident OVPhysX binding, and the raw ``root_view.set_attribute`` restore + below forbids cross-device staging, so the raw COM tensors stay on the host. + """ + articulation = scene.islands["com_order"] + device = scene.device + assert articulation.body_ordering is not None - gravity = torch.tensor(sim.cfg.gravity, device=device) - body_weight = articulation.data.body_mass.torch[:, 1:, None] * gravity - expected_compensation = -torch.einsum( - "nbij,nbi->nj", articulation.data.body_com_jacobian_w.torch[:, :, :3], body_weight + public_body_id = 1 + backend_body_id = articulation.body_ordering.user_to_backend_indices[public_body_id] + assert backend_body_id != public_body_id + + # A partial ordered COM write lands on the selected backend body and preserves every other one. + backend_before = _read_binding_to_torch(articulation, TT.BODY_COM_POSE, "cpu").clone() + assert torch.unique(backend_before[0], dim=0).shape[0] > 1 + articulation.data._body_com_pose_b.timestamp = -1.0 + backend_staging = articulation.data._body_com_pose_b_backend + if backend_staging is not None: + backend_staging.timestamp = -1.0 + selected_com = backend_before[0, backend_body_id].clone() + selected_com[0] += 0.001 + articulation.set_coms_index( + coms=wp.from_torch(selected_com.reshape(1, 1, 7).to(device), dtype=wp.transformf), + env_ids=wp.array([0], dtype=wp.int32, device=device), + body_ids=wp.array([public_body_id], dtype=wp.int32, device=device), ) - assert torch.all(expected_compensation.abs() > 0.1) + backend_after = _read_binding_to_torch(articulation, TT.BODY_COM_POSE, "cpu").clone() + articulation.root_view.set_attribute(TT.BODY_COM_POSE, wp.from_torch(backend_before)) + noop_after = _read_binding_to_torch(articulation, TT.BODY_COM_POSE, "cpu").clone() + unselected_body_mask = torch.ones(backend_before.shape[1], dtype=torch.bool) + unselected_body_mask[backend_body_id] = False + assert torch.equal(noop_after[..., :3], backend_before[..., :3]) + assert torch.equal(backend_after[0, backend_body_id, :3], selected_com[:3]) + assert torch.equal(backend_after[0, unselected_body_mask, :3], backend_before[0, unselected_body_mask, :3]) + # Bound semantic orientation equality by the native setter's float32 no-op normalization. + native_orientation_atol = torch.max( + torch.abs(noop_after[0, unselected_body_mask, 3:7] - backend_before[0, unselected_body_mask, 3:7]) + ).item() + assert native_orientation_atol <= torch.finfo(backend_before.dtype).eps torch.testing.assert_close( - articulation.data.gravity_compensation_forces.torch, expected_compensation, atol=1e-5, rtol=1e-5 + backend_after[0, unselected_body_mask, 3:7], + backend_before[0, unselected_body_mask, 3:7], + rtol=0.0, + atol=native_orientation_atol, ) + assert torch.equal(backend_after[..., 3:7], noop_after[..., 3:7]) + # A COM rotation refreshes the static inertia cache. + coms = articulation.data.body_com_pose_b.torch[:, public_body_id : public_body_id + 1].clone() + coms[..., 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) + articulation.set_coms_index(coms=wp.from_torch(coms, dtype=wp.transformf), body_ids=[public_body_id]) -def test_joint_dof_sign_resolution_traverses_instance_proxies(): - """Resolve reversed joints inside an instanceable articulation.""" - source_stage = Usd.Stage.CreateInMemory() - UsdGeom.Xform.Define(source_stage, "/Robot") - UsdGeom.Xform.Define(source_stage, "/Robot/base") - UsdGeom.Xform.Define(source_stage, "/Robot/link") - joint = UsdPhysics.RevoluteJoint.Define(source_stage, "/Robot/joint") - joint.GetBody0Rel().SetTargets(["/Robot/link"]) - joint.GetBody1Rel().SetTargets(["/Robot/base"]) - stage = Usd.Stage.CreateInMemory() - instance = UsdGeom.Xform.Define(stage, "/World/Robot").GetPrim() - instance.GetReferences().AddReference(source_stage.GetRootLayer().identifier, "/Robot") - instance.SetInstanceable(True) + principal_inertia = torch.tensor([[[1.0, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 3.0]]], device=device) + articulation.set_inertias_index(inertias=principal_inertia, body_ids=[public_body_id]) + torch.testing.assert_close( + articulation.data.body_inertia.torch[:, public_body_id : public_body_id + 1], principal_inertia + ) - articulation = Mock( - cfg=Mock(prim_path="/World/Robot"), - _joint_names=["joint"], - _body_names=["base", "link"], + coms[..., 3:7] = torch.tensor([0.0, 0.0, 0.70710677, 0.70710677], device=device) + articulation.set_coms_index(coms=wp.from_torch(coms, dtype=wp.transformf), body_ids=[public_body_id]) + scene.step(articulation) + expected_rotated_inertia = torch.tensor([[[2.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 3.0]]], device=device) + torch.testing.assert_close( + articulation.data.body_inertia.torch[:, public_body_id : public_body_id + 1], + expected_rotated_inertia, ) - assert Articulation._resolve_joint_dof_signs(articulation, stage) == (-1,) +@pytest.mark.parametrize("scene", _CUDA_SCENES, indirect=True) +def test_loading_gains_from_usd(scene: _ArticulationScene) -> None: + """Adopt the authored drive gains when the actuator config leaves stiffness and damping unset.""" + articulation = scene.islands["usd_gains"] + actuator = articulation.actuators["joints"] + # Angular drive gains are authored per degree and exposed per radian. + expected_stiffness = torch.full_like(actuator.stiffness, math.degrees(_STIFFNESS)) + expected_damping = torch.full_like(actuator.damping, math.degrees(_DAMPING)) + torch.testing.assert_close(actuator.stiffness, expected_stiffness) + torch.testing.assert_close(actuator.damping, expected_damping) + torch.testing.assert_close(_read_binding_to_torch(articulation, TT.DOF_STIFFNESS, scene.device), expected_stiffness) + torch.testing.assert_close(_read_binding_to_torch(articulation, TT.DOF_DAMPING, scene.device), expected_damping) + + +@pytest.mark.parametrize("scene", _CUDA_SCENES, indirect=True) +def test_explicit_articulation_root_on_a_floating_base_body(scene: _ArticulationScene) -> None: + """Initialize and simulate a floating articulation from an explicit root path to a rigid body below the asset.""" + articulation = scene.islands["non_root"] + num_articulations = articulation.num_instances + assert articulation.is_initialized + assert not articulation.is_fixed_base + assert articulation.root_view.prim_paths == [ + f"/World/non_root/Env_{index}/Robot/base" for index in range(num_articulations) + ] + 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, articulation.num_joints) + for tt in (TT.DOF_POSITION, TT.DOF_VELOCITY, TT.DOF_STIFFNESS): + assert articulation.root_view.binding_for(tt).shape[1] == articulation.num_joints + for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): + assert articulation.root_view.binding_for(tt).shape[1] == articulation.num_bodies + for actuator in articulation.actuators.values(): + assert actuator.is_implicit_model + assert actuator.joint_indices == slice(None) + scene.step(articulation, count=10) + assert torch.isfinite(articulation.data.root_link_pose_w.torch).all() -@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) -def test_live_anymal_c_manual_joint_ordering_reorders_joint_targets(sim, device): - """Write nonidentity-ordered joint targets into their intended backend columns.""" - stiffness, damping = 10.0, 2.0 - backend_joint_names = ANYMAL_C_PHYSX_JOINT_NAMES - joint_ordering = (*backend_joint_names[1:], backend_joint_names[0]) - articulation_cfg = replace( - generate_articulation_cfg("anymal"), - joint_ordering=joint_ordering, - actuators={"legs": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=stiffness, damping=damping)}, - ) - articulation, _ = generate_articulation(articulation_cfg, 1, device=device) - sim.reset() +@pytest.mark.parametrize("scene", _CUDA_SCENES, indirect=True) +def test_implicit_drive_targets_submit_feedforward_effort_and_track(scene: _ArticulationScene) -> None: + """Write nonidentity-ordered implicit targets to their backend columns and track them in the solver.""" + articulation = scene.islands["drive"] + device = scene.device ordering = articulation.joint_ordering assert ordering is not None user_to_backend = torch.as_tensor(ordering.user_to_backend_indices, dtype=torch.long, device=device) @@ -860,9 +1213,11 @@ def test_live_anymal_c_manual_joint_ordering_reorders_joint_targets(sim, device) assert not torch.equal(user_to_backend, backend_to_user) joint_index = torch.arange(articulation.num_joints, dtype=torch.float32, device=device).unsqueeze(0) - position_target = -0.25 + 0.031 * joint_index - velocity_target = 0.07 + 0.017 * joint_index - effort_target = torch.where(joint_index.remainder(2) == 0, 0.0, 0.13 * joint_index) + position_target = (-0.25 + 0.031 * joint_index).repeat(articulation.num_instances, 1) + velocity_target = (0.07 + 0.017 * joint_index).repeat(articulation.num_instances, 1) + effort_target = torch.where(joint_index.remainder(2) == 0, 0.0, 0.13 * joint_index).repeat( + articulation.num_instances, 1 + ) joint_pos = articulation.data.joint_pos.torch.clone() joint_vel = articulation.data.joint_vel.torch.clone() articulation.set_joint_position_target_index(target=position_target) @@ -878,224 +1233,141 @@ def test_live_anymal_c_manual_joint_ordering_reorders_joint_targets(sim, device) _read_binding_to_torch(articulation, TT.DOF_ACTUATION_FORCE, device), effort_target[:, backend_to_user] ) computed_effort = ( - stiffness * (position_target - joint_pos) + damping * (velocity_target - joint_vel) + effort_target + _STIFFNESS * (position_target - joint_pos) + _DAMPING * (velocity_target - joint_vel) + effort_target ) limits = articulation.data.joint_effort_limits.torch expected_telemetry = torch.clamp(computed_effort, min=-limits, max=limits) torch.testing.assert_close(articulation.actuators.applied_effort.torch, expected_telemetry) assert not torch.allclose(expected_telemetry, effort_target) - -@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) -def test_reversed_joint_ordering_joint_state_index_writes_backend_order(sim, device): - """Read dynamics and write full, partial, and position-only joint state through a nonidentity joint axis.""" - articulation_cfg = replace( - generate_articulation_cfg("anymal"), joint_ordering=tuple(reversed(ANYMAL_C_PHYSX_JOINT_NAMES)) - ) - articulation, _ = generate_articulation(articulation_cfg, 2, device=device) - sim.reset() - - ordering = articulation.joint_ordering - assert ordering is not None - num_joints = articulation.num_joints - user_to_backend = torch.as_tensor(ordering.user_to_backend_indices, dtype=torch.long, device=device) - backend_to_user = torch.as_tensor(ordering.backend_to_user_indices, dtype=torch.long, device=device) - - # Floating-base columns stay leading while the actuated-joint axes are gathered. - generalized_user_to_backend = torch.cat((torch.arange(6, device=device), 6 + user_to_backend)) - raw_jacobian = _read_binding_to_torch(articulation, TT.JACOBIAN, device).reshape( - 2, articulation.num_bodies, 6, num_joints + 6 - ) - raw_mass_matrix = _read_binding_to_torch(articulation, TT.MASS_MATRIX, device) - raw_gravity = _read_binding_to_torch(articulation, TT.GRAVITY_FORCE, device) - torch.testing.assert_close( - articulation.data.body_com_jacobian_w.torch, - raw_jacobian[:, :, :, generalized_user_to_backend], - ) - torch.testing.assert_close( - articulation.data.mass_matrix.torch, - raw_mass_matrix[:, generalized_user_to_backend, :][:, :, generalized_user_to_backend], - ) - torch.testing.assert_close( - articulation.data.gravity_compensation_forces.torch, - raw_gravity[:, generalized_user_to_backend], + # A partial position target lands on the selected backend cell; the other environment keeps its targets. + articulation.set_joint_velocity_target_index(target=torch.zeros_like(velocity_target)) + articulation.actuators.target_command.set_effort_index(value=torch.zeros_like(effort_target)) + raw_target_before = _read_binding_to_torch(articulation, TT.DOF_POSITION_TARGET, device).clone() + drive_target = position_target[1:, :1] + 0.4 + articulation.actuators.target_command.set_position_index( + value=drive_target, + env_ids=torch.tensor([1], dtype=torch.int32, device=device), + joint_ids=torch.tensor([0], dtype=torch.int32, device=device), ) + articulation.write_data_to_sim() + backend_target = _read_binding_to_torch(articulation, TT.DOF_POSITION_TARGET, device) + torch.testing.assert_close(backend_target[1, user_to_backend[0]], drive_target[0, 0]) + torch.testing.assert_close(backend_target[0], raw_target_before[0]) - backend_pos_before = torch.arange(2 * num_joints, dtype=torch.float32, device=device).reshape(2, num_joints) - backend_vel_before = backend_pos_before + 100.0 - articulation.root_view.set_attribute(TT.DOF_POSITION, wp.from_torch(backend_pos_before.contiguous())) - articulation.root_view.set_attribute(TT.DOF_VELOCITY, wp.from_torch(backend_vel_before.contiguous())) - for buffer in ( - articulation.data._joint_pos_buf, - articulation.data._joint_vel_buf, - articulation.data._joint_pos_backend, - articulation.data._joint_vel_backend, - ): - if buffer is not None: - buffer.timestamp = -1.0 - - torch.testing.assert_close(articulation.data.joint_pos.torch, backend_pos_before[:, user_to_backend]) - torch.testing.assert_close(articulation.data.joint_vel.torch, backend_vel_before[:, user_to_backend]) - - def write_and_check(position, velocity, env_ids, joint_ids, expected_public_pos, expected_public_vel): - articulation.write_joint_state_to_sim_index( - position=position, - velocity=velocity, - env_ids=env_ids, - joint_ids=joint_ids, - ) - torch.testing.assert_close(articulation.data.joint_pos.torch, expected_public_pos) - torch.testing.assert_close(articulation.data.joint_vel.torch, expected_public_vel) - torch.testing.assert_close( - _read_binding_to_torch(articulation, TT.DOF_POSITION, device), - expected_public_pos[:, backend_to_user], - ) - torch.testing.assert_close( - _read_binding_to_torch(articulation, TT.DOF_VELOCITY, device), - expected_public_vel[:, backend_to_user], - ) - - # Full write. - position = torch.arange(2 * num_joints, dtype=torch.float32, device=device).reshape(2, num_joints) + 200.0 - velocity = position + 100.0 - write_and_check(position, velocity, None, None, position, velocity) - - # Partial write keeps every unselected public and backend entry. - joint_ids = [0, 2] - expected_public_pos = position.clone() - expected_public_vel = velocity.clone() - partial_position = torch.tensor([[401.0, 403.0]], device=device) - partial_velocity = torch.tensor([[501.0, 503.0]], device=device) - expected_public_pos[1, joint_ids] = partial_position[0] - expected_public_vel[1, joint_ids] = partial_velocity[0] - write_and_check(partial_position, partial_velocity, [1], joint_ids, expected_public_pos, expected_public_vel) - - # A position-only partial write (a distinct kernel) preserves every unselected backend joint. - backend_before = wp.to_torch(articulation.root_view.get_attribute(TT.DOF_POSITION)).clone() - backend_joint_id = ordering.user_to_backend_indices[0] - selected_value = backend_before[0, backend_joint_id] + 0.001 - articulation.write_joint_position_to_sim_index( - position=selected_value.reshape(1, 1), - env_ids=wp.array([0], dtype=wp.int32, device=device), - joint_ids=wp.array([0], dtype=wp.int32, device=device), - ) - expected_backend = backend_before.clone() - expected_backend[0, backend_joint_id] = selected_value - torch.testing.assert_close( - wp.to_torch(articulation.root_view.get_attribute(TT.DOF_POSITION)), expected_backend, rtol=0.0, atol=0.0 - ) + # The drives track the commanded targets. Without the target write, they would hold the old ones. + expected_position = position_target.clone() + expected_position[1, 0] = drive_target[0, 0] + scene.step(articulation, count=100) + torch.testing.assert_close(articulation.data.joint_pos.torch, expected_position, atol=0.1, rtol=0.0) -@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) -def test_branching_fixture_physx_ordering_is_identity_on_ovphysx(sim, device): - """Take the same-backend identity fast path for ``joint_ordering="physx"`` on OVPhysX. +@pytest.mark.parametrize("scene", _CUDA_SCENES, indirect=True) +def test_computed_dynamics_follow_public_order_and_model_writes(scene: _ArticulationScene) -> None: + """Resolve depth-first MJWarp order for ``"mjwarp"`` on OVPhysX and gather computed dynamics into it. - Live coverage for the same-backend symbolic-convention path: OVPhysX is a - PhysX-family backend whose articulation view is already in PhysX (breadth-first) order, so - requesting ``physx`` must expose the public joint/body axes verbatim in backend order with no - reorder map. + OVPhysX is a PhysX-family backend (native breadth-first order), so requesting ``mjwarp`` triggers a + temporary Newton USD discovery of the depth-first order and reorders the public joint/body axes to it. + Joint-state and model-property writes refresh the gathered dynamics without advancing simulation time. """ - articulation = Articulation( - ArticulationCfg( - prim_path="/World/Robot", - spawn=sim_utils.UsdFileCfg( - usd_path=str(Path(__file__).parent / "data" / "articulation_ordering_branching.usda") - ), - actuators={}, - joint_ordering="physx", - body_ordering="physx", - ) - ) - sim.reset() - assert articulation.is_initialized + articulation = scene.islands["dynamics"] + device = scene.device + num_envs = articulation.num_instances # OVPhysX exposes the native breadth-first PhysX order on the backend axis. assert tuple(articulation.backend_joint_names) == BRANCHING_PHYSX_JOINT_NAMES assert tuple(articulation.backend_body_names) == BRANCHING_PHYSX_BODY_NAMES - # Same-backend preset: the public axis equals the backend axis and no reorder map is created. - assert tuple(articulation.joint_names) == tuple(articulation.backend_joint_names) - assert tuple(articulation.body_names) == tuple(articulation.backend_body_names) - assert tuple(articulation.joint_names) == BRANCHING_PHYSX_JOINT_NAMES - assert tuple(articulation.body_names) == BRANCHING_PHYSX_BODY_NAMES - assert articulation.joint_ordering is None - assert articulation.body_ordering is None - - -@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) -def test_articulation_dynamics_fixed_base_match_raw_ovphysx_bindings(sim, device): - """Expose fixed-base computed dynamics through the backend-agnostic data API.""" - articulation, _ = generate_articulation(generate_articulation_cfg("panda"), 1, device=device) - sim.reset() + # Cross-backend Newton discovery resolves the depth-first MJWarp order and reorders the public axis. + assert get_articulation_name_ordering(articulation, "mjwarp", kind="joint") == BRANCHING_MJWARP_JOINT_NAMES + assert get_articulation_name_ordering(articulation, "mjwarp", kind="body") == BRANCHING_MJWARP_BODY_NAMES + assert tuple(articulation.joint_names) == BRANCHING_MJWARP_JOINT_NAMES + assert tuple(articulation.body_names) == BRANCHING_MJWARP_BODY_NAMES - num_generalized_dofs = articulation.num_joints - num_jacobian_bodies = articulation.num_bodies - 1 - raw_jacobian = _read_binding_to_torch(articulation, TT.JACOBIAN, device).reshape( - 1, num_jacobian_bodies, 6, num_generalized_dofs + joint_ordering = articulation.joint_ordering + body_ordering = articulation.body_ordering + assert joint_ordering is not None + assert body_ordering is not None + joint_user_to_backend = torch.as_tensor(joint_ordering.user_to_backend_indices, dtype=torch.long, device=device) + body_offset = 1 if articulation.is_fixed_base else 0 + body_user_to_backend = torch.as_tensor( + [ + backend_body_id - body_offset + for backend_body_id in body_ordering.user_to_backend_indices + if not body_offset or backend_body_id != 0 + ], + dtype=torch.long, + device=device, + ) + generalized_user_to_backend = torch.cat( + ( + torch.arange(articulation.num_base_dofs, device=device), + articulation.num_base_dofs + joint_user_to_backend, + ) ) - raw_mass_matrix = _read_binding_to_torch(articulation, TT.MASS_MATRIX, device) - raw_gravity = _read_binding_to_torch(articulation, TT.GRAVITY_FORCE, device) - body_com_jacobian = articulation.data.body_com_jacobian_w - mass_matrix = articulation.data.mass_matrix - gravity = articulation.data.gravity_compensation_forces - assert body_com_jacobian is articulation.data.body_com_jacobian_w - assert mass_matrix is articulation.data.mass_matrix - assert gravity is articulation.data.gravity_compensation_forces - - torch.testing.assert_close(body_com_jacobian.torch, raw_jacobian) - torch.testing.assert_close(mass_matrix.torch, raw_mass_matrix) - torch.testing.assert_close(gravity.torch, raw_gravity) - assert articulation.data.body_link_jacobian_w.torch.shape == raw_jacobian.shape - assert torch.isfinite(articulation.data.body_link_jacobian_w.torch).all() - torch.testing.assert_close(mass_matrix.torch, mass_matrix.torch.transpose(-1, -2), rtol=1e-5, atol=1e-5) + def expected_dynamics() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + raw_jacobian = _read_binding_to_torch(articulation, TT.JACOBIAN, device).reshape( + num_envs, + articulation.num_bodies - body_offset, + 6, + articulation.num_joints + articulation.num_base_dofs, + ) + raw_mass_matrix = _read_binding_to_torch(articulation, TT.MASS_MATRIX, device) + raw_gravity = _read_binding_to_torch(articulation, TT.GRAVITY_FORCE, device) + return ( + raw_jacobian[:, body_user_to_backend, :, :][:, :, :, generalized_user_to_backend], + raw_mass_matrix[:, generalized_user_to_backend, :][:, :, generalized_user_to_backend], + raw_gravity[:, generalized_user_to_backend], + ) + + def assert_dynamics_match(jacobian: bool = True, gravity: bool = True) -> None: + expected_jacobian, expected_mass_matrix, expected_gravity = expected_dynamics() + if jacobian: + torch.testing.assert_close(articulation.data.body_com_jacobian_w.torch, expected_jacobian) + torch.testing.assert_close(articulation.data.mass_matrix.torch, expected_mass_matrix) + if gravity: + torch.testing.assert_close(articulation.data.gravity_compensation_forces.torch, expected_gravity) - joint_position = articulation.data.joint_pos.torch.clone() + data = articulation.data + body_com_jacobian = data.body_com_jacobian_w + mass_matrix = data.mass_matrix + gravity = data.gravity_compensation_forces + assert body_com_jacobian is data.body_com_jacobian_w + assert mass_matrix is data.mass_matrix + assert gravity is data.gravity_compensation_forces + assert_dynamics_match() + assert data.body_link_jacobian_w.torch.shape == body_com_jacobian.torch.shape + assert torch.isfinite(data.body_link_jacobian_w.torch).all() + torch.testing.assert_close(mass_matrix.torch, mass_matrix.torch.transpose(-1, -2), rtol=1e-5, atol=1e-5) + + raw_jacobian_before = expected_dynamics()[0] + joint_position = data.joint_pos.torch.clone() joint_position[:, 1] += 0.2 articulation.write_joint_position_to_sim_index(position=joint_position) - updated_raw_jacobian = _read_binding_to_torch(articulation, TT.JACOBIAN, device).reshape( - 1, num_jacobian_bodies, 6, num_generalized_dofs - ) - updated_raw_mass_matrix = _read_binding_to_torch(articulation, TT.MASS_MATRIX, device) - updated_raw_gravity = _read_binding_to_torch(articulation, TT.GRAVITY_FORCE, device) - assert not torch.allclose(updated_raw_jacobian, raw_jacobian) - assert not torch.allclose(updated_raw_mass_matrix, raw_mass_matrix) - torch.testing.assert_close(articulation.data.body_com_jacobian_w.torch, updated_raw_jacobian) - torch.testing.assert_close(articulation.data.mass_matrix.torch, updated_raw_mass_matrix) - torch.testing.assert_close(articulation.data.gravity_compensation_forces.torch, updated_raw_gravity) - - joint_velocity = torch.linspace(-0.2, 0.2, articulation.num_joints, dtype=torch.float32, device=device).unsqueeze(0) + assert not torch.allclose(expected_dynamics()[0], raw_jacobian_before) + assert_dynamics_match() + + joint_velocity = torch.linspace(-0.2, 0.2, articulation.num_joints, dtype=torch.float32, device=device) + joint_velocity = joint_velocity.repeat(num_envs, 1) articulation.write_joint_velocity_to_sim_index(velocity=joint_velocity) - expected_com_velocity = torch.einsum("nbij,nj->nbi", articulation.data.body_com_jacobian_w.torch, joint_velocity) - expected_link_velocity = torch.einsum("nbij,nj->nbi", articulation.data.body_link_jacobian_w.torch, joint_velocity) - torch.testing.assert_close( - expected_com_velocity, articulation.data.body_com_vel_w.torch[:, 1:], atol=1e-5, rtol=1e-4 - ) - torch.testing.assert_close( - expected_link_velocity, articulation.data.body_link_vel_w.torch[:, 1:], atol=1e-5, rtol=1e-4 - ) + expected_com_velocity = torch.einsum("nbij,nj->nbi", data.body_com_jacobian_w.torch, joint_velocity) + expected_link_velocity = torch.einsum("nbij,nj->nbi", data.body_link_jacobian_w.torch, joint_velocity) + torch.testing.assert_close(expected_com_velocity, data.body_com_vel_w.torch[:, 1:], atol=1e-5, rtol=1e-4) + torch.testing.assert_close(expected_link_velocity, data.body_link_vel_w.torch[:, 1:], atol=1e-5, rtol=1e-4) # Model-property writes refresh the computed dynamics without advancing simulation time. - data = articulation.data # COM offsets affect COM Jacobians, mass matrices, and gravity forces. data.body_com_jacobian_w data.mass_matrix data.gravity_compensation_forces - coms = _read_binding_to_torch(articulation, TT.BODY_COM_POSE, device) + coms = data.body_com_pose_b.torch.clone() coms[:, -1, 0] += 0.01 - articulation.set_coms_index(coms=wp.from_torch(coms.contiguous(), dtype=wp.transformf)) + articulation.set_coms_index(coms=wp.from_torch(coms, dtype=wp.transformf)) assert data._body_com_jacobian_w.timestamp < data._sim_timestamp assert data._mass_matrix.timestamp < data._sim_timestamp assert data._gravity_compensation_forces.timestamp < data._sim_timestamp - raw_jacobian = _read_binding_to_torch(articulation, TT.JACOBIAN, device).reshape( - 1, num_jacobian_bodies, 6, num_generalized_dofs - ) - torch.testing.assert_close(data.body_com_jacobian_w.torch, raw_jacobian) - torch.testing.assert_close(data.mass_matrix.torch, _read_binding_to_torch(articulation, TT.MASS_MATRIX, device)) - torch.testing.assert_close( - data.gravity_compensation_forces.torch, - _read_binding_to_torch(articulation, TT.GRAVITY_FORCE, device), - ) + assert_dynamics_match() # Mass affects the mass matrix and gravity forces, but not the kinematic Jacobian. masses = data.body_mass.torch.clone() @@ -1103,1572 +1375,333 @@ def test_articulation_dynamics_fixed_base_match_raw_ovphysx_bindings(sim, device articulation.set_masses_index(masses=masses) assert data._mass_matrix.timestamp < data._sim_timestamp assert data._gravity_compensation_forces.timestamp < data._sim_timestamp - torch.testing.assert_close(data.mass_matrix.torch, _read_binding_to_torch(articulation, TT.MASS_MATRIX, device)) - torch.testing.assert_close( - data.gravity_compensation_forces.torch, - _read_binding_to_torch(articulation, TT.GRAVITY_FORCE, device), - ) + assert_dynamics_match(jacobian=False) # Inertia and armature each affect only the generalized mass matrix. inertias = data.body_inertia.torch.clone() inertias[:, -1, [0, 4, 8]] *= 1.1 articulation.set_inertias_index(inertias=inertias) assert data._mass_matrix.timestamp < data._sim_timestamp - torch.testing.assert_close(data.mass_matrix.torch, _read_binding_to_torch(articulation, TT.MASS_MATRIX, device)) + assert_dynamics_match(jacobian=False, gravity=False) armature = data.joint_armature.torch.clone() armature[:, -1] += 0.01 articulation.write_joint_armature_to_sim_index(armature=armature) assert data._mass_matrix.timestamp < data._sim_timestamp - torch.testing.assert_close(data.mass_matrix.torch, _read_binding_to_torch(articulation, TT.MASS_MATRIX, device)) - - -@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) -def test_articulation_dynamics_reorder_body_rows_and_joint_axes(sim, device): - """Resolve depth-first MJWarp order for ``"mjwarp"`` on OVPhysX and gather computed dynamics into it. - - OVPhysX is a PhysX-family backend (native breadth-first order), so requesting ``mjwarp`` triggers a - temporary Newton USD discovery of the depth-first order and reorders the public joint/body axes to it. - The MJWarp/DFS ground truth is the same tuple isaaclab_newton's - ``test_mjwarp_ordering_resolver_matches_newton_backend_names`` pins for its live Newton backend. - """ - articulation = Articulation( - ArticulationCfg( - prim_path="/World/Robot", - spawn=sim_utils.UsdFileCfg( - usd_path=str(Path(__file__).parent / "data" / "articulation_ordering_branching.usda") - ), - actuators={}, - joint_ordering="mjwarp", - body_ordering="mjwarp", - ) - ) - sim.reset() - assert articulation.is_initialized - - # OVPhysX exposes the native breadth-first PhysX order on the backend axis. - assert tuple(articulation.backend_joint_names) == BRANCHING_PHYSX_JOINT_NAMES - assert tuple(articulation.backend_body_names) == BRANCHING_PHYSX_BODY_NAMES - - # Cross-backend Newton discovery resolves the depth-first MJWarp order and reorders the public axis. - assert get_articulation_name_ordering(articulation, "mjwarp", kind="joint") == BRANCHING_MJWARP_JOINT_NAMES - assert get_articulation_name_ordering(articulation, "mjwarp", kind="body") == BRANCHING_MJWARP_BODY_NAMES - assert tuple(articulation.joint_names) == BRANCHING_MJWARP_JOINT_NAMES - assert tuple(articulation.body_names) == BRANCHING_MJWARP_BODY_NAMES - - joint_ordering = articulation.joint_ordering - body_ordering = articulation.body_ordering - assert joint_ordering is not None - assert body_ordering is not None - joint_user_to_backend = torch.as_tensor(joint_ordering.user_to_backend_indices, dtype=torch.long, device=device) - body_offset = 1 if articulation.is_fixed_base else 0 - body_user_to_backend = torch.as_tensor( - [ - backend_body_id - body_offset - for backend_body_id in body_ordering.user_to_backend_indices - if not body_offset or backend_body_id != 0 - ], - dtype=torch.long, - device=device, - ) - generalized_user_to_backend = torch.cat( - ( - torch.arange(articulation.num_base_dofs, device=device), - articulation.num_base_dofs + joint_user_to_backend, - ) - ) - - raw_jacobian = _read_binding_to_torch(articulation, TT.JACOBIAN, device).reshape( - 1, - articulation.num_bodies - body_offset, - 6, - articulation.num_joints + articulation.num_base_dofs, - ) - raw_mass_matrix = _read_binding_to_torch(articulation, TT.MASS_MATRIX, device) - raw_gravity = _read_binding_to_torch(articulation, TT.GRAVITY_FORCE, device) - - expected_jacobian = raw_jacobian[:, body_user_to_backend, :, :][:, :, :, generalized_user_to_backend] - expected_mass_matrix = raw_mass_matrix[:, generalized_user_to_backend, :][:, :, generalized_user_to_backend] - expected_gravity = raw_gravity[:, generalized_user_to_backend] - torch.testing.assert_close(articulation.data.body_com_jacobian_w.torch, expected_jacobian) - torch.testing.assert_close(articulation.data.mass_matrix.torch, expected_mass_matrix) - torch.testing.assert_close(articulation.data.gravity_compensation_forces.torch, expected_gravity) - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("add_ground_plane", [True]) -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 an ``OvPhysxView`` over the per-tensor-type bindings, 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): - binding = articulation.root_view.try_binding_for(tt) - if binding is not None: - assert binding.shape[1] == articulation.num_joints - for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): - binding = articulation.root_view.try_binding_for(tt) - if binding is not None: - assert binding.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 - assert actuator.joint_indices == slice(None) - - # Simulate physics - for _ in range(10): - # perform rendering - sim.step() - # update articulation - articulation.update(sim.cfg.dt) - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("add_ground_plane", [True]) -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 an ``OvPhysxView`` over the per-tensor-type bindings, 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): - binding = articulation.root_view.try_binding_for(tt) - if binding is not None: - assert binding.shape[1] == articulation.num_joints - for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): - binding = articulation.root_view.try_binding_for(tt) - if binding is not None: - assert binding.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", [2]) -@pytest.mark.parametrize("device", test_devices()) -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 an ``OvPhysxView`` over the per-tensor-type bindings, 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): - binding = articulation.root_view.try_binding_for(tt) - if binding is not None: - assert binding.shape[1] == articulation.num_joints - for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): - binding = articulation.root_view.try_binding_for(tt) - if binding is not None: - assert binding.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 - assert isinstance(actuator.joint_indices, torch.Tensor) - assert actuator.joint_indices.dtype == torch.int32 - assert actuator.joint_indices.device == torch.device(device) - - # 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", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("add_ground_plane", [True]) -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 an ``OvPhysxView`` over the per-tensor-type bindings, 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): - binding = articulation.root_view.try_binding_for(tt) - if binding is not None: - assert binding.shape[1] == articulation.num_joints - for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): - binding = articulation.root_view.try_binding_for(tt) - if binding is not None: - assert binding.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", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_fixed_tendon_position_target_writes_offset(sim, num_articulations, device): - """Initialize the fixed-base tendon hand and land tendon length targets on the selected cells only. - - A target lands in the simulation as ``rest_length - target``. The index form commands every tendon of - environment 0; the mask form commands tendon 0 of environment 1. Every other cell must keep its initial offset. - """ - 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 - - sim.reset() - assert articulation.is_initialized - assert articulation.is_fixed_base - 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) - for tt in (TT.DOF_POSITION, TT.DOF_VELOCITY, TT.DOF_STIFFNESS): - binding = articulation.root_view.try_binding_for(tt) - if binding is not None: - assert binding.shape[1] == articulation.num_joints - for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): - binding = articulation.root_view.try_binding_for(tt) - if binding is not None: - assert binding.shape[1] == articulation.num_bodies - 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 - - num_tendons = articulation.num_fixed_tendons - assert num_tendons > 0 - rest_length = articulation.data.fixed_tendon_rest_length.torch.clone() - initial_offset = articulation.data.fixed_tendon_offset.torch.clone() - - index_target = torch.full((1, num_tendons), 0.3, dtype=torch.float32, device=device) - articulation.set_fixed_tendon_position_target_index(target=index_target, env_ids=[0]) - # Distinct per-cell values: a uniform target cannot catch the mask form reading the wrong - # cell, because every wrong read returns the same number. - mask_target = ( - 0.7 - + 0.1 * torch.arange(num_articulations, dtype=torch.float32, device=device).unsqueeze(1) - + 0.01 * torch.arange(num_tendons, dtype=torch.float32, device=device).unsqueeze(0) - ) - env_mask = wp.array([False, True], dtype=wp.bool, device=device) - tendon_mask = wp.array([i == 0 for i in range(num_tendons)], dtype=wp.bool, device=device) - articulation.set_fixed_tendon_position_target_mask( - target=mask_target, fixed_tendon_mask=tendon_mask, env_mask=env_mask - ) - - articulation.write_data_to_sim() - sim.step() - articulation.update(sim.cfg.dt) - - expected = initial_offset.clone() - expected[0] = rest_length[0] - 0.3 - expected[1, 0] = rest_length[1, 0] - mask_target[1, 0] - torch.testing.assert_close(articulation.data.fixed_tendon_offset.torch, expected) - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("add_ground_plane", [True]) -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 = clone(generate_articulation_cfg(articulation_type="anymal")) - # Fix root link by making it kinematic - articulation_cfg.spawn.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 an ``OvPhysxView`` over the per-tensor-type bindings, 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): - binding = articulation.root_view.try_binding_for(tt) - if binding is not None: - assert binding.shape[1] == articulation.num_joints - for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): - binding = articulation.root_view.try_binding_for(tt) - if binding is not None: - assert binding.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("device", ["cpu"]) -def test_fragment_fix_root_reenables_existing_joint(sim, device): - """The fragment path must normalize OVPhysX topology even when a disabled fixed joint exists.""" - articulation_cfg = clone(generate_articulation_cfg(articulation_type="anymal")) - articulation_cfg.spawn.articulation_props = [] - articulation_cfg.spawn.fix_root_link = None - articulation, _ = generate_articulation(articulation_cfg, num_articulations=1, device=device) - - stage = sim.stage - asset_path = "/World/Env_0/Robot" - root = sim_utils.get_first_matching_child_prim( - asset_path, lambda prim: prim.HasAPI(UsdPhysics.ArticulationRootAPI), stage=stage - ) - assert root is not None and root.HasAPI(UsdPhysics.RigidBodyAPI) - old_root_path = root.GetPath().pathString - final_root = root.GetParent() - - joint = UsdPhysics.FixedJoint.Define(stage, f"{old_root_path}/PreAuthoredFixedJoint") - joint.CreateBody1Rel().SetTargets([root.GetPath()]) - joint.CreateJointEnabledAttr().Set(False) - - assert sim_utils.apply_articulation_root_properties(f"{asset_path}(/.*)?", [], stage, fix_root_link=True) - assert joint.GetJointEnabledAttr().Get() is True - world_joints = [] - for prim in sim_utils.get_all_matching_child_prims( - asset_path, lambda prim: prim.IsA(UsdPhysics.FixedJoint), stage=stage - ): - usd_joint = UsdPhysics.Joint(prim) - has_body_0 = bool(usd_joint.GetBody0Rel().GetTargets()) - has_body_1 = bool(usd_joint.GetBody1Rel().GetTargets()) - if has_body_0 != has_body_1: - world_joints.append(prim) - assert world_joints == [joint.GetPrim()] - assert final_root.HasAPI(UsdPhysics.ArticulationRootAPI) - assert not root.HasAPI(UsdPhysics.ArticulationRootAPI) - - sim.reset() - assert articulation.is_initialized - assert articulation.is_fixed_base - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("add_ground_plane", [True]) -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 - """ - # Copy so the shared panda cfg stays fixed-base for later tests. - articulation_cfg = clone(generate_articulation_cfg(articulation_type="panda")) - # Unfix root link by making it non-kinematic - articulation_cfg.spawn.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 an ``OvPhysxView`` over the per-tensor-type bindings, 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): - binding = articulation.root_view.try_binding_for(tt) - if binding is not None: - assert binding.shape[1] == articulation.num_joints - for tt in (TT.BODY_MASS, TT.BODY_COM_POSE): - binding = articulation.root_view.try_binding_for(tt) - if binding is not None: - assert binding.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", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("add_ground_plane", [True]) -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 = clone(generate_articulation_cfg(articulation_type="panda")) - 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", test_devices()) -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 = replace(FRANKA_PANDA_CFG, 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", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("add_ground_plane", [True]) -def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane, caplog): - """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) - - # Logging level changes must not change clamping or reuse an old violation count. - limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) - limits[..., 1] = 0.5 - articulation_logger = type(articulation).__module__ - for level in (logging.WARNING, logging.INFO): - articulation.data.default_joint_pos.torch.fill_(1.0) - caplog.clear() - with caplog.at_level(level, logger=articulation_logger): - articulation.write_joint_position_limit_to_sim_index(limits=limits, warn_limit_violation=False) - assert bool(caplog.records) == (level == logging.INFO) - torch.testing.assert_close(articulation.data.default_joint_pos.torch, torch.full_like(limits[..., 1], 0.5)) - caplog.clear() - with caplog.at_level(logging.INFO, logger=articulation_logger): - articulation.write_joint_position_limit_to_sim_mask(limits=limits, warn_limit_violation=False) - assert not caplog.records - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -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", [2]) -@pytest.mark.parametrize("device", test_devices()) -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 are calculated and composed correctly - 3. The forces affect the articulation's motion correctly - 4. The articulation responds to the forces as expected - - The global-frame position path is covered by :func:`test_external_force_on_multiple_bodies_at_position`; - a fixed world-frame force here would lift the robot instead of tipping it over. - - 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", [2]) -@pytest.mark.parametrize("device", test_devices()) -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", [2]) -@pytest.mark.parametrize("device", test_devices()) -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, translations = 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): - # Preserve environment separation when converting the default root pose to world coordinates. - root_pose = articulation.data.default_root_pose.torch.clone() - root_pose[:, :3] += translations - 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() - - 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): - # the response axis depends on the link frames, so check that the articulation rotates - assert torch.linalg.vector_norm(articulation.data.root_ang_vel_w.torch[i]).item() > 0.1 - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -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("device", test_devices()) -@pytest.mark.parametrize("joint_limit", [1e5, None]) -def test_setting_velocity_and_effort_limits_write_to_solver(sim, device, joint_limit): - """Test that the resolved joint velocity and effort limits reach the PhysX tensor-API solver. - - The full limit-resolution matrix (config override vs. USD default, implicit and explicit - actuators, actuator-limit soft fallback) is covered on the Newton backend and at unit - level. This smoke test only verifies the PhysX tensor-API write path: the configured limits (or the - USD-authored defaults when unset) land in the native solver buffers and match - ``data.joint_vel_limits`` and ``data.joint_effort_limits``. - """ - joint_velocity_limit = joint_effort_limit = joint_limit - articulation_cfg = generate_articulation_cfg( - articulation_type="single_joint_implicit", - joint_velocity_limit=joint_velocity_limit, - joint_effort_limit=joint_effort_limit, - ) - articulation, _ = generate_articulation( - articulation_cfg=articulation_cfg, - num_articulations=1, - device=device, - ) - # Play sim - 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_vel_limits.torch, physx_vel_limit) - # the solver clamp comes from joint_velocity_limit when set, otherwise the USD-authored value - if joint_velocity_limit is None: - limit = next( - p.max_joint_velocity for p in articulation_cfg.spawn.joint_drive_props if isinstance(p, PhysxJointCfg) - ) - else: - limit = joint_velocity_limit - expected_velocity_limit = torch.full_like(physx_vel_limit, limit) - torch.testing.assert_close(physx_vel_limit, expected_velocity_limit) + assert_dynamics_match(jacobian=False, gravity=False) - # obtain the physx effort limits - physx_effort_limit = _read_binding_to_torch(articulation, TT.DOF_MAX_FORCE, device) - # check data buffer - torch.testing.assert_close(articulation.data.joint_effort_limits.torch, physx_effort_limit) - # the solver keeps the USD-authored limit unless the user overrides it explicitly - if joint_effort_limit is None: - limit = next( - p.max_force for p in articulation_cfg.spawn.joint_drive_props if isinstance(p, sim_utils.UsdPhysicsDriveCfg) - ) - else: - limit = joint_effort_limit - expected_effort_limit = torch.full_like(physx_effort_limit, limit) - torch.testing.assert_close(physx_effort_limit, expected_effort_limit) +@pytest.mark.parametrize("scene", _CUDA_SCENES, indirect=True) +@pytest.mark.parametrize("island", ["reversed", "reversed_ordered"]) +def test_reversed_joint_dynamics_use_public_joint_basis(scene: _ArticulationScene, island: str) -> None: + """Check velocity, kinetic energy and gravity in the public joint basis.""" + articulation = scene.islands[island] + device = scene.device + assert (articulation.joint_ordering is not None) == (island == "reversed_ordered") + + articulation.write_joint_position_to_sim_index(position=torch.zeros_like(articulation.data.joint_pos.torch)) + velocity = torch.zeros((1, articulation.num_joints), device=device) + velocity[:, articulation.find_joints("left_shoulder")[0][0]] = 0.4 + velocity[:, articulation.find_joints("left_elbow")[0][0]] = 0.7 + articulation.write_joint_velocity_to_sim_index(velocity=velocity) -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -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 - ) + joint_velocity = articulation.data.joint_vel.torch + predicted_velocity = torch.einsum("nbij,nj->nbi", articulation.data.body_com_jacobian_w.torch, joint_velocity) + torch.testing.assert_close(predicted_velocity, articulation.data.body_com_vel_w.torch[:, 1:], atol=1e-5, rtol=1e-5) - # Play the simulator - sim.reset() + generalized_energy = 0.5 * torch.einsum( + "ni,nij,nj->n", joint_velocity, articulation.data.mass_matrix.torch, joint_velocity + ) + body_velocity = articulation.data.body_com_vel_w.torch + body_inertia = articulation.data.body_inertia.torch.reshape(1, articulation.num_bodies, 3, 3) + body_energy = 0.5 * ( + (articulation.data.body_mass.torch.unsqueeze(-1) * body_velocity[..., :3].square()).sum((-1, -2)) + + torch.einsum("nbi,nbij,nbj->n", body_velocity[..., 3:], body_inertia, body_velocity[..., 3:]) + ) + torch.testing.assert_close(generalized_energy, body_energy, atol=1e-5, rtol=1e-5) - # Now we are ready! - # reset articulation - articulation.reset() + gravity = torch.tensor(scene.sim.cfg.gravity, device=device) + body_weight = articulation.data.body_mass.torch[:, 1:, None] * gravity + expected_compensation = -torch.einsum( + "nbij,nbi->nj", articulation.data.body_com_jacobian_w.torch[:, :, :3], body_weight + ) + assert torch.all(expected_compensation.abs() > 0.1) + torch.testing.assert_close( + articulation.data.gravity_compensation_forces.torch, expected_compensation, atol=1e-5, rtol=1e-5 + ) - # 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 +@pytest.mark.parametrize("scene", _CUDA_SCENES, indirect=True) +def test_floating_articulation_root_state_dynamics_wrench_and_reset(scene: _ArticulationScene) -> None: + """Prove floating-root state, generalized dynamics, reset, and a real OVPhysX external-wrench response.""" + articulation = scene.islands["floating"] + device = scene.device + num_articulations = articulation.num_instances 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), + num_joints = articulation.num_joints + assert articulation.is_initialized + assert not articulation.is_fixed_base + assert articulation.data.root_link_pose_w.torch.shape == (num_articulations, 7) + assert articulation.data.root_com_pose_w.torch.shape == (num_articulations, 7) + assert articulation.data.body_link_pose_w.torch.shape == (num_articulations, num_bodies, 7) + assert articulation.data.body_com_pose_w.torch.shape == (num_articulations, num_bodies, 7) + + # Floating-base columns stay leading while the actuated-joint axes are gathered. + user_to_backend = torch.as_tensor(articulation.joint_ordering.user_to_backend_indices, device=device) + generalized_user_to_backend = torch.cat((torch.arange(6, device=device), 6 + user_to_backend)) + raw_jacobian = _read_binding_to_torch(articulation, TT.JACOBIAN, device).reshape( + num_articulations, num_bodies, 6, num_joints + 6 + ) + raw_mass_matrix = _read_binding_to_torch(articulation, TT.MASS_MATRIX, device) + raw_gravity = _read_binding_to_torch(articulation, TT.GRAVITY_FORCE, device) + torch.testing.assert_close( + articulation.data.body_com_jacobian_w.torch, + raw_jacobian[:, :, :, generalized_user_to_backend], ) - 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), + torch.testing.assert_close( + articulation.data.mass_matrix.torch, + raw_mass_matrix[:, generalized_user_to_backend, :][:, :, generalized_user_to_backend], ) - 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", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("add_ground_plane", [True]) -def test_apply_joint_command(sim, num_articulations, device, add_ground_plane): - """Arm targets track commands while the passive finger follows the driven finger.""" - articulation_cfg = generate_articulation_cfg(articulation_type="panda") - articulation, _ = generate_articulation( - articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=device + torch.testing.assert_close( + articulation.data.gravity_compensation_forces.torch, + raw_gravity[:, generalized_user_to_backend], ) - # 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 - leader_id = articulation.find_joints("panda_finger_joint1")[0][0] - follower_id = articulation.find_joints("panda_finger_joint2")[0][0] - initial_leader_pos = articulation.data.joint_pos.torch[:, leader_id].clone() - joint_pos[:, [leader_id, follower_id]] = 0.01 - - # apply action to the articulation - articulation.set_joint_position_target_index(target=joint_pos) - articulation.write_data_to_sim() + # A partial reset clears the wrenches of the selected environment only; a full reset clears them all. + composers = (articulation.instantaneous_wrench_composer, articulation.permanent_wrench_composer) + ones = torch.ones((num_articulations, num_bodies, 3), device=device) + articulation.permanent_wrench_composer.set_forces_and_torques_index(forces=ones, torques=ones) + articulation.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones) + articulation.reset(env_ids=torch.tensor([0], device=device)) + for composer in composers: + assert composer.active + assert torch.count_nonzero(composer.composed_force.torch) == num_bodies * 3 + assert torch.count_nonzero(composer.composed_torque.torch) == num_bodies * 3 + articulation.reset() + for composer in composers: + assert not composer.active + assert torch.count_nonzero(composer.composed_force.torch) == 0 + assert torch.count_nonzero(composer.composed_torque.torch) == 0 - for _ in range(100): - # perform step - sim.step() - # update buffers - articulation.update(sim.cfg.dt) - - # The elbow, wrist, and finger joints track the commanded target. The shoulder joints sag under - # gravity with these gains, so they are not checked. Without the target write, the drives pull - # every joint toward zero instead (the wrist ends near 0.3 rad instead of 3.0 rad). - torch.testing.assert_close(articulation.data.joint_pos.torch[:, 3:], joint_pos[:, 3:], atol=0.1, rtol=0.0) - - leader_pos = articulation.data.joint_pos.torch[:, leader_id] - follower_pos = articulation.data.joint_pos.torch[:, follower_id] - assert torch.all(torch.abs(leader_pos - initial_leader_pos) > 0.005) - torch.testing.assert_close(follower_pos, leader_pos, rtol=0.0, atol=5.0e-4) - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_body_root_state(sim, num_articulations, device): - """Test for reading the `body_state_w` property. - - This test verifies that: - 1. Body states can be read correctly - 2. States are correct with a COM offset from the link frame - 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 - """ - sim._app_control_on_stop_handle = None - articulation_cfg = generate_articulation_cfg(articulation_type="single_joint_implicit") - articulation, _ = 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 - offset = [0.5, 0.0, 0.0] - - # create com offsets — apply offset to the Arm body - num_bodies = articulation.num_bodies + # A written root frame moves the other frame through the root-body COM offset. 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), + com[:, 0, :3] = torch.tensor([1.0, 0.0, 0.0], device=device) + articulation.set_coms_index(coms=wp.from_torch(com, dtype=wp.transformf)) + com_pos_b, com_quat_b = com[:, 0, :3], com[:, 0, 3:7] + link_pos_in_com, link_quat_in_com = math_utils.subtract_frame_transforms(com_pos_b, com_quat_b) + rand_state = torch.zeros(num_articulations, 13, device=device) + rand_state[..., :3] = articulation.data.root_link_pos_w.torch + torch.tensor([0.1, 0.2, 0.3], device=device) + rand_state[..., 3:7] = torch.nn.functional.normalize(torch.tensor([0.1, 0.2, 0.3, 0.9], device=device), dim=-1) + for env_ids in (None, torch.arange(num_articulations, dtype=torch.int32, device=device)): + articulation.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_ids) + articulation.write_root_com_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_ids) + 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) + # The link frame is the written COM frame composed with the inverse COM offset. + expected_link_pos, expected_link_quat = math_utils.combine_frame_transforms( + rand_state[..., :3], rand_state[..., 3:7], link_pos_in_com, link_quat_in_com + ) + torch.testing.assert_close(articulation.data.root_link_pos_w.torch, expected_link_pos) + torch.testing.assert_close(articulation.data.root_link_quat_w.torch, expected_link_quat) + + articulation.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_ids) + articulation.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_ids) + 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) + # The COM frame is the written link frame composed with the COM offset. + expected_com_pos, expected_com_quat = math_utils.combine_frame_transforms( + rand_state[..., :3], rand_state[..., 3:7], com_pos_b, com_quat_b + ) + torch.testing.assert_close(articulation.data.root_com_pos_w.torch, expected_com_pos) + torch.testing.assert_close(articulation.data.root_com_quat_w.torch, expected_com_quat) + + # A partial root write keeps the other environment, and a wrench on it accelerates only that one. + env_ids = torch.tensor([1], dtype=torch.int32, device=device) + initial_pose = articulation.data.root_link_pose_w.torch.clone() + target_pose = initial_pose[env_ids].clone() + target_pose[:, :3] += torch.tensor([0.2, -0.1, 0.3], device=device) + articulation.write_root_link_pose_to_sim_index(root_pose=target_pose, env_ids=env_ids) + torch.testing.assert_close(articulation.data.root_link_pose_w.torch[env_ids], target_pose) + torch.testing.assert_close(articulation.data.root_link_pose_w.torch[:1], initial_pose[:1]) + articulation.write_root_velocity_to_sim_index(root_velocity=torch.zeros_like(rand_state[..., 7:])) + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=torch.tensor([[[8.0, 0.0, 0.0]]], device=device), + torques=torch.zeros((1, 1, 3), device=device), + env_ids=env_ids, + body_ids=torch.tensor([0], dtype=torch.int32, device=device), ) + scene.step(articulation) + # Gravity acts along z, so only the pushed environment gains velocity along x. + assert articulation.data.root_com_lin_vel_w.torch[1, 0] > 1e-3 + torch.testing.assert_close( + articulation.data.root_com_lin_vel_w.torch[0, :2], torch.zeros(2, device=device), atol=1e-6, rtol=0 + ) + articulation.reset() - # 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_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 - - # 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 - # OVStage determines the runtime link pose from the joint frames, which may differ - # from the authored USD Xform. Verify the COM offset relative to that runtime pose. - pos_gt = body_link_pose_w[..., :3].clone() - px = offset[0] * torch.cos(joint_pos) - py = torch.zeros(num_articulations, 1, 1, device=device) - pz = offset[0] * torch.sin(joint_pos) - pos_gt[:, arm_idx, :] += torch.cat([px, py, pz], dim=-1).squeeze(-2) - 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:]) - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("state_location", ["com", "link"]) -@pytest.mark.parametrize("gravity_enabled", [False]) -def test_write_root_state(sim, num_articulations, device, 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. The other frame follows the written one through the COM offset - 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 - 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) +@pytest.mark.parametrize("scene", _CUDA_SCENES, indirect=True) +def test_tendon_properties_and_position_targets(scene: _ArticulationScene) -> None: + """Prove locally authored spatial and fixed tendons are discovered and writable. - # Play sim - sim.reset() + A fixed-tendon target lands in the simulation as ``rest_length - target``. The index form commands + every tendon of environment 0; the mask form commands tendon 0 of environment 1. Every other cell must + keep its initial offset. + """ + articulation = scene.islands["tendon"] + device = scene.device + num_articulations = articulation.num_instances + assert articulation.is_initialized + assert articulation.is_fixed_base + assert articulation.num_spatial_tendons == 1 + env_ids = torch.tensor([1], dtype=torch.int32, device=device) + initial_stiffness = articulation.data.spatial_tendon_stiffness.torch.clone() + stiffness = torch.tensor([[12.0]], device=device) + damping = torch.tensor([[1.5]], device=device) + limit_stiffness = torch.tensor([[3.0]], device=device) + offset = torch.tensor([[0.1]], device=device) + articulation.set_spatial_tendon_stiffness_index(stiffness=stiffness, env_ids=env_ids) + articulation.set_spatial_tendon_damping_index(damping=damping, env_ids=env_ids) + articulation.set_spatial_tendon_limit_stiffness_index(limit_stiffness=limit_stiffness, env_ids=env_ids) + articulation.set_spatial_tendon_offset_index(offset=offset, env_ids=env_ids) + torch.testing.assert_close(articulation.data.spatial_tendon_stiffness.torch[env_ids], stiffness) + torch.testing.assert_close(articulation.data.spatial_tendon_stiffness.torch[:1], initial_stiffness[:1]) + torch.testing.assert_close(articulation.data.spatial_tendon_damping.torch[env_ids], damping) + torch.testing.assert_close(articulation.data.spatial_tendon_limit_stiffness.torch[env_ids], limit_stiffness) + torch.testing.assert_close(articulation.data.spatial_tendon_offset.torch[env_ids], offset) - # change center of mass offset from link frame - offset = torch.tensor([1.0, 0.0, 0.0]).repeat(num_articulations, 1, 1) + num_tendons = articulation.num_fixed_tendons + assert num_tendons == 2 + rest_length = articulation.data.fixed_tendon_rest_length.torch.clone() + initial_offset = articulation.data.fixed_tendon_offset.torch.clone() - # 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), + index_target = torch.full((1, num_tendons), 0.3, dtype=torch.float32, device=device) + articulation.set_fixed_tendon_position_target_index(target=index_target, env_ids=[0]) + # Distinct per-cell values: a uniform target cannot catch the mask form reading the wrong + # cell, because every wrong read returns the same number. + mask_target = ( + 0.7 + + 0.1 * torch.arange(num_articulations, dtype=torch.float32, device=device).unsqueeze(1) + + 0.01 * torch.arange(num_tendons, dtype=torch.float32, device=device).unsqueeze(0) + ) + env_mask = wp.array([False, True], dtype=wp.bool, device=device) + tendon_mask = wp.from_torch(torch.arange(num_tendons, device=device) == 0) + articulation.set_fixed_tendon_position_target_mask( + target=mask_target, fixed_tendon_mask=tendon_mask, env_mask=env_mask ) - # check they are set - torch.testing.assert_close(_read_binding_to_torch(articulation, TT.BODY_COM_POSE, device), com) + scene.step(articulation) - 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) + expected = initial_offset.clone() + expected[0] = rest_length[0] - 0.3 + expected[1, 0] = rest_length[1, 0] - mask_target[1, 0] + torch.testing.assert_close(articulation.data.fixed_tendon_offset.torch, expected) - # Root-body COM frame in the link frame, from the values written above. - com_pos_b, com_quat_b = com[:, 0, :3], com[:, 0, 3:7] - link_pos_in_com, link_quat_in_com = math_utils.subtract_frame_transforms(com_pos_b, com_quat_b) - 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) - # The link frame is the written COM frame composed with the inverse COM offset. - expected_link_pos, expected_link_quat = math_utils.combine_frame_transforms( - rand_state[..., :3], rand_state[..., 3:7], link_pos_in_com, link_quat_in_com - ) - torch.testing.assert_close(articulation.data.root_link_pos_w.torch, expected_link_pos) - torch.testing.assert_close(articulation.data.root_link_quat_w.torch, expected_link_quat) - 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) - # The COM frame is the written link frame composed with the COM offset. - expected_com_pos, expected_com_quat = math_utils.combine_frame_transforms( - rand_state[..., :3], rand_state[..., 3:7], com_pos_b, com_quat_b - ) - torch.testing.assert_close(articulation.data.root_com_pos_w.torch, expected_com_pos) - torch.testing.assert_close(articulation.data.root_com_quat_w.torch, expected_com_quat) +@pytest.mark.parametrize("scene", _CUDA_SCENES, indirect=True) +def test_native_actuator_submits_real_effort(scene: _ArticulationScene) -> None: + """Run a Newton-native explicit actuator through the current OVPhysX state and effort binding.""" + articulation = scene.islands["native"] + device = scene.device + assert articulation._actuator_control.native_actuator_path_active + assert articulation.newton_actuator_adapter is not None + backend_to_user = list(articulation.joint_ordering.backend_to_user_indices) + user_to_backend = list(articulation.joint_ordering.user_to_backend_indices) + + initial_pos = articulation.data.joint_pos.torch.clone() + target = initial_pos + 0.2 + articulation.actuators.target_command.set_position_index(value=target) + articulation.write_data_to_sim() + assert torch.any(articulation.actuators.computed_effort.torch != 0.0) + assert torch.any(articulation.actuators.applied_effort.torch != 0.0) + raw_effort = wp.to_torch(articulation._physx_actuator_wrapper.joint_f_2d).clone() + assert torch.any(raw_effort != 0.0) + backend_effort = _read_binding_to_torch(articulation, TT.DOF_ACTUATION_FORCE, device) + torch.testing.assert_close(backend_effort, raw_effort[:, backend_to_user]) + torch.testing.assert_close(backend_effort, articulation.actuators.applied_effort.torch[:, backend_to_user]) -@pytest.mark.parametrize("device", ["cpu"]) -def test_com_orientation_write_invalidates_static_inertia_cache_with_body_ordering(sim, device): - """Partial COM writes and COM rotations follow non-identity body order. + scene.step(articulation, count=8) + # Use raw OV bindings so the observation cannot refresh the public state shadow. + current_pos = _read_binding_to_torch(articulation, TT.DOF_POSITION, device)[:, user_to_backend] + current_vel = _read_binding_to_torch(articulation, TT.DOF_VELOCITY, device)[:, user_to_backend] + assert not torch.allclose(current_pos, initial_pos) - COM pose is a CPU-resident OVPhysX binding, and the raw ``root_view.set_attribute`` restore - below forbids cross-device staging, so this test is CPU-only. - """ - sim._app_control_on_stop_handle = None - articulation_cfg = replace(FRANKA_PANDA_CFG, body_ordering=PANDA_ROOT_PRESERVING_REVERSED_BODY_NAMES) - articulation, _ = generate_articulation(articulation_cfg, 1, device=device) + articulation.write_data_to_sim() + expected_effort = torch.clamp(_STIFFNESS * (target - current_pos) - _DAMPING * current_vel, -_MAX_FORCE, _MAX_FORCE) + torch.testing.assert_close(articulation.actuators.applied_effort.torch, expected_effort) + torch.testing.assert_close( + _read_binding_to_torch(articulation, TT.DOF_ACTUATION_FORCE, device), + wp.to_torch(articulation._physx_actuator_wrapper.joint_f_2d)[:, backend_to_user], + ) - sim.reset() - articulation.update(sim.cfg.dt) - assert articulation.body_ordering is not None - public_body_id = 1 - backend_body_id = articulation.body_ordering.user_to_backend_indices[public_body_id] - assert backend_body_id != public_body_id +@pytest.mark.parametrize("scene", _CUDA_SCENES, indirect=True) +def test_native_actuator_reset_and_gain_event_are_environment_selective(scene: _ArticulationScene) -> None: + """Reset and randomize only the selected OVPhysX native-controller environment.""" + # Newton is only needed by the native actuator path. + from isaaclab.actuators.newton import read_group_parameter # noqa: PLC0415 - # A partial ordered COM write lands on the selected backend body and preserves every other one. - backend_before = _read_binding_to_torch(articulation, TT.BODY_COM_POSE, device).clone() - assert torch.unique(backend_before[0], dim=0).shape[0] > 1 - articulation.data._body_com_pose_b.timestamp = -1.0 - backend_staging = articulation.data._body_com_pose_b_backend - if backend_staging is not None: - backend_staging.timestamp = -1.0 - selected_com = backend_before[0, backend_body_id].clone() - selected_com[0] += 0.001 - articulation.set_coms_index( - coms=wp.from_torch(selected_com.reshape(1, 1, 7).contiguous(), dtype=wp.transformf), - env_ids=wp.array([0], dtype=wp.int32, device=device), - body_ids=wp.array([public_body_id], dtype=wp.int32, device=device), - ) - backend_after = _read_binding_to_torch(articulation, TT.BODY_COM_POSE, device).clone() - articulation.root_view.set_attribute(TT.BODY_COM_POSE, wp.from_torch(backend_before.contiguous())) - noop_after = _read_binding_to_torch(articulation, TT.BODY_COM_POSE, device).clone() - unselected_body_mask = torch.ones(backend_before.shape[1], dtype=torch.bool, device=device) - unselected_body_mask[backend_body_id] = False - assert torch.equal(noop_after[..., :3], backend_before[..., :3]) - assert torch.equal(backend_after[0, backend_body_id, :3], selected_com[:3]) - assert torch.equal(backend_after[0, unselected_body_mask, :3], backend_before[0, unselected_body_mask, :3]) - # Bound semantic orientation equality by the native setter's float32 no-op normalization. - native_orientation_atol = torch.max( - torch.abs(noop_after[0, unselected_body_mask, 3:7] - backend_before[0, unselected_body_mask, 3:7]) - ).item() - assert native_orientation_atol <= torch.finfo(backend_before.dtype).eps - torch.testing.assert_close( - backend_after[0, unselected_body_mask, 3:7], - backend_before[0, unselected_body_mask, 3:7], - rtol=0.0, - atol=native_orientation_atol, - ) - assert torch.equal(backend_after[..., 3:7], noop_after[..., 3:7]) + class Env: + def __init__(self, asset): + self.scene = self + self.num_envs = asset.num_instances + self.device = asset.device + self._asset = asset - # A COM rotation refreshes the static inertia cache. - coms = articulation.data.body_com_pose_b.torch[:, public_body_id : public_body_id + 1].clone() - coms[..., 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) - articulation.set_coms_index(coms=wp.from_torch(coms.contiguous(), dtype=wp.transformf), body_ids=[public_body_id]) + def __getitem__(self, name): + assert name == "robot" + return self._asset - principal_inertia = torch.tensor([[[1.0, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 3.0]]], device=device) - articulation.set_inertias_index(inertias=principal_inertia, body_ids=[public_body_id]) - torch.testing.assert_close( - articulation.data.body_inertia.torch[:, public_body_id : public_body_id + 1], principal_inertia - ) + articulation = scene.islands["delayed"] + device = scene.device + num_joints = articulation.num_joints + scene.step(articulation, count=3) + + adapter = articulation.newton_actuator_adapter + stateful_pairs = [ + state + for actuator, state in zip(adapter.actuators, adapter._states_a) + if state is not None and getattr(state, "delay_state", None) is not None + ] + assert len(stateful_pairs) == 1 + articulation.reset(env_ids=torch.tensor([0], device=device, dtype=torch.long)) + # The delay state is flattened environment-major over the group's joints. + assert stateful_pairs[0].delay_state.num_pushes.numpy().tolist() == [0] * num_joints + [1] * num_joints + + env = Env(articulation) + asset_cfg = SceneEntityCfg("robot") + event_params = { + "asset_cfg": asset_cfg, + "stiffness_distribution_params": (101.0, 101.0), + "damping_distribution_params": (3.0, 3.0), + "operation": "abs", + "distribution": "uniform", + } + event = randomize_actuator_gains(EventTermCfg(func=randomize_actuator_gains, params=event_params), env) + event(env, env_ids=torch.tensor([0], device=device), **event_params) - coms[..., 3:7] = torch.tensor([0.0, 0.0, 0.70710677, 0.70710677], device=device) - articulation.set_coms_index(coms=wp.from_torch(coms.contiguous(), dtype=wp.transformf), body_ids=[public_body_id]) - sim.step() - articulation.update(sim.cfg.dt) - expected_rotated_inertia = torch.tensor([[[2.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 3.0]]], device=device) - torch.testing.assert_close( - articulation.data.body_inertia.torch[:, public_body_id : public_body_id + 1], - expected_rotated_inertia, - ) + stiffness = read_group_parameter(articulation.actuators, "joint", "controller", "kp") + damping = read_group_parameter(articulation.actuators, "joint", "controller", "kd") + torch.testing.assert_close(stiffness, torch.tensor([[101.0], [20.0]], device=device).repeat(1, num_joints)) + torch.testing.assert_close(damping, torch.tensor([[3.0], [1.0]], device=device).repeat(1, num_joints)) -@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) -def test_root_link_vel_w_refreshes_fk_before_body_com_vel_w_read(sim, device, monkeypatch): +@pytest.mark.parametrize("scene", _CUDA_SCENES, indirect=True) +def test_root_link_vel_w_refreshes_fk_before_body_com_vel_w_read(scene: _ArticulationScene, monkeypatch) -> None: """Reading ``root_link_vel_w`` must run FK before ``body_com_vel_w`` sees a "fresh" buffer. Native tensor reads may refresh internally, so also verify that asset and rendering reads share one explicit FK update after each write, regardless of which reader comes first. """ - sim._app_control_on_stop_handle = None - articulation_cfg = generate_articulation_cfg(articulation_type="single_joint_implicit") - articulation, _ = generate_articulation(articulation_cfg, 2, device=device) - - sim.reset() - articulation.update(sim.cfg.dt) + articulation = scene.islands["ordered"] + articulation.update(scene.sim.cfg.dt) # Prime the derived buffers before the write so their TimestampedBuffers are populated; otherwise # the reads below would trivially be "first reads" regardless of the cache-invalidation bug. articulation.data.root_link_vel_w articulation.data.body_com_vel_w - joint_vel = torch.full((2, articulation.num_joints), 3.0, device=device) + joint_vel = torch.full((2, articulation.num_joints), 3.0, device=scene.device) for render_first in (False, True): articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) articulation.write_joint_position_to_sim_index(position=articulation.data.joint_pos.torch.clone()) @@ -2684,309 +1717,6 @@ def test_root_link_vel_w_refreshes_fk_before_body_com_vel_w_read(sim, device, mo assert torch.linalg.norm(body_com_vel_w[:, 1, :]) > 1e-3 -@pytest.mark.parametrize("device", test_devices()) -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", test_devices()) -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("device", test_devices()) -def test_write_joint_state_to_sim_index_partial(sim, device, mocker): - """Test fused joint-state writes with partial environment and joint indices. - - The first write uses unsorted int64 selectors; the deprecated combined writer must delegate - to the public index writers. - """ - sim._app_control_on_stop_handle = None - articulation_cfg = generate_articulation_cfg(articulation_type="anymal") - articulation, _ = generate_articulation(articulation_cfg, 2, device) - sim.reset() - - original_joint_pos = articulation.data.joint_pos.torch.clone() - original_joint_vel = articulation.data.joint_vel.torch.clone() - _ = articulation.data.body_link_pose_w - _ = articulation.data.body_com_vel_w - pose_timestamp = articulation.data._body_link_pose_w.timestamp - velocity_timestamp = articulation.data._body_com_vel_w.timestamp - - previous_joint_vel = wp.to_torch(articulation.data._previous_joint_vel) - joint_acc = wp.to_torch(articulation.data._joint_acc.data) - previous_joint_vel.fill_(3.0) - joint_acc.fill_(4.0) - articulation.data._joint_acc.timestamp = -1.0 - - position = torch.tensor([[0.1, -0.1]], device=device) - velocity = torch.tensor([[0.2, -0.2]], device=device) - # int64 selectors with the joints in reverse order; the payload follows the selector order. - articulation.write_joint_state_to_sim_index( - position=position[:, [1, 0]], - velocity=velocity[:, [1, 0]], - env_ids=torch.tensor([1], dtype=torch.int64, device=device), - joint_ids=torch.tensor([2, 0], dtype=torch.int64, device=device), - skip_forward=True, - ) - - expected_joint_pos = original_joint_pos.clone() - expected_joint_vel = original_joint_vel.clone() - expected_previous_joint_vel = torch.full_like(previous_joint_vel, 3.0) - expected_joint_acc = torch.full_like(joint_acc, 4.0) - expected_joint_pos[1, [0, 2]] = position[0] - expected_joint_vel[1, [0, 2]] = velocity[0] - expected_previous_joint_vel[1, [0, 2]] = velocity[0] - expected_joint_acc[1, [0, 2]] = 0.0 - - torch.testing.assert_close(articulation.data.joint_pos.torch, expected_joint_pos) - torch.testing.assert_close(articulation.data.joint_vel.torch, expected_joint_vel) - torch.testing.assert_close(previous_joint_vel, expected_previous_joint_vel) - torch.testing.assert_close(joint_acc, expected_joint_acc) - assert articulation.data._joint_acc.timestamp == articulation.data._sim_timestamp - assert articulation.data._body_link_pose_w.timestamp == pose_timestamp - assert articulation.data._body_com_vel_w.timestamp == velocity_timestamp - torch.testing.assert_close(_read_binding_to_torch(articulation, TT.DOF_POSITION, device), expected_joint_pos) - torch.testing.assert_close(_read_binding_to_torch(articulation, TT.DOF_VELOCITY, device), expected_joint_vel) - - articulation.write_joint_state_to_sim_index(position=position, velocity=velocity, env_ids=[1], joint_ids=[0, 2]) - assert articulation.data._body_link_pose_w.timestamp < articulation.data._sim_timestamp - assert articulation.data._body_com_vel_w.timestamp < articulation.data._sim_timestamp - - # The deprecated combined API stays a thin composition of the public writers. - position = torch.tensor([[0.1]], device=device) - velocity = torch.tensor([[0.2]], device=device) - position_writer = mocker.patch.object(articulation, "write_joint_position_to_sim_index") - velocity_writer = mocker.patch.object(articulation, "write_joint_velocity_to_sim_index") - with pytest.warns(DeprecationWarning): - articulation.write_joint_state_to_sim(position=position, velocity=velocity, joint_ids=[0], env_ids=[0]) - position_writer.assert_called_once_with(position=position, joint_ids=[0], env_ids=[0]) - velocity_writer.assert_called_once_with(velocity=velocity, joint_ids=[0], env_ids=[0]) - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("gravity_enabled", [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_state_to_sim_index(position=rand_joint_pos, 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 - # 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_vel_w, body_com_vel_w_fresh) - - -@pytest.mark.parametrize("add_ground_plane", [True]) -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_write_joint_frictions_to_sim(sim, num_articulations, device, add_ground_plane): - """Write joint friction coefficients and read them back from the simulation and public getters.""" - 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() - - # 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() - sim.step() - 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()) - - # The public getters expose the same coefficients. - torch.testing.assert_close(articulation.data.joint_friction_coeff.torch, friction) - torch.testing.assert_close(articulation.data.joint_dynamic_friction_coeff.torch, dynamic_friction) - torch.testing.assert_close(articulation.data.joint_viscous_friction_coeff.torch, viscous_friction) - - -@pytest.mark.parametrize("add_ground_plane", [True]) -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_set_material_properties(sim, num_articulations, device, add_ground_plane): - """Randomize per-shape articulation materials through the OVPhysX material event. - - OVPhysX exposes per-collision-shape material as the - ``articulation_shape_friction_and_restitution`` tensor binding (shape ``[N, S, 3]`` = - static friction, dynamic friction, restitution), addressed through the - :class:`~isaaclab_ov.sim.views.OvPhysxView`. The binding is CPU-native, so the - buffer lives in host memory. - """ - from isaaclab.envs.mdp import randomize_rigid_body_material # noqa: PLC0415 - - 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() - - # The ranges exclude the asset's default material, so values inside them prove the write happened. - static_range, dynamic_range, restitution_range = (1.5, 2.0), (1.2, 1.4), (0.6, 0.8) - view = articulation.root_view - materials_before = wp.to_torch(view.get_attribute(TT.SHAPE_FRICTION_AND_RESTITUTION)).clone() - assert (materials_before[..., 0] < static_range[0]).all(), f"default material overlaps: {materials_before}" - params = { - "static_friction_range": static_range, - "dynamic_friction_range": dynamic_range, - "restitution_range": restitution_range, - "num_buckets": 16, - "asset_cfg": SceneEntityCfg("robot"), - } - env = SimpleNamespace(scene={"robot": articulation}, sim=sim) - cfg = EventTermCfg(func=randomize_rigid_body_material, mode="reset", params=params) - randomize = randomize_rigid_body_material(cfg, env) - - # Randomize only the last environment; the others keep their materials. - randomize(env, torch.tensor([num_articulations - 1], device=device), **cfg.params) - sim.step() - articulation.update(sim.cfg.dt) - - materials = wp.to_torch(view.get_attribute(TT.SHAPE_FRICTION_AND_RESTITUTION)) - torch.testing.assert_close(materials[:-1], materials_before[:-1]) - eps = 1e-5 - for component, (lo, hi) in enumerate((static_range, dynamic_range, restitution_range)): - values = materials[-1, :, component] - assert ((values >= lo - eps) & (values <= hi + eps)).all(), values - - @wp.kernel def _occupy_stream_kernel(iterations: int, out: wp.array(dtype=wp.float32)): total = float(0.0) @@ -2995,8 +1725,8 @@ def _occupy_stream_kernel(iterations: int, out: wp.array(dtype=wp.float32)): out[0] = total -@pytest.mark.parametrize("device", ["cuda:0"]) -def test_cpu_only_property_writes_wait_for_pinned_host_staging(sim, device): +@pytest.mark.parametrize("scene", _CUDA_SCENES, indirect=True) +def test_cpu_only_property_writes_wait_for_pinned_host_staging(scene: _ArticulationScene) -> None: """Land CPU-only property writes while the device stream is busy. OVPhysX keeps joint and body properties on the host even for a GPU simulation, so the writers @@ -3005,18 +1735,8 @@ def test_cpu_only_property_writes_wait_for_pinned_host_staging(sim, device): device stream, so a long kernel queued ahead of the copy must not let the setter consume the previous contents of the pinned buffers. """ - articulation, _ = generate_articulation( - ArticulationCfg( - prim_path="/World/Robot", - spawn=sim_utils.UsdFileCfg( - usd_path=str(Path(__file__).parent / "data" / "articulation_ordering_branching.usda") - ), - actuators={}, - ), - num_articulations=4, - device=device, - ) - sim.reset() + articulation = scene.islands["staging"] + device = scene.device num_envs, num_joints = articulation.num_instances, articulation.num_joints stiffness_before = _read_binding_to_torch(articulation, TT.DOF_STIFFNESS, device) @@ -3058,5 +1778,46 @@ def occupy_device_stream(iterations: int = 500_000): torch.testing.assert_close(_read_binding_to_torch(articulation, TT.BODY_MASS, device), expected_masses) -if __name__ == "__main__": - pytest.main([__file__, "-v", "--maxfail=1"]) +@pytest.mark.parametrize("scene", _ALL_DEVICES, indirect=True) +def test_joint_position_limit_writes_clamp_default_joint_positions(scene: _ArticulationScene, caplog) -> None: + """Clamp only the selected default joint positions into written limits; logging never changes clamping. + + The logging level controls reporting without reusing a stale violation count. + """ + articulation = scene.islands["ordered"] + device = scene.device + logger = type(articulation).__module__ + original_limits = articulation.data.joint_pos_limits.torch.clone() + original_defaults = articulation.data.default_joint_pos.torch.clone() + try: + # The default positions are zero, so both new ranges exclude them. + env_ids = torch.tensor([1], dtype=torch.int32, device=device) + joint_ids = torch.tensor([0, 2], dtype=torch.int32, device=device) + partial_limits = torch.tensor([[[0.2, 0.3], [0.25, 0.35]]], device=device) + articulation.write_joint_position_limit_to_sim_index( + limits=partial_limits, env_ids=env_ids, joint_ids=joint_ids + ) + torch.testing.assert_close(articulation.data.joint_pos_limits.torch[env_ids][:, joint_ids], partial_limits) + default_joint_pos = articulation.data.default_joint_pos.torch + selected = default_joint_pos[env_ids][:, joint_ids] + assert torch.all((selected >= partial_limits[..., 0]) & (selected <= partial_limits[..., 1])) + unselected = torch.ones_like(default_joint_pos, dtype=torch.bool) + unselected[env_ids[:, None].long(), joint_ids.long()] = False + torch.testing.assert_close(default_joint_pos[unselected], original_defaults[unselected]) + + limits = torch.zeros_like(original_limits) + limits[..., 1] = 0.5 + for level in (logging.WARNING, logging.INFO): + articulation.data.default_joint_pos.torch.fill_(1.0) + caplog.clear() + with caplog.at_level(level, logger=logger): + articulation.write_joint_position_limit_to_sim_index(limits=limits, warn_limit_violation=False) + assert [record.levelno for record in caplog.records] == ([logging.INFO] if level == logging.INFO else []) + torch.testing.assert_close(articulation.data.default_joint_pos.torch, torch.full_like(limits[..., 1], 0.5)) + caplog.clear() + with caplog.at_level(logging.INFO, logger=logger): + articulation.write_joint_position_limit_to_sim_mask(limits=limits, warn_limit_violation=False) + assert not caplog.records + finally: + articulation.write_joint_position_limit_to_sim_index(limits=original_limits) + articulation.data.default_joint_pos.torch.copy_(original_defaults) diff --git a/source/isaaclab_ov/test/assets/test_articulation_helpers.py b/source/isaaclab_ov/test/assets/test_articulation_helpers.py deleted file mode 100644 index ca43d1b1c46b..000000000000 --- a/source/isaaclab_ov/test/assets/test_articulation_helpers.py +++ /dev/null @@ -1,119 +0,0 @@ -# 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) 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 OVPhysX runtime wheel is optional. Skip gracefully when it is not installed; -# CI jobs that need OVPhysX coverage install it explicitly. -pytest.importorskip("ovphysx.types", reason="ovphysx wheel not installed") - -from isaaclab_ov.assets.articulation.articulation import Articulation # noqa: E402 -from isaaclab_ov.physics import OvPhysxManager # noqa: E402 -from isaaclab_ov.test.fixtures.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_usda() -> str: - """Serialize one relevant articulation subtree and unrelated joints in memory.""" - 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", - ) - - return stage.Flatten().ExportToString() - - -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, - ) - # The migrated Articulation reads tendon counts off its OvPhysxView; inject the mock - # view over these bindings so the metadata passthrough resolves without a real view. - object.__setattr__(articulation, "_root_view", bindings.view) - 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(): - """Tendon discovery should ignore joints that live outside the current articulation subtree.""" - articulation = _make_articulation_shell() - stage_usda = _make_articulation_root_stage_usda() - old_stage_usda = OvPhysxManager._stage_usda - OvPhysxManager._stage_usda = stage_usda - try: - articulation._process_tendons() - finally: - OvPhysxManager._stage_usda = old_stage_usda - - # the tendon is reported by its schema INSTANCE name, matching PhysX; scope leakage would - # add the identically-named instance from /World/unrelated, giving two entries - assert articulation.fixed_tendon_names == ["inst0"] - assert articulation.spatial_tendon_names == ["spatial_joint"] diff --git a/source/isaaclab_ov/test/assets/test_articulation_kernels.py b/source/isaaclab_ov/test/assets/test_articulation_kernels.py deleted file mode 100644 index b1da63857647..000000000000 --- a/source/isaaclab_ov/test/assets/test_articulation_kernels.py +++ /dev/null @@ -1,50 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Tests for OVPhysX articulation Warp kernels.""" - -import numpy as np -import pytest -import warp as wp -from isaaclab_ov.assets import kernels - - -def _selector(values: list[int], dtype: type) -> wp.array: - return wp.array(values, dtype=dtype, device="cpu") - - -@pytest.mark.parametrize("env_dtype", [wp.int32, wp.int64]) -def test_root_worker_accepts_selector_widths(env_dtype: type) -> None: - env_ids = _selector([1, 0], env_dtype) - data = wp.array( - [[11.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0], [21.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]], - dtype=wp.transformf, - device="cpu", - ) - output = wp.zeros(2, dtype=wp.transformf, device="cpu") - sim_env_ids = wp.empty(2, dtype=wp.int32, device="cpu") - kernel = kernels.set_root_link_pose_to_sim_index - if env_dtype == wp.int64: - kernel = kernels.set_root_link_pose_to_sim_index_kernel(env_ids) - - wp.launch(kernel, dim=2, inputs=[data, env_ids], outputs=[output, sim_env_ids], device="cpu") - - np.testing.assert_array_equal(output.numpy(), data.numpy()[[1, 0]]) - np.testing.assert_array_equal(sim_env_ids.numpy(), [1, 0]) - - -@pytest.mark.parametrize(("env_dtype", "item_dtype"), [(wp.int32, wp.int32), (wp.int64, wp.int64)]) -def test_item_worker_accepts_selector_widths(env_dtype: type, item_dtype: type) -> None: - env_ids = _selector([1, 0], env_dtype) - item_ids = _selector([2, 0], item_dtype) - data = wp.array([[11.0, 12.0], [21.0, 22.0]], dtype=wp.float32, device="cpu") - output = wp.full((2, 3), value=-1.0, dtype=wp.float32, device="cpu") - kernel = kernels.write_2d_data_to_buffer_with_indices - if env_dtype != wp.int32 or item_dtype != wp.int32: - kernel = kernels.write_2d_data_to_buffer_with_indices_kernel(env_ids, item_ids) - - wp.launch(kernel, dim=(2, 2), inputs=[data, env_ids, item_ids], outputs=[output], device="cpu") - - np.testing.assert_array_equal(output.numpy(), [[22.0, -1.0, 21.0], [12.0, -1.0, 11.0]]) diff --git a/source/isaaclab_ov/test/assets/test_deformable_object.py b/source/isaaclab_ov/test/assets/test_deformable_object.py index 75453d4e0cc2..3e46afa80da8 100644 --- a/source/isaaclab_ov/test/assets/test_deformable_object.py +++ b/source/isaaclab_ov/test/assets/test_deformable_object.py @@ -6,15 +6,13 @@ # ignore private usage of variables warning # pyright: reportPrivateUsage=none -"""Real-backend tests for the OVPhysX deformable object.""" +"""OVPhysX deformable object tests: host-side asset shells first, then real-backend scenes.""" from __future__ import annotations -import multiprocessing -import queue +import re import sys -import traceback -from typing import Any +from types import SimpleNamespace import ovphysx.types # noqa: F401 import pytest @@ -22,7 +20,12 @@ import warp as wp from flaky import flaky from isaaclab_ov import tensor_types as TT # noqa: E402 +from isaaclab_ov.assets import DeformableObject as OvPhysxDeformableObject # noqa: E402 +from isaaclab_ov.assets.deformable_object.deformable_object_data import DeformableObjectData # noqa: E402 +from isaaclab_ov.assets.deformable_object.kernels import vec6f # noqa: E402 +from isaaclab_ov.assets.deformable_object.views import OvPhysxDeformableBodyView # noqa: E402 from isaaclab_ov.physics import OvPhysxCfg, OvPhysxManager # noqa: E402 +from isaaclab_ov.physics.ovphysx_manager import OvPhysxSceneDataBackend # noqa: E402 from isaaclab_physx.sim.schemas import PhysxRigidBodyCfg # noqa: E402 from isaaclab_physx.sim.spawners.materials import PhysxDeformableBodyMaterialCfg # noqa: E402 @@ -130,42 +133,394 @@ def _assert_rest_positions_match_authored( ) -def _run_cpu_deformable_initialization(result_queue: Any) -> None: - """Run the CPU initialization contract in a fresh spawned process.""" - try: - with _ovphysx_sim_context(device="cpu") as sim: - deformable = _generate_deformable_scene(pre_tetrahedralized_deformable_spawn_cfg(), num_objects=5) - assert sys.getrefcount(deformable) < 10 - try: - sim.reset() - except RuntimeError as error: - result = ("runtime_error", str(error), deformable.is_initialized) - else: - result = ("no_error", "", deformable.is_initialized) - except BaseException: - result = ("child_error", traceback.format_exc(), None) - result_queue.put(result) +## +# Host-side units on asset shells over a fake body view. They need no simulation and run first. +## + + +@pytest.fixture +def host_shell(monkeypatch): + """Run an asset shell on CPU with a scene-data backend, without a simulation.""" + monkeypatch.setattr(OvPhysxManager, "_scene_data_backend", OvPhysxSceneDataBackend()) + with wp.ScopedDevice("cpu"): + yield + + +class _FakeBodyView: + """Minimal generic OVPhysX view used by the asset shell.""" + + def __init__(self, num_instances: int, num_vertices: int) -> None: + self.count = num_instances + self.max_simulation_nodes_per_body = num_vertices + self.max_simulation_elements_per_body = 2 + self.max_collision_elements_per_body = 3 + self.max_collision_nodes_per_body = 0 + self.positions = wp.zeros((num_instances, num_vertices, 3), dtype=wp.float32, device="cpu") + self.velocities = wp.zeros((num_instances, num_vertices, 3), dtype=wp.float32, device="cpu") + self.targets = wp.zeros((num_instances, num_vertices, 4), dtype=wp.float32, device="cpu") + self.position_reads = 0 + self.velocity_reads = 0 + self.position_write_count = 0 + self.velocity_write_count = 0 + + def binding_for(self, tensor_type): + if tensor_type in (TT.DEFORMABLE_SIM_ELEMENT_INDICES, TT.SURFACE_DEFORMABLE_SIM_ELEMENT_INDICES): + shape = (self.count, self.max_simulation_elements_per_body, 4) + elif tensor_type == TT.DEFORMABLE_COLLISION_ELEMENT_INDICES: + shape = (self.count, self.max_collision_elements_per_body, 4) + elif tensor_type == TT.DEFORMABLE_SIM_KINEMATIC_TARGET: + shape = (self.count, self.max_simulation_nodes_per_body, 4) + else: + shape = (self.count, self.max_simulation_nodes_per_body, 3) + return SimpleNamespace(shape=shape, count=self.count) + + def read_into(self, tensor_type, values: wp.array) -> None: + if tensor_type in (TT.DEFORMABLE_SIM_NODAL_POSITION, TT.SURFACE_DEFORMABLE_SIM_POSITION): + self.position_reads += 1 + wp.copy(values, self.positions) + elif tensor_type in (TT.DEFORMABLE_SIM_NODAL_VELOCITY, TT.SURFACE_DEFORMABLE_SIM_VELOCITY): + self.velocity_reads += 1 + wp.copy(values, self.velocities) + else: + raise AssertionError(f"Unexpected tensor read: {tensor_type}") + + def get_attribute(self, tensor_type) -> wp.array: + if tensor_type == TT.DEFORMABLE_SIM_KINEMATIC_TARGET: + return self.targets + if tensor_type in (TT.DEFORMABLE_SIM_NODAL_POSITION, TT.SURFACE_DEFORMABLE_SIM_POSITION): + return self.positions + if tensor_type in (TT.DEFORMABLE_SIM_NODAL_VELOCITY, TT.SURFACE_DEFORMABLE_SIM_VELOCITY): + return self.velocities + binding = self.binding_for(tensor_type) + return wp.zeros(binding.shape, dtype=wp.int32, device="cpu") + + def set_attribute( + self, + tensor_type, + values: wp.array, + indices: wp.array(dtype=wp.int32) | None = None, + mask: wp.array(dtype=wp.bool) | None = None, + ) -> None: + assert indices is None or indices.is_contiguous # Required by the OVPhysX DLPack binding. + self.last_indices = indices + if tensor_type in (TT.DEFORMABLE_SIM_NODAL_POSITION, TT.SURFACE_DEFORMABLE_SIM_POSITION): + self.position_write_count += 1 + elif tensor_type in (TT.DEFORMABLE_SIM_NODAL_VELOCITY, TT.SURFACE_DEFORMABLE_SIM_VELOCITY): + self.velocity_write_count += 1 + elif tensor_type != TT.DEFORMABLE_SIM_KINEMATIC_TARGET: + raise AssertionError(f"Unexpected tensor write: {tensor_type}") + + +class _FakeVisualizer: + """Capture marker positions from the debug visualization callback.""" + + def __init__(self) -> None: + self.positions: torch.Tensor | None = None + + def visualize(self, positions: torch.Tensor) -> None: + self.positions = positions + + +def _make_asset_shell( + *, + deformable_type: str, + num_instances: int = 2, + num_vertices: int = 4, +) -> OvPhysxDeformableObject: + asset = object.__new__(OvPhysxDeformableObject) + asset._device = "cpu" + asset._check_shapes = True + asset._DTYPE_TO_TORCH_TRAILING_DIMS = {**asset._DTYPE_TO_TORCH_TRAILING_DIMS, vec6f: (6,)} + asset._deformable_type = deformable_type + if deformable_type == "volume": + asset._sim_nodal_position_type = TT.DEFORMABLE_SIM_NODAL_POSITION + asset._sim_nodal_velocity_type = TT.DEFORMABLE_SIM_NODAL_VELOCITY + asset._sim_kinematic_target_type = TT.DEFORMABLE_SIM_KINEMATIC_TARGET + else: + asset._sim_nodal_position_type = TT.SURFACE_DEFORMABLE_SIM_POSITION + asset._sim_nodal_velocity_type = TT.SURFACE_DEFORMABLE_SIM_VELOCITY + asset._sim_kinematic_target_type = None + asset._root_physx_view = _FakeBodyView(num_instances, num_vertices) + asset._material_physx_view = None + asset._data = DeformableObjectData( + asset._root_physx_view, + asset._device, + position_tensor_type=asset._sim_nodal_position_type, + velocity_tensor_type=asset._sim_nodal_velocity_type, + ) + asset._ALL_INDICES = wp.array(range(num_instances), dtype=wp.int32, device=asset.device) + asset._nodal_pos_w_f32 = None + asset._nodal_vel_w_f32 = None + asset._is_initialized = True + asset._debug_vis_handle = None + asset._initialize_handle = None + asset._invalidate_initialize_handle = None + asset._prim_deletion_handle = None + return asset + + +@pytest.mark.usefixtures("host_shell") +def test_indexed_state_write_refreshes_unselected_stale_cache_rows(): + asset = _make_asset_shell(deformable_type="volume", num_instances=3, num_vertices=2) + initial_positions = torch.full((3, 2, 3), -1.0, device=asset.device) + initial_velocities = torch.full((3, 2, 3), -2.0, device=asset.device) + asset.root_view.positions = wp.from_torch(initial_positions.contiguous(), dtype=wp.float32) + asset.root_view.velocities = wp.from_torch(initial_velocities.contiguous(), dtype=wp.float32) + + asset.data.nodal_pos_w + asset.data.nodal_vel_w + + latest_positions = torch.tensor( + [ + [[10.0, 11.0, 12.0], [13.0, 14.0, 15.0]], + [[20.0, 21.0, 22.0], [23.0, 24.0, 25.0]], + [[30.0, 31.0, 32.0], [33.0, 34.0, 35.0]], + ], + device=asset.device, + ) + latest_velocities = torch.tensor( + [ + [[-10.0, -11.0, -12.0], [-13.0, -14.0, -15.0]], + [[-20.0, -21.0, -22.0], [-23.0, -24.0, -25.0]], + [[-30.0, -31.0, -32.0], [-33.0, -34.0, -35.0]], + ], + device=asset.device, + ) + asset.root_view.positions = wp.from_torch(latest_positions.contiguous(), dtype=wp.float32) + asset.root_view.velocities = wp.from_torch(latest_velocities.contiguous(), dtype=wp.float32) + asset.update(0.1) + + # Prime the derived buffers at the current timestamp so the write itself must invalidate them. + asset.data.nodal_state_w + asset.data.root_pos_w + asset.data.root_vel_w + + selected_state = torch.cat( + ( + torch.full((1, 2, 3), 100.0, device=asset.device), + torch.full((1, 2, 3), -100.0, device=asset.device), + ), + dim=-1, + ) + asset.write_nodal_state_to_sim_index(selected_state, env_ids=[0]) + + expected_positions = latest_positions.clone() + expected_positions[0] = selected_state[0, :, :3] + expected_velocities = latest_velocities.clone() + expected_velocities[0] = selected_state[0, :, 3:] + + assert asset.data._nodal_pos_w.timestamp == asset.data._sim_timestamp + assert asset.data._nodal_vel_w.timestamp == asset.data._sim_timestamp + torch.testing.assert_close(asset.data.nodal_pos_w.torch, expected_positions) + torch.testing.assert_close(asset.data.nodal_vel_w.torch, expected_velocities) + torch.testing.assert_close( + asset.data.nodal_state_w.torch, torch.cat((expected_positions, expected_velocities), dim=-1) + ) + torch.testing.assert_close(asset.data.root_pos_w.torch, expected_positions.mean(dim=1)) + torch.testing.assert_close(asset.data.root_vel_w.torch, expected_velocities.mean(dim=1)) +@pytest.mark.usefixtures("host_shell") @pytest.mark.parametrize( - "num_objects, material_path", + ("property_name", "write_method_name", "simulator_attribute", "command_value"), [ - (2, None), - # A shared absolute material whose path only textually prefixes a sibling asset path. - (2, "/World/Table_0/ObjectSiblingMaterial"), + ("nodal_pos_w", "write_nodal_pos_to_sim_index", "positions", 100.0), + ("nodal_vel_w", "write_nodal_velocity_to_sim_index", "velocities", -100.0), ], ) +def test_indexed_full_data_write_preserves_retained_aliased_edits( + property_name: str, write_method_name: str, simulator_attribute: str, command_value: float +) -> None: + """A retained public buffer preserves selected edits while stale rows hydrate.""" + asset = _make_asset_shell(deformable_type="volume", num_instances=3, num_vertices=2) + initial = torch.full((3, 2, 3), -1.0, device=asset.device) + setattr(asset.root_view, simulator_attribute, wp.from_torch(initial.contiguous(), dtype=wp.float32)) + retained = getattr(asset.data, property_name) + retained_torch = retained.torch + + latest = torch.arange(18, dtype=torch.float32, device=asset.device).reshape(3, 2, 3) + setattr(asset.root_view, simulator_attribute, wp.from_torch(latest.contiguous(), dtype=wp.float32)) + asset.update(0.1) + retained_torch[1].fill_(command_value) + + getattr(asset, write_method_name)(retained, env_ids=[1], full_data=True) + + expected = latest.clone() + expected[1].fill_(command_value) + assert getattr(asset.data, property_name) is retained + torch.testing.assert_close(retained.torch, expected) + + +@pytest.mark.usefixtures("host_shell") +@pytest.mark.parametrize( + ("property_name", "write_method_name", "simulator_attribute", "command_value"), + [ + ("nodal_pos_w", "write_nodal_pos_to_sim_index", "positions", 100.0), + ("nodal_vel_w", "write_nodal_velocity_to_sim_index", "velocities", -100.0), + ], +) +@pytest.mark.parametrize("env_ids", [slice(1, 2), slice(None, None, 2)], ids=["contiguous", "strided"]) +def test_indexed_partial_write_preserves_retained_aliased_slice( + property_name: str, write_method_name: str, simulator_attribute: str, command_value: float, env_ids: slice +) -> None: + """A retained selected slice survives hydration of stale unselected rows.""" + asset = _make_asset_shell(deformable_type="volume", num_instances=3, num_vertices=2) + initial = torch.full((3, 2, 3), -1.0, device=asset.device) + setattr(asset.root_view, simulator_attribute, wp.from_torch(initial.contiguous(), dtype=wp.float32)) + retained = getattr(asset.data, property_name) + + latest = torch.arange(18, dtype=torch.float32, device=asset.device).reshape(3, 2, 3) + setattr(asset.root_view, simulator_attribute, wp.from_torch(latest.contiguous(), dtype=wp.float32)) + asset.update(0.1) + selected = retained.torch[env_ids] + selected.fill_(command_value) + + getattr(asset, write_method_name)(selected, env_ids=env_ids) + # Only the selected environments are written back to the simulator. + assert asset.root_view.last_indices.numpy().tolist() == list(range(3)[env_ids]) + + expected = latest.clone() + expected[env_ids].fill_(command_value) + assert getattr(asset.data, property_name) is retained + torch.testing.assert_close(retained.torch, expected) + + +@pytest.mark.usefixtures("host_shell") +def test_full_overwrite_stale_cache_does_not_read_simulator() -> None: + asset = _make_asset_shell(deformable_type="volume", num_instances=3, num_vertices=2) + + asset.data.nodal_pos_w + asset.data.nodal_vel_w + position_reads = asset.root_view.position_reads + velocity_reads = asset.root_view.velocity_reads + asset.update(0.1) + + full_state = torch.arange(36, dtype=torch.float32, device=asset.device).reshape(3, 2, 6) + asset.write_nodal_state_to_sim_index(full_state, full_data=True) + + assert asset.root_view.position_reads == position_reads + assert asset.root_view.velocity_reads == velocity_reads + torch.testing.assert_close(asset.data.nodal_state_w.torch, full_state) + + +@pytest.mark.usefixtures("host_shell") +def test_malformed_state_write_fails_before_mutating_or_writing(): + asset = _make_asset_shell(deformable_type="volume", num_instances=2, num_vertices=4) + original_positions = asset.data.nodal_pos_w.torch.clone() + original_velocities = asset.data.nodal_vel_w.torch.clone() + malformed_state = torch.ones((1, 4, 7), device=asset.device) + + with pytest.raises(AssertionError, match="nodal_state.*Shape mismatch"): + asset.write_nodal_state_to_sim_index(malformed_state, env_ids=[1]) + + torch.testing.assert_close(asset.data.nodal_pos_w.torch, original_positions) + torch.testing.assert_close(asset.data.nodal_vel_w.torch, original_velocities) + assert asset.root_view.position_write_count == 0 + assert asset.root_view.velocity_write_count == 0 + + +@pytest.mark.usefixtures("host_shell") +def test_lazy_root_means_refresh_in_place_after_update(): + root_view = _FakeBodyView(num_instances=2, num_vertices=2) + root_view.positions = wp.array( + [[[0.0, 0.0, 0.0], [2.0, 4.0, 6.0]], [[1.0, 3.0, 5.0], [3.0, 5.0, 7.0]]], + dtype=wp.float32, + ) + data = DeformableObjectData( + root_view, + "cpu", + position_tensor_type=TT.DEFORMABLE_SIM_NODAL_POSITION, + velocity_tensor_type=TT.DEFORMABLE_SIM_NODAL_VELOCITY, + ) + + root_pos = data.root_pos_w + torch.testing.assert_close(root_pos.torch, torch.tensor([[1.0, 2.0, 3.0], [2.0, 4.0, 6.0]])) + assert root_view.position_reads == 1 + + root_view.positions = wp.array( + [[[2.0, 2.0, 2.0], [4.0, 4.0, 4.0]], [[6.0, 6.0, 6.0], [8.0, 8.0, 8.0]]], + dtype=wp.float32, + ) + data.update(0.1) + + assert data.root_pos_w is root_pos + torch.testing.assert_close(root_pos.torch, torch.tensor([[3.0, 3.0, 3.0], [7.0, 7.0, 7.0]])) + assert root_view.position_reads == 2 + + +@pytest.mark.usefixtures("host_shell") +def test_surface_debug_visualization_uses_below_ground_sentinel(): + asset = _make_asset_shell(deformable_type="surface") + asset.target_visualizer = _FakeVisualizer() + + asset._debug_vis_callback(None) + + torch.testing.assert_close(asset.target_visualizer.positions, torch.tensor([[0.0, 0.0, -10.0]])) + + +class _MixedTopologyBinding: + """Padded simulation connectivity with five nodes in one body and four in another.""" + + _SHAPES = { + TT.DEFORMABLE_SIM_NODAL_POSITION: (2, 5, 3), + TT.DEFORMABLE_SIM_ELEMENT_INDICES: (2, 2, 4), + TT.DEFORMABLE_COLLISION_ELEMENT_INDICES: (2, 6, 4), + } + + def __init__(self, tensor_type: int): + self._tensor_type = tensor_type + self.shape = self._SHAPES[tensor_type] + float_dtype = tensor_type == TT.DEFORMABLE_SIM_NODAL_POSITION + self.dtype = SimpleNamespace(code=2 if float_dtype else 0, bits=32, lanes=1) + self.count = self.shape[0] + self.prim_paths = [f"/World/env_{index}/Soft" for index in range(self.count)] + + def read(self, values: wp.array) -> None: + if self._tensor_type == TT.DEFORMABLE_SIM_ELEMENT_INDICES: + connectivity = [[[0, 1, 2, 3], [1, 2, 3, 4]], [[0, 1, 2, 3], [0, 0, 0, 0]]] + wp.copy(values, wp.array(connectivity, dtype=wp.int32, device="cpu")) + elif self._tensor_type == TT.DEFORMABLE_COLLISION_ELEMENT_INDICES: + wp.copy(values, wp.full(self.shape, value=4, dtype=wp.int32, device="cpu")) + + +@pytest.mark.usefixtures("host_shell") +def test_volume_view_rejects_mixed_simulation_node_counts(): + physx = SimpleNamespace( + create_tensor_binding=lambda *, tensor_type, pattern=None: _MixedTopologyBinding(tensor_type) + ) + with pytest.raises(ValueError, match=r"uniform simulation-node counts.*\[5, 4\]"): + OvPhysxDeformableBodyView( + physx, + pattern="/World/env_*/Soft", + device="cpu", + tensor_types=list(_MixedTopologyBinding._SHAPES), + eager=True, + simulation_nodal_position_type=TT.DEFORMABLE_SIM_NODAL_POSITION, + simulation_element_indices_type=TT.DEFORMABLE_SIM_ELEMENT_INDICES, + collision_element_indices_type=TT.DEFORMABLE_COLLISION_ELEMENT_INDICES, + ) + + +## +# Real-backend tests. +## + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="OVPhysX deformables require CUDA") -def test_initialization(num_objects: int, material_path: str | None): - """Test volume deformable initialization, optional and shared materials, and public buffer shapes.""" +def test_initialization_with_shared_material(): + """Test volume deformable initialization with a shared material, and public buffer shapes. + + The shared absolute material path only textually prefixes a sibling asset path. The no-material case is covered + by :func:`test_set_nodal_state_with_applied_transform`. + """ + num_objects, material_path = 2, "/World/Table_0/ObjectSiblingMaterial" with _ovphysx_sim_context(device="cuda:0") as sim: deformable = _generate_deformable_scene( pre_tetrahedralized_deformable_spawn_cfg(material_path=material_path), num_objects=num_objects ) - if material_path is not None: - # A same-named material under the other table must not be matched by prefix expansion. - distractor_cfg = PhysxDeformableBodyMaterialCfg() - distractor_cfg.func("/World/Table_1/ObjectSiblingMaterial", distractor_cfg) + # A same-named material under the other table must not be matched by prefix expansion. + distractor_cfg = PhysxDeformableBodyMaterialCfg() + distractor_cfg.func("/World/Table_1/ObjectSiblingMaterial", distractor_cfg) assert sys.getrefcount(deformable) < 10 sim.reset() @@ -174,11 +529,8 @@ def test_initialization(num_objects: int, material_path: str | None): assert deformable.num_instances == num_objects assert deformable.num_bodies == 1 assert deformable.root_view.count == num_objects - if material_path is None: - assert deformable.material_physx_view is None - else: - assert deformable.material_physx_view is not None - assert deformable.material_physx_view.count == 1 + assert deformable.material_physx_view is not None + assert deformable.material_physx_view.count == 1 assert deformable.data.nodal_state_w.torch.shape == ( num_objects, deformable.max_sim_vertices_per_body, @@ -199,41 +551,29 @@ def test_initialization(num_objects: int, material_path: str | None): def test_initialization_on_device_cpu(): - """Test that OVPhysX deformable initialization rejects a CPU simulation.""" - context = multiprocessing.get_context("spawn") - result_queue = context.Queue() - process = context.Process(target=_run_cpu_deformable_initialization, args=(result_queue,)) - process.start() - process.join(timeout=30.0) - - if process.is_alive(): - process.terminate() - process.join() - pytest.fail("CPU deformable initialization child process timed out.") - assert process.exitcode == 0 - - try: - result_kind, message, is_initialized = result_queue.get(timeout=5.0) - except queue.Empty: - pytest.fail("CPU deformable initialization child process returned no result.") - finally: - result_queue.close() - result_queue.join_thread() - - assert result_kind == "runtime_error", message - assert message == "OVPhysX deformable tensors require a CUDA simulation device; received 'cpu'." - assert is_initialized is False + """Test that OVPhysX deformable initialization rejects a CPU simulation. + + CPU and CUDA OVPhysX simulations can follow each other in one process, so this runs beside the CUDA tests. + """ + message = "OVPhysX deformable tensors require a CUDA simulation device; received 'cpu'." + with _ovphysx_sim_context(device="cpu") as sim: + deformable = _generate_deformable_scene(pre_tetrahedralized_deformable_spawn_cfg(), num_objects=5) + assert sys.getrefcount(deformable) < 10 + with pytest.raises(RuntimeError, match=f"^{re.escape(message)}$"): + sim.reset() + assert deformable.is_initialized is False @flaky(max_runs=3, min_passes=1) @pytest.mark.skipif(not torch.cuda.is_available(), reason="OVPhysX deformables require CUDA") def test_set_nodal_state_with_applied_transform(): - """Test combined nodal state writes after applying rigid transforms.""" + """Test combined nodal state writes after applying rigid transforms, on deformables without a material.""" with _ovphysx_sim_context(device="cuda:0", gravity_enabled=False) as sim: - deformable = _generate_deformable_scene(pre_tetrahedralized_deformable_spawn_cfg(), num_objects=2) + deformable = _generate_deformable_scene(pre_tetrahedralized_deformable_spawn_cfg(material_path=None)) sim.reset() + assert deformable.material_physx_view is None - for _ in range(5): + for _ in range(2): nodal_state = deformable.data.default_nodal_state_w.torch.clone() mean_nodal_pos_default = nodal_state[..., :3].mean(dim=1) @@ -248,7 +588,7 @@ def test_set_nodal_state_with_applied_transform(): deformable.write_nodal_state_to_sim_index(nodal_state) deformable.reset() - for _ in range(50): + for _ in range(10): sim.step() deformable.update(sim.cfg.dt) @@ -266,7 +606,7 @@ def test_set_kinematic_targets(): deformable.root_view.get_attribute(TT.DEFORMABLE_SIM_KINEMATIC_TARGET) ).clone() - for _ in range(5): + for _ in range(2): deformable.write_nodal_state_to_sim_index(deformable.data.default_nodal_state_w.torch) default_root_pos = deformable.data.default_nodal_state_w.torch[..., :3].mean(dim=1) deformable.reset() @@ -278,7 +618,7 @@ def test_set_kinematic_targets(): nodal_kinematic_targets[0:1], env_ids=torch.tensor([0], device=sim.device) ) - for _ in range(20): + for _ in range(10): sim.step() deformable.update(sim.cfg.dt) diff --git a/source/isaaclab_ov/test/assets/test_deformable_object_helpers.py b/source/isaaclab_ov/test/assets/test_deformable_object_helpers.py deleted file mode 100644 index e60b62e7a7ae..000000000000 --- a/source/isaaclab_ov/test/assets/test_deformable_object_helpers.py +++ /dev/null @@ -1,345 +0,0 @@ -# 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 - -"""Unit tests for the kitless OVPhysX deformable object asset.""" - -from __future__ import annotations - -import importlib -from types import SimpleNamespace - -import pytest - -pytest.importorskip("ovphysx.types", reason="ovphysx wheel not installed") - -import torch # noqa: E402 -import warp as wp # noqa: E402 -from isaaclab_ov import tensor_types as TT # noqa: E402 -from isaaclab_ov.assets.deformable_object.deformable_object import DeformableObject # noqa: E402 -from isaaclab_ov.assets.deformable_object.deformable_object_data import ( # noqa: E402 - DeformableObjectData, -) -from isaaclab_ov.assets.deformable_object.kernels import vec6f # noqa: E402 -from isaaclab_ov.physics.ovphysx_manager import OvPhysxManager, OvPhysxSceneDataBackend # noqa: E402 - -wp.init() -wp.set_device("cpu") - - -@pytest.fixture(autouse=True) -def scene_data_backend(monkeypatch): - monkeypatch.setattr(OvPhysxManager, "_scene_data_backend", OvPhysxSceneDataBackend()) - - -class _FakeBodyView: - """Minimal generic OVPhysX view used by the asset shell.""" - - def __init__(self, num_instances: int, num_vertices: int) -> None: - self.count = num_instances - self.max_simulation_nodes_per_body = num_vertices - self.max_simulation_elements_per_body = 2 - self.max_collision_elements_per_body = 3 - self.max_collision_nodes_per_body = 0 - self.positions = wp.zeros((num_instances, num_vertices, 3), dtype=wp.float32, device="cpu") - self.velocities = wp.zeros((num_instances, num_vertices, 3), dtype=wp.float32, device="cpu") - self.targets = wp.zeros((num_instances, num_vertices, 4), dtype=wp.float32, device="cpu") - self.position_reads = 0 - self.velocity_reads = 0 - self.position_write_count = 0 - self.velocity_write_count = 0 - - def binding_for(self, tensor_type): - if tensor_type in (TT.DEFORMABLE_SIM_ELEMENT_INDICES, TT.SURFACE_DEFORMABLE_SIM_ELEMENT_INDICES): - shape = (self.count, self.max_simulation_elements_per_body, 4) - elif tensor_type == TT.DEFORMABLE_COLLISION_ELEMENT_INDICES: - shape = (self.count, self.max_collision_elements_per_body, 4) - elif tensor_type == TT.DEFORMABLE_SIM_KINEMATIC_TARGET: - shape = (self.count, self.max_simulation_nodes_per_body, 4) - else: - shape = (self.count, self.max_simulation_nodes_per_body, 3) - return SimpleNamespace(shape=shape, count=self.count) - - def read_into(self, tensor_type, values: wp.array) -> None: - if tensor_type in (TT.DEFORMABLE_SIM_NODAL_POSITION, TT.SURFACE_DEFORMABLE_SIM_POSITION): - self.position_reads += 1 - wp.copy(values, self.positions) - elif tensor_type in (TT.DEFORMABLE_SIM_NODAL_VELOCITY, TT.SURFACE_DEFORMABLE_SIM_VELOCITY): - self.velocity_reads += 1 - wp.copy(values, self.velocities) - else: - raise AssertionError(f"Unexpected tensor read: {tensor_type}") - - def get_attribute(self, tensor_type) -> wp.array: - if tensor_type == TT.DEFORMABLE_SIM_KINEMATIC_TARGET: - return self.targets - if tensor_type in (TT.DEFORMABLE_SIM_NODAL_POSITION, TT.SURFACE_DEFORMABLE_SIM_POSITION): - return self.positions - if tensor_type in (TT.DEFORMABLE_SIM_NODAL_VELOCITY, TT.SURFACE_DEFORMABLE_SIM_VELOCITY): - return self.velocities - binding = self.binding_for(tensor_type) - return wp.zeros(binding.shape, dtype=wp.int32, device="cpu") - - def set_attribute( - self, - tensor_type, - values: wp.array, - indices: wp.array(dtype=wp.int32) | None = None, - mask: wp.array(dtype=wp.bool) | None = None, - ) -> None: - assert indices is None or indices.is_contiguous # Required by the OVPhysX DLPack binding. - self.last_indices = indices - if tensor_type in (TT.DEFORMABLE_SIM_NODAL_POSITION, TT.SURFACE_DEFORMABLE_SIM_POSITION): - self.position_write_count += 1 - elif tensor_type in (TT.DEFORMABLE_SIM_NODAL_VELOCITY, TT.SURFACE_DEFORMABLE_SIM_VELOCITY): - self.velocity_write_count += 1 - elif tensor_type != TT.DEFORMABLE_SIM_KINEMATIC_TARGET: - raise AssertionError(f"Unexpected tensor write: {tensor_type}") - - -class _FakeVisualizer: - """Capture marker positions from the debug visualization callback.""" - - def __init__(self) -> None: - self.positions: torch.Tensor | None = None - - def visualize(self, positions: torch.Tensor) -> None: - self.positions = positions - - -def _make_asset_shell( - *, - deformable_type: str, - num_instances: int = 2, - num_vertices: int = 4, -) -> DeformableObject: - asset = object.__new__(DeformableObject) - asset._device = "cpu" - asset._check_shapes = True - asset._DTYPE_TO_TORCH_TRAILING_DIMS = {**asset._DTYPE_TO_TORCH_TRAILING_DIMS, vec6f: (6,)} - asset._deformable_type = deformable_type - if deformable_type == "volume": - asset._sim_nodal_position_type = TT.DEFORMABLE_SIM_NODAL_POSITION - asset._sim_nodal_velocity_type = TT.DEFORMABLE_SIM_NODAL_VELOCITY - asset._sim_kinematic_target_type = TT.DEFORMABLE_SIM_KINEMATIC_TARGET - else: - asset._sim_nodal_position_type = TT.SURFACE_DEFORMABLE_SIM_POSITION - asset._sim_nodal_velocity_type = TT.SURFACE_DEFORMABLE_SIM_VELOCITY - asset._sim_kinematic_target_type = None - asset._root_physx_view = _FakeBodyView(num_instances, num_vertices) - asset._material_physx_view = None - asset._data = DeformableObjectData( - asset._root_physx_view, - asset._device, - position_tensor_type=asset._sim_nodal_position_type, - velocity_tensor_type=asset._sim_nodal_velocity_type, - ) - asset._ALL_INDICES = wp.array(range(num_instances), dtype=wp.int32, device=asset.device) - asset._nodal_pos_w_f32 = None - asset._nodal_vel_w_f32 = None - asset._is_initialized = True - asset._debug_vis_handle = None - asset._initialize_handle = None - asset._invalidate_initialize_handle = None - asset._prim_deletion_handle = None - return asset - - -def test_indexed_state_write_refreshes_unselected_stale_cache_rows(): - asset = _make_asset_shell(deformable_type="volume", num_instances=3, num_vertices=2) - initial_positions = torch.full((3, 2, 3), -1.0, device=asset.device) - initial_velocities = torch.full((3, 2, 3), -2.0, device=asset.device) - asset.root_view.positions = wp.from_torch(initial_positions.contiguous(), dtype=wp.float32) - asset.root_view.velocities = wp.from_torch(initial_velocities.contiguous(), dtype=wp.float32) - - asset.data.nodal_pos_w - asset.data.nodal_vel_w - - latest_positions = torch.tensor( - [ - [[10.0, 11.0, 12.0], [13.0, 14.0, 15.0]], - [[20.0, 21.0, 22.0], [23.0, 24.0, 25.0]], - [[30.0, 31.0, 32.0], [33.0, 34.0, 35.0]], - ], - device=asset.device, - ) - latest_velocities = torch.tensor( - [ - [[-10.0, -11.0, -12.0], [-13.0, -14.0, -15.0]], - [[-20.0, -21.0, -22.0], [-23.0, -24.0, -25.0]], - [[-30.0, -31.0, -32.0], [-33.0, -34.0, -35.0]], - ], - device=asset.device, - ) - asset.root_view.positions = wp.from_torch(latest_positions.contiguous(), dtype=wp.float32) - asset.root_view.velocities = wp.from_torch(latest_velocities.contiguous(), dtype=wp.float32) - asset.update(0.1) - - # Prime the derived buffers at the current timestamp so the write itself must invalidate them. - asset.data.nodal_state_w - asset.data.root_pos_w - asset.data.root_vel_w - - selected_state = torch.cat( - ( - torch.full((1, 2, 3), 100.0, device=asset.device), - torch.full((1, 2, 3), -100.0, device=asset.device), - ), - dim=-1, - ) - asset.write_nodal_state_to_sim_index(selected_state, env_ids=[0]) - - expected_positions = latest_positions.clone() - expected_positions[0] = selected_state[0, :, :3] - expected_velocities = latest_velocities.clone() - expected_velocities[0] = selected_state[0, :, 3:] - - assert asset.data._nodal_pos_w.timestamp == asset.data._sim_timestamp - assert asset.data._nodal_vel_w.timestamp == asset.data._sim_timestamp - torch.testing.assert_close(asset.data.nodal_pos_w.torch, expected_positions) - torch.testing.assert_close(asset.data.nodal_vel_w.torch, expected_velocities) - torch.testing.assert_close( - asset.data.nodal_state_w.torch, torch.cat((expected_positions, expected_velocities), dim=-1) - ) - torch.testing.assert_close(asset.data.root_pos_w.torch, expected_positions.mean(dim=1)) - torch.testing.assert_close(asset.data.root_vel_w.torch, expected_velocities.mean(dim=1)) - - -@pytest.mark.parametrize( - ("property_name", "write_method_name", "simulator_attribute", "command_value"), - [ - ("nodal_pos_w", "write_nodal_pos_to_sim_index", "positions", 100.0), - ("nodal_vel_w", "write_nodal_velocity_to_sim_index", "velocities", -100.0), - ], -) -def test_indexed_full_data_write_preserves_retained_aliased_edits( - property_name: str, write_method_name: str, simulator_attribute: str, command_value: float -) -> None: - """A retained public buffer preserves selected edits while stale rows hydrate.""" - asset = _make_asset_shell(deformable_type="volume", num_instances=3, num_vertices=2) - initial = torch.full((3, 2, 3), -1.0, device=asset.device) - setattr(asset.root_view, simulator_attribute, wp.from_torch(initial.contiguous(), dtype=wp.float32)) - retained = getattr(asset.data, property_name) - retained_torch = retained.torch - - latest = torch.arange(18, dtype=torch.float32, device=asset.device).reshape(3, 2, 3) - setattr(asset.root_view, simulator_attribute, wp.from_torch(latest.contiguous(), dtype=wp.float32)) - asset.update(0.1) - retained_torch[1].fill_(command_value) - - getattr(asset, write_method_name)(retained, env_ids=[1], full_data=True) - - expected = latest.clone() - expected[1].fill_(command_value) - assert getattr(asset.data, property_name) is retained - torch.testing.assert_close(retained.torch, expected) - - -@pytest.mark.parametrize( - ("property_name", "write_method_name", "simulator_attribute", "command_value"), - [ - ("nodal_pos_w", "write_nodal_pos_to_sim_index", "positions", 100.0), - ("nodal_vel_w", "write_nodal_velocity_to_sim_index", "velocities", -100.0), - ], -) -@pytest.mark.parametrize("env_ids", [slice(1, 2), slice(None, None, 2)], ids=["contiguous", "strided"]) -def test_indexed_partial_write_preserves_retained_aliased_slice( - property_name: str, write_method_name: str, simulator_attribute: str, command_value: float, env_ids: slice -) -> None: - """A retained selected slice survives hydration of stale unselected rows.""" - asset = _make_asset_shell(deformable_type="volume", num_instances=3, num_vertices=2) - initial = torch.full((3, 2, 3), -1.0, device=asset.device) - setattr(asset.root_view, simulator_attribute, wp.from_torch(initial.contiguous(), dtype=wp.float32)) - retained = getattr(asset.data, property_name) - - latest = torch.arange(18, dtype=torch.float32, device=asset.device).reshape(3, 2, 3) - setattr(asset.root_view, simulator_attribute, wp.from_torch(latest.contiguous(), dtype=wp.float32)) - asset.update(0.1) - selected = retained.torch[env_ids] - selected.fill_(command_value) - - getattr(asset, write_method_name)(selected, env_ids=env_ids) - # Only the selected environments are written back to the simulator. - assert asset.root_view.last_indices.numpy().tolist() == list(range(3)[env_ids]) - - expected = latest.clone() - expected[env_ids].fill_(command_value) - assert getattr(asset.data, property_name) is retained - torch.testing.assert_close(retained.torch, expected) - - -def test_full_overwrite_stale_cache_does_not_read_simulator() -> None: - asset = _make_asset_shell(deformable_type="volume", num_instances=3, num_vertices=2) - - asset.data.nodal_pos_w - asset.data.nodal_vel_w - position_reads = asset.root_view.position_reads - velocity_reads = asset.root_view.velocity_reads - asset.update(0.1) - - full_state = torch.arange(36, dtype=torch.float32, device=asset.device).reshape(3, 2, 6) - asset.write_nodal_state_to_sim_index(full_state, full_data=True) - - assert asset.root_view.position_reads == position_reads - assert asset.root_view.velocity_reads == velocity_reads - torch.testing.assert_close(asset.data.nodal_state_w.torch, full_state) - - -@pytest.mark.parametrize("trailing_dimension", [7]) -def test_malformed_state_write_fails_before_mutating_or_writing(trailing_dimension: int): - asset = _make_asset_shell(deformable_type="volume", num_instances=2, num_vertices=4) - original_positions = asset.data.nodal_pos_w.torch.clone() - original_velocities = asset.data.nodal_vel_w.torch.clone() - malformed_state = torch.ones((1, 4, trailing_dimension), device=asset.device) - - with pytest.raises(AssertionError, match="nodal_state.*Shape mismatch"): - asset.write_nodal_state_to_sim_index(malformed_state, env_ids=[1]) - - torch.testing.assert_close(asset.data.nodal_pos_w.torch, original_positions) - torch.testing.assert_close(asset.data.nodal_vel_w.torch, original_velocities) - assert asset.root_view.position_write_count == 0 - assert asset.root_view.velocity_write_count == 0 - - -def test_lazy_root_means_refresh_in_place_after_update(): - root_view = _FakeBodyView(num_instances=2, num_vertices=2) - root_view.positions = wp.array( - [[[0.0, 0.0, 0.0], [2.0, 4.0, 6.0]], [[1.0, 3.0, 5.0], [3.0, 5.0, 7.0]]], - dtype=wp.float32, - ) - data = DeformableObjectData( - root_view, - "cpu", - position_tensor_type=TT.DEFORMABLE_SIM_NODAL_POSITION, - velocity_tensor_type=TT.DEFORMABLE_SIM_NODAL_VELOCITY, - ) - - root_pos = data.root_pos_w - torch.testing.assert_close(root_pos.torch, torch.tensor([[1.0, 2.0, 3.0], [2.0, 4.0, 6.0]])) - assert root_view.position_reads == 1 - - root_view.positions = wp.array( - [[[2.0, 2.0, 2.0], [4.0, 4.0, 4.0]], [[6.0, 6.0, 6.0], [8.0, 8.0, 8.0]]], - dtype=wp.float32, - ) - data.update(0.1) - - assert data.root_pos_w is root_pos - torch.testing.assert_close(root_pos.torch, torch.tensor([[3.0, 3.0, 3.0], [7.0, 7.0, 7.0]])) - assert root_view.position_reads == 2 - - -def test_surface_debug_visualization_uses_below_ground_sentinel(): - asset = _make_asset_shell(deformable_type="surface") - asset.target_visualizer = _FakeVisualizer() - - asset._debug_vis_callback(None) - - torch.testing.assert_close(asset.target_visualizer.positions, torch.tensor([[0.0, 0.0, -10.0]])) - - -def test_factory_export_is_present(): - assets_module = importlib.import_module("isaaclab_ov.assets") - - assert assets_module.DeformableObject is DeformableObject diff --git a/source/isaaclab_ov/test/assets/test_deformable_views.py b/source/isaaclab_ov/test/assets/test_deformable_views.py deleted file mode 100644 index f3f9d1c1a98a..000000000000 --- a/source/isaaclab_ov/test/assets/test_deformable_views.py +++ /dev/null @@ -1,119 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Tests for the thin OVPhysX deformable body view.""" - -from __future__ import annotations - -from dataclasses import dataclass - -import pytest - -pytest.importorskip("ovphysx.types", reason="ovphysx wheel not installed") - -import warp as wp # noqa: E402 -from isaaclab_ov.assets.deformable_object.views import OvPhysxDeformableBodyView # noqa: E402 -from ovphysx.types import TensorType # noqa: E402 - -wp.init() -wp.set_device("cpu") - - -@dataclass(frozen=True) -class _FakeDtype: - code: int - bits: int - lanes: int = 1 - - -_FLOAT32 = _FakeDtype(code=2, bits=32) -_INT32 = _FakeDtype(code=0, bits=32) - -_SPECS = { - TensorType.DEFORMABLE_SIM_NODAL_POSITION: ((2, 5, 3), wp.float32), - TensorType.DEFORMABLE_SIM_ELEMENT_INDICES: ((2, 2, 4), wp.int32), - TensorType.DEFORMABLE_COLLISION_ELEMENT_INDICES: ((2, 6, 4), wp.int32), -} - - -class _FakeBinding: - """Minimal tensor binding exposing dimensions and OVPhysX metadata.""" - - def __init__(self, tensor_type: TensorType): - self._tensor_type = tensor_type - self.shape, dtype = _SPECS[tensor_type] - self.dtype = _FLOAT32 if dtype == wp.float32 else _INT32 - self.count = self.shape[0] - self.prim_paths = [f"/World/env_{index}/Soft" for index in range(self.count)] - - def read(self, values: wp.array) -> None: - if self._tensor_type == TensorType.DEFORMABLE_SIM_ELEMENT_INDICES: - connectivity = wp.array( - [ - [[0, 1, 2, 3], [1, 2, 3, 4]], - [[0, 1, 2, 3], [1, 2, 3, 4]], - ], - dtype=wp.int32, - device="cpu", - ) - wp.copy(values, connectivity) - elif self._tensor_type == TensorType.DEFORMABLE_COLLISION_ELEMENT_INDICES: - wp.copy(values, wp.full(self.shape, value=4, dtype=wp.int32, device="cpu")) - - def write(self, values: wp.array, indices: wp.array | None = None, mask: wp.array | None = None) -> None: - pass - - -class _FakePhysX: - """Fake PhysX instance returning the requested binding type.""" - - def create_tensor_binding(self, *, tensor_type: TensorType, pattern: str | None = None) -> _FakeBinding: - return _FakeBinding(tensor_type) - - -class _MixedTopologyBinding(_FakeBinding): - """Simulation connectivity with five nodes in one body and four in another.""" - - def read(self, values: wp.array) -> None: - if self._tensor_type == TensorType.DEFORMABLE_SIM_ELEMENT_INDICES: - connectivity = wp.array( - [ - [[0, 1, 2, 3], [1, 2, 3, 4]], - [[0, 1, 2, 3], [0, 0, 0, 0]], - ], - dtype=wp.int32, - device="cpu", - ) - wp.copy(values, connectivity) - else: - super().read(values) - - -class _MixedTopologyPhysX(_FakePhysX): - """Fake PhysX instance exposing padded mixed-topology bindings.""" - - def create_tensor_binding(self, *, tensor_type: TensorType, pattern: str | None = None) -> _FakeBinding: - return _MixedTopologyBinding(tensor_type) - - -def test_volume_view_rejects_mixed_simulation_node_counts(): - with pytest.raises( - ValueError, - match=r"uniform simulation-node counts.*\[5, 4\]", - ): - OvPhysxDeformableBodyView( - _MixedTopologyPhysX(), - pattern="/World/env_*/Soft", - device="cpu", - tensor_types=[ - TensorType.DEFORMABLE_SIM_NODAL_POSITION, - TensorType.DEFORMABLE_SIM_ELEMENT_INDICES, - TensorType.DEFORMABLE_COLLISION_ELEMENT_INDICES, - ], - eager=True, - simulation_nodal_position_type=TensorType.DEFORMABLE_SIM_NODAL_POSITION, - simulation_element_indices_type=TensorType.DEFORMABLE_SIM_ELEMENT_INDICES, - collision_element_indices_type=TensorType.DEFORMABLE_COLLISION_ELEMENT_INDICES, - ) diff --git a/source/isaaclab_ov/test/assets/test_rigid_object.py b/source/isaaclab_ov/test/assets/test_rigid_object.py index 63f6dc712fd4..f0d2deb35bbf 100644 --- a/source/isaaclab_ov/test/assets/test_rigid_object.py +++ b/source/isaaclab_ov/test/assets/test_rigid_object.py @@ -6,110 +6,104 @@ # ignore private usage of variables warning # pyright: reportPrivateUsage=none +"""Real OVPhysX rigid-object and rigid-object-collection coverage on one module-scoped scene per device. -"""Real-backend tests for the OVPhysX RigidObject. +Each device builds one scene of locally authored cuboid pairs and ``N=2, B=3`` collections, resets it once, +and keeps it alive for every test that uses it. Each asset belongs to one test, so those tests do not depend +on each other's order. The CPU scene covers host-resident property bindings; the CUDA scene covers their +pinned-host staging. -Run via ``./scripts/run_ovphysx.sh -m pytest`` (kitless, no the Kit launcher). - -The OVPhysX runtime fixes device mode (CPU vs GPU) when the process creates -its first ``ovphysx.PhysX`` instance and cannot switch it without a process -restart. Full coverage therefore requires two separate pytest -invocations -- once with ``-k 'cpu'`` and once with ``-k 'cuda:0'``. The -``_ovphysx_skip_other_device`` autouse fixture below preempts the manager's -:exc:`RuntimeError` by ``pytest.skip``-ing on the unlocked device so -single-device runs finish cleanly. +Initialization failures need their own scenes. A second simulation context cannot start while a shared +scene is alive, so these tests must run before any shared-scene test in the same session; they are +defined first and pytest runs them before it creates the module-scoped scenes. """ from __future__ import annotations -import logging -import sys -from typing import Literal +from collections.abc import Iterator +from dataclasses import dataclass import pytest import torch import warp as wp -from flaky import flaky from pxr import UsdPhysics -from isaaclab.test.utils import test_devices - -# The OVPhysX runtime wheel is optional. Skip gracefully when it is not installed; -# CI jobs that need OVPhysX coverage install it explicitly. pytest.importorskip("ovphysx.types", reason="ovphysx wheel not installed") from isaaclab_ov import tensor_types as TT # noqa: E402 -from isaaclab_ov.assets import RigidObject # noqa: E402 +from isaaclab_ov.assets import RigidObject, RigidObjectCollection # noqa: E402 from isaaclab_ov.physics import OvPhysxCfg, OvPhysxManager # noqa: E402 import isaaclab.sim as sim_utils # noqa: E402 from isaaclab import cloner # noqa: E402 -from isaaclab.assets import AssetBaseCfg, RigidObjectCfg # noqa: E402 +from isaaclab.assets import AssetBaseCfg, RigidObjectCfg, RigidObjectCollectionCfg # noqa: E402 from isaaclab.scene import InteractiveScene, InteractiveSceneCfg # noqa: E402 -from isaaclab.sim import SimulationCfg, build_simulation_context # noqa: E402 -from isaaclab.sim.spawners import materials # noqa: E402 +from isaaclab.sim import SimulationCfg, SimulationContext, build_simulation_context # noqa: E402 +from isaaclab.test.utils import DeviceScope, test_devices # noqa: E402 from isaaclab.utils import configclass # noqa: E402 -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR # noqa: E402 -from isaaclab.utils.math import ( # noqa: E402 - combine_frame_transforms, - default_orientation, - quat_apply_inverse, - quat_inv, - quat_mul, - quat_rotate, - random_orientation, -) - -wp.init() - -pytestmark = pytest.mark.device_split - -_logger = logging.getLogger(__name__) - - -_LOCKED_DEVICE: list[str | None] = [None] -"""Device the session pins to on the first parametrized test that runs.""" - +from isaaclab.utils.math import quat_apply_inverse, quat_mul # noqa: E402 + +pytestmark = pytest.mark.integration + +_NUM_CUBES = 2 +_NUM_ENVS, _NUM_BODIES = 2, 3 + + +def _spawn_cubes(name: str, y_offset: float, **rigid_props) -> RigidObject: + """Author one pair of local cuboids.""" + for index in range(_NUM_CUBES): + sim_utils.create_prim(f"/World/{name}/Env_{index}", "Xform", translation=(2.0 * index, y_offset, 0.0)) + return RigidObject( + RigidObjectCfg( + prim_path=f"/World/{name}/Env_[^/]*/Cube", + spawn=sim_utils.CuboidCfg( + size=(0.2, 0.2, 0.2), + rigid_props=sim_utils.RigidBodyPropertiesCfg(**rigid_props), + mass_props=sim_utils.MassPropertiesCfg(mass=1.0), + collision_props=sim_utils.CollisionPropertiesCfg(), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 1.0)), + ) + ) -@pytest.fixture(autouse=True) -def _ovphysx_skip_other_device(request): - """Skip parametrized tests on the device the session is not pinned to. - See the module docstring for the wheel's process-global device-mode lock. - """ - callspec = getattr(request.node, "callspec", None) - device = callspec.params.get("device") if callspec is not None else None - if device is None: - 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 _spawn_static_colliders() -> RigidObject: + """Author a collider without a rigid body at the rigid-object path.""" + for index in range(_NUM_CUBES): + sim_utils.create_prim(f"/World/Env_{index}", "Xform", translation=(2.0 * index, 0.0, 0.0)) + return RigidObject( + RigidObjectCfg( + prim_path="/World/Env_[^/]*/Cube", + spawn=sim_utils.CuboidCfg(size=(0.1, 0.1, 0.1), collision_props=sim_utils.UsdPhysicsCollisionCfg()), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 1.0)), ) + ) -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 _spawn_collection(name: str, y_offset: float, spawn: sim_utils.SpawnerCfg | None = None) -> RigidObjectCollection: + """Author one canonical ``N=2, B=3`` local collection.""" + if spawn is None: + spawn = sim_utils.CuboidCfg( + size=(0.2, 0.2, 0.2), + rigid_props=sim_utils.RigidBodyPropertiesCfg(disable_gravity=True), + mass_props=sim_utils.MassPropertiesCfg(mass=1.0), + collision_props=sim_utils.CollisionPropertiesCfg(), + ) + for env_index in range(_NUM_ENVS): + sim_utils.create_prim(f"/World/{name}/Env_{env_index}", "Xform", translation=(3.0 * env_index, y_offset, 0.0)) + return RigidObjectCollection( + RigidObjectCollectionCfg( + rigid_objects={ + body_name: RigidObjectCfg( + prim_path=f"/World/{name}/Env_[^/]*/Object_{body_index}", + spawn=spawn, + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, float(body_index), 1.0)), + ) + for body_index, body_name in enumerate(("left", "middle", "right")) + } + ) + ) @configclass @@ -132,16 +126,37 @@ class HeterogeneousRigidSceneCfg(InteractiveSceneCfg): collision_props=sim_utils.CollisionBaseCfg(), mass_props=sim_utils.MassPropertiesCfg(mass=1.0), ), - init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 1.0)), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 0.4)), ) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("filter_collisions", [False, True]) -@pytest.mark.parametrize("support_cloning", ["native", "usd_and_native", "usd", "usd_nested"]) -def test_heterogeneous_clone_contacts(device, filter_collisions, support_cloning): - """Sources and clones from different variants contact their own support.""" - with _ovphysx_sim_context(device=device, dt=1.0 / 120.0) as sim: +_CPU, _CUDA = test_devices(DeviceScope.CPU), test_devices(DeviceScope.CUDA) +# CUDA-only scene tests keep the CPU row as a skip so that pytest groups every test of one scene device together +# and builds each module-scoped scene once. +_CUDA_SCENES = [ + device if device in _CUDA else pytest.param(device, marks=pytest.mark.skip(reason="CUDA-only path")) + for device in test_devices() +] + + +@pytest.mark.parametrize( + ("device", "support_cloning", "filter_collisions"), + # CUDA replays native clones, so it covers every cloning path. CPU serializes the full stage, one path. + [(device, "native", False) for device in _CPU] + + [ + (device, *row) + for device in _CUDA + for row in (("native", False), ("usd_and_native", True), ("usd", False), ("usd_nested", True)) + ], +) +def test_heterogeneous_clone_contacts_and_indexed_state(device, support_cloning, filter_collisions): + """Variants contact their own support, and indexed writes reach the environment-ordered bodies. + + Collision filtering is an independent scene setting, so it alternates across the cloning paths. + """ + with build_simulation_context( + device=device, sim_cfg=SimulationCfg(physics=OvPhysxCfg(), device=device, dt=1.0 / 120.0) + ) as sim: cfg = HeterogeneousRigidSceneCfg(num_envs=6, env_spacing=2.0, filter_collisions=filter_collisions) if support_cloning != "native": cfg.clone_cfg.clone_template = "/Scenes/World_{}" @@ -156,966 +171,445 @@ def test_heterogeneous_clone_contacts(device, filter_collisions, support_cloning cfg.clone_cfg.clone_combinations = [cloner.InclusionSet(assets=names) for names in (["assembly"], [])] scene = InteractiveScene(cfg) sim.reset() - for _ in range(240): + obj = scene["object"] + expected_paths = [f"{cfg.clone_cfg.clone_template.format(i)}/Object" for i in range(scene.num_envs)] + assert obj.root_view.prim_paths == expected_paths + for _ in range(60): sim.step() scene.update(sim.get_physics_dt()) - positions = scene["object"].data.root_pos_w.torch + positions = obj.data.root_pos_w.torch torch.testing.assert_close(positions[:, 2], torch.full_like(positions[:, 2], 0.2), atol=0.02, rtol=0.0) torch.testing.assert_close(positions[:, :2], scene.env_origins[:, :2], atol=0.02, rtol=0.0) - -@pytest.mark.parametrize("device", test_devices()) -def test_heterogeneous_clone_collision_isolation(device): - """Collision groups isolate overlapping environments, including both retained sources.""" - with _ovphysx_sim_context(device=device, dt=1.0 / 120.0) as sim: - # Without GPU environment-ID filtering, 2048 overlapping worlds exhaust broadphase pairs. - num_envs = 2048 if device.startswith("cuda") else 6 - scene = InteractiveScene(HeterogeneousRigidSceneCfg(num_envs=num_envs, env_spacing=0.0)) - sim.reset() - for _ in range(240): - sim.step() - scene.update(sim.get_physics_dt()) - positions = scene["object"].data.root_pos_w.torch - expected = scene.env_origins.clone() - expected[:, 2] = 0.2 - torch.testing.assert_close(positions, expected, atol=0.02, rtol=0.0) - - -@pytest.mark.parametrize("device", test_devices()) -def test_heterogeneous_clone_indexed_state(device): - """Indexed reads and writes use environment order across six interleaved variants.""" - with _ovphysx_sim_context(device=device, gravity_enabled=False) as sim: - scene = InteractiveScene(HeterogeneousRigidSceneCfg(num_envs=6, env_spacing=2.0)) - sim.reset() - obj = scene["object"] - expected_paths = [f"/World/envs/env_{i}/Object" for i in range(scene.num_envs)] - assert obj.root_view.prim_paths == expected_paths - torch.testing.assert_close(obj.data.root_pos_w.torch[:, :2], scene.env_origins[:, :2]) + # Lift one resting body; an independent exact-path binding detects a write to the wrong body. selected = torch.tensor([3], device=device) pose = obj.data.root_pose_w.torch[selected].clone() pose[:, 2] = 3.0 obj.write_root_pose_to_sim(pose, env_ids=selected) sim.step() scene.update(sim.get_physics_dt()) - # An independent exact-path binding detects a write to the wrong physical object. binding = OvPhysxManager.get_physx_instance().create_tensor_binding( prim_paths=expected_paths, tensor_type=TT.RIGID_BODY_POSE ) try: actual = torch.empty(binding.shape, device=device) binding.read(actual) - expected_height = torch.ones(scene.num_envs, device=device) - expected_height[selected] = 3.0 - torch.testing.assert_close(actual[:, 2], expected_height) - torch.testing.assert_close(obj.data.root_pos_w.torch, actual[:, :3]) finally: binding.destroy() + expected_height = torch.full((scene.num_envs,), 0.2, device=device) + expected_height[selected] = 3.0 + torch.testing.assert_close(actual[:, 2], expected_height, atol=0.02, rtol=0.0) + torch.testing.assert_close(obj.data.root_pos_w.torch, actual[:, :3]) -def generate_cubes_scene( - num_cubes: int = 1, - height=1.0, - api: Literal["none", "rigid_body"] = "rigid_body", - kinematic_enabled: bool = False, - device: str = "cuda:0", -) -> tuple[RigidObject, torch.Tensor]: - """Generate a scene with the provided number of cubes. - - Args: - num_cubes: Number of cubes to generate. - height: Height of the cubes. - api: The type of API that the cubes should have. - kinematic_enabled: Whether the cubes are kinematic. - device: Device to use for the simulation. - - Returns: - A tuple containing the rigid object representing the cubes and the origins of the cubes. - - """ - origins = torch.tensor([(i * 1.0, 0, height) for i in range(num_cubes)]).to(device) - # Create Top-level Xforms, one for each cube - for i, origin in enumerate(origins): - sim_utils.create_prim(f"/World/Table_{i}", "Xform", translation=origin) - - # Resolve spawn configuration - if api == "none": - # since no rigid body properties defined, this is just a static collider - spawn_cfg = sim_utils.CuboidCfg( - size=(0.1, 0.1, 0.1), - collision_props=sim_utils.UsdPhysicsCollisionCfg(), - ) - elif api == "rigid_body": - spawn_cfg = sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", - rigid_props=sim_utils.UsdPhysicsRigidBodyCfg(kinematic_enabled=kinematic_enabled), - ) - else: - raise ValueError(f"Unknown api: {api}") - - # Create rigid object. OVPhysX matches prim paths via fnmatch globs (not regex), - # so use ``Table_*`` rather than the PhysX ``Table_.*`` form. - cube_object_cfg = RigidObjectCfg( - prim_path="/World/Table_[^/]+/Object", - spawn=spawn_cfg, - init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, height)), - ) - cube_object = RigidObject(cfg=cube_object_cfg) - - return cube_object, origins - - -# --------------------------------------------------------------------------- -# Per-shape material helpers (friction / restitution) -# -# OVPhysX exposes per-collision-shape material as the -# ``rigid_body_shape_friction_and_restitution`` tensor binding (shape ``[N, S, 3]`` = -# static friction, dynamic friction, restitution), addressed through the OvPhysxView. -# (The PhysX backend instead uses ``root_view.get/set_material_properties`` / ``max_shapes``.) -# --------------------------------------------------------------------------- - - -def _num_shapes(cube_object) -> int: - """Per-body collision-shape count, from the material binding's ``[N, S, 3]`` shape.""" - return cube_object.root_view.binding_for(TT.RIGID_BODY_SHAPE_FRICTION_AND_RESTITUTION).shape[1] - - -def _write_shape_material(cube_object, materials_nS3: torch.Tensor) -> None: - """Write a full ``[N, S, 3]`` (static friction, dynamic friction, restitution) tensor via the view.""" - cube_object.root_view.set_attribute( - TT.RIGID_BODY_SHAPE_FRICTION_AND_RESTITUTION, wp.from_torch(materials_nS3.contiguous(), dtype=wp.float32) - ) - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_initialization_with_kinematic_enabled(num_cubes, device): - """Test that initialization for prim with kinematic flag enabled.""" - with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: - # Generate cubes scene - cube_object, origins = generate_cubes_scene(num_cubes=num_cubes, kinematic_enabled=True, device=device) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(cube_object) < 10 - - # Play sim +@pytest.mark.parametrize("device", _CUDA) +def test_heterogeneous_clone_collision_isolation(device): + """Collision groups isolate overlapping environments, including both retained sources.""" + with build_simulation_context( + device=device, sim_cfg=SimulationCfg(physics=OvPhysxCfg(), device=device, dt=1.0 / 120.0) + ) as sim: + # Without GPU environment-ID filtering, 2048 overlapping worlds exhaust broadphase pairs. + scene = InteractiveScene(HeterogeneousRigidSceneCfg(num_envs=2048, env_spacing=0.0)) sim.reset() - - # SDP bindings must be ready on CPU and GPU before any asset pose read. - provider = sim.get_scene_data_provider() - assert provider.transform_count == num_cubes - assert set(provider.backend.transform_paths) == {f"/World/Table_{i}/Object" for i in range(num_cubes)} - - # Check if object is initialized - assert cube_object.is_initialized - assert len(cube_object.body_names) == 1 - - # Check buffers that exists and have correct shapes - assert cube_object.data.root_pos_w.torch.shape == (num_cubes, 3) - assert cube_object.data.root_quat_w.torch.shape == (num_cubes, 4) - - # Simulate physics - for _ in range(2): - # perform rendering + for _ in range(60): sim.step() - # update object - cube_object.update(sim.cfg.dt) - # check that the object is kinematic - default_root_pose = cube_object.data.default_root_pose.torch.clone() - default_root_vel = cube_object.data.default_root_vel.torch.clone() - default_root_pose[:, :3] += origins - torch.testing.assert_close(cube_object.data.root_link_pose_w.torch, default_root_pose) - torch.testing.assert_close(cube_object.data.root_com_vel_w.torch, default_root_vel) - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_initialization_with_no_rigid_body(num_cubes, device): - """Test that initialization fails when no rigid body is found at the provided prim path.""" - with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: - # Generate cubes scene - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, api="none", device=device) + scene.update(sim.get_physics_dt()) + positions = scene["object"].data.root_pos_w.torch + expected = scene.env_origins.clone() + expected[:, 2] = 0.2 + torch.testing.assert_close(positions, expected, atol=0.02, rtol=0.0) - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(cube_object) < 10 - # Play sim +@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) +def test_initialization_with_no_rigid_body(device: str) -> None: + """Test that initialization fails when no rigid body is found at the provided prim path.""" + with build_simulation_context(device=device, sim_cfg=SimulationCfg(physics=OvPhysxCfg(), device=device)) as sim: + # Keep the asset alive: only a live asset initializes, and fails, on reset. + cube_object = _spawn_static_colliders() with pytest.raises(RuntimeError, match="Expected 1 prims"): sim.reset() + assert not cube_object.is_initialized -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) @pytest.mark.xfail( strict=True, raises=pytest.fail.Exception, reason="OVPhysX RigidObject has no ArticulationRootAPI guard; the rigid body initializes normally.", ) -def test_initialization_with_articulation_root(num_cubes, device): +def test_initialization_with_articulation_root(device: str) -> None: """Test that initialization fails when an articulation root is found at the provided prim path.""" - with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: - # Generate cubes scene and mark each rigid body as an articulation root. - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, device=device) + with build_simulation_context(device=device, sim_cfg=SimulationCfg(physics=OvPhysxCfg(), device=device)) as sim: + # Mark each rigid body as an articulation root; the asset must stay alive for reset. + _cube_object = _spawn_cubes("Rooted", y_offset=0.0) stage = sim_utils.get_current_stage() - for i in range(num_cubes): - UsdPhysics.ArticulationRootAPI.Apply(stage.GetPrimAtPath(f"/World/Table_{i}/Object")) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(cube_object) < 10 - - # Play sim + for index in range(_NUM_CUBES): + UsdPhysics.ArticulationRootAPI.Apply(stage.GetPrimAtPath(f"/World/Rooted/Env_{index}/Cube")) with pytest.raises(RuntimeError): sim.reset() -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_external_force_on_single_body(num_cubes, device): - """Test application of external force on the base of the object. - - In this test, we apply a force equal to the weight of an object on the base of - one of the objects. We check that the object does not move. For the other object, - we do not apply any force and check that it falls down. - - We validate that this works when we apply the force in the global frame and in the local frame. - """ - # Generate cubes scene - with _ovphysx_sim_context(device=device, add_ground_plane=True, auto_add_lighting=True) as sim: - cube_object, origins = generate_cubes_scene(num_cubes=num_cubes, device=device) - - # Play the simulator - sim.reset() - - # Find bodies to apply the force - body_ids, body_names = cube_object.find_bodies(".*") - - # Sample a force equal to the weight of the object. PhysX reads the mass - # from ``root_view.get_masses()``; OVPhysX exposes the same value via - # ``cube_object.data.body_mass`` (shape ``(N, 1)``). - external_wrench_b = torch.zeros(cube_object.num_instances, len(body_ids), 6, device=sim.device) - # Every 2nd cube should have a force applied to it - external_wrench_b[0::2, :, 2] = 9.81 * cube_object.data.body_mass.torch[0] - - # Now we are ready! - for i in range(5): - # reset root state - root_pose = cube_object.data.default_root_pose.torch.clone() - root_vel = cube_object.data.default_root_vel.torch.clone() - - # need to shift the position of the cubes otherwise they will be on top of each other - root_pose[:, :3] = origins - cube_object.write_root_pose_to_sim_index(root_pose=root_pose) - cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) - - # reset object - cube_object.reset() - - is_global = False - if i % 2 == 0: - is_global = True - positions = cube_object.data.body_com_pos_w.torch[:, body_ids, :3] - else: - positions = None - - # apply force - cube_object.permanent_wrench_composer.set_forces_and_torques_index( - forces=external_wrench_b[..., :3], - torques=external_wrench_b[..., 3:], - positions=positions, - body_ids=body_ids, - is_global=is_global, - ) - # perform simulation - for _ in range(5): - # apply action to the object - cube_object.write_data_to_sim() - - # perform step - sim.step() - - # update buffers - cube_object.update(sim.cfg.dt) - - # First object should still be at the same Z position (1.0) - torch.testing.assert_close( - cube_object.data.root_pos_w.torch[0::2, 2], torch.ones(num_cubes // 2, device=sim.device) - ) - # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(cube_object.data.root_pos_w.torch[1::2, 2] < 1.0) - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_external_force_on_single_body_at_position(num_cubes, device): - """Test application of external force on the base of the object at a specific position. - - In this test, we apply a force equal to the weight of an object on the base of - one of the objects at 1m in the Y direction, we check that the object rotates around it's X axis. - For the other object, we do not apply any force and check that it falls down. - - We validate that this works when we apply the force in the global frame and in the local frame. - """ - # Generate cubes scene - with _ovphysx_sim_context(device=device, add_ground_plane=True, auto_add_lighting=True) as sim: - cube_object, origins = generate_cubes_scene(num_cubes=num_cubes, device=device) - - # Play the simulator - sim.reset() - - # Find bodies to apply the force - body_ids, body_names = cube_object.find_bodies(".*") - - # Sample a force equal to the weight of the object - external_wrench_b = torch.zeros(cube_object.num_instances, len(body_ids), 6, device=sim.device) - external_wrench_positions_b = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=sim.device) - # Every 2nd cube should have a force applied to it - external_wrench_b[0::2, :, 2] = 500.0 - external_wrench_positions_b[0::2, :, 1] = 1.0 - - # Desired force and torque - desired_force = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=sim.device) - desired_force[0::2, :, 2] = 1000.0 - desired_torque = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=sim.device) - desired_torque[0::2, :, 0] = 1000.0 - # Now we are ready! - for i in range(5): - # reset root state - root_pose = cube_object.data.default_root_pose.torch.clone() - root_vel = cube_object.data.default_root_vel.torch.clone() - - # need to shift the position of the cubes otherwise they will be on top of each other - root_pose[:, :3] = origins - cube_object.write_root_pose_to_sim_index(root_pose=root_pose) - cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) - - # reset object - cube_object.reset() - - is_global = False - if i % 2 == 0: - is_global = True - body_com_pos_w = cube_object.data.body_com_pos_w.torch[:, body_ids, :3] - external_wrench_positions_b[..., 0] = 0.0 - external_wrench_positions_b[..., 1] = 1.0 - external_wrench_positions_b[..., 2] = 0.0 - external_wrench_positions_b += body_com_pos_w - else: - external_wrench_positions_b[..., 0] = 0.0 - external_wrench_positions_b[..., 1] = 1.0 - external_wrench_positions_b[..., 2] = 0.0 - - # apply force - cube_object.permanent_wrench_composer.set_forces_and_torques_index( - forces=external_wrench_b[..., :3], - torques=external_wrench_b[..., 3:], - positions=external_wrench_positions_b, - body_ids=body_ids, - is_global=is_global, - ) - cube_object.permanent_wrench_composer.add_forces_and_torques_index( - forces=external_wrench_b[..., :3], - torques=external_wrench_b[..., 3:], - positions=external_wrench_positions_b, - body_ids=body_ids, - is_global=is_global, - ) - torch.testing.assert_close( - cube_object._permanent_wrench_composer.composed_force.torch[:, 0, :], - desired_force[:, 0, :], - rtol=1e-6, - atol=1e-7, - ) - torch.testing.assert_close( - cube_object._permanent_wrench_composer.composed_torque.torch[:, 0, :], - desired_torque[:, 0, :], - rtol=1e-6, - atol=1e-7, - ) - # perform simulation - for _ in range(5): - # apply action to the object - cube_object.write_data_to_sim() - - # perform step - sim.step() - - # update buffers - cube_object.update(sim.cfg.dt) - - # The first object should be rotating around it's X axis - assert torch.all(torch.abs(cube_object.data.root_ang_vel_b.torch[0::2, 0]) > 0.1) - # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(cube_object.data.root_pos_w.torch[1::2, 2] < 1.0) - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_set_rigid_object_state(num_cubes, device): - """Test setting the state of the rigid object. - - In this test, we set the state of the rigid object to a random state and check - that the object is in that state after simulation. We set gravity to zero as - we don't want any external forces acting on the object to ensure state remains static. - """ - # Turn off gravity for this test as we don't want any external forces acting on the object - # to ensure state remains static - with _ovphysx_sim_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: - # Generate cubes scene - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, device=device) - - # Play the simulator - sim.reset() - - state_types = ["root_pos_w", "root_quat_w", "root_lin_vel_w", "root_ang_vel_w"] - - # Set each state type individually as they are dependent on each other - for state_type_to_randomize in state_types: - state_dict = { - "root_pos_w": torch.zeros_like(cube_object.data.root_pos_w.torch, device=sim.device), - "root_quat_w": default_orientation(num=num_cubes, device=sim.device), - "root_lin_vel_w": torch.zeros_like(cube_object.data.root_lin_vel_w.torch, device=sim.device), - "root_ang_vel_w": torch.zeros_like(cube_object.data.root_ang_vel_w.torch, device=sim.device), - } - - # Now we are ready! - for _ in range(5): - # reset object - cube_object.reset() - - # Set random state - if state_type_to_randomize == "root_quat_w": - state_dict[state_type_to_randomize] = random_orientation(num=num_cubes, device=sim.device) - else: - state_dict[state_type_to_randomize] = torch.randn(num_cubes, 3, device=sim.device) - - # perform simulation - for _ in range(5): - root_pose = torch.cat( - [state_dict["root_pos_w"], state_dict["root_quat_w"]], - dim=-1, - ) - root_vel = torch.cat( - [state_dict["root_lin_vel_w"], state_dict["root_ang_vel_w"]], - dim=-1, - ) - # reset root state - cube_object.write_root_pose_to_sim_index(root_pose=root_pose) - cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) - - sim.step() - - # assert that set root quantities are equal to the ones set in the state_dict - for key, expected_value in state_dict.items(): - value = getattr(cube_object.data, key).torch - torch.testing.assert_close(value, expected_value, rtol=1e-3, atol=1e-3) - - cube_object.update(sim.cfg.dt) - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_reset_rigid_object(num_cubes, device): - """Test resetting the state of the rigid object.""" - with _ovphysx_sim_context(device=device, gravity_enabled=True, auto_add_lighting=True) as sim: - # Generate cubes scene - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, device=device) - - # Play the simulator - sim.reset() - - for i in range(5): - # perform rendering - sim.step() - - # update object - cube_object.update(sim.cfg.dt) - - # Move the object to a random position - root_pose = cube_object.data.default_root_pose.torch.clone() - root_pose[:, :3] = torch.randn(num_cubes, 3, device=sim.device) - - # Random orientation - root_pose[:, 3:7] = random_orientation(num=num_cubes, device=sim.device) - cube_object.write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = cube_object.data.default_root_vel.torch.clone() - cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) - - if i % 2 == 0: - # Make both wrench composers active so the reset has something to clear. - cube_object.permanent_wrench_composer.set_forces_and_torques_index( - forces=torch.ones((num_cubes, 1, 3), device=device), - torques=torch.ones((num_cubes, 1, 3), device=device), - ) - cube_object.instantaneous_wrench_composer.add_forces_and_torques_index( - forces=torch.ones((num_cubes, 1, 3), device=device), - torques=torch.ones((num_cubes, 1, 3), device=device), - ) - assert cube_object._instantaneous_wrench_composer.active - assert cube_object._permanent_wrench_composer.active - - # reset object - cube_object.reset() - - # Reset should zero external forces and torques - assert not cube_object._instantaneous_wrench_composer.active - assert not cube_object._permanent_wrench_composer.active - assert torch.count_nonzero(cube_object._instantaneous_wrench_composer.composed_force.torch) == 0 - assert torch.count_nonzero(cube_object._instantaneous_wrench_composer.composed_torque.torch) == 0 - assert torch.count_nonzero(cube_object._permanent_wrench_composer.composed_force.torch) == 0 - assert torch.count_nonzero(cube_object._permanent_wrench_composer.composed_torque.torch) == 0 - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_rigid_body_no_friction(num_cubes, device): - """Test that a rigid object with no friction maintains its tangential velocity on a plane.""" - with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=0.0, device=device) - - # Create a ground plane with no friction - cfg = sim_utils.GroundPlaneCfg( - physics_material=materials.RigidBodyMaterialBaseCfg( - static_friction=0.0, dynamic_friction=0.0, restitution=0.0 - ) - ) - cfg.func("/World/GroundPlane", cfg) - - # Play sim - sim.reset() - - # Set the cubes' friction (and restitution) to zero. This test isolates friction, so a zero - # restitution keeps the resting contact clean instead of letting the cube bounce vertically. - num_shapes = _num_shapes(cube_object) - cube_materials = torch.zeros(num_cubes, num_shapes, 3, device="cpu") - _write_shape_material(cube_object, cube_materials) - - # Establish ground contact before measuring tangential sliding velocity. - for _ in range(30): - sim.step() - cube_object.update(sim.cfg.dt) - - # Initial velocity in X to get the block moving. - initial_velocity = torch.zeros((num_cubes, 6), device=device) - initial_velocity[:, 0] = 0.1 - cube_object.write_root_velocity_to_sim_index(root_velocity=initial_velocity) - - # Non-deterministic on GPU, so use a looser tolerance there. - tolerance = 1e-2 if device == "cuda:0" else 1e-5 - for _ in range(5): - sim.step() - cube_object.update(sim.cfg.dt) - torch.testing.assert_close( - cube_object.data.root_lin_vel_w.torch[:, :2], initial_velocity[:, :2], rtol=1e-5, atol=tolerance - ) - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_rigid_body_with_static_friction(num_cubes, device): - """Test that static friction applied to rigid object works as expected. - - This test works by applying a force to the object and checking if the object moves or not based on the - mu (coefficient of static friction) value set for the object. We set the static friction to be non-zero and - apply a force to the object. When the force applied is below mu, the object should not move. When the force - applied is above mu, the object should move. - """ - with _ovphysx_sim_context(device=device, dt=0.01, auto_add_lighting=True) as sim: - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=0.03125, device=device) - - # Create ground plane. Dynamic friction is set equal to static friction to work around a PhysX bug. - mu = 0.5 - cfg = sim_utils.GroundPlaneCfg( - physics_material=materials.RigidBodyMaterialBaseCfg(static_friction=mu, dynamic_friction=mu) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) +def test_collection_initialization_with_no_rigid_body(device: str) -> None: + """Test that collection initialization fails when no rigid body is found at a body prim path.""" + with build_simulation_context( + device=device, sim_cfg=SimulationCfg(physics=OvPhysxCfg(), device=device, gravity=(0.0, 0.0, 0.0)) + ) as sim: + # Keep the asset alive: only a live asset initializes, and fails, on reset. + object_collection = _spawn_collection( + "Static", + y_offset=0.0, + spawn=sim_utils.CuboidCfg(size=(0.1, 0.1, 0.1), collision_props=sim_utils.UsdPhysicsCollisionCfg()), ) - cfg.func("/World/GroundPlane", cfg) - - # Play sim - sim.reset() - - # Set the cubes' static (and, per the PhysX bug, dynamic) friction to mu, restitution zero. - num_shapes = _num_shapes(cube_object) - cube_materials = torch.zeros(num_cubes, num_shapes, 3, device="cpu") - cube_materials[..., 0] = mu - cube_materials[..., 1] = mu - _write_shape_material(cube_object, cube_materials) - - # Let everything settle. - for _ in range(100): - sim.step() - cube_object.update(sim.cfg.dt) - cube_object.write_root_velocity_to_sim_index(root_velocity=torch.zeros((num_cubes, 6), device=device)) - - cube_mass = cube_object.data.body_mass.torch[:, 0] # [N], PhysX reads this via root_view.get_masses() - gravity_magnitude = abs(sim.cfg.gravity[2]) - # below mu: block should not move (applied force <= mu); above mu: block should move. - for force in "below_mu", "above_mu": - cube_object.write_root_velocity_to_sim_index(root_velocity=torch.zeros((num_cubes, 6), device=device)) - - external_wrench_b = torch.zeros((num_cubes, 1, 6), device=device) - factor = 0.99 if force == "below_mu" else 1.01 - external_wrench_b[:, 0, 0] = mu * cube_mass * gravity_magnitude * factor - - cube_object.permanent_wrench_composer.set_forces_and_torques_index( - forces=external_wrench_b[..., :3], - torques=external_wrench_b[..., 3:], - ) - - initial_root_pos = cube_object.data.root_pos_w.torch.clone() - for _ in range(200): - cube_object.write_data_to_sim() - sim.step() - cube_object.update(sim.cfg.dt) - if force == "below_mu": - torch.testing.assert_close( - cube_object.data.root_pos_w.torch, initial_root_pos, rtol=2e-3, atol=2e-3 - ) - if force == "above_mu": - assert (cube_object.data.root_pos_w.torch[..., 0] - initial_root_pos[..., 0] > 0.02).all() - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_rigid_body_with_restitution(num_cubes, device): - """Test that restitution when applied to rigid object works as expected. - - This test works by dropping a block from a height and checking if the block bounces or not based on the - restitution value set for the object. We set the restitution to be non-zero and drop the block from a height. - When the restitution is 0, the block should not bounce. When the restitution is between 0 and 1, the block - should bounce with less energy. - """ - for expected_collision_type in "partially_elastic", "inelastic": - with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=1.0, device=device) - - restitution_coefficient = 0.0 if expected_collision_type == "inelastic" else 0.5 - - # Create ground plane with the matching restitution. - cfg = sim_utils.GroundPlaneCfg( - physics_material=materials.RigidBodyMaterialBaseCfg(restitution=restitution_coefficient) - ) - cfg.func("/World/GroundPlane", cfg) - - # Play sim + with pytest.raises(RuntimeError, match="Expected 1 prims"): sim.reset() - - # Drop the cubes from a height with an initial downward velocity. - root_pose = torch.zeros(num_cubes, 7, device=device) - root_pose[:, 3] = 1.0 # unit quaternion - for i in range(num_cubes): - root_pose[i, 1] = 1.0 * i - root_pose[:, 2] = 1.0 - root_vel = torch.zeros(num_cubes, 6, device=device) - root_vel[:, 2] = -1.0 - cube_object.write_root_pose_to_sim_index(root_pose=root_pose) - cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) - - # Frictionless cubes with the matching restitution. - num_shapes = _num_shapes(cube_object) - cube_materials = torch.zeros(num_cubes, num_shapes, 3, device="cpu") - cube_materials[..., 2] = restitution_coefficient - _write_shape_material(cube_object, cube_materials) - - curr_z_velocity = cube_object.data.root_lin_vel_w.torch[:, 2].clone() - prev_z_velocity = curr_z_velocity - for _ in range(100): - sim.step() - cube_object.update(sim.cfg.dt) - curr_z_velocity = cube_object.data.root_lin_vel_w.torch[:, 2].clone() - - if expected_collision_type == "inelastic": - # The block must not bounce: its z velocity stays <= 0. - assert (curr_z_velocity <= 0.0).all() - - if torch.all(curr_z_velocity <= 0.0): - prev_z_velocity = curr_z_velocity # still falling - else: - break # collision happened (now moving up) - - if expected_collision_type == "partially_elastic": - # The block bounced but lost energy. - assert torch.all(torch.le(abs(curr_z_velocity), abs(prev_z_velocity))) - assert (curr_z_velocity > 0.0).all() - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_rigid_body_set_mass(num_cubes, device): - """Test getting and setting mass of rigid object.""" - with _ovphysx_sim_context( - device=device, gravity_enabled=False, add_ground_plane=True, auto_add_lighting=True - ) as sim: - # Create a scene with random cubes - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=1.0, device=device) - - # Play sim + assert not object_collection.is_initialized + + +@dataclass +class _RigidScene: + """Cuboid pairs and collections that share one real OVPhysX lifecycle.""" + + sim: SimulationContext + device: str + dynamic: RigidObject + kinematic: RigidObject + spinning: RigidObject + fused: RigidObjectCollection + stepped: RigidObjectCollection + wrenched: RigidObjectCollection + + +@pytest.fixture(scope="module") +def scene(request: pytest.FixtureRequest) -> Iterator[_RigidScene]: + """Initialize every cuboid pair and collection for one device once for this module.""" + device = request.param + with build_simulation_context(device=device, sim_cfg=SimulationCfg(physics=OvPhysxCfg(), device=device)) as sim: + dynamic = _spawn_cubes("Dynamic", y_offset=0.0, disable_gravity=True) + kinematic = _spawn_cubes("Kinematic", y_offset=2.0, kinematic_enabled=True) + spinning = _spawn_cubes("Spinning", y_offset=4.0, disable_gravity=True) + # Collection bodies disable gravity and span three units along y. + fused = _spawn_collection("Fused", y_offset=7.0) + stepped = _spawn_collection("Stepped", y_offset=11.0) + wrenched = _spawn_collection("Wrenched", y_offset=15.0) sim.reset() + yield _RigidScene(sim, device, dynamic, kinematic, spinning, fused, stepped, wrenched) + + +@pytest.mark.parametrize("scene", test_devices(), indirect=True) +def test_rigid_object_real_ovphysx_seams(scene: _RigidScene) -> None: + """Prove partial state, inertial properties, one real wrench delivery, and the wrench reset.""" + rigid_object, sim, device = scene.dynamic, scene.sim, scene.device + assert rigid_object.is_initialized + assert rigid_object.num_instances == _NUM_CUBES + assert rigid_object.data.body_mass.torch.shape == (_NUM_CUBES, 1) + assert rigid_object.data.body_com_pose_b.torch.shape == (_NUM_CUBES, 1, 7) + assert rigid_object.data.body_inertia.torch.shape == (_NUM_CUBES, 1, 9) + + initial_pose = rigid_object.data.root_link_pose_w.torch.clone() + target_pose = initial_pose[1:2].clone() + target_pose[0, :3] += torch.tensor([0.25, -0.1, 0.2], device=device) + rigid_object.write_root_link_pose_to_sim_index(root_pose=target_pose, env_ids=[1]) + torch.testing.assert_close(rigid_object.data.root_link_pose_w.torch[1:2], target_pose) + torch.testing.assert_close(rigid_object.data.root_link_pose_w.torch[0:1], initial_pose[0:1]) + raw_pose = wp.to_torch(rigid_object.root_view.get_attribute(TT.RIGID_BODY_POSE)).to(device) + torch.testing.assert_close(raw_pose, torch.cat((initial_pose[0:1], target_pose))) + + raw_mass_before = wp.to_torch(rigid_object.root_view.get_attribute(TT.RIGID_BODY_MASS)).clone() + rigid_object.set_masses_index(masses=wp.array([[3.0]], dtype=wp.float32, device=device), env_ids=[1]) + expected_raw_mass = raw_mass_before.clone() + expected_raw_mass[1] = 3.0 + torch.testing.assert_close(wp.to_torch(rigid_object.root_view.get_attribute(TT.RIGID_BODY_MASS)), expected_raw_mass) + + raw_com_before = wp.to_torch(rigid_object.root_view.get_attribute(TT.RIGID_BODY_COM_POSE)).clone() + coms = rigid_object.data.body_com_pose_b.torch[1:2].clone() + coms[0, 0, :3] = torch.tensor([0.01, -0.02, 0.03], device=device) + rigid_object.set_coms_index(coms=wp.from_torch(coms, dtype=wp.transformf), env_ids=[1]) + expected_raw_com = raw_com_before.clone() + expected_raw_com[1] = coms[0, 0].cpu() + torch.testing.assert_close( + wp.to_torch(rigid_object.root_view.get_attribute(TT.RIGID_BODY_COM_POSE)), expected_raw_com + ) - # Get masses before increasing - original_masses = cube_object.data.body_mass.torch.clone() - - assert original_masses.shape == (num_cubes, 1) - - # Randomize mass of the object - masses = original_masses + torch.FloatTensor(num_cubes, 1).uniform_(4, 8).to(sim.device) - - indices = torch.tensor(range(num_cubes), dtype=torch.int32) - - # Set the new masses via the OVPhysX writer (matches PhysX/Newton). - cube_object.set_masses_index( - masses=wp.from_torch(masses.contiguous(), dtype=wp.float32), - env_ids=wp.from_torch(indices, dtype=wp.int32), - ) - - torch.testing.assert_close(cube_object.data.body_mass.torch, masses) - - # Simulate physics - # perform rendering + raw_inertia_before = wp.to_torch(rigid_object.root_view.get_attribute(TT.RIGID_BODY_INERTIA)).clone() + inertias = rigid_object.data.body_inertia.torch[1:2].clone() + inertias[0, 0, 0] *= 1.5 + rigid_object.set_inertias_index(inertias=wp.from_torch(inertias, dtype=wp.float32), env_ids=[1]) + expected_raw_inertia = raw_inertia_before.clone() + expected_raw_inertia[1] = inertias[0, 0].cpu() + torch.testing.assert_close( + wp.to_torch(rigid_object.root_view.get_attribute(TT.RIGID_BODY_INERTIA)), expected_raw_inertia + ) + torch.testing.assert_close(rigid_object.data.body_mass.torch[:, 0], torch.tensor([1.0, 3.0], device=device)) + torch.testing.assert_close(rigid_object.data.body_com_pose_b.torch[1:2], coms) + torch.testing.assert_close(rigid_object.data.body_inertia.torch[1:2], inertias) + + initial_velocity = rigid_object.data.root_com_vel_w.torch.clone() + forces = torch.zeros((_NUM_CUBES, 1, 3), device=device) + forces[1, 0, 0] = 20.0 + rigid_object.instantaneous_wrench_composer.set_forces_and_torques_index(forces=forces, is_global=True) + rigid_object.write_data_to_sim() + sim.step() + rigid_object.update(sim.cfg.dt) + + assert rigid_object.data.root_com_vel_w.torch[1, 0] > initial_velocity[1, 0] + torch.testing.assert_close(rigid_object.data.root_com_vel_w.torch[0], initial_velocity[0], atol=1e-6, rtol=0.0) + + # Reset clears both wrench composers. + ones = torch.ones((_NUM_CUBES, 1, 3), device=device) + rigid_object.permanent_wrench_composer.set_forces_and_torques_index(forces=ones, torques=ones) + rigid_object.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones) + rigid_object.reset() + for composer in (rigid_object.instantaneous_wrench_composer, rigid_object.permanent_wrench_composer): + assert not composer.active + assert torch.count_nonzero(composer.composed_force.torch) == 0 + assert torch.count_nonzero(composer.composed_torque.torch) == 0 + + +@pytest.mark.parametrize("scene", test_devices(), indirect=True) +def test_initialization_with_kinematic_enabled(scene: _RigidScene) -> None: + """Test that kinematic bodies publish transforms and hold their pose under gravity.""" + cube_object, sim = scene.kinematic, scene.sim + + # SDP bindings must be ready on CPU and GPU before any asset pose read. + provider = sim.get_scene_data_provider() + names = ("Dynamic", "Kinematic", "Spinning") + expected_paths = {f"/World/{name}/Env_{i}/Cube" for name in names for i in (0, 1)} + expected_paths |= { + f"/World/{name}/Env_{i}/Object_{body}" + for name in ("Fused", "Stepped", "Wrenched") + for i in (0, 1) + for body in (0, 1, 2) + } + assert provider.transform_count == len(expected_paths) + assert set(provider.backend.transform_paths) == expected_paths + + assert cube_object.is_initialized + assert len(cube_object.body_names) == 1 + default_root_pose = cube_object.data.default_root_pose.torch.clone() + default_root_pose[:, :3] += torch.tensor([[0.0, 2.0, 0.0], [2.0, 2.0, 0.0]], device=scene.device) + default_root_vel = cube_object.data.default_root_vel.torch.clone() + for _ in range(2): sim.step() - # update object cube_object.update(sim.cfg.dt) - - # Check the simulation received the masses; the data cache is written by the setter itself. - masses_to_check = wp.to_torch(cube_object.root_view.get_attribute(TT.RIGID_BODY_MASS)) - torch.testing.assert_close(masses_to_check.reshape(num_cubes, 1), masses.cpu()) + torch.testing.assert_close(cube_object.data.root_link_pose_w.torch, default_root_pose) + torch.testing.assert_close(cube_object.data.root_com_vel_w.torch, default_root_vel) -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("gravity_enabled", [True, False]) -def test_gravity_vec_w(num_cubes, device, gravity_enabled): - """Test initialization and that gravity vector direction is set correctly for the rigid object.""" - with _ovphysx_sim_context(device=device, gravity_enabled=gravity_enabled) as sim: - # Create a scene with random cubes - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, device=device) +@pytest.mark.parametrize("scene", _CUDA_SCENES, indirect=True) +def test_body_root_state_properties(scene: _RigidScene) -> None: + """Test the root and body link and COM states of cubes spinning about an offset center of mass.""" + cube_object, sim, device = scene.spinning, scene.sim, scene.device + env_pos = cube_object.data.root_link_pos_w.torch.clone() - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(cube_object) < 10 - - # Obtain gravity direction - if gravity_enabled: - gravity_dir = (0.0, 0.0, -1.0) - else: - gravity_dir = (0.0, 0.0, 0.0) - - # Play sim - sim.reset() + # Offset the center of mass along the link x-axis. + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(_NUM_CUBES, 1) + com = cube_object.data.body_com_pose_b.torch.clone() # shape (N, 1, 7) + com[..., :3] = offset.unsqueeze(1) + cube_object.set_coms_index(coms=wp.from_torch(com, dtype=wp.transformf)) + torch.testing.assert_close(cube_object.data.body_com_pose_b.torch, com) - # Check if object is initialized - assert cube_object.is_initialized - assert len(cube_object.body_names) == 1 + spin_twist = torch.zeros(6, device=device) + spin_twist[5] = 2.0 - # Check buffers that exists and have correct shapes - assert cube_object.data.root_pos_w.torch.shape == (num_cubes, 3) - assert cube_object.data.root_quat_w.torch.shape == (num_cubes, 4) - assert cube_object.data.body_mass.torch.shape == (num_cubes, 1) - assert cube_object.data.body_inertia.torch.shape == (num_cubes, 1, 9) - - # Check that gravity is set correctly - assert cube_object.data.GRAVITY_VEC_W.torch[0, 0] == gravity_dir[0] - assert cube_object.data.GRAVITY_VEC_W.torch[0, 1] == gravity_dir[1] - assert cube_object.data.GRAVITY_VEC_W.torch[0, 2] == gravity_dir[2] - - # Simulate physics - for _ in range(2): - # perform rendering - sim.step() - # update object - cube_object.update(sim.cfg.dt) - - # Expected gravity value is the acceleration of the body - gravity = torch.zeros(num_cubes, 1, 6, device=device) - if gravity_enabled: - gravity[:, :, 2] = -9.81 - # Check the body accelerations are correct - torch.testing.assert_close(cube_object.data.body_acc_w.torch, gravity) - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@flaky(max_runs=3, min_passes=1) -def test_body_root_state_properties(num_cubes, device): - """Test the root_com_state_w, root_link_state_w, body_com_state_w, and body_link_state_w properties.""" - with _ovphysx_sim_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: - # Create a scene with random cubes - cube_object, env_pos = generate_cubes_scene(num_cubes=num_cubes, height=0.0, device=device) - env_idx = torch.tensor([x for x in range(num_cubes)], dtype=torch.int32) - - # Play sim - sim.reset() + for _ in range(10): + # Keep spinning about the z-axis through the center of mass. + cube_object.write_root_velocity_to_sim_index(root_velocity=spin_twist.repeat(_NUM_CUBES, 1)) + sim.step() + cube_object.update(sim.cfg.dt) - # Check if cube_object is initialized - assert cube_object.is_initialized - - # change center of mass offset from link frame - offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_cubes, 1) - - # Read current COMs, mutate the translation, write back via the OVPhysX - # ``set_coms_index`` setter (PhysX uses ``root_view.set_coms`` for the same - # operation; OVPhysX wraps the wheel ``RIGID_BODY_COM_POSE`` write in - # :meth:`set_coms_index`, which follows the PhysX ``wp.transformf`` contract). - com = cube_object.data.body_com_pose_b.torch.clone() # shape (N, 1, 7) - com[..., :3] = offset.to(com.device).unsqueeze(1) - cube_object.set_coms_index( - coms=wp.from_torch(com.contiguous(), dtype=wp.transformf), - env_ids=wp.from_torch(env_idx, dtype=wp.int32), + root_link_pose_w = cube_object.data.root_link_pose_w.torch + root_link_vel_w = cube_object.data.root_link_vel_w.torch + root_com_pose_w = cube_object.data.root_com_pose_w.torch + root_com_vel_w = cube_object.data.root_com_vel_w.torch + body_link_pose_w = cube_object.data.body_link_pose_w.torch + body_link_vel_w = cube_object.data.body_link_vel_w.torch + body_com_pose_w = cube_object.data.body_com_pose_w.torch + body_com_vel_w = cube_object.data.body_com_vel_w.torch + + # The center of mass stays fixed while the cube spins about it. + torch.testing.assert_close(env_pos + offset, root_com_pose_w[..., :3]) + torch.testing.assert_close(env_pos + offset, body_com_pose_w[..., :3].squeeze(-2)) + # The link origin stays at the negated COM offset in the link frame. + root_link_state_pos_rel_com = quat_apply_inverse( + root_link_pose_w[..., 3:], + root_link_pose_w[..., :3] - root_com_pose_w[..., :3], ) - - # check ceter of mass has been set - torch.testing.assert_close(cube_object.data.body_com_pose_b.torch, com) - - # random z spin velocity - spin_twist = torch.zeros(6, device=device) - spin_twist[5] = torch.randn(1, device=device) - - # Simulate physics - for _ in range(100): - # spin the object around Z axis (com) - cube_object.write_root_velocity_to_sim_index(root_velocity=spin_twist.repeat(num_cubes, 1)) - # perform rendering - sim.step() - # update object - cube_object.update(sim.cfg.dt) - - # get state properties - root_link_pose_w = cube_object.data.root_link_pose_w.torch - root_link_vel_w = cube_object.data.root_link_vel_w.torch - root_com_pose_w = cube_object.data.root_com_pose_w.torch - root_com_vel_w = cube_object.data.root_com_vel_w.torch - body_link_pose_w = cube_object.data.body_link_pose_w.torch - body_link_vel_w = cube_object.data.body_link_vel_w.torch - body_com_pose_w = cube_object.data.body_com_pose_w.torch - body_com_vel_w = cube_object.data.body_com_vel_w.torch - - # cubes are spinning around center of mass - # position will not match - # center of mass position will be constant (i.e. spinning around com) - torch.testing.assert_close(env_pos + offset, root_com_pose_w[..., :3]) - torch.testing.assert_close(env_pos + offset, body_com_pose_w[..., :3].squeeze(-2)) - # link position will be moving but should stay constant away from center of mass - root_link_state_pos_rel_com = quat_apply_inverse( - root_link_pose_w[..., 3:], - root_link_pose_w[..., :3] - root_com_pose_w[..., :3], - ) - torch.testing.assert_close(-offset, root_link_state_pos_rel_com) - body_link_state_pos_rel_com = quat_apply_inverse( - body_link_pose_w[..., 3:], - body_link_pose_w[..., :3] - body_com_pose_w[..., :3], - ) - torch.testing.assert_close(-offset, body_link_state_pos_rel_com.squeeze(-2)) - - # orientation of com will be a constant rotation from link orientation - com_quat_b = cube_object.data.body_com_quat_b.torch - com_quat_w = quat_mul(body_link_pose_w[..., 3:], com_quat_b) - torch.testing.assert_close(com_quat_w, body_com_pose_w[..., 3:]) - torch.testing.assert_close(com_quat_w.squeeze(-2), root_com_pose_w[..., 3:]) - - # lin_vel will not match - # center of mass vel will be constant (i.e. spinning around com) - torch.testing.assert_close(torch.zeros_like(root_com_vel_w[..., :3]), root_com_vel_w[..., :3]) - torch.testing.assert_close(torch.zeros_like(body_com_vel_w[..., :3]), body_com_vel_w[..., :3]) - # link frame will be moving, and should account for the reported COM velocity and offset - lin_vel_rel_root_gt = quat_apply_inverse(root_link_pose_w[..., 3:], root_link_vel_w[..., :3]) - lin_vel_rel_body_gt = quat_apply_inverse(body_link_pose_w[..., 3:], body_link_vel_w[..., :3]) - com_lin_vel_rel_gt = quat_apply_inverse(root_link_pose_w[..., 3:], root_com_vel_w[..., :3]) - com_ang_vel_rel_gt = quat_apply_inverse(root_link_pose_w[..., 3:], root_com_vel_w[..., 3:]) - lin_vel_rel_gt = com_lin_vel_rel_gt + torch.linalg.cross(com_ang_vel_rel_gt, -offset) - torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_root_gt, atol=1e-4, rtol=1e-4) - torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_body_gt.squeeze(-2), atol=1e-4, rtol=1e-4) - - # ang_vel will always match - torch.testing.assert_close(root_com_vel_w[..., 3:], root_link_vel_w[..., 3:]) - torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("state_location", ["com", "link", "root"]) -def test_write_root_state(num_cubes, device, state_location): - """Test the root-state setters in the link, center-of-mass, and root reference frames. - - The written frame reads back unchanged and the other frame follows it through the COM offset. - """ - with _ovphysx_sim_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: - # Create a scene with random cubes - cube_object, env_pos = generate_cubes_scene(num_cubes=num_cubes, height=0.0, device=device) - env_idx = torch.tensor([x for x in range(num_cubes)], dtype=torch.int32) - - # Play sim - sim.reset() - - # Check if cube_object is initialized - assert cube_object.is_initialized - - # change center of mass offset from link frame - offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_cubes, 1) - - com = cube_object.data.body_com_pose_b.torch.clone() # shape (N, 1, 7) - com[..., :3] = offset.to(com.device).unsqueeze(1) - cube_object.set_coms_index( - coms=wp.from_torch(com.contiguous(), dtype=wp.transformf), - env_ids=wp.from_torch(env_idx, dtype=wp.int32), + torch.testing.assert_close(-offset, root_link_state_pos_rel_com) + body_link_state_pos_rel_com = quat_apply_inverse( + body_link_pose_w[..., 3:], + body_link_pose_w[..., :3] - body_com_pose_w[..., :3], ) + torch.testing.assert_close(-offset, body_link_state_pos_rel_com.squeeze(-2)) + + # The COM orientation is a constant rotation of the link orientation. + com_quat_b = cube_object.data.body_com_quat_b.torch + com_quat_w = quat_mul(body_link_pose_w[..., 3:], com_quat_b) + torch.testing.assert_close(com_quat_w, body_com_pose_w[..., 3:]) + torch.testing.assert_close(com_quat_w.squeeze(-2), root_com_pose_w[..., 3:]) + + # The center of mass does not translate. + torch.testing.assert_close(torch.zeros_like(root_com_vel_w[..., :3]), root_com_vel_w[..., :3]) + torch.testing.assert_close(torch.zeros_like(body_com_vel_w[..., :3]), body_com_vel_w[..., :3]) + # The link velocity adds the rotation of the offset to the COM velocity. + lin_vel_rel_root_gt = quat_apply_inverse(root_link_pose_w[..., 3:], root_link_vel_w[..., :3]) + lin_vel_rel_body_gt = quat_apply_inverse(body_link_pose_w[..., 3:], body_link_vel_w[..., :3]) + com_lin_vel_rel_gt = quat_apply_inverse(root_link_pose_w[..., 3:], root_com_vel_w[..., :3]) + com_ang_vel_rel_gt = quat_apply_inverse(root_link_pose_w[..., 3:], root_com_vel_w[..., 3:]) + lin_vel_rel_gt = com_lin_vel_rel_gt + torch.linalg.cross(com_ang_vel_rel_gt, -offset) + torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_root_gt, atol=1e-4, rtol=1e-4) + torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_body_gt.squeeze(-2), atol=1e-4, rtol=1e-4) + + # Link and COM frames share the angular velocity. + 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:]) + + +@pytest.mark.parametrize("scene", test_devices(), indirect=True) +def test_rigid_object_collection_real_ovphysx_seams(scene: _RigidScene) -> None: + """Prove fused remapping through partial state, inertial, and material writes.""" + collection, device = scene.fused, scene.device + assert collection.is_initialized + assert collection.num_instances == _NUM_ENVS + assert collection.body_names == ["left", "middle", "right"] + assert collection.data.body_mass.torch.shape == (_NUM_ENVS, _NUM_BODIES) + + env_ids = torch.tensor([1, 0], dtype=torch.int32, device=device) + body_ids = torch.tensor([2, 0], dtype=torch.int32, device=device) + initial_pose = collection.data.body_link_pose_w.torch.clone() + target_pose = initial_pose[env_ids][:, body_ids].clone() + target_pose[0, 0, :3] += torch.tensor([0.2, 0.3, 0.4], device=device) + target_pose[1, 1, :3] += torch.tensor([-0.1, -0.2, 0.1], device=device) + collection.write_body_link_pose_to_sim_index(body_poses=target_pose, env_ids=env_ids, body_ids=body_ids) + torch.testing.assert_close(collection.data.body_link_pose_w.torch[env_ids][:, body_ids], target_pose) + torch.testing.assert_close(collection.data.body_link_pose_w.torch[:, 1], initial_pose[:, 1]) + + # The fused OVPhysX bindings are body-major; the public data is environment-major. + cpu_env_ids, cpu_body_ids = env_ids.cpu(), body_ids.cpu() + raw_pose = wp.to_torch(collection.root_view.get_attribute(TT.RIGID_BODY_POSE)) + expected_pose = initial_pose.clone() + expected_pose[env_ids[:, None].long(), body_ids[None, :].long()] = target_pose + torch.testing.assert_close(raw_pose.reshape(_NUM_BODIES, _NUM_ENVS, 7).transpose(0, 1).to(device), expected_pose) + initial_mass = collection.data.body_mass.torch.clone() + raw_mass_before = wp.to_torch(collection.root_view.get_attribute(TT.BODY_MASS)).reshape(3, 2).T.clone() + masses = torch.tensor([[5.0, 6.0], [7.0, 8.0]], device=device) + collection.set_masses_index(masses=masses, env_ids=env_ids, body_ids=body_ids) + expected_raw_mass = raw_mass_before.clone() + expected_raw_mass[cpu_env_ids[:, None], cpu_body_ids[None, :]] = masses.cpu() + raw_mass = wp.to_torch(collection.root_view.get_attribute(TT.BODY_MASS)).reshape(3, 2).T + torch.testing.assert_close(raw_mass, expected_raw_mass) + torch.testing.assert_close(collection.data.body_mass.torch[env_ids][:, body_ids], masses) + torch.testing.assert_close(collection.data.body_mass.torch[:, 1], initial_mass[:, 1]) + + raw_com_before = ( + wp.to_torch(collection.root_view.get_attribute(TT.BODY_COM_POSE)).reshape(3, 2, 7).transpose(0, 1).clone() + ) + coms = collection.data.body_com_pose_b.torch[env_ids][:, body_ids].clone() + coms[..., :3] = torch.tensor( + [[[0.01, 0.02, 0.03], [-0.01, 0.03, 0.02]], [[0.02, -0.01, 0.01], [0.03, 0.01, -0.02]]], device=device + ) + collection.set_coms_index(coms=coms, env_ids=env_ids, body_ids=body_ids) + expected_raw_com = raw_com_before.clone() + expected_raw_com[cpu_env_ids[:, None], cpu_body_ids[None, :]] = coms.cpu() + raw_com = wp.to_torch(collection.root_view.get_attribute(TT.BODY_COM_POSE)).reshape(3, 2, 7).transpose(0, 1) + torch.testing.assert_close(raw_com, expected_raw_com) + torch.testing.assert_close(collection.data.body_com_pose_b.torch[env_ids][:, body_ids], coms) + + raw_inertia_before = ( + wp.to_torch(collection.root_view.get_attribute(TT.BODY_INERTIA)).reshape(3, 2, 9).transpose(0, 1).clone() + ) + inertias = collection.data.body_inertia.torch[env_ids][:, body_ids].clone() + inertias[..., 0] *= 1.2 + inertias[..., 4] *= 1.3 + collection.set_inertias_index(inertias=inertias, env_ids=env_ids, body_ids=body_ids) + expected_raw_inertia = raw_inertia_before.clone() + expected_raw_inertia[cpu_env_ids[:, None], cpu_body_ids[None, :]] = inertias.cpu() + raw_inertia = wp.to_torch(collection.root_view.get_attribute(TT.BODY_INERTIA)).reshape(3, 2, 9).transpose(0, 1) + torch.testing.assert_close(raw_inertia, expected_raw_inertia) + torch.testing.assert_close(collection.data.body_inertia.torch[env_ids][:, body_ids], inertias) + + materials = torch.tensor( + [ + [[0.9, 0.4, 0.1], [0.8, 0.3, 0.2], [0.7, 0.2, 0.3]], + [[0.6, 0.1, 0.4], [0.5, 0.2, 0.1], [0.4, 0.3, 0.2]], + ] + ) + fused_materials = collection.reshape_data_to_view_3d(wp.from_torch(materials, dtype=wp.float32), 3, device="cpu") + collection.root_view.set_attribute(TT.RIGID_BODY_SHAPE_FRICTION_AND_RESTITUTION, fused_materials) + raw_materials = wp.to_torch(collection.root_view.get_attribute(TT.RIGID_BODY_SHAPE_FRICTION_AND_RESTITUTION)) + torch.testing.assert_close(raw_materials.reshape(3, 2, 3).transpose(0, 1), materials) - # check center of mass has been set - torch.testing.assert_close(cube_object.data.body_com_pose_b.torch, com) - rand_state = torch.rand(num_cubes, 13, device=device) - rand_state[..., :3] += env_pos - # make quaternion a unit vector - rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) +@pytest.mark.parametrize("scene", _CUDA_SCENES, indirect=True) +def test_wrench_reaches_only_the_selected_collection_body(scene: _RigidScene) -> None: + """Deliver an external force through the fused body-major binding to the selected body only, then reset it. - env_idx = env_idx.to(device) - for i in range(10): - # perform step - sim.step() - # update buffers - cube_object.update(sim.cfg.dt) - - write_env_ids = None if i % 2 == 0 else env_idx - if state_location == "com": - cube_object.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=write_env_ids) - cube_object.write_root_com_velocity_to_sim_index( - root_velocity=rand_state[..., 7:], env_ids=write_env_ids - ) - elif state_location == "link": - cube_object.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=write_env_ids) - cube_object.write_root_link_velocity_to_sim_index( - root_velocity=rand_state[..., 7:], env_ids=write_env_ids - ) - elif state_location == "root": - cube_object.write_root_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=write_env_ids) - cube_object.write_root_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=write_env_ids) - - root_link_pose_w = cube_object.data.root_link_pose_w.torch - root_link_vel_w = cube_object.data.root_link_vel_w.torch - root_com_pose_w = cube_object.data.root_com_pose_w.torch - root_com_vel_w = cube_object.data.root_com_vel_w.torch - body_com_pose_b = cube_object.data.body_com_pose_b.torch - if state_location == "com": - torch.testing.assert_close(rand_state[..., :7], root_com_pose_w) - torch.testing.assert_close(rand_state[..., 7:], root_com_vel_w) - expected_root_link_pos, expected_root_link_quat = combine_frame_transforms( - root_com_pose_w[:, :3], - root_com_pose_w[:, 3:], - quat_rotate(quat_inv(body_com_pose_b[:, 0, 3:7]), -body_com_pose_b[:, 0, :3]), - quat_inv(body_com_pose_b[:, 0, 3:7]), - ) - expected_root_link_pose = torch.cat((expected_root_link_pos, expected_root_link_quat), dim=1) - # test root_link successfully updated when root_com updates - torch.testing.assert_close(expected_root_link_pose, root_link_pose_w) - else: - # the "root" setters write the link pose and the COM velocity - torch.testing.assert_close(rand_state[..., :7], root_link_pose_w) - written_vel = root_link_vel_w if state_location == "link" else root_com_vel_w - torch.testing.assert_close(rand_state[..., 7:], written_vel) - expected_com_pos, expected_com_quat = combine_frame_transforms( - root_link_pose_w[:, :3], - root_link_pose_w[:, 3:], - body_com_pose_b[:, 0, :3], - body_com_pose_b[:, 0, 3:7], - ) - expected_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1) - # test root_com successfully updated when root_link updates - torch.testing.assert_close(expected_com_pose, root_com_pose_w) - # skip lin_vel because it differs from link frame, this should be fine because we are only checking - # if velocity update is triggered, which can be determined by comparing angular velocity - torch.testing.assert_close(root_com_vel_w[:, 3:], root_link_vel_w[:, 3:]) + Environment 1, body 1 has different flat indices in body-major and environment-major layouts. + """ + collection, sim, device = scene.wrenched, scene.sim, scene.device + velocity_before = collection.data.body_com_vel_w.torch.clone() + collection.permanent_wrench_composer.set_forces_and_torques_index( + forces=torch.tensor([[[20.0, 0.0, 0.0]]], device=device), + torques=torch.zeros((1, 1, 3), device=device), + env_ids=torch.tensor([1], dtype=torch.int32, device=device), + body_ids=torch.tensor([1], dtype=torch.int32, device=device), + ) + collection.write_data_to_sim() + sim.step() + collection.update(sim.cfg.dt) + + velocity = collection.data.body_com_vel_w.torch + assert velocity[1, 1, 0] > velocity_before[1, 1, 0] + 1e-3 + unselected = torch.ones((_NUM_ENVS, _NUM_BODIES), dtype=torch.bool, device=device) + unselected[1, 1] = False + torch.testing.assert_close(velocity[unselected], velocity_before[unselected], atol=1e-6, rtol=0.0) + + # Reset clears both wrench composers. + ones = torch.ones((_NUM_ENVS, _NUM_BODIES, 3), device=device) + collection.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones) + collection.reset() + for composer in (collection.instantaneous_wrench_composer, collection.permanent_wrench_composer): + assert not composer.active + assert torch.count_nonzero(composer.composed_force.torch) == 0 + assert torch.count_nonzero(composer.composed_torque.torch) == 0 + + +@pytest.mark.parametrize("scene", _CUDA_SCENES, indirect=True) +@pytest.mark.xfail( + strict=True, + raises=AssertionError, + reason=( + "OVPhysX RigidObjectCollection partial body pose writes push whole environment rows from a pose buffer " + "that is not refreshed first, so unselected bodies of a selected environment are rewound to stale poses." + ), +) +def test_partial_body_pose_write_preserves_unselected_bodies_after_steps(scene: _RigidScene) -> None: + """A partial (environment, body) pose write must not rewind the unselected bodies of that environment.""" + collection, sim, device = scene.stepped, scene.sim, scene.device + + def read_backend_poses() -> torch.Tensor: + raw = wp.to_torch(collection.root_view.get_attribute(TT.RIGID_BODY_POSE)) + return raw.reshape(_NUM_BODIES, _NUM_ENVS, 7).transpose(0, 1).to(device) + + # Read the public poses once, then let a force move every body for a few steps. + collection.data.body_link_pose_w + collection.permanent_wrench_composer.set_forces_and_torques_index( + forces=torch.tensor([10.0, 0.0, 0.0], device=device).repeat(_NUM_ENVS, _NUM_BODIES, 1), + torques=torch.zeros((_NUM_ENVS, _NUM_BODIES, 3), device=device), + ) + for _ in range(5): + collection.write_data_to_sim() + sim.step() + collection.update(sim.cfg.dt) + collection.reset() + + expected = read_backend_poses().clone() + body_pose = expected[1:, 1:2].clone() + body_pose[..., 2] += 0.5 + expected[1, 1] = body_pose[0, 0] + collection.write_body_link_pose_to_sim_index( + body_poses=body_pose, + env_ids=torch.tensor([1], dtype=torch.int32, device=device), + body_ids=torch.tensor([1], dtype=torch.int32, device=device), + ) + torch.testing.assert_close(read_backend_poses(), expected) diff --git a/source/isaaclab_ov/test/assets/test_rigid_object_collection.py b/source/isaaclab_ov/test/assets/test_rigid_object_collection.py deleted file mode 100644 index 34c7a8848222..000000000000 --- a/source/isaaclab_ov/test/assets/test_rigid_object_collection.py +++ /dev/null @@ -1,779 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -# ignore private usage of variables warning -# pyright: reportPrivateUsage=none - - -"""Real-backend tests for the OVPhysX RigidObjectCollection. - -Run via ``./scripts/run_ovphysx.sh -m pytest`` (kitless, no the Kit launcher). - -The OVPhysX runtime fixes device mode (CPU vs GPU) when the process creates -its first ``ovphysx.PhysX`` instance and cannot switch it without a process -restart. Full coverage therefore requires two separate pytest -invocations -- once with ``-k 'cpu'`` and once with ``-k 'cuda:0'``. The -``_ovphysx_skip_other_device`` autouse fixture below preempts the manager's -:exc:`RuntimeError` by ``pytest.skip``-ing on the unlocked device so -single-device runs finish cleanly. -""" - -from __future__ import annotations - -import sys - -import pytest -import torch -import warp as wp - -from isaaclab.test.utils import test_devices - -# The OVPhysX runtime wheel is optional. Skip gracefully when it is not installed; -# CI jobs that need OVPhysX coverage install it explicitly. -pytest.importorskip("ovphysx.types", reason="ovphysx wheel not installed") - -from isaaclab_ov import tensor_types as TT # noqa: E402 -from isaaclab_ov.assets import RigidObjectCollection # noqa: E402 -from isaaclab_ov.physics import OvPhysxCfg # noqa: E402 - -import isaaclab.sim as sim_utils # noqa: E402 -from isaaclab.assets import RigidObjectCfg, RigidObjectCollectionCfg # noqa: E402 -from isaaclab.sim import SimulationCfg, build_simulation_context # noqa: E402 -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR # noqa: E402 -from isaaclab.utils.math import ( # noqa: E402 - combine_frame_transforms, - default_orientation, - quat_apply_inverse, - quat_inv, - quat_mul, - quat_rotate, - random_orientation, - subtract_frame_transforms, -) - -wp.init() - -pytestmark = pytest.mark.device_split - - -_LOCKED_DEVICE: list[str | None] = [None] -"""Device the session pins to on the first parametrized test that runs.""" - - -@pytest.fixture(autouse=True) -def _ovphysx_skip_other_device(request): - """Skip parametrized tests on the device the session is not pinned to. - - See the module docstring for the wheel's process-global device-mode lock. - """ - callspec = getattr(request.node, "callspec", None) - device = callspec.params.get("device") if callspec is not None else None - if device is None: - # Test does not parametrize on device. - return - locked = _LOCKED_DEVICE[0] - if locked is None: - _LOCKED_DEVICE[0] = device - return - if device != locked: - pytest.skip( - f"ovphysx process-global device lock is held by '{locked}'; cannot run '{device}' " - "tests in the same session. Run pytest twice (once per device) for full coverage." - ) - - -def _ovphysx_sim_context(device: str, **kwargs): - """Wrapper around :func:`build_simulation_context` that injects OVPhysX cfg. - - PhysX tests pass ``device=device`` directly and let - :func:`build_simulation_context` build a default :class:`SimulationCfg`. - OVPhysX needs ``physics=OvPhysxCfg()`` set on the cfg so the manager - dispatches to OVPhysX rather than PhysX, so we build the cfg here and - pass it through. ``gravity_enabled`` is consumed locally (it is ignored - by ``build_simulation_context`` once a ``sim_cfg`` is provided). - ``add_ground_plane``, ``auto_add_lighting``, and other kwargs continue - to flow through ``build_simulation_context`` as before. - """ - dt = kwargs.pop("dt", 1.0 / 60.0) - gravity_enabled = kwargs.pop("gravity_enabled", True) - gravity = (0.0, 0.0, -9.81) if gravity_enabled else (0.0, 0.0, 0.0) - sim_cfg = SimulationCfg(physics=OvPhysxCfg(), device=device, dt=dt, gravity=gravity) - return build_simulation_context(device=device, sim_cfg=sim_cfg, **kwargs) - - -def generate_cubes_scene( - num_envs: int = 1, - num_cubes: int = 1, - height=1.0, - has_api: bool = True, - kinematic_enabled: bool = False, - device: str = "cuda:0", -) -> tuple[RigidObjectCollection, torch.Tensor]: - """Generate a scene with the provided number of cubes. - - Args: - num_envs: Number of envs to generate. - num_cubes: Number of cubes to generate. - height: Height of the cubes. - has_api: Whether the cubes have a rigid body API on them. - kinematic_enabled: Whether the cubes are kinematic. - device: Device to use for the simulation. - - Returns: - A tuple containing the rigid object representing the cubes and the origins of the cubes. - - """ - origins = torch.tensor([(i * 3.0, 0, height) for i in range(num_envs)]).to(device) - # Create Top-level Xforms, one for each cube - for i, origin in enumerate(origins): - sim_utils.create_prim(f"/World/Table_{i}", "Xform", translation=origin) - - # Resolve spawn configuration - if has_api: - spawn_cfg = sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", - rigid_props=sim_utils.UsdPhysicsRigidBodyCfg(kinematic_enabled=kinematic_enabled), - ) - else: - # since no rigid body properties defined, this is just a static collider - spawn_cfg = sim_utils.CuboidCfg( - size=(0.1, 0.1, 0.1), - collision_props=sim_utils.UsdPhysicsCollisionCfg(), - ) - - # create the rigid object configs. OVPhysX matches prim paths via fnmatch globs (not regex), - # so use ``Table_*`` rather than the PhysX ``Table_.*`` form. - cube_config_dict = {} - for i in range(num_cubes): - cube_object_cfg = RigidObjectCfg( - prim_path=f"/World/Table_[^/]+/Object_{i}", - spawn=spawn_cfg, - init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 3 * i, height)), - ) - cube_config_dict[f"cube_{i}"] = cube_object_cfg - # create the rigid object collection - cube_object_collection_cfg = RigidObjectCollectionCfg(rigid_objects=cube_config_dict) - cube_object_colection = RigidObjectCollection(cfg=cube_object_collection_cfg) - - return cube_object_colection, origins - - -@pytest.mark.parametrize("device", test_devices()) -def test_id_conversion(device): - """Test environment and object index conversion to physics view indices.""" - with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: - object_collection, _ = generate_cubes_scene(num_envs=2, num_cubes=3, device=device) - - # Play sim - sim.reset() - - expected = [ - torch.tensor([4, 5], device=device, dtype=torch.int32), - torch.tensor([4], device=device, dtype=torch.int32), - torch.tensor([0, 2, 4], device=device, dtype=torch.int32), - torch.tensor([1, 3, 5], device=device, dtype=torch.int32), - ] - - torch_all_env_indices = wp.to_torch(object_collection._ALL_ENV_INDICES) - torch_all_body_indices = wp.to_torch(object_collection._ALL_BODY_INDICES) - - view_ids = object_collection._env_body_ids_to_view_ids( - torch_all_env_indices, torch_all_body_indices[None, 2], device=device - ) - assert (wp.to_torch(view_ids) == expected[0]).all() - view_ids = object_collection._env_body_ids_to_view_ids( - torch_all_env_indices[None, 0], torch_all_body_indices[None, 2], device=device - ) - assert (wp.to_torch(view_ids) == expected[1]).all() - view_ids = object_collection._env_body_ids_to_view_ids( - torch_all_env_indices[None, 0], torch_all_body_indices, device=device - ) - assert (wp.to_torch(view_ids) == expected[2]).all() - view_ids = object_collection._env_body_ids_to_view_ids( - torch_all_env_indices[None, 1], torch_all_body_indices, device=device - ) - assert (wp.to_torch(view_ids) == expected[3]).all() - - -@pytest.mark.parametrize("num_envs", [2]) -@pytest.mark.parametrize("num_cubes", [3]) -@pytest.mark.parametrize("device", test_devices()) -def test_initialization_with_kinematic_enabled(num_envs, num_cubes, device): - """Test that initialization for prim with kinematic flag enabled.""" - with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: - object_collection, origins = generate_cubes_scene( - num_envs=num_envs, num_cubes=num_cubes, kinematic_enabled=True, device=device - ) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(object_collection) < 10 - - # Play sim - sim.reset() - - # Check if object is initialized - assert object_collection.is_initialized - assert len(object_collection.body_names) == num_cubes - - # Check buffers that exist and have correct shapes - assert object_collection.data.body_link_pos_w.torch.shape == (num_envs, num_cubes, 3) - assert object_collection.data.body_link_quat_w.torch.shape == (num_envs, num_cubes, 4) - - # Simulate physics - for _ in range(2): - sim.step() - object_collection.update(sim.cfg.dt) - # check that the object is kinematic - default_body_pose = object_collection.data.default_body_pose.torch.clone() - default_body_vel = object_collection.data.default_body_vel.torch.clone() - default_body_pose[..., :3] += origins.unsqueeze(1) - torch.testing.assert_close(object_collection.data.body_link_pose_w.torch, default_body_pose) - torch.testing.assert_close(object_collection.data.body_link_vel_w.torch, default_body_vel) - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_initialization_with_no_rigid_body(num_cubes, device): - """Test that initialization fails when no rigid body is found at the provided prim path.""" - with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: - object_collection, _ = generate_cubes_scene(num_cubes=num_cubes, has_api=False, device=device) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(object_collection) < 10 - - # Play sim - with pytest.raises(RuntimeError, match="Expected 1 prims"): - sim.reset() - - -@pytest.mark.parametrize("num_envs", [2]) -@pytest.mark.parametrize("num_cubes", [4]) -@pytest.mark.parametrize("device", test_devices()) -def test_external_force_on_single_body(num_envs, num_cubes, device): - """Test application of external force on the base of the object.""" - with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: - object_collection, origins = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) - sim.reset() - - # find objects to apply the force - object_ids, object_names = object_collection.find_bodies(".*") - - # Sample a force equal to the weight of the object - external_wrench_b = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=sim.device) - # Every 2nd cube should have a force applied to it - external_wrench_b[:, 0::2, 2] = 9.81 * object_collection.data.body_mass.torch[:, 0::2] - - for i in range(5): - # reset object state - body_pose = object_collection.data.default_body_pose.torch.clone() - body_vel = object_collection.data.default_body_vel.torch.clone() - # need to shift the position of the cubes otherwise they will be on top of each other - body_pose[..., :2] += origins.unsqueeze(1)[..., :2] - object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) - object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) - # reset object - object_collection.reset() - - is_global = False - if i % 2 == 0: - positions = object_collection.data.body_link_pos_w.torch[:, object_ids, :3] - is_global = True - else: - positions = None - - # apply force - object_collection.permanent_wrench_composer.set_forces_and_torques_index( - forces=external_wrench_b[..., :3], - torques=external_wrench_b[..., 3:], - positions=positions, - body_ids=object_ids, - env_ids=None, - is_global=is_global, - ) - for _ in range(10): - # write data to sim - object_collection.write_data_to_sim() - # step sim - sim.step() - # update object collection - object_collection.update(sim.cfg.dt) - - # First object should still be at the same Z position (1.0) - torch.testing.assert_close( - object_collection.data.body_link_pos_w.torch[:, 0::2, 2], - torch.ones_like(object_collection.data.body_link_pos_w.torch[:, 0::2, 2]), - ) - # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(object_collection.data.body_link_pos_w.torch[:, 1::2, 2] < 1.0) - - -@pytest.mark.parametrize("num_envs", [2]) -@pytest.mark.parametrize("num_cubes", [4]) -@pytest.mark.parametrize("device", test_devices()) -def test_external_force_on_single_body_at_position(num_envs, num_cubes, device): - """Test application of external force on the base of the object at a specific position. - - In this test, we apply a force equal to the weight of an object on the base of - one of the objects at 1m in the Y direction, we check that the object rotates around it's X axis. - For the other object, we do not apply any force and check that it falls down. - """ - with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: - object_collection, origins = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) - sim.reset() - - # find objects to apply the force - object_ids, object_names = object_collection.find_bodies(".*") - - # Sample a force equal to the weight of the object - external_wrench_b = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=sim.device) - external_wrench_positions_b = torch.zeros( - object_collection.num_instances, len(object_ids), 3, device=sim.device - ) - # Every 2nd cube should have a force applied to it - external_wrench_b[:, 0::2, 2] = 500.0 - external_wrench_positions_b[:, 0::2, 1] = 1.0 - - # Desired force and torque - for i in range(5): - # reset object state - body_pose = object_collection.data.default_body_pose.torch.clone() - body_vel = object_collection.data.default_body_vel.torch.clone() - # need to shift the position of the cubes otherwise they will be on top of each other - body_pose[..., :2] += origins.unsqueeze(1)[..., :2] - object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) - object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) - # reset object - object_collection.reset() - - is_global = False - if i % 2 == 0: - body_com_pos_w = object_collection.data.body_link_pos_w.torch[:, object_ids, :3] - external_wrench_positions_b[..., 0] = 0.0 - external_wrench_positions_b[..., 1] = 1.0 - external_wrench_positions_b[..., 2] = 0.0 - external_wrench_positions_b += body_com_pos_w - is_global = True - else: - external_wrench_positions_b[..., 0] = 0.0 - external_wrench_positions_b[..., 1] = 1.0 - external_wrench_positions_b[..., 2] = 0.0 - - # apply force - object_collection.permanent_wrench_composer.set_forces_and_torques_index( - forces=external_wrench_b[..., :3], - torques=external_wrench_b[..., 3:], - positions=external_wrench_positions_b, - body_ids=object_ids, - env_ids=None, - is_global=is_global, - ) - object_collection.permanent_wrench_composer.add_forces_and_torques_index( - forces=external_wrench_b[..., :3], - torques=external_wrench_b[..., 3:], - positions=external_wrench_positions_b, - body_ids=object_ids, - is_global=is_global, - ) - - for _ in range(10): - # write data to sim - object_collection.write_data_to_sim() - # step sim - sim.step() - # update object collection - object_collection.update(sim.cfg.dt) - - # First object should be rotating around it's X axis - assert torch.all(object_collection.data.body_com_ang_vel_b.torch[:, 0::2, 0] > 0.1) - # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(object_collection.data.body_link_pos_w.torch[:, 1::2, 2] < 1.0) - - -@pytest.mark.parametrize("num_envs", [3]) -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("gravity_enabled", [False]) -def test_set_object_state(num_envs, num_cubes, device, gravity_enabled): - """Test setting the state of the object. - - .. note:: - Turn off gravity for this test as we don't want any external forces acting on the object - to ensure state remains static - """ - with _ovphysx_sim_context(device=device, gravity_enabled=gravity_enabled, auto_add_lighting=True) as sim: - object_collection, origins = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) - sim.reset() - - state_types = ["body_link_pos_w", "body_link_quat_w", "body_com_lin_vel_w", "body_com_ang_vel_w"] - - # Set each state type individually as they are dependent on each other - for state_type_to_randomize in state_types: - state_dict = { - "body_link_pos_w": torch.zeros_like(object_collection.data.body_link_pos_w.torch, device=sim.device), - "body_link_quat_w": default_orientation(num=num_cubes * num_envs, device=sim.device).view( - num_envs, num_cubes, 4 - ), - "body_com_lin_vel_w": torch.zeros_like( - object_collection.data.body_com_lin_vel_w.torch, device=sim.device - ), - "body_com_ang_vel_w": torch.zeros_like( - object_collection.data.body_com_ang_vel_w.torch, device=sim.device - ), - } - - for _ in range(5): - # reset object - object_collection.reset() - - # Set random state - if state_type_to_randomize == "body_link_quat_w": - state_dict[state_type_to_randomize] = random_orientation( - num=num_cubes * num_envs, device=sim.device - ).view(num_envs, num_cubes, 4) - else: - state_dict[state_type_to_randomize] = torch.randn(num_envs, num_cubes, 3, device=sim.device) - # make sure objects do not overlap - if state_type_to_randomize == "body_link_pos_w": - state_dict[state_type_to_randomize][..., :2] += origins.unsqueeze(1)[..., :2] - - # perform simulation - for _ in range(5): - body_pose = torch.cat( - [state_dict["body_link_pos_w"], state_dict["body_link_quat_w"]], - dim=-1, - ) - body_vel = torch.cat( - [state_dict["body_com_lin_vel_w"], state_dict["body_com_ang_vel_w"]], - dim=-1, - ) - # reset object state - object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) - object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) - sim.step() - - # assert that set object quantities are equal to the ones set in the state_dict - for key, expected_value in state_dict.items(): - value = getattr(object_collection.data, key).torch - torch.testing.assert_close(value, expected_value, rtol=1e-5, atol=1e-5) - - object_collection.update(sim.cfg.dt) - - -@pytest.mark.parametrize("num_envs", [4]) -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("gravity_enabled", [False]) -def test_object_state_properties(num_envs, num_cubes, device, gravity_enabled): - """Test the object_com_state_w and object_link_state_w properties.""" - with _ovphysx_sim_context(device=device, gravity_enabled=gravity_enabled, auto_add_lighting=True) as sim: - cube_object, env_pos = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, height=0.0, device=device) - env_ids = torch.tensor([x for x in range(num_envs)], dtype=torch.int32) - - sim.reset() - - # check if cube_object is initialized - assert cube_object.is_initialized - - # change center of mass offset from link frame - offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) - - # Read current COMs, mutate the translation, write back via the OVPhysX - # ``set_coms_index`` setter (PhysX uses ``root_view.set_coms`` + reshape helpers - # for the same operation; OVPhysX wraps the wheel RIGID_BODY_COM_POSE write in - # :meth:`set_coms_index`). - com = cube_object.data.body_com_pose_b.torch.clone() # shape (num_envs, num_cubes, 7) - com[..., :3] = offset.to(com.device) - cube_object.set_coms_index( - coms=wp.from_torch(com.contiguous(), dtype=wp.transformf), - env_ids=wp.from_torch(env_ids, dtype=wp.int32), - ) - - # check center of mass has been set - torch.testing.assert_close(cube_object.data.body_com_pose_b.torch, com) - - # random z spin velocity - spin_twist = torch.zeros(6, device=device) - spin_twist[5] = torch.randn(1, device=device) - - # initial spawn point - init_com = cube_object.data.body_com_pose_w.torch[..., :3] - - for i in range(10): - # spin the object around Z axis (com) - cube_object.write_body_com_velocity_to_sim_index(body_velocities=spin_twist.repeat(num_envs, num_cubes, 1)) - sim.step() - cube_object.update(sim.cfg.dt) - - # get state properties - object_link_pose_w = cube_object.data.body_link_pose_w.torch - object_link_vel_w = cube_object.data.body_link_vel_w.torch - object_com_pose_w = cube_object.data.body_com_pose_w.torch - object_com_vel_w = cube_object.data.body_com_vel_w.torch - - # cubes are spinning around center of mass - # position will not match - # center of mass position will be constant (i.e. spinning around com) - torch.testing.assert_close(init_com, object_com_pose_w[..., :3]) - - # link position will be moving but should stay constant away from center of mass - object_link_state_pos_rel_com = quat_apply_inverse( - object_link_pose_w[..., 3:], - object_link_pose_w[..., :3] - object_com_pose_w[..., :3], - ) - - torch.testing.assert_close(-offset, object_link_state_pos_rel_com) - - # orientation of com will be a constant rotation from link orientation - com_quat_b = cube_object.data.body_com_quat_b.torch - com_quat_w = quat_mul(object_link_pose_w[..., 3:], com_quat_b) - torch.testing.assert_close(com_quat_w, object_com_pose_w[..., 3:]) - - # lin_vel will not match - # center of mass vel will be constant (i.e. spinning around com) - torch.testing.assert_close( - torch.zeros_like(object_com_vel_w[..., :3]), - object_com_vel_w[..., :3], - ) - - # link frame will be moving, and should be equal to input angular velocity cross offset - lin_vel_rel_object_gt = quat_apply_inverse(object_link_pose_w[..., 3:], object_link_vel_w[..., :3]) - lin_vel_rel_gt = torch.linalg.cross(spin_twist.repeat(num_envs, num_cubes, 1)[..., 3:], -offset) - torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_object_gt, atol=1e-4, rtol=1e-3) - - # ang_vel will always match - torch.testing.assert_close(object_com_vel_w[..., 3:], object_link_vel_w[..., 3:]) - - -@pytest.mark.parametrize("num_envs", [3]) -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("state_location", ["com", "link", "root"]) -@pytest.mark.parametrize("gravity_enabled", [False]) -def test_write_object_state(num_envs, num_cubes, device, state_location, gravity_enabled): - """Test the object-state setters in the link, center-of-mass, and root reference frames. - - The written frame reads back unchanged and the other frame follows it through the link-to-COM - transform measured from the simulation before the write. - """ - with _ovphysx_sim_context(device=device, gravity_enabled=gravity_enabled, auto_add_lighting=True) as sim: - # Create a scene with random cubes - cube_object, env_pos = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, height=0.0, device=device) - env_ids = torch.tensor([x for x in range(num_envs)], dtype=torch.int32) - object_ids = torch.tensor([x for x in range(num_cubes)], dtype=torch.int32) - - sim.reset() - - # Check if cube_object is initialized - assert cube_object.is_initialized - - # change center of mass offset from link frame - offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) - - com = cube_object.data.body_com_pose_b.torch.clone() # shape (num_envs, num_cubes, 7) - com[..., :3] = offset.to(com.device) - cube_object.set_coms_index( - coms=wp.from_torch(com.contiguous(), dtype=wp.transformf), - env_ids=wp.from_torch(env_ids, dtype=wp.int32), - ) - # check center of mass has been set - torch.testing.assert_close(cube_object.data.body_com_pose_b.torch, com) - - rand_state = torch.rand(num_envs, num_cubes, 13, device=device) - rand_state[..., :3] += cube_object.data.body_link_pos_w.torch - # make quaternion a unit vector - rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) - - env_ids = env_ids.to(device) - object_ids = object_ids.to(device) - for i in range(10): - sim.step() - cube_object.update(sim.cfg.dt) - - body_link_pose_w = cube_object.data.body_link_pose_w.torch - body_com_pose_w = cube_object.data.body_com_pose_w.torch - object_link_to_com_pos, object_link_to_com_quat = subtract_frame_transforms( - body_link_pose_w[..., :3].view(-1, 3), - body_link_pose_w[..., 3:7].view(-1, 4), - body_com_pose_w[..., :3].view(-1, 3), - body_com_pose_w[..., 3:7].view(-1, 4), - ) - - ids = {} if i % 2 == 0 else {"env_ids": env_ids, "body_ids": object_ids} - if state_location == "com": - cube_object.write_body_com_pose_to_sim_index(body_poses=rand_state[..., :7], **ids) - cube_object.write_body_com_velocity_to_sim_index(body_velocities=rand_state[..., 7:], **ids) - elif state_location == "link": - cube_object.write_body_link_pose_to_sim_index(body_poses=rand_state[..., :7], **ids) - cube_object.write_body_link_velocity_to_sim_index(body_velocities=rand_state[..., 7:], **ids) - elif state_location == "root": - cube_object.write_body_link_pose_to_sim_index(body_poses=rand_state[..., :7], **ids) - cube_object.write_body_com_velocity_to_sim_index(body_velocities=rand_state[..., 7:], **ids) - - link_pose_w = cube_object.data.body_link_pose_w.torch - link_vel_w = cube_object.data.body_link_vel_w.torch - com_pose_w = cube_object.data.body_com_pose_w.torch - com_vel_w = cube_object.data.body_com_vel_w.torch - if state_location == "com": - torch.testing.assert_close(rand_state[..., :7], com_pose_w) - torch.testing.assert_close(rand_state[..., 7:], com_vel_w) - expected_link_pos, expected_link_quat = combine_frame_transforms( - com_pose_w[..., :3].view(-1, 3), - com_pose_w[..., 3:].view(-1, 4), - quat_rotate(quat_inv(object_link_to_com_quat), -object_link_to_com_pos), - quat_inv(object_link_to_com_quat), - ) - expected_link_pose = torch.cat((expected_link_pos, expected_link_quat), dim=1).view(num_envs, -1, 7) - # test body_link successfully updated when body_com updates - torch.testing.assert_close(expected_link_pose, link_pose_w) - else: - torch.testing.assert_close(rand_state[..., :7], link_pose_w) - written_vel = link_vel_w if state_location == "link" else com_vel_w - torch.testing.assert_close(rand_state[..., 7:], written_vel) - expected_com_pos, expected_com_quat = combine_frame_transforms( - link_pose_w[..., :3].view(-1, 3), - link_pose_w[..., 3:].view(-1, 4), - object_link_to_com_pos, - object_link_to_com_quat, - ) - expected_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1).view(num_envs, -1, 7) - # test body_com successfully updated when body_link updates - torch.testing.assert_close(expected_com_pose, com_pose_w) - # skip lin_vel because it differs from link frame, this should be fine because we are only checking - # if velocity update is triggered, which can be determined by comparing angular velocity - torch.testing.assert_close(com_vel_w[..., 3:], link_vel_w[..., 3:]) - - -@pytest.mark.parametrize("num_envs", [3]) -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_reset_object_collection(num_envs, num_cubes, device): - """Test resetting the state of the rigid object.""" - with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: - object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) - sim.reset() - - for i in range(5): - sim.step() - object_collection.update(sim.cfg.dt) - - # Move the object to a random position - body_pose = object_collection.data.default_body_pose.torch.clone() - body_pose[..., :3] = torch.randn(num_envs, num_cubes, 3, device=sim.device) - # Random orientation - body_pose[..., 3:7] = random_orientation(num=num_cubes, device=sim.device) - object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) - body_vel = object_collection.data.default_body_vel.torch.clone() - object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) - - if i % 2 == 0: - # Make both wrench composers active so the reset has something to clear. - object_collection.permanent_wrench_composer.set_forces_and_torques_index( - forces=torch.ones((num_envs, num_cubes, 3), device=device), - torques=torch.ones((num_envs, num_cubes, 3), device=device), - ) - object_collection.instantaneous_wrench_composer.add_forces_and_torques_index( - forces=torch.ones((num_envs, num_cubes, 3), device=device), - torques=torch.ones((num_envs, num_cubes, 3), device=device), - ) - assert object_collection._instantaneous_wrench_composer.active - assert object_collection._permanent_wrench_composer.active - - object_collection.reset() - - # Reset should zero external forces and torques - assert not object_collection._instantaneous_wrench_composer.active - assert not object_collection._permanent_wrench_composer.active - assert torch.count_nonzero(object_collection._instantaneous_wrench_composer.composed_force.torch) == 0 - assert torch.count_nonzero(object_collection._instantaneous_wrench_composer.composed_torque.torch) == 0 - assert torch.count_nonzero(object_collection._permanent_wrench_composer.composed_force.torch) == 0 - assert torch.count_nonzero(object_collection._permanent_wrench_composer.composed_torque.torch) == 0 - - -@pytest.mark.parametrize("num_envs", [3]) -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_set_material_properties(num_envs, num_cubes, device): - """Test getting and setting per-shape material properties of a rigid object collection. - - OVPhysX exposes per-collision-shape material as the - ``rigid_body_shape_friction_and_restitution`` tensor binding, addressed through the - fused :class:`~isaaclab_ov.sim.views.OvPhysxView`. The binding is body-major flat, - so data ``[num_envs, num_cubes, 3]`` is mapped to/from the view layout with - :meth:`~isaaclab_ov.assets.RigidObjectCollection.reshape_data_to_view_3d` and its - inverse. (The PhysX backend uses ``root_view.get/set_material_properties``.) - """ - with _ovphysx_sim_context(device=device, add_ground_plane=True, auto_add_lighting=True) as sim: - object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) - - # Play sim - sim.reset() - - # Random material per object: (static_friction, dynamic_friction, restitution), on the CPU. - static_friction = torch.empty(num_envs, num_cubes, 1, device="cpu").uniform_(0.4, 0.8) - dynamic_friction = torch.empty(num_envs, num_cubes, 1, device="cpu").uniform_(0.4, 0.8) - restitution = torch.empty(num_envs, num_cubes, 1, device="cpu").uniform_(0.0, 0.2) - materials = torch.cat([static_friction, dynamic_friction, restitution], dim=-1) # [num_envs, num_cubes, 3] - - # Map data -> body-major view layout, write through the view. - view_materials = object_collection.reshape_data_to_view_3d(wp.from_torch(materials, dtype=wp.float32), 3) - object_collection.root_view.set_attribute(TT.RIGID_BODY_SHAPE_FRICTION_AND_RESTITUTION, view_materials) - - # Simulate physics - sim.step() - object_collection.update(sim.cfg.dt) - - # Read back and map view -> data layout (inverse of reshape_data_to_view_3d), then verify. - view_check = wp.to_torch( - object_collection.root_view.get_attribute(TT.RIGID_BODY_SHAPE_FRICTION_AND_RESTITUTION) - ) - materials_check = view_check.reshape(num_cubes, num_envs, 3).transpose(0, 1).contiguous() - torch.testing.assert_close(materials_check, materials) - - -@pytest.mark.parametrize("num_envs", [3]) -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("gravity_enabled", [True, False]) -def test_gravity_vec_w(num_envs, num_cubes, device, gravity_enabled): - """Test initialization and that gravity vector direction is set correctly for the rigid object.""" - with _ovphysx_sim_context(device=device, gravity_enabled=gravity_enabled, auto_add_lighting=True) as sim: - object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(object_collection) < 10 - - # Obtain gravity direction - gravity_dir = (0.0, 0.0, -1.0) if gravity_enabled else (0.0, 0.0, 0.0) - - sim.reset() - - # Check if object is initialized - assert object_collection.is_initialized - assert len(object_collection.body_names) == num_cubes - - # Check buffers that exist and have correct shapes - assert object_collection.data.body_link_pos_w.torch.shape == (num_envs, num_cubes, 3) - assert object_collection.data.body_link_quat_w.torch.shape == (num_envs, num_cubes, 4) - assert object_collection.data.body_mass.torch.shape == (num_envs, num_cubes) - assert object_collection.data.body_inertia.torch.shape == (num_envs, num_cubes, 9) - - # Check if gravity vector is set correctly - gravity_vec = object_collection.data.GRAVITY_VEC_W.torch - assert gravity_vec[0, 0, 0] == gravity_dir[0] - assert gravity_vec[0, 0, 1] == gravity_dir[1] - assert gravity_vec[0, 0, 2] == gravity_dir[2] - - # Perform simulation - for _ in range(2): - sim.step() - object_collection.update(sim.cfg.dt) - - # Expected gravity value is the acceleration of the body - gravity = torch.zeros(num_envs, num_cubes, 6, device=device) - if gravity_enabled: - gravity[..., 2] = -9.81 - - # Check the body accelerations are correct - torch.testing.assert_close(object_collection.data.body_com_acc_w.torch, gravity) diff --git a/source/isaaclab_ov/test/physics/test_ovphysx_manager_lifecycle.py b/source/isaaclab_ov/test/physics/test_ovphysx_manager_lifecycle.py index 9b2094b9d665..10fba55987d9 100644 --- a/source/isaaclab_ov/test/physics/test_ovphysx_manager_lifecycle.py +++ b/source/isaaclab_ov/test/physics/test_ovphysx_manager_lifecycle.py @@ -16,8 +16,13 @@ import pytest +from isaaclab.test.utils import DeviceScope, test_devices + pytest.importorskip("ovphysx.types", reason="ovphysx wheel not installed") +_CPU_DEVICES = test_devices(DeviceScope.CPU) +_CUDA_DEVICES = test_devices(DeviceScope.CUDA) + class _FakePhysXConfig: def __init__(self, num_threads=None, cooked_collider_cache_dir=None, carbonite_overrides=None): @@ -27,10 +32,6 @@ def __init__(self, num_threads=None, cooked_collider_cache_dir=None, carbonite_o class _FakePhysX: - @classmethod - def set_cpu_mode(cls, enabled): - pass - def __init__(self, active_cuda_gpus=None, config=None): self.active_cuda_gpus = active_cuda_gpus self.config = config @@ -63,7 +64,6 @@ def manager_module(monkeypatch): "_next_control_ordinal": 2, "_warmup_done": False, "_requires_full_stage": False, - "_locked_device": None, "_clone_recipes": [], "_atexit_registered": False, "_scene_data_backend": None, @@ -459,6 +459,66 @@ def report_unraisable(unraisable): ) +def _device_sequence_script(devices: tuple[str, ...]) -> str: + return f"DEVICES = {devices!r}\n" + textwrap.dedent( + """ + import torch + from ovphysx.dlpack import DLDeviceType + + import isaaclab.sim as sim_utils + from isaaclab.assets import RigidObjectCfg + from isaaclab.sim import SimulationCfg, build_simulation_context + from isaaclab_ov import tensor_types as TT + from isaaclab_ov.assets import RigidObject + from isaaclab_ov.physics import OvPhysxCfg + + def drop_cube(device): + sim_cfg = SimulationCfg(physics=OvPhysxCfg(), device=device, dt=1.0 / 60.0) + with build_simulation_context(device=device, sim_cfg=sim_cfg) as sim: + cube = RigidObject( + RigidObjectCfg( + prim_path="/World/Cube", + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 2.0)), + spawn=sim_utils.CuboidCfg( + size=(0.5, 0.5, 0.5), + rigid_props=sim_utils.RigidBodyBaseCfg(), + mass_props=sim_utils.MassPropertiesCfg(mass=1.0), + collision_props=sim_utils.CollisionBaseCfg(), + ), + ) + ) + sim.reset() + + # A CPU scene after a CUDA scene must not inherit its DirectGPU state bindings. + native_device = cube.root_view.binding_for(TT.RIGID_BODY_POSE).native_device + expected_type = DLDeviceType.kDLCUDA if device.startswith("cuda") else DLDeviceType.kDLCPU + assert native_device.device_type.value == expected_type, (device, str(native_device)) + + root_pose = cube.data.root_link_pose_w.torch.clone() + assert root_pose.device == torch.device(device) + root_pose[:, 0] = 0.25 + cube.write_root_link_pose_to_sim_index(root_pose=root_pose) + cube.update(sim.get_physics_dt()) + torch.testing.assert_close(cube.data.root_link_pose_w.torch, root_pose) + + heights = [] + for _ in range(10): + sim.step() + cube.update(sim.get_physics_dt()) + heights.append(cube.data.root_link_pose_w.torch[0, 2].item()) + return heights + + trajectories = {} + for device in DEVICES: + heights = drop_cube(device) + assert all(later < earlier for earlier, later in zip([2.0] + heights, heights)), (device, heights) + # A device's scene is reproducible regardless of the scenes that ran before it. + assert trajectories.setdefault(device, heights) == heights, (device, trajectories[device], heights) + print("DEVICE_SEQUENCE_OK", flush=True) + """ + ) + + def _run_child(script: str) -> tuple[subprocess.CompletedProcess[str], str]: completed = subprocess.run( [sys.executable, "-c", script], @@ -510,6 +570,31 @@ def test_scene_settings(monkeypatch, manager_module, device, rate, override, exp assert prim.GetAttribute("physxScene:enableSceneQuerySupport").Get() is sim_cfg.enable_scene_query_support +@pytest.mark.parametrize(("device", "gpu_dynamics", "broadphase"), [("cpu", False, "MBP"), ("gpu", True, "GPU")]) +def test_scenes_author_device_dynamics_and_broadphase(monkeypatch, manager_module, device, gpu_dynamics, broadphase): + """Every physics scene selects the simulation device, since CPU and GPU scenes share one process.""" + from isaaclab_ov.physics.ovphysx_manager_cfg import OvPhysxCfg + + from pxr import Sdf, Usd, UsdPhysics + + from isaaclab.physics import PhysicsManager + + monkeypatch.setattr( + PhysicsManager, "_sim", SimpleNamespace(cfg=SimpleNamespace(dt=1.0 / 60.0, enable_scene_query_support=False)) + ) + stage = Usd.Stage.CreateInMemory() + UsdPhysics.Scene.Define(stage, "/World/PhysicsScene") + # An asset's own scene, authored for the other device. + asset_scene = UsdPhysics.Scene.Define(stage, "/World/Asset/PhysicsScene").GetPrim() + asset_scene.CreateAttribute("physxScene:enableGPUDynamics", Sdf.ValueTypeNames.Bool).Set(not gpu_dynamics) + manager_module.OvPhysxManager._configure_physics_scenes(stage, "/World/PhysicsScene", OvPhysxCfg(), device) + for path in ("/World/PhysicsScene", "/World/Asset/PhysicsScene"): + prim = stage.GetPrimAtPath(path) + assert "PhysxSceneAPI" in prim.GetMetadata("apiSchemas").GetAddedOrExplicitItems() + assert prim.GetAttribute("physxScene:enableGPUDynamics").Get() is gpu_dynamics + assert prim.GetAttribute("physxScene:broadphaseType").Get() == broadphase + + def test_construct_physx_forwards_cooked_collider_cache_dir(monkeypatch, manager_module, tmp_path): """Configured, default, and unset cache directories reach ``PhysXConfig`` unchanged.""" from isaaclab_ov.physics.ovphysx_manager_cfg import DEFAULT_COOKED_COLLIDER_CACHE_DIR, OvPhysxBackendCfg, OvPhysxCfg @@ -559,3 +644,13 @@ def test_default_cache_dir_rejects_a_directory_owned_by_another_user(manager_mod with pytest.raises(RuntimeError, match="owned"): manager_module._prepare_default_cache_dir(str(target)) + + +@pytest.mark.skipif(not (_CPU_DEVICES and _CUDA_DEVICES), reason="The device sequence requires a CPU and a CUDA device") +def test_cpu_and_cuda_scenes_run_sequentially_in_one_process(): + """A CPU scene must not prevent a later CUDA scene in the same process, or the reverse.""" + cpu, cuda = _CPU_DEVICES[0], _CUDA_DEVICES[0] + completed, output = _run_child(_device_sequence_script((cpu, cuda, cpu))) + + assert completed.returncode == 0, output[-8000:] + assert "DEVICE_SEQUENCE_OK" in output, output[-8000:] diff --git a/source/isaaclab_ov/test/physics/test_ovphysx_scene_data_backend.py b/source/isaaclab_ov/test/physics/test_ovphysx_scene_data_backend.py index 9130b51116cc..5228f92205bc 100644 --- a/source/isaaclab_ov/test/physics/test_ovphysx_scene_data_backend.py +++ b/source/isaaclab_ov/test/physics/test_ovphysx_scene_data_backend.py @@ -81,14 +81,12 @@ def test_manager_forced_rewarm_invalidates_bindings_before_loading(monkeypatch): assert OvPhysxManager.kinematics_dirty -@pytest.mark.parametrize( - ("device", "expected_cpu_mode", "expected_active_cuda_gpus"), - [("cpu", True, None), ("cuda:2", False, "2")], -) -def test_manager_supports_pinned_runtime_api( - monkeypatch, tmp_path, device, expected_cpu_mode, expected_active_cuda_gpus -): - """The pinned OVPhysX wheel keeps its constructor, step, and reset API.""" +@pytest.mark.parametrize(("device", "expected_active_cuda_gpus"), [("cpu", None), ("cuda:2", "2")]) +def test_manager_supports_pinned_runtime_api(monkeypatch, tmp_path, device, expected_active_cuda_gpus): + """The pinned OVPhysX wheel keeps its constructor, step, and reset API. + + The runtime never enters the wheel's sticky process-wide CPU-only mode, which would block later CUDA scenes. + """ import isaaclab_ov.physics.ovphysx_manager as module from isaaclab_ov.physics import OvPhysxBackendCfg, OvPhysxManager @@ -141,11 +139,11 @@ def pinned_config(*, num_threads=None, cooked_collider_cache_dir=None, carbonite OvPhysxManager.step() OvPhysxManager._prepare_physx_for_stage_reuse() - assert PinnedPhysX.cpu_mode is expected_cpu_mode + assert PinnedPhysX.cpu_mode is None assert physx.constructor["active_cuda_gpus"] == expected_active_cuda_gpus assert physx.constructor["config"].num_threads == 8 assert physx.constructor["config"].cooked_collider_cache_dir == cache_dir - assert physx.constructor["config"].carbonite_overrides["/ovphysx/clone/useEnvIds"] is not expected_cpu_mode + assert physx.constructor["config"].carbonite_overrides["/ovphysx/clone/useEnvIds"] is device.startswith("cuda") updates = [("step_sync", 0.02), ("update_articulations_kinematic",)] assert physx.calls == updates + [("destroy_view",), ("reset_stage",), ("wait_op", 23)] assert OvPhysxManager.backend.rigid_body_view is None diff --git a/source/isaaclab_ov/test/sensors/test_contact_sensor.py b/source/isaaclab_ov/test/sensors/test_contact_sensor.py index 5dabbbbee0f9..63632b477fae 100644 --- a/source/isaaclab_ov/test/sensors/test_contact_sensor.py +++ b/source/isaaclab_ov/test/sensors/test_contact_sensor.py @@ -10,14 +10,6 @@ Run via ``uv run --extra ovphysx --extra test python -m pytest`` without a Kit runtime. -The OVPhysX runtime fixes device mode (CPU vs GPU) when the process creates -its first ``ovphysx.PhysX`` instance and cannot switch it without a process -restart. Full coverage therefore requires two separate pytest -invocations -- once with ``-k 'cpu'`` and once with ``-k 'cuda:0'``. The -``_ovphysx_skip_other_device`` autouse fixture below preempts the manager's -:exc:`RuntimeError` by ``pytest.skip``-ing on the unlocked device so -single-device runs finish cleanly. - Detailed-contact tests cover filtered positions and friction, optional tracking modes, force history, selective reset, multi-body indexing, and insufficient contact capacity. @@ -61,37 +53,6 @@ wp.init() -pytestmark = pytest.mark.device_split - -# --------------------------------------------------------------------------- -# Device-lock autouse fixture -# --------------------------------------------------------------------------- - -_LOCKED_DEVICE: list[str | None] = [None] -"""Device the session pins to on the first parametrized test that runs.""" - - -@pytest.fixture(autouse=True) -def _ovphysx_skip_other_device(request): - """Skip parametrized tests on the device the session is not pinned to. - - See the module docstring for the wheel's process-global device-mode lock. - """ - callspec = getattr(request.node, "callspec", None) - device = callspec.params.get("device") if callspec is not None else None - if device is None: - # Test does not parametrize on device. - return - locked = _LOCKED_DEVICE[0] - if locked is None: - _LOCKED_DEVICE[0] = device - return - if device != locked: - pytest.skip( - f"ovphysx process-global device lock is held by '{locked}'; cannot run '{device}' " - "tests in the same session. Run pytest twice (once per device) for full coverage." - ) - # --------------------------------------------------------------------------- # Simulation context helper diff --git a/source/isaaclab_ov/test/sensors/test_frame_transformer.py b/source/isaaclab_ov/test/sensors/test_frame_transformer.py index 185fb16143ea..d5cf816f469c 100644 --- a/source/isaaclab_ov/test/sensors/test_frame_transformer.py +++ b/source/isaaclab_ov/test/sensors/test_frame_transformer.py @@ -55,39 +55,6 @@ wp.init() -pytestmark = pytest.mark.device_split - - -# --------------------------------------------------------------------------- -# Device-lock autouse fixture (mirrors test_contact_sensor.py) # -# --------------------------------------------------------------------------- - -_LOCKED_DEVICE: list[str | None] = [None] -"""Device the session pins to on the first parametrized test that runs.""" - - -@pytest.fixture(autouse=True) -def _ovphysx_skip_other_device(request): - """Skip parametrized tests on the device the session is not pinned to. - - The OVPhysX runtime fixes device mode when the process creates its first - ``ovphysx.PhysX`` instance and cannot switch modes without a process - restart. Pin the session to whichever device is first used. - """ - callspec = getattr(request.node, "callspec", None) - device = callspec.params.get("device") if callspec is not None else None - if device is None: - 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." - ) - # --------------------------------------------------------------------------- # Simulation context helper # diff --git a/source/isaaclab_ov/test/sensors/test_imu.py b/source/isaaclab_ov/test/sensors/test_imu.py index 296149e3c67a..8ea15cae95c7 100644 --- a/source/isaaclab_ov/test/sensors/test_imu.py +++ b/source/isaaclab_ov/test/sensors/test_imu.py @@ -68,7 +68,6 @@ wp.init() -pytestmark = pytest.mark.device_split # --------------------------------------------------------------------------- # Constants @@ -182,33 +181,6 @@ class _StaleResetSceneCfg(InteractiveSceneCfg): imu_cube: ImuCfg = ImuCfg(prim_path="{ENV_REGEX_NS}/cube") -# --------------------------------------------------------------------------- -# Process-global device-mode lock (matches the rigid-object and contact-sensor -# tests). The ovphysx wheel can only run one device per process; parametrized -# tests skip on the unlocked device so single-device runs finish cleanly. -# --------------------------------------------------------------------------- - -_LOCKED_DEVICE: list[str | None] = [None] - - -@pytest.fixture(autouse=True) -def _ovphysx_skip_other_device(request): - """Skip parametrized tests on the device the session is not pinned to.""" - callspec = getattr(request.node, "callspec", None) - device = callspec.params.get("device") if callspec is not None else None - if device is None: - 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." - ) - - # --------------------------------------------------------------------------- # Sim context fixture (real OVPhysX backend, device-parametrized) # --------------------------------------------------------------------------- diff --git a/source/isaaclab_ov/test/sensors/test_joint_wrench_sensor.py b/source/isaaclab_ov/test/sensors/test_joint_wrench_sensor.py index e31bc732052f..ac4dc305ec98 100644 --- a/source/isaaclab_ov/test/sensors/test_joint_wrench_sensor.py +++ b/source/isaaclab_ov/test/sensors/test_joint_wrench_sensor.py @@ -11,13 +11,6 @@ Wrench values and frames are checked against analytic loads by the shared ``test_joint_wrench_frame`` contract imported below; the local tests cover initialization, body resolution, and reset behavior. - -The OVPhysX runtime fixes device mode (CPU vs GPU) when the process creates -its first ``ovphysx.PhysX`` instance. Full coverage therefore requires two -pytest runs -- once with ``-k 'cpu'`` and once with -``-k 'cuda:0'``. The ``_ovphysx_skip_other_device`` autouse fixture below -preempts the manager's :exc:`RuntimeError` by ``pytest.skip``-ing on the -unlocked device so single-device runs finish cleanly. """ from __future__ import annotations @@ -54,33 +47,6 @@ wp.init() -pytestmark = pytest.mark.device_split - -# --------------------------------------------------------------------------- -# Device-lock autouse fixture (copied from test_contact_sensor.py) -# --------------------------------------------------------------------------- - -_LOCKED_DEVICE: list[str | None] = [None] -"""Device the session pins to on the first parametrized test that runs.""" - - -@pytest.fixture(autouse=True) -def _ovphysx_skip_other_device(request): - """Skip parametrized tests on the device the session is not pinned to.""" - callspec = getattr(request.node, "callspec", None) - device = callspec.params.get("device") if callspec is not None else None - if device is None: - 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." - ) - # --------------------------------------------------------------------------- # Simulation context helper (mirrors test_contact_sensor.py) diff --git a/source/isaaclab_ov/test/sensors/test_pva.py b/source/isaaclab_ov/test/sensors/test_pva.py index c4e8cf1339f4..39dff9597416 100644 --- a/source/isaaclab_ov/test/sensors/test_pva.py +++ b/source/isaaclab_ov/test/sensors/test_pva.py @@ -55,7 +55,6 @@ wp.init() -pytestmark = pytest.mark.device_split # --------------------------------------------------------------------------- # Constants @@ -163,33 +162,6 @@ class _StaleResetSceneCfg(InteractiveSceneCfg): pva_cube: PvaCfg = PvaCfg(prim_path="{ENV_REGEX_NS}/cube") -# --------------------------------------------------------------------------- -# Process-global device-mode lock (matches the IMU and contact-sensor tests). -# The ovphysx wheel can only run one device per process; parametrized tests -# skip on the unlocked device so single-device runs finish cleanly. -# --------------------------------------------------------------------------- - -_LOCKED_DEVICE: list[str | None] = [None] - - -@pytest.fixture(autouse=True) -def _ovphysx_skip_other_device(request): - """Skip parametrized tests on the device the session is not pinned to.""" - callspec = getattr(request.node, "callspec", None) - device = callspec.params.get("device") if callspec is not None else None - if device is None: - 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." - ) - - # --------------------------------------------------------------------------- # Sim context fixture (real OVPhysX backend, device-parametrized) # --------------------------------------------------------------------------- diff --git a/source/isaaclab_ov/test/sim/test_views_xform_prim_ovphysx.py b/source/isaaclab_ov/test/sim/test_views_xform_prim_ovphysx.py index 826f490328e4..b08cbd83c0a6 100644 --- a/source/isaaclab_ov/test/sim/test_views_xform_prim_ovphysx.py +++ b/source/isaaclab_ov/test/sim/test_views_xform_prim_ovphysx.py @@ -27,8 +27,6 @@ OVPHYSX_SIM_CFG = SimulationCfg(physics=OvPhysxCfg()) -pytestmark = pytest.mark.device_split - def test_view_raises_before_physics_ready(): """A view constructed before PHYSICS_READY raises a clear error on pose-method calls.""" diff --git a/source/isaaclab_ov/test/test_randomize_rigid_body_collider_offsets_mdp.py b/source/isaaclab_ov/test/test_randomize_rigid_body_collider_offsets_mdp.py index 0238c5d498ad..d84bd48f851f 100644 --- a/source/isaaclab_ov/test/test_randomize_rigid_body_collider_offsets_mdp.py +++ b/source/isaaclab_ov/test/test_randomize_rigid_body_collider_offsets_mdp.py @@ -12,8 +12,7 @@ ``env`` / ``asset_cfg`` inputs are stubbed: the term only reads ``cfg.params["asset_cfg"]``, ``env.scene[...]``, ``env.scene.num_envs`` and ``env.sim.physics_manager``. -Kitless; run once per device (``-k cpu`` / ``-k 'cuda:0'``) -- the ovphysx runtime binds the -device mode process-globally (see the asset tests' module docstring). +Kitless; the CPU and CUDA cases run in one process. """ from __future__ import annotations @@ -41,33 +40,12 @@ wp.init() -pytestmark = pytest.mark.device_split - -_LOCKED_DEVICE: list[str | None] = [None] REST_RANGE = (0.0002, 0.0006) # below the 1 mm default contact offset: PhysX rejects rest >= contact CONTACT_RANGE = (0.03, 0.05) EPS = 1e-6 -@pytest.fixture(autouse=True) -def _ovphysx_skip_other_device(request): - """Skip parametrized tests on the device the session is not pinned to (process-global lock).""" - callspec = getattr(request.node, "callspec", None) - device = callspec.params.get("device") if callspec is not None else None - if device is None: - 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): """Build a simulation context that dispatches to the OVPhysX manager.""" sim_cfg = SimulationCfg(physics=OvPhysxCfg(), device=device, dt=1.0 / 60.0, gravity=(0.0, 0.0, -9.81)) diff --git a/source/isaaclab_ov/test/test_randomize_rigid_body_material_mdp.py b/source/isaaclab_ov/test/test_randomize_rigid_body_material_mdp.py index 40819ee8c417..f895ef175524 100644 --- a/source/isaaclab_ov/test/test_randomize_rigid_body_material_mdp.py +++ b/source/isaaclab_ov/test/test_randomize_rigid_body_material_mdp.py @@ -8,8 +8,7 @@ Constructs the public term with an ``EventTermCfg`` and verifies per-shape friction/restitution writes through a real OVPhysX rigid object's ``OvPhysxView``. -Kitless; run once per device (``-k cpu`` / ``-k 'cuda:0'``) -- the ovphysx runtime binds the -device mode process-globally (see the asset tests' module docstring). +Kitless; the CPU and CUDA cases run in one process. """ from __future__ import annotations @@ -34,28 +33,6 @@ wp.init() -pytestmark = pytest.mark.device_split - -_LOCKED_DEVICE: list[str | None] = [None] - - -@pytest.fixture(autouse=True) -def _ovphysx_skip_other_device(request): - """Skip parametrized tests on the device the session is not pinned to (process-global lock).""" - callspec = getattr(request.node, "callspec", None) - device = callspec.params.get("device") if callspec is not None else None - if device is None: - 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): """Build a simulation context that dispatches to the OVPhysX manager.""" diff --git a/source/isaaclab_physx/changelog.d/antoiner-fix-physx-collection-com-inertia.fixed.rst b/source/isaaclab_physx/changelog.d/antoiner-fix-physx-collection-com-inertia.fixed.rst new file mode 100644 index 000000000000..e8ee1bc3b0a2 --- /dev/null +++ b/source/isaaclab_physx/changelog.d/antoiner-fix-physx-collection-com-inertia.fixed.rst @@ -0,0 +1,10 @@ +* Fixed :meth:`~isaaclab_physx.assets.RigidObjectCollection.set_coms_index` and + :meth:`~isaaclab_physx.assets.RigidObjectCollection.set_coms_mask` raising a ``KeyError`` because the + center-of-mass poses were passed to the PhysX tensor API as ``wp.transformf`` instead of ``float32``. +* Fixed :meth:`~isaaclab_physx.assets.RigidObjectCollection.set_coms_index` and + :meth:`~isaaclab_physx.assets.RigidObjectCollection.set_coms_mask` leaving the unselected entries of + :attr:`~isaaclab_physx.assets.RigidObjectCollectionData.body_com_pose_b` zeroed when they were called for a + subset of bodies or environments before the property was first read. +* Fixed :meth:`~isaaclab_physx.assets.RigidObjectCollection.set_inertias_index` and + :meth:`~isaaclab_physx.assets.RigidObjectCollection.set_inertias_mask` failing in the PhysX backend because the + inertias were passed to the PhysX tensor API with the wrong layout. diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py index 341d60b62bbb..fc5f1245589d 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py @@ -1032,13 +1032,13 @@ def set_coms_index( ], device=self.device, ) - # Invalidate the cached buffer - self.data._body_com_pose_b.timestamp = self.data._sim_timestamp + # Invalidate the cached buffers. The next read refreshes from the view, so that a partial write does not + # expose entries of the buffer that were never read back from the simulation. self.data._reset_body_com_pose_b_dependents() # Set into simulation, note that when updating "model" properties with PhysX we need to do it on CPU. # Convert from instance order (num_instances, num_bodies, 7) to view order (num_bodies*num_instances, 7) for # PhysX. - com_view_order = self.reshape_data_to_view_2d(self.data._body_com_pose_b.data, device="cpu") # (B*I, 7) + com_view_order = self.reshape_data_to_view_2d(self.data._body_com_pose_b.data, device="cpu").view(wp.float32) view_ids = self._env_body_ids_to_view_ids(env_ids, body_ids, device="cpu") self.root_view.set_coms(com_view_order, indices=view_ids) @@ -1128,8 +1128,9 @@ def set_inertias_index( device=self.device, ) # Set into simulation, note that when updating "model" properties with PhysX we need to do it on CPU. - # Convert from instance order (num_instances, num_bodies) to view order for PhysX. - inertia_view_order = self.reshape_data_to_view_2d(self.data._body_inertia, device="cpu") + # Convert from instance order (num_instances, num_bodies, 9) to view order (num_bodies*num_instances, 9) for + # PhysX. + inertia_view_order = self.reshape_data_to_view_3d(self.data._body_inertia, 9, device="cpu") view_ids = self._env_body_ids_to_view_ids(env_ids, body_ids, device="cpu") self.root_view.set_inertias(inertia_view_order, indices=view_ids) diff --git a/source/isaaclab_physx/test/assets/test_articulation.py b/source/isaaclab_physx/test/assets/test_articulation.py index e0a81ef3aa97..5cb2a55b479f 100644 --- a/source/isaaclab_physx/test/assets/test_articulation.py +++ b/source/isaaclab_physx/test/assets/test_articulation.py @@ -3,6 +3,15 @@ # # SPDX-License-Identifier: BSD-3-Clause +"""Real PhysX articulation coverage. + +Most checks run against one module-scoped composite scene built from the local branching fixture. Each +articulation island holds two environments so that partial writes can target environment 1 and prove that +environment 0 is preserved in the real PhysX state. Tests that need their own simulation context (failure +modes and asset-specific seams) are defined first: pytest runs them before the composite scene is created, +and a new simulation context would replace the composite stage. +""" + # ignore private usage of variables warning # pyright: reportPrivateUsage=none @@ -12,2258 +21,1252 @@ BRANCHING_MJWARP_JOINT_NAMES, BRANCHING_PHYSX_BODY_NAMES, BRANCHING_PHYSX_JOINT_NAMES, - PANDA_BODY_NAMES, - PANDA_JOINT_NAMES, - PANDA_ROOT_PRESERVING_REVERSED_BODY_NAMES, ) -from isaaclab.utils import clone, replace - -HEADLESS = True +from isaaclab.utils import replace launch_test_simulation() import logging +import math import sys +from collections.abc import Iterator +from dataclasses import dataclass from pathlib import Path +import numpy as np import pytest import torch import warp as wp from isaaclab_physx.assets import Articulation -from isaaclab_physx.sim.schemas import PhysxJointCfg +from isaaclab_physx.assets.articulation.kernels import write_joint_state_data, write_joint_state_data_kernel -from pxr import UsdPhysics +from pxr import Gf, PhysxSchema, UsdGeom, UsdPhysics import isaaclab.sim as sim_utils -import isaaclab.utils.math as math_utils -import isaaclab.utils.string as string_utils -from isaaclab.actuators import IdealPDActuatorCfg, ImplicitActuatorCfg +from isaaclab.actuators import ImplicitActuatorCfg from isaaclab.assets import ArticulationCfg, get_articulation_name_ordering -from isaaclab.controllers import DifferentialIKController, DifferentialIKControllerCfg -from isaaclab.envs.mdp.terminations import joint_effort_out_of_limit -from isaaclab.managers import SceneEntityCfg -from isaaclab.sim import build_simulation_context -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR -from isaaclab.utils.math import compute_pose_error, matrix_from_quat, quat_inv, subtract_frame_transforms -from isaaclab.utils.version import get_isaac_sim_version, has_kit - -## -# Pre-defined configs -## -from isaaclab_assets import ( # isort:skip - ANYMAL_C_CFG, - FRANKA_PANDA_CFG, - FRANKA_PANDA_HIGH_PD_CFG, -) -from isaaclab_assets.robots.shadow_hand import SHADOW_HAND_PHYSX_CFG +from isaaclab.sim import SimulationContext, build_simulation_context +from isaaclab.utils.math import combine_frame_transforms, matrix_from_quat, quat_apply + +_FIXTURE = Path(__file__).parent / "data" / "articulation_ordering_branching.usda" +_NUM_ENVS = 2 +_ENV_SPACING = 3.0 + +# The branching fixture authors no geometry. Place the links along the x axis so that joints, Jacobians, +# and center-of-mass offsets act on non-degenerate lever arms. Each joint is anchored at its child link origin, +# offset from the parent link origin by the listed vector. +_JOINT_ANCHORS = { + "left_shoulder": (0.2, 0.0, 0.0), + "left_elbow": (0.5, 0.0, 0.0), + "right_shoulder": (-0.2, 0.0, 0.0), + "right_elbow": (-0.5, 0.0, 0.0), +} +_LINK_POSITIONS = { + "left_upper": (0.2, 0.0, 0.0), + "left_tip": (0.7, 0.0, 0.0), + "right_upper": (-0.2, 0.0, 0.0), + "right_tip": (-0.7, 0.0, 0.0), +} +# Drive and limit values authored in USD, keyed by joint name. USD angular drives use degree units, so the +# authored values are converted such that PhysX reports the SI values listed here. +_USD_STIFFNESS = {"left_shoulder": 4.0, "left_elbow": 5.0, "right_shoulder": 6.0, "right_elbow": 7.0} +_USD_DAMPING = {"left_shoulder": 0.4, "left_elbow": 0.5, "right_shoulder": 0.6, "right_elbow": 0.7} +_USD_MAX_FORCE = {"left_shoulder": 40.0, "left_elbow": 41.0, "right_shoulder": 42.0, "right_elbow": 43.0} +_USD_MAX_VELOCITY = {"left_shoulder": 2.0, "left_elbow": 3.0, "right_shoulder": 4.0, "right_elbow": 5.0} +# Configured gains and limits shared by both floating islands, so that their responses are comparable. +_FLOATING_ACTUATORS = { + "joints": ImplicitActuatorCfg( + joint_names_expr=[".*"], + stiffness={".*_shoulder": 6.0, ".*_elbow": 4.0}, + damping={".*_shoulder": 0.6, ".*_elbow": 0.4}, + joint_velocity_limit=7.0, + joint_effort_limit=30.0, + ) +} -def generate_articulation_cfg( - articulation_type: str, - stiffness: float | None = 10.0, - damping: float | None = 2.0, - actuator_velocity_limit: float | None = None, - actuator_effort_limit: float | None = None, - joint_velocity_limit: float | None = None, - joint_effort_limit: float | None = None, -) -> ArticulationCfg: - """Generate an articulation configuration. +def _branching_cfg(prim_path: str = "/World/Robot", **kwargs) -> ArticulationCfg: + """Create a configuration of the local branching fixture with implicit drives on every joint.""" + kwargs.setdefault("actuators", {"joints": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=5.0, damping=0.5)}) + return ArticulationCfg(prim_path=prim_path, spawn=sim_utils.UsdFileCfg(usd_path=str(_FIXTURE)), **kwargs) - 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. - actuator_velocity_limit: Velocity limit for the actuators. Only currently used for "single_joint_implicit" - and "single_joint_explicit". - actuator_effort_limit: Effort limit for explicit actuators. Only currently used for - "single_joint_explicit". - joint_velocity_limit: Velocity limit for the actuators (set into the simulation). - Only currently used for "single_joint_implicit" and "single_joint_explicit". - joint_effort_limit: 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. +def _author_branching_robot( + robot_path: str, + *, + fixed_base: bool = False, + root_on_base: bool = False, + reversed_left_elbow: bool = False, + spatial_tendon: bool = False, + fixed_tendon: bool = False, + joint_limits_deg: tuple[float, float] = (-170.0, 170.0), +) -> None: + """Author geometry, drives, and optional structure on one spawned branching robot. + Args: + robot_path: Prim path of the spawned robot. + fixed_base: Whether to fix the base link to the world with a fixed joint. + root_on_base: Whether to move the articulation root API from the robot prim onto the base link. + reversed_left_elbow: Whether to swap the bodies of the left elbow joint. + spatial_tendon: Whether to attach a spatial tendon from the base to the left tip. + fixed_tendon: Whether to couple the right shoulder and elbow with a fixed tendon. + joint_limits_deg: Lower and upper joint limits [deg] applied to every joint. PhysX ignores limit writes + on joints authored without limits. """ - 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_PHYSX_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.UsdPhysicsDriveCfg(max_force=80.0), PhysxJointCfg(max_joint_velocity=5.0)], - ), - actuators={ - "joint": ImplicitActuatorCfg( - joint_names_expr=[".*"], - joint_effort_limit=joint_effort_limit, - joint_velocity_limit=joint_velocity_limit, - actuator_velocity_limit=actuator_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.UsdPhysicsDriveCfg(max_force=80.0), PhysxJointCfg(max_joint_velocity=5.0)], - ), - actuators={ - "joint": IdealPDActuatorCfg( - joint_names_expr=[".*"], - joint_effort_limit=joint_effort_limit, - joint_velocity_limit=joint_velocity_limit, - actuator_effort_limit=actuator_effort_limit, - actuator_velocity_limit=actuator_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 + stage = sim_utils.get_current_stage() + for body_name, position in _LINK_POSITIONS.items(): + UsdGeom.Xformable(stage.GetPrimAtPath(f"{robot_path}/{body_name}")).AddTranslateOp().Set(Gf.Vec3d(*position)) + for joint_name, anchor in _JOINT_ANCHORS.items(): + joint = UsdPhysics.RevoluteJoint.Get(stage, f"{robot_path}/{joint_name}") + if reversed_left_elbow and joint_name == "left_elbow": + body0, body1 = joint.GetBody0Rel().GetTargets(), joint.GetBody1Rel().GetTargets() + joint.GetBody0Rel().SetTargets(body1) + joint.GetBody1Rel().SetTargets(body0) + joint.CreateLocalPos0Attr(Gf.Vec3f(0.0)) + joint.CreateLocalPos1Attr(Gf.Vec3f(*anchor)) + else: + joint.CreateLocalPos0Attr(Gf.Vec3f(*anchor)) + joint.CreateLocalPos1Attr(Gf.Vec3f(0.0)) + joint.CreateLowerLimitAttr(joint_limits_deg[0]) + joint.CreateUpperLimitAttr(joint_limits_deg[1]) + drive = UsdPhysics.DriveAPI.Apply(joint.GetPrim(), "angular") + drive.CreateStiffnessAttr(math.radians(_USD_STIFFNESS[joint_name])) + drive.CreateDampingAttr(math.radians(_USD_DAMPING[joint_name])) + drive.CreateMaxForceAttr(_USD_MAX_FORCE[joint_name]) + physx_joint = PhysxSchema.PhysxJointAPI.Apply(joint.GetPrim()) + physx_joint.CreateMaxJointVelocityAttr(math.degrees(_USD_MAX_VELOCITY[joint_name])) + # The base carries one collision shape so that its inertia follows from real geometry. + collision = UsdGeom.Cube.Define(stage, f"{robot_path}/base/collision") + collision.CreateSizeAttr(0.1) + UsdPhysics.CollisionAPI.Apply(collision.GetPrim()) + if fixed_base: + fixed_joint = UsdPhysics.FixedJoint.Define(stage, f"{robot_path}/fixed_root") + fixed_joint.GetBody1Rel().SetTargets([f"{robot_path}/base"]) + if root_on_base: + stage.GetPrimAtPath(robot_path).RemoveAPI(UsdPhysics.ArticulationRootAPI) + UsdPhysics.ArticulationRootAPI.Apply(stage.GetPrimAtPath(f"{robot_path}/base")) + if spatial_tendon: + root_prim = stage.GetPrimAtPath(f"{robot_path}/base") + root_attachment = PhysxSchema.PhysxTendonAttachmentAPI(root_prim, "root") + root_attachment.CreateLocalPosAttr(Gf.Vec3f(0.0)) + root = PhysxSchema.PhysxTendonAttachmentRootAPI.Apply(root_prim, "root") + root.CreateStiffnessAttr(5.0) + root.CreateDampingAttr(0.5) + root.CreateLimitStiffnessAttr(1.0) + root.CreateOffsetAttr(0.0) + leaf_prim = stage.GetPrimAtPath(f"{robot_path}/left_tip") + leaf_attachment = PhysxSchema.PhysxTendonAttachmentAPI(leaf_prim, "leaf") + leaf_attachment.CreateLocalPosAttr(Gf.Vec3f(0.0)) + leaf_attachment.CreateParentAttachmentAttr("root") + leaf_attachment.CreateParentLinkRel().SetTargets([root_prim.GetPath()]) + leaf = PhysxSchema.PhysxTendonAttachmentLeafAPI.Apply(leaf_prim, "leaf") + leaf.CreateRestLengthAttr(0.5) + leaf.CreateLowerLimitAttr(0.0) + leaf.CreateUpperLimitAttr(2.0) + if fixed_tendon: + # The root joint carries the tendon and its own axis; the elbow adds the child axis. + root = PhysxSchema.PhysxTendonAxisRootAPI.Apply(stage.GetPrimAtPath(f"{robot_path}/right_shoulder"), "fixed") + root.CreateStiffnessAttr(5.0) + root.CreateDampingAttr(0.5) + root.CreateRestLengthAttr(0.4) + root.CreateOffsetAttr(0.0) + for joint_name in ("right_shoulder", "right_elbow"): + axis = PhysxSchema.PhysxTendonAxisAPI.Apply(stage.GetPrimAtPath(f"{robot_path}/{joint_name}"), "fixed") + axis.CreateGearingAttr([1.0]) + axis.CreateForceCoefficientAttr([1.0]) + + +def _in_user_order(values: dict[str, float], names: list[str], device: str) -> torch.Tensor: + """Return per-name values as a ``(num_envs, len(names))`` tensor in the given name order.""" + return torch.tensor([values[name] for name in names], device=device).repeat(_NUM_ENVS, 1) + + +def _yaw_quat(angle: float) -> tuple[float, float, float, float]: + """Return a unit quaternion ``(x, y, z, w)`` for a rotation of ``angle`` [rad] about the world z axis.""" + return (0.0, 0.0, math.sin(0.5 * angle), math.cos(0.5 * angle)) -def generate_articulation( - articulation_cfg: ArticulationCfg, num_articulations: int, device: str -) -> tuple[Articulation, torch.tensor]: - """Generate an articulation from a configuration. +@pytest.fixture +def sim(device: str) -> Iterator[SimulationContext]: + """Create a function-scoped simulation context for tests that own their scene.""" + # A new context would replace the stage of a live composite scene, so these tests must run before it. + assert SimulationContext.instance() is None, "define tests that own a simulation above the composite scene" + with build_simulation_context(device=device, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + yield sim - 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. +## +# Tests that own their simulation context. Keep them above the composite scene. +## - 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 +@pytest.mark.parametrize( + ("init_state", "message"), + [ + ({"joint_pos": {"left_shoulder": 1.0}}, "default positions out of the limits"), + ({"joint_vel": {"left_elbow": 10.0}}, "default velocities out of the limits"), + ], +) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +def test_out_of_range_default_joint_state_raises(sim, device, init_state, message) -> None: + """Initialization rejects configured default joint states outside the solver limits.""" + articulation = Articulation(_branching_cfg(init_state=ArticulationCfg.InitialStateCfg(**init_state))) + # Limits of +-30 deg and the USD velocity limits keep the configured defaults out of range. + _author_branching_robot("/World/Robot", fixed_base=True, joint_limits_deg=(-30.0, 30.0)) - # 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(replace(articulation_cfg, prim_path="/World/Env_[^/]*/Robot")) + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 - return articulation, translations + with pytest.raises(ValueError, match=message): + sim.reset() -# --------------------------------------------------------------------------- -# Franka task-space tracking helpers for the IK test. -# Mirrors the helpers in ``isaaclab_newton/test/assets/test_articulation.py``. -# --------------------------------------------------------------------------- +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +def test_setting_invalid_articulation_root_prim_path(sim, device) -> None: + """A configured articulation root path that does not exist fails initialization.""" + articulation = Articulation(_branching_cfg(articulation_root_prim_path="/non_existing_prim_path")) + _author_branching_robot("/World/Robot") + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 -def _setup_franka_at_home_pose(sim): - """Build a Franka articulation at its configured home pose. + with pytest.raises(RuntimeError, match="Failed to create articulation at"): + sim.reset() - See the Newton-side mirror for full docs. Standalone tests skip the - env reset path that normally pushes ``default_joint_pos`` to sim, - so we teleport explicitly to avoid the URDF-neutral - near-singular pose where the Franka wrist axes nearly align. - Args: - sim: The simulation context to use. +@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) +def test_gravity_compensation_holds_static_equilibrium(sim, device) -> None: + """Gravity compensation efforts hold the arms still; gravity compensation and the mass matrix follow mass writes. - Returns: - Tuple of ``(robot, ee_frame_idx, ee_jacobi_idx, arm_joint_ids)``. + Setting ``tau = g(q)`` at rest gives zero joint acceleration, so sign, frame, and DoF-ordering errors in + :attr:`~isaaclab.assets.BaseArticulationData.gravity_compensation_forces` surface as joint drift. The robot is + rolled 90 deg about x so that its joint axes are horizontal. Inertial writes are checked on CPU, where PhysX + applies them immediately (see ``test_inertial_writes_refresh_the_mass_matrix`` for the GPU pipeline). """ - cfg = replace(clone(FRANKA_PANDA_HIGH_PD_CFG), prim_path="/World/Env_[^/]*/Robot") - sim_utils.create_prim("/World/Env_0", "Xform", translation=(0.0, 0.0, 0.0)) - robot = Articulation(cfg) + articulation = Articulation( + _branching_cfg( + actuators={"joints": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=0.0, damping=0.0)}, + joint_ordering="mjwarp", + init_state=ArticulationCfg.InitialStateCfg(rot=(math.sqrt(0.5), 0.0, 0.0, math.sqrt(0.5))), + ) + ) + _author_branching_robot("/World/Robot", fixed_base=True) + # Offset the tip centers of mass from the elbow axes so that gravity also loads the elbows. + for tip in ("left_tip", "right_tip"): + UsdPhysics.MassAPI(sim.stage.GetPrimAtPath(f"/World/Robot/{tip}")).CreateCenterOfMassAttr(Gf.Vec3f(0.2, 0, 0)) sim.reset() - assert robot.is_initialized + joint_pos = torch.tensor([[0.3, -0.5, 0.4, 0.6]], device=device) + articulation.write_joint_state_to_sim_index(position=joint_pos, velocity=torch.zeros_like(joint_pos)) - ee_frame_idx = robot.find_bodies("panda_hand")[0][0] - ee_jacobi_idx = ee_frame_idx - 1 - arm_joint_ids = robot.find_joints(["panda_joint.*"])[0] + for _ in range(100): + gravity_compensation = articulation.data.gravity_compensation_forces.torch[:, articulation.num_base_dofs :] + assert gravity_compensation.abs().min() > 0.1, "every joint must carry a gravity load" + articulation.set_joint_effort_target_index(target=gravity_compensation) + articulation.write_data_to_sim() + sim.step() + articulation.update(sim.cfg.dt) + torch.testing.assert_close(articulation.data.joint_pos.torch, joint_pos, atol=5e-3, rtol=0.0) - robot.write_joint_state_to_sim( - position=robot.data.default_joint_pos.torch[:, :].clone(), - velocity=robot.data.default_joint_vel.torch[:, :].clone(), - ) - return robot, ee_frame_idx, ee_jacobi_idx, arm_joint_ids + # Mass and inertia writes refresh gravity compensation and a mass matrix already read in this step. + data = articulation.data + for write, kwarg, values in ( + (articulation.set_masses_index, "masses", data.body_mass), + (articulation.set_inertias_index, "inertias", data.body_inertia), + ): + primed = (data.gravity_compensation_forces.torch.clone(), data.mass_matrix.torch.clone()) + write(**{kwarg: 2.0 * values.torch.clone() + 0.1}) + if kwarg == "masses": + assert not torch.allclose(data.gravity_compensation_forces.torch, primed[0], atol=1e-3) + assert not torch.allclose(data.mass_matrix.torch, primed[1], atol=1e-4), f"{kwarg} write left a stale matrix" -def _compute_ee_pose_root(robot, ee_frame_idx): - """Return ``(ee_pos_b, ee_quat_b, root_pose_w)`` in the root frame.""" - ee_pose_w = robot.data.body_pose_w.torch[:, ee_frame_idx] - root_pose_w = robot.data.root_pose_w.torch - ee_pos_b, ee_quat_b = subtract_frame_transforms( - root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_pose_w[:, 0:3], ee_pose_w[:, 3:7] - ) - return ee_pos_b, ee_quat_b, root_pose_w +## +# Composite scene shared by the remaining tests. +## -def _compute_jacobian_root_frame(robot, ee_jacobi_idx, arm_joint_ids): - """Return the EE Jacobian sliced to ``arm_joint_ids`` and rotated to the root frame.""" - jacobian = robot.data.body_link_jacobian_w.torch[:, ee_jacobi_idx, :, :][:, :, arm_joint_ids] - base_rot_matrix = matrix_from_quat(quat_inv(robot.data.root_pose_w.torch[:, 3:7])) - jacobian[:, :3, :] = torch.bmm(base_rot_matrix, jacobian[:, :3, :]) - jacobian[:, 3:, :] = torch.bmm(base_rot_matrix, jacobian[:, 3:, :]) - return jacobian +@dataclass +class _ArticulationScene: + """Articulation islands that share one real PhysX lifecycle.""" + + sim: SimulationContext + device: str + ordered: Articulation + """Fixed base, MJWarp joint and body ordering, gains and limits loaded from USD.""" + floating: Articulation + """Floating base discovered on the base link below the spawned prim, gains and limits from the config.""" + reordered: Articulation + """Floating base with an explicit root path and a body ordering that moves the root link.""" + tendon: Articulation + """Fixed base with a reversed left elbow, a spatial tendon, and a fixed tendon.""" + origins: dict[str, torch.Tensor] + """World position of every environment of each island, keyed by island name.""" + refcounts: dict[str, int] + """Reference count of each articulation right after construction.""" + cold_inertial_view_values: tuple[torch.Tensor, torch.Tensor] + """Masses and inertias written through the floating island's view before any read.""" + cold_inertial_reads: tuple[torch.Tensor, torch.Tensor] + """Masses and inertias of the floating island on their first read after that write.""" + + @property + def islands(self) -> dict[str, Articulation]: + """All articulations keyed by island name.""" + return {"ordered": self.ordered, "floating": self.floating, "reordered": self.reordered, "tendon": self.tendon} + + def step(self, num_steps: int = 1) -> None: + """Write, step, and update every articulation.""" + for _ in range(num_steps): + for articulation in self.islands.values(): + articulation.write_data_to_sim() + self.sim.step() + for articulation in self.islands.values(): + articulation.update(self.sim.cfg.dt) + + +def _spawn_island(name: str, y_offset: float, cfg: ArticulationCfg, **authoring) -> tuple[Articulation, torch.Tensor]: + """Spawn one two-environment island and return the articulation and its environment origins.""" + island_path = f"/World/{name}" + sim_utils.create_prim(island_path, "Xform", translation=(0.0, y_offset, 0.0)) + origins = [] + for env_index in range(_NUM_ENVS): + sim_utils.create_prim(f"{island_path}/Env_{env_index}", "Xform", translation=(_ENV_SPACING * env_index, 0, 0)) + origins.append((_ENV_SPACING * env_index, y_offset, 0.0)) + articulation = Articulation(replace(cfg, prim_path=f"{island_path}/Env_[^/]*/Robot")) + for env_index in range(_NUM_ENVS): + _author_branching_robot(f"{island_path}/Env_{env_index}/Robot", **authoring) + return articulation, torch.tensor(origins) + + +@pytest.fixture(scope="module", params=test_devices(DeviceScope.CUDA)) +def articulation_scene(request) -> Iterator[_ArticulationScene]: + """Initialize every composite-scene articulation once for this module.""" + device = request.param + with build_simulation_context(device=device, gravity_enabled=False) as sim: + sim._app_control_on_stop_handle = None + islands, origins, refcounts = {}, {}, {} + island_specs = { + "ordered": ( + _branching_cfg( + actuators={"joints": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=None, damping=None)}, + joint_ordering="mjwarp", + body_ordering="mjwarp", + ), + {"fixed_base": True}, + ), + "floating": (_branching_cfg(actuators=_FLOATING_ACTUATORS), {"root_on_base": True}), + "reordered": ( + _branching_cfg( + actuators=_FLOATING_ACTUATORS, + articulation_root_prim_path="/base", + body_ordering=("left_tip", "base", "right_upper", "right_tip", "left_upper"), + ), + {"root_on_base": True}, + ), + "tendon": ( + _branching_cfg(), + {"fixed_base": True, "reversed_left_elbow": True, "spatial_tendon": True, "fixed_tendon": True}, + ), + } + for index, (name, (cfg, authoring)) in enumerate(island_specs.items()): + islands[name], origins[name] = _spawn_island(name.capitalize(), 3.0 * index, cfg, **authoring) + refcounts[name] = sys.getrefcount(islands[name]) + sim.reset() + # Write inertial properties through the tensor view of the floating island, whose identity body order keeps + # the construction-time buffer timestamps, and read them once before any test steps the scene or reads these + # buffers. Restore the view afterwards. + floating_view = islands["floating"].root_view + cpu_env_ids = wp.array(list(range(_NUM_ENVS)), dtype=wp.int32, device="cpu") + view_masses = wp.to_torch(floating_view.get_masses()).clone() + view_inertias = wp.to_torch(floating_view.get_inertias()).clone() + cold_masses = view_masses + 0.293 + cold_inertias = view_inertias.clone() + cold_inertias[..., [0, 4, 8]] += 0.023 + floating_view.set_masses(wp.from_torch(cold_masses, dtype=wp.float32), indices=cpu_env_ids) + floating_view.set_inertias(wp.from_torch(cold_inertias, dtype=wp.float32), indices=cpu_env_ids) + cold_reads = ( + islands["floating"].data.body_mass.torch.clone(), + islands["floating"].data.body_inertia.torch.clone(), + ) + floating_view.set_masses(wp.from_torch(view_masses, dtype=wp.float32), indices=cpu_env_ids) + floating_view.set_inertias(wp.from_torch(view_inertias, dtype=wp.float32), indices=cpu_env_ids) + yield _ArticulationScene( + sim=sim, + device=device, + origins={name: value.to(device) for name, value in origins.items()}, + refcounts=refcounts, + cold_inertial_view_values=(cold_masses, cold_inertias), + cold_inertial_reads=cold_reads, + **islands, + ) + +def test_articulation_initialization_and_ordering(articulation_scene: _ArticulationScene) -> None: + """Resolve the configured structure, orderings, gains, and limits of every real island.""" + scene = articulation_scene + device = scene.device + for name, articulation in scene.islands.items(): + assert articulation.is_initialized, name + assert articulation.num_instances == _NUM_ENVS, name + # Check that the framework doesn't hold excessive strong references. + assert scene.refcounts[name] < 10, name + 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 + # The composite scene runs without gravity. + torch.testing.assert_close(articulation.data.GRAVITY_VEC_W.torch, torch.zeros((_NUM_ENVS, 3), device=device)) + assert scene.ordered.is_fixed_base and scene.tendon.is_fixed_base + assert not scene.floating.is_fixed_base and not scene.reordered.is_fixed_base + # The floating root is discovered on the base link below the spawned robot prim; the reordered root is + # resolved through the configured root path. + assert scene.floating.root_view.prim_paths[0] == "/World/Floating/Env_0/Robot/base" + assert scene.reordered.root_view.prim_paths[0] == "/World/Reordered/Env_0/Robot/base" + + ordered, floating = scene.ordered, scene.floating + assert tuple(ordered.backend_joint_names) == BRANCHING_PHYSX_JOINT_NAMES + assert tuple(ordered.backend_body_names) == BRANCHING_PHYSX_BODY_NAMES + assert get_articulation_name_ordering(ordered, "mjwarp", "joint") == BRANCHING_MJWARP_JOINT_NAMES + assert get_articulation_name_ordering(ordered, "mjwarp", "body") == BRANCHING_MJWARP_BODY_NAMES + assert tuple(ordered.joint_names) == BRANCHING_MJWARP_JOINT_NAMES + assert tuple(ordered.body_names) == BRANCHING_MJWARP_BODY_NAMES + assert ordered.joint_ordering is not None + assert ordered.body_ordering is not None + # Without a custom ordering the public body order is the PhysX link order. + assert floating.body_ordering is None + assert [path.split("/")[-1] for path in floating.root_view.link_paths[0]] == floating.body_names + for articulation in (ordered, floating): + assert articulation.data.root_pos_w.torch.shape == (_NUM_ENVS, 3) + assert articulation.data.root_quat_w.torch.shape == (_NUM_ENVS, 4) + assert articulation.data.joint_pos.torch.shape == (_NUM_ENVS, 4) + assert articulation.data.body_mass.torch.shape == (_NUM_ENVS, 5) + assert articulation.data.body_inertia.torch.shape == (_NUM_ENVS, 5, 9) + + # Unset actuator gains and limits are loaded from the USD drives, in public joint order. + joint_backend_to_user = list(ordered.joint_ordering.backend_to_user_indices) + for public, usd_values, raw in ( + (ordered.actuators["joints"].stiffness, _USD_STIFFNESS, ordered.root_view.get_dof_stiffnesses()), + (ordered.actuators["joints"].damping, _USD_DAMPING, ordered.root_view.get_dof_dampings()), + (ordered.data.joint_effort_limits.torch, _USD_MAX_FORCE, ordered.root_view.get_dof_max_forces()), + (ordered.data.joint_vel_limits.torch, _USD_MAX_VELOCITY, ordered.root_view.get_dof_max_velocities()), + ): + expected = _in_user_order(usd_values, ordered.joint_names, device) + torch.testing.assert_close(public, expected) + torch.testing.assert_close(wp.to_torch(raw).to(device), expected[:, joint_backend_to_user]) + torch.testing.assert_close(ordered.data.joint_stiffness.torch, ordered.actuators["joints"].stiffness) + torch.testing.assert_close(ordered.data.joint_damping.torch, ordered.actuators["joints"].damping) + # Configured gains and limits override USD and reach the solver. + for public, raw, values in ( + (floating.actuators["joints"].stiffness, floating.root_view.get_dof_stiffnesses(), (6.0, 4.0)), + (floating.actuators["joints"].damping, floating.root_view.get_dof_dampings(), (0.6, 0.4)), + ): + per_joint = {name: values[0] if name.endswith("shoulder") else values[1] for name in floating.joint_names} + expected = _in_user_order(per_joint, floating.joint_names, device) + torch.testing.assert_close(public, expected) + torch.testing.assert_close(wp.to_torch(raw).to(device), expected) + for public, raw, limit in ( + (floating.data.joint_vel_limits.torch, floating.root_view.get_dof_max_velocities(), 7.0), + (floating.data.joint_effort_limits.torch, floating.root_view.get_dof_max_forces(), 30.0), + ): + torch.testing.assert_close(public, torch.full((_NUM_ENVS, 4), limit, device=device)) + torch.testing.assert_close(wp.to_torch(raw).to(device), public) + + +def test_articulation_joint_state_writes_follow_ordering(articulation_scene: _ArticulationScene) -> None: + """Joint state writes reach PhysX in backend order and reads return public order.""" + scene = articulation_scene + device = scene.device + articulation = scene.ordered + joint_user_to_backend = list(articulation.joint_ordering.user_to_backend_indices) + joint_backend_to_user = list(articulation.joint_ordering.backend_to_user_indices) + body_user_to_backend = list(articulation.body_ordering.user_to_backend_indices) -def _build_relative_pose_target(robot, ee_frame_idx, delta_xyz, device): - """Build a target pose = (current EE pose) + ``delta_xyz``, preserving orientation.""" - initial_ee_pos_b, initial_ee_quat_b, _ = _compute_ee_pose_root(robot, ee_frame_idx) - target_pos_b = initial_ee_pos_b + torch.tensor([list(delta_xyz)], device=device, dtype=initial_ee_pos_b.dtype) - return torch.cat([target_pos_b, initial_ee_quat_b], dim=-1) + # Full-data writes land in backend order. + joint_pos = torch.linspace(-0.3, 0.3, 4, device=device).repeat(_NUM_ENVS, 1) + joint_vel = torch.linspace(0.05, 0.13, 4, device=device).repeat(_NUM_ENVS, 1) + initial_stiffness = articulation.data.joint_stiffness.torch.clone() + joint_stiffness = 10.0 + torch.arange(4, device=device, dtype=torch.float32).repeat(_NUM_ENVS, 1) + original_body_link_pose_w = articulation.data.body_link_pose_w.torch.clone() + articulation.write_joint_stiffness_to_sim_index(stiffness=joint_stiffness, full_data=True) + articulation.write_joint_position_to_sim_index(position=joint_pos, full_data=True) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel, full_data=True) + for raw, written in ( + (articulation.root_view.get_dof_positions(), joint_pos), + (articulation.root_view.get_dof_velocities(), joint_vel), + (articulation.root_view.get_dof_stiffnesses(), joint_stiffness), + ): + torch.testing.assert_close(wp.to_torch(raw).to(device), written[:, joint_backend_to_user]) + + # A joint write refreshes the body state without stepping. + body_link_pose_w = articulation.data.body_link_pose_w.torch + assert torch.all((original_body_link_pose_w[:, 1:] != body_link_pose_w[:, 1:]).any(dim=-1)) + body_com_pose_b = articulation.data.body_com_pose_b.torch + expected_com_pos, expected_com_quat = combine_frame_transforms( + body_link_pose_w[..., :3].reshape(-1, 3), + body_link_pose_w[..., 3:].reshape(-1, 4), + body_com_pose_b[..., :3].reshape(-1, 3), + body_com_pose_b[..., 3:].reshape(-1, 4), + ) + torch.testing.assert_close(expected_com_pos.view(_NUM_ENVS, -1, 3), articulation.data.body_com_pos_w.torch) + torch.testing.assert_close(expected_com_quat.view(_NUM_ENVS, -1, 4), articulation.data.body_com_quat_w.torch) + body_com_vel_w = articulation.data.body_com_vel_w.torch + torch.testing.assert_close(body_com_vel_w[..., 3:], articulation.data.body_link_vel_w.torch[..., 3:]) + torch.testing.assert_close(body_com_vel_w[..., :3], articulation.data.body_com_lin_vel_w.torch) + torch.testing.assert_close(body_com_vel_w[..., 3:], articulation.data.body_com_ang_vel_w.torch) + + # After a step, public reads are the backend state in public order. + scene.step() + for public, raw, user_to_backend in ( + (articulation.data.joint_pos.torch, articulation.root_view.get_dof_positions(), joint_user_to_backend), + (articulation.data.joint_vel.torch, articulation.root_view.get_dof_velocities(), joint_user_to_backend), + (articulation.data.joint_stiffness.torch, articulation.root_view.get_dof_stiffnesses(), joint_user_to_backend), + (articulation.data.body_link_pose_w.torch, articulation.root_view.get_link_transforms(), body_user_to_backend), + (articulation.data.body_com_pose_b.torch, articulation.root_view.get_coms(), body_user_to_backend), + ): + torch.testing.assert_close(public, wp.to_torch(raw).to(device)[:, user_to_backend]) + # The split accessors must be sliced from the reordered public poses, not from a stale backend-order cache. + data = articulation.data + torch.testing.assert_close(data.body_com_pos_b.torch, data.body_com_pose_b.torch[..., :3]) + torch.testing.assert_close(data.body_com_quat_b.torch, data.body_com_pose_b.torch[..., 3:]) + torch.testing.assert_close(data.body_pos_w.torch, data.body_link_pose_w.torch[..., :3]) + torch.testing.assert_close(data.body_quat_w.torch, data.body_link_pose_w.torch[..., 3:]) + + # Partial writes with int64 selectors in non-sorted order reach only the selected cells. + env_ids = torch.tensor([1], dtype=torch.int64, device=device) + joint_ids = torch.tensor([3, 0], dtype=torch.int64, device=device) + position = torch.tensor([[0.21, -0.13]], device=device) + velocity = torch.tensor([[0.41, -0.23]], device=device) + expected_position = articulation.data.joint_pos.torch.clone() + expected_velocity = articulation.data.joint_vel.torch.clone() + expected_position[env_ids[:, None], joint_ids] = position + expected_velocity[env_ids[:, None], joint_ids] = velocity + articulation.write_joint_state_to_sim_index( + position=position, velocity=velocity, env_ids=env_ids, joint_ids=joint_ids + ) + torch.testing.assert_close(articulation.data.joint_pos.torch, expected_position) + torch.testing.assert_close(articulation.data.joint_vel.torch, expected_velocity) + torch.testing.assert_close( + wp.to_torch(articulation.root_view.get_dof_positions()).to(device), + expected_position[:, joint_backend_to_user], + ) + torch.testing.assert_close( + wp.to_torch(articulation.root_view.get_dof_velocities()).to(device), + expected_velocity[:, joint_backend_to_user], + ) + # Restore the USD gains for the following tests. + articulation.write_joint_stiffness_to_sim_index(stiffness=initial_stiffness, full_data=True) -def _summarize_history(history, tail: int = 200): - """Return ``(min, mean)`` over the last ``tail`` samples.""" - tail_slice = history[-tail:] - return min(tail_slice), sum(tail_slice) / len(tail_slice) +def test_articulation_joint_and_body_properties_round_trip(articulation_scene: _ArticulationScene) -> None: + """Selected joint and body properties reach the selected PhysX entries and read back in public order.""" + scene = articulation_scene + device = scene.device + # Direct tensor-view mass and inertia writes become visible on the lazy read: on the first read of a cold + # buffer, taken by the fixture before any test step, and after the next update once the buffer is primed. + for cold_read, view_values in zip(scene.cold_inertial_reads, scene.cold_inertial_view_values): + torch.testing.assert_close(cold_read, view_values.to(device)) + cpu_env_ids = wp.array(list(range(_NUM_ENVS)), dtype=wp.int32, device="cpu") -def _to_device_tensor(array: wp.array, device: str) -> torch.Tensor: - """Convert a Warp array to a torch tensor on :paramref:`device`.""" - return wp.to_torch(array).to(device=device) + def write_backend_mass_inertia( + articulation: Articulation, delta_mass: float, delta_inertia: float + ) -> tuple[torch.Tensor, torch.Tensor]: + """Write distinct backend-order masses/inertias straight through the tensor view.""" + backend_masses = wp.to_torch(articulation.root_view.get_masses()).clone() + delta_mass + backend_inertias = wp.to_torch(articulation.root_view.get_inertias()).clone() + backend_inertias[..., [0, 4, 8]] += delta_inertia + articulation.root_view.set_masses(wp.from_torch(backend_masses, dtype=wp.float32), indices=cpu_env_ids) + articulation.root_view.set_inertias(wp.from_torch(backend_inertias, dtype=wp.float32), indices=cpu_env_ids) + return backend_masses.to(device), backend_inertias.to(device) + articulation = scene.ordered + body_user_to_backend = list(articulation.body_ordering.user_to_backend_indices) + body_backend_to_user = list(articulation.body_ordering.backend_to_user_indices) + joint_backend_to_user = list(articulation.joint_ordering.backend_to_user_indices) + # Reading the buffers primes them; the later tests drive the island with these initial properties. + initial_properties = { + "masses": articulation.data.body_mass.torch.clone(), + "coms": articulation.data.body_com_pose_b.torch.clone(), + "inertias": articulation.data.body_inertia.torch.clone(), + } + backend_masses, backend_inertias = write_backend_mass_inertia(articulation, 0.137, 0.011) + articulation.update(scene.sim.cfg.dt) + torch.testing.assert_close(articulation.data.body_mass.torch, backend_masses[:, body_user_to_backend]) + torch.testing.assert_close(articulation.data.body_inertia.torch, backend_inertias[:, body_user_to_backend]) + + # Joint friction coefficients reach their TensorAPI slots for the selected cells only. + env_ids = torch.tensor([1], dtype=torch.int32, device=device) + joint_ids = torch.tensor([3, 0], dtype=torch.int32, device=device) + raw_friction_before = wp.to_torch(articulation.root_view.get_dof_friction_properties()).clone() + initial_frictions = { + "joint_friction_coeff": articulation.data.joint_friction_coeff.torch.clone(), + "joint_dynamic_friction_coeff": articulation.data.joint_dynamic_friction_coeff.torch.clone(), + "joint_viscous_friction_coeff": articulation.data.joint_viscous_friction_coeff.torch.clone(), + } + frictions = { + "joint_friction_coeff": torch.tensor([[0.9, 0.7]], device=device), + "joint_dynamic_friction_coeff": torch.tensor([[0.4, 0.3]], device=device), + "joint_viscous_friction_coeff": torch.tensor([[0.11, 0.22]], device=device), + } + friction_readers = { + "joint_friction_coeff": lambda: articulation.data.joint_friction_coeff.torch, + "joint_dynamic_friction_coeff": lambda: articulation.data.joint_dynamic_friction_coeff.torch, + "joint_viscous_friction_coeff": lambda: articulation.data.joint_viscous_friction_coeff.torch, + } + expected_frictions = {name: read().clone() for name, read in friction_readers.items()} + articulation.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff=frictions["joint_friction_coeff"], + joint_dynamic_friction_coeff=frictions["joint_dynamic_friction_coeff"], + joint_viscous_friction_coeff=frictions["joint_viscous_friction_coeff"], + env_ids=env_ids, + joint_ids=joint_ids, + ) + expected_raw_friction = raw_friction_before.clone() + for component, (name, values) in enumerate(frictions.items()): + expected_frictions[name][env_ids[:, None], joint_ids] = values + torch.testing.assert_close(friction_readers[name](), expected_frictions[name]) + expected_raw_friction[..., component] = expected_frictions[name][:, joint_backend_to_user].cpu() + torch.testing.assert_close(wp.to_torch(articulation.root_view.get_dof_friction_properties()), expected_raw_friction) + + # Joint position limits reach the solver; limits that exclude a default position clamp it. + raw_limits_before = wp.to_torch(articulation.root_view.get_dof_limits()).clone() + initial_limits = articulation.data.joint_pos_limits.torch.clone() + default_joint_pos = articulation.data.default_joint_pos.torch.clone() + limits = torch.tensor([[[-1.5, 1.25], [-1.0, 0.75]]], device=device) + articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) + expected_limits = initial_limits.clone() + expected_limits[env_ids[:, None], joint_ids] = limits + torch.testing.assert_close(articulation.data.joint_pos_limits.torch, expected_limits) + torch.testing.assert_close(articulation.data.default_joint_pos.torch, default_joint_pos) + expected_raw_limits = raw_limits_before.clone() + expected_raw_limits[1] = expected_limits[1, joint_backend_to_user].cpu() + torch.testing.assert_close(wp.to_torch(articulation.root_view.get_dof_limits()), expected_raw_limits) + articulation.write_joint_position_limit_to_sim_index( + limits=torch.tensor([[[0.05, 0.2], [-0.2, -0.1]]], device=device), env_ids=env_ids, joint_ids=joint_ids + ) + expected_defaults = default_joint_pos.clone() + expected_defaults[1, joint_ids] = torch.tensor([0.05, -0.1], device=device) + torch.testing.assert_close(articulation.data.default_joint_pos.torch, expected_defaults) + articulation.write_joint_position_limit_to_sim_index(limits=initial_limits, full_data=True) + + # Inertial properties of selected bodies reach the backend entries of the reordered bodies. + body_ids = torch.tensor([4, 1], dtype=torch.int32, device=device) + masses = torch.tensor([[2.5, 3.5]], device=device) + coms = articulation.data.body_com_pose_b.torch[env_ids][:, body_ids].clone() + coms[0, 0, :3] = torch.tensor([0.02, -0.01, 0.03], device=device) + coms[0, 1, :3] = torch.tensor([-0.03, 0.01, 0.02], device=device) + inertias = articulation.data.body_inertia.torch[env_ids][:, body_ids].clone() + inertias[0, 0, 0] *= 1.2 + inertias[0, 1, 4] *= 1.3 + expected_masses = articulation.data.body_mass.torch.clone() + expected_masses[env_ids[:, None], body_ids] = masses + articulation.set_masses_index(masses=masses, env_ids=env_ids, body_ids=body_ids) + torch.testing.assert_close(articulation.data.body_mass.torch, expected_masses) + expected_coms = articulation.data.body_com_pose_b.torch.clone() + expected_coms[env_ids[:, None], body_ids] = coms + articulation.set_coms_index(coms=coms, env_ids=env_ids, body_ids=body_ids) + torch.testing.assert_close(articulation.data.body_com_pose_b.torch, expected_coms) + expected_inertias = articulation.data.body_inertia.torch.clone() + expected_inertias[env_ids[:, None], body_ids] = inertias + articulation.set_inertias_index(inertias=inertias, env_ids=env_ids, body_ids=body_ids) + torch.testing.assert_close(articulation.data.body_inertia.torch, expected_inertias) + for public, raw in ( + (articulation.data.body_mass.torch, articulation.root_view.get_masses()), + (articulation.data.body_com_pose_b.torch, articulation.root_view.get_coms()), + (articulation.data.body_inertia.torch, articulation.root_view.get_inertias()), + ): + torch.testing.assert_close(wp.to_torch(raw).to(device), public[:, body_backend_to_user]) + + # Restore the initial properties for the tests that drive this island. + articulation.set_masses_index(masses=initial_properties["masses"], full_data=True) + articulation.set_coms_index(coms=initial_properties["coms"], full_data=True) + articulation.set_inertias_index(inertias=initial_properties["inertias"], full_data=True) + articulation.write_joint_friction_coefficient_to_sim_index(**initial_frictions, full_data=True) + + +def _rotate_z(angle: torch.Tensor, vector: tuple[float, float, float] | torch.Tensor) -> torch.Tensor: + """Rotate vectors about the world z axis. -def _assert_backend_to_user( - public_tensor: torch.Tensor, backend_tensor: torch.Tensor, user_to_backend: list[int] -) -> None: - """Assert a public tensor equals a backend tensor reordered to user order.""" - torch.testing.assert_close(public_tensor, backend_tensor.to(device=public_tensor.device)[:, user_to_backend]) + Args: + angle: Rotation angles [rad], shape [N]. + vector: Vector to rotate, shape [3] or [N, 3]. + Returns: + Rotated vectors, shape [N, 3]. + """ + vector = torch.as_tensor(vector, dtype=angle.dtype, device=angle.device).expand(*angle.shape, 3) + cos, sin = torch.cos(angle), torch.sin(angle) + return torch.stack( + (cos * vector[..., 0] - sin * vector[..., 1], sin * vector[..., 0] + cos * vector[..., 1], vector[..., 2]), -1 + ) -def _assert_user_write_reaches_backend( - user_tensor: torch.Tensor, backend_tensor: torch.Tensor, backend_to_user: list[int] -) -> None: - """Assert a user-order write reached backend storage in backend order.""" - torch.testing.assert_close(backend_tensor.to(device=user_tensor.device), user_tensor[:, backend_to_user]) +def _cross_z(rate: torch.Tensor, vector: torch.Tensor) -> torch.Tensor: + """Return ``(rate * z) x vector``, the velocity induced by a rotation about the world z axis. -@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 build_simulation_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 + Args: + rate: Angular rates about the world z axis [rad/s], shape [N]. + vector: Lever arms in the world frame [m], shape [N, 3]. + Returns: + Induced linear velocities [m/s], shape [N, 3]. + """ + return torch.stack((-rate * vector[..., 1], rate * vector[..., 0], torch.zeros_like(rate)), -1) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("gravity_enabled", [False]) -def test_live_manual_root_preserving_ordering_reorders_backend_reads_and_writes(sim, device, gravity_enabled): - """Smoke-test non-identity joint/body ordering through a live PhysX articulation.""" - articulation_cfg = replace( - FRANKA_PANDA_CFG, - prim_path="/World/Robot", - joint_ordering=tuple(reversed(PANDA_JOINT_NAMES)), - body_ordering=PANDA_ROOT_PRESERVING_REVERSED_BODY_NAMES, - ) - articulation = Articulation(articulation_cfg) - sim.reset() - assert articulation.is_initialized - assert articulation.backend_joint_names == list(PANDA_JOINT_NAMES) - assert articulation.backend_body_names == list(PANDA_BODY_NAMES) - assert articulation.joint_ordering is not None - assert articulation.body_ordering is not None +def _expected_fixed_branching_state( + articulation: Articulation, origins: torch.Tensor, com_b: torch.Tensor +) -> dict[str, torch.Tensor]: + """Compute link and center-of-mass kinematics of a fixed branching island from its joint state. - joint_user_to_backend = list(articulation.joint_ordering.user_to_backend_indices) - joint_backend_to_user = list(articulation.joint_ordering.backend_to_user_indices) - body_user_to_backend = list(articulation.body_ordering.user_to_backend_indices) + Every joint rotates about +z and is anchored at its child link origin, so the expected state follows from + the literal fixture geometry, the joint state, and the body-frame center-of-mass offsets. - joint_index = torch.arange(articulation.num_joints, device=device, dtype=torch.float32).unsqueeze(0) - joint_pos = torch.linspace(-0.3, 0.3, articulation.num_joints, device=device).unsqueeze(0) - joint_vel = torch.linspace(0.05, 0.13, articulation.num_joints, device=device).unsqueeze(0) - joint_stiffness = 10.0 + joint_index + Args: + articulation: Fixed-base branching articulation. + origins: World position of each environment [m], shape [N, 3]. + com_b: Center-of-mass offsets in the link frames [m], shape [N, B, 3], in public body order. - articulation.write_joint_stiffness_to_sim_index(stiffness=joint_stiffness, full_data=True) - articulation.write_joint_state_to_sim_index(position=joint_pos, velocity=joint_vel, full_data=True) - articulation.write_data_to_sim() + Returns: + World-frame link x axes (``heading``), link and center-of-mass positions [m], angular velocities + [rad/s], and link and center-of-mass linear velocities [m/s], each of shape [N, B, 3]. + """ + q = dict(zip(articulation.joint_names, articulation.data.joint_pos.torch.unbind(-1))) + qd = dict(zip(articulation.joint_names, articulation.data.joint_vel.torch.unbind(-1))) + zero = torch.zeros_like(q["left_shoulder"]) + links = {"base": (zero, origins, zero, torch.zeros_like(origins))} + for side in ("left", "right"): + shoulder, elbow = f"{side}_shoulder", f"{side}_elbow" + upper_pos = origins + torch.tensor(_JOINT_ANCHORS[shoulder], device=origins.device) + tip_lever = _rotate_z(q[shoulder], _JOINT_ANCHORS[elbow]) + links[f"{side}_upper"] = (q[shoulder], upper_pos, qd[shoulder], torch.zeros_like(origins)) + links[f"{side}_tip"] = ( + q[shoulder] + q[elbow], + upper_pos + tip_lever, + qd[shoulder] + qd[elbow], + _cross_z(qd[shoulder], tip_lever), + ) + expected = {key: [] for key in ("heading", "link_pos", "ang_vel", "link_lin_vel", "com_pos", "com_lin_vel")} + for body_index, body_name in enumerate(articulation.body_names): + yaw, link_pos, yaw_rate, link_lin_vel = links[body_name] + com_lever = _rotate_z(yaw, com_b[:, body_index]) + expected["heading"].append(_rotate_z(yaw, (1.0, 0.0, 0.0))) + expected["link_pos"].append(link_pos) + expected["ang_vel"].append(torch.stack((zero, zero, yaw_rate), -1)) + expected["link_lin_vel"].append(link_lin_vel) + expected["com_pos"].append(link_pos + com_lever) + expected["com_lin_vel"].append(link_lin_vel + _cross_z(yaw_rate, com_lever)) + return {key: torch.stack(values, dim=1) for key, values in expected.items()} + + +def _assert_jacobian_contract(articulation: Articulation, generalized_velocity: torch.Tensor) -> None: + """Check the link and center-of-mass Jacobians against body velocities and the mass-matrix contract. - _assert_user_write_reaches_backend( - joint_pos, _to_device_tensor(articulation.root_view.get_dof_positions(), device), joint_backend_to_user - ) - _assert_user_write_reaches_backend( - joint_vel, _to_device_tensor(articulation.root_view.get_dof_velocities(), device), joint_backend_to_user - ) - _assert_user_write_reaches_backend( - joint_stiffness, - _to_device_tensor(articulation.root_view.get_dof_stiffnesses(), device), - joint_backend_to_user, + Args: + articulation: Articulation whose state was just written. + generalized_velocity: Base and joint velocities, shape [N, num_base_dofs + num_joints]. + """ + data = articulation.data + num_dofs = articulation.num_joints + articulation.num_base_dofs + first_body = 1 if articulation.is_fixed_base else 0 + num_jacobian_bodies = articulation.num_bodies - first_body + # Read the body state first: it refreshes forward kinematics after the joint write. + lin_vels = (data.body_link_lin_vel_w.torch.clone(), data.body_com_lin_vel_w.torch.clone()) + ang_vel = data.body_link_ang_vel_w.torch.clone() + for jacobian, lin_vel in zip((data.body_link_jacobian_w.torch, data.body_com_jacobian_w.torch), lin_vels): + # Shape contract: fixed-base Jacobians omit the root body; floating-base ones prepend the base DoFs. + assert jacobian.shape == (_NUM_ENVS, num_jacobian_bodies, 6, num_dofs) + predicted = torch.einsum("nbij,nj->nbi", jacobian, generalized_velocity) + torch.testing.assert_close(predicted[..., :3], lin_vel[:, first_body:], atol=1e-4, rtol=1e-4) + torch.testing.assert_close(predicted[..., 3:], ang_vel[:, first_body:], atol=1e-4, rtol=1e-4) + # The joint-space mass matrix is square over all DoFs, symmetric, and positive definite. + mass_matrix = data.mass_matrix.torch + assert mass_matrix.shape == (_NUM_ENVS, num_dofs, num_dofs) + assert (mass_matrix.diagonal(dim1=-2, dim2=-1) > 1e-6).all() + asymmetry = (mass_matrix - mass_matrix.transpose(-1, -2)).abs().max().item() + assert asymmetry < 1e-4, f"|M - M^T|_max = {asymmetry:.3e} — mass matrix is not symmetric" + eye = torch.eye(num_dofs, device=mass_matrix.device, dtype=mass_matrix.dtype).expand_as(mass_matrix) + torch.linalg.cholesky(mass_matrix + 1e-6 * eye) + # The generalized kinetic energy equals the kinetic energy of the bodies. + generalized_energy = 0.5 * torch.einsum("ni,nij,nj->n", generalized_velocity, mass_matrix, generalized_velocity) + body_velocity = data.body_com_vel_w.torch + body_inertia = data.body_inertia.torch.reshape(_NUM_ENVS, articulation.num_bodies, 3, 3) + body_rotation = matrix_from_quat(data.body_com_quat_w.torch) + world_inertia = body_rotation @ body_inertia @ body_rotation.transpose(-1, -2) + body_energy = 0.5 * ( + (data.body_mass.torch.unsqueeze(-1) * body_velocity[..., :3].square()).sum((-1, -2)) + + torch.einsum("nbi,nbij,nbj->n", body_velocity[..., 3:], world_inertia, body_velocity[..., 3:]) ) + torch.testing.assert_close(generalized_energy, body_energy, atol=1e-5, rtol=1e-4) - sim.step() - articulation.update(sim.cfg.dt) - _assert_backend_to_user( - articulation.data.joint_pos.torch, - _to_device_tensor(articulation.root_view.get_dof_positions(), device), - joint_user_to_backend, - ) - _assert_backend_to_user( - articulation.data.joint_vel.torch, - _to_device_tensor(articulation.root_view.get_dof_velocities(), device), - joint_user_to_backend, - ) - _assert_backend_to_user( - articulation.data.joint_stiffness.torch, - _to_device_tensor(articulation.root_view.get_dof_stiffnesses(), device), - joint_user_to_backend, +@pytest.mark.isaacsim_ci +def test_articulation_drive_and_dynamics(articulation_scene: _ArticulationScene) -> None: + """Drive targets and efforts move the selected joints; dynamics quantities match the live state.""" + scene = articulation_scene + device = scene.device + articulation = scene.ordered + joint_backend_to_user = list(articulation.joint_ordering.backend_to_user_indices) + left_shoulder = articulation.find_joints("left_shoulder")[0][0] + right_shoulder = articulation.find_joints("right_shoulder")[0][0] + + # Start at rest on the drive targets, then command one position target and one effort in environment 1. + zeros = torch.zeros((_NUM_ENVS, articulation.num_joints), device=device) + articulation.write_joint_state_to_sim_index(position=zeros, velocity=zeros, full_data=True) + target_command = articulation.actuators.target_command + target_command.set_position_index(value=zeros) + target_command.set_velocity_index(value=zeros) + target_command.set_effort_index(value=zeros) + target_command.set_position_index( + value=torch.tensor([[0.4]], device=device), joint_ids=[left_shoulder], env_ids=[1] ) - _assert_backend_to_user( - articulation.data.body_link_pose_w.torch, - _to_device_tensor(articulation.root_view.get_link_transforms(), device), - body_user_to_backend, + target_command.set_effort_index( + value=torch.tensor([[10.0]], device=device), joint_ids=[right_shoulder], env_ids=[1] ) - _assert_backend_to_user( - articulation.data.body_com_pose_b.torch, - _to_device_tensor(articulation.root_view.get_coms(), device), - body_user_to_backend, + articulation.write_data_to_sim() + expected_targets = zeros.clone() + expected_targets[1, left_shoulder] = 0.4 + expected_efforts = zeros.clone() + expected_efforts[1, right_shoulder] = 10.0 + torch.testing.assert_close( + wp.to_torch(articulation.root_view.get_dof_position_targets()).to(device), + expected_targets[:, joint_backend_to_user], ) - - # The split accessors must be sliced from the reordered public poses, not from a stale backend-order cache. - torch.testing.assert_close(articulation.data.body_com_pos_b.torch, articulation.data.body_com_pose_b.torch[..., :3]) torch.testing.assert_close( - articulation.data.body_com_quat_b.torch, articulation.data.body_com_pose_b.torch[..., 3:] + wp.to_torch(articulation.root_view.get_dof_actuation_forces()).to(device), + expected_efforts[:, joint_backend_to_user], ) - torch.testing.assert_close(articulation.data.body_pos_w.torch, articulation.data.body_link_pose_w.torch[..., :3]) - torch.testing.assert_close(articulation.data.body_quat_w.torch, articulation.data.body_link_pose_w.torch[..., 3:]) - - -@pytest.mark.parametrize("device", ["cpu"]) -@pytest.mark.parametrize("gravity_enabled", [False]) -def test_reversed_joint_dynamics_use_public_joint_basis(sim, device, gravity_enabled): - """Keep dynamics tensors consistent with public joint velocity.""" - articulation = Articulation( - ArticulationCfg( - prim_path="/World/Robot", - spawn=sim_utils.UsdFileCfg( - usd_path=str(Path(__file__).parent / "data" / "articulation_ordering_branching.usda") - ), - actuators={}, + # The effort must reach the solver, not only its staging buffer. + scene.step() + assert articulation.data.joint_vel.torch[1, right_shoulder] > 1e-3, "a positive effort must accelerate the joint" + scene.step(19) + # The commanded joint moves toward its target; environment 0 stays at rest on its targets. + assert 0.05 < articulation.data.joint_pos.torch[1, left_shoulder] < 0.4 + torch.testing.assert_close(articulation.data.joint_pos.torch[0], zeros[0], atol=1e-6, rtol=0.0) + torch.testing.assert_close(articulation.data.joint_vel.torch[0], zeros[0], atol=1e-6, rtol=0.0) + # The fixed root stays at its default state. + for fixed_articulation, origins in ( + (articulation, scene.origins["ordered"]), + (scene.tendon, scene.origins["tendon"]), + ): + default_root_pose = fixed_articulation.data.default_root_pose.torch.clone() + default_root_pose[:, :3] += origins + torch.testing.assert_close(fixed_articulation.data.root_link_pose_w.torch, default_root_pose) + torch.testing.assert_close( + fixed_articulation.data.root_com_vel_w.torch, fixed_articulation.data.default_root_vel.torch ) - ) - UsdPhysics.FixedJoint.Define(sim.stage, "/World/Robot/fixed_root").GetBody1Rel().SetTargets(["/World/Robot/base"]) - joint = UsdPhysics.RevoluteJoint.Get(sim.stage, "/World/Robot/left_elbow") - body0, body1 = joint.GetBody0Rel().GetTargets(), joint.GetBody1Rel().GetTargets() - joint.GetBody0Rel().SetTargets(body1) - joint.GetBody1Rel().SetTargets(body0) - sim.reset() - - velocity = torch.zeros((1, articulation.num_joints), device=device) - velocity[:, articulation.find_joints("left_shoulder")[0][0]] = 0.4 - velocity[:, articulation.find_joints("left_elbow")[0][0]] = 0.7 - articulation.write_joint_velocity_to_sim_index(velocity=velocity) - sim.step() - articulation.update(sim.cfg.dt) - - joint_velocity = articulation.data.joint_vel.torch - predicted_velocity = torch.einsum("nbij,nj->nbi", articulation.data.body_com_jacobian_w.torch, joint_velocity) - torch.testing.assert_close(predicted_velocity, articulation.data.body_com_vel_w.torch[:, 1:], atol=1e-5, rtol=1e-5) - generalized_energy = 0.5 * torch.einsum( - "ni,nij,nj->n", joint_velocity, articulation.data.mass_matrix.torch, joint_velocity + # With known center-of-mass offsets, the link and center-of-mass kinematics follow the joint state. + com_b = torch.tensor( + [[0.0, 0.0, 0.0], [0.03, 0.02, -0.01], [0.05, -0.02, 0.01], [-0.04, 0.01, 0.02], [0.02, 0.03, -0.02]], + device=device, + ).repeat(_NUM_ENVS, 1, 1) + coms = torch.cat((com_b, torch.tensor([0.0, 0.0, 0.0, 1.0], device=device).expand(_NUM_ENVS, 5, 4)), dim=-1) + articulation.set_coms_index(coms=coms, full_data=True) + # PhysX applies new mass properties to the kinematic state on the next step. + scene.step() + joint_pos = torch.tensor([[0.3, -0.4, 0.2, 0.5], [-0.2, 0.6, -0.5, 0.1]], device=device) + joint_vel = torch.tensor([[0.8, -0.6, 0.5, 0.7], [-0.4, 0.9, 0.3, -0.5]], device=device) + articulation.write_joint_state_to_sim_index(position=joint_pos, velocity=joint_vel, full_data=True) + expected = _expected_fixed_branching_state(articulation, scene.origins["ordered"], com_b) + data = articulation.data + torch.testing.assert_close( + quat_apply(data.body_link_quat_w.torch, torch.tensor([1.0, 0.0, 0.0], device=device).expand(_NUM_ENVS, 5, 3)), + expected["heading"], + atol=1e-4, + rtol=1e-4, ) - body_velocity = articulation.data.body_com_vel_w.torch - body_inertia = articulation.data.body_inertia.torch.reshape(1, articulation.num_bodies, 3, 3) - body_energy = 0.5 * ( - (articulation.data.body_mass.torch.unsqueeze(-1) * body_velocity[..., :3].square()).sum((-1, -2)) - + torch.einsum("nbi,nbij,nbj->n", body_velocity[..., 3:], body_inertia, body_velocity[..., 3:]) + for public, key in ( + (data.body_link_pos_w.torch, "link_pos"), + (data.body_link_lin_vel_w.torch, "link_lin_vel"), + (data.body_link_ang_vel_w.torch, "ang_vel"), + (data.body_com_pos_w.torch, "com_pos"), + (data.body_com_lin_vel_w.torch, "com_lin_vel"), + (data.body_com_ang_vel_w.torch, "ang_vel"), + ): + torch.testing.assert_close(public, expected[key], atol=1e-4, rtol=1e-4, msg=key) + _assert_jacobian_contract(articulation, data.joint_vel.torch) + + # A reversed joint keeps the dynamics in the public joint basis; floating Jacobians prepend the base DoFs and, + # with center-of-mass offsets on the moving links, shift their linear rows to the link origins. + scene.floating.set_coms_index(coms=coms, full_data=True) + scene.step() + for island in (scene.tendon, scene.floating): + island.write_joint_state_to_sim_index(position=joint_pos, velocity=joint_vel, full_data=True) + root_velocity = torch.tensor([[0.3, -0.2, 0.1, 0.2, 0.1, -0.3], [-0.1, 0.2, 0.3, -0.2, 0.3, 0.1]], device=device) + scene.floating.write_root_link_velocity_to_sim_index(root_velocity=root_velocity) + _assert_jacobian_contract(scene.tendon, scene.tendon.data.joint_vel.torch) + floating_velocity = torch.cat( + (scene.floating.data.root_link_vel_w.torch, scene.floating.data.joint_vel.torch), dim=-1 ) - torch.testing.assert_close(generalized_energy, body_energy, atol=1e-5, rtol=1e-5) + _assert_jacobian_contract(scene.floating, floating_velocity) + + # Each Jacobian and mass-matrix read independently reflects a manual joint write without stepping. + for island in (articulation, scene.floating): + q_initial = island.data.joint_pos.torch.clone()[:1] + env_ids = wp.array([0], dtype=wp.int32, device=device) + for property_name in ("body_link_jacobian_w", "body_com_jacobian_w", "mass_matrix"): + island.write_joint_position_to_sim_index(position=q_initial, env_ids=env_ids) + before = getattr(island.data, property_name).torch.clone() + # A separate write for each getter prevents another getter from refreshing FK on its behalf. + island.write_joint_position_to_sim_index(position=q_initial + 0.5, env_ids=env_ids) + after = getattr(island.data, property_name).torch.clone() + assert not torch.allclose(before, after, atol=1e-3), f"{property_name} stayed stale after a joint write" + + +def _place_at_rest(articulation: Articulation, root_pose: torch.Tensor) -> None: + """Teleport a floating island to ``root_pose`` with zero root and joint velocities and no external wrench.""" + zeros = torch.zeros((_NUM_ENVS, articulation.num_joints), device=root_pose.device) + articulation.write_root_link_pose_to_sim_index(root_pose=root_pose) + articulation.write_root_velocity_to_sim_index(root_velocity=torch.zeros((_NUM_ENVS, 6), device=root_pose.device)) + articulation.write_joint_state_to_sim_index(position=zeros, velocity=zeros, full_data=True) + articulation.reset() -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("gravity_enabled", [False]) -def test_live_floating_root_writers_match_identity_after_body_reordering(sim, device, gravity_enabled): - """Keep floating-base root writes invariant when public body order moves the root.""" - floating_spawn = replace(FRANKA_PANDA_CFG.spawn, fix_root_link=False) - identity = Articulation( - replace(FRANKA_PANDA_CFG, prim_path="/World/IdentityRobot", spawn=floating_spawn, body_ordering=None) - ) - ordered = Articulation( - replace( - FRANKA_PANDA_CFG, - prim_path="/World/OrderedRobot", - spawn=floating_spawn, - body_ordering=tuple(reversed(PANDA_BODY_NAMES)), - ) +def test_floating_articulation_root_writes_and_wrenches(articulation_scene: _ArticulationScene, monkeypatch) -> None: + """Root writes and external wrenches reach the selected floating roots in the frames they are given in.""" + scene = articulation_scene + device = scene.device + floating, reordered = scene.floating, scene.reordered + rest_pose = torch.cat( + (scene.origins["floating"], torch.tensor(_yaw_quat(0.5 * math.pi), device=device).expand(_NUM_ENVS, 4)), -1 ) - - sim.reset() - assert identity.is_initialized and ordered.is_initialized - assert not identity.is_fixed_base and not ordered.is_fixed_base - assert identity.body_ordering is None - assert ordered.body_ordering is not None - assert ordered.body_ordering.backend_to_user_indices[0] != 0 - - backend_coms = torch.zeros((1, len(PANDA_BODY_NAMES), 7), device=device) - body_index = torch.arange(len(PANDA_BODY_NAMES), device=device, dtype=torch.float32) - backend_coms[0, :, 0] = 0.05 + 0.01 * body_index - backend_coms[0, :, 1] = -0.03 - 0.02 * body_index - backend_coms[0, :, 2] = 0.02 + 0.03 * body_index + _place_at_rest(floating, rest_pose) + reordered_rest_pose = rest_pose.clone() + reordered_rest_pose[:, :3] = scene.origins["reordered"] + _place_at_rest(reordered, reordered_rest_pose) + + # A partial root pose write reaches only the selected root. + target_pose = rest_pose[1:].clone() + target_pose[:, :3] += torch.tensor([0.2, -0.1, 0.3], device=device) + target_pose[:, 3:] = torch.tensor(_yaw_quat(0.7), device=device) + floating.write_root_link_pose_to_sim_index(root_pose=target_pose, env_ids=[1]) + expected_pose = torch.cat((rest_pose[:1], target_pose)) + torch.testing.assert_close(floating.data.root_link_pose_w.torch, expected_pose) + torch.testing.assert_close(wp.to_torch(floating.root_view.get_root_transforms()).to(device), expected_pose) + + # Root writers are invariant to a public body order that moves the root link. + assert reordered.body_ordering is not None + assert reordered.body_ordering.backend_to_user_indices[0] != 0 + backend_coms = torch.zeros((_NUM_ENVS, 5, 7), device=device) + body_index = torch.arange(5, device=device, dtype=torch.float32) + backend_coms[:, :, 0] = 0.05 + 0.01 * body_index + backend_coms[:, :, 1] = -0.03 - 0.02 * body_index + backend_coms[:, :, 2] = 0.02 + 0.03 * body_index backend_coms[..., 6] = 1.0 - identity.set_coms_index( - coms=wp.from_torch(backend_coms.contiguous(), dtype=wp.transformf), - full_data=True, - ) - ordered_user_to_backend = list(ordered.body_ordering.user_to_backend_indices) - ordered.set_coms_index( - coms=wp.from_torch(backend_coms[:, ordered_user_to_backend].contiguous(), dtype=wp.transformf), - full_data=True, - ) - torch.testing.assert_close(_to_device_tensor(identity.root_view.get_coms(), device), backend_coms) - torch.testing.assert_close(_to_device_tensor(ordered.root_view.get_coms(), device), backend_coms) - - root_com_pose = torch.tensor([[1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 1.0]], device=device) - root_link_velocity = torch.tensor([[0.4, -0.3, 0.2, 1.1, -0.7, 0.9]], device=device) - for articulation in (identity, ordered): - articulation.write_root_com_pose_to_sim_index(root_pose=root_com_pose) + reordered_user_to_backend = list(reordered.body_ordering.user_to_backend_indices) + floating.set_coms_index(coms=backend_coms, full_data=True) + reordered.set_coms_index(coms=backend_coms[:, reordered_user_to_backend], full_data=True) + torch.testing.assert_close(wp.to_torch(floating.root_view.get_coms()).to(device), backend_coms) + torch.testing.assert_close(wp.to_torch(reordered.root_view.get_coms()).to(device), backend_coms) + root_com_pose = rest_pose.clone() + root_com_pose[:, :3] += torch.tensor([0.5, 0.25, 0.75], device=device) + root_com_pose[:, 3:] = torch.tensor(_yaw_quat(0.6), device=device) + island_offset = scene.origins["reordered"] - scene.origins["floating"] + root_link_velocity = torch.tensor([[0.4, -0.3, 0.2, 1.1, -0.7, 0.9]], device=device).repeat(_NUM_ENVS, 1) + for articulation, offset in ((floating, 0.0), (reordered, island_offset)): + pose = root_com_pose.clone() + pose[:, :3] += offset + articulation.write_root_com_pose_to_sim_index(root_pose=pose) articulation.write_root_link_velocity_to_sim_index(root_velocity=root_link_velocity) - + torch.testing.assert_close(articulation.data.root_com_pose_w.torch, pose) + torch.testing.assert_close(articulation.data.root_link_vel_w.torch, root_link_velocity) + reordered_transforms = wp.to_torch(reordered.root_view.get_root_transforms()).to(device).clone() + reordered_transforms[:, :3] -= island_offset + torch.testing.assert_close(reordered_transforms, wp.to_torch(floating.root_view.get_root_transforms()).to(device)) torch.testing.assert_close( - _to_device_tensor(ordered.root_view.get_root_transforms(), device), - _to_device_tensor(identity.root_view.get_root_transforms(), device), + wp.to_torch(reordered.root_view.get_root_velocities()).to(device), + wp.to_torch(floating.root_view.get_root_velocities()).to(device), + ) + torch.testing.assert_close(reordered.data.root_com_vel_w.torch, floating.data.root_com_vel_w.torch) + # The derived root link pose carries the written center-of-mass pose at the root's center-of-mass offset. + for articulation, offset in ((floating, 0.0), (reordered, island_offset)): + link_pose = articulation.data.root_link_pose_w.torch + root_com_b = articulation.data.body_com_pose_b.torch[:, articulation.body_names.index("base")] + com_pos, com_quat = combine_frame_transforms( + link_pose[:, :3], link_pose[:, 3:], root_com_b[:, :3], root_com_b[:, 3:] + ) + expected_com_pose = root_com_pose.clone() + expected_com_pose[:, :3] += offset + torch.testing.assert_close(torch.cat((com_pos, com_quat), -1), expected_com_pose) + + # Link-pose and center-of-mass velocity writes keep the other frame consistent with the offset. + floating.write_root_link_pose_to_sim_index(root_pose=root_com_pose) + floating.write_root_com_velocity_to_sim_index(root_velocity=root_link_velocity) + torch.testing.assert_close(floating.data.root_link_pose_w.torch, root_com_pose) + torch.testing.assert_close(floating.data.root_com_vel_w.torch, root_link_velocity) + expected_com_pos, expected_com_quat = combine_frame_transforms( + root_com_pose[:, :3], root_com_pose[:, 3:], backend_coms[:, 0, :3], backend_coms[:, 0, 3:] ) torch.testing.assert_close( - _to_device_tensor(ordered.root_view.get_root_velocities(), device), - _to_device_tensor(identity.root_view.get_root_velocities(), device), + floating.data.root_com_pose_w.torch, torch.cat((expected_com_pos, expected_com_quat), -1) ) - torch.testing.assert_close(ordered.data.root_com_vel_w.torch, identity.data.root_com_vel_w.torch) - for articulation in (identity, ordered): - torch.testing.assert_close(articulation.data.root_com_pose_w.torch, root_com_pose) - torch.testing.assert_close(articulation.data.root_link_vel_w.torch, root_link_velocity) - - -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("gravity_enabled", [False]) -@pytest.mark.parametrize("body_ordering", ["identity", "reversed"]) -def test_live_direct_view_mass_inertia_writes_become_visible(sim, device, gravity_enabled, body_ordering): - """Direct tensor-view writes to body masses/inertias become visible on the lazy read. - - Develop reads these properties directly from the tensor view on every access. The timestamp-lazy - implementation could instead hide direct ``root_view.set_masses`` / ``set_inertias`` writes after - its first read. This regression requires those writes to become visible through ``data.body_mass`` / - ``data.body_inertia`` when the lazy buffer is next eligible to refresh: - - - Case A: after a primed read and a subsequent simulation update (the lazy gate opens once per - step). - - Case B: on the very first read of a cold buffer. - - Case B runs first so that the buffers are still cold. Both identity and non-identity (reversed) body - ordering are covered; under ordering the public buffers must equal the backend-order view gathered through - ``user_to_backend``. - """ - body_ordering_arg = None if body_ordering == "identity" else PANDA_ROOT_PRESERVING_REVERSED_BODY_NAMES - articulation = Articulation(replace(FRANKA_PANDA_CFG, prim_path="/World/Robot", body_ordering=body_ordering_arg)) - sim.reset() - assert articulation.is_initialized - - if body_ordering == "identity": - assert articulation.body_ordering is None - body_user_to_backend = list(range(articulation.num_bodies)) - else: - assert articulation.body_ordering is not None - body_user_to_backend = list(articulation.body_ordering.user_to_backend_indices) - - cpu_env_ids = wp.array(list(range(articulation.num_instances)), dtype=wp.int32, device="cpu") + torch.testing.assert_close(floating.data.root_com_vel_w.torch[:, 3:], floating.data.root_link_vel_w.torch[:, 3:]) - def write_backend_mass_inertia(delta_mass: float, delta_inertia: float) -> tuple[torch.Tensor, torch.Tensor]: - """Write distinct backend-order masses/inertias straight through the tensor view.""" - backend_masses = wp.to_torch(articulation.root_view.get_masses()).clone() + delta_mass - backend_inertias = wp.to_torch(articulation.root_view.get_inertias()).clone() + delta_inertia - articulation.root_view.set_masses( - wp.from_torch(backend_masses.contiguous(), dtype=wp.float32), indices=cpu_env_ids - ) - articulation.root_view.set_inertias( - wp.from_torch(backend_inertias.contiguous(), dtype=wp.float32), indices=cpu_env_ids + # A body-frame force on environment 1 accelerates only that root, along the rotated force direction. + base = [floating.find_bodies("base")[0][0]] + _place_at_rest(floating, rest_pose) + floating.permanent_wrench_composer.set_forces_and_torques_index( + forces=torch.tensor([[[8.0, 0.0, 0.0]]], device=device), body_ids=base, env_ids=[1] + ) + scene.step() + root_velocity = floating.data.root_com_lin_vel_w.torch + assert root_velocity[1, 1] > 0.01, f"body-frame force did not act along the rotated axis: {root_velocity[1]}" + torch.testing.assert_close(root_velocity[0], torch.zeros(3, device=device), atol=1e-6, rtol=0.0) + + # A world-frame force, and a force at a world-frame position, act like their body-frame counterparts. + for local_wrench, global_wrench, response in ( + ({"forces": [[[8.0, 0.0, 0.0]]]}, {"forces": [[[0.0, 8.0, 0.0]]]}, lambda: floating.data.root_com_lin_vel_w), + ( + {"forces": [[[0.0, 0.0, 8.0]]], "positions": [[[0.0, 1.0, 0.0]]]}, + {"forces": [[[0.0, 0.0, 8.0]]], "positions": [[[-1.0, 0.0, 0.0]]]}, + lambda: floating.data.root_ang_vel_b, + ), + ): + _place_at_rest(floating, rest_pose) + local_wrench = {key: torch.tensor(value, device=device) for key, value in local_wrench.items()} + global_wrench = {key: torch.tensor(value, device=device) for key, value in global_wrench.items()} + if "positions" in global_wrench: + # World-frame positions add the rotated body-frame lever arm to the base center of mass. + global_wrench["positions"] = global_wrench["positions"] + floating.data.body_com_pos_w.torch[:1, base] + floating.permanent_wrench_composer.set_forces_and_torques_index(body_ids=base, env_ids=[1], **local_wrench) + floating.permanent_wrench_composer.set_forces_and_torques_index( + body_ids=base, env_ids=[0], is_global=True, **global_wrench ) - return backend_masses, backend_inertias - - # Case B: the buffers have not been read since reset, so the first read must reflect the write. - backend_masses, backend_inertias = write_backend_mass_inertia(0.293, 0.023) - _assert_backend_to_user(articulation.data.body_mass.torch, backend_masses, body_user_to_backend) - _assert_backend_to_user(articulation.data.body_inertia.torch, backend_inertias, body_user_to_backend) - - # Case A: the buffers are now primed; a new write becomes visible after the sim update opens the lazy gate. - backend_masses, backend_inertias = write_backend_mass_inertia(0.137, 0.011) - articulation.update(sim.cfg.dt) - _assert_backend_to_user(articulation.data.body_mass.torch, backend_masses, body_user_to_backend) - _assert_backend_to_user(articulation.data.body_inertia.torch, backend_inertias, body_user_to_backend) - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("gravity_enabled", [False]) -def test_branching_fixture_resolves_distinct_conventions(sim, device, gravity_enabled): - """Resolve concrete breadth-first PhysX and depth-first MJWarp name orders.""" - fixture_path = Path(__file__).parent / "data" / "articulation_ordering_branching.usda" - articulation = Articulation( - ArticulationCfg( - prim_path="/World/Robot", - spawn=sim_utils.UsdFileCfg(usd_path=str(fixture_path)), - actuators={}, - joint_ordering="mjwarp", - body_ordering="mjwarp", + scene.step() + response_value = response().torch + torch.testing.assert_close(response_value[0], response_value[1], atol=1e-4, rtol=1e-3) + # Applied 1 m along the base y-axis instead of at the center of mass, the upward force also rolls the base + # about its x-axis. + _place_at_rest(floating, rest_pose) + upward_force = torch.tensor([[[0.0, 0.0, 8.0]]], device=device) + floating.permanent_wrench_composer.set_forces_and_torques_index(forces=upward_force, body_ids=base, env_ids=[0]) + floating.permanent_wrench_composer.set_forces_and_torques_index( + forces=upward_force, positions=torch.tensor([[[0.0, 1.0, 0.0]]], device=device), body_ids=base, env_ids=[1] + ) + scene.step() + roll_rate = floating.data.root_ang_vel_b.torch[:, 0] + assert roll_rate[1] - roll_rate[0] > 0.1, roll_rate + + # Forces on several bodies push each selected body along its own force direction. Under a public body order + # that differs from the PhysX link order, every body responds as in the identity-ordered island. + tip_forces = torch.tensor([[[0.0, 4.0, 0.0], [0.0, -4.0, 0.0]]], device=device) + for articulation, pose in ((floating, rest_pose), (reordered, reordered_rest_pose)): + _place_at_rest(articulation, pose) + tips = articulation.find_bodies(["left_tip", "right_tip"], preserve_order=True)[0] + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=tip_forces, body_ids=tips, env_ids=[1] ) + # Refresh the link transforms after the teleport: on the GPU pipeline, the step applies link-frame wrenches + # with the link transforms of the last kinematic update. + articulation.data.body_link_pose_w.torch + scene.step() + tips = reordered.find_bodies(["left_tip", "right_tip"], preserve_order=True)[0] + tip_velocity = reordered.data.body_link_lin_vel_w.torch[:, tips] + tip_direction = quat_apply(reordered.data.body_link_quat_w.torch[1, tips], tip_forces[0]) + assert torch.all((tip_velocity[1] * tip_direction).sum(-1) > 1e-3), tip_velocity[1] + torch.testing.assert_close(tip_velocity[0], torch.zeros_like(tip_velocity[0]), atol=1e-6, rtol=0.0) + floating_order = [reordered.body_names.index(name) for name in floating.body_names] + torch.testing.assert_close( + reordered.data.body_link_lin_vel_w.torch[:, floating_order], floating.data.body_link_lin_vel_w.torch ) - sim.reset() - assert articulation.is_initialized - - assert tuple(articulation.backend_joint_names) == BRANCHING_PHYSX_JOINT_NAMES - assert tuple(articulation.backend_body_names) == BRANCHING_PHYSX_BODY_NAMES - assert get_articulation_name_ordering(articulation, "mjwarp", "joint") == BRANCHING_MJWARP_JOINT_NAMES - assert get_articulation_name_ordering(articulation, "mjwarp", "body") == BRANCHING_MJWARP_BODY_NAMES - assert tuple(articulation.joint_names) == BRANCHING_MJWARP_JOINT_NAMES - assert tuple(articulation.body_names) == BRANCHING_MJWARP_BODY_NAMES - assert articulation.joint_ordering is not None - assert articulation.body_ordering is not None + # A partial reset clears the selected environment of the actuators and wrench composers; a full reset resets + # every actuator environment and clears all external forces and torques. + actuator = next(iter(floating.actuators.values())) + actuator_reset = actuator.reset + reset_env_ids = [] -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("add_ground_plane", [True]) -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 - - 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) + def record_actuator_reset(env_ids=None) -> None: + reset_env_ids.append(env_ids) + actuator_reset(env_ids) - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(articulation) < 10 + monkeypatch.setattr(actuator, "reset", record_actuator_reset) + composers = (floating.instantaneous_wrench_composer, floating.permanent_wrench_composer) + ones = torch.ones((_NUM_ENVS, floating.num_bodies, 3), device=device) + floating.permanent_wrench_composer.set_forces_and_torques_index(forces=ones, torques=ones) + floating.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones) + floating.reset(env_ids=torch.tensor([0], device=device)) + assert [env_ids.tolist() for env_ids in reset_env_ids] == [[0]] + for composer in composers: + assert composer.active + for buffer in (composer.out_force_b.torch, composer.out_torque_b.torch): + assert torch.count_nonzero(buffer[0]) == 0 + assert torch.count_nonzero(buffer[1:]) == buffer[1:].numel() + reset_env_ids.clear() + floating.reset() + assert reset_env_ids == [None] + for composer in composers: + assert not composer.active + assert torch.count_nonzero(composer.out_force_b.torch) == 0 + assert torch.count_nonzero(composer.out_torque_b.torch) == 0 - # 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) - - # -- link names (check within articulation ordering is correct) - prim_path_body_names = [path.split("/")[-1] for path in articulation.root_view.link_paths[0]] - assert prim_path_body_names == articulation.body_names - # -- actuator type - for actuator_name, actuator in articulation.actuators.items(): - is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) - assert getattr(actuator, "is_implicit_model", False) == is_implicit_model_cfg - - -@pytest.mark.parametrize("num_articulations", [1, 2]) -@pytest.mark.parametrize("device", test_devices()) -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) +def _set_spatial_tendon_properties( + articulation: Articulation, values: dict[str, torch.Tensor], env_ids: torch.Tensor | None = None +) -> None: + """Set every spatial tendon property of the selected environments without writing it to the simulation.""" + articulation.set_spatial_tendon_stiffness_index(stiffness=values["stiffness"], env_ids=env_ids) + articulation.set_spatial_tendon_limit_stiffness_index(limit_stiffness=values["limit_stiffness"], env_ids=env_ids) + articulation.set_spatial_tendon_damping_index(damping=values["damping"], env_ids=env_ids) + articulation.set_spatial_tendon_offset_index(offset=values["offset"], env_ids=env_ids) - # 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 +def test_spatial_tendon_properties_round_trip(articulation_scene: _ArticulationScene) -> None: + """A locally authored spatial tendon is discovered and its selected properties reach PhysX.""" + device = articulation_scene.device + articulation = articulation_scene.tendon 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) - - # -- link names (check within articulation ordering is correct) - prim_path_body_names = [path.split("/")[-1] for path in articulation.root_view.link_paths[0]] - assert prim_path_body_names == articulation.body_names - # -- actuator type - for actuator_name, actuator in articulation.actuators.items(): - is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) - assert getattr(actuator, "is_implicit_model", False) == 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", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_fixed_tendon_position_target_writes_offset(sim, num_articulations, device): - """A tendon length target lands in the simulation as ``rest_length - target`` on the selected cells only. + assert articulation.num_spatial_tendons == 1 + data = articulation.data + root_view = articulation.root_view + # Readers of the Lab data and of the solver for each property. + readers = { + "stiffness": (lambda: data.spatial_tendon_stiffness.torch, root_view.get_spatial_tendon_stiffnesses), + "limit_stiffness": ( + lambda: data.spatial_tendon_limit_stiffness.torch, + root_view.get_spatial_tendon_limit_stiffnesses, + ), + "damping": (lambda: data.spatial_tendon_damping.torch, root_view.get_spatial_tendon_dampings), + "offset": (lambda: data.spatial_tendon_offset.torch, root_view.get_spatial_tendon_offsets), + } - The index form commands every tendon of environment 0; the mask form commands tendon 0 of - environment 1. Every other cell must keep its initial offset. + # Distinct per-environment values so that an environment mix-up cannot pass + env_offset = torch.arange(_NUM_ENVS, dtype=torch.float32, device=device).unsqueeze(1) + values = {"stiffness": 10.0 + env_offset, "limit_stiffness": 20.0 + env_offset, "damping": 3.0 + env_offset} + values["offset"] = env_offset + _set_spatial_tendon_properties(articulation, values) + articulation.write_spatial_tendon_properties_to_sim_index() + for name, (_, read_view) in readers.items(): + torch.testing.assert_close(wp.to_torch(read_view()).to(device), values[name]) + + # Partial writes, first to environment 1 alone and then through unsorted int64 selectors with distinct rows, + # change the selected environments in the data and in the solver and preserve the others. + for env_ids, row_offsets in ( + (torch.tensor([1], dtype=torch.int32, device=device), (0.0,)), + (torch.tensor([1, 0], dtype=torch.int64, device=device), (1.0, 2.0)), + ): + partial = { + name: torch.tensor([[value + row_offset] for row_offset in row_offsets], device=device) + for name, value in (("stiffness", 12.0), ("limit_stiffness", 3.0), ("damping", 1.5), ("offset", 0.1)) + } + _set_spatial_tendon_properties(articulation, partial, env_ids=env_ids) + for name, (read_data, _) in readers.items(): + values[name][env_ids.long()] = partial[name] + torch.testing.assert_close(read_data(), values[name]) + articulation.write_spatial_tendon_properties_to_sim_index(env_ids=env_ids) + for name, (_, read_view) in readers.items(): + torch.testing.assert_close(wp.to_torch(read_view()).to(device), values[name]) + + +def test_fixed_tendon_position_target_writes_offset(articulation_scene: _ArticulationScene) -> None: + """A tendon length target reaches PhysX as ``rest_length - target`` in the selected environments only. + + The index form commands environment 0 and the mask form environment 1, with distinct values. """ - articulation_cfg = generate_articulation_cfg(articulation_type="shadow_hand") - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device) - - sim.reset() - assert articulation.is_initialized - assert articulation.is_fixed_base - 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) - for actuator_name, actuator in articulation.actuators.items(): - is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) - assert getattr(actuator, "is_implicit_model", False) == is_implicit_model_cfg - num_tendons = articulation.num_fixed_tendons - assert num_tendons > 0 + device = articulation_scene.device + articulation = articulation_scene.tendon + assert articulation.num_fixed_tendons == 1 rest_length = articulation.data.fixed_tendon_rest_length.torch.clone() - initial_offset = articulation.data.fixed_tendon_offset.torch.clone() - - index_target = torch.full((1, num_tendons), 0.3, dtype=torch.float32, device=device) - articulation.set_fixed_tendon_position_target_index(target=index_target, env_ids=[0]) - # Distinct per-cell values: a uniform target cannot catch the mask form reading the wrong - # cell, because every wrong read returns the same number. - mask_target = ( - 0.7 - + 0.1 * torch.arange(num_articulations, dtype=torch.float32, device=device).unsqueeze(1) - + 0.01 * torch.arange(num_tendons, dtype=torch.float32, device=device).unsqueeze(0) - ) - env_mask = wp.array([False, True], dtype=wp.bool, device=device) - tendon_mask = wp.array([i == 0 for i in range(num_tendons)], dtype=wp.bool, device=device) - articulation.set_fixed_tendon_position_target_mask( - target=mask_target, fixed_tendon_mask=tendon_mask, env_mask=env_mask - ) + torch.testing.assert_close(rest_length, torch.full((_NUM_ENVS, 1), 0.4, device=device)) + offset = articulation.data.fixed_tendon_offset.torch.clone() + articulation.set_fixed_tendon_position_target_index(target=torch.tensor([[0.3]], device=device), env_ids=[0]) articulation.write_data_to_sim() - sim.step() - articulation.update(sim.cfg.dt) - - expected = initial_offset.clone() - expected[0] = rest_length[0] - 0.3 - expected[1, 0] = rest_length[1, 0] - mask_target[1, 0] - torch.testing.assert_close(articulation.data.fixed_tendon_offset.torch, expected) - + offset[0] = rest_length[0] - 0.3 + torch.testing.assert_close(articulation.data.fixed_tendon_offset.torch, offset) + torch.testing.assert_close(wp.to_torch(articulation.root_view.get_fixed_tendon_offsets()).to(device), offset) -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("add_ground_plane", [True]) -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. + articulation.set_fixed_tendon_position_target_mask( + target=torch.tensor([[0.7], [0.8]], device=device), + env_mask=wp.array([False, True], dtype=wp.bool, device=device), + ) + articulation.write_data_to_sim() + offset[1] = rest_length[1] - 0.8 + torch.testing.assert_close(articulation.data.fixed_tendon_offset.torch, offset) + torch.testing.assert_close(wp.to_torch(articulation.root_view.get_fixed_tendon_offsets()).to(device), offset) + + +@pytest.mark.xfail(strict=True, reason="GPU PhysX applies mass and inertia writes on the next step") +def test_inertial_writes_refresh_the_mass_matrix(articulation_scene: _ArticulationScene) -> None: + """Mass and inertia writes refresh a mass matrix that was already read in the same step.""" + articulation = articulation_scene.ordered + data = articulation.data + writes = ( + (articulation.set_masses_index, "masses", data.body_mass), + (articulation.set_inertias_index, "inertias", data.body_inertia), + ) + for write, kwarg, values in writes: + initial = values.torch.clone() + before = data.mass_matrix.torch.clone() + write(**{kwarg: 2.0 * initial + 0.1}) + after = data.mass_matrix.torch.clone() + write(**{kwarg: initial}) + assert not torch.allclose(after, before, atol=1e-4), f"{kwarg} write left a stale matrix" + + +def test_joint_position_limit_clamping_respects_logging(articulation_scene: _ArticulationScene, caplog) -> None: + """Logging level controls reporting without changing clamping or reusing a stale violation count.""" + articulation = articulation_scene.ordered + logger = type(articulation).__module__ + original_limits = articulation.data.joint_pos_limits.torch.clone() + original_defaults = articulation.data.default_joint_pos.torch.clone() + limits = torch.zeros_like(original_limits) + limits[..., 1] = 0.5 + try: + for level in (logging.WARNING, logging.INFO): + articulation.data.default_joint_pos.torch.fill_(1.0) + caplog.clear() + with caplog.at_level(level, logger=logger): + articulation.write_joint_position_limit_to_sim_index(limits=limits, warn_limit_violation=False) + assert [record.levelno for record in caplog.records] == ([logging.INFO] if level == logging.INFO else []) + torch.testing.assert_close(articulation.data.default_joint_pos.torch, torch.full_like(limits[..., 1], 0.5)) + caplog.clear() + with caplog.at_level(logging.INFO, logger=logger): + articulation.write_joint_position_limit_to_sim_mask(limits=limits, warn_limit_violation=False) + assert not caplog.records + finally: + articulation.write_joint_position_limit_to_sim_index(limits=original_limits) + articulation.data.default_joint_pos.torch.copy_(original_defaults) - 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 = clone(generate_articulation_cfg(articulation_type="anymal")) - # Fix root link by making it kinematic - articulation_cfg.spawn.fix_root_link = True - articulation, translations = generate_articulation(articulation_cfg, num_articulations, device=device) +## +# Kernels, without a simulation. +## - # 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) - - # -- link names (check within articulation ordering is correct) - prim_path_body_names = [path.split("/")[-1] for path in articulation.root_view.link_paths[0]] - assert prim_path_body_names == articulation.body_names - - # 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", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("add_ground_plane", [True]) -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 - - Args: - sim: The simulation fixture - num_articulations: Number of articulations to test - """ - articulation_cfg = clone(generate_articulation_cfg(articulation_type="panda")) - # Unfix root link by making it non-kinematic - articulation_cfg.spawn.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) - - # -- link names (check within articulation ordering is correct) - prim_path_body_names = [path.split("/")[-1] for path in articulation.root_view.link_paths[0]] - assert prim_path_body_names == articulation.body_names - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("add_ground_plane", [True]) -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 = clone(generate_articulation_cfg(articulation_type="panda")) - 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", test_devices(DeviceScope.CUDA)) -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 = replace(FRANKA_PANDA_CFG, 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", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("add_ground_plane", [True]) -def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane, caplog): - """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] +def _selector(values: list[int], dtype: type) -> wp.array: + """Create a CPU Warp selector with the requested integer width.""" + return wp.array(values, dtype=dtype, device="cpu") + + +@pytest.mark.parametrize("env_dtype", [wp.int32, wp.int64]) +@pytest.mark.parametrize("joint_dtype", [wp.int32, wp.int64]) +def test_write_joint_state_data_scatters_nonidentity_selectors(env_dtype: type, joint_dtype: type) -> None: + """Scatter compact joint state for every supported selector-width combination.""" + pos_data = wp.array(np.asarray([[11.0, 12.0], [21.0, 22.0]], dtype=np.float32), device="cpu") + vel_data = wp.array(np.asarray([[111.0, 112.0], [121.0, 122.0]], dtype=np.float32), device="cpu") + env_ids = _selector([1, 0], env_dtype) + joint_ids = _selector([2, 0], joint_dtype) + joint_pos = wp.full((2, 3), value=-1.0, dtype=wp.float32, device="cpu") + joint_vel = wp.full((2, 3), value=-1.0, dtype=wp.float32, device="cpu") + prev_joint_vel = wp.full((2, 3), value=-1.0, dtype=wp.float32, device="cpu") + joint_acc = wp.full((2, 3), value=-1.0, dtype=wp.float32, device="cpu") + kernel = write_joint_state_data + if env_dtype != wp.int32 or joint_dtype != wp.int32: + kernel = write_joint_state_data_kernel(env_ids, joint_ids) + + wp.launch( + kernel, + dim=(2, 2), + inputs=[pos_data, vel_data, env_ids, joint_ids, False], + outputs=[joint_pos, joint_vel, prev_joint_vel, joint_acc], + device="cpu", ) - assert torch.all(within_bounds) - - # Logging level changes must not change clamping or reuse an old violation count. - limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) - limits[..., 1] = 0.5 - articulation_logger = type(articulation).__module__ - for level in (logging.WARNING, logging.INFO): - articulation.data.default_joint_pos.torch.fill_(1.0) - caplog.clear() - with caplog.at_level(level, logger=articulation_logger): - articulation.write_joint_position_limit_to_sim_index(limits=limits, warn_limit_violation=False) - assert bool(caplog.records) == (level == logging.INFO) - torch.testing.assert_close(articulation.data.default_joint_pos.torch, torch.full_like(limits[..., 1], 0.5)) - caplog.clear() - with caplog.at_level(logging.INFO, logger=articulation_logger): - articulation.write_joint_position_limit_to_sim_mask(limits=limits, warn_limit_violation=False) - assert not caplog.records - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@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", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("add_ground_plane", [True]) -def test_external_force_on_single_body(sim, num_articulations, device, add_ground_plane): - """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, translations = 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 - root_pose = articulation.data.default_root_pose.torch.clone() - root_pose[:, :3] += translations - 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 - 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) - # the lateral force along the base y-axis must push the standing robot sideways - lateral_displacement = articulation.data.root_pos_w.torch[:, 1] - translations[:, 1] - assert torch.all(lateral_displacement > 0.5), f"lateral displacement {lateral_displacement.tolist()}" - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("add_ground_plane", [True]) -def test_external_force_on_single_body_at_position(sim, num_articulations, device, add_ground_plane): - """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, translations = 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 - - # Now we are ready! - for i in range(5): - # reset root state - root_pose = articulation.data.default_root_pose.torch.clone() - root_pose[:, :3] += translations - - 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) - # the upward force applied 1 m along the base y-axis must roll the robot about its x-axis - roll_rate = articulation.data.root_ang_vel_b.torch[:, 0] - assert torch.all(roll_rate > 0.1), f"roll rate {roll_rate.tolist()} (is_global={is_global})" - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -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", [2]) -@pytest.mark.parametrize("device", test_devices()) -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 - - # 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): - # the response axis depends on the link frames, so check that the articulation rotates - assert torch.linalg.vector_norm(articulation.data.root_ang_vel_w.torch[i]).item() > 0.1 - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -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", [2]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("add_ground_plane", [True]) -@pytest.mark.parametrize(("stiffness", "damping"), [(10.0, 2.0), ({".*": 10.0}, {".*": 2.0})]) -def test_setting_gains_from_cfg(sim, num_articulations, device, add_ground_plane, stiffness, damping): - """Test that gains are loaded from the configuration correctly. - - This test verifies that: - 1. Gains are loaded correctly from a scalar or a per-joint dictionary configuration - 2. The gains match the expected values - 3. The gains are applied correctly to the actuators - - The humanoid has its articulation root on a rigid body below the spawned prim, so this test also checks - that the floating-base root is discovered there. - - Args: - sim: The simulation fixture - num_articulations: Number of articulations to test - """ - articulation_cfg = generate_articulation_cfg(articulation_type="humanoid", stiffness=stiffness, damping=damping) - articulation, _ = generate_articulation( - articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=sim.device - ) - - # Play sim - sim.reset() - - # Floating-base root discovered on a rigid body below the spawned prim - assert articulation.is_initialized - assert not articulation.is_fixed_base - assert articulation.data.joint_pos.torch.shape == (num_articulations, 21) - prim_path_body_names = [path.split("/")[-1] for path in articulation.root_view.link_paths[0]] - assert prim_path_body_names == articulation.body_names - - # 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("device", test_devices()) -@pytest.mark.parametrize("joint_velocity_limit", [1e5, None]) -def test_setting_velocity_limit_writes_to_solver(sim, device, joint_velocity_limit): - """Test that the resolved joint velocity limit reaches the PhysX solver. - - The full limit-resolution matrix (config override vs. USD default, implicit and explicit - actuators, actuator-limit soft fallback) is covered on the Newton backend and at unit - level. This smoke test only verifies the PhysX write path: the configured limit (or the - USD-authored default when unset) lands in the native solver buffers and matches - ``data.joint_vel_limits``. - """ - articulation_cfg = generate_articulation_cfg( - articulation_type="single_joint_implicit", - joint_velocity_limit=joint_velocity_limit, - ) - articulation, _ = generate_articulation( - articulation_cfg=articulation_cfg, - num_articulations=1, - device=device, - ) - # Play sim - sim.reset() - - # read the values set into the simulation - physx_vel_limit = wp.to_torch(articulation.root_view.get_dof_max_velocities()).to(device) - # check data buffer - torch.testing.assert_close(articulation.data.joint_vel_limits.torch, physx_vel_limit) - # the solver clamp comes from joint_velocity_limit when set, otherwise the USD-authored value - if joint_velocity_limit is None: - limit = next( - p.max_joint_velocity for p in articulation_cfg.spawn.joint_drive_props if isinstance(p, PhysxJointCfg) - ) - else: - limit = joint_velocity_limit - expected_velocity_limit = torch.full_like(physx_vel_limit, limit) - torch.testing.assert_close(physx_vel_limit, expected_velocity_limit) - - -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("joint_effort_limit", [1e5, None]) -@pytest.mark.parametrize("gravity_enabled", [False]) -def test_setting_effort_limit_writes_to_solver(sim, device, joint_effort_limit, gravity_enabled): - """Test that the resolved joint effort limit and a commanded effort reach the PhysX solver. - - The full limit-resolution matrix (config override vs. USD default, implicit and explicit - actuators, actuator-limit soft fallback) is covered on the Newton backend and at unit - level. This smoke test verifies the PhysX write path: the configured limit (or the USD-authored - default when unset) lands in the native solver buffers, and a commanded effort produces motion. - """ - articulation_cfg = generate_articulation_cfg( - articulation_type="single_joint_implicit", - joint_effort_limit=joint_effort_limit, - ) - articulation, _ = generate_articulation( - articulation_cfg=articulation_cfg, - num_articulations=1, - device=device, - ) - # Play sim - sim.reset() - - # obtain the physx effort limits - physx_effort_limit = wp.to_torch(articulation.root_view.get_dof_max_forces()).to(device=device) - # check data buffer - torch.testing.assert_close(articulation.data.joint_effort_limits.torch, physx_effort_limit) - # the solver keeps the USD-authored limit unless the user overrides it explicitly - if joint_effort_limit is None: - limit = next( - p.max_force for p in articulation_cfg.spawn.joint_drive_props if isinstance(p, sim_utils.UsdPhysicsDriveCfg) - ) - else: - limit = joint_effort_limit - expected_effort_limit = torch.full_like(physx_effort_limit, limit) - torch.testing.assert_close(physx_effort_limit, expected_effort_limit) - - # Exercise the command path as well as the property readback. The Isaac Sim 6.0 tensor backend - # accepted this write and updated its staging buffer without applying the effort to the joint. - initial_position = articulation.data.default_joint_pos.torch.clone() - articulation.write_joint_state_to_sim_index( - position=initial_position, - velocity=torch.zeros_like(initial_position), - full_data=True, - ) - effort_target = torch.full_like(initial_position, 10.0) - articulation.actuators.target_command.set_position_index(value=initial_position, full_data=True) - articulation.actuators.target_command.set_velocity_index(value=torch.zeros_like(initial_position), full_data=True) - articulation.actuators.target_command.set_effort_index(value=effort_target, full_data=True) - articulation.write_data_to_sim() - sim.step() - articulation.update(sim.cfg.dt) - assert torch.all(articulation.data.joint_vel.torch > 1e-3), ( - "a positive effort target must accelerate the commanded joint" - ) - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_reset(sim, num_articulations, device, monkeypatch): - """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! - actuator = next(iter(articulation.actuators.values())) - actuator_reset = actuator.reset - reset_env_ids = [] - - def record_actuator_reset(env_ids=None): - reset_env_ids.append(env_ids) - actuator_reset(env_ids) - - monkeypatch.setattr(actuator, "reset", record_actuator_reset) - - num_bodies = articulation.num_bodies - composers = (articulation.instantaneous_wrench_composer, articulation.permanent_wrench_composer) - ones = torch.ones((num_articulations, num_bodies, 3), device=device) - articulation.permanent_wrench_composer.set_forces_and_torques_index(forces=ones, torques=ones) - articulation.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones) - - # A partial reset clears only the selected environment - articulation.reset(env_ids=torch.tensor([0], device=device)) - for composer in composers: - assert composer.active - for buffer in (composer.out_force_b.torch, composer.out_torque_b.torch): - assert torch.count_nonzero(buffer[0]) == 0 - assert torch.count_nonzero(buffer[1:]) == buffer[1:].numel() - - # A full reset resets every actuator environment and clears all external forces and torques - reset_env_ids.clear() - articulation.reset() - assert reset_env_ids == [None] - for composer in composers: - assert not composer.active - assert torch.count_nonzero(composer.out_force_b.torch) == 0 - assert torch.count_nonzero(composer.out_torque_b.torch) == 0 - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("add_ground_plane", [True]) -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) - - # The commanded joint must move toward its target. We can't check that it reached the target as the gains - # are not properly tuned - default_joint_pos = articulation.data.default_joint_pos.torch - assert torch.all(articulation.data.joint_pos.torch[:, 3].abs() < default_joint_pos[:, 3].abs() - 0.5) - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("with_offset", [True, False]) -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" - assert articulation.data.joint_pos.torch.shape == (num_articulations, 1) - prim_path_body_names = [path.split("/")[-1] for path in articulation.root_view.link_paths[0]] - assert prim_path_body_names == articulation.body_names - - # 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 = wp.to_torch(articulation.root_view.get_coms()) - 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) - articulation.set_coms_index( - coms=wp.from_torch(com.to(device).contiguous(), dtype=wp.transformf), - env_ids=wp.from_torch(env_idx, dtype=wp.int32), - ) - - # check they are set - torch.testing.assert_close(wp.to_torch(articulation.root_view.get_coms()), com.cpu()) - - 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 - - # the fixed root stays at its default state - default_root_pose = articulation.data.default_root_pose.torch.clone() - default_root_pose[:, :3] += env_pos - torch.testing.assert_close(root_link_pose_w, default_root_pose) - torch.testing.assert_close(root_com_vel_w, articulation.data.default_root_vel.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", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("with_offset", [True, False]) -@pytest.mark.parametrize("state_location", ["com", "link"]) -@pytest.mark.parametrize("gravity_enabled", [False]) -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 = wp.to_torch(articulation.root_view.get_coms()) - new_com = offset - com[:, 0, :3] = new_com.squeeze(-2) - articulation.set_coms_index( - coms=wp.from_torch(com.to(device).contiguous(), dtype=wp.transformf), - env_ids=wp.from_torch(env_idx, dtype=wp.int32), - ) - - # check they are set - torch.testing.assert_close(wp.to_torch(articulation.root_view.get_coms()), 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("device", test_devices(DeviceScope.CUDA)) -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", test_devices(DeviceScope.CUDA)) -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", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("gravity_enabled", [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) - # every non-root body must move with the joint write - assert torch.all((original_body_states[:, 1:] != body_state_w[:, 1:]).any(dim=-1)) - # validate body - link consistency - body_link_vel_w = articulation.data.body_link_vel_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) - - -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -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. Spatial tendon properties reach the simulation for the right environments - 5. Joint state writes accept int64 environment and joint selectors - - 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 - - # Distinct per-environment values so that an environment mix-up cannot pass - env_offset = torch.arange(num_articulations, dtype=torch.float32, device=device).unsqueeze(1) - stiffness, limit_stiffness, damping, offset = (10.0 + env_offset, 20.0 + env_offset, 3.0 + env_offset, env_offset) - articulation.set_spatial_tendon_stiffness_index(stiffness=stiffness) - articulation.set_spatial_tendon_limit_stiffness_index(limit_stiffness=limit_stiffness) - articulation.set_spatial_tendon_damping_index(damping=damping) - articulation.set_spatial_tendon_offset_index(offset=offset) - articulation.write_spatial_tendon_properties_to_sim_index() - - root_view = articulation.root_view - torch.testing.assert_close(_to_device_tensor(root_view.get_spatial_tendon_stiffnesses(), device), stiffness) - torch.testing.assert_close( - _to_device_tensor(root_view.get_spatial_tendon_limit_stiffnesses(), device), limit_stiffness - ) - torch.testing.assert_close(_to_device_tensor(root_view.get_spatial_tendon_dampings(), device), damping) - torch.testing.assert_close(_to_device_tensor(root_view.get_spatial_tendon_offsets(), device), offset) - - # Joint state writes accept int64 selectors (regression) - env_ids = torch.tensor([1, 0], dtype=torch.int64, device=device) - joint_ids = torch.tensor([articulation.num_joints - 1, 0], dtype=torch.int64, device=device) - position = torch.tensor([[0.21, 0.11], [0.22, 0.12]], device=device) - velocity = torch.tensor([[1.21, 1.11], [1.22, 1.12]], device=device) - - expected_position = articulation.data.joint_pos.torch.clone() - expected_velocity = articulation.data.joint_vel.torch.clone() - - articulation.write_joint_state_to_sim_index( - position=position, velocity=velocity, env_ids=env_ids, joint_ids=joint_ids - ) - expected_position[env_ids[:, None], joint_ids[None, :]] = position - expected_velocity[env_ids[:, None], joint_ids[None, :]] = velocity - torch.testing.assert_close(articulation.data.joint_pos.torch, expected_position) - torch.testing.assert_close(articulation.data.joint_vel.torch, expected_velocity) - - -@pytest.mark.parametrize("add_ground_plane", [True]) -@pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_write_joint_frictions_to_sim(sim, num_articulations, device, add_ground_plane): - """Test that joint friction coefficients written through the combined API reach the simulation.""" - 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() - - # 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 = wp.to_torch(articulation.root_view.get_dof_friction_properties()) - 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()) - - -## -# Shape-contract regression tests for the new BaseArticulation accessors. -# Mirror the Newton-side tests so both backends can be diffed against the -# same documented contract. These are PhysX's reference shapes — when the -# Newton-side tests pass with the same expected_shape formulas, the -# cross-backend contract holds. -## - - -@pytest.mark.parametrize("num_articulations", [4]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("articulation_type", ["panda", "anymal"]) -@pytest.mark.parametrize("gravity_enabled", [False]) -@pytest.mark.isaacsim_ci -def test_get_jacobians_link_origin_contract(sim, num_articulations, device, articulation_type, gravity_enabled): - """PhysX reference: ``J · q_dot`` matches ``[body_link_lin_vel_w; body_link_ang_vel_w]``. - - The cross-backend contract on - :attr:`~isaaclab.assets.BaseArticulationData.body_link_jacobian_w` says - the Jacobian's linear rows reference each body's link origin. PhysX's - raw ``_root_view.get_jacobians()`` returns COM-referenced linear rows; - the IsaacLab wrapper applies the COM→origin shift kernel so the contract - holds. This test pins the identity from the PhysX side and parametrizes - on Anymal so the (non-trivial) shift surfaces if it ever regresses. - - Scene gravity is disabled (``gravity_enabled=False``) so the only source - of a J · q_dot ↔ body_*_w mismatch is the reference-point contract (or a - regression). The tolerance ``5e-2`` is loose enough to absorb the small - PhysX state-propagation lag between the Jacobian and the velocity - buffers (~2% on max angular speed) but well below the - COM-vs-link-origin bug magnitude (panda hand COM offset ≈ 3 cm × ω at - typical motion ≈ several rad/s gives a 0.1+ m/s linear-row residual, - 2× the tolerance). - """ - articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device) - sim.reset() - assert articulation.is_initialized - - torch.manual_seed(0) - qdot = torch.randn(num_articulations, articulation.num_joints, device=device) * 0.5 - articulation.write_joint_velocity_to_sim(velocity=qdot) - sim.step() - articulation.update(sim.cfg.dt) - - # body_link_jacobian_w prepends ``num_base_dofs`` floating-base columns; slice past - # them so the joint axis aligns with joint_vel (actuated-only). - J = articulation.data.body_link_jacobian_w.torch[..., articulation.num_base_dofs :] - qdot_view = articulation.data.joint_vel.torch - v_pred = torch.einsum("nbij,nj->nbi", J, qdot_view) - - body_lin_w = articulation.data.body_link_lin_vel_w.torch - body_ang_w = articulation.data.body_link_ang_vel_w.torch - if articulation.is_fixed_base: - body_lin_w = body_lin_w[:, 1:] - body_ang_w = body_ang_w[:, 1:] - - torch.testing.assert_close(v_pred[..., 3:6], body_ang_w, atol=1.5e-1, rtol=5e-2) - torch.testing.assert_close(v_pred[..., 0:3], body_lin_w, atol=1.5e-1, rtol=5e-2) - - # Shape contract: fixed-base Jacobians omit the root body; floating-base ones prepend the base DoFs. - num_dofs = articulation.num_joints + articulation.num_base_dofs - num_jacobian_bodies = articulation.num_bodies - 1 if articulation.is_fixed_base else articulation.num_bodies - J_full = articulation.data.body_link_jacobian_w.torch - assert J_full.shape == (num_articulations, num_jacobian_bodies, 6, num_dofs) - - # The joint-space mass matrix is square over all DoFs, symmetric, and positive definite. - M = articulation.data.mass_matrix.torch - assert M.shape == (num_articulations, num_dofs, num_dofs) - assert (M.diagonal(dim1=-2, dim2=-1) > 1e-6).all() - asym = (M - M.transpose(-1, -2)).abs().max().item() - assert asym < 1e-4, f"|M - M^T|_max = {asym:.3e} — mass matrix is not symmetric" - eye = torch.eye(M.shape[-1], device=M.device, dtype=M.dtype).expand_as(M) - torch.linalg.cholesky(M + 1e-6 * eye) - - -@pytest.mark.parametrize("num_articulations", [1]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("articulation_type", ["panda", "anymal"]) -@pytest.mark.parametrize("gravity_enabled", [False]) -@pytest.mark.isaacsim_ci -def test_dynamics_refresh_after_manual_joint_write(sim, num_articulations, device, articulation_type, gravity_enabled): - """Each Jacobian and mass-matrix read independently reflects a manual joint write without stepping.""" - articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device) - sim.reset() - sim.step() - articulation.update(sim.cfg.dt) - - q_initial = articulation.data.joint_pos.torch.clone() - env_ids = wp.array([0], dtype=wp.int32, device=device) - for property_name in ("body_link_jacobian_w", "body_com_jacobian_w", "mass_matrix"): - articulation.write_joint_position_to_sim_index(position=q_initial, env_ids=env_ids) - before = getattr(articulation.data, property_name).torch.clone() - # A separate write for each getter prevents another getter from refreshing FK on its behalf. - articulation.write_joint_position_to_sim_index(position=q_initial + 0.5, env_ids=env_ids) - after = getattr(articulation.data, property_name).torch.clone() - assert not torch.allclose(before, after, atol=1e-3), f"{property_name} stayed stale after a joint write" - - -@pytest.mark.parametrize("num_articulations", [1]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("articulation_type", ["panda"]) -@pytest.mark.isaacsim_ci -def test_get_gravity_compensation_forces_static_equilibrium(sim, num_articulations, device, articulation_type): - """PhysX accuracy: ``τ_gc`` must hold the manipulator in static equilibrium. - - The contract is the EOM identity ``M(q) q̈ + C(q,q̇) q̇ + g(q) = τ_input``. - Setting ``τ_input = g(q)`` at ``q̇ = 0`` gives ``q̈ = 0`` — the arm should - not move. This pins - :attr:`~isaaclab.assets.BaseArticulationData.gravity_compensation_forces` - in isolation: sign errors, frame errors, and DoF-ordering errors all - surface as joint drift, while a controller-level test would have those - bugs averaged out by PD damping. - - Newton-side variant of the same name lives in - ``isaaclab_newton/test/assets/test_articulation.py`` (backend parity). - """ - base_cfg = generate_articulation_cfg(articulation_type=articulation_type) - # Replace default Franka actuators with a passthrough implicit actuator - # (stiffness = 0, damping = 0). With both gains zero the effort target - # we set IS the joint torque applied — no PD spring-damper masks the - # gravity-comp signal. Default Franka cfg has stiffness=80 / damping=4 - # which would absorb gravity through PD bias and hide accessor bugs. - cfg = replace(base_cfg, actuators={"all": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=0.0, damping=0.0)}) - # FRANKA_PANDA_CFG has rigid_props.disable_gravity=False already, but be - # defensive — gravity must be ON for τ_gc to have anything to cancel. - cfg = replace(cfg, spawn=replace(cfg.spawn, rigid_props=replace(cfg.spawn.rigid_props, disable_gravity=False))) - - articulation, _ = generate_articulation(cfg, num_articulations, device=device) - sim.reset() - assert articulation.is_initialized - - # Force a clean static state: default joint positions, zero velocities. - # ``sim.reset`` may leave residual ``q_dot`` from solver settling under - # gravity, so we pin it explicitly here. - default_q = articulation.data.default_joint_pos.torch.clone() - default_qd = torch.zeros_like(default_q) - articulation.write_joint_state_to_sim(default_q, default_qd) - articulation.update(sim.cfg.dt) - - # Default joint pose from FRANKA_PANDA_CFG bends the elbow - # (joint2=-0.569, joint4=-2.81, joint6=3.04) so several links carry a - # gravity load — τ_gc is non-trivial in this configuration. A natural- - # hang pose (all zeros) would produce near-zero τ_gc and make this - # test uninformative. - init_q = articulation.data.joint_pos.torch.clone() - - # Step 100 times applying only τ_gc as joint efforts. - for _ in range(100): - # ``gravity_compensation_forces`` shape is ``(N, num_joints + num_base_dofs)`` - # — leading ``num_base_dofs`` floating-base entries (0 on fixed-base) followed - # by the actuated-joint entries. Slice past the floating-base entries so the - # remaining tensor aligns with ``set_joint_effort_target`` (actuated only). - tau_gc = articulation.data.gravity_compensation_forces.torch[:, articulation.num_base_dofs :] - articulation.set_joint_effort_target(tau_gc) - articulation.write_data_to_sim() - sim.step() - articulation.update(sim.cfg.dt) - - final_q = articulation.data.joint_pos.torch - drift = (final_q - init_q).abs().max() - # Tight bound: 5e-3 rad ≈ 0.3°. Numerical integration over 100 steps will - # accumulate some floor (sub-millirad on Franka), but a sign or frame bug - # in τ_gc produces drift of at least a degree per step on bent-elbow - # poses. This bound separates "correct" from "broken" cleanly. - assert drift < 5e-3, ( - f"max joint drift {drift:.5f} rad after 100 gravity-comp-only steps —" - " τ_gc did not hold static equilibrium. Check sign, DoF ordering, and" - " whether gravity_compensation_forces returns g(q) (positive) or" - " its negation." - ) - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("articulation_type", ["panda"]) -@pytest.mark.parametrize("gravity_enabled", [False]) -@pytest.mark.isaacsim_ci -def test_franka_ik_tracking_accuracy(sim, device, articulation_type, gravity_enabled): - """PhysX-side IK convergence sentinel — backend parity with the Newton test. - - Mirrors :func:`isaaclab_newton.test.assets.test_articulation.test_franka_ik_tracking_accuracy` - so both backends are pinned by the same IK trajectory. With the - robot teleported to its configured init_state home pose and scene - gravity off, PhysX's IK converges to ~mm precision on this 5 cm - Cartesian step. A bridge regression (wrong J shape, wrong DoF - ordering) would push the steady-state error well past the - threshold. - """ - robot, ee_frame_idx, ee_jacobi_idx, arm_joint_ids = _setup_franka_at_home_pose(sim) - - sim.step() - robot.update(sim.cfg.dt) - target_pose_b = _build_relative_pose_target(robot, ee_frame_idx, (0.05, 0.0, 0.0), device) - - ik = DifferentialIKController( - DifferentialIKControllerCfg(command_type="pose", use_relative_mode=False, ik_method="dls"), - num_envs=1, - device=device, - ) - ik.set_command(target_pose_b) - - pos_history: list[float] = [] - rot_history: list[float] = [] - for _ in range(800): - jacobian = _compute_jacobian_root_frame(robot, ee_jacobi_idx, arm_joint_ids) - ee_pos_b, ee_quat_b, _ = _compute_ee_pose_root(robot, ee_frame_idx) - joint_pos = robot.data.joint_pos.torch[:, arm_joint_ids] - - joint_pos_des = ik.compute(ee_pos_b, ee_quat_b, jacobian, joint_pos) - - robot.set_joint_position_target(joint_pos_des, joint_ids=arm_joint_ids) - robot.write_data_to_sim() - sim.step() - robot.update(sim.cfg.dt) - - pos_error, rot_error = compute_pose_error(ee_pos_b, ee_quat_b, target_pose_b[:, 0:3], target_pose_b[:, 3:7]) - pos_history.append(pos_error.norm(dim=-1).max().item()) - rot_history.append(rot_error.norm(dim=-1).max().item()) - - pos_min, pos_mean = _summarize_history(pos_history) - rot_min, rot_mean = _summarize_history(rot_history) - - print(f"IK_METRIC pos_min={pos_min:.5f} pos_mean={pos_mean:.5f} rot_min={rot_min:.5f} rot_mean={rot_mean:.5f}") - - # Assert on tail mean (not min) so an oscillating envelope can't - # squeeze through. Threshold matched to the Newton-side test - # (5 mm / 0.05 rad). - assert pos_mean < 5e-3, f"IK pos_mean {pos_mean:.5f} > 5 mm — bridge regression?" - assert rot_mean < 5e-2, f"IK rot_mean {rot_mean:.5f} > 0.05 rad — bridge regression?" - -if __name__ == "__main__": - pytest.main([__file__, "-v", "--maxfail=1"]) + expected_position = np.asarray([[22.0, -1.0, 21.0], [12.0, -1.0, 11.0]], dtype=np.float32) + expected_velocity = np.asarray([[122.0, -1.0, 121.0], [112.0, -1.0, 111.0]], dtype=np.float32) + expected_acceleration = np.asarray([[0.0, -1.0, 0.0], [0.0, -1.0, 0.0]], dtype=np.float32) + np.testing.assert_array_equal(joint_pos.numpy(), expected_position) + np.testing.assert_array_equal(joint_vel.numpy(), expected_velocity) + np.testing.assert_array_equal(prev_joint_vel.numpy(), expected_velocity) + np.testing.assert_array_equal(joint_acc.numpy(), expected_acceleration) diff --git a/source/isaaclab_physx/test/assets/test_articulation_kernels.py b/source/isaaclab_physx/test/assets/test_articulation_kernels.py deleted file mode 100644 index 93d2368b68ba..000000000000 --- a/source/isaaclab_physx/test/assets/test_articulation_kernels.py +++ /dev/null @@ -1,160 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Tests for PhysX articulation Warp kernels.""" - -import sys -import warnings -from types import ModuleType, SimpleNamespace -from unittest.mock import patch - -import numpy as np -import pytest -import warp as wp -from isaaclab_physx.assets.articulation.kernels import ( - write_joint_state_data, - write_joint_state_data_kernel, -) - - -def _selector(values: list[int], dtype: type) -> wp.array: - """Create a CPU Warp selector with the requested integer width.""" - return wp.array(values, dtype=dtype, device="cpu") - - -def _articulation_class(): - """Import PhysX Articulation while suppressing only unavailable lazy Kit exports.""" - cloner = ModuleType("isaaclab_physx.cloner") - cloner.queue_physx_replication = lambda cfg: None - physics = ModuleType("isaaclab_physx.physics") - physics.PhysxManager = type("PhysxManager", (), {}) - with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", - message=( - "^'RigidBodyMaterialCfg' is deprecated and will be removed in 5[.]0[.] Use " - "'isaaclab_physx[.]sim[.]spawners[.]materials[.]PhysxRigidBodyMaterialCfg' for PhysX properties, " - "or 'isaaclab[.]sim[.]spawners[.]materials[.]RigidBodyMaterialBaseCfg' for solver-common properties " - "only[.]$" - ), - category=DeprecationWarning, - ) - warnings.filterwarnings( - "ignore", - message=( - "^`torch[.]jit[.]script` is deprecated[.] Please switch to `torch[.]compile` or `torch[.]export`[.]$" - ), - category=DeprecationWarning, - ) - with patch.dict(sys.modules, {"isaaclab_physx.cloner": cloner, "isaaclab_physx.physics": physics}): - from isaaclab_physx.assets.articulation.articulation import Articulation - - return Articulation - - -def _model_writer_articulation(Articulation) -> tuple[SimpleNamespace, SimpleNamespace]: - """Create a minimal receiver for PhysX model-property writer methods.""" - ArticulationData = Articulation.__init__.__globals__["ArticulationData"] - env_ids = _selector([0], wp.int32) - joint_ids = _selector([0], wp.int32) - body_ids = _selector([0], wp.int32) - data = SimpleNamespace( - _sim_timestamp=3.0, - _joint_armature=wp.zeros((1, 1), dtype=wp.float32, device="cpu"), - _joint_armature_backend=None, - _body_mass=SimpleNamespace(data=wp.zeros((1, 1), dtype=wp.float32, device="cpu"), timestamp=3.0), - _body_mass_backend=None, - _body_inertia=SimpleNamespace(data=wp.zeros((1, 1, 9), dtype=wp.float32, device="cpu"), timestamp=3.0), - _body_inertia_backend=None, - _body_com_jacobian_w=SimpleNamespace(timestamp=3.0), - _mass_matrix=SimpleNamespace(timestamp=3.0), - _gravity_compensation_forces=SimpleNamespace(timestamp=3.0), - has_body_ordering=False, - ) - data._reset_dynamics = lambda **kwargs: ArticulationData._reset_dynamics(data, **kwargs) - articulation = SimpleNamespace( - data=data, - device="cpu", - num_instances=1, - num_joints=1, - num_bodies=1, - root_view=SimpleNamespace( - set_dof_armatures=lambda *args, **kwargs: None, - set_masses=lambda *args, **kwargs: None, - set_inertias=lambda *args, **kwargs: None, - ), - _resolve_env_ids=lambda ids: env_ids, - _resolve_joint_ids=lambda ids: joint_ids, - _resolve_body_ids=lambda ids: body_ids, - _sim_env_ids_view=lambda count: env_ids, - _get_cpu_env_ids=lambda ids, sim_ids=None: env_ids, - _get_backend_ordered_joint_buffer=lambda user, backend: user, - _body_user_to_backend_map=lambda: body_ids, - assert_shape_and_dtype=lambda *args, **kwargs: None, - ) - return articulation, data - - -@pytest.mark.parametrize( - ("writer_name", "kwargs", "invalidates_gravity"), - [ - ("write_joint_armature_to_sim_index", {"armature": 1.0}, False), - ("set_masses_index", {"masses": wp.ones((1, 1), dtype=wp.float32, device="cpu")}, True), - ("set_inertias_index", {"inertias": wp.ones((1, 1, 9), dtype=wp.float32, device="cpu")}, False), - ], -) -def test_model_property_writers_invalidate_same_timestamp_dynamics( - writer_name: str, kwargs: dict, invalidates_gravity: bool -) -> None: - """Invalidate every computed-dynamics cache affected by a PhysX model write.""" - Articulation = _articulation_class() - articulation, data = _model_writer_articulation(Articulation) - method = getattr(Articulation, writer_name) - globals_ = method.__globals__ - - with ( - patch.object(globals_["wp"], "launch"), - patch.object(globals_["ordering_kernels"], "write_float_user_to_backend_with_indices_and_sim_ids"), - patch.object(globals_["ordering_kernels"], "write_3d_user_to_backend_with_indices_and_sim_ids"), - ): - method(articulation, **kwargs) - - assert data._mass_matrix.timestamp == -1.0 - expected_gravity_timestamp = -1.0 if invalidates_gravity else 3.0 - assert data._gravity_compensation_forces.timestamp == expected_gravity_timestamp - assert data._body_com_jacobian_w.timestamp == 3.0 - - -@pytest.mark.parametrize("env_dtype", [wp.int32, wp.int64]) -@pytest.mark.parametrize("joint_dtype", [wp.int32, wp.int64]) -def test_write_joint_state_data_scatters_nonidentity_selectors(env_dtype: type, joint_dtype: type) -> None: - """Scatter compact joint state for every supported selector-width combination.""" - pos_data = wp.array(np.asarray([[11.0, 12.0], [21.0, 22.0]], dtype=np.float32), device="cpu") - vel_data = wp.array(np.asarray([[111.0, 112.0], [121.0, 122.0]], dtype=np.float32), device="cpu") - env_ids = _selector([1, 0], env_dtype) - joint_ids = _selector([2, 0], joint_dtype) - joint_pos = wp.full((2, 3), value=-1.0, dtype=wp.float32, device="cpu") - joint_vel = wp.full((2, 3), value=-1.0, dtype=wp.float32, device="cpu") - prev_joint_vel = wp.full((2, 3), value=-1.0, dtype=wp.float32, device="cpu") - joint_acc = wp.full((2, 3), value=-1.0, dtype=wp.float32, device="cpu") - kernel = write_joint_state_data - if env_dtype != wp.int32 or joint_dtype != wp.int32: - kernel = write_joint_state_data_kernel(env_ids, joint_ids) - - wp.launch( - kernel, - dim=(2, 2), - inputs=[pos_data, vel_data, env_ids, joint_ids, False], - outputs=[joint_pos, joint_vel, prev_joint_vel, joint_acc], - device="cpu", - ) - - expected_position = np.asarray([[22.0, -1.0, 21.0], [12.0, -1.0, 11.0]], dtype=np.float32) - expected_velocity = np.asarray([[122.0, -1.0, 121.0], [112.0, -1.0, 111.0]], dtype=np.float32) - expected_acceleration = np.asarray([[0.0, -1.0, 0.0], [0.0, -1.0, 0.0]], dtype=np.float32) - np.testing.assert_array_equal(joint_pos.numpy(), expected_position) - np.testing.assert_array_equal(joint_vel.numpy(), expected_velocity) - np.testing.assert_array_equal(prev_joint_vel.numpy(), expected_velocity) - np.testing.assert_array_equal(joint_acc.numpy(), expected_acceleration) diff --git a/source/isaaclab_physx/test/assets/test_deformable_object.py b/source/isaaclab_physx/test/assets/test_deformable_object.py index bde56f791c8f..cfd389fa99ef 100644 --- a/source/isaaclab_physx/test/assets/test_deformable_object.py +++ b/source/isaaclab_physx/test/assets/test_deformable_object.py @@ -3,337 +3,268 @@ # # SPDX-License-Identifier: BSD-3-Clause +"""Real PhysX deformable-object coverage on locally authored surface and volume meshes. + +The meshes are authored directly with their simulation topology, so no mesher or collision cooking runs. Each +deformable holds two environments so that partial writes can target environment 1 and prove that environment 0 +is preserved in the real PhysX state. The CPU failure test owns its simulation context and is defined first: +pytest runs it before the composite scene is created, and a new simulation context would replace that stage. +""" + # ignore private usage of variables warning # pyright: reportPrivateUsage=none -from isaaclab.test.utils import launch_test_simulation +from isaaclab.test.utils import DeviceScope, launch_test_simulation, test_devices launch_test_simulation() import sys +from collections.abc import Iterator +from dataclasses import dataclass import pytest import torch import warp as wp -from flaky import flaky from isaaclab_physx.assets import DeformableObject -from isaaclab_physx.sim import ( - PhysxDeformableBodyMaterialCfg, - PhysxDeformableBodyPropertiesCfg, - PhysxSurfaceDeformableBodyMaterialCfg, -) +from isaaclab_physx.sim import PhysxDeformableBodyMaterialCfg, PhysxSurfaceDeformableBodyMaterialCfg -import carb +from pxr import Gf, Sdf, Usd, UsdGeom, UsdShade import isaaclab.sim as sim_utils -import isaaclab.utils.math as math_utils from isaaclab.assets import DeformableObjectCfg from isaaclab.sim import build_simulation_context +from isaaclab.utils.math import quat_apply -# Temporarily disabled: this suite intermittently aborts with SIGABRT on CI. -# Re-enable once the underlying crash is fixed. -pytestmark = pytest.mark.skip(reason="Temporarily disabled due to intermittent crash on CI.") +_NUM_ENVS = 2 +_VOLUME_POINTS = [(0.0, 0.0, 0.0), (0.2, 0.0, 0.0), (0.0, 0.2, 0.0), (0.0, 0.0, 0.2), (0.2, 0.2, 0.2)] +_VOLUME_TETS = [(0, 1, 2, 3), (1, 2, 3, 4)] +_VOLUME_FACES = [(0, 2, 1), (0, 1, 3), (0, 3, 2), (1, 2, 4), (2, 3, 4), (3, 1, 4)] +_SURFACE_POINTS = [(0.0, 0.0, 0.0), (0.2, 0.0, 0.0), (0.2, 0.2, 0.0), (0.0, 0.2, 0.0)] +_SURFACE_TRIANGLES = [(0, 1, 2), (0, 2, 3)] -def generate_cubes_scene( - num_cubes: int = 1, - height: float = 1.0, - initial_rot: tuple[float, ...] = (0.0, 0.0, 0.0, 1.0), - has_api: bool = True, - material_path: str | None = "material", - kinematic_enabled: bool = False, - deformable_type: str = "volume", - device: str = "cuda:0", -) -> DeformableObject: - """Generate a scene with the provided number of cubes. +def _author_deformable_body( + body_prim: Usd.Prim, + points: list[Gf.Vec3f], + sim_api: str, + material_cfg: PhysxDeformableBodyMaterialCfg | PhysxSurfaceDeformableBodyMaterialCfg, +) -> None: + """Author the deformable body schemas, rest shape, and a bound physics material on a mesh prim. Args: - num_cubes: Number of cubes to generate. - height: Height of the cubes. Default is 1.0. - initial_rot: Initial rotation of the cubes (xyzw format). Default is (0.0, 0.0, 0.0, 1.0). - has_api: Whether the cubes have a deformable body API on them. - material_path: Path to the material file. If None, no material is added. Default is "material", - which is path relative to the spawned object prim path. - kinematic_enabled: Whether the cubes are kinematic. - deformable_type: The type of deformable body to spawn. Supported values are "volume" and "surface". - device: Device to use for the simulation. - - Returns: - The deformable object representing the cubes. - + body_prim: Mesh or tetrahedral mesh prim that becomes the deformable body. + points: Rest positions of the simulation nodes in the prim frame [m]. + sim_api: Surface or volume deformable simulation API schema to apply. + material_cfg: Physics material to create below the prim and bind to it. """ - origins = torch.tensor([(i * 1.0, 0, height) for i in range(num_cubes)]).to(device) - # Create Top-level Xforms, one for each cube - for i, origin in enumerate(origins): - sim_utils.create_prim(f"/World/Table_{i}", "Xform", translation=origin) - - # Resolve spawn configuration - if has_api: - spawn_cfg = sim_utils.MeshCuboidCfg( - size=(0.2, 0.2, 0.2), - deformable_props=PhysxDeformableBodyPropertiesCfg(kinematic_enabled=kinematic_enabled), - ) - # Add physics material if provided - if material_path is not None: - if deformable_type == "surface": - spawn_cfg.physics_material = PhysxSurfaceDeformableBodyMaterialCfg() - else: - spawn_cfg.physics_material = PhysxDeformableBodyMaterialCfg() - spawn_cfg.physics_material_path = material_path - else: - spawn_cfg.physics_material = None - else: - # since no deformable body properties defined, this is just a static collider - spawn_cfg = sim_utils.MeshCuboidCfg( - size=(0.2, 0.2, 0.2), - collision_props=sim_utils.CollisionPropertiesCfg(), - ) - # Create deformable object - cube_object_cfg = DeformableObjectCfg( - prim_path="/World/Table_[^/]*/Object", - spawn=spawn_cfg, - init_state=DeformableObjectCfg.InitialStateCfg(pos=(0.0, 0.0, height), rot=initial_rot), + applied_schemas = Sdf.TokenListOp() + applied_schemas.explicitItems = [ + "OmniPhysicsDeformableBodyAPI", + sim_api, + "OmniPhysicsDeformablePoseAPI:default", + "PhysicsCollisionAPI", + ] + body_prim.SetMetadata("apiSchemas", applied_schemas) + body_prim.CreateAttribute("deformablePose:default:omniphysics:points", Sdf.ValueTypeNames.Point3fArray).Set(points) + body_prim.CreateAttribute("deformablePose:default:omniphysics:purposes", Sdf.ValueTypeNames.TokenArray).Set( + ["bindPose"] ) - cube_object = DeformableObject(cfg=cube_object_cfg) - - return cube_object - - -@pytest.fixture -def sim(): - """Create simulation context.""" - with build_simulation_context(auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - yield sim - - -@pytest.mark.parametrize( - "num_cubes, material_path", - [ - (1, "material"), - (2, None), - (2, "/World/SoftMaterial"), - (2, "material"), - ], -) -def test_initialization(sim, num_cubes, material_path): - """Test initialization for prim with deformable body API at the provided prim path.""" - cube_object = generate_cubes_scene(num_cubes=num_cubes, material_path=material_path) - - # Check that the framework doesn't hold excessive strong references. - # sys.getrefcount() adds 1 for its own argument. The baseline is 2 (local var + - # getrefcount arg) but Omniverse event-bus subscriptions and Python/torch runtime - # internals may legitimately add a few more. We use a threshold to catch real leaks. - assert sys.getrefcount(cube_object) < 10 - - # Play sim - sim.reset() - - # Check if object is initialized - assert cube_object.is_initialized - - # Check correct number of cubes - assert cube_object.num_instances == num_cubes - assert cube_object.root_view.count == num_cubes - - # Check correct number of materials in the view - if material_path: - if material_path.startswith("/"): - assert cube_object.material_physx_view.count == 1 - else: - assert cube_object.material_physx_view.count == num_cubes - else: - assert cube_object.material_physx_view is None - - # Check buffers that exist and have correct shapes - # nodal_state_w is (N, V) vec6f -> wp.to_torch gives (N, V, 6) - assert cube_object.data.nodal_state_w.torch.shape == (num_cubes, cube_object.max_sim_vertices_per_body, 6) - # nodal_kinematic_target is (N, V) vec4f -> .torch gives (N, V, 4) - assert cube_object.data.nodal_kinematic_target.torch.shape == ( - num_cubes, - cube_object.max_sim_vertices_per_body, - 4, + body_prim.CreateAttribute("omniphysics:restShapePoints", Sdf.ValueTypeNames.Point3fArray).Set(points) + body_prim.CreateAttribute("velocities", Sdf.ValueTypeNames.Vector3fArray).Set([Gf.Vec3f()] * len(points)) + material_prim = material_cfg.func(f"{body_prim.GetPath()}/material", material_cfg) + UsdShade.MaterialBindingAPI.Apply(body_prim) + UsdShade.MaterialBindingAPI(body_prim).Bind( + UsdShade.Material(material_prim), + bindingStrength=UsdShade.Tokens.weakerThanDescendants, + materialPurpose="physics", ) - # root_pos_w is (N,) vec3f -> wp.to_torch gives (N, 3) - assert cube_object.data.root_pos_w.torch.shape == (num_cubes, 3) - assert cube_object.data.root_vel_w.torch.shape == (num_cubes, 3) -@pytest.mark.isaacsim_ci -def test_initialization_surface_deformable(sim): - """Test initialization of a surface deformable body.""" - num_cubes = 2 - cube_object = generate_cubes_scene(num_cubes=num_cubes, deformable_type="surface") - - # Play sim - sim.reset() - - # Check if object is initialized - assert cube_object.is_initialized - assert cube_object._deformable_type == "surface" - - # Check correct number of instances - assert cube_object.num_instances == num_cubes - assert cube_object.root_view.count == num_cubes - - # Check material view is created - assert cube_object.material_physx_view is not None - assert cube_object.material_physx_view.count == num_cubes - - # Check nodal state buffers have correct shapes - assert cube_object.data.nodal_state_w.torch.shape == (num_cubes, cube_object.max_sim_vertices_per_body, 6) - assert cube_object.data.root_pos_w.torch.shape == (num_cubes, 3) - assert cube_object.data.root_vel_w.torch.shape == (num_cubes, 3) - - # Kinematic targets are not allocated for surface deformables - assert cube_object.data.nodal_kinematic_target is None - - # Writing kinematic targets should raise ValueError - dummy_targets = torch.zeros(num_cubes, cube_object.max_sim_vertices_per_body, 4, device=sim.device) - with pytest.raises(ValueError, match="Kinematic targets can only be set for volume deformable bodies"): - cube_object.write_nodal_kinematic_target_to_sim_index(dummy_targets) +def _spawn_volume_deformables(root: str, y_offset: float) -> DeformableObject: + """Author one five-node, two-tetrahedron volume deformable per environment.""" + stage = sim_utils.get_current_stage() + points = [Gf.Vec3f(*point) for point in _VOLUME_POINTS] + for env_index in range(_NUM_ENVS): + object_path = f"{root}/Env_{env_index}/Object" + sim_utils.create_prim(f"{root}/Env_{env_index}", "Xform", translation=(env_index * 1.0, y_offset, 1.0)) + UsdGeom.Xform.Define(stage, object_path) + tet_mesh = UsdGeom.TetMesh.Define(stage, f"{object_path}/simulation") + tet_mesh.CreatePointsAttr(points) + tet_mesh.CreateTetVertexIndicesAttr([Gf.Vec4i(*tet) for tet in _VOLUME_TETS]) + tet_mesh.CreateSurfaceFaceVertexIndicesAttr([Gf.Vec3i(*face) for face in _VOLUME_FACES]) + body_prim = tet_mesh.GetPrim() + _author_deformable_body( + body_prim, points, "OmniPhysicsVolumeDeformableSimAPI", PhysxDeformableBodyMaterialCfg() + ) + body_prim.CreateAttribute("omniphysics:restTetVtxIndices", Sdf.ValueTypeNames.Int4Array).Set( + [Gf.Vec4i(*tet) for tet in _VOLUME_TETS] + ) + visual = UsdGeom.Mesh.Define(stage, f"{object_path}/visual") + visual.CreatePointsAttr(points) + visual.CreateFaceVertexCountsAttr([3] * len(_VOLUME_FACES)) + visual.CreateFaceVertexIndicesAttr([index for face in _VOLUME_FACES for index in face]) + return DeformableObject(DeformableObjectCfg(prim_path=f"{root}/Env_[^/]*/Object")) + + +def _spawn_surface_deformables(root: str, y_offset: float) -> DeformableObject: + """Author one four-node, two-triangle surface deformable per environment.""" + stage = sim_utils.get_current_stage() + points = [Gf.Vec3f(*point) for point in _SURFACE_POINTS] + for env_index in range(_NUM_ENVS): + sim_utils.create_prim(f"{root}/Env_{env_index}", "Xform", translation=(env_index * 1.0, y_offset, 1.0)) + mesh = UsdGeom.Mesh.Define(stage, f"{root}/Env_{env_index}/Object") + mesh.CreatePointsAttr(points) + mesh.CreateFaceVertexCountsAttr([3] * len(_SURFACE_TRIANGLES)) + mesh.CreateFaceVertexIndicesAttr([index for triangle in _SURFACE_TRIANGLES for index in triangle]) + body_prim = mesh.GetPrim() + _author_deformable_body( + body_prim, + points, + "OmniPhysicsSurfaceDeformableSimAPI", + PhysxSurfaceDeformableBodyMaterialCfg(density=900.0, youngs_modulus=2000.0, surface_thickness=0.02), + ) + body_prim.CreateAttribute("omniphysics:restTriVtxIndices", Sdf.ValueTypeNames.Int3Array).Set( + [Gf.Vec3i(*triangle) for triangle in _SURFACE_TRIANGLES] + ) + return DeformableObject(DeformableObjectCfg(prim_path=f"{root}/Env_[^/]*/Object")) @pytest.mark.isaacsim_ci -def test_initialization_on_device_cpu(): +def test_initialization_on_device_cpu() -> None: """Test that initialization fails with deformable body API on the CPU.""" - with build_simulation_context(device="cpu", auto_add_lighting=True) as sim: + with build_simulation_context(device="cpu", gravity_enabled=False) as sim: sim._app_control_on_stop_handle = None - cube_object = generate_cubes_scene(num_cubes=5, device="cpu") + deformable = _spawn_volume_deformables("/World/Volume", 0.0) # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(cube_object) < 10 + assert sys.getrefcount(deformable) < 10 - # Play sim - with pytest.raises(RuntimeError): + with pytest.raises(RuntimeError, match="Failed to create deformable body at"): sim.reset() -@pytest.mark.isaacsim_ci -def test_set_nodal_state(sim): - """Test setting the state of the deformable object.""" - num_cubes = 2 - cube_object = generate_cubes_scene(num_cubes=num_cubes) - - # Play the simulator - sim.reset() - - for state_type_to_randomize in ["nodal_pos_w", "nodal_vel_w"]: - state_dict = { - "nodal_pos_w": torch.zeros_like(cube_object.data.nodal_pos_w.torch), - "nodal_vel_w": torch.zeros_like(cube_object.data.nodal_vel_w.torch), - } - - for _ in range(5): - cube_object.reset() - - state_dict[state_type_to_randomize] = torch.randn( - num_cubes, cube_object.max_sim_vertices_per_body, 3, device=sim.device - ) - - for _ in range(5): - nodal_state = torch.cat( - [ - state_dict["nodal_pos_w"], - state_dict["nodal_vel_w"], - ], - dim=-1, - ) - cube_object.write_nodal_state_to_sim_index(nodal_state) - - torch.testing.assert_close(cube_object.data.nodal_state_w.torch, nodal_state, rtol=1e-5, atol=1e-5) - - sim.step() - cube_object.update(sim.cfg.dt) - - -@pytest.mark.parametrize( - "num_cubes, randomize_pos, randomize_rot", - [ - (1, False, False), - (1, True, False), - (1, False, True), - (2, True, True), - ], -) -@flaky(max_runs=3, min_passes=1) -@pytest.mark.isaacsim_ci -def test_set_nodal_state_with_applied_transform(num_cubes, randomize_pos, randomize_rot): - """Test setting the state of the deformable object with applied transform.""" - carb_settings_iface = carb.settings.get_settings() - carb_settings_iface.set_bool("/physics/cooking/ujitsoCollisionCooking", False) - - # Create simulation context with gravity disabled (no fixture needed) - with build_simulation_context(auto_add_lighting=True, gravity_enabled=False) as sim: - sim._app_control_on_stop_handle = None - cube_object = generate_cubes_scene(num_cubes=num_cubes) - sim.reset() - - for _ in range(5): - nodal_state = cube_object.data.default_nodal_state_w.torch.clone() - mean_nodal_pos_default = nodal_state[..., :3].mean(dim=1) - - if randomize_pos: - pos_w = 0.5 * torch.rand(cube_object.num_instances, 3, device=sim.device) - pos_w[:, 2] += 0.5 - else: - pos_w = None - if randomize_rot: - quat_w = math_utils.random_orientation(cube_object.num_instances, device=sim.device) - else: - quat_w = None +@dataclass +class _DeformableScene: + """Volume and surface deformables that share one real PhysX lifecycle.""" - nodal_state[..., :3] = cube_object.transform_nodal_pos(nodal_state[..., :3], pos_w, quat_w) - mean_nodal_pos_init = nodal_state[..., :3].mean(dim=1) + sim: sim_utils.SimulationContext + device: str + volume: DeformableObject + surface: DeformableObject + refcounts: dict[str, int] - if pos_w is None: - torch.testing.assert_close(mean_nodal_pos_init, mean_nodal_pos_default, rtol=1e-5, atol=1e-5) - else: - torch.testing.assert_close(mean_nodal_pos_init, mean_nodal_pos_default + pos_w, rtol=1e-5, atol=1e-5) + def step(self, num_steps: int = 1) -> None: + """Write, step, and update both deformables.""" + for _ in range(num_steps): + for deformable in (self.volume, self.surface): + deformable.write_data_to_sim() + self.sim.step() + for deformable in (self.volume, self.surface): + deformable.update(self.sim.cfg.dt) - cube_object.write_nodal_state_to_sim_index(nodal_state) - cube_object.reset() - for _ in range(50): - sim.step() - cube_object.update(sim.cfg.dt) - - torch.testing.assert_close(cube_object.data.root_pos_w.torch, mean_nodal_pos_init, rtol=1e-4, atol=1e-4) +@pytest.fixture(scope="module", params=test_devices(DeviceScope.CUDA)) +def deformable_scene(request) -> Iterator[_DeformableScene]: + """Initialize the deformables once for this module; PhysX deformables require CUDA.""" + device = request.param + with build_simulation_context(device=device, gravity_enabled=False) as sim: + sim._app_control_on_stop_handle = None + volume = _spawn_volume_deformables("/World/Volume", 0.0) + surface = _spawn_surface_deformables("/World/Surface", 2.0) + refcounts = {"volume": sys.getrefcount(volume), "surface": sys.getrefcount(surface)} + sim.reset() + yield _DeformableScene(sim=sim, device=device, volume=volume, surface=surface, refcounts=refcounts) @pytest.mark.isaacsim_ci -def test_set_kinematic_targets(sim): - """Test setting kinematic targets for the deformable object.""" - num_cubes = 2 - cube_object = generate_cubes_scene(num_cubes=num_cubes, height=1.0) - - sim.reset() - - nodal_kinematic_targets = wp.to_torch(cube_object.root_view.get_simulation_nodal_kinematic_targets()).clone() - - for _ in range(5): - cube_object.write_nodal_state_to_sim_index(cube_object.data.default_nodal_state_w.torch) +def test_deformable_initialization(deformable_scene: _DeformableScene) -> None: + """Resolve the volume and surface views, their materials, and their nodal buffers.""" + scene = deformable_scene + for name, deformable, num_vertices in ( + ("volume", scene.volume, len(_VOLUME_POINTS)), + ("surface", scene.surface, len(_SURFACE_POINTS)), + ): + # Check that the framework doesn't hold excessive strong references. + assert scene.refcounts[name] < 10 + assert deformable.is_initialized + assert deformable.num_instances == _NUM_ENVS + assert deformable.root_view.count == _NUM_ENVS + # Each environment binds its own material. + assert deformable.material_physx_view is not None + assert deformable.material_physx_view.count == _NUM_ENVS + assert deformable.max_sim_vertices_per_body == num_vertices + assert deformable.data.nodal_state_w.torch.shape == (_NUM_ENVS, num_vertices, 6) + assert deformable.data.root_pos_w.torch.shape == (_NUM_ENVS, 3) + assert deformable.data.root_vel_w.torch.shape == (_NUM_ENVS, 3) + # Only volume deformables carry kinematic targets. + assert scene.volume.data.nodal_kinematic_target.torch.shape == (_NUM_ENVS, len(_VOLUME_POINTS), 4) + assert scene.surface.data.nodal_kinematic_target is None + dummy_targets = torch.zeros(_NUM_ENVS, len(_SURFACE_POINTS), 4, device=scene.device) + with pytest.raises(ValueError, match="Kinematic targets can only be set for volume deformable bodies"): + scene.surface.write_nodal_kinematic_target_to_sim_index(dummy_targets) - default_root_pos = cube_object.data.default_nodal_state_w.torch.mean(dim=1) - cube_object.reset() +@pytest.mark.isaacsim_ci +def test_deformable_nodal_state_writes(deformable_scene: _DeformableScene) -> None: + """Partial nodal state writes reach only the selected environment of each real view.""" + scene = deformable_scene + device = scene.device + env_ids = torch.tensor([1], dtype=torch.int32, device=device) + for deformable in (scene.volume, scene.surface): + initial_state = deformable.data.nodal_state_w.torch.clone() + state = initial_state[1:].clone() + state[:, 0, :3] += torch.tensor([0.01, 0.02, 0.03], device=device) + state[:, 0, 3] = 0.25 + deformable.write_nodal_state_to_sim_index(state, env_ids=env_ids) + expected = torch.cat((initial_state[:1], state)) + torch.testing.assert_close(deformable.data.nodal_state_w.torch, expected) + for getter, component in ( + (deformable.root_view.get_simulation_nodal_positions, slice(0, 3)), + (deformable.root_view.get_simulation_nodal_velocities, slice(3, 6)), + ): + raw = wp.to_torch(getter()).to(device).reshape(expected[..., component].shape) + torch.testing.assert_close(raw, expected[..., component]) + + # A transformed default state keeps its shape and moves its mean to the requested position. + volume = scene.volume + nodal_state = volume.data.default_nodal_state_w.torch.clone() + pos_w = torch.tensor([[0.1, 0.2, 1.3], [0.4, -0.3, 1.1]], device=device) + quat_w = torch.tensor([[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.7071068, 0.7071068]], device=device) + nodal_pos = volume.transform_nodal_pos(nodal_state[..., :3], pos_w, quat_w) + default_offsets = nodal_state[..., :3] - nodal_state[..., :3].mean(dim=1, keepdim=True) + torch.testing.assert_close(nodal_pos.mean(dim=1), nodal_state[..., :3].mean(dim=1) + pos_w) + torch.testing.assert_close( + nodal_pos - nodal_pos.mean(dim=1, keepdim=True), + quat_apply(quat_w[:, None].expand_as(nodal_state[..., :4]), default_offsets), + ) + nodal_state[..., :3] = nodal_pos + nodal_state[..., 3:] = 0.0 + volume.write_nodal_state_to_sim_index(nodal_state) + raw_nodal_pos = wp.to_torch(volume.root_view.get_simulation_nodal_positions()).to(device) + torch.testing.assert_close(raw_nodal_pos.reshape(nodal_pos.shape), nodal_pos) - nodal_kinematic_targets[1:, :, 3] = 1.0 - nodal_kinematic_targets[0, :, 3] = 0.0 - nodal_kinematic_targets[0, :, :3] = cube_object.data.default_nodal_state_w.torch[0, :, :3] - cube_object.write_nodal_kinematic_target_to_sim_index( - nodal_kinematic_targets[0:1], env_ids=torch.tensor([0], device=sim.device) - ) - for _ in range(20): - sim.step() - cube_object.update(sim.cfg.dt) - - torch.testing.assert_close( - cube_object.data.nodal_pos_w.torch[0], - nodal_kinematic_targets[0, :, :3], - rtol=1e-5, - atol=1e-5, - ) - root_pos_w = cube_object.data.root_pos_w.torch - assert torch.all(root_pos_w[1:, 2] < default_root_pos[1:, 2]) +@pytest.mark.isaacsim_ci +def test_volume_kinematic_targets_hold_selected_nodes(deformable_scene: _DeformableScene) -> None: + """Kinematic targets of environment 0 hold its nodes while environment 1 keeps moving freely.""" + scene = deformable_scene + device = scene.device + volume = scene.volume + targets = volume.data.nodal_kinematic_target.torch.clone() + raw_targets_before = wp.to_torch(volume.root_view.get_simulation_nodal_kinematic_targets()).clone() + targets[0, :, :3] = volume.data.nodal_pos_w.torch[0] + torch.tensor([0.01, 0.02, 0.03], device=device) + targets[0, :, 3] = 0.0 + volume.write_nodal_kinematic_target_to_sim_index(targets[:1], env_ids=torch.tensor([0], device=device)) + torch.testing.assert_close(volume.data.nodal_kinematic_target.torch, targets) + raw_targets = wp.to_torch(volume.root_view.get_simulation_nodal_kinematic_targets()).reshape(targets.shape) + torch.testing.assert_close(raw_targets[0].to(device), targets[0]) + torch.testing.assert_close(raw_targets[1], raw_targets_before.reshape(targets.shape)[1]) + + # Environment 1 moves with a uniform nodal velocity; environment 0 stays on its targets. + state = volume.data.nodal_state_w.torch.clone() + state[1, :, 3:] = torch.tensor([0.0, 0.0, 0.5], device=device) + volume.write_nodal_state_to_sim_index(state[1:], env_ids=torch.tensor([1], device=device)) + free_root_pos = volume.data.root_pos_w.torch[1].clone() + scene.step(5) + torch.testing.assert_close(volume.data.nodal_pos_w.torch[0], targets[0, :, :3], rtol=1e-5, atol=1e-5) + assert volume.data.root_pos_w.torch[1, 2] > free_root_pos[2] + 1e-2 diff --git a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py index 266cd318d3a6..2895693cab1d 100644 --- a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py +++ b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py @@ -3,47 +3,48 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""PD actuator equivalence tests on ANYmal-C (floating-base quadruped) — PhysX backend. - -Compares IsaacLab-native actuators against Newton-native actuators (created -from the same Lab configs via USD authoring, stepped via -:class:`PhysxActuatorWrapper`) on the PhysX physics backend. Both paths -must produce identical joint trajectories within tolerance. - -Using ANYmal-C — a 12-DOF quadruped on a floating base — exercises the -full Lab-to-Newton config translation pipeline on a real-world robot. +"""Newton-native actuator seams on the PhysX backend. + +One rollout steps every robot with Isaac Lab actuators; a second rollout steps the same robots, plus a few that +only exist on the Newton path, with Newton-native actuators (authored from the same Lab configs and stepped through +:class:`PhysxActuatorWrapper`). The robots are local two-leg and cartpole fixtures without collision geometry, so +their dynamics do not depend on where they are placed. The translation of each actuator model is backend independent +and owned by the Newton backend suite; these tests keep the PhysX seams: wrapper efforts next to implicit PhysX +drives, stateful actuators with ping-pong CUDA graphs, non-identity joint ordering, per-articulation adapters on +heterogeneous robots, per-environment resets, gain randomization, network actuators, and native-path preparation. """ -from isaaclab.test.utils import launch_test_simulation +from isaaclab.test.utils import DeviceScope, launch_test_simulation, test_devices from isaaclab.utils import replace launch_test_simulation() -import functools import os -import unittest +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path from types import SimpleNamespace import pytest import torch import warp as wp from isaaclab_physx.assets import Articulation +from isaaclab_physx.assets.articulation import actuator_control from isaaclab_physx.assets.articulation.actuator_control import PhysxActuatorControl from isaaclab_physx.physics import PhysxCfg import isaaclab.sim as sim_utils from isaaclab.actuators import IdealPDActuatorCfg, ImplicitActuatorCfg -from isaaclab.actuators.newton import read_group_parameter -from isaaclab.sim import SimulationCfg, build_simulation_context +from isaaclab.actuators.actuator_net_cfg import ActuatorNetLSTMCfg, ActuatorNetMLPCfg +from isaaclab.actuators.newton import NewtonActuatorAdapter, PhysxActuatorWrapper, read_group_parameter +from isaaclab.assets import ArticulationCfg +from isaaclab.sim import SimulationCfg, SimulationContext, build_simulation_context from isaaclab.test.utils.actuator_equivalence import ( CARTPOLE_EXPLICIT_ACTUATORS, - DC_MOTOR_ACTUATORS, DELAYED_PD_ACTUATORS, IDEAL_PD_ACTUATORS, IMPLICIT_ONLY_ACTUATORS, MIXED_WITH_IMPLICIT_ACTUATORS, - ActuatorStateResetBase, - EquivalenceAssertionsMixin, MockEnv, build_dr_term, make_dummy_lstm_checkpoint, @@ -51,39 +52,339 @@ ) from isaaclab.test.utils.articulation_ordering import assert_articulation_ordering_trace_matches -from isaaclab_assets import ANYMAL_C_CFG -from isaaclab_assets.robots.spot import joint_parameter_lookup as SPOT_KNEE_LOOKUP - -# --------------------------------------------------------------------------- -# Constants -# --------------------------------------------------------------------------- - NUM_ENVS = 2 NUM_STEPS = 10 DT = 1.0 / 120.0 -TARGET_OFFSET = 0.1 # [rad] added to initial joint positions -_ANYMAL_C_PHYSX_JOINT_NAMES = ( - "LF_HAA", - "LH_HAA", - "RF_HAA", - "RH_HAA", - "LF_HFE", - "LH_HFE", - "RF_HFE", - "RH_HFE", - "LF_KFE", - "LH_KFE", - "RF_KFE", - "RH_KFE", -) +_DATA = Path(__file__).resolve().parents[3] / "isaaclab_newton" / "test" / "assets" / "data" +"""Local robot fixtures shared with the Newton backend tests.""" +_PERMUTED_JOINT_NAMES = ("RH_KFE", "LF_HAA", "RH_HAA", "LF_HFE", "RH_HFE", "LF_KFE") +"""A public order of the two-leg joints that is not its own inverse, so a swapped index map cannot pass.""" +_LEG_JOINT_POS = {".*HAA": 0.0, ".*HFE": 0.4, ".*KFE": -0.8} +# Absolute and relative tolerance of each trace in the Lab-versus-Newton comparison. +_TRACE_TOLERANCES = { + "joint_pos": (2e-3, 1e-3), + "joint_vel": (1e-2, 1e-2), + "computed_effort": (1e-3, 1e-3), + "applied_effort": (1e-3, 1e-3), +} + +_LAB_ROBOTS = ("delayed", "mixed", "ideal", "cartpole") +"""Robots of the Lab-actuator rollout, which the Newton-actuator rollout reproduces.""" +_NEWTON_ROBOTS = (*_LAB_ROBOTS, "permuted", "implicit", "mlp", "lstm") +"""Robots of the Newton-actuator rollout.""" + + +def _robot_cfgs(names: tuple[str, ...], networks: dict[str, str]) -> dict[str, ArticulationCfg]: + """Return the configurations of the named robots. + Args: + names: Robot names. ``permuted`` is the ``ideal`` robot with a permuted public joint order. + networks: Paths of the ``mlp`` and ``lstm`` actuator checkpoints. + """ + network_limits = {"saturation_effort": 120.0, "actuator_effort_limit": 80.0, "actuator_velocity_limit": 7.5} + pd_legs = IdealPDActuatorCfg( + joint_names_expr=[".*HFE", ".*KFE"], stiffness=40.0, damping=5.0, actuator_effort_limit=80.0 + ) + leg_actuators = { + "delayed": DELAYED_PD_ACTUATORS, + "mixed": MIXED_WITH_IMPLICIT_ACTUATORS, + "ideal": IDEAL_PD_ACTUATORS, + "implicit": IMPLICIT_ONLY_ACTUATORS, + "mlp": { + "mlp_legs": ActuatorNetMLPCfg( + joint_names_expr=[".*HAA"], + network_file=networks["mlp"], + pos_scale=-1.0, + vel_scale=1.0, + torque_scale=1.0, + input_order="pos_vel", + input_idx=[0, 1, 2], + **network_limits, + ), + "pd_legs": pd_legs, + }, + "lstm": { + "lstm_legs": ActuatorNetLSTMCfg( + joint_names_expr=[".*HAA"], network_file=networks["lstm"], **network_limits + ), + "pd_legs": pd_legs, + }, + } + leg_actuators["permuted"] = IDEAL_PD_ACTUATORS + cfgs = {} + for row, name in enumerate(names): + init_state = ArticulationCfg.InitialStateCfg(pos=(0.0, 3.0 * row, 1.0)) + if name == "cartpole": + usd, actuators = "fixed_cartpole.usda", CARTPOLE_EXPLICIT_ACTUATORS + else: + usd, actuators = "floating_two_leg.usda", leg_actuators[name] + init_state = replace(init_state, joint_pos=_LEG_JOINT_POS) + cfgs[name] = ArticulationCfg( + prim_path=f"/World/Env_[^/]*/{name}", + # The fixtures author no drives; give every joint a drive with an effort limit for the implicit groups. + spawn=sim_utils.UsdFileCfg( + usd_path=str(_DATA / usd), + joint_drive_props={"/.*": [sim_utils.UsdPhysicsDriveCfg(max_force=80.0)]}, + joint_drive_props_create_if_missing=True, + ), + init_state=init_state, + actuators=actuators, + joint_ordering=_PERMUTED_JOINT_NAMES if name == "permuted" else None, + ) + return cfgs -def test_prepare_native_actuators_does_not_zero_solver_gains(monkeypatch): - """Leave solver gains untouched until collection construction resolves actuator defaults.""" - from isaaclab_physx.assets.articulation import actuator_control - from isaaclab.actuators.newton import NewtonActuatorAdapter, PhysxActuatorWrapper +def _rollout(sim: SimulationContext, cfgs: dict[str, ArticulationCfg]) -> tuple[dict[str, Articulation], dict]: + """Spawn the robots side by side, command permutation-sensitive targets, and record ``NUM_STEPS`` steps. + + Args: + sim: Simulation context the robots are spawned in. + cfgs: Robot configurations keyed by name. + + Returns: + The articulations and their recorded traces, both keyed by robot name. + """ + for env_index in range(NUM_ENVS): + sim_utils.create_prim(f"/World/Env_{env_index}", "Xform", translation=(3.0 * env_index, 0.0, 0.0)) + articulations = {name: Articulation(cfg) for name, cfg in cfgs.items()} + sim.reset() + + traces = {} + for name, articulation in articulations.items(): + assert articulation.is_initialized, name + joint_names = tuple(articulation.joint_names) + backend_joint_names = tuple(articulation.backend_joint_names) + ordering = articulation.joint_ordering + # Distinct commands per physical joint, so that a joint-order mix-up cannot pass. + scale = torch.tensor([backend_joint_names.index(joint) + 1.0 for joint in joint_names], device=sim.device) + scale = scale.expand(NUM_ENVS, -1) + target_pos = articulation.data.default_joint_pos.torch + 0.01 * scale + articulation.set_joint_position_target_index(target=target_pos) + articulation.set_joint_velocity_target_index(target=0.001 * scale) + articulation.set_joint_effort_target_index(target=0.1 * scale) + traces[name] = { + "joint_names": joint_names, + "backend_joint_names": backend_joint_names, + "adapter_joint_names": joint_names, + "joint_ordering": None + if ordering is None + else { + "user_names": joint_names, + "backend_names": backend_joint_names, + "user_to_backend_indices": ordering.user_to_backend_indices, + "backend_to_user_indices": ordering.backend_to_user_indices, + }, + "user_to_backend": [backend_joint_names.index(joint) for joint in joint_names], + "target_pos": target_pos, + "target_vel": 0.001 * scale, + "effort_target": 0.1 * scale, + "raw_joint_pos": [], + "raw_joint_vel": [], + "adapter_applied_effort": [], + **{field: [] for field in _TRACE_TOLERANCES}, + } + + for _ in range(NUM_STEPS): + for name, articulation in articulations.items(): + trace = traces[name] + # Read the raw PhysX state first: it bypasses the data's joint-state shadow, so it cannot refresh a stale + # shadow before the actuators read it. + for key, read in (("raw_joint_pos", "get_dof_positions"), ("raw_joint_vel", "get_dof_velocities")): + trace[key].append(wp.to_torch(getattr(articulation.root_view, read)())[:, trace["user_to_backend"]]) + articulation.write_data_to_sim() + # Record the state the actuators acted on and their outputs. Reading the public state only after the + # actuators ran keeps a stale shadow observable on the next step. + trace["joint_pos"].append(articulation.data.joint_pos.torch.clone()) + trace["joint_vel"].append(articulation.data.joint_vel.torch.clone()) + trace["computed_effort"].append(articulation.actuators.computed_effort.torch.clone()) + trace["applied_effort"].append(articulation.actuators.applied_effort.torch.clone()) + if articulation._physx_actuator_wrapper is not None: + wrapper_effort = wp.to_torch(articulation._physx_actuator_wrapper.joint_f_2d) + trace["adapter_applied_effort"].append(wrapper_effort.clone()) + sim.step() + for articulation in articulations.values(): + articulation.update(DT) + return articulations, traces + + +@dataclass +class _Rollouts: + """Recorded traces of both rollouts and the live robots of the Newton-actuator rollout.""" + + lab: dict[str, dict] + """Traces of the Lab-actuator rollout.""" + newton: dict[str, dict] + """Traces of the Newton-actuator rollout.""" + robots: dict[str, Articulation] + """Robots of the Newton-actuator rollout, whose simulation stays alive for the tests.""" + + +def _simulation(device: str, use_newton_actuators: bool): + """Return a simulation context with gravity on the PhysX backend.""" + sim_cfg = SimulationCfg(dt=DT, physics=PhysxCfg(), use_newton_actuators=use_newton_actuators) + return build_simulation_context(device=device, gravity_enabled=True, sim_cfg=sim_cfg) + + +@pytest.fixture(scope="module", params=test_devices(DeviceScope.DEFAULT_CUDA)) +def rollouts(request) -> Iterator[_Rollouts]: + """Record the Lab-actuator rollout, then keep the Newton-actuator rollout alive.""" + device = request.param + networks = {"mlp": make_dummy_mlp_checkpoint(), "lstm": make_dummy_lstm_checkpoint()} + try: + with _simulation(device, use_newton_actuators=False) as sim: + sim._app_control_on_stop_handle = None + lab_traces = _rollout(sim, _robot_cfgs(_LAB_ROBOTS, networks))[1] + with _simulation(device, use_newton_actuators=True) as sim: + sim._app_control_on_stop_handle = None + robots, newton_traces = _rollout(sim, _robot_cfgs(_NEWTON_ROBOTS, networks)) + yield _Rollouts(lab_traces, newton_traces, robots) + finally: + for path in networks.values(): + os.unlink(path) + + +@pytest.mark.parametrize("robot", ["delayed", "mixed", "ideal", "cartpole"]) +def test_newton_actuators_match_lab_actuators(rollouts: _Rollouts, robot: str) -> None: + """Newton-native actuators reproduce the Lab trajectories and efforts next to implicit PhysX drives.""" + for field, (atol, rtol) in _TRACE_TOLERANCES.items(): + for step, (lab, newton) in enumerate( + zip(rollouts.lab[robot][field], rollouts.newton[robot][field], strict=True) + ): + torch.testing.assert_close( + newton, + lab, + atol=atol, + rtol=rtol, + msg=lambda detail: f"{robot} {field} diverged at step {step}: {detail}", + ) + + +def test_stateful_actuators_capture_ping_pong_graphs(rollouts: _Rollouts) -> None: + """Graphable stateful Newton actuators capture the two ping-pong CUDA graphs.""" + assert len(rollouts.robots["delayed"]._actuator_control._native_actuator_graphs) == 2 + + +def test_permuted_joint_ordering_matches_identity_ordering(rollouts: _Rollouts) -> None: + """A permuted public joint order commands and observes the same physical joints as the backend order.""" + assert_articulation_ordering_trace_matches( + rollouts.newton["ideal"], rollouts.newton["permuted"], _PERMUTED_JOINT_NAMES + ) + + +def test_permuted_joint_ordering_uses_current_joint_state(rollouts: _Rollouts) -> None: + """Regression: under a permuted joint order, each step's effort follows the IdealPD law on that step's state.""" + trace = rollouts.newton["permuted"] + kp, kd, effort_limit = 40.0, 5.0, 80.0 + for step, (applied, position, velocity) in enumerate( + zip(trace["applied_effort"], trace["raw_joint_pos"], trace["raw_joint_vel"], strict=True) + ): + demand = kp * (trace["target_pos"] - position) + kd * (trace["target_vel"] - velocity) + trace["effort_target"] + expected = demand.clamp(-effort_limit, effort_limit) + torch.testing.assert_close(applied, expected, atol=1e-3, rtol=1e-3, msg=f"stale joint state at step {step}") + + +@pytest.mark.parametrize("robot", ["mlp", "lstm"]) +def test_network_actuators_drive_haa_joints(rollouts: _Rollouts, robot: str) -> None: + """Network actuators produce finite, non-zero HAA efforts that follow the moving joint state.""" + trace = rollouts.newton[robot] + haa_ids = [index for index, name in enumerate(trace["joint_names"]) if name.endswith("HAA")] + assert len(haa_ids) == 2 + for step, (position, effort) in enumerate(zip(trace["joint_pos"], trace["applied_effort"])): + assert torch.isfinite(position).all(), f"non-finite positions at step {step}" + assert torch.all(effort[:, haa_ids] != 0.0), f"zero HAA effort at step {step}" + # A constant output (e.g. only the output-layer bias) would mean the network is not fed the joint state. + assert not torch.allclose(trace["applied_effort"][0][:, haa_ids], trace["applied_effort"][-1][:, haa_ids]) + + +def test_reset_clears_actuator_state_of_selected_envs(rollouts: _Rollouts) -> None: + """Resetting environment 0 clears the delay history of its DoFs only.""" + articulation = rollouts.robots["delayed"] + adapter = articulation.newton_actuator_adapter + delay_states = [ + (actuator, state.delay_state) + for actuator, state in zip(adapter.actuators, adapter._states_a) + if state is not None and getattr(state, "delay_state", None) is not None + ] + assert delay_states + assert all((delay_state.num_pushes.numpy() > 0).all() for _, delay_state in delay_states) + + articulation.reset(env_ids=torch.tensor([0], device=articulation.device)) + for actuator, delay_state in delay_states: + env_of_dof = actuator.indices.numpy() // adapter.num_joints + pushes = delay_state.num_pushes.numpy() + assert (pushes[env_of_dof == 0] == 0).all() + assert (pushes[env_of_dof != 0] > 0).all() + + +def test_gain_randomization_reuses_payload_for_implicit_storage(rollouts: _Rollouts) -> None: + """Randomized implicit gains reach the actuator and the solver as one payload, in the selected env only.""" + robot = rollouts.robots["implicit"] + actuator = robot.actuators["legs"] + stiffness_before, damping_before = actuator.stiffness.clone(), actuator.damping.clone() + env = MockEnv({"robot": robot}, NUM_ENVS, robot.device) + term, asset_cfg = build_dr_term(env, "robot") + env_ids = torch.tensor([0], device=robot.device) + torch.manual_seed(12345) + + term( + env, + env_ids=env_ids, + asset_cfg=asset_cfg, + stiffness_distribution_params=(25.0, 75.0), + damping_distribution_params=(1.0, 9.0), + operation="abs", + distribution="uniform", + ) + + assert torch.unique(actuator.stiffness[env_ids]).numel() > 1 + assert torch.unique(actuator.damping[env_ids]).numel() > 1 + torch.testing.assert_close(actuator.stiffness[env_ids], robot.data.joint_stiffness.torch[env_ids]) + torch.testing.assert_close(actuator.damping[env_ids], robot.data.joint_damping.torch[env_ids]) + torch.testing.assert_close(actuator.stiffness[1:], stiffness_before[1:]) + torch.testing.assert_close(actuator.damping[1:], damping_before[1:]) + + +def test_gain_randomization_reaches_only_the_selected_articulation(rollouts: _Rollouts) -> None: + """Each PhysX articulation owns its adapter, so randomizing the cartpole leaves the two-leg controllers alone.""" + legs, cartpole = rollouts.robots["ideal"], rollouts.robots["cartpole"] + assert legs.newton_actuator_adapter is not None and cartpole.newton_actuator_adapter is not None + assert legs.newton_actuator_adapter is not cartpole.newton_actuator_adapter + groups = {"legs": (legs, "legs"), "cartpole": (cartpole, "all_joints")} + + def read_gains() -> dict[tuple[str, str], torch.Tensor]: + return { + (name, gain): read_group_parameter(articulation.actuators, group, "controller", gain).clone() + for name, (articulation, group) in groups.items() + for gain in ("kp", "kd") + } + + expected = read_gains() + env = MockEnv({"legs": legs, "cartpole": cartpole}, NUM_ENVS, legs.device) + term, asset_cfg = build_dr_term(env, "cartpole") + + term( + env, + env_ids=torch.tensor([0], device=legs.device), + asset_cfg=asset_cfg, + stiffness_distribution_params=(100.0, 100.0), + damping_distribution_params=(5.0, 5.0), + operation="abs", + distribution="uniform", + ) + + expected["cartpole", "kp"][0] = 100.0 + expected["cartpole", "kd"][0] = 5.0 + for key, gains in read_gains().items(): + torch.testing.assert_close(gains, expected[key], msg=str(key)) + +## +# Native-path preparation, without a simulation. +## + + +def test_prepare_native_actuators_does_not_zero_solver_gains(monkeypatch): + """Leave solver gains untouched until collection construction resolves actuator defaults.""" joint_buffer = SimpleNamespace(warp=wp.zeros((1, 1), dtype=wp.float32, device="cpu")) collection = SimpleNamespace( target_command=SimpleNamespace(position=joint_buffer, velocity=joint_buffer, effort=joint_buffer) @@ -117,8 +418,6 @@ def test_prepare_native_actuators_does_not_zero_solver_gains(monkeypatch): def test_prepare_native_actuators_leaves_implicit_only_articulation_on_standard_path(monkeypatch): """Keep implicit-only articulations on the unchanged solver-drive path.""" - from isaaclab_physx.assets.articulation import actuator_control - runtime_prepare_calls = [] runtime = SimpleNamespace( prepare=lambda *args, **kwargs: runtime_prepare_calls.append(True), wrapper=None, adapter=None @@ -140,719 +439,3 @@ def test_prepare_native_actuators_leaves_implicit_only_articulation_on_standard_ assert not control.native_actuator_path_active assert not articulation._has_newton_actuators assert runtime_prepare_calls == [] - - -# --------------------------------------------------------------------------- -# Simulation runner -# --------------------------------------------------------------------------- - - -def _run_simulation( - actuators: dict, - use_newton_actuators: bool, - *, - num_steps: int = NUM_STEPS, - feedforward: float | None = None, - joint_ordering: tuple[str, ...] | None = None, - permutation_sensitive_commands: bool = False, - capture_first_compute: bool = False, -) -> dict: - """Run ANYmal-C on PhysX and return recorded trajectories + telemetry. - - Always records public joint state/telemetry and the Newton adapter outputs. - - Args: - actuators: Actuator configuration replacing ANYmal-C defaults. - use_newton_actuators: Whether to use the Newton actuator fast path. - num_steps: Number of simulation steps to record. - feedforward: Optional constant effort target for every joint. - joint_ordering: Optional explicit public joint-name order. - permutation_sensitive_commands: Whether to command distinct position, velocity, and effort values by - physical joint name. - capture_first_compute: Whether to invoke the first actuator computation inside an outer CUDA capture. - - Returns: - Recorded joint-name metadata, commands, public trajectories and torque telemetry, and adapter effort traces. - """ - sim_cfg = SimulationCfg(dt=DT, physics=PhysxCfg(), use_newton_actuators=use_newton_actuators) - with build_simulation_context( - device="cuda:0", - gravity_enabled=True, - add_ground_plane=True, - sim_cfg=sim_cfg, - ) as sim: - sim._app_control_on_stop_handle = None - for i in range(NUM_ENVS): - sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0)) - art_cfg = replace( - ANYMAL_C_CFG, actuators=actuators, prim_path="/World/Env_[^/]*/Robot", joint_ordering=joint_ordering - ) - articulation = Articulation(art_cfg) - sim.reset() - assert articulation.is_initialized - - joint_names = tuple(articulation.joint_names) - backend_joint_names = tuple(articulation.backend_joint_names) - installed_ordering = articulation.joint_ordering - joint_ordering_state = ( - None - if installed_ordering is None - else { - "user_names": joint_names, - "backend_names": backend_joint_names, - "user_to_backend_indices": installed_ordering.user_to_backend_indices, - "backend_to_user_indices": installed_ordering.backend_to_user_indices, - } - ) - init_pos = wp.to_torch(articulation.data.joint_pos).clone() - if permutation_sensitive_commands: - scale_by_name = {name: index + 1 for index, name in enumerate(backend_joint_names)} - joint_scale = torch.tensor( - [scale_by_name[name] for name in joint_names], - device=articulation.device, - dtype=init_pos.dtype, - ).unsqueeze(0) - joint_scale = joint_scale.expand_as(init_pos) - target_pos = init_pos + 0.01 * joint_scale - target_vel = 0.001 * joint_scale - effort_target = 0.1 * joint_scale - else: - target_pos = init_pos + TARGET_OFFSET - target_vel = torch.zeros_like(init_pos) - effort_target = None if feedforward is None else torch.full_like(init_pos, feedforward) - - articulation.set_joint_position_target_index(target=target_pos) - articulation.set_joint_velocity_target_index(target=target_vel) - if effort_target is not None: - articulation.set_joint_effort_target_index(target=effort_target) - - recorded_pos, recorded_vel = [], [] - recorded_computed_effort, recorded_applied_effort = [], [] - recorded_adapter_applied = [] - if capture_first_compute: - with wp.ScopedCapture(device=articulation.device, force_module_load=True): - articulation.actuators.compute(DT) - for _ in range(num_steps): - articulation.write_data_to_sim() - sim.step() - articulation.update(DT) - recorded_pos.append(wp.to_torch(articulation.data.joint_pos).clone()) - recorded_vel.append(wp.to_torch(articulation.data.joint_vel).clone()) - recorded_computed_effort.append(articulation.actuators.computed_effort.torch.clone()) - recorded_applied_effort.append(articulation.actuators.applied_effort.torch.clone()) - if articulation._physx_actuator_wrapper is not None: - recorded_adapter_applied.append(wp.to_torch(articulation._physx_actuator_wrapper.joint_f_2d).clone()) - native_actuator_graph_count = len(getattr(articulation._actuator_control, "_native_actuator_graphs", ()) or ()) - - return { - "joint_names": joint_names, - "backend_joint_names": backend_joint_names, - "joint_ordering": joint_ordering_state, - "adapter_joint_names": joint_names, - "joint_pos": recorded_pos, - "joint_vel": recorded_vel, - "computed_effort": recorded_computed_effort, - "applied_effort": recorded_applied_effort, - "adapter_applied_effort": recorded_adapter_applied, - "target_pos": target_pos.clone(), - "target_vel": target_vel.clone(), - "effort_target": None if effort_target is None else effort_target.clone(), - "native_actuator_graph_count": native_actuator_graph_count, - } - - -def test_newton_actuator_graph_capture_failure_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) -> None: - class FailingCapture: - def __init__(self, *args, **kwargs): - pass - - def __enter__(self): - raise RuntimeError("capture unavailable") - - def __exit__(self, exc_type, exc_value, traceback): - return False - - monkeypatch.setattr(wp, "ScopedCapture", FailingCapture) - - result = _run_simulation( - DC_MOTOR_ACTUATORS, - use_newton_actuators=True, - num_steps=2, - feedforward=1.0, - ) - - assert result["native_actuator_graph_count"] == 0 - assert len(result["joint_pos"]) == 2 - assert all(torch.isfinite(joint_pos).all() for joint_pos in result["joint_pos"]) - assert all(torch.any(effort != 0.0) for effort in result["applied_effort"]) - assert all(torch.any(effort != 0.0) for effort in result["adapter_applied_effort"]) - - -def test_stateful_newton_actuators_reject_outer_cuda_capture() -> None: - with pytest.raises(RuntimeError, match="stateful Newton actuators cannot run inside an outer CUDA graph capture"): - _run_simulation( - DELAYED_PD_ACTUATORS, - use_newton_actuators=True, - num_steps=0, - capture_first_compute=True, - ) - - -def test_newton_actuator_rollout_matches_reversed_joint_ordering() -> None: - """Match PhysX Newton-actuator traces under reversed public joint ordering.""" - identity_result = _run_simulation( - IDEAL_PD_ACTUATORS, - use_newton_actuators=True, - permutation_sensitive_commands=True, - ) - requested_joint_names = tuple(reversed(identity_result["joint_names"])) - reversed_result = _run_simulation( - IDEAL_PD_ACTUATORS, - use_newton_actuators=True, - joint_ordering=requested_joint_names, - permutation_sensitive_commands=True, - ) - - assert_articulation_ordering_trace_matches(identity_result, reversed_result, requested_joint_names) - - -def _assert_newton_actuator_uses_current_joint_state( - joint_ordering: tuple[str, ...] | None, *, num_steps: int = NUM_STEPS -) -> None: - """Check that ``applied_effort`` always matches the IdealPD formula on *this* step's true state. - - Ground truth is read every step via ``root_view.get_dof_positions()``/``get_dof_velocities()`` -- - the raw PhysX view, bypassing :class:`ArticulationData`'s cached ``joint_pos``/``joint_vel`` shadow - entirely -- so the read itself cannot refresh (and thereby mask staleness in) the shadow under test. - - Args: - joint_ordering: Optional explicit public joint-name order to install on the articulation. - num_steps: Number of simulation steps to check. - """ - kp, kd, effort_limit = 40.0, 5.0, 80.0 - actuators = { - "legs": IdealPDActuatorCfg( - joint_names_expr=[".*HAA", ".*HFE", ".*KFE"], - stiffness=kp, - damping=kd, - actuator_effort_limit=effort_limit, - ), - } - sim_cfg = SimulationCfg(dt=DT, physics=PhysxCfg(), use_newton_actuators=True) - with build_simulation_context( - device="cuda:0", - gravity_enabled=True, - add_ground_plane=True, - sim_cfg=sim_cfg, - ) as sim: - sim._app_control_on_stop_handle = None - for i in range(NUM_ENVS): - sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0)) - art_cfg = replace( - ANYMAL_C_CFG, actuators=actuators, prim_path="/World/Env_[^/]*/Robot", joint_ordering=joint_ordering - ) - articulation = Articulation(art_cfg) - sim.reset() - assert articulation.is_initialized - - ordering = articulation.joint_ordering - if ordering is not None: - user_to_backend = torch.tensor( - ordering.user_to_backend_indices, dtype=torch.long, device=articulation.device - ) - - def to_user_order(raw_backend: wp.array) -> torch.Tensor: - return wp.to_torch(raw_backend).index_select(1, user_to_backend) - - else: - - def to_user_order(raw_backend: wp.array) -> torch.Tensor: - return wp.to_torch(raw_backend) - - init_pos = to_user_order(articulation.root_view.get_dof_positions()).clone() - target_pos = init_pos + TARGET_OFFSET - target_vel = torch.zeros_like(init_pos) - articulation.set_joint_position_target_index(target=target_pos) - articulation.set_joint_velocity_target_index(target=target_vel) - - for step in range(num_steps): - # Ground truth for *this* step, independent of ArticulationData's joint_pos/joint_vel shadow. - true_pos = to_user_order(articulation.root_view.get_dof_positions()).clone() - true_vel = to_user_order(articulation.root_view.get_dof_velocities()).clone() - - articulation.write_data_to_sim() - applied = articulation.actuators.applied_effort.torch.clone() - - expected = torch.clamp(kp * (target_pos - true_pos) - kd * true_vel, -effort_limit, effort_limit) - torch.testing.assert_close( - applied, - expected, - atol=1e-3, - rtol=1e-3, - msg=( - f"applied_effort at step {step} does not match the IdealPD formula evaluated on this" - " step's true PhysX joint state -- the Newton actuator likely used a stale" - " joint_pos/joint_vel shadow" - ), - ) - - sim.step() - articulation.update(DT) - - -def test_newton_actuator_reversed_ordering_uses_current_joint_state() -> None: - """Regression test: non-identity joint ordering must preserve current-state torque evaluation. - - The adapter must resolve joint state in articulation order before evaluating the actuator model, - regardless of the backend view ordering. - """ - reversed_joint_names = tuple(reversed(_ANYMAL_C_PHYSX_JOINT_NAMES)) - _assert_newton_actuator_uses_current_joint_state(reversed_joint_names) - - -# --------------------------------------------------------------------------- -# Base test class -# --------------------------------------------------------------------------- - - -class _EquivalenceTestBase(EquivalenceAssertionsMixin, unittest.TestCase): - """Base for Lab-vs-Newton equivalence tests on the PhysX backend. - - Subclasses set ``actuators`` to the config under test. ``setUpClass`` - runs the simulation with both ``use_newton_actuators=False`` (Lab path) - and ``True`` (Newton via PhysxActuatorWrapper) and stores the results. - The ``test_*_match`` oracles come from :class:`EquivalenceAssertionsMixin`. - """ - - __test__ = False - actuators: dict = {} - - @classmethod - def setUpClass(cls): - cls.lab_result = _run_simulation(cls.actuators, use_newton_actuators=False) - cls.newton_result = _run_simulation(cls.actuators, use_newton_actuators=True) - - -# --------------------------------------------------------------------------- -# Equivalence tests with different actuator types -# --------------------------------------------------------------------------- - - -class TestIdealPDEquivalence(_EquivalenceTestBase): - """IdealPDActuator on all 12 joints: Lab vs Newton (PhysX backend).""" - - __test__ = True - actuators = IDEAL_PD_ACTUATORS - - -class TestDCMotorEquivalence(_EquivalenceTestBase): - """DCMotor actuator on all 12 joints: Lab vs Newton (PhysX backend).""" - - __test__ = True - actuators = DC_MOTOR_ACTUATORS - - -class TestDelayedPDEquivalence(_EquivalenceTestBase): - """DelayedPDActuator on all 12 joints: Lab vs Newton (PhysX). - - Verifies that actuator command delays are correctly authored and - produce matching trajectories on the PhysX backend. - """ - - __test__ = True - actuators = DELAYED_PD_ACTUATORS - - def test_native_graph_count(self): - """Graphable Newton actuators capture the two ping-pong graphs.""" - self.assertEqual(self.newton_result["native_actuator_graph_count"], 2) - - -class TestMixedWithImplicitEquivalence(_EquivalenceTestBase): - """Implicit HAA + IdealPD HFE + DCMotor KFE: Lab vs Newton (PhysX). - - Verifies that implicit actuators (handled by PhysX joint drives) - coexist correctly with explicit Newton actuators via PhysxActuatorWrapper. - """ - - __test__ = True - actuators = MIXED_WITH_IMPLICIT_ACTUATORS - - -# --------------------------------------------------------------------------- -# Heterogeneous multi-articulation (ANYmal floating-base + Cartpole fixed-base) -# --------------------------------------------------------------------------- - - -def _run_anymal_and_cartpole(use_newton_actuators: bool, *, num_steps: int = NUM_STEPS) -> dict: - """Spawn ANYmal-C + Cartpole per env on PhysX (different DOF counts, base types).""" - from isaaclab_assets import CARTPOLE_CFG # noqa: PLC0415 - - sim_cfg = SimulationCfg(dt=DT, physics=PhysxCfg(), use_newton_actuators=use_newton_actuators) - with build_simulation_context( - device="cuda:0", - gravity_enabled=True, - add_ground_plane=True, - sim_cfg=sim_cfg, - ) as sim: - sim._app_control_on_stop_handle = None - - for i in range(NUM_ENVS): - sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 6.0, 0, 0)) - - anymal_cfg = replace(ANYMAL_C_CFG, actuators=IDEAL_PD_ACTUATORS, prim_path="/World/Env_[^/]*/Anymal") - cartpole_cfg = replace( - CARTPOLE_CFG, actuators=CARTPOLE_EXPLICIT_ACTUATORS, prim_path="/World/Env_[^/]*/Cartpole" - ) - cartpole_cfg.init_state = replace(cartpole_cfg.init_state, pos=(0.0, 3.0, 2.0)) - - anymal = Articulation(anymal_cfg) - cartpole = Articulation(cartpole_cfg) - sim.reset() - assert anymal.is_initialized and cartpole.is_initialized - - init_anymal = wp.to_torch(anymal.data.joint_pos).clone() - init_cartpole = wp.to_torch(cartpole.data.joint_pos).clone() - anymal.set_joint_position_target_index(target=init_anymal + TARGET_OFFSET) - anymal.set_joint_velocity_target_index(target=torch.zeros_like(init_anymal)) - cartpole.set_joint_position_target_index(target=init_cartpole + TARGET_OFFSET) - cartpole.set_joint_velocity_target_index(target=torch.zeros_like(init_cartpole)) - - pos_anymal, pos_cartpole = [], [] - for _ in range(num_steps): - anymal.write_data_to_sim() - cartpole.write_data_to_sim() - sim.step() - anymal.update(DT) - cartpole.update(DT) - pos_anymal.append(wp.to_torch(anymal.data.joint_pos).clone()) - pos_cartpole.append(wp.to_torch(cartpole.data.joint_pos).clone()) - - return {"joint_pos_anymal": pos_anymal, "joint_pos_cartpole": pos_cartpole} - - -class TestHeterogeneousMultiArticulationPhysx(unittest.TestCase): - """Two structurally-different articulations (ANYmal floating + Cartpole fixed) on PhysX. - - Each PhysX articulation owns its own :class:`PhysxActuatorWrapper` - and per-art :class:`NewtonActuatorAdapter`. Heterogeneous DOF counts - (12 vs 2) and base types (floating vs fixed) verify the - per-articulation authoring + adapter construction works for varied - structures. Equivalence against the Lab actuator path is the - meaningful end-to-end check. - """ - - @classmethod - def setUpClass(cls): - cls.lab_result = _run_anymal_and_cartpole(use_newton_actuators=False) - cls.newton_result = _run_anymal_and_cartpole(use_newton_actuators=True) - - def test_anymal_matches_lab(self): - for step_i, (lab, newton) in enumerate( - zip(self.lab_result["joint_pos_anymal"], self.newton_result["joint_pos_anymal"]) - ): - torch.testing.assert_close( - newton, - lab, - atol=2e-3, - rtol=1e-3, - msg=f"ANYmal joint_pos diverged from Lab path at step {step_i}", - ) - - def test_cartpole_matches_lab(self): - for step_i, (lab, newton) in enumerate( - zip(self.lab_result["joint_pos_cartpole"], self.newton_result["joint_pos_cartpole"]) - ): - torch.testing.assert_close( - newton, - lab, - atol=2e-3, - rtol=1e-3, - msg=f"Cartpole joint_pos diverged from Lab path at step {step_i}", - ) - - -# --------------------------------------------------------------------------- -# Domain randomization via events.py — PhysX backend -# --------------------------------------------------------------------------- - - -class TestRandomizeActuatorGainsViaEventsPhysx(unittest.TestCase): - """End-to-end DR test for the PhysX backend. - - Drives ``randomize_actuator_gains`` (events.py) and verifies the new - kp/kd values reach the controllers of the articulation's Newton - actuators — exercising the full path: events → the actuator adapter → - write_stiffness/damping → propagation to controllers. The assertions - read the controllers back via the public - ``read_group_parameter``. - - The native-controller tests use degenerate ranges for exact expected values. - The implicit-storage regression instead seeds the generator and uses - non-degenerate ranges to verify one sampled payload reaches every storage. - """ - - def test_implicit_storage_reuses_randomized_payload(self): - """Keep actuator-owned and implicit-solver gains identical after randomization.""" - sim_cfg = SimulationCfg(dt=DT, physics=PhysxCfg(), use_newton_actuators=False) - with build_simulation_context( - device="cuda:0", - gravity_enabled=True, - add_ground_plane=True, - sim_cfg=sim_cfg, - ) as sim: - sim._app_control_on_stop_handle = None - for i in range(NUM_ENVS): - sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0)) - art_cfg = replace(ANYMAL_C_CFG, actuators=IMPLICIT_ONLY_ACTUATORS, prim_path="/World/Env_.*/Robot") - anymal = Articulation(art_cfg) - sim.reset() - - actuator = anymal.actuators["legs"] - stiffness_before = actuator.stiffness.clone() - damping_before = actuator.damping.clone() - env = MockEnv({"robot": anymal}, NUM_ENVS, anymal.device) - term, asset_cfg = build_dr_term(env, "robot") - env_ids = torch.tensor([0], device=anymal.device, dtype=torch.long) - torch.manual_seed(12345) - - term( - env, - env_ids=env_ids, - asset_cfg=asset_cfg, - stiffness_distribution_params=(25.0, 75.0), - damping_distribution_params=(1.0, 9.0), - operation="abs", - distribution="uniform", - ) - - randomized_stiffness = actuator.stiffness[env_ids] - randomized_damping = actuator.damping[env_ids] - self.assertGreater(torch.unique(randomized_stiffness).numel(), 1) - self.assertGreater(torch.unique(randomized_damping).numel(), 1) - torch.testing.assert_close(randomized_stiffness, anymal.data.joint_stiffness.torch[env_ids]) - torch.testing.assert_close(randomized_damping, anymal.data.joint_damping.torch[env_ids]) - torch.testing.assert_close(actuator.stiffness[1:], stiffness_before[1:]) - torch.testing.assert_close(actuator.damping[1:], damping_before[1:]) - - def test_two_articulations(self): - from isaaclab_assets import CARTPOLE_CFG # noqa: PLC0415 - - sim_cfg = SimulationCfg(dt=DT, physics=PhysxCfg(), use_newton_actuators=True) - with build_simulation_context( - device="cuda:0", - gravity_enabled=True, - add_ground_plane=True, - sim_cfg=sim_cfg, - ) as sim: - sim._app_control_on_stop_handle = None - for i in range(NUM_ENVS): - sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 6.0, 0, 0)) - - anymal_cfg = replace(ANYMAL_C_CFG, actuators=IDEAL_PD_ACTUATORS, prim_path="/World/Env_[^/]*/Anymal") - cartpole_cfg = replace( - CARTPOLE_CFG, actuators=CARTPOLE_EXPLICIT_ACTUATORS, prim_path="/World/Env_[^/]*/Cartpole" - ) - cartpole_cfg.init_state = replace(cartpole_cfg.init_state, pos=(0.0, 3.0, 2.0)) - anymal = Articulation(anymal_cfg) - cartpole = Articulation(cartpole_cfg) - sim.reset() - - # On PhysX each articulation owns its own adapter — they are distinct objects. - anymal_adapter = anymal.newton_actuator_adapter - cartpole_adapter = cartpole.newton_actuator_adapter - self.assertIsNotNone(anymal_adapter) - self.assertIsNotNone(cartpole_adapter) - self.assertIsNot(anymal_adapter, cartpole_adapter) - - anymal_read = functools.partial(read_group_parameter, anymal.actuators) - cartpole_read = functools.partial(read_group_parameter, cartpole.actuators) - n_cp = cartpole.num_joints - anymal_kp_before = anymal_read("legs", "controller", "kp").clone() - anymal_kd_before = anymal_read("legs", "controller", "kd").clone() - cp_kp_before = cartpole_read("all_joints", "controller", "kp").clone() - cp_kd_before = cartpole_read("all_joints", "controller", "kd").clone() - - env = MockEnv({"anymal": anymal, "cartpole": cartpole}, NUM_ENVS, anymal.device) - term, asset_cfg = build_dr_term(env, "cartpole") - env_ids = torch.tensor([0], device=anymal.device, dtype=torch.long) - - term( - env, - env_ids=env_ids, - asset_cfg=asset_cfg, - stiffness_distribution_params=(100.0, 100.0), - damping_distribution_params=(5.0, 5.0), - operation="abs", - distribution="uniform", - ) - - cp_kp_after = cartpole_read("all_joints", "controller", "kp") - cp_kd_after = cartpole_read("all_joints", "controller", "kd") - torch.testing.assert_close(cp_kp_after[0], torch.full((n_cp,), 100.0, device=anymal.device)) - torch.testing.assert_close(cp_kd_after[0], torch.full((n_cp,), 5.0, device=anymal.device)) - # Cartpole's other envs are untouched (env_ids=[0] only). - for env_idx in range(1, NUM_ENVS): - torch.testing.assert_close(cp_kp_after[env_idx], cp_kp_before[env_idx]) - torch.testing.assert_close(cp_kd_after[env_idx], cp_kd_before[env_idx]) - - # ANYmal's controllers are fully untouched — DR was scoped to cartpole. - torch.testing.assert_close(anymal_read("legs", "controller", "kp"), anymal_kp_before) - torch.testing.assert_close(anymal_read("legs", "controller", "kd"), anymal_kd_before) - - -# --------------------------------------------------------------------------- -# Per-env reset: actuator state isolation -# --------------------------------------------------------------------------- - - -class TestActuatorStateReset(ActuatorStateResetBase, unittest.TestCase): - """Per-env actuator state reset isolation on the PhysX backend. - - The scenario and assertions live in :class:`ActuatorStateResetBase`; - this subclass provides the PhysX sim config and the per-articulation - adapter (``articulation.newton_actuator_adapter``). - """ - - def _make_sim_cfg(self, use_newton_actuators: bool) -> SimulationCfg: - return SimulationCfg(dt=DT, physics=PhysxCfg(), use_newton_actuators=use_newton_actuators) - - def _make_articulation(self) -> Articulation: - return Articulation(replace(ANYMAL_C_CFG, actuators=DELAYED_PD_ACTUATORS, prim_path="/World/Env_.*/Robot")) - - def _get_adapter(self, articulation): - return articulation.newton_actuator_adapter - - -# --------------------------------------------------------------------------- -# RemotizedPD equivalence: PD + delay + position-based clamping lookup table -# --------------------------------------------------------------------------- - - -class TestRemotizedPDEquivalence(_EquivalenceTestBase): - """RemotizedPD (PD + delay + position-based clamping): Lab vs Newton (PhysX). - - Uses the Spot knee lookup table on ANYmal's KFE joints with IdealPD - on HAA and HFE. - """ - - __test__ = True - - @classmethod - def setUpClass(cls): - from isaaclab.actuators.actuator_pd_cfg import RemotizedPDActuatorCfg # noqa: PLC0415 - - cls.actuators = { - "hips": IdealPDActuatorCfg( - joint_names_expr=[".*HAA", ".*HFE"], - stiffness=40.0, - damping=5.0, - actuator_effort_limit=80.0, - ), - "knees": RemotizedPDActuatorCfg( - joint_names_expr=[".*KFE"], - stiffness=60.0, - damping=1.5, - actuator_effort_limit=80.0, - max_delay=3, - joint_parameter_lookup=SPOT_KNEE_LOOKUP, - ), - } - super().setUpClass() - - -# --------------------------------------------------------------------------- -# Neural network actuator authoring: MLP and LSTM -# --------------------------------------------------------------------------- - - -def _assert_network_drives_haa_joints(test: unittest.TestCase, result: dict) -> None: - """Assert that the network actuator produces finite, non-zero HAA efforts that follow the joint state.""" - haa_ids = [i for i, name in enumerate(result["joint_names"]) if name.endswith("HAA")] - test.assertEqual(len(haa_ids), 4) - for step_i, (pos, effort) in enumerate(zip(result["joint_pos"], result["applied_effort"])): - test.assertTrue(torch.isfinite(pos).all(), f"Non-finite positions at step {step_i}") - haa_effort = effort[:, haa_ids] - test.assertTrue(torch.all(haa_effort != 0.0), f"Network produced zero HAA effort at step {step_i}") - # A constant output (e.g. only the output-layer bias) would mean the network is not fed the moving joint state. - first_effort = result["applied_effort"][0][:, haa_ids] - last_effort = result["applied_effort"][-1][:, haa_ids] - test.assertFalse(torch.allclose(first_effort, last_effort), "Network HAA effort did not follow the joint state") - - -class TestNeuralMLPFunctional(unittest.TestCase): - """Verify ActuatorNetMLPCfg runs on PhysX with Newton actuators.""" - - @classmethod - def setUpClass(cls): - from isaaclab.actuators.actuator_net_cfg import ActuatorNetMLPCfg # noqa: PLC0415 - - cls.mlp_path = make_dummy_mlp_checkpoint() - cls.result = _run_simulation( - { - "mlp_legs": ActuatorNetMLPCfg( - joint_names_expr=[".*HAA"], - network_file=cls.mlp_path, - saturation_effort=120.0, - actuator_effort_limit=80.0, - actuator_velocity_limit=7.5, - pos_scale=-1.0, - vel_scale=1.0, - torque_scale=1.0, - input_order="pos_vel", - input_idx=[0, 1, 2], - ), - "pd_legs": IdealPDActuatorCfg( - joint_names_expr=[".*HFE", ".*KFE"], - stiffness=40.0, - damping=5.0, - actuator_effort_limit=80.0, - ), - }, - use_newton_actuators=True, - ) - - @classmethod - def tearDownClass(cls): - os.unlink(cls.mlp_path) - - def test_network_drives_haa_joints(self): - _assert_network_drives_haa_joints(self, self.result) - - -class TestNeuralLSTMFunctional(unittest.TestCase): - """Verify ActuatorNetLSTMCfg runs on PhysX with Newton actuators.""" - - @classmethod - def setUpClass(cls): - from isaaclab.actuators.actuator_net_cfg import ActuatorNetLSTMCfg # noqa: PLC0415 - - cls.lstm_path = make_dummy_lstm_checkpoint() - cls.result = _run_simulation( - { - "lstm_legs": ActuatorNetLSTMCfg( - joint_names_expr=[".*HAA"], - network_file=cls.lstm_path, - saturation_effort=120.0, - actuator_effort_limit=80.0, - actuator_velocity_limit=7.5, - ), - "pd_legs": IdealPDActuatorCfg( - joint_names_expr=[".*HFE", ".*KFE"], - stiffness=40.0, - damping=5.0, - actuator_effort_limit=80.0, - ), - }, - use_newton_actuators=True, - ) - - @classmethod - def tearDownClass(cls): - os.unlink(cls.lstm_path) - - def test_network_drives_haa_joints(self): - _assert_network_drives_haa_joints(self, self.result) - - -if __name__ == "__main__": - unittest.main() diff --git a/source/isaaclab_physx/test/assets/test_rigid_object.py b/source/isaaclab_physx/test/assets/test_rigid_object.py index ac84e269f1b6..8a0d87c96c6e 100644 --- a/source/isaaclab_physx/test/assets/test_rigid_object.py +++ b/source/isaaclab_physx/test/assets/test_rigid_object.py @@ -3,774 +3,483 @@ # # SPDX-License-Identifier: BSD-3-Clause +"""Real PhysX rigid-object and rigid-object-collection coverage. + +Most checks run against one module-scoped scene of locally authored cubes, holding both asset families. Each asset +holds two environments so that partial writes can target one environment and prove that the other is preserved in +the real PhysX state; PhysX stores collection bodies in body-major view order, so collection writes select +non-sorted environment and body subsets and read the view back. Tests that need their own simulation context are +defined first: pytest runs them before the composite scene is created, and a new simulation context would replace +the composite stage. +""" + # ignore private usage of variables warning # pyright: reportPrivateUsage=none -from isaaclab.test.utils import launch_test_simulation, test_devices +from isaaclab.test.utils import DeviceScope, launch_test_simulation, test_devices launch_test_simulation() +import math import sys +from collections.abc import Iterator +from dataclasses import dataclass from typing import Literal +from unittest.mock import MagicMock, patch import pytest import torch import warp as wp -from flaky import flaky -from isaaclab_physx.assets import RigidObject +from isaaclab_physx.assets import RigidObject, RigidObjectCollection +from isaaclab_physx.sim.schemas import PhysxArticulationCfg, PhysxRigidBodyCfg -import isaaclab.sim as sim_utils -from isaaclab.assets import RigidObjectCfg -from isaaclab.sim import build_simulation_context -from isaaclab.sim.spawners import materials -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR -from isaaclab.utils.math import ( - combine_frame_transforms, - default_orientation, - quat_apply_inverse, - quat_mul, - random_orientation, -) +import omni.kit.app +import omni.physx +from pxr import UsdPhysics +import isaaclab.sim as sim_utils +from isaaclab.assets import RigidObjectCfg, RigidObjectCollectionCfg +from isaaclab.sim import SimulationContext, build_simulation_context +from isaaclab.sim.schemas import apply_articulation_root_properties +from isaaclab.utils.math import combine_frame_transforms, quat_apply, quat_apply_inverse, quat_mul + +_NUM_ENVS = 2 +_NUM_CUBES = 3 +"""Cubes per environment of the dynamic collection.""" + + +def _cube_spawn(*, kinematic: bool = False, disable_gravity: bool = False) -> sim_utils.CuboidCfg: + """Create the spawn configuration of a local 1 kg collision cube.""" + return sim_utils.CuboidCfg( + size=(0.2, 0.2, 0.2), + rigid_props=[ + sim_utils.UsdPhysicsRigidBodyCfg(kinematic_enabled=kinematic), + PhysxRigidBodyCfg(disable_gravity=disable_gravity), + ], + mass_props=sim_utils.MassCfg(mass=1.0), + collision_props=sim_utils.UsdPhysicsCollisionCfg(), + ) -def generate_cubes_scene( - num_cubes: int = 1, - height=1.0, - api: Literal["none", "rigid_body", "articulation_root"] = "rigid_body", - kinematic_enabled: bool = False, - device: str = "cuda:0", -) -> tuple[RigidObject, torch.Tensor]: - """Generate a scene with the provided number of cubes. - Args: - num_cubes: Number of cubes to generate. - height: Height of the cubes. - api: The type of API that the cubes should have. - kinematic_enabled: Whether the cubes are kinematic. - device: Device to use for the simulation. +_STATIC_COLLIDER = sim_utils.CuboidCfg(size=(0.1, 0.1, 0.1), collision_props=sim_utils.UsdPhysicsCollisionCfg()) +"""A cube without rigid body properties, which PhysX treats as a static collider.""" - Returns: - A tuple containing the rigid object representing the cubes and the origins of the cubes. - """ - origins = torch.tensor([(i * 1.0, 0, height) for i in range(num_cubes)]).to(device) - # Create Top-level Xforms, one for each cube - for i, origin in enumerate(origins): - sim_utils.create_prim(f"/World/Table_{i}", "Xform", translation=origin) - - # Resolve spawn configuration - if api == "none": - # since no rigid body properties defined, this is just a static collider - spawn_cfg = sim_utils.CuboidCfg( - size=(0.1, 0.1, 0.1), - collision_props=sim_utils.UsdPhysicsCollisionCfg(), - ) - elif api == "rigid_body": - spawn_cfg = sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", - rigid_props=sim_utils.UsdPhysicsRigidBodyCfg(kinematic_enabled=kinematic_enabled), - ) - elif api == "articulation_root": - spawn_cfg = sim_utils.UsdFileCfg( - usd_path=f"{ISAACLAB_NUCLEUS_DIR}/Tests/RigidObject/Cube/dex_cube_instanceable_with_articulation_root.usd", - # Only tune existing bodies; do not create one on this invalid articulation fixture. - rigid_props={"(/.*)?": [sim_utils.UsdPhysicsRigidBodyCfg(kinematic_enabled=kinematic_enabled)]}, - ) - else: - raise ValueError(f"Unknown api: {api}") - - # Create rigid object - cube_object_cfg = RigidObjectCfg( - prim_path="/World/Table_[^/]*/Object", - spawn=spawn_cfg, - init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, height)), +def _cube_cfg(prim_path: str, **spawn_kwargs) -> RigidObjectCfg: + """Create a local cube one meter above its environment origin.""" + return RigidObjectCfg( + prim_path=prim_path, + spawn=_cube_spawn(**spawn_kwargs), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 1.0)), ) - cube_object = RigidObject(cfg=cube_object_cfg) - - return cube_object, origins -@pytest.mark.parametrize("num_cubes", [1, 2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.isaacsim_ci -def test_initialization(num_cubes, device): - """Test initialization for prim with rigid body API at the provided prim path.""" - with build_simulation_context(device=device, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - # Generate cubes scene - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, device=device) +def _spawn_envs(root: str, y_offset: float = 0.0) -> torch.Tensor: + """Create two environment prims 2 m apart and return their origins.""" + origins = torch.tensor([(2.0 * index, y_offset, 0.0) for index in range(_NUM_ENVS)]) + for index, origin in enumerate(origins.tolist()): + sim_utils.create_prim(f"{root}/Env_{index}", "Xform", translation=origin) + return origins - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(cube_object) < 10 - # Play sim - sim.reset() +def _collection( + root: str, num_envs: int, num_cubes: int, spawn: sim_utils.CuboidCfg, y_offset: float = 0.0 +) -> tuple[RigidObjectCollection, torch.Tensor]: + """Create a collection of cubes 3 m apart along y within environments 3 m apart along x, 1 m above ground. - # Check if object is initialized - assert cube_object.is_initialized - assert len(cube_object.body_names) == 1 + Returns: + The collection and the origins of its environments. + """ + origins = torch.tensor([(3.0 * index, y_offset, 1.0) for index in range(num_envs)]) + for index, origin in enumerate(origins.tolist()): + sim_utils.create_prim(f"{root}/Table_{index}", "Xform", translation=origin) + rigid_objects = { + f"cube_{index}": RigidObjectCfg( + prim_path=f"{root}/Table_[^/]*/Object_{index}", + spawn=spawn, + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 3.0 * index, 1.0)), + ) + for index in range(num_cubes) + } + return RigidObjectCollection(cfg=RigidObjectCollectionCfg(rigid_objects=rigid_objects)), origins - # Check buffers that exists and have correct shapes - assert cube_object.data.root_pos_w.torch.shape == (num_cubes, 3) - assert cube_object.data.root_quat_w.torch.shape == (num_cubes, 4) - assert cube_object.data.body_mass.torch.shape == (num_cubes, 1) - assert cube_object.data.body_inertia.torch.shape == (num_cubes, 1, 9) - # Simulate physics - for _ in range(2): - # perform rendering - sim.step() - # update object - cube_object.update(sim.cfg.dt) +def _yaw_quat(angle: float) -> tuple[float, float, float, float]: + """Return a unit quaternion ``(x, y, z, w)`` for a rotation of ``angle`` [rad] about the world z axis.""" + return (0.0, 0.0, math.sin(0.5 * angle), math.cos(0.5 * angle)) -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.isaacsim_ci -def test_initialization_with_kinematic_enabled(num_cubes, device): - """Test that initialization for prim with kinematic flag enabled.""" +@pytest.fixture +def sim(device: str) -> Iterator[SimulationContext]: + """Create a function-scoped simulation context for tests that own their scene.""" + # A new context would replace the stage of a live composite scene, so these tests must run before it. + assert SimulationContext.instance() is None, "define tests that own a simulation above the composite scene" with build_simulation_context(device=device, auto_add_lighting=True) as sim: sim._app_control_on_stop_handle = None - # Generate cubes scene - cube_object, origins = generate_cubes_scene(num_cubes=num_cubes, kinematic_enabled=True, device=device) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(cube_object) < 10 - - # Play sim - sim.reset() + yield sim - # Check if object is initialized - assert cube_object.is_initialized - assert len(cube_object.body_names) == 1 - - # Check buffers that exists and have correct shapes - assert cube_object.data.root_pos_w.torch.shape == (num_cubes, 3) - assert cube_object.data.root_quat_w.torch.shape == (num_cubes, 4) - - # Simulate physics - for _ in range(2): - # perform rendering - sim.step() - # update object - cube_object.update(sim.cfg.dt) - # check that the object is kinematic - default_root_pose = cube_object.data.default_root_pose.torch.clone() - default_root_vel = cube_object.data.default_root_vel.torch.clone() - default_root_pose[:, :3] += origins - torch.testing.assert_close(cube_object.data.root_link_pose_w.torch, default_root_pose) - torch.testing.assert_close(cube_object.data.root_com_vel_w.torch, default_root_vel) - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.isaacsim_ci -def test_initialization_with_no_rigid_body(num_cubes, device): - """Test that initialization fails when no rigid body is found at the provided prim path.""" - with build_simulation_context(device=device, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - # Generate cubes scene - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, api="none", device=device) - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(cube_object) < 10 +## +# Tests that own their simulation context. Keep them above the composite scene. +## - # Play sim - with pytest.raises(RuntimeError): - sim.reset() +@pytest.mark.xfail( + strict=True, raises=AssertionError, reason="The first CPU simulation after a failed initialization reads zero poses" +) +def test_collection_without_rigid_bodies_fails_without_leaking_into_the_next_cpu_simulation() -> None: + """A collection of static colliders fails to initialize; a kinematic cube then reports its spawn pose in the next + CPU simulation. -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.isaacsim_ci -def test_initialization_with_articulation_root(num_cubes, device): - """Test that initialization fails when an articulation root is found at the provided prim path.""" - with build_simulation_context(device=device, auto_add_lighting=True) as sim: + The failure only reproduces before any CUDA simulation ran in the process, so this test comes first. + """ + with build_simulation_context(device="cpu") as sim: sim._app_control_on_stop_handle = None - # Generate cubes scene - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, api="articulation_root", device=device) - + invalid_collection, _ = _collection("/World", 1, 2, _STATIC_COLLIDER) # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(cube_object) < 10 - - # Play sim + assert sys.getrefcount(invalid_collection) < 10 with pytest.raises(RuntimeError): sim.reset() - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.isaacsim_ci -def test_set_rigid_object_state(num_cubes, device): - """Test setting the state of the rigid object. - - In this test, we set the state of the rigid object to a random state and check - that the object is in that state after simulation. We set gravity to zero as - we don't want any external forces acting on the object to ensure state remains static. - """ - # Turn off gravity for this test as we don't want any external forces acting on the object - # to ensure state remains static - with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + with build_simulation_context(device="cpu") as sim: sim._app_control_on_stop_handle = None - # Generate cubes scene - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, device=device) - - # Play the simulator + kinematic, origins = _collection("/World", 1, 1, _cube_spawn(kinematic=True)) sim.reset() - - state_types = ["root_pos_w", "root_quat_w", "root_lin_vel_w", "root_ang_vel_w"] - - # Set each state type individually as they are dependent on each other - for state_type_to_randomize in state_types: - state_dict = { - "root_pos_w": torch.zeros_like(cube_object.data.root_pos_w.torch, device=sim.device), - "root_quat_w": default_orientation(num=num_cubes, device=sim.device), - "root_lin_vel_w": torch.zeros_like(cube_object.data.root_lin_vel_w.torch, device=sim.device), - "root_ang_vel_w": torch.zeros_like(cube_object.data.root_ang_vel_w.torch, device=sim.device), - } - - # Now we are ready! - for _ in range(5): - # reset object - cube_object.reset() - - # Set random state - if state_type_to_randomize == "root_quat_w": - state_dict[state_type_to_randomize] = random_orientation(num=num_cubes, device=sim.device) - else: - state_dict[state_type_to_randomize] = torch.randn(num_cubes, 3, device=sim.device) - - # perform simulation - for _ in range(5): - root_pose = torch.cat( - [state_dict["root_pos_w"], state_dict["root_quat_w"]], - dim=-1, - ) - root_vel = torch.cat( - [state_dict["root_lin_vel_w"], state_dict["root_ang_vel_w"]], - dim=-1, - ) - # reset root state - cube_object.write_root_pose_to_sim_index(root_pose=root_pose) - cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) - - sim.step() - - # assert that set root quantities are equal to the ones set in the state_dict - for key, expected_value in state_dict.items(): - value = getattr(cube_object.data, key).torch - torch.testing.assert_close(value, expected_value, rtol=1e-3, atol=1e-3) - - cube_object.update(sim.cfg.dt) - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.isaacsim_ci -def test_reset_rigid_object(num_cubes, device): - """Test resetting the state of the rigid object.""" - with build_simulation_context(device=device, gravity_enabled=True, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - # Generate cubes scene - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, device=device) - - # Play the simulator - sim.reset() - sim.step() - cube_object.update(sim.cfg.dt) - - composers = (cube_object.instantaneous_wrench_composer, cube_object.permanent_wrench_composer) - ones = torch.ones((num_cubes, 1, 3), device=device) - cube_object.permanent_wrench_composer.set_forces_and_torques_index(forces=ones, torques=ones) - cube_object.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones) - - # A partial reset clears only the selected environment - cube_object.reset(env_ids=torch.tensor([0], device=device)) - for composer in composers: - assert composer.active - for buffer in (composer.out_force_b.torch, composer.out_torque_b.torch): - assert torch.count_nonzero(buffer[0]) == 0 - assert torch.count_nonzero(buffer[1:]) == buffer[1:].numel() - - # A full reset clears every environment and deactivates the composers - cube_object.reset() - for composer in composers: - assert not composer.active - assert torch.count_nonzero(composer.out_force_b.torch) == 0 - assert torch.count_nonzero(composer.out_torque_b.torch) == 0 - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) + torch.testing.assert_close(kinematic.data.body_link_pos_w.torch[:, 0], origins + torch.tensor([0.0, 0.0, 1.0])) + + +@pytest.mark.parametrize( + "api", + [ + "none", + pytest.param( + "articulation_root", + marks=pytest.mark.xfail( + strict=True, + raises=pytest.fail.Exception, + reason="The PhysX rigid object does not reject an enabled articulation root.", + ), + ), + ], +) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) @pytest.mark.isaacsim_ci -def test_rigid_body_no_friction(num_cubes, device): - """Test that a rigid object with no friction maintains its tangential velocity on a plane.""" - with build_simulation_context(device=device, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - # Generate cubes scene - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=0.0, device=device) - - # Create ground plane with no friction - cfg = sim_utils.GroundPlaneCfg( - physics_material=materials.RigidBodyMaterialCfg( - static_friction=0.0, - dynamic_friction=0.0, - restitution=0.0, - ) - ) - cfg.func("/World/GroundPlane", cfg) - - # Play sim - sim.reset() - - # Set material friction properties to be all zero - static_friction = torch.zeros(num_cubes, 1) - dynamic_friction = torch.zeros(num_cubes, 1) - restitution = torch.FloatTensor(num_cubes, 1).uniform_(0.0, 0.2) - - cube_object_materials = torch.cat([static_friction, dynamic_friction, restitution], dim=-1) - indices = torch.tensor(range(num_cubes), dtype=torch.int32) - - cube_object.root_view.set_material_properties( - wp.from_torch(cube_object_materials, dtype=wp.float32), wp.from_torch(indices, dtype=wp.int32) +def test_initialization_rejects_invalid_rigid_body(sim, device, api: Literal["none", "articulation_root"]) -> None: + """Initialization fails without a rigid body and when the rigid body is an articulation root.""" + _spawn_envs("/World") + if api == "none": + # Without rigid body properties the cube is a static collider. + cube_object = RigidObject(cfg=RigidObjectCfg(prim_path="/World/Env_[^/]*/Object", spawn=_STATIC_COLLIDER)) + else: + cube_object = RigidObject(cfg=_cube_cfg("/World/Env_[^/]*/Object")) + # An enabled articulation root on the rigid body turns the body into an articulation link. + assert apply_articulation_root_properties( + "/World/Env_[^/]*/Object", + [PhysxArticulationCfg(articulation_enabled=True)], + stage=sim.stage, + create_if_missing=True, ) + for index in range(_NUM_ENVS): + prim = sim.stage.GetPrimAtPath(f"/World/Env_{index}/Object") + assert prim.HasAPI(UsdPhysics.ArticulationRootAPI) + assert prim.GetAttribute("physxArticulation:articulationEnabled").Get() - # Set initial velocity - # Initial velocity in X to get the block moving - initial_velocity = torch.zeros((num_cubes, 6), device=sim.cfg.device) - initial_velocity[:, 0] = 0.1 + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(cube_object) < 10 - cube_object.write_root_velocity_to_sim_index(root_velocity=initial_velocity) - - # Simulate physics - for _ in range(5): - # perform rendering - sim.step() - # update object - cube_object.update(sim.cfg.dt) - - # Non-deterministic when on GPU, so we use different tolerances - if device.startswith("cuda"): - tolerance = 1e-2 - else: - tolerance = 1e-5 - - torch.testing.assert_close( - cube_object.data.root_lin_vel_w.torch[:, :2], initial_velocity[:, :2], rtol=1e-5, atol=tolerance - ) - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.isaacsim_ci -def test_rigid_body_with_static_friction(num_cubes, device): - """Test that static friction applied to rigid object works as expected. - - This test works by applying a force to the object and checking if the object moves or not based on the - mu (coefficient of static friction) value set for the object. We set the static friction to be non-zero and - apply a force to the object. When the force applied is below mu, the object should not move. When the force - applied is above mu, the object should move. - """ - with build_simulation_context(device=device, dt=0.01, add_ground_plane=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=0.03125, device=device) - - # Create ground plane - static_friction_coefficient = 0.5 - cfg = sim_utils.GroundPlaneCfg( - physics_material=materials.RigidBodyMaterialCfg( - static_friction=static_friction_coefficient, - dynamic_friction=static_friction_coefficient, # This shouldn't be required but is due to a bug in PhysX - ) - ) - cfg.func("/World/GroundPlane", cfg) - - # Play sim + with pytest.raises(RuntimeError): sim.reset() - # Set static friction to be non-zero - # Dynamic friction also needs to be zero due to a bug in PhysX - static_friction = torch.Tensor([[static_friction_coefficient]] * num_cubes) - dynamic_friction = torch.Tensor([[static_friction_coefficient]] * num_cubes) - restitution = torch.zeros(num_cubes, 1) - - cube_object_materials = torch.cat([static_friction, dynamic_friction, restitution], dim=-1) - - indices = torch.tensor(range(num_cubes), dtype=torch.int32) - - # Add friction to cube - cube_object.root_view.set_material_properties( - wp.from_torch(cube_object_materials, dtype=wp.float32), wp.from_torch(indices, dtype=wp.int32) - ) - - # let everything settle - for _ in range(100): - sim.step() - cube_object.update(sim.cfg.dt) - cube_object.write_root_velocity_to_sim_index(root_velocity=torch.zeros((num_cubes, 6), device=sim.device)) - cube_mass = wp.to_torch(cube_object.root_view.get_masses()) - gravity_magnitude = abs(sim.cfg.gravity[2]) - # 2 cases: force applied is below and above mu - # below mu: block should not move as the force applied is <= mu - # above mu: block should move as the force applied is > mu - for force in "below_mu", "above_mu": - # set initial velocity to zero - cube_object.write_root_velocity_to_sim_index(root_velocity=torch.zeros((num_cubes, 6), device=sim.device)) - - external_wrench_b = torch.zeros((num_cubes, 1, 6), device=sim.device) - if force == "below_mu": - external_wrench_b[..., 0] = static_friction_coefficient * cube_mass * gravity_magnitude * 0.99 - else: - external_wrench_b[..., 0] = static_friction_coefficient * cube_mass * gravity_magnitude * 1.01 - - cube_object.permanent_wrench_composer.set_forces_and_torques_index( - forces=external_wrench_b[..., :3], - torques=external_wrench_b[..., 3:], - ) - # Get root state - initial_root_pos = cube_object.data.root_pos_w.torch.clone() - # Simulate physics - for _ in range(200): - # apply the wrench - cube_object.write_data_to_sim() - sim.step() - # update object - cube_object.update(sim.cfg.dt) - if force == "below_mu": - # Assert that the block has not moved - torch.testing.assert_close( - cube_object.data.root_pos_w.torch, initial_root_pos, rtol=2e-3, atol=2e-3 - ) - if force == "above_mu": - assert (cube_object.data.root_pos_w.torch[..., 0] - initial_root_pos[..., 0] > 0.02).all() - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) @pytest.mark.isaacsim_ci -def test_rigid_body_with_restitution(num_cubes, device): - """Test that restitution when applied to rigid object works as expected. - - This test works by dropping a block from a height and checking if the block bounces or not based on the - restitution value set for the object. We set the restitution to be non-zero and drop the block from a height. - When the restitution is 0, the block should not bounce. When the restitution is between 0 and 1, the block - should bounce with less energy. - """ - for expected_collision_type in "partially_elastic", "inelastic": - with build_simulation_context(device=device, add_ground_plane=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=1.0, device=device) - - # Set static friction to be non-zero - if expected_collision_type == "inelastic": - restitution_coefficient = 0.0 - elif expected_collision_type == "partially_elastic": - restitution_coefficient = 0.5 - - # Create ground plane such that has a restitution of 1.0 (perfectly elastic collision) - cfg = sim_utils.GroundPlaneCfg( - physics_material=materials.RigidBodyMaterialCfg( - restitution=restitution_coefficient, - ) - ) - cfg.func("/World/GroundPlane", cfg) - - indices = torch.tensor(range(num_cubes), dtype=torch.int32) - - # Play sim - sim.reset() - - root_pose = torch.zeros(num_cubes, 7, device=sim.device) - root_pose[:, 3] = 1.0 # To make orientation a quaternion - for i in range(num_cubes): - root_pose[i, 1] = 1.0 * i - root_pose[:, 2] = 1.0 # Set an initial drop height - root_vel = torch.zeros(num_cubes, 6, device=sim.device) - root_vel[:, 2] = -1.0 # Set an initial downward velocity - - cube_object.write_root_pose_to_sim_index(root_pose=root_pose) - cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) - - static_friction = torch.zeros(num_cubes, 1) - dynamic_friction = torch.zeros(num_cubes, 1) - restitution = torch.Tensor([[restitution_coefficient]] * num_cubes) - - cube_object_materials = torch.cat([static_friction, dynamic_friction, restitution], dim=-1) - - # Add restitution to cube - cube_object.root_view.set_material_properties( - wp.from_torch(cube_object_materials, dtype=wp.float32), wp.from_torch(indices, dtype=wp.int32) - ) - - curr_z_velocity = cube_object.data.root_lin_vel_w.torch[:, 2].clone() - - for _ in range(100): - sim.step() - - # update object - cube_object.update(sim.cfg.dt) - curr_z_velocity = cube_object.data.root_lin_vel_w.torch[:, 2].clone() - - if expected_collision_type == "inelastic": - # Allow a small contact separation velocity while ensuring that the block does not bounce. - assert (curr_z_velocity <= 1e-3).all() - - if torch.all(curr_z_velocity <= 0.0): - # Still in the air - prev_z_velocity = curr_z_velocity - else: - # collision has happened, exit the for loop - break - - if expected_collision_type == "partially_elastic": - # Assert that the block has lost some energy by checking that the z velocity is less - assert torch.all(torch.le(abs(curr_z_velocity), abs(prev_z_velocity))) - assert (curr_z_velocity > 0.0).all() - - -@pytest.mark.parametrize("num_cubes", [2]) @pytest.mark.parametrize("device", test_devices()) -@pytest.mark.isaacsim_ci -def test_rigid_body_set_mass(num_cubes, device): - """Test getting and setting mass of rigid object.""" - with build_simulation_context( - device=device, gravity_enabled=False, add_ground_plane=True, auto_add_lighting=True - ) as sim: - sim._app_control_on_stop_handle = None - # Create a scene with random cubes - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=1.0, device=device) - - # Play sim +def test_warmup_loads_physics_once(sim, device) -> None: + """Attach on GPU or force-load on CPU, without destroying and rebuilding native objects.""" + sim_utils.create_prim("/World/Table_0", "Xform", translation=(0.0, 0.0, 1.0)) + cube = RigidObject(cfg=_cube_cfg("/World/Table_[^/]*/Object")) + + # Wrap read-only native interfaces through their accessors; real calls still execute. + physx_spy = MagicMock(wraps=omni.physx.get_physx_interface()) + physx_sim_spy = MagicMock(wraps=omni.physx.get_physx_simulation_interface()) + with ( + patch("omni.physx.get_physx_interface", return_value=physx_spy), + patch("omni.physx.get_physx_simulation_interface", return_value=physx_sim_spy), + ): sim.reset() - # Get masses before increasing - original_masses = wp.to_torch(cube_object.root_view.get_masses()) - - assert original_masses.shape == (num_cubes, 1) - - # Randomize mass of the object through the public writer - masses = original_masses + torch.FloatTensor(num_cubes, 1).uniform_(4, 8) - cube_object.set_masses_index(masses=masses.to(device)) - - torch.testing.assert_close(cube_object.data.body_mass.torch, masses.to(device)) - torch.testing.assert_close(wp.to_torch(cube_object.root_view.get_masses()), masses) - - # The written mass must survive a simulation step - sim.step() - cube_object.update(sim.cfg.dt) + extension_manager = omni.kit.app.get_app().get_extension_manager() + assert extension_manager.is_extension_enabled("omni.physics.physx"), ( + "The omni.physics.physx bridge must register PhysX with the unified physics API." + ) + assert physx_sim_spy.attach_stage.call_count == int(device.startswith("cuda")) + assert physx_spy.force_load_physics_from_usd.call_count == int(device == "cpu") + assert cube.is_initialized and cube.num_instances == 1 + initial = cube.data.root_pos_w.torch.clone() + for _ in range(10): + sim.step(render=False) + cube.update(sim.get_physics_dt()) + assert torch.all(cube.data.root_pos_w.torch[:, 2] < initial[:, 2]) + + +## +# Composite scene shared by the remaining tests. +## + + +@dataclass +class _Scene: + """Rigid objects and collections that share one real PhysX lifecycle under gravity.""" + + sim: SimulationContext + device: str + cubes: RigidObject + """Dynamic cubes that ignore gravity.""" + kinematic: RigidObject + """Kinematic cubes.""" + falling: RigidObject + """Dynamic cubes under gravity.""" + collection: RigidObjectCollection + """Two environments of three dynamic cubes that ignore gravity.""" + kinematic_collection: RigidObjectCollection + """One environment of one kinematic cube.""" + falling_collection: RigidObjectCollection + """Two environments of two dynamic cubes under gravity.""" + origins: dict[str, torch.Tensor] + """Environment origins keyed by asset name.""" + refcounts: dict[str, int] + """Reference count of each asset right after construction.""" + + @property + def assets(self) -> tuple[RigidObject | RigidObjectCollection, ...]: + return ( + self.cubes, + self.kinematic, + self.falling, + self.collection, + self.kinematic_collection, + self.falling_collection, + ) - torch.testing.assert_close(cube_object.data.body_mass.torch, masses.to(device)) - torch.testing.assert_close(wp.to_torch(cube_object.root_view.get_masses()), masses) + def step(self, num_steps: int = 1) -> None: + """Write, step, and update every asset.""" + for _ in range(num_steps): + for asset in self.assets: + asset.write_data_to_sim() + self.sim.step() + for asset in self.assets: + asset.update(self.sim.cfg.dt) + + def place_cubes_at_rest(self, root_pose: torch.Tensor) -> None: + """Teleport the dynamic cubes to ``root_pose`` at rest and clear their external wrenches.""" + self.cubes.write_root_link_pose_to_sim_index(root_pose=root_pose) + self.cubes.write_root_com_velocity_to_sim_index(root_velocity=torch.zeros((_NUM_ENVS, 6), device=self.device)) + self.cubes.reset() + + def place_collection_at_rest(self, yaw: float = 0.0) -> torch.Tensor: + """Teleport the dynamic collection to its default poses with a given yaw, at rest and without wrenches. + + Returns: + The body poses the collection was placed at. + """ + poses = self.collection.data.default_body_pose.torch.clone() + poses[..., :2] += self.origins["collection"].unsqueeze(1)[..., :2] + poses[..., 3:] = torch.tensor(_yaw_quat(yaw), device=self.device) + self.collection.write_body_link_pose_to_sim_index(body_poses=poses) + self.collection.write_body_com_velocity_to_sim_index( + body_velocities=torch.zeros((_NUM_ENVS, _NUM_CUBES, 6), device=self.device) + ) + self.collection.permanent_wrench_composer.reset() + self.collection.instantaneous_wrench_composer.reset() + return poses -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("gravity_enabled", [True, False]) -@pytest.mark.isaacsim_ci -def test_gravity_vec_w(num_cubes, device, gravity_enabled): - """Test that gravity vector direction is set correctly for the rigid object.""" - with build_simulation_context(device=device, gravity_enabled=gravity_enabled) as sim: +@pytest.fixture(scope="module", params=test_devices(DeviceScope.CUDA)) +def scene(request) -> Iterator[_Scene]: + """Initialize the composite scene once.""" + device = request.param + with build_simulation_context(device=device, gravity_enabled=True) as sim: sim._app_control_on_stop_handle = None - # Create a scene with random cubes - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, device=device) - - # Obtain gravity direction - if gravity_enabled: - gravity_dir = (0.0, 0.0, -1.0) - else: - gravity_dir = (0.0, 0.0, 0.0) - - # Play sim + origins = { + "cubes": _spawn_envs("/World/Cubes"), + "kinematic": _spawn_envs("/World/Kinematic", 3.0), + "falling": _spawn_envs("/World/Falling", 6.0), + } + assets = { + "cubes": RigidObject(cfg=_cube_cfg("/World/Cubes/Env_[^/]*/Object", disable_gravity=True)), + "kinematic": RigidObject(cfg=_cube_cfg("/World/Kinematic/Env_[^/]*/Object", kinematic=True)), + "falling": RigidObject(cfg=_cube_cfg("/World/Falling/Env_[^/]*/Object")), + } + for name, num_envs, num_cubes, spawn, y_offset in ( + ("collection", _NUM_ENVS, _NUM_CUBES, _cube_spawn(disable_gravity=True), 10.0), + ("kinematic_collection", 1, 1, _cube_spawn(kinematic=True), 22.0), + ("falling_collection", _NUM_ENVS, 2, _cube_spawn(), 26.0), + ): + assets[name], origins[name] = _collection(f"/World/{name}", num_envs, num_cubes, spawn, y_offset) + refcounts = {name: sys.getrefcount(asset) for name, asset in assets.items()} sim.reset() - - # Check that gravity is set correctly - assert cube_object.data.GRAVITY_VEC_W.torch[0, 0] == gravity_dir[0] - assert cube_object.data.GRAVITY_VEC_W.torch[0, 1] == gravity_dir[1] - assert cube_object.data.GRAVITY_VEC_W.torch[0, 2] == gravity_dir[2] - - # Simulate physics - for _ in range(2): - # perform rendering - sim.step() - # update object - cube_object.update(sim.cfg.dt) - - # Expected gravity value is the acceleration of the body - gravity = torch.zeros(num_cubes, 1, 6, device=device) - if gravity_enabled: - gravity[:, :, 2] = -9.81 - # Check the body accelerations are correct - torch.testing.assert_close(cube_object.data.body_acc_w.torch, gravity) + yield _Scene( + sim=sim, + device=device, + origins={name: value.to(device) for name, value in origins.items()}, + refcounts=refcounts, + **assets, + ) -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("with_offset", [True, False]) @pytest.mark.isaacsim_ci -@flaky(max_runs=3, min_passes=1) -def test_body_root_state_properties(num_cubes, device, with_offset): - """Test the root_com_state_w, root_link_state_w, body_com_state_w, and body_link_state_w properties.""" - with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - # Create a scene with random cubes - cube_object, env_pos = generate_cubes_scene(num_cubes=num_cubes, height=0.0, device=device) - - # Play sim - sim.reset() - - # Check if cube_object is initialized - assert cube_object.is_initialized - - # change center of mass offset from link frame - if with_offset: - offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_cubes, 1) - else: - offset = torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_cubes, 1) - - com = wp.to_torch(cube_object.root_view.get_coms()) - com[..., :3] = offset.to("cpu") - cube_object.set_coms_index(coms=wp.from_torch(com.unsqueeze(1).to(device).contiguous(), dtype=wp.transformf)) - - # check center of mass has been set - torch.testing.assert_close(wp.to_torch(cube_object.root_view.get_coms()), com) +def test_initialization(scene: _Scene) -> None: + """Initialize local rigid objects and collections, including a single-cube collection; under gravity, kinematic + cubes hold their default pose and dynamic cubes accelerate downward on every step.""" + for name, asset in zip( + ("cubes", "kinematic", "collection", "kinematic_collection"), scene.assets[:2] + scene.assets[3:5] + ): + # Check that the framework doesn't hold excessive strong references. + assert scene.refcounts[name] < 10, name + assert asset.is_initialized, name + for rigid_object in (scene.cubes, scene.kinematic): + assert rigid_object.num_instances == _NUM_ENVS + assert len(rigid_object.body_names) == 1 + assert rigid_object.data.root_pos_w.torch.shape == (_NUM_ENVS, 3) + assert rigid_object.data.root_quat_w.torch.shape == (_NUM_ENVS, 4) + assert rigid_object.data.body_mass.torch.shape == (_NUM_ENVS, 1) + assert rigid_object.data.body_inertia.torch.shape == (_NUM_ENVS, 1, 9) + torch.testing.assert_close( + rigid_object.data.GRAVITY_VEC_W.torch, torch.tensor([[0.0, 0.0, -1.0]] * _NUM_ENVS, device=scene.device) + ) + for collection, num_envs, num_cubes in ( + (scene.collection, _NUM_ENVS, _NUM_CUBES), + (scene.kinematic_collection, 1, 1), + ): + assert collection.num_instances == num_envs + assert len(collection.body_names) == num_cubes + assert collection.data.body_link_pos_w.torch.shape == (num_envs, num_cubes, 3) + assert collection.data.body_link_quat_w.torch.shape == (num_envs, num_cubes, 4) + assert collection.data.body_mass.torch.shape == (num_envs, num_cubes) + assert collection.data.body_inertia.torch.shape == (num_envs, num_cubes, 9) + torch.testing.assert_close( + collection.data.GRAVITY_VEC_W.torch[..., 2], torch.full((num_envs, num_cubes), -1.0, device=scene.device) + ) - # random z spin velocity - spin_twist = torch.zeros(6, device=device) - spin_twist[5] = torch.randn(1, device=device) + gravity_acceleration = torch.tensor([0.0, 0.0, -9.81, 0.0, 0.0, 0.0], device=scene.device) + kinematic, kinematic_collection = scene.kinematic, scene.kinematic_collection + for _ in range(2): + scene.step() + torch.testing.assert_close(scene.falling.data.body_acc_w.torch, gravity_acceleration.expand(_NUM_ENVS, 1, 6)) + torch.testing.assert_close( + scene.falling_collection.data.body_com_acc_w.torch, gravity_acceleration.expand(_NUM_ENVS, 2, 6) + ) + default_root_pose = kinematic.data.default_root_pose.torch.clone() + default_root_pose[:, :3] += scene.origins["kinematic"] + torch.testing.assert_close(kinematic.data.root_link_pose_w.torch, default_root_pose) + torch.testing.assert_close(kinematic.data.root_com_vel_w.torch, kinematic.data.default_root_vel.torch) + default_body_pose = kinematic_collection.data.default_body_pose.torch.clone() + default_body_pose[..., :3] += scene.origins["kinematic_collection"].unsqueeze(1) + torch.testing.assert_close(kinematic_collection.data.body_link_pose_w.torch, default_body_pose) + torch.testing.assert_close( + kinematic_collection.data.body_link_vel_w.torch, kinematic_collection.data.default_body_vel.torch + ) - # Simulate physics - for _ in range(100): - # spin the object around Z axis (com) - cube_object.write_root_velocity_to_sim_index(root_velocity=spin_twist.repeat(num_cubes, 1)) - # perform rendering - sim.step() - # update object - cube_object.update(sim.cfg.dt) - # get state properties - root_link_pose_w = cube_object.data.root_link_pose_w.torch - root_link_vel_w = cube_object.data.root_link_vel_w.torch - root_com_pose_w = cube_object.data.root_com_pose_w.torch - root_com_vel_w = cube_object.data.root_com_vel_w.torch - body_link_pose_w = cube_object.data.body_link_pose_w.torch - body_link_vel_w = cube_object.data.body_link_vel_w.torch - body_com_pose_w = cube_object.data.body_com_pose_w.torch - body_com_vel_w = cube_object.data.body_com_vel_w.torch - - # if offset is [0,0,0] all root_state_%_w will match and all body_%_w will match - if not with_offset: - torch.testing.assert_close(root_link_pose_w, root_com_pose_w) - torch.testing.assert_close(root_com_vel_w, root_link_vel_w) - torch.testing.assert_close(body_link_pose_w, body_com_pose_w) - torch.testing.assert_close(body_com_vel_w, body_link_vel_w) - else: - # cubes are spinning around center of mass - # position will not match - # center of mass position will be constant (i.e. spinning around com) - torch.testing.assert_close(env_pos + offset, root_com_pose_w[..., :3]) - torch.testing.assert_close(env_pos + offset, body_com_pose_w[..., :3].squeeze(-2)) - # link position will be moving but should stay constant away from center of mass - root_link_state_pos_rel_com = quat_apply_inverse( - root_link_pose_w[..., 3:], - root_link_pose_w[..., :3] - root_com_pose_w[..., :3], - ) - torch.testing.assert_close(-offset, root_link_state_pos_rel_com) - body_link_state_pos_rel_com = quat_apply_inverse( - body_link_pose_w[..., 3:], - body_link_pose_w[..., :3] - body_com_pose_w[..., :3], - ) - torch.testing.assert_close(-offset, body_link_state_pos_rel_com.squeeze(-2)) - - # orientation of com will be a constant rotation from link orientation - com_quat_b = cube_object.data.body_com_quat_b.torch - com_quat_w = quat_mul(body_link_pose_w[..., 3:], com_quat_b) - torch.testing.assert_close(com_quat_w, body_com_pose_w[..., 3:]) - torch.testing.assert_close(com_quat_w.squeeze(-2), root_com_pose_w[..., 3:]) - - # lin_vel will not match - # center of mass vel will be constant (i.e. spinning around com) - torch.testing.assert_close(torch.zeros_like(root_com_vel_w[..., :3]), root_com_vel_w[..., :3]) - torch.testing.assert_close(torch.zeros_like(body_com_vel_w[..., :3]), body_com_vel_w[..., :3]) - # link frame will be moving, and should be equal to input angular velocity cross offset - lin_vel_rel_root_gt = quat_apply_inverse(root_link_pose_w[..., 3:], root_link_vel_w[..., :3]) - lin_vel_rel_body_gt = quat_apply_inverse(body_link_pose_w[..., 3:], body_link_vel_w[..., :3]) - lin_vel_rel_gt = torch.linalg.cross(spin_twist.repeat(num_cubes, 1)[..., 3:], -offset) - torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_root_gt, atol=1e-4, rtol=1e-4) - torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_body_gt.squeeze(-2), atol=1e-4, rtol=1e-4) - - # ang_vel will always match - torch.testing.assert_close(root_com_vel_w[..., 3:], root_link_vel_w[..., 3:]) - torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("with_offset", [True, False]) -@pytest.mark.parametrize("state_location", ["com", "link", "root"]) @pytest.mark.isaacsim_ci -def test_write_root_state(num_cubes, device, with_offset, state_location): - """Test the root state setters in the link, center-of-mass, and mixed root frames. - - Every write must round-trip through the written frame and keep the other frame consistent with the - center-of-mass offset. - """ - with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - # Create a scene with random cubes - cube_object, env_pos = generate_cubes_scene(num_cubes=num_cubes, height=0.0, device=device) - env_idx = torch.tensor([x for x in range(num_cubes)], dtype=torch.int32) - - # Play sim - sim.reset() - - # Check if cube_object is initialized - assert cube_object.is_initialized - - # change center of mass offset from link frame - if with_offset: - offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_cubes, 1) - else: - offset = torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_cubes, 1) - - com = wp.to_torch(cube_object.root_view.get_coms()) - com[..., :3] = offset.to("cpu") - cube_object.root_view.set_coms(wp.from_torch(com, dtype=wp.float32), wp.from_torch(env_idx, dtype=wp.int32)) +def test_rigid_object_inertial_properties(scene: _Scene) -> None: + """Mass, center-of-mass, and inertia writes reach only the selected PhysX entries and survive a step.""" + device = scene.device + cube_object = scene.cubes + + # Full-data mass writes reach the view and survive a simulation step. + original_masses = wp.to_torch(cube_object.root_view.get_masses()).clone() + assert original_masses.shape == (_NUM_ENVS, 1) + masses = original_masses + torch.tensor([[4.5], [6.5]]) + cube_object.set_masses_index(masses=masses.to(device)) + torch.testing.assert_close(cube_object.data.body_mass.torch, masses.to(device)) + torch.testing.assert_close(wp.to_torch(cube_object.root_view.get_masses()), masses) + scene.step() + torch.testing.assert_close(cube_object.data.body_mass.torch, masses.to(device)) + torch.testing.assert_close(wp.to_torch(cube_object.root_view.get_masses()), masses) + + # Partial writes with an int64 selector change environment 1 only. + env_ids = torch.tensor([1], dtype=torch.int64, device=device) + body_ids = torch.tensor([0], dtype=torch.int32, device=device) + expected = { + "body_mass": cube_object.data.body_mass.torch.clone(), + "body_com_pose_b": cube_object.data.body_com_pose_b.torch.clone(), + "body_inertia": cube_object.data.body_inertia.torch.clone(), + } + expected["body_mass"][1, 0] = 3.0 + expected["body_com_pose_b"][1, 0, :3] = torch.tensor([0.03, -0.02, 0.01], device=device) + expected["body_inertia"][1, 0, [0, 4, 8]] *= torch.tensor([1.2, 1.3, 1.4], device=device) + cube_object.set_masses_index(masses=expected["body_mass"][1:], env_ids=env_ids, body_ids=body_ids) + cube_object.set_coms_index(coms=expected["body_com_pose_b"][1:], env_ids=env_ids, body_ids=body_ids) + cube_object.set_inertias_index(inertias=expected["body_inertia"][1:], env_ids=env_ids, body_ids=body_ids) + for name, data, raw in ( + ("body_mass", cube_object.data.body_mass, cube_object.root_view.get_masses()), + ("body_com_pose_b", cube_object.data.body_com_pose_b, cube_object.root_view.get_coms().view(wp.float32)), + ("body_inertia", cube_object.data.body_inertia, cube_object.root_view.get_inertias()), + ): + torch.testing.assert_close(data.torch, expected[name]) + torch.testing.assert_close(wp.to_torch(raw).to(device).reshape(expected[name].shape), expected[name]) - # check center of mass has been set - torch.testing.assert_close(wp.to_torch(cube_object.root_view.get_coms()), com) - # random pose and velocity so frame conversions see a non-trivial rotation - rand_state = torch.rand(num_cubes, 13, device=device) - rand_state[..., :3] += env_pos - # make quaternion a unit vector - rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) - - env_idx = env_idx.to(device) - for i in range(10): - # perform step - sim.step() - # update buffers - cube_object.update(sim.cfg.dt) - - env_ids = None if i % 2 == 0 else env_idx +@pytest.mark.isaacsim_ci +def test_rigid_object_root_state_writes(scene: _Scene) -> None: + """Root writes round-trip through the written frame and keep the other frame consistent with the offset.""" + device = scene.device + cube_object = scene.cubes + rest_pose = torch.cat( + (scene.origins["cubes"], torch.tensor([0.0, 0.0, 0.0, 1.0], device=device).repeat(_NUM_ENVS, 1)), -1 + ) + rest_pose[:, 2] += 1.0 + scene.place_cubes_at_rest(rest_pose) + + # A partial pose and velocity write reaches only the selected root. + initial_pose = cube_object.data.root_link_pose_w.torch.clone() + initial_velocity = cube_object.data.root_link_vel_w.torch.clone() + initial_com_velocity = cube_object.data.root_com_vel_w.torch.clone() + target_pose = initial_pose[1:].clone() + target_pose[:, :3] += torch.tensor([0.25, -0.1, 0.3], device=device) + target_pose[:, 3:] = torch.tensor(_yaw_quat(0.4), device=device) + target_velocity = torch.tensor([[0.0, 0.2, 0.0, 0.0, 0.0, 0.1]], device=device) + cube_object.write_root_link_pose_to_sim_index(root_pose=target_pose, env_ids=[1]) + cube_object.write_root_link_velocity_to_sim_index(root_velocity=target_velocity, env_ids=[1]) + torch.testing.assert_close(cube_object.data.root_link_pose_w.torch[1:], target_pose) + torch.testing.assert_close(cube_object.data.root_link_vel_w.torch[1:], target_velocity) + torch.testing.assert_close(cube_object.data.root_link_pose_w.torch[:1], initial_pose[:1]) + torch.testing.assert_close(cube_object.data.root_link_vel_w.torch[:1], initial_velocity[:1]) + raw_pose = wp.to_torch(cube_object.root_view.get_transforms()).to(device) + torch.testing.assert_close(raw_pose, torch.cat((initial_pose[:1], target_pose))) + # PhysX holds the center-of-mass velocity: the link velocity plus the angular velocity crossed with the + # world-frame center-of-mass offset. + com_offset_w = quat_apply(target_pose[:, 3:], cube_object.data.body_com_pos_b.torch[1:, 0]) + target_com_velocity = target_velocity.clone() + target_com_velocity[:, :3] += torch.linalg.cross(target_velocity[:, 3:], com_offset_w) + raw_velocity = wp.to_torch(cube_object.root_view.get_velocities()).to(device) + torch.testing.assert_close(raw_velocity, torch.cat((initial_com_velocity[:1], target_com_velocity))) + + # With a center-of-mass offset, every writer round-trips through its own frame. + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(_NUM_ENVS, 1) + coms = cube_object.data.body_com_pose_b.torch.clone() + coms[:, 0, :3] = offset + coms[:, 0, 3:] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) + cube_object.set_coms_index(coms=coms) + # random pose and velocity so frame conversions see a non-trivial rotation + rand_state = torch.rand(_NUM_ENVS, 13, device=device) + rand_state[..., :3] += rest_pose[:, :3] + rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) + env_idx = torch.arange(_NUM_ENVS, dtype=torch.int32, device=device) + for state_location in ("com", "link", "root"): + for env_ids in (None, env_idx): if state_location == "com": cube_object.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_ids) cube_object.write_root_com_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_ids) - elif state_location == "link": - cube_object.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_ids) - cube_object.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_ids) - elif state_location == "root": - cube_object.write_root_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_ids) - cube_object.write_root_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_ids) - - if state_location == "com": torch.testing.assert_close(rand_state[..., :7], cube_object.data.root_com_pose_w.torch) torch.testing.assert_close(rand_state[..., 7:], cube_object.data.root_com_vel_w.torch) elif state_location == "link": + cube_object.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_ids) + cube_object.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_ids) torch.testing.assert_close(rand_state[..., :7], cube_object.data.root_link_pose_w.torch) torch.testing.assert_close(rand_state[..., 7:], cube_object.data.root_link_vel_w.torch) - elif state_location == "root": + else: + cube_object.write_root_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_ids) + cube_object.write_root_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_ids) torch.testing.assert_close(rand_state[..., :7], cube_object.data.root_link_pose_w.torch) torch.testing.assert_close(rand_state[..., 7:], cube_object.data.root_com_vel_w.torch) @@ -789,39 +498,339 @@ def test_write_root_state(num_cubes, device, with_offset, state_location): torch.testing.assert_close( cube_object.data.root_com_vel_w.torch[:, 3:], cube_object.data.root_link_vel_w.torch[:, 3:] ) + # Move away from the written state so that the next pass writes it again. + scene.step() + + # Spinning about the center of mass keeps it in place while the link frame orbits it. + scene.place_cubes_at_rest(rest_pose) + spin_twist = torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0, 1.3], device=device).repeat(_NUM_ENVS, 1) + initial_com_pos = cube_object.data.root_com_pose_w.torch[:, :3].clone() + for _ in range(10): + cube_object.write_root_com_velocity_to_sim_index(root_velocity=spin_twist) + scene.step() + root_link_pose_w = cube_object.data.root_link_pose_w.torch + root_link_vel_w = cube_object.data.root_link_vel_w.torch + root_com_pose_w = cube_object.data.root_com_pose_w.torch + root_com_vel_w = cube_object.data.root_com_vel_w.torch + body_link_pose_w = cube_object.data.body_link_pose_w.torch + body_com_pose_w = cube_object.data.body_com_pose_w.torch + # center of mass position will be constant (i.e. spinning around com) + torch.testing.assert_close(root_com_pose_w[:, :3], initial_com_pos) + torch.testing.assert_close(body_com_pose_w[:, 0, :3], initial_com_pos) + # link position will be moving but should stay constant away from center of mass + torch.testing.assert_close( + -offset, quat_apply_inverse(root_link_pose_w[:, 3:], root_link_pose_w[:, :3] - root_com_pose_w[:, :3]) + ) + # orientation of com will be a constant rotation from link orientation + com_quat_w = quat_mul(body_link_pose_w[..., 3:], cube_object.data.body_com_quat_b.torch) + torch.testing.assert_close(com_quat_w, body_com_pose_w[..., 3:]) + torch.testing.assert_close(com_quat_w[:, 0], root_com_pose_w[:, 3:]) + # center of mass is at rest while the link frame moves with angular velocity cross offset + torch.testing.assert_close(torch.zeros_like(root_com_vel_w[:, :3]), root_com_vel_w[:, :3]) + lin_vel_rel_root = quat_apply_inverse(root_link_pose_w[:, 3:], root_link_vel_w[:, :3]) + lin_vel_rel_gt = torch.linalg.cross(spin_twist[:, 3:], -offset) + torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_root, atol=1e-4, rtol=1e-4) + # ang_vel will always match + torch.testing.assert_close(root_com_vel_w[:, 3:], root_link_vel_w[:, 3:]) @pytest.mark.isaacsim_ci -@pytest.mark.parametrize("device", test_devices()) -def test_warmup_loads_physics_once(device): - """Attach on GPU or force-load on CPU, without destroying and rebuilding native objects.""" - from unittest.mock import MagicMock, patch - - import omni.kit.app - import omni.physx +def test_rigid_object_wrench_delivery_and_reset(scene: _Scene) -> None: + """External wrenches act in the frame they are given in on the selected root; reset clears them.""" + device = scene.device + cube_object = scene.cubes + rest_pose = torch.cat( + (scene.origins["cubes"], torch.tensor(_yaw_quat(0.5 * math.pi), device=device).expand(_NUM_ENVS, 4)), -1 + ) + rest_pose[:, 2] += 1.0 + # Give both environments the inertial properties of a uniform 1 kg cube so that their responses are comparable. + cube_object.set_masses_index(masses=torch.ones((_NUM_ENVS, 1), device=device)) + coms = torch.zeros((_NUM_ENVS, 1, 7), device=device) + coms[..., 6] = 1.0 + cube_object.set_coms_index(coms=coms) + inertias = torch.zeros((_NUM_ENVS, 1, 9), device=device) + inertias[..., [0, 4, 8]] = 1.0 * (0.2**2 + 0.2**2) / 12.0 + cube_object.set_inertias_index(inertias=inertias) + + # A body-frame force on environment 1 accelerates only that cube, along the rotated force direction. + scene.place_cubes_at_rest(rest_pose) + cube_object.permanent_wrench_composer.set_forces_and_torques_index( + forces=torch.tensor([[[6.0, 0.0, 0.0]]], device=device), env_ids=[1] + ) + scene.step() + velocity = cube_object.data.root_com_lin_vel_w.torch + assert velocity[1, 1] > 1e-2, velocity + torch.testing.assert_close(velocity[0], torch.zeros(3, device=device), atol=1e-5, rtol=0.0) + + # A world-frame force, and a force at a position, match their body-frame counterparts. + for local_wrench, global_wrench, response in ( + ({"forces": [[[6.0, 0.0, 0.0]]]}, {"forces": [[[0.0, 6.0, 0.0]]]}, lambda: cube_object.data.root_com_lin_vel_w), + ( + {"forces": [[[0.0, 0.0, 6.0]]], "positions": [[[0.0, 0.1, 0.0]]]}, + {"forces": [[[0.0, 0.0, 6.0]]], "positions": [[[-0.1, 0.0, 0.0]]]}, + lambda: cube_object.data.root_com_ang_vel_b, + ), + ): + scene.place_cubes_at_rest(rest_pose) + local_wrench = {key: torch.tensor(value, device=device) for key, value in local_wrench.items()} + global_wrench = {key: torch.tensor(value, device=device) for key, value in global_wrench.items()} + if "positions" in global_wrench: + # World-frame positions add the rotated body-frame lever arm to the center of mass. + global_wrench["positions"] = global_wrench["positions"] + cube_object.data.body_com_pos_w.torch[:1] + cube_object.permanent_wrench_composer.set_forces_and_torques_index(env_ids=[1], **local_wrench) + cube_object.permanent_wrench_composer.set_forces_and_torques_index(env_ids=[0], is_global=True, **global_wrench) + scene.step() + response_value = response().torch + torch.testing.assert_close(response_value[0], response_value[1], atol=1e-4, rtol=1e-3) + # An upward force 0.1 m along the body y-axis rolls the cube about its x-axis. + assert torch.all(cube_object.data.root_com_ang_vel_b.torch[:, 0] > 0.1) + assert torch.all( + quat_apply( + cube_object.data.root_link_quat_w.torch, torch.tensor([1.0, 0.0, 0.0], device=device).repeat(_NUM_ENVS, 1) + )[:, 1] + > 0.99 + ) - with build_simulation_context(device=device, add_ground_plane=True, dt=0.01, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - cube, _ = generate_cubes_scene(num_cubes=1, height=1.0, device=device) - - # Wrap read-only native interfaces through their accessors; real calls still execute. - physx_spy = MagicMock(wraps=omni.physx.get_physx_interface()) - physx_sim_spy = MagicMock(wraps=omni.physx.get_physx_simulation_interface()) - with ( - patch("omni.physx.get_physx_interface", return_value=physx_spy), - patch("omni.physx.get_physx_simulation_interface", return_value=physx_sim_spy), - ): - sim.reset() + # A partial reset clears only the selected environment; a full reset deactivates the composers. + composers = (cube_object.instantaneous_wrench_composer, cube_object.permanent_wrench_composer) + ones = torch.ones((_NUM_ENVS, 1, 3), device=device) + cube_object.permanent_wrench_composer.set_forces_and_torques_index(forces=ones, torques=ones) + cube_object.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones) + cube_object.reset(env_ids=torch.tensor([0], device=device)) + for composer in composers: + assert composer.active + for buffer in (composer.out_force_b.torch, composer.out_torque_b.torch): + assert torch.count_nonzero(buffer[0]) == 0 + assert torch.count_nonzero(buffer[1:]) == buffer[1:].numel() + cube_object.reset() + for composer in composers: + assert not composer.active + assert torch.count_nonzero(composer.out_force_b.torch) == 0 + assert torch.count_nonzero(composer.out_torque_b.torch) == 0 + + +def test_collection_state_writes(scene: _Scene) -> None: + """Body state writes reach the selected body-major view entries in the frame they are given in.""" + device = scene.device + collection = scene.collection + scene.place_collection_at_rest() + + # A non-sorted partial write reaches only the selected view entries. + env_ids = torch.tensor([1, 0], dtype=torch.int32, device=device) + body_ids = torch.tensor([2], dtype=torch.int32, device=device) + initial_pose = collection.data.body_link_pose_w.torch.clone() + initial_velocity = collection.data.body_com_vel_w.torch.clone() + target_pose = initial_pose[env_ids][:, body_ids].clone() + target_pose[0, 0, :3] += torch.tensor([0.2, 0.3, 0.4], device=device) + target_pose[1, 0, :3] += torch.tensor([-0.1, -0.2, 0.1], device=device) + target_pose[..., 3:] = torch.tensor(_yaw_quat(0.3), device=device) + target_velocity = torch.tensor([[[0.1, 0.2, 0.3, 0.4, 0.5, 0.6]], [[-0.6, -0.5, -0.4, -0.3, -0.2, -0.1]]]) + target_velocity = target_velocity.to(device) + collection.write_body_link_pose_to_sim_index(body_poses=target_pose, env_ids=env_ids, body_ids=body_ids) + collection.write_body_com_velocity_to_sim_index(body_velocities=target_velocity, env_ids=env_ids, body_ids=body_ids) + expected_pose = initial_pose.clone() + expected_pose[env_ids[:, None], body_ids] = target_pose + expected_velocity = initial_velocity.clone() + expected_velocity[env_ids[:, None], body_ids] = target_velocity + torch.testing.assert_close(collection.data.body_link_pose_w.torch, expected_pose) + torch.testing.assert_close(collection.data.body_com_vel_w.torch, expected_velocity) + for raw, expected in ( + (collection.root_view.get_transforms(), expected_pose), + (collection.root_view.get_velocities(), expected_velocity), + ): + raw_by_env = wp.to_torch(raw).to(device).reshape(_NUM_CUBES, _NUM_ENVS, -1).transpose(0, 1) + torch.testing.assert_close(raw_by_env, expected) + + # With a center-of-mass offset, every writer round-trips through its own frame. + scene.place_collection_at_rest() + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(_NUM_ENVS, _NUM_CUBES, 1) + coms = collection.data.body_com_pose_b.torch.clone() + coms[..., :3] = offset + coms[..., 3:] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) + collection.set_coms_index(coms=coms) + # random pose and velocity so frame conversions see a non-trivial rotation + rand_state = torch.rand(_NUM_ENVS, _NUM_CUBES, 13, device=device) + rand_state[..., :3] += collection.data.body_link_pos_w.torch + rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) + all_ids = {"env_ids": torch.arange(_NUM_ENVS, device=device), "body_ids": torch.arange(_NUM_CUBES, device=device)} + for state_location in ("com", "link", "root"): + for ids in ({}, all_ids): + if state_location == "com": + collection.write_body_com_pose_to_sim_index(body_poses=rand_state[..., :7], **ids) + collection.write_body_com_velocity_to_sim_index(body_velocities=rand_state[..., 7:], **ids) + torch.testing.assert_close(rand_state[..., :7], collection.data.body_com_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], collection.data.body_com_vel_w.torch) + elif state_location == "link": + collection.write_body_link_pose_to_sim_index(body_poses=rand_state[..., :7], **ids) + collection.write_body_link_velocity_to_sim_index(body_velocities=rand_state[..., 7:], **ids) + torch.testing.assert_close(rand_state[..., :7], collection.data.body_link_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], collection.data.body_link_vel_w.torch) + # PhysX holds the center-of-mass velocity: the link velocity plus the angular velocity crossed + # with the world-frame center-of-mass offset. + com_offset_w = quat_apply(rand_state[..., 3:7], offset) + expected_com_velocity = rand_state[..., 7:].clone() + expected_com_velocity[..., :3] += torch.linalg.cross(rand_state[..., 10:], com_offset_w, dim=-1) + raw_velocity = wp.to_torch(collection.root_view.get_velocities()).to(device) + raw_velocity = raw_velocity.reshape(_NUM_CUBES, _NUM_ENVS, 6).transpose(0, 1) + torch.testing.assert_close(raw_velocity, expected_com_velocity) + else: + collection.write_body_link_pose_to_sim_index(body_poses=rand_state[..., :7], **ids) + collection.write_body_com_velocity_to_sim_index(body_velocities=rand_state[..., 7:], **ids) + torch.testing.assert_close(rand_state[..., :7], collection.data.body_link_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], collection.data.body_com_vel_w.torch) - extension_manager = omni.kit.app.get_app().get_extension_manager() - assert extension_manager.is_extension_enabled("omni.physics.physx"), ( - "The omni.physics.physx bridge must register PhysX with the unified physics API." + # the frame that was not written must follow through the center-of-mass offset + link_pose_w = collection.data.body_link_pose_w.torch + body_com_pose_b = collection.data.body_com_pose_b.torch + expected_com_pos, expected_com_quat = combine_frame_transforms( + link_pose_w[..., :3].reshape(-1, 3), + link_pose_w[..., 3:].reshape(-1, 4), + body_com_pose_b[..., :3].reshape(-1, 3), + body_com_pose_b[..., 3:].reshape(-1, 4), + ) + expected_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1).view(_NUM_ENVS, _NUM_CUBES, 7) + torch.testing.assert_close(expected_com_pose, collection.data.body_com_pose_w.torch) + torch.testing.assert_close( + collection.data.body_com_vel_w.torch[..., 3:], collection.data.body_link_vel_w.torch[..., 3:] + ) + # Move away from the written state so that the next pass writes it again. + scene.step() + + # Spinning about the center of mass keeps it in place while the link frame orbits it. + scene.place_collection_at_rest() + spin_twist = torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0, 1.3], device=device).repeat(_NUM_ENVS, _NUM_CUBES, 1) + init_com = collection.data.body_com_pose_w.torch[..., :3].clone() + for _ in range(10): + collection.write_body_com_velocity_to_sim_index(body_velocities=spin_twist) + scene.step() + object_link_pose_w = collection.data.body_link_pose_w.torch + object_link_vel_w = collection.data.body_link_vel_w.torch + object_com_pose_w = collection.data.body_com_pose_w.torch + object_com_vel_w = collection.data.body_com_vel_w.torch + # center of mass position will be constant (i.e. spinning around com) + torch.testing.assert_close(init_com, object_com_pose_w[..., :3]) + # link position will be moving but should stay constant away from center of mass + object_link_state_pos_rel_com = quat_apply_inverse( + object_link_pose_w[..., 3:], object_link_pose_w[..., :3] - object_com_pose_w[..., :3] ) - assert physx_sim_spy.attach_stage.call_count == int(device.startswith("cuda")) - assert physx_spy.force_load_physics_from_usd.call_count == int(device == "cpu") - assert cube.is_initialized and cube.num_instances == 1 - initial = cube.data.root_pos_w.torch.clone() - for _ in range(10): - sim.step(render=False) - cube.update(sim.get_physics_dt()) - assert torch.all(cube.data.root_pos_w.torch[:, 2] < initial[:, 2]) + torch.testing.assert_close(-offset, object_link_state_pos_rel_com) + # orientation of com will be a constant rotation from link orientation + com_quat_w = quat_mul(object_link_pose_w[..., 3:], collection.data.body_com_quat_b.torch) + torch.testing.assert_close(com_quat_w, object_com_pose_w[..., 3:]) + # center of mass is at rest while the link frame moves with angular velocity cross offset + torch.testing.assert_close(torch.zeros_like(object_com_vel_w[..., :3]), object_com_vel_w[..., :3]) + lin_vel_rel_object_gt = quat_apply_inverse(object_link_pose_w[..., 3:], object_link_vel_w[..., :3]) + lin_vel_rel_gt = torch.linalg.cross(spin_twist[..., 3:], -offset) + torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_object_gt, atol=1e-4, rtol=1e-3) + # ang_vel will always match + torch.testing.assert_close(object_com_vel_w[..., 3:], object_link_vel_w[..., 3:]) + + +def test_collection_inertial_properties_reach_selected_view_entries(scene: _Scene) -> None: + """Mass, center-of-mass, and inertia writes to a non-sorted subset reach only the selected body-major entries.""" + device = scene.device + collection = scene.collection + env_ids = torch.tensor([1, 0], dtype=torch.int32, device=device) + body_ids = torch.tensor([1], dtype=torch.int32, device=device) + initial = { + "masses": collection.data.body_mass.torch.clone(), + "coms": collection.data.body_com_pose_b.torch.clone(), + "inertias": collection.data.body_inertia.torch.clone(), + } + values = {name: value[env_ids][:, body_ids].clone() for name, value in initial.items()} + values["masses"][:, 0] = torch.tensor([5.0, 7.0], device=device) + values["coms"][:, 0, :3] = torch.tensor([[0.02, 0.03, 0.04], [-0.01, 0.01, 0.02]], device=device) + values["inertias"][0, 0, [0, 4, 8]] *= 1.5 + values["inertias"][1, 0, [0, 4, 8]] *= 2.0 + collection.set_masses_index(masses=values["masses"], env_ids=env_ids, body_ids=body_ids) + collection.set_coms_index(coms=values["coms"], env_ids=env_ids, body_ids=body_ids) + collection.set_inertias_index(inertias=values["inertias"], env_ids=env_ids, body_ids=body_ids) + for name, data, raw in ( + ("masses", collection.data.body_mass, collection.root_view.get_masses()), + ("coms", collection.data.body_com_pose_b, collection.root_view.get_coms().view(wp.float32)), + ("inertias", collection.data.body_inertia, collection.root_view.get_inertias()), + ): + expected = initial[name].clone() + expected[env_ids[:, None], body_ids] = values[name] + torch.testing.assert_close(data.torch, expected) + # The view is body-major: (num_cubes * num_envs, ...). + raw = wp.to_torch(raw).to(device).reshape(_NUM_CUBES, _NUM_ENVS, -1).transpose(0, 1) + torch.testing.assert_close(raw, expected.reshape(_NUM_ENVS, _NUM_CUBES, -1)) + # Restore the initial properties for the wrench checks. + collection.set_masses_index(masses=initial["masses"]) + collection.set_coms_index(coms=initial["coms"]) + collection.set_inertias_index(inertias=initial["inertias"]) + + +def test_collection_wrench_delivery_and_reset(scene: _Scene) -> None: + """External wrenches act on the selected bodies in the frame they are given in; reset clears them.""" + device = scene.device + collection = scene.collection + # Give every body a 1 kg mass and the center of mass and inertia of the first body so that their responses are + # comparable. + collection.set_masses_index(masses=torch.ones((_NUM_ENVS, _NUM_CUBES), device=device)) + coms = collection.data.body_com_pose_b.torch[:1, :1] + collection.set_coms_index(coms=coms.expand(_NUM_ENVS, _NUM_CUBES, 7).contiguous()) + inertias = collection.data.body_inertia.torch[:1, :1] + collection.set_inertias_index(inertias=inertias.expand(_NUM_ENVS, _NUM_CUBES, 9).contiguous()) + + # A body-frame force on bodies 0 and 2 accelerates only those bodies, along the rotated force direction; the + # same force given in the world frame produces the same response. + object_ids, _ = collection.find_bodies(".*") + for local_wrench, global_wrench, response, is_linear in ( + ( + {"forces": [[6.0, 0.0, 0.0]]}, + {"forces": [[0.0, 6.0, 0.0]]}, + lambda: collection.data.body_com_lin_vel_w, + True, + ), + ( + {"forces": [[0.0, 0.0, 6.0]], "positions": [[0.0, 0.1, 0.0]]}, + {"forces": [[0.0, 0.0, 6.0]], "positions": [[-0.1, 0.0, 0.0]]}, + lambda: collection.data.body_com_ang_vel_b, + False, + ), + ): + scene.place_collection_at_rest(yaw=0.5 * math.pi) + for env_id, wrench, is_global in ((1, local_wrench, False), (0, global_wrench, True)): + wrench = {key: torch.tensor(value, device=device).expand(1, 2, 3).clone() for key, value in wrench.items()} + if is_global and "positions" in wrench: + # World-frame positions add the rotated body-frame lever arm to the center of mass. + wrench["positions"] += collection.data.body_com_pos_w.torch[env_id : env_id + 1, 0::2] + collection.permanent_wrench_composer.set_forces_and_torques_index( + body_ids=object_ids[0::2], env_ids=[env_id], is_global=is_global, **wrench + ) + scene.step() + response_value = response().torch + torch.testing.assert_close(response_value[0], response_value[1], atol=1e-4, rtol=1e-3) + torch.testing.assert_close(response_value[:, 1], torch.zeros_like(response_value[:, 1]), atol=1e-5, rtol=0) + if is_linear: + assert torch.all(response_value[:, 0::2, 1] > 1e-2), response_value + # The permanent force persists across steps, so a second step doubles the velocity. + first_step_value = response_value.clone() + scene.step() + torch.testing.assert_close(response().torch, 2.0 * first_step_value, atol=1e-4, rtol=1e-3) + # An upward force 0.1 m along the body y-axis rolls the cube about its x-axis. + assert torch.all(collection.data.body_com_ang_vel_b.torch[:, 0::2, 0] > 0.1) + heading = quat_apply( + collection.data.body_link_quat_w.torch, torch.tensor([1.0, 0.0, 0.0], device=device).expand(2, 3, 3) + ) + assert torch.all(heading[..., 1] > 0.99) + + # A partial reset clears the external wrenches of the selected environments only. + composers = (collection.instantaneous_wrench_composer, collection.permanent_wrench_composer) + ones = torch.ones((_NUM_ENVS, _NUM_CUBES, 3), device=device) + collection.permanent_wrench_composer.set_forces_and_torques_index(forces=ones, torques=ones) + collection.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones) + collection.reset(env_ids=torch.tensor([0], device=device)) + for composer in composers: + assert composer.active + for buffer in (composer.out_force_b.torch, composer.out_torque_b.torch): + assert torch.count_nonzero(buffer[0]) == 0 + assert torch.count_nonzero(buffer[1:]) == buffer[1:].numel() + + # A full reset clears every environment + collection.reset() + for composer in composers: + assert torch.count_nonzero(composer.out_force_b.torch) == 0 + assert torch.count_nonzero(composer.out_torque_b.torch) == 0 diff --git a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py deleted file mode 100644 index 2e0af9a4be8b..000000000000 --- a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py +++ /dev/null @@ -1,650 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -# ignore private usage of variables warning -# pyright: reportPrivateUsage=none - -from isaaclab.test.utils import launch_test_simulation, test_devices - -launch_test_simulation() - -import sys - -import pytest -import torch -import warp as wp -from isaaclab_physx.assets import RigidObjectCollection - -import isaaclab.sim as sim_utils -from isaaclab.assets import RigidObjectCfg, RigidObjectCollectionCfg -from isaaclab.sim import build_simulation_context -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR -from isaaclab.utils.math import ( - combine_frame_transforms, - default_orientation, - quat_apply_inverse, - quat_mul, - random_orientation, -) - - -def generate_cubes_scene( - num_envs: int = 1, - num_cubes: int = 1, - height=1.0, - has_api: bool = True, - kinematic_enabled: bool = False, - device: str = "cuda:0", -) -> tuple[RigidObjectCollection, torch.Tensor]: - """Generate a scene with the provided number of cubes. - - Args: - num_envs: Number of envs to generate. - num_cubes: Number of cubes to generate. - height: Height of the cubes. - has_api: Whether the cubes have a rigid body API on them. - kinematic_enabled: Whether the cubes are kinematic. - device: Device to use for the simulation. - - Returns: - A tuple containing the rigid object representing the cubes and the origins of the cubes. - - """ - origins = torch.tensor([(i * 3.0, 0, height) for i in range(num_envs)]).to(device) - # Create Top-level Xforms, one for each cube - for i, origin in enumerate(origins): - sim_utils.create_prim(f"/World/Table_{i}", "Xform", translation=origin) - - # Resolve spawn configuration - if has_api: - spawn_cfg = sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", - rigid_props=sim_utils.UsdPhysicsRigidBodyCfg(kinematic_enabled=kinematic_enabled), - ) - else: - # since no rigid body properties defined, this is just a static collider - spawn_cfg = sim_utils.CuboidCfg( - size=(0.1, 0.1, 0.1), - collision_props=sim_utils.UsdPhysicsCollisionCfg(), - ) - - # create the rigid object configs - cube_config_dict = {} - for i in range(num_cubes): - cube_object_cfg = RigidObjectCfg( - prim_path=f"/World/Table_[^/]*/Object_{i}", - spawn=spawn_cfg, - init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 3 * i, height)), - ) - cube_config_dict[f"cube_{i}"] = cube_object_cfg - # create the rigid object collection - cube_object_collection_cfg = RigidObjectCollectionCfg(rigid_objects=cube_config_dict) - cube_object_colection = RigidObjectCollection(cfg=cube_object_collection_cfg) - - return cube_object_colection, origins - - -@pytest.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 - with build_simulation_context(device=device, auto_add_lighting=True, gravity_enabled=gravity_enabled) as sim: - sim._app_control_on_stop_handle = None - yield sim - - -@pytest.mark.parametrize(("num_envs", "num_cubes"), [(1, 1), (2, 3)]) -@pytest.mark.parametrize("device", test_devices()) -def test_initialization(sim, num_envs, num_cubes, device): - """Test initialization for prim with rigid body API at the provided prim path.""" - object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(object_collection) < 10 - - # Play sim - sim.reset() - - # Check if object is initialized - assert object_collection.is_initialized - assert len(object_collection.body_names) == num_cubes - - # Check buffers that exist and have correct shapes - assert object_collection.data.body_link_pos_w.torch.shape == (num_envs, num_cubes, 3) - assert object_collection.data.body_link_quat_w.torch.shape == (num_envs, num_cubes, 4) - assert object_collection.data.body_mass.torch.shape == (num_envs, num_cubes) - assert object_collection.data.body_inertia.torch.shape == (num_envs, num_cubes, 9) - - # Simulate physics - for _ in range(2): - sim.step() - object_collection.update(sim.cfg.dt) - - -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("gravity_enabled", [False]) -def test_subset_write_reaches_selected_view_entry(sim, device, gravity_enabled): - """A write to one (env, body) cell must move only that object in the simulation.""" - object_collection, _ = generate_cubes_scene(num_envs=2, num_cubes=3, device=device) - - # Play sim - sim.reset() - sim.step() - object_collection.update(sim.cfg.dt) - - initial_pose = object_collection.data.body_link_pose_w.torch.clone() - new_pose = initial_pose[1:2, 2:3].clone() - new_pose[..., 2] += 0.5 - object_collection.write_body_link_pose_to_sim_index(body_poses=new_pose, env_ids=[1], body_ids=[2]) - - # Read the pose back from the simulation so that a wrong view index cannot hide in the data buffer - sim.step() - object_collection.update(sim.cfg.dt) - - expected_pose = initial_pose.clone() - expected_pose[1, 2] = new_pose[0, 0] - torch.testing.assert_close(object_collection.data.body_link_pose_w.torch, expected_pose) - - -@pytest.mark.parametrize("num_envs", [2]) -@pytest.mark.parametrize("num_cubes", [3]) -@pytest.mark.parametrize("device", test_devices()) -def test_initialization_with_kinematic_enabled(sim, num_envs, num_cubes, device): - """Test that initialization for prim with kinematic flag enabled.""" - object_collection, origins = generate_cubes_scene( - num_envs=num_envs, num_cubes=num_cubes, kinematic_enabled=True, device=device - ) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(object_collection) < 10 - - # Play sim - sim.reset() - - # Check if object is initialized - assert object_collection.is_initialized - assert len(object_collection.body_names) == num_cubes - - # Check buffers that exist and have correct shapes - assert object_collection.data.body_link_pos_w.torch.shape == (num_envs, num_cubes, 3) - assert object_collection.data.body_link_quat_w.torch.shape == (num_envs, num_cubes, 4) - - # Simulate physics - for _ in range(2): - sim.step() - object_collection.update(sim.cfg.dt) - # check that the object is kinematic - default_body_pose = object_collection.data.default_body_pose.torch.clone() - default_body_vel = object_collection.data.default_body_vel.torch.clone() - default_body_pose[..., :3] += origins.unsqueeze(1) - torch.testing.assert_close(object_collection.data.body_link_pose_w.torch, default_body_pose) - torch.testing.assert_close(object_collection.data.body_link_vel_w.torch, default_body_vel) - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_initialization_with_no_rigid_body(sim, num_cubes, device): - """Test that initialization fails when no rigid body is found at the provided prim path.""" - object_collection, _ = generate_cubes_scene(num_cubes=num_cubes, has_api=False, device=device) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(object_collection) < 10 - - # Play sim - with pytest.raises(RuntimeError): - sim.reset() - - -@pytest.mark.parametrize("num_envs", [2]) -@pytest.mark.parametrize("num_cubes", [4]) -@pytest.mark.parametrize("device", test_devices()) -def test_external_force_on_single_body(sim, num_envs, num_cubes, device): - """Test application of external force on the base of the object.""" - object_collection, origins = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) - sim.reset() - - # find objects to apply the force - object_ids, object_names = object_collection.find_bodies(".*") - - # Sample a force equal to the weight of the object - external_wrench_b = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=sim.device) - # Every 2nd cube should have a force applied to it - external_wrench_b[:, 0::2, 2] = 9.81 * object_collection.data.body_mass.torch[:, 0::2] - - for i in range(5): - # reset object state - body_pose = object_collection.data.default_body_pose.torch.clone() - body_vel = object_collection.data.default_body_vel.torch.clone() - # need to shift the position of the cubes otherwise they will be on top of each other - body_pose[..., :2] += origins.unsqueeze(1)[..., :2] - object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) - object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) - # reset object - object_collection.reset() - - is_global = False - if i % 2 == 0: - positions = object_collection.data.body_link_pos_w.torch[:, object_ids, :3] - is_global = True - else: - positions = None - - # apply force - object_collection.permanent_wrench_composer.set_forces_and_torques_index( - forces=external_wrench_b[..., :3], - torques=external_wrench_b[..., 3:], - positions=positions, - body_ids=object_ids, - env_ids=None, - is_global=is_global, - ) - for _ in range(10): - # write data to sim - object_collection.write_data_to_sim() - # step sim - sim.step() - # update object collection - object_collection.update(sim.cfg.dt) - - # First object should still be at the same Z position (1.0) - torch.testing.assert_close( - object_collection.data.body_link_pos_w.torch[:, 0::2, 2], - torch.ones_like(object_collection.data.body_link_pos_w.torch[:, 0::2, 2]), - ) - # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(object_collection.data.body_link_pos_w.torch[:, 1::2, 2] < 1.0) - - -@pytest.mark.parametrize("num_envs", [2]) -@pytest.mark.parametrize("num_cubes", [4]) -@pytest.mark.parametrize("device", test_devices()) -def test_external_force_on_single_body_at_position(sim, num_envs, num_cubes, device): - """Test application of external force on the base of the object at a specific position. - - In this test, we apply a force equal to the weight of an object on the base of - one of the objects at 1m in the Y direction, we check that the object rotates around it's X axis. - For the other object, we do not apply any force and check that it falls down. - """ - object_collection, origins = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) - sim.reset() - - # find objects to apply the force - object_ids, object_names = object_collection.find_bodies(".*") - - # Sample a force equal to the weight of the object - external_wrench_b = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=sim.device) - external_wrench_positions_b = torch.zeros(object_collection.num_instances, len(object_ids), 3, device=sim.device) - # Every 2nd cube should have a force applied to it - external_wrench_b[:, 0::2, 2] = 500.0 - external_wrench_positions_b[:, 0::2, 1] = 1.0 - - # Desired force and torque - for i in range(5): - # reset object state - body_pose = object_collection.data.default_body_pose.torch.clone() - body_vel = object_collection.data.default_body_vel.torch.clone() - # need to shift the position of the cubes otherwise they will be on top of each other - body_pose[..., :2] += origins.unsqueeze(1)[..., :2] - object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) - object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) - # reset object - object_collection.reset() - - is_global = False - if i % 2 == 0: - body_com_pos_w = object_collection.data.body_link_pos_w.torch[:, object_ids, :3] - external_wrench_positions_b[..., 0] = 0.0 - external_wrench_positions_b[..., 1] = 1.0 - external_wrench_positions_b[..., 2] = 0.0 - external_wrench_positions_b += body_com_pos_w - is_global = True - else: - external_wrench_positions_b[..., 0] = 0.0 - external_wrench_positions_b[..., 1] = 1.0 - external_wrench_positions_b[..., 2] = 0.0 - - # apply force - object_collection.permanent_wrench_composer.set_forces_and_torques_index( - forces=external_wrench_b[..., :3], - torques=external_wrench_b[..., 3:], - positions=external_wrench_positions_b, - body_ids=object_ids, - env_ids=None, - is_global=is_global, - ) - object_collection.permanent_wrench_composer.add_forces_and_torques_index( - forces=external_wrench_b[..., :3], - torques=external_wrench_b[..., 3:], - positions=external_wrench_positions_b, - body_ids=object_ids, - is_global=is_global, - ) - - for _ in range(10): - # write data to sim - object_collection.write_data_to_sim() - # step sim - sim.step() - # update object collection - object_collection.update(sim.cfg.dt) - - # First object should be rotating around it's X axis - assert torch.all(object_collection.data.body_com_ang_vel_b.torch[:, 0::2, 0] > 0.1) - # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(object_collection.data.body_link_pos_w.torch[:, 1::2, 2] < 1.0) - - -@pytest.mark.parametrize("num_envs", [3]) -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("gravity_enabled", [False]) -def test_set_object_state(sim, num_envs, num_cubes, device, gravity_enabled): - """Test setting the state of the object. - - .. note:: - Turn off gravity for this test as we don't want any external forces acting on the object - to ensure state remains static - """ - object_collection, origins = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) - sim.reset() - - state_types = ["body_link_pos_w", "body_link_quat_w", "body_com_lin_vel_w", "body_com_ang_vel_w"] - - # Set each state type individually as they are dependent on each other - for state_type_to_randomize in state_types: - state_dict = { - "body_link_pos_w": torch.zeros_like(object_collection.data.body_link_pos_w.torch, device=sim.device), - "body_link_quat_w": default_orientation(num=num_cubes * num_envs, device=sim.device).view( - num_envs, num_cubes, 4 - ), - "body_com_lin_vel_w": torch.zeros_like(object_collection.data.body_com_lin_vel_w.torch, device=sim.device), - "body_com_ang_vel_w": torch.zeros_like(object_collection.data.body_com_ang_vel_w.torch, device=sim.device), - } - - for _ in range(5): - # reset object - object_collection.reset() - - # Set random state - if state_type_to_randomize == "body_link_quat_w": - state_dict[state_type_to_randomize] = random_orientation( - num=num_cubes * num_envs, device=sim.device - ).view(num_envs, num_cubes, 4) - else: - state_dict[state_type_to_randomize] = torch.randn(num_envs, num_cubes, 3, device=sim.device) - # make sure objects do not overlap - if state_type_to_randomize == "body_link_pos_w": - state_dict[state_type_to_randomize][..., :2] += origins.unsqueeze(1)[..., :2] - - # perform simulation - for _ in range(5): - body_pose = torch.cat( - [state_dict["body_link_pos_w"], state_dict["body_link_quat_w"]], - dim=-1, - ) - body_vel = torch.cat( - [state_dict["body_com_lin_vel_w"], state_dict["body_com_ang_vel_w"]], - dim=-1, - ) - # reset object state - object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) - object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) - sim.step() - - # assert that set object quantities are equal to the ones set in the state_dict - for key, expected_value in state_dict.items(): - value = getattr(object_collection.data, key).torch - torch.testing.assert_close(value, expected_value, rtol=1e-5, atol=1e-5) - - object_collection.update(sim.cfg.dt) - - -@pytest.mark.parametrize("num_envs", [4]) -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("with_offset", [True, False]) -@pytest.mark.parametrize("gravity_enabled", [False]) -def test_object_state_properties(sim, num_envs, num_cubes, device, with_offset, gravity_enabled): - """Test the object_com_state_w and object_link_state_w properties.""" - cube_object, env_pos = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, height=0.0, device=device) - view_ids = torch.tensor([x for x in range(num_cubes * num_envs)], dtype=torch.int32) - - sim.reset() - - # check if cube_object is initialized - assert cube_object.is_initialized - - # change center of mass offset from link frame - offset = ( - torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) - if with_offset - else torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) - ) - - com = wp.to_torch(cube_object.reshape_view_to_data_2d(cube_object.root_view.get_coms().view(wp.transformf))) - com[..., :3] = offset.to("cpu") - cube_object.root_view.set_coms( - cube_object.reshape_data_to_view_2d(wp.from_torch(com.clone(), dtype=wp.transformf)).view(wp.float32), - wp.from_torch(view_ids, dtype=wp.int32), - ) - - # check center of mass has been set - torch.testing.assert_close( - wp.to_torch(cube_object.reshape_view_to_data_2d(cube_object.root_view.get_coms().view(wp.transformf))), com - ) - - # random z spin velocity - spin_twist = torch.zeros(6, device=device) - spin_twist[5] = torch.randn(1, device=device) - - # initial spawn point - init_com = cube_object.data.body_com_pose_w.torch[..., :3] - - for i in range(10): - # spin the object around Z axis (com) - cube_object.write_body_com_velocity_to_sim_index(body_velocities=spin_twist.repeat(num_envs, num_cubes, 1)) - sim.step() - cube_object.update(sim.cfg.dt) - - # get state properties - object_link_pose_w = cube_object.data.body_link_pose_w.torch - object_link_vel_w = cube_object.data.body_link_vel_w.torch - object_com_pose_w = cube_object.data.body_com_pose_w.torch - object_com_vel_w = cube_object.data.body_com_vel_w.torch - - # if offset is [0,0,0] all object_state_%_w will match and all body_%_w will match - if not with_offset: - torch.testing.assert_close(object_link_pose_w, object_com_pose_w) - torch.testing.assert_close(object_com_vel_w, object_link_vel_w) - else: - # cubes are spinning around center of mass - # position will not match - # center of mass position will be constant (i.e. spinning around com) - torch.testing.assert_close(init_com, object_com_pose_w[..., :3]) - - # link position will be moving but should stay constant away from center of mass - object_link_state_pos_rel_com = quat_apply_inverse( - object_link_pose_w[..., 3:], - object_link_pose_w[..., :3] - object_com_pose_w[..., :3], - ) - - torch.testing.assert_close(-offset, object_link_state_pos_rel_com) - - # orientation of com will be a constant rotation from link orientation - com_quat_b = cube_object.data.body_com_quat_b.torch - com_quat_w = quat_mul(object_link_pose_w[..., 3:], com_quat_b) - torch.testing.assert_close(com_quat_w, object_com_pose_w[..., 3:]) - - # lin_vel will not match - # center of mass vel will be constant (i.e. spinning around com) - torch.testing.assert_close( - torch.zeros_like(object_com_vel_w[..., :3]), - object_com_vel_w[..., :3], - ) - - # link frame will be moving, and should be equal to input angular velocity cross offset - lin_vel_rel_object_gt = quat_apply_inverse(object_link_pose_w[..., 3:], object_link_vel_w[..., :3]) - lin_vel_rel_gt = torch.linalg.cross(spin_twist.repeat(num_envs, num_cubes, 1)[..., 3:], -offset) - torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_object_gt, atol=1e-4, rtol=1e-3) - - # ang_vel will always match - torch.testing.assert_close(object_com_vel_w[..., 3:], object_link_vel_w[..., 3:]) - - -@pytest.mark.parametrize("num_envs", [3]) -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("with_offset", [True, False]) -@pytest.mark.parametrize("state_location", ["com", "link", "root"]) -@pytest.mark.parametrize("gravity_enabled", [False]) -def test_write_object_state(sim, num_envs, num_cubes, device, with_offset, state_location, gravity_enabled): - """Test the body state setters in the link, center-of-mass, and mixed root frames. - - Every write must round-trip through the written frame and keep the other frame consistent with the - center-of-mass offset. - """ - # Create a scene with random cubes - cube_object, env_pos = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, height=0.0, device=device) - view_ids = torch.tensor([x for x in range(num_cubes * num_envs)], dtype=torch.int32) - env_ids = torch.tensor([x for x in range(num_envs)], dtype=torch.int32) - object_ids = torch.tensor([x for x in range(num_cubes)], dtype=torch.int32) - - sim.reset() - - # Check if cube_object is initialized - assert cube_object.is_initialized - - # change center of mass offset from link frame - offset = ( - torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) - if with_offset - else torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) - ) - - com = wp.to_torch(cube_object.reshape_view_to_data_2d(cube_object.root_view.get_coms().view(wp.transformf))) - com[..., :3] = offset.to("cpu") - cube_object.root_view.set_coms( - cube_object.reshape_data_to_view_2d(wp.from_torch(com.clone(), dtype=wp.transformf)).view(wp.float32), - wp.from_torch(view_ids, dtype=wp.int32), - ) - # check center of mass has been set - torch.testing.assert_close( - wp.to_torch(cube_object.reshape_view_to_data_2d(cube_object.root_view.get_coms().view(wp.transformf))), com - ) - - # random pose and velocity so frame conversions see a non-trivial rotation - rand_state = torch.rand(num_envs, num_cubes, 13, device=device) - rand_state[..., :3] += cube_object.data.body_link_pos_w.torch - # make quaternion a unit vector - rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) - - env_ids = env_ids.to(device) - object_ids = object_ids.to(device) - for i in range(10): - sim.step() - cube_object.update(sim.cfg.dt) - - ids = {} if i % 2 == 0 else {"env_ids": env_ids, "body_ids": object_ids} - if state_location == "com": - cube_object.write_body_com_pose_to_sim_index(body_poses=rand_state[..., :7], **ids) - cube_object.write_body_com_velocity_to_sim_index(body_velocities=rand_state[..., 7:], **ids) - elif state_location == "link": - cube_object.write_body_link_pose_to_sim_index(body_poses=rand_state[..., :7], **ids) - cube_object.write_body_link_velocity_to_sim_index(body_velocities=rand_state[..., 7:], **ids) - elif state_location == "root": - cube_object.write_body_link_pose_to_sim_index(body_poses=rand_state[..., :7], **ids) - cube_object.write_body_com_velocity_to_sim_index(body_velocities=rand_state[..., 7:], **ids) - - if state_location == "com": - torch.testing.assert_close(rand_state[..., :7], cube_object.data.body_com_pose_w.torch) - torch.testing.assert_close(rand_state[..., 7:], cube_object.data.body_com_vel_w.torch) - elif state_location == "link": - torch.testing.assert_close(rand_state[..., :7], cube_object.data.body_link_pose_w.torch) - torch.testing.assert_close(rand_state[..., 7:], cube_object.data.body_link_vel_w.torch) - elif state_location == "root": - torch.testing.assert_close(rand_state[..., :7], cube_object.data.body_link_pose_w.torch) - torch.testing.assert_close(rand_state[..., 7:], cube_object.data.body_com_vel_w.torch) - - # the frame that was not written must follow through the center-of-mass offset - link_pose_w = cube_object.data.body_link_pose_w.torch - body_com_pose_b = cube_object.data.body_com_pose_b.torch - expected_com_pos, expected_com_quat = combine_frame_transforms( - link_pose_w[..., :3].reshape(-1, 3), - link_pose_w[..., 3:].reshape(-1, 4), - body_com_pose_b[..., :3].reshape(-1, 3), - body_com_pose_b[..., 3:].reshape(-1, 4), - ) - expected_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1).view(num_envs, num_cubes, 7) - torch.testing.assert_close(expected_com_pose, cube_object.data.body_com_pose_w.torch) - torch.testing.assert_close( - cube_object.data.body_com_vel_w.torch[..., 3:], cube_object.data.body_link_vel_w.torch[..., 3:] - ) - - -@pytest.mark.parametrize("num_envs", [3]) -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_reset_object_collection(sim, num_envs, num_cubes, device): - """Test that reset clears the external wrenches of the selected environments only.""" - object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) - sim.reset() - sim.step() - object_collection.update(sim.cfg.dt) - - composers = (object_collection.instantaneous_wrench_composer, object_collection.permanent_wrench_composer) - ones = torch.ones((num_envs, num_cubes, 3), device=device) - object_collection.permanent_wrench_composer.set_forces_and_torques_index(forces=ones, torques=ones) - object_collection.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones) - - # A partial reset clears only the selected environment - object_collection.reset(env_ids=torch.tensor([0], device=device)) - for composer in composers: - assert composer.active - for buffer in (composer.out_force_b.torch, composer.out_torque_b.torch): - assert torch.count_nonzero(buffer[0]) == 0 - assert torch.count_nonzero(buffer[1:]) == buffer[1:].numel() - - # A full reset clears every environment - object_collection.reset() - for composer in composers: - assert torch.count_nonzero(composer.out_force_b.torch) == 0 - assert torch.count_nonzero(composer.out_torque_b.torch) == 0 - - -@pytest.mark.parametrize("num_envs", [3]) -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("gravity_enabled", [True, False]) -def test_gravity_vec_w(sim, num_envs, num_cubes, device, gravity_enabled): - """Test that gravity vector direction is set correctly for the rigid object.""" - object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) - - # Obtain gravity direction - gravity_dir = (0.0, 0.0, -1.0) if gravity_enabled else (0.0, 0.0, 0.0) - - sim.reset() - - # Check if gravity vector is set correctly - gravity_vec = object_collection.data.GRAVITY_VEC_W.torch - assert gravity_vec[0, 0, 0] == gravity_dir[0] - assert gravity_vec[0, 0, 1] == gravity_dir[1] - assert gravity_vec[0, 0, 2] == gravity_dir[2] - - # Perform simulation - for _ in range(2): - sim.step() - object_collection.update(sim.cfg.dt) - - # Expected gravity value is the acceleration of the body - gravity = torch.zeros(num_envs, num_cubes, 6, device=device) - if gravity_enabled: - gravity[..., 2] = -9.81 - - # Check the body accelerations are correct - torch.testing.assert_close(object_collection.data.body_com_acc_w.torch, gravity) diff --git a/source/isaaclab_physx/test/assets/test_surface_gripper.py b/source/isaaclab_physx/test/assets/test_surface_gripper.py index c2d02b226f56..729282e55e5a 100644 --- a/source/isaaclab_physx/test/assets/test_surface_gripper.py +++ b/source/isaaclab_physx/test/assets/test_surface_gripper.py @@ -3,243 +3,194 @@ # # SPDX-License-Identifier: BSD-3-Clause -# ignore private usage of variables warning -# pyright: reportPrivateUsage=none +"""Real PhysX surface-gripper coverage on a locally authored gripper.""" import os -from isaaclab.test.utils import launch_test_simulation -from isaaclab.utils import replace +from isaaclab.test.utils import DeviceScope, launch_test_simulation, test_devices launch_test_simulation() +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np import pytest import torch import warp as wp from isaaclab_physx.assets import SurfaceGripper, SurfaceGripperCfg +from isaaclab.sim.utils import enable_extension + +enable_extension("isaacsim.robot.surface_gripper") + +from usd.schema.isaac import robot_schema + +from isaacsim.robot.surface_gripper import create_surface_gripper +from pxr import Gf, Sdf, UsdGeom, UsdPhysics + import isaaclab.sim as sim_utils -from isaaclab.actuators import ImplicitActuatorCfg -from isaaclab.assets import ( - Articulation, - ArticulationCfg, - RigidObject, - RigidObjectCfg, -) from isaaclab.sim import build_simulation_context -from isaaclab.utils.assets import ISAACLAB_NUCLEUS_DIR -from isaaclab.utils.version import get_isaac_sim_version, has_kit - -# from isaacsim.robot.surface_gripper import GripperView _RUNNING_CI = bool( os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true" or os.environ.get("GITLAB_CI") ) -def generate_surface_gripper_cfgs( - kinematic_enabled: bool = False, - max_grip_distance: float = 0.1, - coaxial_force_limit: float = 100.0, - shear_force_limit: float = 100.0, - retry_interval: float = 0.1, - reset_xform_op_properties: bool = False, -) -> tuple[SurfaceGripperCfg, ArticulationCfg]: - """Generate a surface gripper cfg and an articulation cfg. - - Args: - max_grip_distance: The maximum grip distance of the surface gripper. - coaxial_force_limit: The coaxial force limit of the surface gripper. - shear_force_limit: The shear force limit of the surface gripper. - retry_interval: The retry interval of the surface gripper. - reset_xform_op_properties: Whether to reset the xform op properties of the surface gripper. - - Returns: - A tuple containing the surface gripper cfg and the articulation cfg. - """ - articulation_cfg = ArticulationCfg( - spawn=sim_utils.UsdFileCfg( - usd_path=f"{ISAACLAB_NUCLEUS_DIR}/Tests/SurfaceGripper/test_gripper.usd", - rigid_props=sim_utils.UsdPhysicsRigidBodyCfg(kinematic_enabled=kinematic_enabled), - ), - init_state=ArticulationCfg.InitialStateCfg( - pos=(0.0, 0.0, 0.5), - rot=(0.0, 0.0, 0.0, 1.0), - joint_pos={ - ".*": 0.0, - }, - ), - actuators={ - "dummy": ImplicitActuatorCfg( - joint_names_expr=[".*"], - stiffness=0.0, - damping=0.0, - ), - }, +def _create_rigid_cube(path: str, position: tuple[float, float, float]) -> None: + """Author one local 1 kg rigid collision cube.""" + stage = sim_utils.get_current_stage() + cube = UsdGeom.Cube.Define(stage, path) + cube.CreateSizeAttr(0.1) + cube.AddTranslateOp().Set(Gf.Vec3f(*position)) + prim = cube.GetPrim() + UsdPhysics.CollisionAPI.Apply(prim) + UsdPhysics.RigidBodyAPI.Apply(prim) + UsdPhysics.MassAPI.Apply(prim).CreateMassAttr(1.0) + + +def _author_surface_gripper() -> SurfaceGripper: + """Author a gripper on the upper of two stacked cubes, with its attachment point facing the lower cube.""" + stage = sim_utils.get_current_stage() + env_path = "/World/Env_0" + UsdGeom.Xform.Define(stage, env_path) + _create_rigid_cube(f"{env_path}/box0", (0.0, 0.0, 0.05)) + _create_rigid_cube(f"{env_path}/box1", (0.0, 0.0, 0.15)) + create_surface_gripper(stage, env_path) + + gripper_prim = stage.GetPrimAtPath(f"{env_path}/SurfaceGripper") + gripper_prim.GetAttribute(robot_schema.Attributes.COAXIAL_FORCE_LIMIT.name).Set(100.0) + gripper_prim.GetAttribute(robot_schema.Attributes.SHEAR_FORCE_LIMIT.name).Set(100.0) + gripper_prim.GetAttribute(robot_schema.Attributes.MAX_GRIP_DISTANCE.name).Set(0.1) + + joint_path = Sdf.Path(f"{env_path}/box1/attachment") + joint = UsdPhysics.Joint.Define(stage, joint_path) + robot_schema.ApplyAttachmentPointAPI(joint.GetPrim()) + joint.GetPrim().CreateAttribute( + robot_schema.Attributes.FORWARD_AXIS.name, robot_schema.Attributes.FORWARD_AXIS.type + ).Set(UsdPhysics.Tokens.x) + joint.GetPrim().CreateAttribute( + robot_schema.Attributes.CLEARANCE_OFFSET.name, robot_schema.Attributes.CLEARANCE_OFFSET.type + ).Set(0.0) + for limit in ["rotX", "rotY", "rotZ", "transX", "transY", "transZ"]: + limit_api = UsdPhysics.LimitAPI.Apply(joint.GetPrim(), limit) + limit_api.CreateHighAttr().Set(-1.0) + limit_api.CreateLowAttr().Set(1.0) + joint.CreateBody0Rel().SetTargets([f"{env_path}/box1"]) + joint.CreateLocalPos0Attr().Set(Gf.Vec3f(0.0, 0.0, -0.0499)) + joint.CreateLocalRot0Attr().Set(Gf.Quatf(0.5, -0.5, 0.5, 0.5)) + gripper_prim.GetRelationship(robot_schema.Relations.ATTACHMENT_POINTS.name).SetTargets([joint_path]) + + return SurfaceGripper( + SurfaceGripperCfg( + prim_path="/World/Env_[^/]*/SurfaceGripper", + max_grip_distance=0.1, + coaxial_force_limit=100.0, + shear_force_limit=100.0, + retry_interval=0.1, + ) ) - surface_gripper_cfg = SurfaceGripperCfg( - max_grip_distance=max_grip_distance, - coaxial_force_limit=coaxial_force_limit, - shear_force_limit=shear_force_limit, - retry_interval=retry_interval, - ) - - return surface_gripper_cfg, articulation_cfg - - -def generate_surface_gripper( - surface_gripper_cfg: SurfaceGripperCfg, - articulation_cfg: ArticulationCfg, - num_surface_grippers: int, - device: str, -) -> tuple[SurfaceGripper, Articulation, torch.Tensor]: - """Generate a surface gripper and an articulation. - Args: - surface_gripper_cfg: The surface gripper cfg. - articulation_cfg: The articulation cfg. - num_surface_grippers: The number of surface grippers to generate. - device: The device to run the test on. - - Returns: - A tuple containing the surface gripper, the articulation, and the translations of the surface grippers. - """ - # Generate translations of 2.5 m in x for each articulation - translations = torch.zeros(num_surface_grippers, 3, device=device) - translations[:, 0] = torch.arange(num_surface_grippers) * 2.5 - - # Create Top-level Xforms, one for each articulation - for i in range(num_surface_grippers): - sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=translations[i][:3]) - articulation = Articulation(replace(articulation_cfg, prim_path="/World/Env_[^/]*/Robot")) - surface_gripper_cfg = replace(surface_gripper_cfg, prim_path="/World/Env_[^/]*/Robot/Gripper/SurfaceGripper") - surface_gripper = SurfaceGripper(surface_gripper_cfg) - - return surface_gripper, articulation, translations - - -def generate_grippable_object(sim, num_grippable_objects: int): - object_cfg = RigidObjectCfg( - prim_path="/World/Env_[^/]*/Object", - spawn=sim_utils.CuboidCfg( - size=(1.0, 1.0, 1.0), - rigid_props=sim_utils.UsdPhysicsRigidBodyCfg(), - mass_props=sim_utils.MassCfg(mass=1.0), - collision_props=sim_utils.UsdPhysicsCollisionCfg(), - visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 1.0, 0.0)), - ), - init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 0.5)), - ) - grippable_object = RigidObject(object_cfg) - - return grippable_object - - -@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 build_simulation_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]) -@pytest.mark.parametrize("device", ["cpu"]) -@pytest.mark.parametrize("add_ground_plane", [True]) @pytest.mark.isaacsim_ci @pytest.mark.skipif( _RUNNING_CI, reason="Isaac Sim SurfaceGripperView initialization can deadlock in CI; keep CUDA fail-fast coverage only.", ) -def test_close_and_open_command(sim, num_articulations, device, add_ground_plane) -> None: +def test_close_and_open_command() -> None: """Test that the close/open commands actually drive the surface gripper status. This is a regression test for the command plumbing: a single ``close`` command must move the gripper out of the *open* state into a *closing* (or *closed*) state, and a subsequent ``open`` command must bring it back to the *open* state. - - .. note:: - The shared ``test_gripper.usd`` does not contain a separate grippable rigid body (the cube - at the attachment point is a link of the gripper articulation), so the gripper cannot latch - onto anything and will not reach the *closed* state. We therefore only assert that the - command takes effect (the status leaves *open*), which is the deterministic behavior. - - Args: - num_articulations: The number of articulations to initialize. - device: The device to run the test on. - add_ground_plane: Whether to add a ground plane to the simulation. """ - if has_kit() and get_isaac_sim_version().major < 5: - return - surface_gripper_cfg, articulation_cfg = generate_surface_gripper_cfgs(kinematic_enabled=False) - surface_gripper, articulation, _ = generate_surface_gripper( - surface_gripper_cfg, articulation_cfg, num_articulations, device - ) + with build_simulation_context(device="cpu", gravity_enabled=False) as sim: + sim._app_control_on_stop_handle = None + surface_gripper = _author_surface_gripper() - sim.reset() - - assert articulation.is_initialized - assert surface_gripper.is_initialized - assert surface_gripper.command.shape == (num_articulations,) - assert surface_gripper.state.shape == (num_articulations,) - # after a reset the gripper is idle (0.0) and open (-1.0) - assert torch.all(wp.to_torch(surface_gripper.command) == 0.0) - assert torch.all(wp.to_torch(surface_gripper.state) == -1.0) - - # send a single close command (the action term is edge-triggered, so commands are sent once) - close_cmd = wp.array([1.0] * num_articulations, dtype=wp.float32, device=device) - surface_gripper.set_grippers_command_index(close_cmd) - surface_gripper.write_data_to_sim() - # step the simulation so the gripper reacts to the command - for _ in range(3): - sim.step() - articulation.update(sim.cfg.dt) - surface_gripper.update(sim.cfg.dt) - # the close command must take effect: status is "closing" (0.0) or "closed" (1.0), never "open" (-1.0) - state_after_close = wp.to_torch(surface_gripper.state) - assert torch.all(state_after_close >= 0.0), f"close command had no effect, state={state_after_close.tolist()}" - - # send a single open command; the gripper must return to the open state - open_cmd = wp.array([-1.0] * num_articulations, dtype=wp.float32, device=device) - surface_gripper.set_grippers_command_index(open_cmd) - surface_gripper.write_data_to_sim() - for _ in range(3): - sim.step() - articulation.update(sim.cfg.dt) - surface_gripper.update(sim.cfg.dt) - # the open command must take effect: status is back to "open" (-1.0) - state_after_open = wp.to_torch(surface_gripper.state) - assert torch.all(state_after_open == -1.0), f"open command had no effect, state={state_after_open.tolist()}" - - -@pytest.mark.parametrize("device", ["cuda:0"]) -@pytest.mark.parametrize("add_ground_plane", [True]) + sim.reset() + + assert surface_gripper.is_initialized + assert surface_gripper.command.shape == (1,) + assert surface_gripper.state.shape == (1,) + # after a reset the gripper is idle (0.0) and open (-1.0) + assert torch.all(wp.to_torch(surface_gripper.command) == 0.0) + assert torch.all(wp.to_torch(surface_gripper.state) == -1.0) + + # send a single close command (the action term is edge-triggered, so commands are sent once) + surface_gripper.set_grippers_command_index(wp.array([1.0], dtype=wp.float32, device="cpu")) + surface_gripper.write_data_to_sim() + for _ in range(3): + sim.step() + surface_gripper.update(sim.cfg.dt) + # the close command must take effect: status is "closing" (0.0) or "closed" (1.0), never "open" (-1.0) + state_after_close = wp.to_torch(surface_gripper.state) + assert torch.all(state_after_close >= 0.0), f"close command had no effect, state={state_after_close.tolist()}" + + # send a single open command; the gripper must return to the open state + surface_gripper.set_grippers_command_index(wp.array([-1.0], dtype=wp.float32, device="cpu")) + surface_gripper.write_data_to_sim() + for _ in range(3): + sim.step() + surface_gripper.update(sim.cfg.dt) + # the open command must take effect: status is back to "open" (-1.0) + state_after_open = wp.to_torch(surface_gripper.state) + assert torch.all(state_after_open == -1.0), f"open command had no effect, state={state_after_open.tolist()}" + + +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) @pytest.mark.isaacsim_ci -def test_raise_error_if_not_cpu(sim, device, add_ground_plane) -> None: +def test_raise_error_if_not_cpu(device) -> None: """Test that the SurfaceGripper raises an error if the device is not CPU.""" - if has_kit() and get_isaac_sim_version().major < 5: - return - num_articulations = 1 - surface_gripper_cfg, articulation_cfg = generate_surface_gripper_cfgs(kinematic_enabled=False) - surface_gripper, articulation, translations = generate_surface_gripper( - surface_gripper_cfg, articulation_cfg, num_articulations, device + with build_simulation_context(device=device, gravity_enabled=False) as sim: + sim._app_control_on_stop_handle = None + surface_gripper = _author_surface_gripper() + assert not surface_gripper.is_initialized + + with pytest.raises(Exception, match="only supported on CPU"): + sim.reset() + + +## +# View payloads, without a simulation. +## + + +def test_command_filter_and_partial_property_update_use_literal_view_payloads() -> None: + """Submit only open/close commands and preserve property selector ordering.""" + # Skip USD initialization: a three-environment CPU gripper around a recording view. + gripper = object.__new__(SurfaceGripper) + gripper._device = "cpu" + gripper._num_envs = 3 + gripper._ALL_INDICES = wp.array([0, 1, 2], dtype=wp.int32, device="cpu") + gripper._gripper_command = wp.zeros(3, dtype=wp.float32, device="cpu") + gripper._max_grip_distance = wp.zeros(3, dtype=wp.float32, device="cpu") + gripper._coaxial_force_limit = wp.zeros(3, dtype=wp.float32, device="cpu") + gripper._shear_force_limit = wp.zeros(3, dtype=wp.float32, device="cpu") + gripper._retry_interval = wp.zeros(3, dtype=wp.float32, device="cpu") + gripper._gripper_view = SimpleNamespace( + apply_gripper_action=Mock(), + set_surface_gripper_properties=Mock(), ) - with pytest.raises(Exception, match="only supported on CPU"): - sim.reset() + gripper.set_grippers_command_index(wp.array([0.5, 0.0, -0.5], dtype=wp.float32, device="cpu")) + gripper.write_data_to_sim() + + gripper.gripper_view.apply_gripper_action.assert_called_once_with([0.5, 0.0, -0.5], [[0], [2]]) + + env_ids = wp.array([2, 0], dtype=wp.int32, device="cpu") + gripper.update_gripper_properties_index( + max_grip_distance=wp.array([0.2, 0.4], dtype=wp.float32, device="cpu"), + env_ids=env_ids, + ) -if __name__ == "__main__": - pytest.main([__file__, "-v", "--maxfail=1"]) + np.testing.assert_array_equal(gripper._max_grip_distance.numpy(), np.asarray([0.4, 0.0, 0.2], dtype=np.float32)) + properties = gripper.gripper_view.set_surface_gripper_properties.call_args.kwargs + np.testing.assert_array_equal(properties.pop("max_grip_distance"), np.asarray([0.4, 0.0, 0.2], dtype=np.float32)) + assert properties == { + "coaxial_force_limit": [0.0, 0.0, 0.0], + "shear_force_limit": [0.0, 0.0, 0.0], + "retry_interval": [0.0, 0.0, 0.0], + "indices": [2, 0], + } diff --git a/tools/_device_split.py b/tools/_device_split.py deleted file mode 100644 index e8bad4069133..000000000000 --- a/tools/_device_split.py +++ /dev/null @@ -1,65 +0,0 @@ -# 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 - -"""Helpers for detecting and driving the ``device_split`` pytest marker. - -Test files that declare ``pytestmark = pytest.mark.device_split`` at module -scope must be re-invoked once per device (CPU and GPU) in separate processes -to work around process-global device locks such as ``ovphysx<=0.3.7`` gap G5. -The :func:`is_device_split_file` predicate lets the per-file CI runner in -``tools/conftest.py`` detect this without importing the test module. -""" - -from __future__ import annotations - -import re -from pathlib import Path - -_DEVICE_SPLIT_MARK_RE = re.compile(r"^\s*pytestmark\b.*\bdevice_split\b", re.MULTILINE) -"""Match a module-level ``pytestmark`` assignment that mentions ``device_split``. - -Recognises both single-mark and single-line list forms: - -* ``pytestmark = pytest.mark.device_split`` -* ``pytestmark = [pytest.mark.device_split, pytest.mark.foo]`` - -Multi-line list forms are not supported (currently no test file uses one); if -a future test needs that, expand the parsing rule. -""" - -# Per-pass pytest ``-k`` selectors used by ``tools/conftest.py`` when a file -# declares the ``device_split`` marker. Each entry is ``(suffix, k_expr)``: -# - ``suffix`` is appended to the JUnit report filename to keep both passes' XML. -# - ``k_expr`` is the ``-k`` keyword expression. ``"cpu or not cuda"`` keeps -# non-parametrized tests in the CPU pass; ``"cuda"`` catches GPU-parametrized -# tests only. -DEVICE_SPLIT_PASSES: list[tuple[str, str]] = [ - ("-cpu", "cpu or not cuda"), - ("-cuda", "cuda"), -] - - -def is_device_split_file(path: Path | str, source: str | None = None) -> bool: - """Return whether the test file at ``path`` declares the ``device_split`` marker. - - Matches :data:`_DEVICE_SPLIT_MARK_RE` against ``source`` when supplied. - Otherwise, reads the file source from ``path``. A missing or unreadable - file returns ``False`` so the caller falls back to the default single-pass - invocation. - - Args: - path: Filesystem path to a candidate test file. - source: Optional preloaded source text to inspect. - - Returns: - ``True`` when the file's module-level ``pytestmark`` mentions - ``device_split``; ``False`` otherwise. - """ - if source is None: - try: - source = Path(path).read_text(encoding="utf-8", errors="replace") - except OSError: - return False - return bool(_DEVICE_SPLIT_MARK_RE.search(source)) diff --git a/tools/conftest.py b/tools/conftest.py index 71c6d2022285..2297e92d5785 100644 --- a/tools/conftest.py +++ b/tools/conftest.py @@ -27,7 +27,6 @@ import ovrtx_log # isort: skip import test_settings as test_settings # isort: skip from crash_journal import JOURNAL_ENV_VAR, create_crash_report # isort: skip -from _device_split import DEVICE_SPLIT_PASSES, is_device_split_file # isort: skip from _file_scheduler import JobContext, TestFileJob, run_test_files # isort: skip logging.basicConfig(level=logging.INFO, format="%(message)s") @@ -1510,13 +1509,6 @@ def _run_test_file( timeout = test_settings.PER_TEST_TIMEOUTS.get(file_name, test_settings.DEFAULT_TIMEOUT) - # Read the test file once for device-split detection. - try: - with open(test_file) as fh: - test_content = fh.read() - except OSError: - test_content = "" - # The first renderer in a fresh container compiles shaders (~600 s). # Give it extra time so that doesn't look like a test timeout. is_cold_cache_test = context.renderer_cold @@ -1539,18 +1531,7 @@ def _run_test_file( capture=_CaptureOptions(echo=echo, workers=workers, on_started=context.mark_started), ) - # On a multi-GPU shard, test_devices() already resolves to this shard's single - # GPU and mgpu_shard_select drops every other variant, so the device_split - # CPU/GPU two-pass (which exists to dodge the process-global device lock when - # CPU and GPU share one container) is unnecessary here — the CPU pass would - # collect zero tests yet still pay full Kit-startup cost. Run once on a shard. - if _inject_shard_select: - passes = [("", None)] - elif is_device_split_file(test_file, source=test_content): - logger.info(f"⚙️ device_split detected — invoking {file_name} once per device (CPU then GPU)") - passes = DEVICE_SPLIT_PASSES - else: - passes = [("", None)] + passes = [("", None)] reports = [] failed = False diff --git a/tools/test_device_split.py b/tools/test_device_split.py deleted file mode 100644 index fb34ce40c4f1..000000000000 --- a/tools/test_device_split.py +++ /dev/null @@ -1,108 +0,0 @@ -# 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 - -"""Unit tests for ``tools/_device_split.py``.""" - -from __future__ import annotations - -import textwrap -from pathlib import Path - -from _device_split import is_device_split_file - - -def _write(tmp_path: Path, content: str) -> Path: - p = tmp_path / "test_foo.py" - p.write_text(textwrap.dedent(content)) - return p - - -def test_single_mark(tmp_path): - f = _write( - tmp_path, - """ - import pytest - - pytestmark = pytest.mark.device_split - - def test_x(): - pass - """, - ) - assert is_device_split_file(f) is True - - -def test_list_form_single_line(tmp_path): - f = _write( - tmp_path, - """ - import pytest - - pytestmark = [pytest.mark.device_split, pytest.mark.foo] - - def test_x(): - pass - """, - ) - assert is_device_split_file(f) is True - - -def test_preloaded_source(tmp_path): - source = textwrap.dedent( - """ - import pytest - - pytestmark = pytest.mark.device_split - """ - ) - assert is_device_split_file(tmp_path / "does_not_exist.py", source=source) is True - - -def test_no_mark(tmp_path): - f = _write( - tmp_path, - """ - import pytest - - def test_x(): - pass - """, - ) - assert is_device_split_file(f) is False - - -def test_word_in_comment_does_not_match(tmp_path): - f = _write( - tmp_path, - """ - import pytest - - # This file mentions device_split in a comment but is not marked. - - def test_x(): - pass - """, - ) - assert is_device_split_file(f) is False - - -def test_unrelated_pytestmark_does_not_match(tmp_path): - f = _write( - tmp_path, - """ - import pytest - - pytestmark = pytest.mark.skipif(False, reason="x") - - def test_x(): - pass - """, - ) - assert is_device_split_file(f) is False - - -def test_missing_file(tmp_path): - # A path that does not exist must not raise; treat as not-marked. - assert is_device_split_file(tmp_path / "does_not_exist.py") is False diff --git a/tools/test_settings.py b/tools/test_settings.py index a77891aecad8..8a49d804387b 100644 --- a/tools/test_settings.py +++ b/tools/test_settings.py @@ -49,7 +49,6 @@ "test_visualizer_golden_newton.py": 6000, "test_visuotactile_sensor.py": 1000, "test_visuotactile_render.py": 1000, - "test_rigid_object_collection.py": 1500, "test_outdated_sensor.py": 1000, "test_multi_tiled_camera.py": 1000, "test_multirotor.py": 1000,