From 556bd814f2382e84610ed7d90659a4daa605c058 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 09:49:17 +0200 Subject: [PATCH 01/33] Fix PhysX rigid-object collection COM and inertia setters RigidObjectCollection.set_coms_index passed wp.transformf data to the PhysX tensor API, which raises a KeyError. It also marked the lazily filled body_com_pose_b cache as fresh after a partial write, so the unselected entries read back as zeros when the cache had not been read before. set_inertias_index flattened the (num_instances, num_bodies, 9) inertias with the 2D reshape, which the backend rejects. Pass the COM poses as float32, let the next read refresh the COM cache from the view as RigidObject does, and reshape the inertias with the 3D helper. A new test writes both properties to a non-sorted subset of (env, body) cells and checks the PhysX view and the data buffers. --- ...oiner-fix-physx-collection-com-inertia.rst | 10 ++++++ .../rigid_object_collection.py | 11 +++--- .../assets/test_rigid_object_collection.py | 36 +++++++++++++++++++ 3 files changed, 52 insertions(+), 5 deletions(-) create mode 100644 source/isaaclab_physx/changelog.d/antoiner-fix-physx-collection-com-inertia.rst diff --git a/source/isaaclab_physx/changelog.d/antoiner-fix-physx-collection-com-inertia.rst b/source/isaaclab_physx/changelog.d/antoiner-fix-physx-collection-com-inertia.rst new file mode 100644 index 000000000000..c985ac498742 --- /dev/null +++ b/source/isaaclab_physx/changelog.d/antoiner-fix-physx-collection-com-inertia.rst @@ -0,0 +1,10 @@ +Fixed +^^^^^ + +* Fixed :meth:`~isaaclab_physx.assets.RigidObjectCollection.set_coms_index` 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` leaving the unselected entries of + :attr:`~isaaclab_physx.assets.RigidObjectCollectionData.body_com_pose_b` zeroed when it was called for a subset of + bodies or environments before the property was first read. +* Fixed :meth:`~isaaclab_physx.assets.RigidObjectCollection.set_inertias_index` 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_rigid_object_collection.py b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py index 2e0af9a4be8b..ef289e6673d2 100644 --- a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py @@ -152,6 +152,42 @@ def test_subset_write_reaches_selected_view_entry(sim, device, gravity_enabled): torch.testing.assert_close(object_collection.data.body_link_pose_w.torch, expected_pose) +@pytest.mark.parametrize("device", test_devices()) +def test_inertial_property_subset_writes_reach_selected_view_entries(sim, device): + """COM and inertia writes to a subset of (env, body) cells must reach only those PhysX view entries.""" + num_envs, num_cubes = 2, 3 + object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) + sim.reset() + + def read_view(values: wp.array, data_dim: int) -> torch.Tensor: + # The view is body-major, (num_cubes * num_envs, data_dim); return it as (num_envs, num_cubes, data_dim) + return wp.to_torch(values).reshape(num_cubes, num_envs, data_dim).transpose(0, 1).cpu() + + # Write in non-sorted env order so that a wrong env/body to view index mapping cannot pass + env_ids = [1, 0] + body_ids = [2] + + initial_coms = read_view(object_collection.root_view.get_coms().view(wp.float32), 7) + coms = initial_coms[env_ids][:, body_ids].clone() + coms[0, 0, :3] = torch.tensor([0.02, 0.03, 0.04]) + coms[1, 0, :3] = torch.tensor([-0.01, 0.01, 0.02]) + object_collection.set_coms_index(coms=coms.to(device), env_ids=env_ids, body_ids=body_ids) + expected_coms = initial_coms.clone() + expected_coms[env_ids, 2] = coms[:, 0] + torch.testing.assert_close(read_view(object_collection.root_view.get_coms().view(wp.float32), 7), expected_coms) + torch.testing.assert_close(object_collection.data.body_com_pose_b.torch.cpu(), expected_coms) + + initial_inertias = read_view(object_collection.root_view.get_inertias(), 9) + inertias = initial_inertias[env_ids][:, body_ids].clone() + inertias[0, 0, [0, 4, 8]] *= 1.5 + inertias[1, 0, [0, 4, 8]] *= 2.0 + object_collection.set_inertias_index(inertias=inertias.to(device), env_ids=env_ids, body_ids=body_ids) + expected_inertias = initial_inertias.clone() + expected_inertias[env_ids, 2] = inertias[:, 0] + torch.testing.assert_close(read_view(object_collection.root_view.get_inertias(), 9), expected_inertias) + torch.testing.assert_close(object_collection.data.body_inertia.torch.cpu(), expected_inertias) + + @pytest.mark.parametrize("num_envs", [2]) @pytest.mark.parametrize("num_cubes", [3]) @pytest.mark.parametrize("device", test_devices()) From 0d61d2cfecf94bc8c50fcd03300b8dd95f3d8416 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 10:32:37 +0200 Subject: [PATCH 02/33] Restructure shared asset tests into a capability-declared contract Move the mocked cross-backend asset interface tests into a contract package with api, data, and write entry modules. An explicit capability table gives every backend a visible row: unavailable backends and unsupported branches collect as skips with their reason instead of silently dropping out. A kitless usdrt/carb stub and an availability guard keep PhysX articulations in the matrix, where they previously were skipped without notice. Factory manager patches are scoped to one test and restored afterwards. PhysX views apply TensorAPI writes so partial writers are checked against unselected backend rows, Newton joins the fixed-tendon contract, the deprecated command aliases are checked against the actuator collection, and one enforced inventory classifies the public base-class surface. --- .../antoiner-asset-tests-v2-shared.skip | 0 .../isaaclab/test/assets/_iface_test_boot.py | 30 -- .../isaaclab/test/assets/contract/__init__.py | 6 + .../_articulation_contract_cases.py} | 320 +++++++++++++----- .../_articulation_contract_utils.py} | 109 +++--- .../_articulation_ordering_contract_cases.py} | 48 +-- .../test/assets/contract/_contract_boot.py | 59 ++++ .../assets/contract/_manager_patch_scope.py | 44 +++ .../test/assets/contract/_mock_backends.py | 70 ++++ ...rigid_object_collection_contract_cases.py} | 143 ++++++-- ...rigid_object_collection_contract_utils.py} | 79 ++--- .../_rigid_object_contract_cases.py} | 133 ++++++-- .../_rigid_object_contract_utils.py} | 74 ++-- .../test/assets/contract/capabilities.py | 152 +++++++++ .../isaaclab/test/assets/contract/conftest.py | 19 ++ .../test/assets/contract/public_surface.py | 263 ++++++++++++++ .../contract/test_asset_contract_api.py | 115 +++++++ .../contract/test_asset_contract_data.py | 30 ++ .../contract/test_asset_contract_writes.py | 39 +++ 19 files changed, 1361 insertions(+), 372 deletions(-) create mode 100644 source/isaaclab/changelog.d/antoiner-asset-tests-v2-shared.skip delete mode 100644 source/isaaclab/test/assets/_iface_test_boot.py create mode 100644 source/isaaclab/test/assets/contract/__init__.py rename source/isaaclab/test/assets/{test_articulation_iface.py => contract/_articulation_contract_cases.py} (84%) rename source/isaaclab/test/assets/{_articulation_iface_test_utils.py => contract/_articulation_contract_utils.py} (87%) rename source/isaaclab/test/assets/{test_articulation_ordering_iface.py => contract/_articulation_ordering_contract_cases.py} (98%) create mode 100644 source/isaaclab/test/assets/contract/_contract_boot.py create mode 100644 source/isaaclab/test/assets/contract/_manager_patch_scope.py create mode 100644 source/isaaclab/test/assets/contract/_mock_backends.py rename source/isaaclab/test/assets/{test_rigid_object_collection_iface.py => contract/_rigid_object_collection_contract_cases.py} (87%) rename source/isaaclab/test/assets/{_rigid_object_collection_iface_test_utils.py => contract/_rigid_object_collection_contract_utils.py} (84%) rename source/isaaclab/test/assets/{test_rigid_object_iface.py => contract/_rigid_object_contract_cases.py} (88%) rename source/isaaclab/test/assets/{_rigid_object_iface_test_utils.py => contract/_rigid_object_contract_utils.py} (84%) create mode 100644 source/isaaclab/test/assets/contract/capabilities.py create mode 100644 source/isaaclab/test/assets/contract/conftest.py create mode 100644 source/isaaclab/test/assets/contract/public_surface.py create mode 100644 source/isaaclab/test/assets/contract/test_asset_contract_api.py create mode 100644 source/isaaclab/test/assets/contract/test_asset_contract_data.py create mode 100644 source/isaaclab/test/assets/contract/test_asset_contract_writes.py diff --git a/source/isaaclab/changelog.d/antoiner-asset-tests-v2-shared.skip b/source/isaaclab/changelog.d/antoiner-asset-tests-v2-shared.skip new file mode 100644 index 000000000000..e69de29bb2d1 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/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/test_articulation_iface.py b/source/isaaclab/test/assets/contract/_articulation_contract_cases.py similarity index 84% rename from source/isaaclab/test/assets/test_articulation_iface.py rename to source/isaaclab/test/assets/contract/_articulation_contract_cases.py index 7c8b169e2249..4d314faa281f 100644 --- a/source/isaaclab/test/assets/test_articulation_iface.py +++ b/source/isaaclab/test/assets/contract/_articulation_contract_cases.py @@ -6,27 +6,25 @@ # 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. +"""Shared articulation contract cases, collected by the ``test_asset_contract_*`` entry modules. -The setup is a bit convoluted so that we can run these tests without requiring Isaac Sim or GPU simulation. +Checks that every articulation backend provides the data and writer behavior the base articulation class advertises. +The backends run on mocked views, so these cases need neither Isaac Sim nor a GPU simulation. """ import math import warnings -from unittest.mock import MagicMock +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 isaaclab.test.utils import DeviceScope, test_devices -pytestmark = pytest.mark.integration - +from ._articulation_contract_utils import get_articulation +from .capabilities import contract_backend, require_backend_capability, requires_backend # --------------------------------------------------------------------------- # Helpers @@ -44,16 +42,13 @@ def _check_proxy_array(arr, *, expected_shape: tuple, expected_dtype: type, name # 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) +_api_backends = contract_backend("api") +_data_backends = contract_backend("data") +_writes_backends = contract_backend("writes") +_index_resolution_backends = contract_backend("index_resolution") _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 -) # --------------------------------------------------------------------------- @@ -64,7 +59,7 @@ def _check_proxy_array(arr, *, expected_shape: tuple, expected_dtype: type, name class TestArticulationIndexResolution: """Test backend-specific index resolution helpers.""" - @_production_backends + @_api_backends @_devices def test_resolve_env_ids_handles_tensor_view_shape(self, backend, device): art, _ = get_articulation(backend, num_instances=4, device=device) @@ -114,7 +109,7 @@ def test_resolve_body_ids_handles_tensor_view_shape(self, backend): class TestArticulationProperties: """Test that articulation properties return the correct types/values.""" - @_backends + @_api_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 @@ -137,7 +132,7 @@ def test_articulation_counts_names_and_finders(self, backend, is_fixed_base): class TestArticulationFinderReturnModes: """Test finder return modes on production articulation backends.""" - @_production_backends + @_api_backends @pytest.mark.parametrize( "finder_name, names_attr, count", [ @@ -157,8 +152,8 @@ def test_finder_returns_legacy_list_or_cached_proxy(self, backend, finder_name, num_spatial_tendons=5, device="cpu", ) - if backend == "newton" and finder_name == "find_spatial_tendons": - pytest.skip("Newton does not support spatial tendons.") + if finder_name == "find_spatial_tendons": + require_backend_capability(backend, "spatial_tendons") finder = getattr(art, finder_name) indices, names = finder(".*") @@ -195,27 +190,23 @@ def _articulation(backend): ) return art - @_production_backends + @contract_backend("fixed_tendon_target_scheduling") 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 + @requires_backend("newton") 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)) + @requires_backend("newton") 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") @@ -231,7 +222,7 @@ def test_a_solver_without_tendon_transmission_refuses_to_build_an_adapter(self): class TestResolveMatchingNamesCache: """Test that resolve_matching_names caching returns correct, isolated results.""" - @_backends + @_api_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") @@ -240,7 +231,7 @@ def test_unmatched_regex_raises(self, backend): with pytest.raises(ValueError): art.find_joints("nonexistent_joint_xyz") - @_backends + @_api_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") @@ -256,7 +247,7 @@ def test_mutating_result_does_not_corrupt_cache(self, backend): assert len(idx2) == expected_len assert "corrupted" not in names2 - @_backends + @_api_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") @@ -337,7 +328,7 @@ def test_find_with_preserve_order(self, backend): class TestArticulationDataProperties: """Test that every data property is a ProxyArray with the advertised shape and dtype.""" - @_backends + @_data_backends @_devices def test_articulation_data_property_contract(self, backend, device): art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device) @@ -379,7 +370,7 @@ def test_articulation_data_property_contract(self, backend, device): assert view.data_ptr() == expected.data_ptr() assert view.stride() == expected.stride() - @pytest.mark.skipif("physx" not in BACKENDS, reason="PhysX backend unavailable") + @requires_backend("physx") 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") @@ -399,7 +390,7 @@ def counted_get_coms(): assert num_get_coms_calls == 0 - @pytest.mark.skipif("physx" not in BACKENDS, reason="PhysX backend unavailable") + @requires_backend("physx") 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") @@ -423,7 +414,7 @@ def counted_get_coms(): assert num_get_coms_calls == 1 - @pytest.mark.parametrize("backend", [name for name in ("physx", "ovphysx") if name in BACKENDS]) + @contract_backend("mocked_dynamics") @_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.""" @@ -442,7 +433,7 @@ def test_dynamics_buffers_allocate_on_read_and_reuse_across_refreshes(self, back art.data.update(dt=0.01) assert art.data.body_link_jacobian_w is wrapper - @_production_backends + @_data_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") @@ -465,6 +456,23 @@ def test_actuator_compatibility_projections_are_stable(self, backend): 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] + @_data_backends + def test_deprecated_command_aliases_return_actuator_collection_buffers(self, backend): + """The deprecated data command and torque properties warn and alias the actuator collection's buffers.""" + art, _ = get_articulation(backend, 2, 4, 5, device="cpu") + actuators = art.actuators + 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 + # --------------------------------------------------------------------------- # Tests: Alias/shorthand properties @@ -491,7 +499,7 @@ def test_actuator_compatibility_projections_are_stable(self, backend): class TestArticulationDataAliases: """Test that alias properties return the values of their canonical counterparts.""" - @_backends + @_data_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") @@ -670,7 +678,7 @@ def _make_item_mask(total: int, selected: list[int], device: str) -> wp.array: class TestArticulationWritersRoot: """Test root pose/velocity writers with all input combinations.""" - @_production_backends + @_writes_backends @pytest.mark.parametrize( "body_ordering", [None, ("body_0", "body_3", "body_2", "body_1")], @@ -705,7 +713,7 @@ def test_pose_write_invalidates_pose_dependent_caches(self, backend, body_orderi art.write_root_link_pose_to_sim_index(root_pose=root_pose) _assert_buffers_stale(art.data, buffers) - @_production_backends + @_writes_backends @pytest.mark.parametrize( "body_ordering", [None, ("body_0", "body_3", "body_2", "body_1")], @@ -733,7 +741,7 @@ def test_velocity_write_invalidates_body_frame_caches(self, backend, body_orderi art.write_root_com_velocity_to_sim_index(root_velocity=root_velocity) _assert_buffers_stale(art.data, buffers) - @_production_backends + @_writes_backends @pytest.mark.parametrize("setter_kind", ["index", "mask"]) @pytest.mark.parametrize( "body_ordering", @@ -797,8 +805,6 @@ def set_coms() -> None: 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"): @@ -809,7 +815,7 @@ def set_coms() -> None: # -- index variants -- - @_backends + @_writes_backends @_devices @pytest.mark.parametrize("method_suffix", _ROOT_POSE_METHODS) def test_write_root_pose_to_sim_index(self, backend, device, method_suffix): @@ -837,7 +843,7 @@ def test_write_root_pose_to_sim_index(self, backend, device, method_suffix): with pytest.raises((AssertionError, RuntimeError)): method(root_pose=_make_bad_data_warp((num_instances,), device, wp.transformf)) - @_backends + @_writes_backends @_devices @pytest.mark.parametrize("method_suffix", _ROOT_VEL_METHODS) def test_write_root_velocity_to_sim_index(self, backend, device, method_suffix): @@ -867,7 +873,7 @@ def test_write_root_velocity_to_sim_index(self, backend, device, method_suffix): # -- mask variants -- - @_backends + @_writes_backends @_devices @pytest.mark.parametrize("method_suffix", _ROOT_POSE_METHODS) def test_write_root_pose_to_sim_mask(self, backend, device, method_suffix): @@ -901,7 +907,7 @@ def test_write_root_pose_to_sim_mask(self, backend, device, method_suffix): with pytest.raises((AssertionError, RuntimeError)): method(root_pose=_make_bad_data_warp((num_instances,), device, wp.transformf)) - @_backends + @_writes_backends @_devices @pytest.mark.parametrize("method_suffix", _ROOT_VEL_METHODS) def test_write_root_velocity_to_sim_mask(self, backend, device, method_suffix): @@ -936,7 +942,7 @@ def test_write_root_velocity_to_sim_mask(self, backend, device, method_suffix): method(root_velocity=_make_bad_data_warp((num_instances,), device, wp.spatial_vectorf)) -@pytest.mark.parametrize("backend", BACKENDS) +@_writes_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 @@ -982,7 +988,7 @@ def _read_joint_writer_target(art, getter: str): class TestArticulationWritersJoint: """Test joint writers/setters with all input combinations.""" - @_backends + @_writes_backends @_devices @pytest.mark.parametrize( "method_base, kwarg, wp_dtype, accepts_float, getter", @@ -1036,7 +1042,7 @@ def test_joint_writer_index(self, backend, device, method_base, kwarg, wp_dtype, with pytest.raises((AssertionError, RuntimeError)): method(**{kwarg: _make_bad_data_warp((num_instances, num_joints), device, wp_dtype)}) - @_backends + @_writes_backends @_devices @pytest.mark.parametrize( "method_base, kwarg, wp_dtype, accepts_float, getter", @@ -1134,7 +1140,7 @@ def _make_body_warp(shape: tuple[int, int], device: str, wp_dtype: type, trailin class TestArticulationWritersBody: """Test body property writers/setters with all input combinations.""" - @_backends + @_writes_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): @@ -1175,7 +1181,7 @@ def test_body_writer_index(self, backend, device, method_base, kwarg): with pytest.raises((AssertionError, RuntimeError)): method(**{kwarg: _make_body_warp((num_instances + 1, num_bodies), device, wp_dtype, trailing)}) - @_backends + @_writes_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): @@ -1216,36 +1222,117 @@ def test_body_writer_mask(self, backend, device, method_base, kwarg): method(**{kwarg: _make_body_warp((num_instances + 1, num_bodies), device, wp_dtype, trailing)}) +# --------------------------------------------------------------------------- +# Tests: partial root and body writes reach only the selected backend rows +# --------------------------------------------------------------------------- + +# quantity -> (writer base name, writer keyword, trailing size). Joint partial writes are covered, with and without +# a joint ordering, by ``TestArticulationOrderingWriteParity``. +_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 == "physx": + getter = {"root_pose": "get_root_transforms", "root_velocity": "get_root_velocities", "mass": "get_masses"} + values = wp.to_torch(getattr(raw_backend, getter[quantity])()) + elif backend == "newton": + if quantity == "mass": + values = wp.to_torch(art.data._sim_bind_body_mass) + else: + getter = {"root_pose": "get_root_transforms", "root_velocity": "get_root_velocities"}[quantity] + values = wp.to_torch(getattr(raw_backend, getter)(None)) + else: + 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) + return values.reshape(art.num_instances, -1).cpu().clone() + + +class TestArticulationPartialWriteRows: + """Test that partial root and body writers reach only the selected rows of backend storage.""" + + @_writes_backends + @pytest.mark.parametrize("selection", ["index", "mask"]) + @pytest.mark.parametrize("quantity", _PARTIAL_WRITES) + def test_partial_write_preserves_unselected_backend_rows(self, backend, selection, quantity): + num_instances, num_bodies = 2, 4 + art, raw_backend = get_articulation(backend, num_instances, 3, num_bodies, device="cpu") + writer, kwarg, trailing = _PARTIAL_WRITES[quantity] + width = num_bodies if quantity == "mass" else trailing + # Literal, per-element distinct payload; poses keep an identity rotation. + values = 100.0 * torch.arange(1, num_instances + 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 + kwargs = {} + if selection == "index": + kwargs[kwarg] = values[env : env + 1, body : body + 1] if quantity == "mass" else values[env : env + 1] + kwargs["env_ids"] = torch.tensor([env], dtype=torch.int32) + if quantity == "mass": + kwargs["body_ids"] = [body] + else: + kwargs[kwarg] = values + kwargs["env_mask"] = _make_item_mask(num_instances, [env], "cpu") + if quantity == "mass": + kwargs["body_mask"] = _make_item_mask(num_bodies, [body], "cpu") + + getattr(art, f"{writer}_{selection}")(**kwargs) + + after = _read_backend_rows(backend, art, raw_backend, quantity) + expected = before.clone() + if quantity == "mass": + expected[env, body] = values[env, body] + else: + expected[env] = values[env] + torch.testing.assert_close(after, expected, rtol=0.0, atol=0.0) + + # --------------------------------------------------------------------------- # 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) +# Every backend supports the core fixed-tendon properties. Newton declares the extended fixed-tendon properties and +# spatial tendons unsupported, so those rows stay visible as reasoned capability skips. +_tendon_backends = contract_backend("fixed_tendons") # 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_DIMS = [ + (1, 2, 2, 1, 0), # fixed only + (2, 6, 7, 3, 2), # both types +] +_TENDON_DIM_NAMES = ("num_instances", "num_joints", "num_bodies", "num_fixed_tendons", "num_spatial_tendons") +_tendon_dims = pytest.mark.parametrize(_TENDON_DIM_NAMES, _TENDON_DIMS) + +# tendon property family -> (tendon kind, backend capability) +_TENDON_FAMILIES = { + "fixed": ("fixed", "fixed_tendons"), + "fixed_extended": ("fixed", "fixed_tendon_extended"), + "spatial": ("spatial", "spatial_tendons"), +} + +# (property, family, 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), + ("fixed_tendon_limit_stiffness", "fixed_extended", wp.float32), + ("fixed_tendon_rest_length", "fixed_extended", wp.float32), + ("fixed_tendon_offset", "fixed_extended", wp.float32), ("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), ] +_TENDON_PROPERTY_FAMILY = {name: family for name, family, _ in _TENDON_DATA_PROPERTIES} class TestArticulationTendons: @@ -1254,28 +1341,27 @@ class TestArticulationTendons: @_tendon_backends @_tendon_dims @_devices + @pytest.mark.parametrize("family", _TENDON_FAMILIES) def test_articulation_tendon_contract( - self, backend, num_instances, num_joints, num_bodies, num_fixed_tendons, num_spatial_tendons, device + self, backend, num_instances, num_joints, num_bodies, num_fixed_tendons, num_spatial_tendons, device, family ): + kind, capability = _TENDON_FAMILIES[family] + require_backend_capability(backend, capability) 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) + num_tendons = num_fixed_tendons if kind == "fixed" else num_spatial_tendons + assert getattr(art, f"num_{kind}_tendons") == num_tendons + names = getattr(art, f"{kind}_tendon_names") + assert isinstance(names, list) + assert len(names) == num_tendons + 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: + for name, property_family, dtype in _TENDON_DATA_PROPERTIES: + if property_family != family or num_tendons == 0: continue - expected_shape = (num_instances, counts[kind]) + expected_shape = (num_instances, num_tendons) 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. @@ -1299,34 +1385,64 @@ def test_articulation_tendon_contract( ("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] +# Writer cases pair each method with the dimension rows that configure at least one tendon of its kind. +_TENDON_WRITER_CASES = [ + pytest.param(*method, *dims, id=f"{method[1]}-{'-'.join(map(str, dims))}") + for method in _TENDON_METHODS + for dims in _TENDON_DIMS + if dims[3 if method[0] == "fixed" else 4] > 0 +] + + +def _push_tendon_properties(art, backend: str, kind: str) -> None: + """Write staged tendon properties to the backend so data read-backs see them on every backend.""" + if backend == "newton": + from isaaclab_newton.physics import NewtonManager + + with patch.object(NewtonManager, "add_model_change"): + getattr(art, f"write_{kind}_tendon_properties_to_sim_index")() + else: + getattr(art, f"write_{kind}_tendon_properties_to_sim_index")() 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]) + @pytest.mark.parametrize( + ("kind", "method_base", "kwarg", "getter", *_TENDON_DIM_NAMES), + _TENDON_WRITER_CASES, + ) def test_tendon_writer( self, + request, backend, - num_instances, - num_joints, - num_bodies, - num_fixed_tendons, - num_spatial_tendons, device, selection, kind, method_base, kwarg, getter, + num_instances, + num_joints, + num_bodies, + num_fixed_tendons, + num_spatial_tendons, ): + require_backend_capability(backend, _TENDON_FAMILIES[_TENDON_PROPERTY_FAMILY[getter]][1]) + if backend == "newton" and selection == "mask": + # Product bug: Newton's fixed-tendon mask setters slice full data without validating its shape, so an + # oversized input is silently truncated. Every assertion before the negative shape checks still runs. + request.applymarker( + pytest.mark.xfail( + raises=pytest.fail.Exception, + strict=True, + reason="Newton fixed-tendon mask setters accept data with extra environments", + ) + ) 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 ) @@ -1350,8 +1466,9 @@ def test_tendon_writer( 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 + # warp, all envs + all tendons: after the push, the matching data property reads the written values back method(**{kwarg: _make_payload_warp(full_shape, device, wp_dtype)}) + _push_tendon_properties(art, backend, kind) _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) @@ -1374,6 +1491,7 @@ def test_tendon_writer( @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.""" + require_backend_capability(backend, _TENDON_FAMILIES[kind][1]) num_instances = 3 art, raw_backend = get_articulation(backend, num_instances, 2, 2, 2, 2, device) art.data.update(dt=0.01) @@ -1384,7 +1502,7 @@ 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: + elif backend == "ovphysx": from isaaclab_ov import tensor_types as TT prefix = "FIXED_TENDON" if kind == "fixed" else "SPATIAL_TENDON" @@ -1404,14 +1522,28 @@ def capture(name, values, *, indices=None, mask=None): 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() if selection == "index": - write( - env_ids=None if len(selected_envs) == num_instances else torch.tensor(selected_envs, device=device) - ) + env_ids = None if len(selected_envs) == num_instances else torch.tensor(selected_envs, device=device) + write_kwargs = {"env_ids": env_ids} 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) + write_kwargs = {"env_mask": mask} + if backend == "newton": + from isaaclab_newton.physics import NewtonManager + + with patch.object(NewtonManager, "add_model_change"): + write(**write_kwargs) + 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])] + ) + else: + write(**write_kwargs) 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/_articulation_iface_test_utils.py b/source/isaaclab/test/assets/contract/_articulation_contract_utils.py similarity index 87% rename from source/isaaclab/test/assets/_articulation_iface_test_utils.py rename to source/isaaclab/test/assets/contract/_articulation_contract_utils.py index 25d10a44ee6e..ccaf1fb211e1 100644 --- a/source/isaaclab/test/assets/_articulation_iface_test_utils.py +++ b/source/isaaclab/test/assets/contract/_articulation_contract_utils.py @@ -6,61 +6,57 @@ # ignore private usage of variables warning # pyright: reportPrivateUsage=none -"""Shared mocked articulation backend factories for interface tests.""" +"""Shared mocked articulation backend factories for the contract 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] = {} +from ._manager_patch_scope import patch_contract_manager +from ._mock_backends import install_physx_recording_setters, patch_ovphysx_manager, patch_physx_manager +from .capabilities import available_backends + +BACKENDS = available_backends() -try: +if "physx" in BACKENDS: 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: + +if "newton" in BACKENDS: 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 +if "ovphysx" in BACKENDS: 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") + +_PHYSX_ARTICULATION_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 create_physx_articulation( @@ -75,6 +71,7 @@ def create_physx_articulation( body_ordering: tuple[str, ...] | None = None, ): """Create a test Articulation instance with mocked dependencies.""" + patch_physx_manager() 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)] @@ -102,7 +99,7 @@ def create_physx_articulation( max_spatial_tendons=num_spatial_tendons, ) mock_view.set_random_mock_data() - mock_view._noop_setters = True + install_physx_recording_setters(mock_view, _PHYSX_ARTICULATION_STORAGE) # Set up the mock view's metatype for accessing names/counts mock_metatype = MagicMock() @@ -119,7 +116,7 @@ def create_physx_articulation( # 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) + # Create ArticulationData instance (the PhysX manager is patched for this test above) data = PhysXArticulationData(mock_view, device) object.__setattr__(articulation, "_data", data) @@ -194,7 +191,9 @@ def create_physx_articulation( 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", wp.empty(N, dtype=wp.int32, device="cpu", pinned=wp.is_cuda_available()) + ) 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")) @@ -224,6 +223,7 @@ def create_ovphysx_articulation( body_ordering: tuple[str, ...] | None = None, ): """Create a test OvPhysX Articulation instance with mocked tensor bindings.""" + patch_ovphysx_manager() joint_names = [f"joint_{i}" for i in range(num_joints)] body_names = [f"body_{i}" for i in range(num_bodies)] @@ -300,8 +300,12 @@ def create_ovphysx_articulation( 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) + 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 @@ -329,6 +333,7 @@ def create_newton_articulation( num_instances=num_instances, num_bodies=num_bodies, num_joints=num_joints, + num_tendons=num_fixed_tendons, device=device, is_fixed_base=is_fixed_base, joint_names=joint_names, @@ -336,7 +341,16 @@ def create_newton_articulation( tendon_names=fixed_tendon_names, ) mock_view.set_random_mock_data() - mock_view._noop_setters = True + # MuJoCo tendon properties the data binds when the articulation has fixed tendons, with distinct values. + tendon_shape = (num_instances, 1, num_fixed_tendons) + tendon_values = np.arange(1, num_instances * num_fixed_tendons + 1, dtype=np.float32).reshape(tendon_shape) + tendon_range = np.stack((-tendon_values, tendon_values), axis=-1) + for name, values, dtype in ( + ("mujoco.tendon_stiffness", tendon_values, wp.float32), + ("mujoco.tendon_damping", tendon_values + 0.5, wp.float32), + ("mujoco.tendon_range", tendon_range, wp.vec2f), + ): + mock_view._attributes[name] = wp.array(values, dtype=dtype, device=device) # Mock NewtonManager (aliased as SimulationManager in Newton modules) mock_model = MagicMock() @@ -365,15 +379,9 @@ def create_newton_articulation( 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 + # Patch SimulationManager in the Newton data module until the test finishes. + patch_contract_manager(newton_data_module, "SimulationManager", mock_manager) + data = NewtonArticulationData(mock_view, device) # Create Articulation shell (bypass __init__) articulation = object.__new__(NewtonArticulation) @@ -449,7 +457,6 @@ def create_newton_articulation( return articulation, mock_view - def get_articulation( backend: str, num_instances: int = 2, diff --git a/source/isaaclab/test/assets/test_articulation_ordering_iface.py b/source/isaaclab/test/assets/contract/_articulation_ordering_contract_cases.py similarity index 98% rename from source/isaaclab/test/assets/test_articulation_ordering_iface.py rename to source/isaaclab/test/assets/contract/_articulation_ordering_contract_cases.py index 7fdbad251e57..ca836b6ad9d3 100644 --- a/source/isaaclab/test/assets/test_articulation_ordering_iface.py +++ b/source/isaaclab/test/assets/contract/_articulation_ordering_contract_cases.py @@ -5,7 +5,7 @@ # pyright: reportPrivateUsage=none -"""Mocked cross-backend articulation ordering interface tests.""" +"""Mocked cross-backend articulation joint- and body-ordering contract cases.""" from unittest.mock import MagicMock, patch @@ -13,12 +13,14 @@ 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 +from ._articulation_contract_utils import get_articulation +from .capabilities import backend_parameters, contract_backend, requires_backend + def _make_body_ordering_backend_data(num_instances: int, num_bodies: int) -> tuple[np.ndarray, ...]: """Create deterministic backend-order body data with identity rotations.""" @@ -895,37 +897,21 @@ def _get_backend_body_property_tensors(backend: str, art, raw_backend) -> dict[s 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))) + """Build an ordering parameter row that carries the backend's declared skip reason.""" + marks = [*kwargs.pop("marks", ()), *backend_parameters("ordering", names=(backend,))[0].marks] 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 -) +_requires_physx = requires_backend("physx") +_requires_ovphysx = requires_backend("ovphysx") +_requires_newton = requires_backend("newton") +_all_backends = contract_backend("ordering") +_physx_ovphysx_backends = contract_backend("ordering", names=("physx", "ovphysx")) +_dynamics_ordering_backends = contract_backend("ordering", names=("physx", "newton")) _non_mock_backends = _all_backends -class TestArticulationDataBodyState: +class TestArticulationOrderingBodyState: """Test data properties for all body states.""" @_non_mock_backends @@ -1859,7 +1845,7 @@ def test_ovphysx_partial_effort_target_write_preserves_unselected_backend_rows(s np.testing.assert_array_equal(pushed_after_partial_write, expected_backend) -class TestArticulationDataJointState: +class TestArticulationOrderingJointState: """Test data properties for joint state and joint properties.""" @_requires_ovphysx @@ -1975,7 +1961,7 @@ def _make_item_mask(total: int, selected: list[int], device: str) -> wp.array: # --------------------------------------------------------------------------- -class TestArticulationOperations: +class TestArticulationOrderingOperations: """Test cross-cutting articulation operations.""" @_non_mock_backends @@ -2230,7 +2216,7 @@ def test_physx_validate_cfg_reports_velocity_limits_in_public_joint_order(self): art._validate_cfg() -class TestArticulationWritersJoint: +class TestArticulationOrderingJointWriters: """Test joint writers/setters with all input combinations.""" @_non_mock_backends @@ -2350,7 +2336,7 @@ def test_reversed_joint_ordering_routes_property_writes_to_backend( ) -class TestArticulationWritersBody: +class TestArticulationOrderingBodyWriters: """Test body property writers/setters with all input combinations.""" @_non_mock_backends diff --git a/source/isaaclab/test/assets/contract/_contract_boot.py b/source/isaaclab/test/assets/contract/_contract_boot.py new file mode 100644 index 000000000000..cdf33bfd3a54 --- /dev/null +++ b/source/isaaclab/test/assets/contract/_contract_boot.py @@ -0,0 +1,59 @@ +# 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 Kit and kitless bootstrap for the asset contract tests.""" + +import os +import sys +from importlib.machinery import ModuleSpec +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: + + def _install_stub(module_name: str, *, is_package: bool = False) -> MagicMock: + """Install one import-boundary module double unless the real module is already loaded.""" + if module_name in sys.modules: + return sys.modules[module_name] + stub = MagicMock() + stub.__spec__ = ModuleSpec(module_name, loader=None, is_package=is_package) + if is_package: + stub.__path__ = [] + sys.modules[module_name] = stub + if "." in module_name: + parent_name, attribute = module_name.rsplit(".", 1) + setattr(sys.modules[parent_name], attribute, stub) + return stub + + # Worktree installs include the real ``isaaclab_physx`` package but not Kit's Python runtime. Stub only that + # external boundary so the contracts still import the real PhysX asset and data classes and their fixture views. + _install_stub("carb") + _install_stub("usdrt", is_package=True) + + # ``omni`` is a real namespace package in kitless runs. Install missing submodules in both ``sys.modules`` and the + # namespace attributes. + import omni as _omni # noqa: F401 + + for _module_name, _is_package in ( + ("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), + ): + _install_stub(_module_name, is_package=_is_package) + sys.modules["omni.kit.app"].get_app.return_value = None diff --git a/source/isaaclab/test/assets/contract/_manager_patch_scope.py b/source/isaaclab/test/assets/contract/_manager_patch_scope.py new file mode 100644 index 000000000000..243991ce51ca --- /dev/null +++ b/source/isaaclab/test/assets/contract/_manager_patch_scope.py @@ -0,0 +1,44 @@ +# 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 + +"""Test-local backend-manager patching for mocked asset factories.""" + +import inspect +from collections.abc import Iterator +from contextlib import contextmanager + +_MISSING = object() +_patch_scopes: list[dict[tuple[object, str], object]] = [] + + +@contextmanager +def contract_manager_patch_scope() -> Iterator[None]: + """Restore backend manager attributes changed by factories in this scope.""" + patches: dict[tuple[object, str], object] = {} + _patch_scopes.append(patches) + try: + yield + finally: + for (owner, name), original in reversed(patches.items()): + if original is _MISSING: + delattr(owner, name) + else: + setattr(owner, name, original) + popped = _patch_scopes.pop() + assert popped is patches + + +def patch_contract_manager(owner: object, name: str, replacement: object) -> None: + """Patch a manager binding until the current contract test finishes.""" + if not _patch_scopes: + raise RuntimeError("contract factories must run inside contract_manager_patch_scope()") + patches = _patch_scopes[-1] + key = (owner, name) + if key not in patches: + try: + patches[key] = inspect.getattr_static(owner, name) + except AttributeError: + patches[key] = _MISSING + setattr(owner, name, replacement) diff --git a/source/isaaclab/test/assets/contract/_mock_backends.py b/source/isaaclab/test/assets/contract/_mock_backends.py new file mode 100644 index 000000000000..77c968ddfa7c --- /dev/null +++ b/source/isaaclab/test/assets/contract/_mock_backends.py @@ -0,0 +1,70 @@ +# 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-manager and tensor-view doubles shared by the mocked contract factories.""" + +from unittest.mock import MagicMock + +import numpy as np +import warp as wp + +from ._manager_patch_scope import patch_contract_manager + + +def patch_physx_manager() -> 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) + patch_contract_manager(PhysxManager, "get_physics_sim_view", MagicMock(return_value=physics_sim_view)) + # Writers bump the scene-data transform version that ``initialize()`` would normally create. + patch_contract_manager(PhysxManager, "_scene_data_backend", PhysxSceneDataBackend()) + + +def patch_ovphysx_manager() -> None: + """Create the OVPhysX scene-data backend whose transform version the writers bump.""" + from isaaclab_ov.physics.ovphysx_manager import OvPhysxManager, OvPhysxSceneDataBackend + + patch_contract_manager(OvPhysxManager, "_scene_data_backend", OvPhysxSceneDataBackend()) + + +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 accept both layouts + 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): + 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: + rows = indices.numpy().astype(np.int64) + if values_np.size == stored_np.size: + stored_np[rows] = values_np.reshape(stored_np.shape)[rows] + else: + stored_np[rows] = values_np.reshape(len(rows), *stored_np.shape[1:]) + 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)) diff --git a/source/isaaclab/test/assets/test_rigid_object_collection_iface.py b/source/isaaclab/test/assets/contract/_rigid_object_collection_contract_cases.py similarity index 87% rename from source/isaaclab/test/assets/test_rigid_object_collection_iface.py rename to source/isaaclab/test/assets/contract/_rigid_object_collection_contract_cases.py index 35604f52d9f3..210f2f43858f 100644 --- a/source/isaaclab/test/assets/test_rigid_object_collection_iface.py +++ b/source/isaaclab/test/assets/contract/_rigid_object_collection_contract_cases.py @@ -6,12 +6,10 @@ # 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. +"""Shared rigid-object-collection contract cases, collected by the ``test_asset_contract_*`` entry modules. -The setup is a bit convoluted so that we can run these tests without requiring Isaac Sim or GPU simulation. +Checks that every collection backend provides the data and writer behavior the base collection class advertises. +The backends run on mocked views, so these cases need neither Isaac Sim nor a GPU simulation. """ import math @@ -21,11 +19,11 @@ 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 +from ._rigid_object_collection_contract_utils import get_rigid_object_collection +from .capabilities import contract_backend # Distinct instance and body counts make swapped axes visible. _NUM_INSTANCES, _NUM_BODIES = 2, 3 @@ -47,17 +45,11 @@ def _check_proxy_array(arr, *, expected_shape: tuple, expected_dtype: type, name # 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) +_api_backends = contract_backend("api") +_data_backends = contract_backend("data") +_writes_backends = contract_backend("writes") +_index_resolution_backends = contract_backend("index_resolution") _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 -) # --------------------------------------------------------------------------- @@ -196,7 +188,7 @@ def _make_item_mask(total: int, selected: list[int], device: str) -> wp.array: class TestCollectionIndexResolution: """Test backend-specific index resolution helpers.""" - @_production_backends + @_api_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) @@ -234,7 +226,7 @@ def test_resolve_body_ids_handles_tensor_view_shape(self, backend): class TestCollectionViewReshape: """Test backend-specific view reshape helpers.""" - @_reshape_3d_backends + @_data_backends @_devices @pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) def test_reshape_data_to_view_3d_accepts_torch_tensor(self, backend, device, dtype): @@ -263,7 +255,7 @@ def test_reshape_data_to_view_3d_accepts_torch_tensor(self, backend, device, dty assert str(warp_view.device) == device torch.testing.assert_close(wp.to_torch(warp_view), view) - @_reshape_3d_backends + @_data_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 @@ -290,7 +282,7 @@ def test_reshape_data_to_view_3d_moves_torch_tensor_to_requested_device(self, ba class TestCollectionProperties: """Test that collection properties return the correct types/values.""" - @_backends + @_api_backends def test_collection_counts_and_names(self, backend): from isaaclab.assets.rigid_object_collection.base_rigid_object_collection_data import ( BaseRigidObjectCollectionData, @@ -310,7 +302,7 @@ def test_collection_counts_and_names(self, backend): class TestCollectionFinderReturnModes: """Test finder return modes on production collection backends.""" - @_production_backends + @_api_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") @@ -331,7 +323,7 @@ def test_find_bodies_returns_legacy_tensor_or_cached_proxy(self, backend): _, first_names = collection.find_bodies(first_body) assert first_names == [first_body] - @_production_backends + @_api_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") @@ -387,7 +379,7 @@ def test_find_objects_forwards_return_mode_with_alias_warning(self, backend): class TestCollectionDataProperties: """Test that every data property is a ProxyArray with the advertised shape and dtype.""" - @_backends + @_data_backends @_devices def test_collection_data_property_contract(self, backend, device): obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device) @@ -417,7 +409,7 @@ def test_collection_data_property_contract(self, backend, device): class TestCollectionCacheInvalidation: - @_production_backends + @_writes_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) @@ -437,7 +429,7 @@ def test_pose_write_invalidates_pose_dependent_caches(self, backend): obj.write_body_link_pose_to_sim_index(body_poses=body_pose) _assert_buffers_stale(obj.data, buffers) - @_production_backends + @_writes_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) @@ -454,7 +446,7 @@ def test_velocity_write_invalidates_body_frame_caches(self, backend): obj.write_body_com_velocity_to_sim_index(body_velocities=body_velocity) _assert_buffers_stale(obj.data, buffers) - @_production_backends + @_writes_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") @@ -499,7 +491,7 @@ class TestCollectionWritersPose: # -- index variants for pose -- - @_backends + @_writes_backends @_devices @pytest.mark.parametrize("method_suffix", _BODY_POSE_METHODS) def test_write_body_pose_to_sim_index(self, backend, device, method_suffix): @@ -542,7 +534,7 @@ def test_write_body_pose_to_sim_index(self, backend, device, method_suffix): # -- index variants for velocity -- - @_backends + @_writes_backends @_devices @pytest.mark.parametrize("method_suffix", _BODY_VEL_METHODS) def test_write_body_velocity_to_sim_index(self, backend, device, method_suffix): @@ -589,7 +581,7 @@ def test_write_body_velocity_to_sim_index(self, backend, device, method_suffix): # -- mask variants for pose -- - @_backends + @_writes_backends @_devices @pytest.mark.parametrize("method_suffix", _BODY_POSE_METHODS) def test_write_body_pose_to_sim_mask(self, backend, device, method_suffix): @@ -635,7 +627,7 @@ def test_write_body_pose_to_sim_mask(self, backend, device, method_suffix): # -- mask variants for velocity -- - @_backends + @_writes_backends @_devices @pytest.mark.parametrize("method_suffix", _BODY_VEL_METHODS) def test_write_body_velocity_to_sim_mask(self, backend, device, method_suffix): @@ -717,7 +709,7 @@ def _make_body_warp(shape: tuple[int, int], device: str, wp_dtype: type, trailin class TestCollectionWritersBody: """Test body property writers/setters with all input combinations.""" - @_backends + @_writes_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): @@ -757,7 +749,7 @@ def test_body_writer_index(self, backend, device, method_base, kwarg): with pytest.raises((AssertionError, RuntimeError)): method(**{kwarg: _make_body_warp((num_instances + 1, num_bodies), device, wp_dtype, trailing)}) - @_backends + @_writes_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): @@ -820,7 +812,7 @@ def test_body_writer_mask(self, backend, device, method_base, kwarg): class TestCollectionDataAliases: """Test that alias properties return the values of their canonical counterparts.""" - @_backends + @_data_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") @@ -832,3 +824,86 @@ def test_aliases_match_canonical_values(self, backend): 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 + + +# --------------------------------------------------------------------------- +# Tests: partial writes reach only the selected backend cells +# --------------------------------------------------------------------------- + +# quantity -> (writer base name, writer keyword, trailing size) +_PARTIAL_WRITES = { + "body_pose": ("write_body_link_pose_to_sim", "body_poses", 7), + "body_velocity": ("write_body_com_velocity_to_sim", "body_velocities", 6), + "mass": ("set_masses", "masses", 1), +} + + +def _read_backend_cells(backend: str, raw_backend, quantity: str, num_instances: int, num_bodies: int): + """Read one collection quantity from backend storage in public ``(env, body, ...)`` order.""" + if backend == "physx": + getter = {"body_pose": "get_transforms", "body_velocity": "get_velocities", "mass": "get_masses"}[quantity] + # PhysX collection views are body-major. + values = wp.to_torch(getattr(raw_backend, getter)()).reshape(num_bodies, num_instances, -1).transpose(0, 1) + elif backend == "newton": + if quantity == "mass": + values = wp.to_torch(raw_backend.get_attribute("body_mass", None)) + else: + getter = {"body_pose": "get_root_transforms", "body_velocity": "get_root_velocities"}[quantity] + values = wp.to_torch(getattr(raw_backend, getter)(None)) + else: + from isaaclab_ov import tensor_types as TT + + binding = {"body_pose": TT.LINK_POSE, "body_velocity": TT.LINK_VELOCITY, "mass": TT.BODY_MASS}[quantity] + values = torch.as_tensor(raw_backend.bindings[binding]._data) + return values.reshape(num_instances, num_bodies, -1).cpu().clone() + + +class TestCollectionPartialWriteCells: + """Test that partial writers reach only the selected cells of backend storage.""" + + @_writes_backends + @pytest.mark.parametrize("selection", ["index", "mask"]) + @pytest.mark.parametrize("quantity", _PARTIAL_WRITES) + def test_partial_write_preserves_unselected_backend_cells(self, request, backend, selection, quantity): + if backend == "ovphysx" 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", + ) + ) + num_instances, num_bodies = _NUM_INSTANCES, _NUM_BODIES + obj, raw_backend = get_rigid_object_collection(backend, num_instances, num_bodies, "cpu") + writer, kwarg, trailing = _PARTIAL_WRITES[quantity] + # Literal, per-cell distinct payload; poses keep an identity rotation. + values = 100.0 * torch.arange(1, num_instances * num_bodies + 1, dtype=torch.float32).reshape( + num_instances, num_bodies, 1 + ) + torch.arange(trailing, dtype=torch.float32) + if quantity == "body_pose": + values[..., 3:6] = 0.0 + values[..., 6] = 1.0 + if quantity == "mass": + values = values[..., 0] + before = _read_backend_cells(backend, raw_backend, quantity, num_instances, num_bodies) + # Select the second environment and the last body so identity or transposed routing cannot pass. + env, body = 1, 2 + if selection == "index": + getattr(obj, f"{writer}_index")( + **{kwarg: values[env : env + 1, body : body + 1]}, + env_ids=torch.tensor([env], dtype=torch.int32), + body_ids=_make_body_ids("cpu", [body]), + ) + else: + getattr(obj, f"{writer}_mask")( + **{kwarg: values}, + env_mask=_make_item_mask(num_instances, [env], "cpu"), + body_mask=_make_item_mask(num_bodies, [body], "cpu"), + ) + + after = _read_backend_cells(backend, raw_backend, quantity, num_instances, num_bodies) + expected = before.clone() + expected[env, body] = values[env, body].reshape(-1) + torch.testing.assert_close(after, expected, rtol=0.0, atol=0.0) diff --git a/source/isaaclab/test/assets/_rigid_object_collection_iface_test_utils.py b/source/isaaclab/test/assets/contract/_rigid_object_collection_contract_utils.py similarity index 84% rename from source/isaaclab/test/assets/_rigid_object_collection_iface_test_utils.py rename to source/isaaclab/test/assets/contract/_rigid_object_collection_contract_utils.py index e7593e8c6b8e..371b3e574d25 100644 --- a/source/isaaclab/test/assets/_rigid_object_collection_iface_test_utils.py +++ b/source/isaaclab/test/assets/contract/_rigid_object_collection_contract_utils.py @@ -6,12 +6,10 @@ # ignore private usage of variables warning # pyright: reportPrivateUsage=none -"""Shared mocked rigid-object-collection backend factories for interface tests.""" +"""Shared mocked rigid-object-collection backend factories for the contract tests.""" from unittest.mock import MagicMock -import _iface_test_boot # noqa: F401 (starts the runtime) - import numpy as np import warp as wp @@ -19,30 +17,22 @@ from isaaclab.assets.rigid_object_collection.rigid_object_collection_cfg import RigidObjectCollectionCfg from isaaclab.utils.wrench_composer import WrenchComposer -BACKENDS: list[str] = [] +from ._manager_patch_scope import patch_contract_manager +from ._mock_backends import install_physx_recording_setters, patch_ovphysx_manager, patch_physx_manager +from .capabilities import available_backends + +BACKENDS = available_backends() -try: +if "physx" in BACKENDS: 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: + +if "newton" in BACKENDS: from isaaclab_newton.assets.rigid_object_collection.rigid_object_collection import ( RigidObjectCollection as NewtonRigidObjectCollection, ) @@ -50,14 +40,8 @@ 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 +if "ovphysx" in BACKENDS: from isaaclab_ov.assets.rigid_object_collection.rigid_object_collection import ( RigidObjectCollection as OvPhysxRigidObjectCollection, ) @@ -65,14 +49,14 @@ 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") + +_PHYSX_RIGID_BODY_STORAGE = { + "set_transforms": "_transforms", + "set_velocities": "_velocities", + "set_masses": "_masses", + "set_coms": "_coms", + "set_inertias": "_inertias", +} def create_physx_rigid_object_collection( @@ -81,6 +65,7 @@ def create_physx_rigid_object_collection( device: str = "cuda:0", ): """Create a test RigidObjectCollection instance with mocked dependencies.""" + patch_physx_manager() collection = object.__new__(PhysXRigidObjectCollection) rigid_objects = {f"object_{i}": RigidObjectCfg(prim_path=f"/World/Object_{i}") for i in range(num_bodies)} @@ -92,7 +77,7 @@ def create_physx_rigid_object_collection( device=device, ) mock_view.set_random_mock_data() - mock_view._noop_setters = True + install_physx_recording_setters(mock_view, _PHYSX_RIGID_BODY_STORAGE) object.__setattr__(collection, "_root_view", mock_view) object.__setattr__(collection, "_device", device) @@ -127,10 +112,12 @@ def create_physx_rigid_object_collection( 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) + cpu_all_view_ids = wp.empty(num_view_ids, dtype=wp.int32, device="cpu", pinned=wp.is_cuda_available()) 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", wp.empty(num_view_ids, dtype=wp.int32, device="cpu", pinned=wp.is_cuda_available()) + ) object.__setattr__(collection, "_cpu_view_ids_views", {}) return collection, mock_view @@ -155,7 +142,6 @@ def create_newton_rigid_object_collection( 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() @@ -174,17 +160,10 @@ def create_newton_rigid_object_collection( 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 + # Patch SimulationManager in both data and collection modules until the test finishes. + patch_contract_manager(newton_data_module, "SimulationManager", mock_manager) + patch_contract_manager(newton_coll_module, "SimulationManager", mock_manager) + data = NewtonRigidObjectCollectionData(mock_view, num_bodies, device) # Create collection shell (bypass __init__) collection = object.__new__(NewtonRigidObjectCollection) @@ -229,6 +208,7 @@ def create_ovphysx_rigid_object_collection( device: str = "cuda:0", ): """Create a test OVPhysX RigidObjectCollection instance with mocked tensor bindings.""" + patch_ovphysx_manager() body_names = [f"object_{i}" for i in range(num_bodies)] collection = object.__new__(OvPhysxRigidObjectCollection) @@ -279,7 +259,6 @@ def create_ovphysx_rigid_object_collection( return collection, mock_bindings - def get_rigid_object_collection( backend: str, num_instances: int = 2, diff --git a/source/isaaclab/test/assets/test_rigid_object_iface.py b/source/isaaclab/test/assets/contract/_rigid_object_contract_cases.py similarity index 88% rename from source/isaaclab/test/assets/test_rigid_object_iface.py rename to source/isaaclab/test/assets/contract/_rigid_object_contract_cases.py index 47d02fefe868..223a4224f039 100644 --- a/source/isaaclab/test/assets/test_rigid_object_iface.py +++ b/source/isaaclab/test/assets/contract/_rigid_object_contract_cases.py @@ -6,11 +6,10 @@ # 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. +"""Shared rigid-object contract cases, collected by the ``test_asset_contract_*`` entry modules. -The setup is a bit convoluted so that we can run these tests without requiring Isaac Sim or GPU simulation. +Checks that every rigid-object backend provides the data and writer behavior the base rigid-object class advertises. +The backends run on mocked views, so these cases need neither Isaac Sim nor a GPU simulation. """ import math @@ -20,12 +19,11 @@ 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 - +from ._rigid_object_contract_utils import get_rigid_object +from .capabilities import contract_backend # --------------------------------------------------------------------------- # Helpers @@ -43,12 +41,11 @@ def _check_proxy_array(arr, *, expected_shape: tuple, expected_dtype: type, name # 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) +_api_backends = contract_backend("api") +_data_backends = contract_backend("data") +_writes_backends = contract_backend("writes") _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 -) # --------------------------------------------------------------------------- @@ -59,7 +56,7 @@ def _check_proxy_array(arr, *, expected_shape: tuple, expected_dtype: type, name class TestRigidObjectIndexResolution: """Test backend-specific index resolution helpers.""" - @_production_backends + @_api_backends @_devices def test_resolve_env_ids_handles_tensor_view_shape(self, backend, device): obj, _ = get_rigid_object(backend, num_instances=4, device=device) @@ -86,7 +83,7 @@ def test_resolve_env_ids_handles_tensor_view_shape(self, backend, device): class TestRigidObjectProperties: """Test that rigid object properties return the correct types/values.""" - @_backends + @_api_backends def test_rigid_object_counts_and_names(self, backend): from isaaclab.assets.rigid_object.base_rigid_object_data import BaseRigidObjectData @@ -104,7 +101,7 @@ def test_rigid_object_counts_and_names(self, backend): class TestRigidObjectFinderReturnModes: """Test finder return modes on production rigid-object backends.""" - @_production_backends + @_api_backends def test_find_bodies_returns_legacy_list_or_cached_proxy(self, backend): obj, _ = get_rigid_object(backend, num_instances=2, device="cpu") @@ -179,7 +176,7 @@ def test_find_bodies_returns_legacy_list_or_cached_proxy(self, backend): class TestRigidObjectDataProperties: """Test that every data property is a ProxyArray with the advertised shape and dtype.""" - @_backends + @_data_backends @_devices def test_rigid_object_data_property_contract(self, backend, device): obj, _ = get_rigid_object(backend, _NUM_INSTANCES, device) @@ -226,7 +223,7 @@ def test_rigid_object_data_property_contract(self, backend, device): class TestRigidObjectDataAliases: """Test that alias properties return the values of their canonical counterparts.""" - @_backends + @_data_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") @@ -383,7 +380,7 @@ def _make_item_mask(total: int, selected: list[int], device: str) -> wp.array: class TestRigidObjectCacheInvalidation: - @_production_backends + @_writes_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) @@ -403,7 +400,7 @@ def test_pose_write_invalidates_pose_dependent_caches(self, backend): obj.write_root_link_pose_to_sim_index(root_pose=root_pose) _assert_buffers_stale(obj.data, buffers) - @_production_backends + @_writes_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) @@ -424,7 +421,7 @@ def test_velocity_write_invalidates_body_frame_caches(self, backend): _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 + @_writes_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") @@ -469,7 +466,7 @@ class TestRigidObjectWritersRoot: # -- index variants -- - @_backends + @_writes_backends @_devices @pytest.mark.parametrize("method_suffix", _ROOT_POSE_METHODS) def test_write_root_pose_to_sim_index(self, backend, device, method_suffix): @@ -497,7 +494,7 @@ def test_write_root_pose_to_sim_index(self, backend, device, method_suffix): with pytest.raises((AssertionError, RuntimeError)): method(root_pose=_make_bad_data_warp((num_instances,), device, wp.transformf)) - @_backends + @_writes_backends @_devices @pytest.mark.parametrize("method_suffix", _ROOT_VEL_METHODS) def test_write_root_velocity_to_sim_index(self, backend, device, method_suffix): @@ -527,7 +524,7 @@ def test_write_root_velocity_to_sim_index(self, backend, device, method_suffix): # -- mask variants -- - @_backends + @_writes_backends @_devices @pytest.mark.parametrize("method_suffix", _ROOT_POSE_METHODS) def test_write_root_pose_to_sim_mask(self, backend, device, method_suffix): @@ -561,7 +558,7 @@ def test_write_root_pose_to_sim_mask(self, backend, device, method_suffix): with pytest.raises((AssertionError, RuntimeError)): method(root_pose=_make_bad_data_warp((num_instances,), device, wp.transformf)) - @_backends + @_writes_backends @_devices @pytest.mark.parametrize("method_suffix", _ROOT_VEL_METHODS) def test_write_root_velocity_to_sim_mask(self, backend, device, method_suffix): @@ -603,7 +600,7 @@ def test_write_root_velocity_to_sim_mask(self, backend, device, method_suffix): _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]: +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" @@ -612,9 +609,6 @@ def _body_writer_layout(backend: str, method_base: str) -> tuple[type, int, str 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" @@ -642,7 +636,7 @@ def _make_body_warp(shape: tuple[int, int], device: str, wp_dtype: type, trailin class TestRigidObjectWritersBody: """Test body property writers/setters with all input combinations.""" - @_production_backends + @_writes_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 @@ -692,7 +686,7 @@ def test_external_wrench_frames(self, backend): 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 + @_writes_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): @@ -716,9 +710,8 @@ def test_body_writer_index(self, backend, device, method_base, kwarg): ) # 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) + 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( **{ @@ -734,7 +727,7 @@ def test_body_writer_index(self, backend, device, method_base, kwarg): with pytest.raises((AssertionError, RuntimeError)): method(**{kwarg: _make_body_warp((num_instances + 1, num_bodies), device, wp_dtype, trailing)}) - @_backends + @_writes_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): @@ -757,9 +750,8 @@ def test_body_writer_mask(self, backend, device, method_base, kwarg): ) # 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) + 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( **{ @@ -774,3 +766,72 @@ def test_body_writer_mask(self, backend, device, method_base, kwarg): # negative: bad warp shape with pytest.raises((AssertionError, RuntimeError)): method(**{kwarg: _make_body_warp((num_instances + 1, num_bodies), device, wp_dtype, trailing)}) + + +# --------------------------------------------------------------------------- +# Tests: partial writes reach only the selected backend rows +# --------------------------------------------------------------------------- + +# quantity -> (writer base name, writer keyword, literal per-environment payload) +_PARTIAL_WRITES = { + "root_pose": ( + "write_root_link_pose_to_sim", + "root_pose", + [[31.0, 32.0, 33.0, 0.0, 0.0, 0.0, 1.0], [41.0, 42.0, 43.0, 0.0, 0.0, 0.0, 1.0]], + ), + "root_velocity": ( + "write_root_com_velocity_to_sim", + "root_velocity", + [[21.0, 22.0, 23.0, 24.0, 25.0, 26.0], [71.0, 72.0, 73.0, 74.0, 75.0, 76.0]], + ), + "mass": ("set_masses", "masses", [[51.0], [61.0]]), +} + + +def _read_backend_rows(backend: str, raw_backend, quantity: str) -> torch.Tensor: + """Read one rigid-object quantity from backend storage as one row per environment.""" + if backend == "physx": + getter = {"root_pose": "get_transforms", "root_velocity": "get_velocities", "mass": "get_masses"}[quantity] + values = getattr(raw_backend, getter)() + elif backend == "newton": + if quantity == "mass": + values = raw_backend.get_attribute("body_mass", None) + else: + getter = {"root_pose": "get_root_transforms", "root_velocity": "get_root_velocities"}[quantity] + values = getattr(raw_backend, getter)(None) + else: + from isaaclab_ov import tensor_types as TT + + binding = {"root_pose": TT.RIGID_BODY_POSE, "root_velocity": TT.RIGID_BODY_VELOCITY, "mass": TT.RIGID_BODY_MASS} + values = raw_backend.bindings[binding[quantity]]._data + values = wp.to_torch(values) if isinstance(values, wp.array) else torch.as_tensor(values) + return values.reshape(2, -1).cpu().clone() + + +class TestRigidObjectPartialWriteRows: + """Test that partial writers reach only the selected rows of backend storage.""" + + @_writes_backends + @pytest.mark.parametrize("selection", ["index", "mask"]) + @pytest.mark.parametrize("quantity", _PARTIAL_WRITES) + def test_partial_write_preserves_unselected_backend_rows(self, backend, selection, quantity): + obj, raw_backend = get_rigid_object(backend, num_instances=2, device="cpu") + writer, kwarg, payload = _PARTIAL_WRITES[quantity] + values = torch.tensor(payload, dtype=torch.float32) + before = _read_backend_rows(backend, raw_backend, quantity) + # Select the second environment so a writer that ignores the selection cannot pass by writing row 0. + selected, unselected = 1, 0 + if selection == "index": + kwargs = {kwarg: values[selected : selected + 1], "env_ids": torch.tensor([selected], dtype=torch.int32)} + if quantity == "mass": + kwargs["body_ids"] = [0] + else: + kwargs = {kwarg: values, "env_mask": _make_item_mask(2, [selected], "cpu")} + if quantity == "mass": + kwargs["body_mask"] = _make_item_mask(1, [0], "cpu") + + getattr(obj, f"{writer}_{selection}")(**kwargs) + + after = _read_backend_rows(backend, raw_backend, quantity) + torch.testing.assert_close(after[selected], values[selected], rtol=0.0, atol=0.0) + torch.testing.assert_close(after[unselected], before[unselected], rtol=0.0, atol=0.0) diff --git a/source/isaaclab/test/assets/_rigid_object_iface_test_utils.py b/source/isaaclab/test/assets/contract/_rigid_object_contract_utils.py similarity index 84% rename from source/isaaclab/test/assets/_rigid_object_iface_test_utils.py rename to source/isaaclab/test/assets/contract/_rigid_object_contract_utils.py index b58eae0b64c6..836e361973a8 100644 --- a/source/isaaclab/test/assets/_rigid_object_iface_test_utils.py +++ b/source/isaaclab/test/assets/contract/_rigid_object_contract_utils.py @@ -6,59 +6,44 @@ # ignore private usage of variables warning # pyright: reportPrivateUsage=none -"""Shared mocked rigid-object backend factories for interface tests.""" +"""Shared mocked rigid-object backend factories for the contract 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] = [] +from ._manager_patch_scope import patch_contract_manager +from ._mock_backends import install_physx_recording_setters, patch_ovphysx_manager, patch_physx_manager +from .capabilities import available_backends + +BACKENDS = available_backends() -try: +if "physx" in BACKENDS: 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: + +if "newton" in BACKENDS: 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 +if "ovphysx" in BACKENDS: 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") + +_PHYSX_RIGID_BODY_STORAGE = { + "set_transforms": "_transforms", + "set_velocities": "_velocities", + "set_masses": "_masses", + "set_coms": "_coms", + "set_inertias": "_inertias", +} def create_physx_rigid_object( @@ -66,6 +51,7 @@ def create_physx_rigid_object( device: str = "cuda:0", ): """Create a test RigidObject instance with mocked dependencies.""" + patch_physx_manager() body_names = ["body_0"] rigid_object = object.__new__(PhysXRigidObject) @@ -78,12 +64,12 @@ def create_physx_rigid_object( device=device, ) mock_view.set_random_mock_data() - mock_view._noop_setters = True + install_physx_recording_setters(mock_view, _PHYSX_RIGID_BODY_STORAGE) object.__setattr__(rigid_object, "_root_view", mock_view) object.__setattr__(rigid_object, "_device", device) - # Create RigidObjectData instance (SimulationManager already mocked at module level) + # Create RigidObjectData instance (the PhysX manager is patched for this test above) data = PhysXRigidObjectData(mock_view, device) object.__setattr__(rigid_object, "_data", data) @@ -116,7 +102,9 @@ def create_physx_rigid_object( 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", wp.empty(N, dtype=wp.int32, device="cpu", pinned=wp.is_cuda_available()) + ) 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")) @@ -145,7 +133,6 @@ def create_newton_rigid_object( 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() @@ -164,14 +151,9 @@ def create_newton_rigid_object( 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 + # Patch SimulationManager in the Newton data module until the test finishes. + patch_contract_manager(newton_data_module, "SimulationManager", mock_manager) + data = NewtonRigidObjectData(mock_view, device) # Create RigidObject shell (bypass __init__) rigid_object = object.__new__(NewtonRigidObject) @@ -210,6 +192,7 @@ def create_ovphysx_rigid_object( device: str = "cuda:0", ): """Create a test OvPhysX RigidObject instance with mocked tensor bindings.""" + patch_ovphysx_manager() body_names = ["base_link"] obj = object.__new__(OvPhysxRigidObject) @@ -262,7 +245,6 @@ def create_ovphysx_rigid_object( return obj, mock_bindings - def get_rigid_object( backend: str, num_instances: int = 2, diff --git a/source/isaaclab/test/assets/contract/capabilities.py b/source/isaaclab/test/assets/contract/capabilities.py new file mode 100644 index 000000000000..8c5dbaf2b07d --- /dev/null +++ b/source/isaaclab/test/assets/contract/capabilities.py @@ -0,0 +1,152 @@ +# 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 + +"""Explicit backend capability declarations for the shared asset contracts. + +Every contract test is parametrized over all declared backends. A backend that is unavailable in this process, or +that declares a capability unsupported, still collects as a skipped case with the declared reason, so a missing +backend or an unsupported branch never silently disappears from the matrix. +""" + +from dataclasses import dataclass, field +from importlib.util import find_spec + +import pytest +import warp as wp +from _pytest.mark.structures import ParameterSet + +# Install the kitless import stubs before resolving the declared backend dependencies. +from . import _contract_boot # noqa: F401 + + +@dataclass(frozen=True) +class BackendDeclaration: + """Dependencies and supported contract capabilities of one backend.""" + + name: str + required_modules: tuple[str, ...] + capabilities: frozenset[str] + unsupported: dict[str, str] = field(default_factory=dict) + requires_cuda_runtime: bool = False + + +@dataclass(frozen=True) +class BackendStatus: + """Whether one declared backend can run the contracts in this process.""" + + declaration: BackendDeclaration + available: bool + reason: str | None + + +_SHARED_CAPABILITIES = frozenset({"api", "data", "writes", "ordering", "fixed_tendons"}) +_PHYSX_FAMILY_CAPABILITIES = frozenset( + {"fixed_tendon_extended", "spatial_tendons", "fixed_tendon_target_scheduling", "mocked_dynamics"} +) + +BACKEND_DECLARATIONS = ( + BackendDeclaration( + name="physx", + required_modules=("carb", "isaaclab_physx"), + capabilities=_SHARED_CAPABILITIES | _PHYSX_FAMILY_CAPABILITIES | {"index_resolution"}, + ), + BackendDeclaration( + name="newton", + required_modules=("isaaclab_newton",), + capabilities=_SHARED_CAPABILITIES | {"index_resolution"}, + unsupported={ + "fixed_tendon_extended": "Newton does not implement fixed-tendon limit stiffness, rest length, or offset", + "spatial_tendons": "Newton does not support spatial tendons", + "fixed_tendon_target_scheduling": ( + "Newton commands tendon targets through a MuJoCo tendon actuator the mocked contract does not build" + ), + "mocked_dynamics": "Newton computes task-space dynamics from a solver model that the mocked contract lacks", + }, + ), + BackendDeclaration( + name="ovphysx", + required_modules=("ovphysx", "isaaclab_ov"), + capabilities=_SHARED_CAPABILITIES | _PHYSX_FAMILY_CAPABILITIES, + unsupported={"index_resolution": "OVPhysX does not expose the joint and body index-resolution helpers"}, + # The mocked bindings allocate pinned host staging buffers even for CPU tensors. + requires_cuda_runtime=True, + ), +) + + +def _module_available(module_name: str) -> bool: + """Return whether an explicitly declared dependency can be resolved.""" + try: + return find_spec(module_name) is not None + except ModuleNotFoundError: + return False + + +def _evaluate_backend(declaration: BackendDeclaration, cuda_available: bool) -> BackendStatus: + """Evaluate one backend declaration against the current process.""" + for module_name in declaration.required_modules: + if not _module_available(module_name): + return BackendStatus(declaration, available=False, reason=f"missing required module: {module_name}") + if declaration.requires_cuda_runtime and not cuda_available: + return BackendStatus(declaration, available=False, reason="CUDA runtime unavailable") + return BackendStatus(declaration, available=True, reason=None) + + +BACKEND_STATUSES = tuple(_evaluate_backend(declaration, wp.is_cuda_available()) for declaration in BACKEND_DECLARATIONS) +_STATUS_BY_NAME = {status.declaration.name: status for status in BACKEND_STATUSES} + + +def _unsupported_reason(declaration: BackendDeclaration, capability: str) -> str | None: + """Return why a backend does not run a capability, or None when it does.""" + if capability in declaration.capabilities: + return None + reason = declaration.unsupported.get(capability) + if reason is None: + raise ValueError(f"{declaration.name} has no declaration for capability {capability!r}") + return reason + + +def backend_parameters(capability: str, *, names: tuple[str, ...] | None = None) -> list[ParameterSet]: + """Return one pytest parameter per declared backend, skipped with its reason when it cannot run. + + Args: + capability: Contract capability the test exercises. + names: Optional subset of backend names for a backend-specific branch. Defaults to every backend. + + Returns: + Backend parameters in declaration order. + """ + parameters = [] + for status in BACKEND_STATUSES: + declaration = status.declaration + if names is not None and declaration.name not in names: + continue + reason = status.reason or _unsupported_reason(declaration, capability) + marks = () if reason is None else pytest.mark.skip(reason=reason) + parameters.append(pytest.param(declaration.name, marks=marks, id=declaration.name)) + return parameters + + +def contract_backend(capability: str, *, names: tuple[str, ...] | None = None) -> pytest.MarkDecorator: + """Parametrize ``backend`` over the declared backends for one capability.""" + return pytest.mark.parametrize("backend", backend_parameters(capability, names=names)) + + +def requires_backend(name: str) -> pytest.MarkDecorator: + """Skip a backend-specific test with the declared reason when that backend is unavailable.""" + reason = _STATUS_BY_NAME[name].reason + return pytest.mark.skipif(reason is not None, reason=f"{name}: {reason}") + + +def require_backend_capability(backend: str, capability: str) -> None: + """Skip the current case with the declared reason when a backend lacks a capability.""" + reason = _unsupported_reason(_STATUS_BY_NAME[backend].declaration, capability) + if reason is not None: + pytest.skip(reason) + + +def available_backends() -> list[str]: + """Return the declared backends that can run in this process.""" + return [status.declaration.name for status in BACKEND_STATUSES if status.available] diff --git a/source/isaaclab/test/assets/contract/conftest.py b/source/isaaclab/test/assets/contract/conftest.py new file mode 100644 index 000000000000..b8dc619426b1 --- /dev/null +++ b/source/isaaclab/test/assets/contract/conftest.py @@ -0,0 +1,19 @@ +# 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 + +"""Fixtures shared by the focused asset contract gate.""" + +from collections.abc import Iterator + +import pytest + +from ._manager_patch_scope import contract_manager_patch_scope + + +@pytest.fixture(autouse=True) +def _restore_contract_manager_patches() -> Iterator[None]: + """Limit factory manager substitutions to one test invocation.""" + with contract_manager_patch_scope(): + yield diff --git a/source/isaaclab/test/assets/contract/public_surface.py b/source/isaaclab/test/assets/contract/public_surface.py new file mode 100644 index 000000000000..acf4d0d3e773 --- /dev/null +++ b/source/isaaclab/test/assets/contract/public_surface.py @@ -0,0 +1,263 @@ +# 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 + +"""Public base-class surface classified by the shared contract that owns it. + +:data:`PUBLIC_SURFACE_CONTRACTS` is the single inventory of the public members declared by the asset base classes. +``test_base_public_surface_has_an_explicit_contract_classification`` diffs it against the classes, so adding a public +member forces its contract classification and removing one is reported as a stale entry. +""" + +from dataclasses import dataclass +from enum import StrEnum + +from isaaclab.assets.articulation.base_articulation import BaseArticulation +from isaaclab.assets.articulation.base_articulation_data import BaseArticulationData +from isaaclab.assets.asset_base import AssetBase +from isaaclab.assets.rigid_object.base_rigid_object import BaseRigidObject +from isaaclab.assets.rigid_object.base_rigid_object_data import BaseRigidObjectData +from isaaclab.assets.rigid_object_collection.base_rigid_object_collection import BaseRigidObjectCollection +from isaaclab.assets.rigid_object_collection.base_rigid_object_collection_data import BaseRigidObjectCollectionData + + +class ContractKind(StrEnum): + """Contract family that owns a public member.""" + + API = "api" + DATA = "data" + WRITE = "write" + OUT_OF_SCOPE = "out_of_scope" + + +@dataclass(frozen=True) +class PublicMemberContract: + """Classify one public member, with the reason for an out-of-scope exclusion.""" + + kind: ContractKind + reason: str | None = None + + +BASE_SURFACE_CLASSES = ( + AssetBase, + BaseRigidObject, + BaseRigidObjectData, + BaseRigidObjectCollection, + BaseRigidObjectCollectionData, + BaseArticulation, + BaseArticulationData, +) + + +def _mapping(class_name: str, members: str, kind: ContractKind, reason: str | None = None): + """Classify a whitespace-separated group of members declared by one class.""" + return {f"{class_name}.{member}": PublicMemberContract(kind, reason) for member in members.split()} + + +PUBLIC_SURFACE_CONTRACTS = { + **_mapping("AssetBase", "data device is_initialized num_instances", ContractKind.API), + **_mapping( + "AssetBase", + "assert_shape_and_dtype assert_shape_and_dtype_mask reset update write_data_to_sim", + ContractKind.WRITE, + ), + **_mapping( + "AssetBase", + "has_debug_vis_implementation set_debug_vis", + ContractKind.OUT_OF_SCOPE, + reason="debug visualization requires a live USD stage and Kit application", + ), + **_mapping( + "BaseRigidObject", + """ + body_names data find_bodies instantaneous_wrench_composer num_bodies num_instances permanent_wrench_composer + root_view + """, + ContractKind.API, + ), + **_mapping( + "BaseRigidObject", + """ + reset set_coms set_coms_index set_coms_mask set_external_force_and_torque set_inertias set_inertias_index + set_inertias_mask set_masses set_masses_index set_masses_mask update write_data_to_sim + write_root_com_pose_to_sim write_root_com_pose_to_sim_index write_root_com_pose_to_sim_mask + write_root_com_state_to_sim write_root_com_velocity_to_sim write_root_com_velocity_to_sim_index + write_root_com_velocity_to_sim_mask write_root_link_pose_to_sim write_root_link_pose_to_sim_index + write_root_link_pose_to_sim_mask write_root_link_state_to_sim write_root_link_velocity_to_sim + write_root_link_velocity_to_sim_index write_root_link_velocity_to_sim_mask write_root_pose_to_sim + write_root_pose_to_sim_index write_root_pose_to_sim_mask write_root_state_to_sim write_root_velocity_to_sim + write_root_velocity_to_sim_index write_root_velocity_to_sim_mask + """, + ContractKind.WRITE, + ), + **_mapping( + "BaseRigidObjectData", + """ + body_acc_w body_ang_acc_w body_ang_vel_w body_com_acc_w body_com_ang_acc_w body_com_ang_vel_w + body_com_lin_acc_w body_com_lin_vel_w body_com_pos_b body_com_pos_w body_com_pose_b body_com_pose_w + body_com_quat_b body_com_quat_w body_com_state_w body_com_vel_w body_inertia body_lin_acc_w body_lin_vel_w + body_link_ang_vel_w body_link_lin_vel_w body_link_pos_w body_link_pose_w body_link_quat_w body_link_state_w + body_link_vel_w body_mass body_names body_pos_w body_pose_w body_quat_w body_state_w body_vel_w com_pos_b + com_quat_b default_inertia default_mass default_root_pose default_root_state default_root_vel heading_w + projected_gravity_b root_ang_vel_b root_ang_vel_w root_com_ang_vel_b root_com_ang_vel_w root_com_lin_vel_b + root_com_lin_vel_w root_com_pos_w root_com_pose_w root_com_quat_w root_com_state_w root_com_vel_w + root_lin_vel_b root_lin_vel_w root_link_ang_vel_b root_link_ang_vel_w root_link_lin_vel_b root_link_lin_vel_w + root_link_pos_w root_link_pose_w root_link_quat_w root_link_state_w root_link_vel_w root_pos_w root_pose_w + root_quat_w root_state_w root_vel_w update + """, + ContractKind.DATA, + ), + **_mapping( + "BaseRigidObjectCollection", + """ + body_names data find_bodies find_objects instantaneous_wrench_composer num_bodies num_instances num_objects + object_names permanent_wrench_composer root_view + """, + ContractKind.API, + ), + **_mapping( + "BaseRigidObjectCollection", + """ + reset set_coms set_coms_index set_coms_mask set_external_force_and_torque set_inertias set_inertias_index + set_inertias_mask set_masses set_masses_index set_masses_mask update write_body_com_pose_to_sim + write_body_com_pose_to_sim_index write_body_com_pose_to_sim_mask write_body_com_state_to_sim + write_body_com_velocity_to_sim write_body_com_velocity_to_sim_index write_body_com_velocity_to_sim_mask + write_body_link_pose_to_sim write_body_link_pose_to_sim_index write_body_link_pose_to_sim_mask + write_body_link_state_to_sim write_body_link_velocity_to_sim write_body_link_velocity_to_sim_index + write_body_link_velocity_to_sim_mask write_body_pose_to_sim write_body_pose_to_sim_index + write_body_pose_to_sim_mask write_body_state_to_sim write_body_velocity_to_sim write_body_velocity_to_sim_index + write_body_velocity_to_sim_mask write_data_to_sim write_object_com_pose_to_sim write_object_com_state_to_sim + write_object_com_velocity_to_sim write_object_link_pose_to_sim write_object_link_state_to_sim + write_object_link_velocity_to_sim write_object_pose_to_sim write_object_state_to_sim + write_object_velocity_to_sim + """, + ContractKind.WRITE, + ), + **_mapping( + "BaseRigidObjectCollectionData", + """ + body_acc_w body_ang_acc_w body_ang_vel_w body_com_acc_w body_com_ang_acc_w body_com_ang_vel_b + body_com_ang_vel_w body_com_lin_acc_w body_com_lin_vel_b body_com_lin_vel_w body_com_pos_b body_com_pos_w + body_com_pose_b body_com_pose_w body_com_quat_b body_com_quat_w body_com_state_w body_com_vel_w body_inertia + body_lin_acc_w body_lin_vel_w body_link_ang_vel_b body_link_ang_vel_w body_link_lin_vel_b body_link_lin_vel_w + body_link_pos_w body_link_pose_w body_link_quat_w body_link_state_w body_link_vel_w body_mass body_names + body_pos_w body_pose_w body_quat_w body_state_w body_vel_w com_pos_b com_quat_b default_body_pose + default_body_state default_body_vel default_inertia default_mass default_object_pose default_object_state + default_object_vel heading_w object_acc_w object_ang_acc_w object_ang_vel_b object_ang_vel_w object_com_acc_w + object_com_ang_acc_w object_com_ang_vel_b object_com_ang_vel_w object_com_lin_acc_w object_com_lin_vel_b + object_com_lin_vel_w object_com_pos_b object_com_pos_w object_com_pose_b object_com_pose_w object_com_quat_b + object_com_quat_w object_com_state_w object_com_vel_w object_lin_acc_w object_lin_vel_b object_lin_vel_w + object_link_ang_vel_b object_link_ang_vel_w object_link_lin_vel_b object_link_lin_vel_w object_link_pos_w + object_link_pose_w object_link_quat_w object_link_state_w object_link_vel_w object_pos_w object_pose_w + object_quat_w object_state_w object_vel_w projected_gravity_b update + """, + ContractKind.DATA, + ), + **_mapping( + "BaseArticulation", + """ + backend_body_names backend_joint_names body_names body_ordering data find_bodies find_fixed_tendons find_joints + find_spatial_tendons fixed_tendon_names instantaneous_wrench_composer is_fixed_base joint_names joint_ordering + map_body_ids_to_backend map_joint_ids_to_backend num_base_dofs num_bodies num_fixed_tendons num_instances + num_joints num_spatial_tendons permanent_wrench_composer root_view spatial_tendon_names + """, + ContractKind.API, + ), + **_mapping( + "BaseArticulation", + """ + reset set_coms set_coms_index set_coms_mask set_external_force_and_torque set_fixed_tendon_damping + set_fixed_tendon_damping_index set_fixed_tendon_damping_mask set_fixed_tendon_limit + set_fixed_tendon_limit_stiffness set_fixed_tendon_limit_stiffness_index set_fixed_tendon_limit_stiffness_mask + set_fixed_tendon_offset set_fixed_tendon_offset_index set_fixed_tendon_offset_mask + set_fixed_tendon_position_limit set_fixed_tendon_position_limit_index set_fixed_tendon_position_limit_mask + set_fixed_tendon_position_target_index set_fixed_tendon_position_target_mask set_fixed_tendon_rest_length + set_fixed_tendon_rest_length_index set_fixed_tendon_rest_length_mask set_fixed_tendon_stiffness + set_fixed_tendon_stiffness_index set_fixed_tendon_stiffness_mask set_inertias set_inertias_index + set_inertias_mask set_joint_effort_target set_joint_effort_target_index set_joint_effort_target_mask + set_joint_position_target set_joint_position_target_index set_joint_position_target_mask + set_joint_velocity_target set_joint_velocity_target_index set_joint_velocity_target_mask set_masses + set_masses_index set_masses_mask set_spatial_tendon_damping set_spatial_tendon_damping_index + set_spatial_tendon_damping_mask set_spatial_tendon_limit_stiffness set_spatial_tendon_limit_stiffness_index + set_spatial_tendon_limit_stiffness_mask set_spatial_tendon_offset set_spatial_tendon_offset_index + set_spatial_tendon_offset_mask set_spatial_tendon_stiffness set_spatial_tendon_stiffness_index + set_spatial_tendon_stiffness_mask update write_data_to_sim write_fixed_tendon_properties_to_sim + write_fixed_tendon_properties_to_sim_index write_fixed_tendon_properties_to_sim_mask + write_joint_armature_to_sim write_joint_armature_to_sim_index write_joint_armature_to_sim_mask + write_joint_damping_to_sim write_joint_damping_to_sim_index write_joint_damping_to_sim_mask + write_joint_effort_limit_to_sim write_joint_effort_limit_to_sim_index write_joint_effort_limit_to_sim_mask + write_joint_friction_coefficient_to_sim write_joint_friction_coefficient_to_sim_index + write_joint_friction_coefficient_to_sim_mask write_joint_friction_to_sim write_joint_limits_to_sim + write_joint_position_limit_to_sim write_joint_position_limit_to_sim_index + write_joint_position_limit_to_sim_mask write_joint_position_to_sim write_joint_position_to_sim_index + write_joint_position_to_sim_mask write_joint_state_to_sim write_joint_stiffness_to_sim + write_joint_stiffness_to_sim_index write_joint_stiffness_to_sim_mask write_joint_velocity_limit_to_sim + write_joint_velocity_limit_to_sim_index write_joint_velocity_limit_to_sim_mask write_joint_velocity_to_sim + write_joint_velocity_to_sim_index write_joint_velocity_to_sim_mask write_root_com_pose_to_sim + write_root_com_pose_to_sim_index write_root_com_pose_to_sim_mask write_root_com_state_to_sim + write_root_com_velocity_to_sim write_root_com_velocity_to_sim_index write_root_com_velocity_to_sim_mask + write_root_link_pose_to_sim write_root_link_pose_to_sim_index write_root_link_pose_to_sim_mask + write_root_link_state_to_sim write_root_link_velocity_to_sim write_root_link_velocity_to_sim_index + write_root_link_velocity_to_sim_mask write_root_pose_to_sim write_root_pose_to_sim_index + write_root_pose_to_sim_mask write_root_state_to_sim write_root_velocity_to_sim write_root_velocity_to_sim_index + write_root_velocity_to_sim_mask write_spatial_tendon_properties_to_sim + write_spatial_tendon_properties_to_sim_index write_spatial_tendon_properties_to_sim_mask + """, + ContractKind.WRITE, + ), + **_mapping( + "BaseArticulationData", + """ + applied_torque bind_actuator_collection body_acc_w body_ang_acc_w body_ang_vel_w body_com_acc_w + body_com_ang_acc_w body_com_ang_vel_w body_com_jacobian_w body_com_lin_acc_w body_com_lin_vel_w body_com_pos_b + body_com_pos_w body_com_pose_b body_com_pose_w body_com_quat_b body_com_quat_w body_com_state_w body_com_vel_w + body_inertia body_lin_acc_w body_lin_vel_w body_link_ang_vel_w body_link_jacobian_w body_link_lin_vel_w + body_link_pos_w body_link_pose_w body_link_quat_w body_link_state_w body_link_vel_w body_mass body_names + body_ordering body_pos_w body_pose_w body_quat_w body_state_w body_vel_w com_pos_b com_quat_b computed_torque + default_fixed_tendon_damping default_fixed_tendon_limit default_fixed_tendon_limit_stiffness + default_fixed_tendon_offset default_fixed_tendon_pos_limits default_fixed_tendon_rest_length + default_fixed_tendon_stiffness default_inertia default_joint_armature default_joint_damping + default_joint_friction default_joint_friction_coeff default_joint_limits default_joint_pos + default_joint_pos_limits default_joint_stiffness default_joint_vel default_joint_viscous_friction_coeff + default_mass default_root_pose default_root_state default_root_vel default_spatial_tendon_damping + default_spatial_tendon_limit_stiffness default_spatial_tendon_offset default_spatial_tendon_stiffness + fixed_tendon_damping fixed_tendon_limit fixed_tendon_limit_stiffness fixed_tendon_names fixed_tendon_offset + fixed_tendon_pos_limits fixed_tendon_rest_length fixed_tendon_stiffness gravity_compensation_forces + has_body_ordering has_joint_ordering heading_w joint_acc joint_armature joint_damping joint_effort_limits + joint_effort_target joint_friction joint_friction_coeff joint_limits joint_names joint_ordering joint_pos + joint_pos_limits joint_pos_target joint_stiffness joint_vel joint_vel_limits joint_vel_target + joint_velocity_limits mass_matrix projected_gravity_b root_ang_vel_b root_ang_vel_w root_com_ang_vel_b + root_com_ang_vel_w root_com_lin_vel_b root_com_lin_vel_w root_com_pos_w root_com_pose_w root_com_quat_w + root_com_state_w root_com_vel_w root_lin_vel_b root_lin_vel_w root_link_ang_vel_b root_link_ang_vel_w + root_link_lin_vel_b root_link_lin_vel_w root_link_pos_w root_link_pose_w root_link_quat_w root_link_state_w + root_link_vel_w root_pos_w root_pose_w root_quat_w root_state_w root_vel_w soft_joint_pos_limits + soft_joint_vel_limits spatial_tendon_damping spatial_tendon_limit_stiffness spatial_tendon_names + spatial_tendon_offset spatial_tendon_stiffness update + """, + ContractKind.DATA, + ), +} + + +def public_surface_mismatches(classes: tuple[type, ...], contracts: dict[str, PublicMemberContract]) -> dict[str, list]: + """Diff the public members the classes declare against their contract classifications. + + Args: + classes: Base classes whose own public members form the audited surface. + contracts: Qualified member name to contract classification. + + Returns: + Sorted ``missing`` (declared but unclassified), ``stale`` (classified but no longer declared), and + ``unreasoned`` (out-of-scope without a reason) member names; all empty when the inventory is exact. + """ + declared = {f"{cls.__name__}.{name}" for cls in classes for name in vars(cls) if not name.startswith("_")} + return { + "missing": sorted(declared - contracts.keys()), + "stale": sorted(contracts.keys() - declared), + "unreasoned": sorted( + name + for name, contract in contracts.items() + if contract.kind is ContractKind.OUT_OF_SCOPE and not contract.reason + ), + } diff --git a/source/isaaclab/test/assets/contract/test_asset_contract_api.py b/source/isaaclab/test/assets/contract/test_asset_contract_api.py new file mode 100644 index 000000000000..7eaf7ef12bba --- /dev/null +++ b/source/isaaclab/test/assets/contract/test_asset_contract_api.py @@ -0,0 +1,115 @@ +# 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 API contract: counts, names, finders, index resolution, and the classified public surface.""" + +import inspect +from importlib.util import find_spec + +import pytest + +from ._articulation_contract_cases import ( # noqa: F401 + TestArticulationFinderReturnModes, + TestArticulationIndexResolution, + TestArticulationProperties, + TestFixedTendonTargetScheduling, + TestResolveMatchingNamesCache, +) +from ._articulation_contract_utils import get_articulation +from ._manager_patch_scope import contract_manager_patch_scope +from ._rigid_object_collection_contract_cases import ( # noqa: F401 + TestCollectionFinderReturnModes, + TestCollectionIndexResolution, + TestCollectionProperties, +) +from ._rigid_object_collection_contract_utils import get_rigid_object_collection +from ._rigid_object_contract_cases import ( # noqa: F401 + TestRigidObjectFinderReturnModes, + TestRigidObjectIndexResolution, + TestRigidObjectProperties, +) +from ._rigid_object_contract_utils import get_rigid_object +from .capabilities import BACKEND_STATUSES, available_backends +from .public_surface import BASE_SURFACE_CLASSES, PUBLIC_SURFACE_CONTRACTS, public_surface_mismatches + +pytestmark = pytest.mark.integration + +_MISSING = object() + + +def test_installed_backend_packages_run_the_contract() -> None: + """An installed backend package must not silently drop out of the contract matrix.""" + unavailable = { + status.declaration.name: status.reason + for status in BACKEND_STATUSES + if not status.available + and status.reason != "CUDA runtime unavailable" + and find_spec(status.declaration.required_modules[-1]) is not None + } + + assert unavailable == {} + + +def _manager_bindings() -> dict[str, object]: + """Return the backend-manager bindings that the mocked factories replace during a test.""" + bindings = {} + backends = available_backends() + if "physx" in backends: + from isaaclab_physx.physics import PhysxManager + + for name in ("get_physics_sim_view", "_scene_data_backend"): + bindings[f"physx.{name}"] = inspect.getattr_static(PhysxManager, name, _MISSING) + if "ovphysx" in backends: + from isaaclab_ov.physics.ovphysx_manager import OvPhysxManager + + bindings["ovphysx._scene_data_backend"] = inspect.getattr_static( + OvPhysxManager, "_scene_data_backend", _MISSING + ) + if "newton" in backends: + import isaaclab_newton.assets.articulation.articulation_data as articulation_data + import isaaclab_newton.assets.rigid_object.rigid_object_data as rigid_object_data + import isaaclab_newton.assets.rigid_object_collection.rigid_object_collection as collection + import isaaclab_newton.assets.rigid_object_collection.rigid_object_collection_data as collection_data + + for module in (articulation_data, rigid_object_data, collection, collection_data): + bindings[f"newton.{module.__name__}"] = module.SimulationManager + return bindings + + +@pytest.mark.parametrize("reverse", [False, True], ids=["forward", "reverse"]) +def test_contract_factories_scope_manager_patches_to_one_test(reverse: bool) -> None: + """Factories patch backend managers only inside a test scope and restore the production bindings after it.""" + original = _manager_bindings() + # Collection alone must not patch a manager: the production bindings are in place before any factory runs. + if "physx" in available_backends(): + assert isinstance(original["physx.get_physics_sim_view"], classmethod) + if "newton" in available_backends(): + from isaaclab_newton.physics import NewtonManager + + assert all(manager is NewtonManager for name, manager in original.items() if name.startswith("newton.")) + + backends = available_backends()[::-1] if reverse else available_backends() + with contract_manager_patch_scope(): + for backend in backends: + rigid_object, _ = get_rigid_object(backend, device="cpu") + collection, _ = get_rigid_object_collection(backend, device="cpu") + articulation, _ = get_articulation(backend, device="cpu") + # Later factories must not break the patched managers the earlier families read through. + assert rigid_object.data.root_link_pose_w.shape == (2,) + assert collection.data.body_link_pose_w.shape == (2, 3) + assert articulation.data.root_link_pose_w.shape == (2,) + + restored = _manager_bindings() + assert restored.keys() == original.keys() + assert all(restored[name] is binding for name, binding in original.items()) + + +def test_base_public_surface_has_an_explicit_contract_classification() -> None: + """Every public member of the asset base classes is classified, and every classified member still exists.""" + assert public_surface_mismatches(BASE_SURFACE_CLASSES, PUBLIC_SURFACE_CONTRACTS) == { + "missing": [], + "stale": [], + "unreasoned": [], + } diff --git a/source/isaaclab/test/assets/contract/test_asset_contract_data.py b/source/isaaclab/test/assets/contract/test_asset_contract_data.py new file mode 100644 index 000000000000..717011671607 --- /dev/null +++ b/source/isaaclab/test/assets/contract/test_asset_contract_data.py @@ -0,0 +1,30 @@ +# 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 data contract: property shapes and dtypes, aliases, lazy buffers, and ordered reads.""" + +import pytest + +from ._articulation_contract_cases import ( # noqa: F401 + TestArticulationDataAliases, + TestArticulationDataProperties, + TestArticulationTendons, +) +from ._articulation_ordering_contract_cases import ( # noqa: F401 + TestArticulationOrderingAllocation, + TestArticulationOrderingBodyState, + TestArticulationOrderingJointState, +) +from ._rigid_object_collection_contract_cases import ( # noqa: F401 + TestCollectionDataAliases, + TestCollectionDataProperties, + TestCollectionViewReshape, +) +from ._rigid_object_contract_cases import ( # noqa: F401 + TestRigidObjectDataAliases, + TestRigidObjectDataProperties, +) + +pytestmark = pytest.mark.integration diff --git a/source/isaaclab/test/assets/contract/test_asset_contract_writes.py b/source/isaaclab/test/assets/contract/test_asset_contract_writes.py new file mode 100644 index 000000000000..4805968d3179 --- /dev/null +++ b/source/isaaclab/test/assets/contract/test_asset_contract_writes.py @@ -0,0 +1,39 @@ +# 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 write contract: writers, read-backs, partial selections, cache invalidation, and ordered writes.""" + +import pytest + +from ._articulation_contract_cases import ( # noqa: F401 + TestArticulationPartialWriteRows, + TestArticulationWritersBody, + TestArticulationWritersJoint, + TestArticulationWritersRoot, + TestArticulationWritersTendon, + test_deprecated_joint_friction_writers, +) +from ._articulation_ordering_contract_cases import ( # noqa: F401 + TestArticulationOrderingBodyWriters, + TestArticulationOrderingComWrites, + TestArticulationOrderingJointWriters, + TestArticulationOrderingOperations, + TestArticulationOrderingRootWriteParity, + TestArticulationOrderingWriteParity, +) +from ._rigid_object_collection_contract_cases import ( # noqa: F401 + TestCollectionCacheInvalidation, + TestCollectionPartialWriteCells, + TestCollectionWritersBody, + TestCollectionWritersPose, +) +from ._rigid_object_contract_cases import ( # noqa: F401 + TestRigidObjectCacheInvalidation, + TestRigidObjectPartialWriteRows, + TestRigidObjectWritersBody, + TestRigidObjectWritersRoot, +) + +pytestmark = pytest.mark.integration From 3bd4e7c1fafd644c2470a53273c68657e157ce06 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 10:32:43 +0200 Subject: [PATCH 03/33] Reduce WrenchComposer tests to literal units and one PhysX delivery test Replace the randomized and device-parametrized composer matrix with literal two-environment, two-body cases that pin frames, fast paths, live CoM and rotation reads, selection, set, reset, merge, lazy composition, and the deprecated API. Keep one real PhysX test: a world force at a world offset must match raw PhysX delivery, and a world wrench at the CoM of a rotated body far from the origin must follow the world axes. The composer-versus-PhysX matrix is removed; its only non-trivial parity case is the retained offset force, and its other cases issued the same PhysX call on both sides. --- .../test/utils/test_wrench_composer.py | 1409 ++++------------- .../utils/test_wrench_composer_integration.py | 680 +------- .../utils/test_wrench_composer_vs_physx.py | 701 -------- 3 files changed, 389 insertions(+), 2401 deletions(-) delete mode 100644 source/isaaclab/test/utils/test_wrench_composer_vs_physx.py diff --git a/source/isaaclab/test/utils/test_wrench_composer.py b/source/isaaclab/test/utils/test_wrench_composer.py index a86d990193e8..c688d4d2574f 100644 --- a/source/isaaclab/test/utils/test_wrench_composer.py +++ b/source/isaaclab/test/utils/test_wrench_composer.py @@ -3,6 +3,13 @@ # # SPDX-License-Identifier: BSD-3-Clause +"""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``. +""" + import ast import inspect from types import SimpleNamespace @@ -14,12 +21,21 @@ 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", +] + def test_wrench_composer_uses_asset_frame_conventions(): """Keep frame selection as an is_global boolean, with no enum or content-classification layer.""" @@ -34,640 +50,202 @@ def test_wrench_composer_uses_asset_frame_conventions(): 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)) - - -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. - - Returns: - Asset fixture with the state required by WrenchComposer. - """ - - # 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() - - return SimpleNamespace( - num_instances=num_envs, num_bodies=num_bodies, device=device, data=_WrenchAssetDataFixture(pos, quat, device) - ) - - -# --- Helper functions for quaternion math --- +# --------------------------------------------------------------------------- +# Literal fixture helpers +# --------------------------------------------------------------------------- -def quat_rotate_inv_np(quat_xyzw: np.ndarray, vec: np.ndarray) -> np.ndarray: - """Rotate a vector by the inverse of a quaternion (numpy). +class _AssetData: + """Body poses that :class:`WrenchComposer` reads when it composes into the body frame.""" - Args: - quat_xyzw: Quaternion in (x, y, z, w) format. Shape: (..., 4) - vec: Vector to rotate. Shape: (..., 3) + def __init__(self, com_pos_w: torch.Tensor, link_quat_w: torch.Tensor) -> None: + self.set_pose(com_pos_w, link_quat_w) - Returns: - Rotated vector. Shape: (..., 3) - """ - # Extract components - xyz = quat_xyzw[..., 0:3] - w = quat_xyzw[..., 3:4] + 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.contiguous(), dtype=wp.vec3f)) + self.body_link_quat_w = ProxyArray(wp.from_torch(link_quat_w.contiguous(), dtype=wp.quatf)) - # 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 _poses(com: dict | None = None, quat: dict | None = None, shape: tuple[int, int] = (2, 2)): + """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 -def random_unit_quaternion_np(rng: np.random.Generator, shape: tuple) -> np.ndarray: - """Generate random unit quaternions in (x, y, z, w) format. +def _make_composer( + com: dict | None = None, quat: dict | 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) - 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 +def _grid(cells: dict, 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()) -@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) +def _vectors(values: torch.Tensor) -> wp.array: + """Convert a ``(..., 3)`` tensor to a Warp vector array.""" + return wp.from_torch(values.contiguous(), dtype=wp.vec3f) -# ============================================================================ -# 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 _mask(values: list[bool]) -> wp.array: + """Convert literal booleans to a Warp mask array with owned storage.""" + return wp.array(np.array(values, dtype=bool), dtype=wp.bool, device="cpu") -@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}" - ) - - # 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}" - ) - - -@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) - - 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 - - # 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) - - -@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) + _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)})) - 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 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) +def test_add_accumulates_local_wrenches() -> None: + composer = _make_composer() - 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) + 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) - # Add local forces first - wrench_composer.add_forces_and_torques_index(forces=forces_local, is_global=False) + _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)})) - # Add global forces - wrench_composer.add_forces_and_torques_index(forces=forces_global, is_global=True) - # 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 +@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 - # 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, 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)})) - 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, - ) - # 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, +@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)) - # 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" + # 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}) ) - 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" - ) - - -@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 -# ============================================================================ - + 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_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]) @@ -676,7 +254,8 @@ def test_index_dtype_combinations_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 +291,226 @@ 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) - - mock_asset = create_mock_asset(num_envs, num_bodies, device) - composer = WrenchComposer(mock_asset) - - # 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), - ) - - # 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 - - # 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("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) - - mock_asset = create_mock_asset(num_envs, num_bodies, device) - composer = WrenchComposer(mock_asset) - - # 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) +def _fill_all_input_buffers(composer: WrenchComposer) -> None: + """Give every cell a distinct contribution in each of the five input buffers.""" + ones = torch.ones((2, 2, 3)) 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), + forces=_vectors(ones * torch.tensor([0.0, 0.0, 1.0])), + positions=_vectors(ones * torch.tensor([0.0, 1.0, 0.0])), 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) + composer.add_forces_and_torques_index(forces=_vectors(ones * torch.tensor([1.0, 0.0, 0.0])), is_global=True) composer.add_forces_and_torques_index( - torques=wp.from_numpy(torques_local_np, dtype=wp.vec3f, device=device), - is_global=False, + forces=_vectors(ones * torch.tensor([0.0, 1.0, 0.0])), torques=_vectors(ones * torch.tensor([0.0, 0.0, 1.0])) ) - # 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, - ) - zeros = np.zeros((num_bodies, 3), dtype=np.float32) +# 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), +} - # 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}" - ) +@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) - # 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" - ) + _apply_to_cell(composer, method, (0, 1), False, forces=(5.0, 0.0, 0.0)) + + 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) -# ============================================================================ -# 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) +@pytest.mark.parametrize("source_is_global", [False, True], ids=["local_source", "global_source"]) +@pytest.mark.parametrize("selector", ["list", "slice", "mask"]) +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( + forces=_vectors(_grid({(0, 0): (1.0, 0.0, 0.0), (1, 0): (1.0, 0.0, 0.0)})), 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, +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)})) - # 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" - ) + # 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 -# ============================================================================ -# Lazy Composition (_ensure_composed) Tests -# ============================================================================ +# --------------------------------------------------------------------------- +# Lazy composition +# --------------------------------------------------------------------------- -@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): +def test_output_property_triggers_lazy_composition(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, - ) + composer = _make_composer(quat={(0, 1): _QUARTER_TURN_Z}) + _apply_to_cell(composer, "add_forces_and_torques_index", (0, 1), True, **{input_name: (2.0, 0.0, 0.0)}) # 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" - ) + assert composer._dirty + _assert_vectors(getattr(composer, output_name), _grid({(0, 1): (0.0, -2.0, 0.0)})) + assert not composer._dirty -@pytest.mark.parametrize("device", test_devices()) -def test_lazy_composition_reflects_later_adds(device: str): +def test_lazy_composition_reflects_later_adds(): """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), - ) + composer = _make_composer() + ones = torch.ones((2, 2, 3)) + composer.add_forces_and_torques_index(forces=_vectors(ones)) # Reading out_force_b composes lazily - assert np.allclose(composer.out_force_b.warp.numpy(), forces_np, atol=1e-4, rtol=1e-5) + _assert_vectors(composer.out_force_b, ones) # 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, - ) - - # First compose + composer.add_forces_and_torques_index(forces=_vectors(ones), torques=_vectors(ones)) + _assert_vectors(composer.out_torque_b, ones) + _assert_vectors(composer.out_force_b, 2.0 * ones) + # Composing again without new input is idempotent. 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) - - 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) - - 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, - ) - - 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, 2.0 * ones) + _assert_vectors(composer.out_torque_b, ones) -@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, - ) - 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() +@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(deprecated_name: str, name: str, input_name: str): + """Test that accessing composed_force/composed_torque emits a DeprecationWarning and aliases the output.""" + composer = _make_composer() + _apply_to_cell(composer, "add_forces_and_torques_index", (1, 0), False, **{input_name: (1.0, 2.0, 3.0)}) - # 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) + with pytest.warns(DeprecationWarning, match=f"{deprecated_name}.*is deprecated"): + result = getattr(composer, deprecated_name) - 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()}" - ) + # Should return the same data as the new output property + assert np.allclose(result.warp.numpy(), getattr(composer, name).warp.numpy(), atol=1e-7) + _assert_vectors(result, _grid({(1, 0): (1.0, 2.0, 3.0)})) -# ============================================================================ -# Deprecated set_forces_and_torques Tests -# ============================================================================ +def test_deprecated_writers_warn_and_keep_add_and_set_behavior() -> None: + composer = _make_composer() + 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)})) -@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) + 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]) - mock_asset = create_mock_asset(num_envs, num_bodies, device) - composer = WrenchComposer(mock_asset) + # 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))) - # 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..d0456e00a837 100644 --- a/source/isaaclab/test/utils/test_wrench_composer_integration.py +++ b/source/isaaclab/test/utils/test_wrench_composer_integration.py @@ -3,15 +3,19 @@ # # 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 @@ -20,629 +24,89 @@ 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 pytestmark = pytest.mark.integration +_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 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(), +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.RigidBodyPropertiesCfg(disable_gravity=True, linear_damping=0.0, angular_damping=0.0), + mass_props=sim_utils.MassPropertiesCfg(mass=1.0), + collision_props=sim_utils.CollisionPropertiesCfg(), ) - - 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, - ) - - -@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}" - ), - ) + return RigidObject(cfg=cfg) @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 + + 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=wp.from_torch(force.view(-1, 3).contiguous(), dtype=wp.float32), + torque_data=wp.from_torch(torch.zeros_like(force).view(-1, 3).contiguous(), dtype=wp.float32), + position_data=wp.from_torch(raw_position.view(-1, 3).contiguous(), dtype=wp.float32), + 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" - ) From 16ac8af05963dba8614872936f8865c05f8baf7f Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 10:45:40 +0200 Subject: [PATCH 04/33] Allow OVPhysX CPU and CUDA simulations in one process OvPhysxManager pinned the first simulation device for the whole process because the backend entered the OVPhysX process-wide CPU-only mode, which OVPhysX 0.6.3 still documents as irreversible. CPU scenes now leave that mode off and select CPU dynamics per scene through enableGPUDynamics=false and the MBP broadphase, as the OVPhysX developer guide recommends for mixed CPU/GPU processes, so the device lock is removed. A child-process test runs CPU, CUDA, CPU and CUDA, CPU scene sequences and checks device placement, state I/O, and that each CPU scene reproduces the same trajectory. A unit test pins the per-device scene attributes. --- .../antoiner-asset-tests-v2-ov.rst | 7 ++ .../isaaclab_ov/physics/ovphysx_manager.py | 62 +++++------ .../physics/ovphysx_manager_cfg.py | 2 +- .../physics/test_ovphysx_manager_lifecycle.py | 100 +++++++++++++++++- .../test_ovphysx_scene_data_backend.py | 16 ++- 5 files changed, 135 insertions(+), 52 deletions(-) create mode 100644 source/isaaclab_ov/changelog.d/antoiner-asset-tests-v2-ov.rst diff --git a/source/isaaclab_ov/changelog.d/antoiner-asset-tests-v2-ov.rst b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-v2-ov.rst new file mode 100644 index 000000000000..56cb1423f640 --- /dev/null +++ b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-v2-ov.rst @@ -0,0 +1,7 @@ +Fixed +^^^^^ + +* Fixed :class:`~isaaclab_ov.physics.OvPhysxManager` refusing to start a CPU simulation after a CUDA one, or the + reverse, in the same process. CPU scenes now select CPU dynamics through their physics-scene attributes instead of + the OVPhysX process-wide CPU-only mode, which cannot be reverted. Deployments that must keep OVPhysX off CUDA can + set ``OVPHYSX_DISABLE_GPU``. diff --git a/source/isaaclab_ov/isaaclab_ov/physics/ovphysx_manager.py b/source/isaaclab_ov/isaaclab_ov/physics/ovphysx_manager.py index 3dd1d9f26269..e9a0b3b78443 100644 --- a/source/isaaclab_ov/isaaclab_ov/physics/ovphysx_manager.py +++ b/source/isaaclab_ov/isaaclab_ov/physics/ovphysx_manager.py @@ -245,7 +245,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) + # Hard CPU-only mode (``PhysX.set_cpu_mode(True)``) is sticky for the process, so CPU scenes + # select CPU dynamics through the scene attributes instead; see _configure_physx_scene_prim. self.physx = ovphysx.PhysX( config=ovphysx.PhysXConfig( num_threads=8, cooked_collider_cache_dir=cache_dir, carbonite_overrides=carbonite_overrides @@ -323,8 +324,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 # Active clone recipes survive the consumable pending queue so a forced # re-warmup can rebuild serialized-stage or runtime-only clones. _active_clone_recipes: ClassVar[list[tuple[str, list[str], list[CloneTransform], list[int] | None]]] = [] @@ -443,7 +442,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() @@ -832,19 +830,15 @@ 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 + 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: @@ -857,13 +851,6 @@ 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._active_clone_recipes: @@ -901,7 +888,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 @@ -969,10 +955,11 @@ def _configure_physx_scene_prim(scene_prim, cfg, device: str) -> None: so we write the apiSchemas list entry and scene attributes directly via raw Sdf metadata manipulation instead of using the high-level USD API. - 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, scene-query-support, solver-determinism/accuracy, and dynamics/broadphase + attributes are applied regardless of device. GPU scenes use ``enableGPUDynamics=true`` and + the GPU broadphase. CPU scenes use ``enableGPUDynamics=false`` and the MBP broadphase; the + schema default enables GPU dynamics, so a CPU scene that omits them would silently simulate + on the GPU. The GPU buffer-capacity attributes are applied only when ``device == "gpu"``. Args: scene_prim: The /World/PhysicsScene prim to configure. @@ -1004,17 +991,18 @@ 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) + scene_prim.CreateAttribute("physxScene:enableGPUDynamics", Sdf.ValueTypeNames.Bool).Set(device == "gpu") + scene_prim.CreateAttribute("physxScene:broadphaseType", Sdf.ValueTypeNames.String).Set( + "GPU" if device == "gpu" else "MBP" + ) + + 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/physics/test_ovphysx_manager_lifecycle.py b/source/isaaclab_ov/test/physics/test_ovphysx_manager_lifecycle.py index 110f864f94ff..2690c821d00f 100644 --- a/source/isaaclab_ov/test/physics/test_ovphysx_manager_lifecycle.py +++ b/source/isaaclab_ov/test/physics/test_ovphysx_manager_lifecycle.py @@ -16,6 +16,8 @@ import pytest +from isaaclab.test.utils import DeviceScope, test_devices + pytest.importorskip("ovphysx.types", reason="ovphysx wheel not installed") @@ -27,10 +29,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 +61,6 @@ def manager_module(monkeypatch): "_next_control_ordinal": 2, "_warmup_done": False, "_requires_full_stage": False, - "_locked_device": None, "_active_clone_recipes": [], "_pending_clones": [], "_atexit_registered": False, @@ -453,6 +450,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], @@ -501,6 +558,23 @@ def test_scene_external_forces_every_iteration(monkeypatch, manager_module, devi assert prim.GetAttribute("physxScene:enableExternalForcesEveryIteration").Get() is expected +@pytest.mark.parametrize(("device", "gpu_dynamics", "broadphase"), [("cpu", False, "MBP"), ("gpu", True, "GPU")]) +def test_scene_authors_device_dynamics_and_broadphase(monkeypatch, manager_module, device, gpu_dynamics, broadphase): + """Each scene selects its own dynamics device, since CPU and GPU scenes share one process.""" + from isaaclab_ov.physics.ovphysx_manager_cfg import OvPhysxCfg + + from pxr import Usd + + from isaaclab.physics import PhysicsManager + + monkeypatch.setattr(PhysicsManager, "_sim", None) + stage = Usd.Stage.CreateInMemory() + prim = stage.DefinePrim("/World/PhysicsScene", "PhysicsScene") + manager_module.OvPhysxManager._configure_physx_scene_prim(prim, OvPhysxCfg(), device) + 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 @@ -550,3 +624,19 @@ 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)) + + +_CPU_DEVICES = test_devices(DeviceScope.CPU) +_CUDA_DEVICES = test_devices(DeviceScope.CUDA) + + +@pytest.mark.skipif(not (_CPU_DEVICES and _CUDA_DEVICES), reason="The device sequence requires a CPU and a CUDA device") +@pytest.mark.parametrize("order", ["cpu-cuda-cpu", "cuda-cpu"]) +def test_cpu_and_cuda_scenes_run_sequentially_in_one_process(order): + """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] + devices = (cpu, cuda, cpu) if order == "cpu-cuda-cpu" else (cuda, cpu) + completed, output = _run_child(_device_sequence_script(devices)) + + 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 21b5137f14d0..5ea903e7593b 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 @@ -333,14 +333,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 @@ -393,7 +391,7 @@ 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 From 4bd517d876e66725ff0c99c602b9323102fc37f7 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 10:45:47 +0200 Subject: [PATCH 05/33] Consolidate OVPhysX articulation tests into composite scenes Replace the per-test Nucleus-robot scenes with one module-scoped scene per device built from locally authored islands of the branching fixture. Each island serves one test, is reset once, and uses two environments where a partial write must prove the other environment keeps its backend state. The CUDA scene keeps the pinned-host staging, native Newton actuator, dynamics, tendon, floating-root, and FK-refresh tests. The CPU scene covers the host-resident property path and the host actuator path, so the file no longer needs one process per device. Unit-level articulation data and actuator-control tests move verbatim to unit modules. --- .../test/assets/test_articulation.py | 3663 +++++------------ .../isaaclab_ov/test/assets/unit/__init__.py | 4 + .../unit/test_ovphysx_actuator_control.py | 71 + .../assets/unit/test_ovphysx_articulation.py | 148 + 4 files changed, 1209 insertions(+), 2677 deletions(-) create mode 100644 source/isaaclab_ov/test/assets/unit/__init__.py create mode 100644 source/isaaclab_ov/test/assets/unit/test_ovphysx_actuator_control.py create mode 100644 source/isaaclab_ov/test/assets/unit/test_ovphysx_articulation.py diff --git a/source/isaaclab_ov/test/assets/test_articulation.py b/source/isaaclab_ov/test/assets/test_articulation.py index f075d4462872..6a259c3981d3 100644 --- a/source/isaaclab_ov/test/assets/test_articulation.py +++ b/source/isaaclab_ov/test/assets/test_articulation.py @@ -6,52 +6,22 @@ # 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 in the module. Each island belongs to one test +unless noted, so the tests do not depend on each other's order. Partial writes target the last +environment and prove the other environments 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. """ from __future__ import annotations -import importlib -import sys +from collections.abc import Iterator +from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace from unittest.mock import Mock @@ -60,943 +30,801 @@ import torch import warp as wp -from pxr import Gf, 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 +from pxr import Gf, Sdf, UsdGeom, 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 import tensor_types as TT # noqa: E402 from isaaclab_ov.assets import Articulation # 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_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.sim import SimulationCfg, build_simulation_context # noqa: E402 -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR # noqa: E402 -from isaaclab.utils.warp.launch_cache import _WarpLaunchCache # 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, +) + +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) +_ROTATED_JOINT_NAMES = (*BRANCHING_PHYSX_JOINT_NAMES[1:], BRANCHING_PHYSX_JOINT_NAMES[0]) +# We set 80.0 for max force because the USD default is 10e10, which makes testing annoying. +_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: + 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 + 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), + "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 __getitem__(self, name: str) -> Articulation: + return self.islands[name] + + 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, + ) + ) + stage = sim_utils.get_current_stage() + for env_index in range(cfg.num_envs): + prim_path = f"{island_path}/Env_{env_index}/Robot" + 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) -## -# Pre-defined configs -## -from isaaclab_assets import ANYMAL_C_CFG, CARTPOLE_CFG, FRANKA_PANDA_CFG, SHADOW_HAND_PHYSX_CFG # isort:skip -wp.init() +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) -pytestmark = pytest.mark.device_split +def _backend_to_user(articulation: Articulation) -> list[int]: + """Return the public joint index of every backend joint column.""" + if articulation.joint_ordering is None: + return list(range(articulation.num_joints)) + return list(articulation.joint_ordering.backend_to_user_indices) -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 +@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["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, ) - articulation = SimpleNamespace( - _sim_cfg=SimpleNamespace(use_newton_actuators=True), - cfg=SimpleNamespace(prim_path="/World/Robot"), + 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 = _backend_to_user(articulation) + torch.testing.assert_close( + wp.to_torch(articulation.root_view.get_attribute(TT.DOF_POSITION)), + expected_position[:, backend_to_user], ) - 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)}, + # 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 ) - - assert native_groups == set() - assert not control.native_actuator_path_active - assert not articulation._has_newton_actuators - assert runtime_prepare_calls == [] -def test_cached_read_launches_reset_on_ordering_and_invalidation(): - """Ordering installation and simulation invalidation should discard recorded reads.""" +@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["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. - class MinimalData(ArticulationData): - def __dir__(self): - return [] + 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 - class Buffer: - timestamp = 1.0 + articulation = scene["ordered"] + num_articulations = articulation.num_instances + device = scene.device + sim = scene.sim - data = MinimalData.__new__(MinimalData) - read_launch_cache = Mock() - data._read_launch_cache = read_launch_cache - data._configure_ordering_buffers = lambda: None - data._make_jacobian_body_user_to_backend = lambda: object() - data.joint_ordering = None - data._body_com_jacobian_w = Buffer() - data._mass_matrix = Buffer() - data._gravity_compensation_forces = Buffer() + # 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 = SimpleNamespace(name="robot", body_ids=slice(None)) + env = SimpleNamespace(scene={"robot": articulation}, sim=sim) + randomize = randomize_rigid_body_material(SimpleNamespace(params={**params, "asset_cfg": asset_cfg}), env) - data._apply_ordering_maps_after_resolve() + # Randomize only the last environment; the others keep their materials. + randomize(env, torch.tensor([num_articulations - 1], device=device), *params.values(), asset_cfg) + scene.step(articulation) - read_launch_cache.clear.assert_called_once_with() - assert data._body_com_jacobian_w.timestamp == -1.0 - assert data._mass_matrix.timestamp == -1.0 - assert data._gravity_compensation_forces.timestamp == -1.0 + 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 - data._is_primed = True - data._sim_timestamp = 1.0 - data._invalidate_initialize_callback(None) - assert read_launch_cache.clear.call_count == 2 - assert data._is_primed is False - assert data._sim_timestamp == 0.0 +@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[island] + device = scene.device + actuator_cfg = _ISLANDS[island].actuators["joints"] + joint_drive_props = _ISLANDS[island].joint_drive_props -def test_static_property_reads_are_not_invalidated_by_simulation_steps(): - """Joint properties and body mass/inertia should be read once per invalidation, not once per step. + # 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 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) - On OVPhysX these are blocking CPU-only binding reads whose cost scales with the number of - environments, so re-reading them every physics step made per-step consumers (such as the - native actuator telemetry sync) host-bound. State buffers must still refresh every step. - """ + # 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 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) - class Buffer: - def __init__(self, shape): - self.data = wp.zeros(shape, dtype=wp.float32, device="cpu") - self.timestamp = -1.0 - - data = ArticulationData.__new__(ArticulationData) - data.device = "cpu" - data.num_instances = 1 - data.num_joints = 2 - data._sim_timestamp = 1.0 - data.body_ordering = None - data._get_binding = lambda tensor_type: object() - reads: list[int] = [] - data._binding_read = lambda tensor_type, dst: reads.append(tensor_type) - - # Joint properties: one read across several steps, one more after explicit invalidation - # (simulation reinitialization). - data.joint_ordering = None - stiffness = Buffer((1, 2)) - for _ in range(3): - data._sim_timestamp += 1.0 - data._read_joint_property_binding(TT.DOF_STIFFNESS, stiffness, None) - assert reads.count(TT.DOF_STIFFNESS) == 1 - stiffness.timestamp = -1.0 - data._read_joint_property_binding(TT.DOF_STIFFNESS, stiffness, None) - assert reads.count(TT.DOF_STIFFNESS) == 2 - - # Body properties behave the same; body state buffers still refresh every step. - mass, link_pose = Buffer((1, 2)), Buffer((1, 2)) - for _ in range(3): - data._sim_timestamp += 1.0 - data._refresh_reordered_body_buffer(mass, None, TT.BODY_MASS, static=True) - data._refresh_reordered_body_buffer(link_pose, None, TT.LINK_POSE) - assert reads.count(TT.BODY_MASS) == 1 - assert reads.count(TT.LINK_POSE) == 3 - mass.timestamp = -1.0 - data._refresh_reordered_body_buffer(mass, None, TT.BODY_MASS, static=True) - assert reads.count(TT.BODY_MASS) == 2 - - # Under a non-identity joint ordering the property is gathered once, then served from cache. - data.joint_ordering = SimpleNamespace(user_to_backend=wp.array([1, 0], dtype=wp.int32, device="cpu")) - data._read_launch_cache = _WarpLaunchCache("cpu") - user_buffer, backend_buffer = Buffer((1, 2)), Buffer((1, 2)) - data._binding_read = lambda tensor_type, dst: (reads.append(tensor_type), dst.assign([[1.0, 2.0]])) - for _ in range(3): - data._sim_timestamp += 1.0 - data._read_joint_property_binding(TT.DOF_DAMPING, user_buffer, backend_buffer) - assert reads.count(TT.DOF_DAMPING) == 1 - torch.testing.assert_close(wp.to_torch(user_buffer.data), torch.tensor([[2.0, 1.0]])) +@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. -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*. + The CUDA scene runs the explicit group on the native Newton actuator path; the CPU scene keeps it + on the host actuator path. """ - arr = articulation.root_view.get_attribute(tensor_type) - return wp.to_torch(arr).to(device) - + articulation = scene["mixed"] + stiffness, effort_limit = 20.0, 400.0 + 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() -# 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] + 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[:, 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 -@pytest.fixture(autouse=True) -def _ovphysx_skip_other_device(request): - """Skip tests whose ``device`` parameter mismatches the session-locked device. +@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. - 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. + 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. """ - 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." - ) + articulation = scene["com_order"] + device = scene.device + 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 + + # 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).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, "cpu").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, "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( + 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.contiguous(), dtype=wp.transformf), body_ids=[public_body_id]) -def _ovphysx_sim_context(device: str, **kwargs): - """Wrapper around :func:`build_simulation_context` that injects OVPhysX cfg. + 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 + ) - 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, + 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]) + 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, ) - return build_simulation_context(device=device, sim_cfg=sim_cfg, **kwargs) - - -def generate_articulation_cfg( - articulation_type: str, - stiffness: float | None = 10.0, - damping: float | None = 2.0, - 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 +@pytest.mark.parametrize("scene", _CUDA_DEVICES, 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["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) + backend_to_user = torch.as_tensor(ordering.backend_to_user_indices, dtype=torch.long, device=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).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) + articulation.set_joint_velocity_target_index(target=velocity_target) + articulation.actuators.target_command.set_effort_index(value=effort_target) + articulation.write_data_to_sim() + backend_position_target = _read_binding_to_torch(articulation, TT.DOF_POSITION_TARGET, device) + backend_velocity_target = _read_binding_to_torch(articulation, TT.DOF_VELOCITY_TARGET, device) + torch.testing.assert_close(backend_position_target, position_target[:, backend_to_user]) + torch.testing.assert_close(backend_velocity_target, velocity_target[:, backend_to_user]) + torch.testing.assert_close( + _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 + ) + 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) -def generate_articulation( - articulation_cfg: ArticulationCfg, num_articulations: int, device: str -) -> tuple[Articulation, torch.tensor]: - """Generate an articulation from a configuration. + # 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]) - Handles the creation of the articulation, the environment prims and the articulation's environment - translations + # 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) - 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", _CUDA_DEVICES, 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. + 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. """ - # Generate translations of 2.5 m in x for each articulation - translations = torch.zeros(num_articulations, 3, device=device) - translations[:, 0] = torch.arange(num_articulations) * 2.5 - - # Create Top-level Xforms, one for each articulation - for i in range(num_articulations): - sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=translations[i][:3]) - articulation = Articulation(replace(articulation_cfg, prim_path="/World/Env_[^/]*/Robot")) + articulation = scene["dynamics"] + device = scene.device + num_envs = articulation.num_instances - return articulation, translations + # 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 -@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, - ) - }, + 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, ) - articulation, _ = generate_articulation(articulation_cfg, 1, device) - sim.reset() + ) - 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, + 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, ) - - 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, + 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], ) - torch.testing.assert_close(articulation.actuators.applied_effort.torch, expected_effort) - -@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 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) - 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) + 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) - monkeypatch.setattr(importlib.util, "find_spec", reject_newton_probe) - importlib.reload(importlib.import_module(module_name)) + 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) + assert not torch.allclose(expected_dynamics()[0], raw_jacobian_before) + assert_dynamics_match() - -@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.""" - 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, SceneEntityCfg # 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 - - -@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, - ) - ) - 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() - - 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) - - 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) - - 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 - ) - 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 - ) - - -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,) - - -@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() - - 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) - backend_to_user = torch.as_tensor(ordering.backend_to_user_indices, dtype=torch.long, device=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) - 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) - articulation.set_joint_velocity_target_index(target=velocity_target) - articulation.actuators.target_command.set_effort_index(value=effort_target) - articulation.write_data_to_sim() - - backend_position_target = _read_binding_to_torch(articulation, TT.DOF_POSITION_TARGET, device) - backend_velocity_target = _read_binding_to_torch(articulation, TT.DOF_VELOCITY_TARGET, device) - torch.testing.assert_close(backend_position_target, position_target[:, backend_to_user]) - torch.testing.assert_close(backend_velocity_target, velocity_target[:, backend_to_user]) - torch.testing.assert_close( - _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 - ) - 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], - ) - - 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 - ) - - -@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. - - 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. - """ - 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 - - # 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() - - 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 - ) - 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) - - joint_position = articulation.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) - 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 - ) + 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", 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)) 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() @@ -1004,1178 +832,100 @@ 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)) - + assert_dynamics_match(jacobian=False, gravity=False) -@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): - """Test write_joint_limits_to_sim API and when default pos falls outside of the new limits. - - This test verifies that: - 1. Joint limits can be set correctly - 2. Default positions are preserved when setting new limits - 3. Joint limits can be set with indexing - 4. Invalid joint positions are properly handled - - Args: - sim: The simulation fixture - num_articulations: Number of articulations to test - """ - # Create articulation - articulation_cfg = generate_articulation_cfg(articulation_type="panda") - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device) - - # Play sim - sim.reset() - # Check if articulation is initialized - assert articulation.is_initialized - - # Get current default joint pos - default_joint_pos = articulation._data.default_joint_pos.torch.clone() - - # Set new joint limits - limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) - limits[..., 0] = (torch.rand(num_articulations, articulation.num_joints, device=device) + 5.0) * -1.0 - limits[..., 1] = torch.rand(num_articulations, articulation.num_joints, device=device) + 5.0 - articulation.write_joint_position_limit_to_sim_index(limits=limits) - - # Check new limits are in place - torch.testing.assert_close(articulation._data.joint_pos_limits.torch, limits) - torch.testing.assert_close(articulation._data.default_joint_pos.torch, default_joint_pos) - - # Set new joint limits with indexing - env_ids = torch.arange(1, device=device, dtype=torch.int32) - joint_ids = torch.arange(2, device=device, dtype=torch.int32) - limits = torch.zeros(env_ids.shape[0], joint_ids.shape[0], 2, device=device) - limits[..., 0] = (torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) + 5.0) * -1.0 - limits[..., 1] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) + 5.0 - articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) - - # Check new limits are in place - torch.testing.assert_close(articulation._data.joint_pos_limits.torch[env_ids][:, joint_ids], limits) - torch.testing.assert_close(articulation._data.default_joint_pos.torch, default_joint_pos) - - # Set new joint limits that invalidate default joint pos - limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) - limits[..., 0] = torch.rand(num_articulations, articulation.num_joints, device=device) * -0.1 - limits[..., 1] = torch.rand(num_articulations, articulation.num_joints, device=device) * 0.1 - articulation.write_joint_position_limit_to_sim_index(limits=limits) - - # Check if all values are within the bounds - default_joint_pos_torch = articulation._data.default_joint_pos.torch - within_bounds = (default_joint_pos_torch >= limits[..., 0]) & (default_joint_pos_torch <= limits[..., 1]) - assert torch.all(within_bounds) - - # Set new joint limits that invalidate default joint pos with indexing - limits = torch.zeros(env_ids.shape[0], joint_ids.shape[0], 2, device=device) - limits[..., 0] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) * -0.1 - limits[..., 1] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) * 0.1 - articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) - - # Check if all values are within the bounds - default_joint_pos_torch = articulation._data.default_joint_pos.torch - within_bounds = (default_joint_pos_torch[env_ids][:, joint_ids] >= limits[..., 0]) & ( - default_joint_pos_torch[env_ids][:, joint_ids] <= limits[..., 1] - ) - assert torch.all(within_bounds) +@pytest.mark.parametrize("scene", _CUDA_DEVICES, 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[island] + device = scene.device + assert (articulation.joint_ordering is not None) == (island == "reversed_ordered") -@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. + 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) - 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 + 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) - 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 + generalized_energy = 0.5 * torch.einsum( + "ni,nij,nj->n", joint_velocity, articulation.data.mass_matrix.torch, joint_velocity ) - 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 + 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:]) ) - 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. + torch.testing.assert_close(generalized_energy, body_energy, atol=1e-5, rtol=1e-5) - 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, + 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 ) - articulation, _ = generate_articulation( - articulation_cfg=articulation_cfg, - num_articulations=1, - device=device, + 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 ) - # 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) - # 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_DEVICES, 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["floating"] + device = scene.device + num_articulations = articulation.num_instances + num_bodies = articulation.num_bodies + 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) -@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 + # 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], + ) + 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], ) - - # Play the simulator - sim.reset() - - # Now we are ready! - # reset articulation - articulation.reset() # Reset should zero external forces and torques + articulation.reset() 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 @@ -2183,7 +933,6 @@ def test_reset(sim, num_articulations, device): assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_force.torch) == 0 assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_torque.torch) == 0 - 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), @@ -2199,352 +948,228 @@ def test_reset(sim, num_articulations, device): 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 + articulation.reset() - -@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 + # 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) + com[:, 0, :3] = torch.tensor([1.0, 0.0, 0.0], device=device) + articulation.set_coms_index(coms=wp.from_torch(com.contiguous(), 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() - # 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 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) - - -@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 - 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), - ) +@pytest.mark.parametrize("scene", _CUDA_DEVICES, indirect=True) +def test_tendon_properties_and_position_targets(scene: _ArticulationScene) -> None: + """Prove locally authored spatial and fixed tendons are discovered and writable. - # 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 + 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. """ - 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() + articulation = scene["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.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 ) - # 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_DEVICES, 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["native"] + device = scene.device + assert articulation._actuator_control.native_actuator_path_active + assert articulation.newton_actuator_adapter is not None + backend_to_user = _backend_to_user(articulation) + 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_DEVICES, 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.""" + from isaaclab.actuators.newton import read_group_parameter # noqa: PLC0415 + from isaaclab.envs.mdp.events import randomize_actuator_gains # noqa: PLC0415 + from isaaclab.managers import EventTermCfg, SceneEntityCfg # 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["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_DEVICES, 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["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()) @@ -2560,308 +1185,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 = SimpleNamespace(name="robot", body_ids=slice(None)) - env = SimpleNamespace(scene={"robot": articulation}, sim=sim) - randomize = randomize_rigid_body_material(SimpleNamespace(params={**params, "asset_cfg": asset_cfg}), env) - - # Randomize only the last environment; the others keep their materials. - randomize(env, torch.tensor([num_articulations - 1], device=device), *params.values(), asset_cfg) - 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) @@ -2870,8 +1193,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_DEVICES, 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 @@ -2880,18 +1203,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["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) @@ -2931,7 +1244,3 @@ def occupy_device_stream(iterations: int = 500_000): torch.testing.assert_close(_read_binding_to_torch(articulation, TT.DOF_STIFFNESS, device), expected_stiffness) torch.testing.assert_close(_read_binding_to_torch(articulation, TT.DOF_DAMPING, device), expected_damping) torch.testing.assert_close(_read_binding_to_torch(articulation, TT.BODY_MASS, device), expected_masses) - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "--maxfail=1"]) diff --git a/source/isaaclab_ov/test/assets/unit/__init__.py b/source/isaaclab_ov/test/assets/unit/__init__.py new file mode 100644 index 000000000000..460a30569089 --- /dev/null +++ b/source/isaaclab_ov/test/assets/unit/__init__.py @@ -0,0 +1,4 @@ +# 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 diff --git a/source/isaaclab_ov/test/assets/unit/test_ovphysx_actuator_control.py b/source/isaaclab_ov/test/assets/unit/test_ovphysx_actuator_control.py new file mode 100644 index 000000000000..38a05f280ec1 --- /dev/null +++ b/source/isaaclab_ov/test/assets/unit/test_ovphysx_actuator_control.py @@ -0,0 +1,71 @@ +# 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 + +"""OVPhysX actuator-control unit tests.""" + +from __future__ import annotations + +import importlib +from types import SimpleNamespace + +import pytest + +pytest.importorskip("ovphysx.types", reason="ovphysx wheel not installed") + +from isaaclab_ov.assets.articulation.actuator_control import OvPhysxActuatorControl # noqa: E402 + +from isaaclab.actuators import ImplicitActuatorCfg # noqa: E402 + +pytestmark = pytest.mark.unit + + +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"), + ) + 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)) diff --git a/source/isaaclab_ov/test/assets/unit/test_ovphysx_articulation.py b/source/isaaclab_ov/test/assets/unit/test_ovphysx_articulation.py new file mode 100644 index 000000000000..c15cd99c69e1 --- /dev/null +++ b/source/isaaclab_ov/test/assets/unit/test_ovphysx_articulation.py @@ -0,0 +1,148 @@ +# 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 + +"""OVPhysX articulation data-cache and joint-direction unit tests.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +import warp as wp + +from pxr import Usd, UsdGeom, UsdPhysics + +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.articulation.articulation_data import ArticulationData # noqa: E402 + +from isaaclab.utils.warp.launch_cache import _WarpLaunchCache # noqa: E402 + +pytestmark = pytest.mark.unit + + +def test_cached_read_launches_reset_on_ordering_and_invalidation(): + """Ordering installation and simulation invalidation should discard recorded reads.""" + + class MinimalData(ArticulationData): + def __dir__(self): + return [] + + class Buffer: + timestamp = 1.0 + + data = MinimalData.__new__(MinimalData) + read_launch_cache = Mock() + data._read_launch_cache = read_launch_cache + data._configure_ordering_buffers = lambda: None + data._make_jacobian_body_user_to_backend = lambda: object() + data.joint_ordering = None + data._body_com_jacobian_w = Buffer() + data._mass_matrix = Buffer() + data._gravity_compensation_forces = Buffer() + + data._apply_ordering_maps_after_resolve() + + read_launch_cache.clear.assert_called_once_with() + assert data._body_com_jacobian_w.timestamp == -1.0 + assert data._mass_matrix.timestamp == -1.0 + assert data._gravity_compensation_forces.timestamp == -1.0 + + data._is_primed = True + data._sim_timestamp = 1.0 + data._invalidate_initialize_callback(None) + + assert read_launch_cache.clear.call_count == 2 + assert data._is_primed is False + assert data._sim_timestamp == 0.0 + + +def test_static_property_reads_are_not_invalidated_by_simulation_steps(): + """Joint properties and body mass/inertia should be read once per invalidation, not once per step. + + On OVPhysX these are blocking CPU-only binding reads whose cost scales with the number of + environments, so re-reading them every physics step made per-step consumers (such as the + native actuator telemetry sync) host-bound. State buffers must still refresh every step. + """ + + class Buffer: + def __init__(self, shape): + self.data = wp.zeros(shape, dtype=wp.float32, device="cpu") + self.timestamp = -1.0 + + data = ArticulationData.__new__(ArticulationData) + data.device = "cpu" + data.num_instances = 1 + data.num_joints = 2 + data._sim_timestamp = 1.0 + data.body_ordering = None + data._get_binding = lambda tensor_type: object() + reads: list[int] = [] + data._binding_read = lambda tensor_type, dst: reads.append(tensor_type) + + # Joint properties: one read across several steps, one more after explicit invalidation + # (simulation reinitialization). + data.joint_ordering = None + stiffness = Buffer((1, 2)) + for _ in range(3): + data._sim_timestamp += 1.0 + data._read_joint_property_binding(TT.DOF_STIFFNESS, stiffness, None) + assert reads.count(TT.DOF_STIFFNESS) == 1 + stiffness.timestamp = -1.0 + data._read_joint_property_binding(TT.DOF_STIFFNESS, stiffness, None) + assert reads.count(TT.DOF_STIFFNESS) == 2 + + # Body properties behave the same; body state buffers still refresh every step. + mass, link_pose = Buffer((1, 2)), Buffer((1, 2)) + for _ in range(3): + data._sim_timestamp += 1.0 + data._refresh_reordered_body_buffer(mass, None, TT.BODY_MASS, static=True) + data._refresh_reordered_body_buffer(link_pose, None, TT.LINK_POSE) + assert reads.count(TT.BODY_MASS) == 1 + assert reads.count(TT.LINK_POSE) == 3 + mass.timestamp = -1.0 + data._refresh_reordered_body_buffer(mass, None, TT.BODY_MASS, static=True) + assert reads.count(TT.BODY_MASS) == 2 + + # Under a non-identity joint ordering the property is gathered once, then served from cache. + data.joint_ordering = SimpleNamespace(user_to_backend=wp.array([1, 0], dtype=wp.int32, device="cpu")) + data._read_launch_cache = _WarpLaunchCache("cpu") + user_buffer, backend_buffer = Buffer((1, 2)), Buffer((1, 2)) + data._binding_read = lambda tensor_type, dst: (reads.append(tensor_type), dst.assign([[1.0, 2.0]])) + for _ in range(3): + data._sim_timestamp += 1.0 + data._read_joint_property_binding(TT.DOF_DAMPING, user_buffer, backend_buffer) + assert reads.count(TT.DOF_DAMPING) == 1 + torch.testing.assert_close(wp.to_torch(user_buffer.data), torch.tensor([[2.0, 1.0]])) + + +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,) From 62c52befc71e1c3b4ec3f08e513c73946928d5f1 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 10:45:53 +0200 Subject: [PATCH 06/33] Consolidate OVPhysX rigid-object tests into composite scenes Rigid-object and collection tests share one module-scoped scene of locally authored cuboids per device instead of rebuilding Nucleus props per test. They keep the partial-state, raw inertial, fused-layout, material, wrench, kinematic, and scene-data-provider checks, and the reset checks that make both wrench composers active first. A strict xfail pins a collection bug: a partial body pose write pushes whole environment rows from a pose buffer that is not refreshed first, so the unselected bodies of that environment are rewound to stale poses. The CPU deformable rejection test no longer needs a spawned child process. --- .../test/assets/test_deformable_object.py | 57 +- .../test/assets/test_rigid_object.py | 1121 +++-------------- .../assets/test_rigid_object_collection.py | 934 +++----------- 3 files changed, 383 insertions(+), 1729 deletions(-) diff --git a/source/isaaclab_ov/test/assets/test_deformable_object.py b/source/isaaclab_ov/test/assets/test_deformable_object.py index 75453d4e0cc2..5b207f4fe459 100644 --- a/source/isaaclab_ov/test/assets/test_deformable_object.py +++ b/source/isaaclab_ov/test/assets/test_deformable_object.py @@ -10,11 +10,8 @@ from __future__ import annotations -import multiprocessing -import queue +import re import sys -import traceback -from typing import Any import ovphysx.types # noqa: F401 import pytest @@ -130,23 +127,6 @@ 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) - - @pytest.mark.parametrize( "num_objects, material_path", [ @@ -199,30 +179,17 @@ 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) diff --git a/source/isaaclab_ov/test/assets/test_rigid_object.py b/source/isaaclab_ov/test/assets/test_rigid_object.py index 4253adb3807f..d58bcaa56c7e 100644 --- a/source/isaaclab_ov/test/assets/test_rigid_object.py +++ b/source/isaaclab_ov/test/assets/test_rigid_object.py @@ -6,37 +6,24 @@ # ignore private usage of variables warning # pyright: reportPrivateUsage=none +"""Real OVPhysX rigid-object coverage on one module-scoped scene per device. -"""Real-backend tests for the OVPhysX RigidObject. - -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. +Each device builds one scene of locally authored cuboid pairs, resets it once, and keeps it alive for +every test that uses it. Each pair belongs to one test, so the tests do not depend on each other's +order. Initialization failures need their own scenes and therefore build them before any shared scene. """ 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 @@ -45,974 +32,206 @@ import isaaclab.sim as sim_utils # noqa: E402 from isaaclab.assets import RigidObjectCfg # noqa: E402 -from isaaclab.sim import SimulationCfg, build_simulation_context # noqa: E402 -from isaaclab.sim.spawners import materials # 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() +from isaaclab.sim import SimulationCfg, SimulationContext, build_simulation_context # noqa: E402 +from isaaclab.test.utils import DeviceScope, test_devices # noqa: E402 -pytestmark = pytest.mark.device_split +pytestmark = pytest.mark.integration -_logger = logging.getLogger(__name__) +_NUM_CUBES = 2 -_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: - 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 _sim_context(device: str): + """Build a local OVPhysX context from an in-memory USD stage.""" + return build_simulation_context(device=device, sim_cfg=SimulationCfg(physics=OvPhysxCfg(), device=device)) -def _ovphysx_sim_context(device: str, **kwargs): - """Wrapper around :func:`build_simulation_context` that injects OVPhysX cfg. - - PhysX tests pass ``device=device`` directly and let - :func:`build_simulation_context` build a default :class:`SimulationCfg`. - OVPhysX needs ``physics=OvPhysxCfg()`` set on the cfg so the manager - dispatches to OVPhysX rather than PhysX, so we build the cfg here and - pass it through. ``gravity_enabled`` is consumed locally (it is ignored - by ``build_simulation_context`` once a ``sim_cfg`` is provided). - ``add_ground_plane``, ``auto_add_lighting``, and other kwargs continue - to flow through ``build_simulation_context`` as before. - """ - dt = kwargs.pop("dt", 1.0 / 60.0) - gravity_enabled = kwargs.pop("gravity_enabled", True) - gravity = (0.0, 0.0, -9.81) if gravity_enabled else (0.0, 0.0, 0.0) - sim_cfg = SimulationCfg(physics=OvPhysxCfg(), device=device, dt=dt, gravity=gravity) - return build_simulation_context(device=device, sim_cfg=sim_cfg, **kwargs) - - -def generate_cubes_scene( - num_cubes: int = 1, - height=1.0, - api: Literal["none", "rigid_body"] = "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(), +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)), ) - 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) +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)), + ) ) -@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 - 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 - 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): +@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) +def test_initialization_with_no_rigid_body(device): """Test that initialization fails when no rigid body is found at the provided prim path.""" - with _ovphysx_sim_context(device=device, auto_add_lighting=True) as sim: - # Generate cubes scene - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, api="none", device=device) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(cube_object) < 10 - - # Play sim + with _sim_context(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): """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 _sim_context(device) as sim: + # Generate cubes and mark each rigid body as an articulation root. Keep the asset 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. +@dataclass +class _RigidScene: + """Cuboid pairs that share one real OVPhysX lifecycle.""" - 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. + sim: SimulationContext + device: str + dynamic: RigidObject + kinematic: RigidObject + resettable: RigidObject - 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 +@pytest.fixture(scope="module") +def scene(request: pytest.FixtureRequest) -> Iterator[_RigidScene]: + """Initialize every cuboid pair for one device once for this module.""" + device = request.param + with _sim_context(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) + resettable = _spawn_cubes("Resettable", y_offset=4.0, disable_gravity=True) sim.reset() + yield _RigidScene(sim=sim, device=device, dynamic=dynamic, kinematic=kinematic, resettable=resettable) + + +@pytest.mark.parametrize("scene", test_devices(), indirect=True) +def test_rigid_object_real_ovphysx_seams(scene: _RigidScene) -> None: + """Prove partial state, inertial properties, and one real wrench delivery.""" + 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_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 + ) - # 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) - ) - 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 - 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 - sim.reset() - - # Get masses before increasing - original_masses = cube_object.data.body_mass.torch.clone() - - assert original_masses.shape == (num_cubes, 1) - - # Randomize mass of the object - masses = original_masses + torch.FloatTensor(num_cubes, 1).uniform_(4, 8).to(sim.device) - - indices = torch.tensor(range(num_cubes), dtype=torch.int32) - - # Set the new masses via the OVPhysX writer (matches PhysX/Newton). - cube_object.set_masses_index( - masses=wp.from_torch(masses.contiguous(), dtype=wp.float32), - env_ids=wp.from_torch(indices, dtype=wp.int32), - ) - - torch.testing.assert_close(cube_object.data.body_mass.torch, masses) - - # Simulate physics - # perform rendering + 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) + + +@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() + expected_paths = {f"/World/{name}/Env_{i}/Cube" for name in ("Dynamic", "Kinematic", "Resettable") for i in (0, 1)} + 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) + 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) - # 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()) +@pytest.mark.parametrize("scene", test_devices(DeviceScope.CUDA), indirect=True) +def test_reset_clears_active_wrench_composers(scene: _RigidScene) -> None: + """Test that resetting the rigid object clears both active wrench composers.""" + cube_object, device = scene.resettable, scene.device -@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) - - # 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() - - # 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) - - # 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() - - # 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), - ) - - # 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), - ) - - # 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) - - 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:]) + # 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 + + 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 diff --git a/source/isaaclab_ov/test/assets/test_rigid_object_collection.py b/source/isaaclab_ov/test/assets/test_rigid_object_collection.py index 34c7a8848222..1acfe0d5b32a 100644 --- a/source/isaaclab_ov/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_ov/test/assets/test_rigid_object_collection.py @@ -6,32 +6,22 @@ # ignore private usage of variables warning # pyright: reportPrivateUsage=none +"""Real OVPhysX rigid-object-collection coverage on one module-scoped scene per device. -"""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. +Each device builds one scene of locally authored ``N=2, B=3`` collections, resets it once, and keeps it +alive for every test that uses it. Each collection belongs to one test, so the tests do not depend on each +other's order. Initialization failures need their own scenes and therefore build them before any shared scene. """ from __future__ import annotations -import sys +from collections.abc import Iterator +from dataclasses import dataclass 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 @@ -40,740 +30,218 @@ 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) +from isaaclab.sim import SimulationCfg, SimulationContext, build_simulation_context # noqa: E402 +from isaaclab.test.utils import DeviceScope, test_devices # noqa: E402 - return cube_object_colection, origins +pytestmark = pytest.mark.integration +_NUM_ENVS, _NUM_BODIES = 2, 3 -@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), - ] +def _sim_context(device: str): + """Build a local OVPhysX context from an in-memory USD stage.""" + return build_simulation_context( + device=device, sim_cfg=SimulationCfg(physics=OvPhysxCfg(), device=device, gravity=(0.0, 0.0, 0.0)) + ) - 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 +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(), ) - 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 + 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")) + } ) + ) - # 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): +@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) +def test_initialization_with_no_rigid_body(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 _sim_context(device) 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()), + ) with pytest.raises(RuntimeError, match="Expected 1 prims"): sim.reset() + assert not object_collection.is_initialized -@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() +@dataclass +class _CollectionScene: + """Collections that share one real OVPhysX lifecycle.""" - # 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() + sim: SimulationContext + device: str + fused: RigidObjectCollection + resettable: RigidObjectCollection + stepped: RigidObjectCollection - # 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) +@pytest.fixture(scope="module") +def scene(request: pytest.FixtureRequest) -> Iterator[_CollectionScene]: + """Initialize every collection for one device once for this module.""" + device = request.param + with _sim_context(device) as sim: + fused = _spawn_collection("Fused", y_offset=0.0) + resettable = _spawn_collection("Resettable", y_offset=4.0) + stepped = _spawn_collection("Stepped", y_offset=8.0) 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 + yield _CollectionScene(sim=sim, device=device, fused=fused, resettable=resettable, stepped=stepped) + + +@pytest.mark.parametrize("scene", test_devices(), indirect=True) +def test_rigid_object_collection_real_ovphysx_seams(scene: _CollectionScene) -> 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() + 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) + + +@pytest.mark.parametrize("scene", test_devices(DeviceScope.CUDA), indirect=True) +def test_reset_clears_active_wrench_composers(scene: _CollectionScene) -> None: + """Test that resetting the collection clears both active wrench composers.""" + object_collection, device = scene.resettable, scene.device + + # 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_BODIES, 3), device=device), + torques=torch.ones((_NUM_ENVS, _NUM_BODIES, 3), device=device), + ) + object_collection.instantaneous_wrench_composer.add_forces_and_torques_index( + forces=torch.ones((_NUM_ENVS, _NUM_BODIES, 3), device=device), + torques=torch.ones((_NUM_ENVS, _NUM_BODIES, 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("scene", test_devices(DeviceScope.CUDA), indirect=True) +@pytest.mark.xfail( + strict=True, + 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: _CollectionScene) -> 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() - 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) + 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) From feff387850265166a2fadcc2cb38c073b2e90729 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 10:45:57 +0200 Subject: [PATCH 07/33] Move OVPhysX asset unit tests under unit --- .../test/assets/{ => unit}/test_articulation_helpers.py | 3 +++ .../test/assets/{ => unit}/test_articulation_kernels.py | 2 ++ .../test/assets/{ => unit}/test_deformable_views.py | 3 +++ .../test_ovphysx_deformable_object.py} | 3 +++ 4 files changed, 11 insertions(+) rename source/isaaclab_ov/test/assets/{ => unit}/test_articulation_helpers.py (99%) rename source/isaaclab_ov/test/assets/{ => unit}/test_articulation_kernels.py (98%) rename source/isaaclab_ov/test/assets/{ => unit}/test_deformable_views.py (99%) rename source/isaaclab_ov/test/assets/{test_deformable_object_helpers.py => unit/test_ovphysx_deformable_object.py} (99%) diff --git a/source/isaaclab_ov/test/assets/test_articulation_helpers.py b/source/isaaclab_ov/test/assets/unit/test_articulation_helpers.py similarity index 99% rename from source/isaaclab_ov/test/assets/test_articulation_helpers.py rename to source/isaaclab_ov/test/assets/unit/test_articulation_helpers.py index ca43d1b1c46b..b2fb56274e71 100644 --- a/source/isaaclab_ov/test/assets/test_articulation_helpers.py +++ b/source/isaaclab_ov/test/assets/unit/test_articulation_helpers.py @@ -26,6 +26,9 @@ from isaaclab_ov.physics import OvPhysxManager # noqa: E402 from isaaclab_ov.test.fixtures.views import MockOvPhysxBindingSet # noqa: E402 +pytestmark = pytest.mark.unit + + wp.init() diff --git a/source/isaaclab_ov/test/assets/test_articulation_kernels.py b/source/isaaclab_ov/test/assets/unit/test_articulation_kernels.py similarity index 98% rename from source/isaaclab_ov/test/assets/test_articulation_kernels.py rename to source/isaaclab_ov/test/assets/unit/test_articulation_kernels.py index b1da63857647..3cb1b7abd9e5 100644 --- a/source/isaaclab_ov/test/assets/test_articulation_kernels.py +++ b/source/isaaclab_ov/test/assets/unit/test_articulation_kernels.py @@ -10,6 +10,8 @@ import warp as wp from isaaclab_ov.assets import kernels +pytestmark = pytest.mark.unit + def _selector(values: list[int], dtype: type) -> wp.array: return wp.array(values, dtype=dtype, device="cpu") diff --git a/source/isaaclab_ov/test/assets/test_deformable_views.py b/source/isaaclab_ov/test/assets/unit/test_deformable_views.py similarity index 99% rename from source/isaaclab_ov/test/assets/test_deformable_views.py rename to source/isaaclab_ov/test/assets/unit/test_deformable_views.py index f3f9d1c1a98a..6211cf9d9d60 100644 --- a/source/isaaclab_ov/test/assets/test_deformable_views.py +++ b/source/isaaclab_ov/test/assets/unit/test_deformable_views.py @@ -17,6 +17,9 @@ from isaaclab_ov.assets.deformable_object.views import OvPhysxDeformableBodyView # noqa: E402 from ovphysx.types import TensorType # noqa: E402 +pytestmark = pytest.mark.unit + + wp.init() wp.set_device("cpu") diff --git a/source/isaaclab_ov/test/assets/test_deformable_object_helpers.py b/source/isaaclab_ov/test/assets/unit/test_ovphysx_deformable_object.py similarity index 99% rename from source/isaaclab_ov/test/assets/test_deformable_object_helpers.py rename to source/isaaclab_ov/test/assets/unit/test_ovphysx_deformable_object.py index e60b62e7a7ae..fc1e2441eff6 100644 --- a/source/isaaclab_ov/test/assets/test_deformable_object_helpers.py +++ b/source/isaaclab_ov/test/assets/unit/test_ovphysx_deformable_object.py @@ -24,6 +24,9 @@ from isaaclab_ov.assets.deformable_object.kernels import vec6f # noqa: E402 from isaaclab_ov.physics.ovphysx_manager import OvPhysxManager, OvPhysxSceneDataBackend # noqa: E402 +pytestmark = pytest.mark.unit + + wp.init() wp.set_device("cpu") From cf825d4a7d58a7d46c7f695b0b1c6892d78f1367 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 10:52:37 +0200 Subject: [PATCH 08/33] Cover partial joint-state writes and the deprecated joint-state writer Take over what the removed OVPhysX real-backend joint-state test checked, on every backend: unsorted int64 environment and joint selectors on the fused writer, the finite-difference baseline it restarts for the written joints, skip_forward leaving body caches to the caller, and the deprecated combined writer producing the same state as the public writers. --- .../contract/_articulation_contract_cases.py | 76 +++++++++++++++++++ 1 file changed, 76 insertions(+) diff --git a/source/isaaclab/test/assets/contract/_articulation_contract_cases.py b/source/isaaclab/test/assets/contract/_articulation_contract_cases.py index 4d314faa281f..d1a0a521f322 100644 --- a/source/isaaclab/test/assets/contract/_articulation_contract_cases.py +++ b/source/isaaclab/test/assets/contract/_articulation_contract_cases.py @@ -24,6 +24,7 @@ from isaaclab.test.utils import DeviceScope, test_devices from ._articulation_contract_utils import get_articulation +from ._articulation_ordering_contract_cases import _read_backend_joint_state from .capabilities import contract_backend, require_backend_capability, requires_backend # --------------------------------------------------------------------------- @@ -988,6 +989,81 @@ def _read_joint_writer_target(art, getter: str): class TestArticulationWritersJoint: """Test joint writers/setters with all input combinations.""" + @_writes_backends + def test_partial_joint_state_write_follows_int64_selectors_and_restarts_acceleration(self, backend): + """A fused joint-state write honors unsorted int64 selectors and restarts the finite-difference baseline.""" + art, raw_backend = get_articulation(backend, 2, 3, 2, device="cpu") + 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. + caches = ("_body_link_vel_w",) if backend == "newton" else ("_body_link_pose_w", "_body_com_vel_w") + for cache in caches: + getattr(art.data, cache[1:]) + # 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 cache in caches: + assert getattr(art.data, cache).timestamp == art.data._sim_timestamp, cache + art.write_joint_state_to_sim_index(position=position, velocity=velocity, env_ids=[1], joint_ids=[0, 2]) + for cache in caches: + assert getattr(art.data, cache).timestamp < art.data._sim_timestamp, cache + + @_writes_backends + def test_deprecated_joint_state_writer_matches_public_writers(self, backend): + """The deprecated combined joint-state writer warns and writes what the public writers would.""" + art, raw_backend = get_articulation(backend, 2, 3, 2, device="cpu") + 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) + @_writes_backends @_devices @pytest.mark.parametrize( From 6eab6fb92fa434c608005da342331898710f0a09 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 11:06:03 +0200 Subject: [PATCH 09/33] Consolidate PhysX articulation tests into one composite scene Run the PhysX articulation checks against one module-scoped scene built from the local branching fixture instead of per-test Nucleus robots. The scene holds two-environment islands (fixed with MJWarp ordering and USD gains, floating with configured gains, floating with a root-moving body order, and fixed with a reversed joint and a spatial tendon), so partial writes target environment 1 and prove environment 0 is preserved in the real PhysX state. Keep the Shadow Hand fixed-tendon, Franka gravity-compensation, and Franka IK tests on their shipped assets, and check initialization failures on the local fixture. Move the kernel unit tests under unit/ and cover the effort termination term with a core unit test. --- .../antoiner-asset-tests-v2-physx.skip | 0 .../test/envs/mdp/test_terminations.py | 31 + .../antoiner-asset-tests-v2-physx.skip | 0 .../test/assets/test_articulation.py | 3142 ++++++----------- .../test_physx_articulation.py} | 0 5 files changed, 1117 insertions(+), 2056 deletions(-) create mode 100644 source/isaaclab/changelog.d/antoiner-asset-tests-v2-physx.skip create mode 100644 source/isaaclab/test/envs/mdp/test_terminations.py create mode 100644 source/isaaclab_physx/changelog.d/antoiner-asset-tests-v2-physx.skip rename source/isaaclab_physx/test/assets/{test_articulation_kernels.py => unit/test_physx_articulation.py} (100%) diff --git a/source/isaaclab/changelog.d/antoiner-asset-tests-v2-physx.skip b/source/isaaclab/changelog.d/antoiner-asset-tests-v2-physx.skip new file mode 100644 index 000000000000..e69de29bb2d1 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..e885a2d029f8 --- /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(): + """Terminate exactly the environments whose applied effort differs from the computed effort on selected joints.""" + computed = torch.tensor([[10.0, 20.0, 30.0], [10.0, 20.0, 30.0], [10.0, 20.0, 30.0], [10.0, 20.0, 30.0]]) + # Environment 0 clips the selected joint, environment 1 clips nothing, environment 2 clips the selected joint + # downward, and environment 3 clips only joints that are not selected. + applied = torch.tensor([[10.0, 19.0, 30.0], [10.0, 20.0, 30.0], [10.0, 17.0, 30.0], [9.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=[1])) + + torch.testing.assert_close(result, torch.tensor([True, False, True, False])) diff --git a/source/isaaclab_physx/changelog.d/antoiner-asset-tests-v2-physx.skip b/source/isaaclab_physx/changelog.d/antoiner-asset-tests-v2-physx.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_physx/test/assets/test_articulation.py b/source/isaaclab_physx/test/assets/test_articulation.py index 7078c528911e..2bb4c02d9073 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,169 +21,182 @@ 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 - launch_test_simulation() +import math import sys +from dataclasses import dataclass from pathlib import Path import pytest import torch import warp as wp from isaaclab_physx.assets import Articulation -from isaaclab_physx.sim.schemas import PhysxJointCfg -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 +from isaaclab.sim import SimulationContext, build_simulation_context +from isaaclab.utils.math import ( + combine_frame_transforms, + compute_pose_error, + matrix_from_quat, + quat_apply, + quat_inv, + subtract_frame_transforms, +) ## # 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_assets import FRANKA_PANDA_CFG, FRANKA_PANDA_HIGH_PD_CFG # isort:skip +from isaaclab_assets.robots.shadow_hand import SHADOW_HAND_PHYSX_CFG # isort:skip + +_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} + + +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) + + +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, + 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. + joint_limits_deg: Lower and upper joint limits [deg] applied to every joint. PhysX ignores limit writes + on joints authored without limits. + """ + 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) -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". +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) - 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, - ), - }, - ) - 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'." - ) +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]] * _NUM_ENVS, device=device) + + +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.fixture +def sim(request): + """Create a function-scoped simulation context for tests that own their scene.""" + device = request.getfixturevalue("device") + gravity_enabled = request.getfixturevalue("gravity_enabled") if "gravity_enabled" in request.fixturenames else True + 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 + - return articulation_cfg +## +# Tests that own their simulation context. Keep them above the composite scene. +## def generate_articulation( articulation_cfg: ArticulationCfg, num_articulations: int, device: str -) -> tuple[Articulation, torch.tensor]: - """Generate an articulation from a configuration. - - Handles the creation of the articulation, the environment prims and the articulation's environment - translations +) -> tuple[Articulation, torch.Tensor]: + """Spawn ``num_articulations`` copies of an articulation 2.5 m apart along x. Args: articulation_cfg: Articulation configuration. @@ -183,20 +205,174 @@ def generate_articulation( Returns: The articulation and environment translations. - """ - # Generate translations of 2.5 m in x for each articulation translations = torch.zeros(num_articulations, 3, device=device) translations[:, 0] = torch.arange(num_articulations) * 2.5 - - # Create Top-level Xforms, one for each articulation for i in range(num_articulations): sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=translations[i][:3]) articulation = Articulation(replace(articulation_cfg, prim_path="/World/Env_[^/]*/Robot")) - return articulation, translations +@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): + """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)) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + with pytest.raises(ValueError, match=message): + sim.reset() + + +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +def test_setting_invalid_articulation_root_prim_path(sim, device): + """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 + + with pytest.raises(RuntimeError, match="Failed to create articulation at"): + sim.reset() + + +@pytest.mark.parametrize("num_articulations", [2]) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +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. + + 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. The Shadow Hand is the shipped asset whose + native PhysX fixed tendons these writers command. + """ + articulation_cfg = SHADOW_HAND_PHYSX_CFG + 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 + 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", [1]) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +@pytest.mark.isaacsim_ci +def test_get_gravity_compensation_forces_static_equilibrium(sim, num_articulations, device): + """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). + """ + # 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( + FRANKA_PANDA_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." + ) + + # --------------------------------------------------------------------------- # Franka task-space tracking helpers for the IK test. # Mirrors the helpers in ``isaaclab_newton/test/assets/test_articulation.py``. @@ -266,1987 +442,841 @@ def _summarize_history(history, tail: int = 200): return min(tail_slice), sum(tail_slice) / len(tail_slice) -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) - +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +@pytest.mark.parametrize("gravity_enabled", [False]) +@pytest.mark.isaacsim_ci +def test_franka_ik_tracking_accuracy(sim, device, gravity_enabled): + """PhysX-side IK convergence sentinel — backend parity with the Newton test. -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]) + 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) -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]) + 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] -@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 + 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) -@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) + 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()) - 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 + pos_min, pos_mean = _summarize_history(pos_history) + rot_min, rot_mean = _summarize_history(rot_history) - 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) + print(f"IK_METRIC pos_min={pos_min:.5f} pos_mean={pos_mean:.5f} rot_min={rot_min:.5f} rot_mean={rot_mean:.5f}") - 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 + # 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?" - 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() - _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, - ) +## +# Composite scene shared by the remaining tests. +## - 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, - ) - _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, - ) - _assert_backend_to_user( - articulation.data.body_com_pose_b.torch, - _to_device_tensor(articulation.root_view.get_coms(), device), - body_user_to_backend, - ) - - # 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:] - ) - 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") +@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 and a spatial 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.""" + + @property + def islands(self) -> dict[str, Articulation]: + """All articulations keyed by island name.""" + return {name: getattr(self, name) for name in ("ordered", "floating", "reordered", "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()) +def articulation_scene(request) -> _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={ + "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, + ) + }, + ), + {"root_on_base": True}, + ), + "reordered": ( + _branching_cfg( + articulation_root_prim_path="/base", + body_ordering=("left_tip", "base", "right_upper", "right_tip", "left_upper"), + ), + {"root_on_base": True}, ), - actuators={}, + "tendon": (_branching_cfg(), {"fixed_base": True, "reversed_left_elbow": True, "spatial_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() + yield _ArticulationScene( + sim=sim, + device=device, + origins={name: value.to(device) for name, value in origins.items()}, + refcounts=refcounts, + **islands, ) - ) - 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 - ) - 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) - - -@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)), - ) - ) - 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 +def test_articulation_initialization_and_ordering(articulation_scene: _ArticulationScene): + """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 getattr(actuator, "is_implicit_model", False) == is_implicit_model_cfg + 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 - assert ordered.body_ordering.backend_to_user_indices[0] != 0 + # 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(_to_device_tensor(raw, 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 attribute, raw, values in ( + ("stiffness", floating.root_view.get_dof_stiffnesses(), (6.0, 4.0)), + ("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(getattr(floating.actuators["joints"], attribute), expected) + torch.testing.assert_close(_to_device_tensor(raw, 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(_to_device_tensor(raw, device), public) + + +def test_articulation_joint_state_writes_follow_ordering(articulation_scene: _ArticulationScene): + """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) - 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 - backend_coms[..., 6] = 1.0 - identity.set_coms_index( - coms=wp.from_torch(backend_coms.contiguous(), dtype=wp.transformf), - full_data=True, + # 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(_to_device_tensor(raw, 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), ) - 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(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, _to_device_tensor(raw, 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(_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) - articulation.write_root_link_velocity_to_sim_index(root_velocity=root_link_velocity) - + 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( - _to_device_tensor(ordered.root_view.get_root_transforms(), device), - _to_device_tensor(identity.root_view.get_root_transforms(), device), + _to_device_tensor(articulation.root_view.get_dof_positions(), device), + expected_position[:, joint_backend_to_user], ) torch.testing.assert_close( - _to_device_tensor(ordered.root_view.get_root_velocities(), device), - _to_device_tensor(identity.root_view.get_root_velocities(), device), + _to_device_tensor(articulation.root_view.get_dof_velocities(), device), + expected_velocity[:, joint_backend_to_user], ) - 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) - + # Restore the USD gains for the following tests. + articulation.write_joint_stiffness_to_sim_index(stiffness=initial_stiffness, full_data=True) -@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) +def test_articulation_joint_and_body_properties_round_trip(articulation_scene: _ArticulationScene): + """Selected joint and body properties reach the selected PhysX entries and read back in public order.""" + scene = articulation_scene + device = scene.device - cpu_env_ids = wp.array(list(range(articulation.num_instances)), dtype=wp.int32, device="cpu") + # Direct tensor-view mass and inertia writes become visible on the lazy read: on the first read of a cold + # buffer (identity order), and after the next update once the buffer is primed (non-identity order). + cpu_env_ids = wp.array(list(range(_NUM_ENVS)), dtype=wp.int32, device="cpu") - def write_backend_mass_inertia(delta_mass: float, delta_inertia: float) -> tuple[torch.Tensor, torch.Tensor]: + def write_backend_mass_inertia(articulation: Articulation, delta_mass: float, delta_inertia: float): """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 + 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.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 ) - return backend_masses, backend_inertias + return backend_masses.to(device), backend_inertias.to(device) - # 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) + backend_masses, backend_inertias = write_backend_mass_inertia(scene.tendon, 0.293, 0.023) + torch.testing.assert_close(scene.tendon.data.body_mass.torch, backend_masses) + torch.testing.assert_close(scene.tendon.data.body_inertia.torch, backend_inertias) - # 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", - ) + 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) + articulation.data.body_mass.torch # prime the lazy buffers + articulation.data.body_inertia.torch + 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() + 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), + } + expected_frictions = {name: getattr(articulation.data, name).torch.clone() for name in frictions} + 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(getattr(articulation.data, name).torch, 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 + for setter, kwarg, value, name in ( + (articulation.set_masses_index, "masses", masses, "body_mass"), + (articulation.set_coms_index, "coms", coms, "body_com_pose_b"), + (articulation.set_inertias_index, "inertias", inertias, "body_inertia"), + ): + expected = getattr(articulation.data, name).torch.clone() + expected[env_ids[:, None], body_ids] = value + setter(**{kwarg: value, "env_ids": env_ids, "body_ids": body_ids}) + torch.testing.assert_close(getattr(articulation.data, name).torch, expected) + 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(_to_device_tensor(raw, device), public[:, body_backend_to_user]) + + +def _rotate_z(angle: torch.Tensor, vector: tuple[float, float, float] | torch.Tensor) -> torch.Tensor: + """Rotate vectors about the world z axis by per-environment angles [rad].""" + 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 ) - 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 - - -@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) - - # 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) - - # -- 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) - - # 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) - - # -- 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) +def _cross_z(rate: torch.Tensor, vector: torch.Tensor) -> torch.Tensor: + """Return ``(rate * z) x vector`` for per-environment angular rates about the world z axis.""" + return torch.stack((-rate * vector[..., 1], rate * vector[..., 0], torch.zeros_like(rate)), -1) -@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. +def _expected_fixed_branching_state(articulation: Articulation, origins: torch.Tensor, com_b: torch.Tensor) -> dict: + """Compute link and center-of-mass kinematics of a fixed branching island from its joint state. - 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. + 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. """ - articulation_cfg = generate_articulation_cfg(articulation_type="shadow_hand") - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device) + 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): + """Check the link and center-of-mass Jacobians against body velocities and the mass-matrix contract.""" + 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. + body_velocities = { + "body_link_jacobian_w": data.body_link_lin_vel_w.torch.clone(), + "body_com_jacobian_w": data.body_com_lin_vel_w.torch.clone(), + } + ang_vel = data.body_link_ang_vel_w.torch.clone() + for jacobian_name, lin_vel in body_velocities.items(): + jacobian = getattr(data, jacobian_name).torch + # 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.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 - 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) +def test_articulation_drive_and_dynamics(articulation_scene: _ArticulationScene): + """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] ) - 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 + target_command.set_effort_index( + value=torch.tensor([[10.0]], device=device), joint_ids=[right_shoulder], env_ids=[1] ) - 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) - - # -- 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 - + 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( + _to_device_tensor(articulation.root_view.get_dof_position_targets(), device), + expected_targets[:, joint_backend_to_user], + ) + torch.testing.assert_close( + _to_device_tensor(articulation.root_view.get_dof_actuation_forces(), device), + expected_efforts[:, joint_backend_to_user], + ) + # 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 + ) -@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. + # 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, + ) + 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. + 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 + ) + _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() - 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, - } +def test_floating_articulation_root_writes_and_wrenches(articulation_scene: _ArticulationScene, monkeypatch): + """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 + ) + _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(_to_device_tensor(floating.root_view.get_root_transforms(), 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 + 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(_to_device_tensor(floating.root_view.get_coms(), device), backend_coms) + torch.testing.assert_close(_to_device_tensor(reordered.root_view.get_coms(), 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 = _to_device_tensor(reordered.root_view.get_root_transforms(), device).clone() + reordered_transforms[:, :3] -= island_offset + torch.testing.assert_close( + reordered_transforms, _to_device_tensor(floating.root_view.get_root_transforms(), device) + ) + torch.testing.assert_close( + _to_device_tensor(reordered.root_view.get_root_velocities(), device), + _to_device_tensor(floating.root_view.get_root_velocities(), 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( + floating.data.root_com_pose_w.torch, torch.cat((expected_com_pos, expected_com_quat), -1) + ) + torch.testing.assert_close(floating.data.root_com_vel_w.torch[:, 3:], floating.data.root_link_vel_w.torch[:, 3:]) - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device) + # 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]]]}, "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]]]}, + "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 + ) + scene.step() + response_value = getattr(floating.data, 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, also when the public body + # order differs from the PhysX link order. + _place_at_rest(reordered, reordered_rest_pose) + tips = reordered.find_bodies(["left_tip", "right_tip"], preserve_order=True)[0] + reordered.permanent_wrench_composer.set_forces_and_torques_index( + forces=torch.tensor([[[0.0, 4.0, 0.0], [0.0, -4.0, 0.0]]], device=device), body_ids=tips, env_ids=[1] + ) + scene.step() + tip_velocity = reordered.data.body_link_lin_vel_w.torch[:, tips] + tip_direction = quat_apply( + reordered.data.body_link_quat_w.torch[1, tips], torch.tensor([[0.0, 1.0, 0.0], [0.0, -1.0, 0.0]], device=device) + ) + 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) - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(articulation) < 10 + # A partial reset clears the selected environment; 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 = [] - # Play sim - with pytest.raises(ValueError): - sim.reset() + def record_actuator_reset(env_ids=None): + reset_env_ids.append(env_ids) + actuator_reset(env_ids) + 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)) + 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 -@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, +def test_spatial_tendon_properties_round_trip(articulation_scene: _ArticulationScene): + """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 + assert articulation.num_spatial_tendons == 1 + root_view = articulation.root_view + getters = { + "stiffness": root_view.get_spatial_tendon_stiffnesses, + "limit_stiffness": root_view.get_spatial_tendon_limit_stiffnesses, + "damping": root_view.get_spatial_tendon_dampings, + "offset": root_view.get_spatial_tendon_offsets, } - 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): - """Test write_joint_limits_to_sim API and when default pos falls outside of the new limits. - - This test verifies that: - 1. Joint limits can be set correctly - 2. Default positions are preserved when setting new limits - 3. Joint limits can be set with indexing - 4. Invalid joint positions are properly handled - - Args: - sim: The simulation fixture - num_articulations: Number of articulations to test - """ - # Create articulation - articulation_cfg = generate_articulation_cfg(articulation_type="panda") - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device) - - # Play sim - sim.reset() - # Check if articulation is initialized - assert articulation.is_initialized - - # Get current default joint pos - default_joint_pos = articulation._data.default_joint_pos.torch.clone() - - # Set new joint limits - limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) - limits[..., 0] = (torch.rand(num_articulations, articulation.num_joints, device=device) + 5.0) * -1.0 - limits[..., 1] = torch.rand(num_articulations, articulation.num_joints, device=device) + 5.0 - articulation.write_joint_position_limit_to_sim_index(limits=limits) - - # Check new limits are in place - torch.testing.assert_close(articulation._data.joint_pos_limits.torch, limits) - torch.testing.assert_close(articulation._data.default_joint_pos.torch, default_joint_pos) - - # Set new joint limits with indexing - env_ids = torch.arange(1, device=device, dtype=torch.int32) - joint_ids = torch.arange(2, device=device, dtype=torch.int32) - limits = torch.zeros(env_ids.shape[0], joint_ids.shape[0], 2, device=device) - limits[..., 0] = (torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) + 5.0) * -1.0 - limits[..., 1] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) + 5.0 - articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) - - # Check new limits are in place - torch.testing.assert_close(articulation._data.joint_pos_limits.torch[env_ids][:, joint_ids], limits) - torch.testing.assert_close(articulation._data.default_joint_pos.torch, default_joint_pos) - - # Set new joint limits that invalidate default joint pos - limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) - limits[..., 0] = torch.rand(num_articulations, articulation.num_joints, device=device) * -0.1 - limits[..., 1] = torch.rand(num_articulations, articulation.num_joints, device=device) * 0.1 - articulation.write_joint_position_limit_to_sim_index(limits=limits) - - # Check if all values are within the bounds - default_joint_pos_torch = articulation._data.default_joint_pos.torch - within_bounds = (default_joint_pos_torch >= limits[..., 0]) & (default_joint_pos_torch <= limits[..., 1]) - assert torch.all(within_bounds) - - # Set new joint limits that invalidate default joint pos with indexing - limits = torch.zeros(env_ids.shape[0], joint_ids.shape[0], 2, device=device) - limits[..., 0] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) * -0.1 - limits[..., 1] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) * 0.1 - articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) - - # Check if all values are within the bounds - default_joint_pos_torch = articulation._data.default_joint_pos.torch - within_bounds = (default_joint_pos_torch[env_ids][:, joint_ids] >= limits[..., 0]) & ( - default_joint_pos_torch[env_ids][:, joint_ids] <= limits[..., 1] - ) - assert torch.all(within_bounds) - - -@pytest.mark.parametrize("num_articulations", [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) + 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 + for name, value in values.items(): + getattr(articulation, f"set_spatial_tendon_{name}_index")(**{name: value}) 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"]) + for name, getter in getters.items(): + torch.testing.assert_close(_to_device_tensor(getter(), device), values[name]) + + # Partial writes change environment 1 in the data and in the solver and preserve environment 0. + env_ids = torch.tensor([1], dtype=torch.int32, device=device) + partial = {"stiffness": 12.0, "limit_stiffness": 3.0, "damping": 1.5, "offset": 0.1} + for name, value in partial.items(): + getattr(articulation, f"set_spatial_tendon_{name}_index")( + **{name: torch.tensor([[value]], device=device), "env_ids": env_ids} + ) + values[name][1] = value + torch.testing.assert_close(getattr(articulation.data, f"spatial_tendon_{name}").torch, values[name]) + articulation.write_spatial_tendon_properties_to_sim_index(env_ids=env_ids) + for name, getter in getters.items(): + torch.testing.assert_close(_to_device_tensor(getter(), device), values[name]) diff --git a/source/isaaclab_physx/test/assets/test_articulation_kernels.py b/source/isaaclab_physx/test/assets/unit/test_physx_articulation.py similarity index 100% rename from source/isaaclab_physx/test/assets/test_articulation_kernels.py rename to source/isaaclab_physx/test/assets/unit/test_physx_articulation.py From 161cf9f1ea0b7b3966a5cd77c78e8d4bbdfe0a6f Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 11:06:15 +0200 Subject: [PATCH 10/33] Consolidate PhysX rigid-object and collection tests into composite scenes Share one module-scoped scene of locally authored cubes per file instead of a Nucleus scene per test. Partial state, inertial, and wrench writes select environment 1 or non-sorted body-major subsets and read the PhysX view back; center-of-mass offsets check every root and body frame writer and the spin kinematics. Failure modes, the gravity vector, physics warmup, and the collection subset regressions keep their own scenes. The rigid-object articulation-root rejection is kept as a strict xfail: the check was removed from the PhysX rigid object, and the previous Nucleus fixture only failed because it resolved no rigid body. --- .../test/assets/test_rigid_object.py | 1082 ++++++----------- .../assets/test_rigid_object_collection.py | 866 ++++++------- 2 files changed, 716 insertions(+), 1232 deletions(-) diff --git a/source/isaaclab_physx/test/assets/test_rigid_object.py b/source/isaaclab_physx/test/assets/test_rigid_object.py index ac84e269f1b6..ee3d89174245 100644 --- a/source/isaaclab_physx/test/assets/test_rigid_object.py +++ b/source/isaaclab_physx/test/assets/test_rigid_object.py @@ -3,577 +3,125 @@ # # SPDX-License-Identifier: BSD-3-Clause +"""Real PhysX rigid-object coverage. + +Most checks run against one module-scoped scene of locally authored cubes. Each rigid object 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 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 dataclasses import dataclass from typing import Literal import pytest import torch import warp as wp -from flaky import flaky from isaaclab_physx.assets import RigidObject +from isaaclab_physx.sim.schemas import PhysxRigidBodyCfg + +from pxr import UsdPhysics 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, -) - - -def generate_cubes_scene( - num_cubes: int = 1, - height=1.0, - api: Literal["none", "rigid_body", "articulation_root"] = "rigid_body", - kinematic_enabled: bool = False, - device: str = "cuda:0", -) -> tuple[RigidObject, torch.Tensor]: - """Generate a scene with the provided number of cubes. - - Args: - num_cubes: Number of cubes to generate. - height: Height of the cubes. - api: The type of API that the cubes should have. - kinematic_enabled: Whether the cubes are kinematic. - device: Device to use for the simulation. - - Returns: - A tuple containing the rigid object representing the cubes and the origins of the cubes. - - """ - origins = torch.tensor([(i * 1.0, 0, height) for i in range(num_cubes)]).to(device) - # Create Top-level Xforms, one for each cube - for i, origin in enumerate(origins): - sim_utils.create_prim(f"/World/Table_{i}", "Xform", translation=origin) - - # Resolve spawn configuration - if api == "none": - # since no rigid body properties defined, this is just a static collider - spawn_cfg = sim_utils.CuboidCfg( - size=(0.1, 0.1, 0.1), +from isaaclab.sim import SimulationContext, build_simulation_context +from isaaclab.utils.math import combine_frame_transforms, quat_apply, quat_apply_inverse, quat_mul + +_NUM_ENVS = 2 + + +def _cube_cfg(prim_path: str, *, kinematic: bool = False, disable_gravity: bool = False) -> RigidObjectCfg: + """Create a local 1 kg collision cube one meter above its environment origin.""" + return RigidObjectCfg( + prim_path=prim_path, + spawn=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(), - ) - 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)), + ), + 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) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(cube_object) < 10 - - # Play sim - sim.reset() - - # Check if object is initialized - assert cube_object.is_initialized - assert len(cube_object.body_names) == 1 - - # Check buffers that exists and have correct shapes - assert cube_object.data.root_pos_w.torch.shape == (num_cubes, 3) - assert cube_object.data.root_quat_w.torch.shape == (num_cubes, 4) - assert cube_object.data.body_mass.torch.shape == (num_cubes, 1) - assert cube_object.data.body_inertia.torch.shape == (num_cubes, 1, 9) - - # Simulate physics - for _ in range(2): - # perform rendering - sim.step() - # update object - cube_object.update(sim.cfg.dt) - - -@pytest.mark.parametrize("num_cubes", [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.""" - 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() +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 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) +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)) - # 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) +## +# Tests that own their simulation context. Keep them above the composite scene. +## -@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 - - # Play sim - with pytest.raises(RuntimeError): - sim.reset() - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize( + "api", + [ + "none", + pytest.param( + "articulation_root", + marks=pytest.mark.xfail( + strict=True, reason="The PhysX rigid object no longer rejects an enabled articulation root." + ), + ), + ], +) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) @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.""" +def test_initialization_rejects_invalid_rigid_body(device, api: Literal["none", "articulation_root"]): + """Initialization fails without a rigid body and when the rigid body is an articulation root.""" 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="articulation_root", device=device) + _spawn_envs("/World") + if api == "none": + # Without rigid body properties the cube is a static collider. + cfg = RigidObjectCfg( + prim_path="/World/Env_[^/]*/Object", + spawn=sim_utils.CuboidCfg(size=(0.1, 0.1, 0.1), collision_props=sim_utils.UsdPhysicsCollisionCfg()), + ) + cube_object = RigidObject(cfg=cfg) + else: + cube_object = RigidObject(cfg=_cube_cfg("/World/Env_[^/]*/Object")) + # The articulation root above the rigid body turns the body into an articulation link. + for index in range(_NUM_ENVS): + UsdPhysics.ArticulationRootAPI.Apply(sim.stage.GetPrimAtPath(f"/World/Env_{index}")) # Check that the framework doesn't hold excessive strong references. assert sys.getrefcount(cube_object) < 10 - # Play sim with pytest.raises(RuntimeError): sim.reset() -@pytest.mark.parametrize("num_cubes", [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: - 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() - - 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()) -@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) - ) - - # 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 - - 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 - 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 - 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) - - 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) - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) @pytest.mark.parametrize("gravity_enabled", [True, False]) @pytest.mark.isaacsim_ci -def test_gravity_vec_w(num_cubes, device, gravity_enabled): +def test_gravity_vec_w(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: sim._app_control_on_stop_handle = None - # Create a scene with random cubes - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, device=device) + _spawn_envs("/World") + cube_object = RigidObject(cfg=_cube_cfg("/World/Env_[^/]*/Object")) # Obtain gravity direction if gravity_enabled: @@ -597,180 +145,221 @@ def test_gravity_vec_w(num_cubes, device, gravity_enabled): cube_object.update(sim.cfg.dt) # Expected gravity value is the acceleration of the body - gravity = torch.zeros(num_cubes, 1, 6, device=device) + gravity = torch.zeros(_NUM_ENVS, 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()) -@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) +@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 - 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)) + import omni.kit.app + import omni.physx - # check center of mass has been set - torch.testing.assert_close(wp.to_torch(cube_object.root_view.get_coms()), com) + 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 + sim_utils.create_prim("/World/Table_0", "Xform", translation=(0.0, 0.0, 1.0)) + cube = RigidObject(cfg=_cube_cfg("/World/Table_[^/]*/Object")) - # random z spin velocity - spin_twist = torch.zeros(6, device=device) - spin_twist[5] = torch.randn(1, 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() - # 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) + 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]) - # 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: +## +# Composite scene shared by the remaining tests. +## + + +@dataclass +class _RigidObjectScene: + """Rigid objects that share one real PhysX lifecycle under gravity.""" + + sim: SimulationContext + device: str + cubes: RigidObject + """Dynamic cubes that ignore gravity.""" + kinematic: RigidObject + """Kinematic cubes.""" + origins: dict[str, torch.Tensor] + """Environment origins keyed by object name.""" + refcounts: dict[str, int] + """Reference count of each object right after construction.""" + + def step(self, num_steps: int = 1) -> None: + """Write, step, and update every object.""" + for _ in range(num_steps): + for rigid_object in (self.cubes, self.kinematic): + rigid_object.write_data_to_sim() + self.sim.step() + for rigid_object in (self.cubes, self.kinematic): + rigid_object.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() + + +@pytest.fixture(scope="module", params=test_devices()) +def rigid_object_scene(request) -> _RigidObjectScene: + """Initialize the composite rigid-object scene once per device.""" + 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, 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 + origins = {"cubes": _spawn_envs("/World/Cubes"), "kinematic": _spawn_envs("/World/Kinematic", 3.0)} + cubes = RigidObject(cfg=_cube_cfg("/World/Cubes/Env_[^/]*/Object", disable_gravity=True)) + kinematic = RigidObject(cfg=_cube_cfg("/World/Kinematic/Env_[^/]*/Object", kinematic=True)) + refcounts = {"cubes": sys.getrefcount(cubes), "kinematic": sys.getrefcount(kinematic)} sim.reset() + yield _RigidObjectScene( + sim=sim, + device=device, + cubes=cubes, + kinematic=kinematic, + origins={name: value.to(device) for name, value in origins.items()}, + refcounts=refcounts, + ) - # 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)) - - # 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 +def test_rigid_object_initialization(rigid_object_scene: _RigidObjectScene): + """Initialize local cubes with the expected buffers; kinematic cubes hold their default pose under gravity.""" + scene = rigid_object_scene + for name, rigid_object in (("cubes", scene.cubes), ("kinematic", scene.kinematic)): + # Check that the framework doesn't hold excessive strong references. + assert scene.refcounts[name] < 10 + assert rigid_object.is_initialized + 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) + + kinematic = scene.kinematic + for _ in range(2): + scene.step() + 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) + + +def test_rigid_object_inertial_properties(rigid_object_scene: _RigidObjectScene): + """Mass, center-of-mass, and inertia writes reach only the selected PhysX entries and survive a step.""" + scene = rigid_object_scene + 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, raw in ( + ("body_mass", cube_object.root_view.get_masses()), + ("body_com_pose_b", cube_object.root_view.get_coms().view(wp.float32)), + ("body_inertia", cube_object.root_view.get_inertias()), + ): + torch.testing.assert_close(getattr(cube_object.data, name).torch, expected[name]) + torch.testing.assert_close(wp.to_torch(raw).to(device).reshape(expected[name].shape), expected[name]) + + +def test_rigid_object_root_state_writes(rigid_object_scene: _RigidObjectScene): + """Root writes round-trip through the written frame and keep the other frame consistent with the offset.""" + scene = rigid_object_scene + device = scene.device + cube_object = scene.cubes + rest_pose = torch.cat((scene.origins["cubes"], torch.tensor([[0.0, 0.0, 0.0, 1.0]] * _NUM_ENVS, device=device)), -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() + 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))) + + # 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) @@ -790,38 +379,107 @@ def test_write_root_state(num_cubes, device, with_offset, state_location): cube_object.data.root_com_vel_w.torch[:, 3:], cube_object.data.root_link_vel_w.torch[:, 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 - - 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() - - 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." + # 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]) ) - 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]) + # 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:]) + + +def test_rigid_object_wrench_delivery_and_reset(rigid_object_scene: _RigidObjectScene): + """External wrenches act in the frame they are given in on the selected root; reset clears them.""" + scene = rigid_object_scene + 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 environment 0 so that their responses are comparable. + for setter, kwarg, name in ( + (cube_object.set_masses_index, "masses", "body_mass"), + (cube_object.set_coms_index, "coms", "body_com_pose_b"), + (cube_object.set_inertias_index, "inertias", "body_inertia"), + ): + value = getattr(cube_object.data, name).torch[:1] + setter(**{kwarg: value.expand(_NUM_ENVS, *value.shape[1:]).contiguous()}) + + # 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]]]}, "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]]]}, + "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 = getattr(cube_object.data, 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]] * 2, device=device))[:, 1] + > 0.99 + ) + + # 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 diff --git a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py index ef289e6673d2..86cfb0972106 100644 --- a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py @@ -3,31 +3,35 @@ # # SPDX-License-Identifier: BSD-3-Clause +"""Real PhysX rigid-object-collection coverage. + +Most checks run against one module-scoped scene of locally authored cubes. PhysX stores collection bodies in +body-major view order, so 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 dataclasses import dataclass import pytest import torch import warp as wp from isaaclab_physx.assets import RigidObjectCollection +from isaaclab_physx.sim.schemas import PhysxRigidBodyCfg 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, -) +from isaaclab.sim import SimulationContext, build_simulation_context +from isaaclab.utils.math import combine_frame_transforms, quat_apply, quat_apply_inverse, quat_mul def generate_cubes_scene( @@ -36,9 +40,11 @@ def generate_cubes_scene( height=1.0, has_api: bool = True, kinematic_enabled: bool = False, - device: str = "cuda:0", + disable_gravity: bool = False, + root: str = "/World", + y_offset: float = 0.0, ) -> tuple[RigidObjectCollection, torch.Tensor]: - """Generate a scene with the provided number of cubes. + """Generate a collection of local 1 kg cubes, 3 m apart along y within each environment. Args: num_envs: Number of envs to generate. @@ -46,92 +52,84 @@ def generate_cubes_scene( 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. + disable_gravity: Whether the cubes ignore gravity. + root: Prim path under which the environments are created. + y_offset: Offset of the environments along y [m]. Returns: - A tuple containing the rigid object representing the cubes and the origins of the cubes. - + A tuple containing the rigid object collection and the origins of the environments. """ - origins = torch.tensor([(i * 3.0, 0, height) for i in range(num_envs)]).to(device) + origins = torch.tensor([(i * 3.0, y_offset, height) for i in range(num_envs)]) # 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) + for i, origin in enumerate(origins.tolist()): + sim_utils.create_prim(f"{root}/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), + size=(0.2, 0.2, 0.2), + rigid_props=[ + sim_utils.UsdPhysicsRigidBodyCfg(kinematic_enabled=kinematic_enabled), + PhysxRigidBodyCfg(disable_gravity=disable_gravity), + ], + mass_props=sim_utils.MassCfg(mass=1.0), collision_props=sim_utils.UsdPhysicsCollisionCfg(), ) + 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}", + cube_config_dict[f"cube_{i}"] = RigidObjectCfg( + prim_path=f"{root}/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 RigidObjectCollection(cfg=RigidObjectCollectionCfg(rigid_objects=cube_config_dict)), origins + - return cube_object_colection, origins +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.fixture def sim(request): - """Create simulation context with the specified device.""" + """Create a function-scoped simulation context for tests that own their scene.""" device = request.getfixturevalue("device") - if "gravity_enabled" in request.fixturenames: - gravity_enabled = request.getfixturevalue("gravity_enabled") - else: - gravity_enabled = True # default to gravity enabled + gravity_enabled = request.getfixturevalue("gravity_enabled") if "gravity_enabled" in request.fixturenames else True 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) +## +# Tests that own their simulation context. Keep them above the composite scene. +## + + +@pytest.mark.parametrize("num_cubes", [2]) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +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) # 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) + with pytest.raises(RuntimeError): + sim.reset() @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) + object_collection, _ = generate_cubes_scene(num_envs=2, num_cubes=3) # Play sim sim.reset() @@ -156,7 +154,7 @@ def test_subset_write_reaches_selected_view_entry(sim, device, gravity_enabled): def test_inertial_property_subset_writes_reach_selected_view_entries(sim, device): """COM and inertia writes to a subset of (env, body) cells must reach only those PhysX view entries.""" num_envs, num_cubes = 2, 3 - object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) + object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes) sim.reset() def read_view(values: wp.array, data_dim: int) -> torch.Tensor: @@ -188,458 +186,319 @@ def read_view(values: wp.array, data_dim: int) -> torch.Tensor: torch.testing.assert_close(object_collection.data.body_inertia.torch.cpu(), expected_inertias) -@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 - ) +@pytest.mark.parametrize("num_envs", [3]) +@pytest.mark.parametrize("num_cubes", [2]) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +@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) - # 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) - # 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) + # 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] - # Simulate physics + # Perform simulation 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) + # 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 -@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() - + # Check the body accelerations are correct + torch.testing.assert_close(object_collection.data.body_com_acc_w.torch, gravity) -@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, +## +# Composite scene shared by the remaining tests. +## + +_NUM_ENVS = 2 +_NUM_CUBES = 3 + + +@dataclass +class _CollectionScene: + """Rigid-object collections that share one real PhysX lifecycle under gravity.""" + + sim: SimulationContext + device: str + cubes: RigidObjectCollection + """Two environments of three dynamic cubes that ignore gravity.""" + kinematic: RigidObjectCollection + """One environment of one kinematic cube.""" + origins: dict[str, torch.Tensor] + """Environment origins keyed by collection name.""" + refcounts: dict[str, int] + """Reference count of each collection right after construction.""" + + def step(self, num_steps: int = 1) -> None: + """Write, step, and update every collection.""" + for _ in range(num_steps): + for collection in (self.cubes, self.kinematic): + collection.write_data_to_sim() + self.sim.step() + for collection in (self.cubes, self.kinematic): + collection.update(self.sim.cfg.dt) + + def place_cubes_at_rest(self, body_poses: torch.Tensor) -> None: + """Teleport the dynamic cubes to ``body_poses`` at rest and clear their external wrenches.""" + self.cubes.write_body_link_pose_to_sim_index(body_poses=body_poses) + self.cubes.write_body_com_velocity_to_sim_index( + body_velocities=torch.zeros((_NUM_ENVS, _NUM_CUBES, 6), device=self.device) ) - 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) + self.cubes.permanent_wrench_composer.reset() + self.cubes.instantaneous_wrench_composer.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_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, +@pytest.fixture(scope="module", params=test_devices()) +def collection_scene(request) -> _CollectionScene: + """Initialize the composite collection scene once per device.""" + device = request.param + with build_simulation_context(device=device, gravity_enabled=True) as sim: + sim._app_control_on_stop_handle = None + cubes, cube_origins = generate_cubes_scene( + num_envs=_NUM_ENVS, num_cubes=_NUM_CUBES, disable_gravity=True, root="/World/Cubes" ) - 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, + kinematic, kinematic_origins = generate_cubes_scene( + num_envs=1, num_cubes=1, kinematic_enabled=True, root="/World/Kinematic", y_offset=12.0 + ) + refcounts = {"cubes": sys.getrefcount(cubes), "kinematic": sys.getrefcount(kinematic)} + sim.reset() + yield _CollectionScene( + sim=sim, + device=device, + cubes=cubes, + kinematic=kinematic, + origins={"cubes": cube_origins.to(device), "kinematic": kinematic_origins.to(device)}, + refcounts=refcounts, ) - 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) +def _rest_poses(scene: _CollectionScene, yaw: float = 0.0) -> torch.Tensor: + """Return the default body poses of the dynamic cubes at their environment origins with a given yaw.""" + poses = scene.cubes.data.default_body_pose.torch.clone() + poses[..., :2] += scene.origins["cubes"].unsqueeze(1)[..., :2] + poses[..., 3:] = torch.tensor(_yaw_quat(yaw), device=scene.device) + return poses + + +def test_collection_initialization(collection_scene: _CollectionScene): + """Initialize local collections, including a single-cube one; kinematic cubes hold their pose under gravity.""" + scene = collection_scene + for name, collection, num_envs, num_cubes in ( + ("cubes", scene.cubes, _NUM_ENVS, _NUM_CUBES), + ("kinematic", scene.kinematic, 1, 1), + ): + # Check that the framework doesn't hold excessive strong references. + assert scene.refcounts[name] < 10 + assert collection.is_initialized + 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) + + kinematic = scene.kinematic + for _ in range(2): + scene.step() + default_body_pose = kinematic.data.default_body_pose.torch.clone() + default_body_vel = kinematic.data.default_body_vel.torch.clone() + default_body_pose[..., :3] += scene.origins["kinematic"].unsqueeze(1) + torch.testing.assert_close(kinematic.data.body_link_pose_w.torch, default_body_pose) + torch.testing.assert_close(kinematic.data.body_link_vel_w.torch, default_body_vel) + + +def test_collection_state_writes(collection_scene: _CollectionScene): + """Body state writes reach the selected body-major view entries in the frame they are given in.""" + scene = collection_scene + device = scene.device + collection = scene.cubes + scene.place_cubes_at_rest(_rest_poses(scene)) + + # 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_cubes_at_rest(_rest_poses(scene)) + 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) 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], + 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) + + # 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), ) - - 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) + 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( - torch.zeros_like(object_com_vel_w[..., :3]), - object_com_vel_w[..., :3], + collection.data.body_com_vel_w.torch[..., 3:], collection.data.body_link_vel_w.torch[..., 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), + # Spinning about the center of mass keeps it in place while the link frame orbits it. + scene.place_cubes_at_rest(_rest_poses(scene)) + 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] ) - 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)) + 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_masses_reach_selected_view_entries(collection_scene: _CollectionScene): + """Mass writes to a non-sorted subset reach only the selected body-major view entries.""" + scene = collection_scene + device = scene.device + collection = scene.cubes + env_ids = torch.tensor([1, 0], dtype=torch.int32, device=device) + body_ids = torch.tensor([1], dtype=torch.int32, device=device) + expected = collection.data.body_mass.torch.clone() + masses = torch.tensor([[5.0], [7.0]], device=device) + expected[env_ids[:, None], body_ids] = masses + collection.set_masses_index(masses=masses, env_ids=env_ids, body_ids=body_ids) + torch.testing.assert_close(collection.data.body_mass.torch, expected) + raw_mass = wp.to_torch(collection.root_view.get_masses()).to(device).reshape(_NUM_CUBES, _NUM_ENVS).T + torch.testing.assert_close(raw_mass, expected) + # Restore uniform masses for the wrench checks. + collection.set_masses_index(masses=torch.ones((_NUM_ENVS, _NUM_CUBES), device=device)) + + +def test_collection_wrench_delivery_and_reset(collection_scene: _CollectionScene): + """External wrenches act on the selected bodies in the frame they are given in; reset clears them.""" + scene = collection_scene + device = scene.device + collection = scene.cubes + rest_poses = _rest_poses(scene, yaw=0.5 * math.pi) + # Give every body the same inertial properties so that their responses are comparable. + for setter, kwarg, name in ( + (collection.set_coms_index, "coms", "body_com_pose_b"), + (collection.set_inertias_index, "inertias", "body_inertia"), + ): + value = getattr(collection.data, name).torch[:1, :1] + setter(**{kwarg: value.expand(_NUM_ENVS, _NUM_CUBES, *value.shape[2:]).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 in ( + ({"forces": [[6.0, 0.0, 0.0]]}, {"forces": [[0.0, 6.0, 0.0]]}, "body_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]]}, + "body_com_ang_vel_b", + ), + ): + scene.place_cubes_at_rest(rest_poses) + 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 = getattr(collection.data, 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 response == "body_com_lin_vel_w": + assert torch.all(response_value[:, 0::2, 1] > 1e-2), response_value + # 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): @@ -647,40 +506,7 @@ def test_reset_object_collection(sim, num_envs, num_cubes, device): assert torch.count_nonzero(buffer[1:]) == buffer[1:].numel() # A full reset clears every environment - object_collection.reset() + 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) From 994646d771096936292d933d29324b594f5b6652 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 11:06:15 +0200 Subject: [PATCH 11/33] Focus PhysX Newton-actuator tests on backend seams Drop the ideal-PD, DC-motor, and remotized-PD equivalence runs, whose actuator translation is backend independent: the mixed implicit/explicit run still exercises the PD and DC-motor authoring on PhysX, and the Newton suite owns the remaining models. Move the native-path preparation checks to a kitless unit file, and replace the graph-capture failure and outer-capture rollouts with unit tests of the shared PhysX actuator runtime. --- .../actuators/test_physx_actuator_runtime.py | 101 ++++++++++ .../assets/test_newton_actuators_physx.py | 182 +----------------- .../unit/test_physx_actuator_control.py | 73 +++++++ 3 files changed, 182 insertions(+), 174 deletions(-) create mode 100644 source/isaaclab/test/actuators/test_physx_actuator_runtime.py create mode 100644 source/isaaclab_physx/test/assets/unit/test_physx_actuator_control.py diff --git a/source/isaaclab/test/actuators/test_physx_actuator_runtime.py b/source/isaaclab/test/actuators/test_physx_actuator_runtime.py new file mode 100644 index 000000000000..6f9a144cbcde --- /dev/null +++ b/source/isaaclab/test/actuators/test_physx_actuator_runtime.py @@ -0,0 +1,101 @@ +# 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 CUDA-graph handling of the host-side PhysX actuator runtime.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import warp as wp + +from isaaclab.actuators.newton.physx_runtime import PhysxActuatorRuntime + + +def _runtime() -> PhysxActuatorRuntime: + """Create a runtime for a CUDA articulation with a mocked logger.""" + return PhysxActuatorRuntime(SimpleNamespace(device="cuda:0"), logger=Mock()) + + +def _cuda_device(*, is_capturing: bool) -> SimpleNamespace: + """Describe a CUDA device with the given capture state.""" + return SimpleNamespace(is_cuda=True, is_capturing=is_capturing) + + +def test_graph_capture_failure_restores_adapter_state_and_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch): + """A failed second capture discards the graphs and restores the state swapped by the first capture.""" + + class _FailingCapture: + capture_count = 0 + + def __init__(self, *args, **kwargs): + self.capture_index = type(self).capture_count + type(self).capture_count += 1 + self.graph = object() + + def __enter__(self): + if self.capture_index == 1: + raise RuntimeError("second capture unavailable") + return self + + def __exit__(self, exc_type, exc_value, traceback): + return False + + 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", _FailingCapture) + 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): + """A graphable adapter whose capture fails still computes the current command eagerly.""" + + 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 + + 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: _cuda_device(is_capturing=False)) + monkeypatch.setattr(wp, "ScopedCapture", _FailingCapture) + 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 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): + """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: _cuda_device(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_physx/test/assets/test_newton_actuators_physx.py b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py index 266cd318d3a6..0e4c0e482602 100644 --- a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py +++ b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py @@ -3,15 +3,13 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""PD actuator equivalence tests on ANYmal-C (floating-base quadruped) — PhysX backend. +"""Newton-native actuator seams on the PhysX backend, driven on ANYmal-C (floating-base quadruped). -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. +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. 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, per-environment resets, gain randomization, and network actuators. """ from isaaclab.test.utils import launch_test_simulation @@ -22,22 +20,18 @@ import functools import os import unittest -from types import SimpleNamespace -import pytest import torch import warp as wp from isaaclab_physx.assets import Articulation -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 import IdealPDActuatorCfg from isaaclab.actuators.newton import read_group_parameter from isaaclab.sim import SimulationCfg, 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, @@ -52,7 +46,6 @@ 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 @@ -78,70 +71,6 @@ ) -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 - - 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) - ) - gain_writes = [] - articulation = SimpleNamespace( - _sim_cfg=SimpleNamespace(use_newton_actuators=True), - cfg=SimpleNamespace(prim_path="/World/Robot"), - joint_names=["joint"], - num_instances=1, - num_joints=1, - device="cpu", - _data=SimpleNamespace(joint_pos=joint_buffer, joint_vel=joint_buffer), - write_joint_stiffness_to_sim_index=lambda **_: gain_writes.append("stiffness"), - write_joint_damping_to_sim_index=lambda **_: gain_writes.append("damping"), - ) - wrapper = SimpleNamespace() - adapter = SimpleNamespace(joint_indices=wp.array([0], dtype=wp.int32), finalize=lambda _: None) - monkeypatch.setattr(actuator_control, "find_first_matching_prim", lambda _: None) - monkeypatch.setattr(PhysxActuatorWrapper, "create", lambda **_: wrapper) - monkeypatch.setattr(NewtonActuatorAdapter, "from_usd", lambda **_: adapter) - - native_groups = PhysxActuatorControl(articulation).prepare_native_actuators( - collection, - {"explicit": IdealPDActuatorCfg(joint_names_expr=["joint"], stiffness=None, damping=None)}, - ) - - assert native_groups == {"explicit"} - assert gain_writes == [] - - -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 - ) - 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 = PhysxActuatorControl(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 == [] - - # --------------------------------------------------------------------------- # Simulation runner # --------------------------------------------------------------------------- @@ -152,10 +81,8 @@ def _run_simulation( 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. @@ -165,11 +92,9 @@ def _run_simulation( 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. @@ -219,7 +144,7 @@ def _run_simulation( 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) + effort_target = None articulation.set_joint_position_target_index(target=target_pos) articulation.set_joint_velocity_target_index(target=target_vel) @@ -229,9 +154,6 @@ def _run_simulation( 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() @@ -261,43 +183,6 @@ def _run_simulation( } -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( @@ -438,20 +323,6 @@ def setUpClass(cls): # --------------------------------------------------------------------------- -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). @@ -723,43 +594,6 @@ 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 # --------------------------------------------------------------------------- diff --git a/source/isaaclab_physx/test/assets/unit/test_physx_actuator_control.py b/source/isaaclab_physx/test/assets/unit/test_physx_actuator_control.py new file mode 100644 index 000000000000..e83e7577af38 --- /dev/null +++ b/source/isaaclab_physx/test/assets/unit/test_physx_actuator_control.py @@ -0,0 +1,73 @@ +# 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 PhysX Newton-actuator path preparation.""" + +from types import SimpleNamespace + +import warp as wp +from isaaclab_physx.assets.articulation import actuator_control +from isaaclab_physx.assets.articulation.actuator_control import PhysxActuatorControl + +from isaaclab.actuators import IdealPDActuatorCfg, ImplicitActuatorCfg +from isaaclab.actuators.newton import NewtonActuatorAdapter, PhysxActuatorWrapper + + +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) + ) + gain_writes = [] + articulation = SimpleNamespace( + _sim_cfg=SimpleNamespace(use_newton_actuators=True), + cfg=SimpleNamespace(prim_path="/World/Robot"), + joint_names=["joint"], + num_instances=1, + num_joints=1, + device="cpu", + _data=SimpleNamespace(joint_pos=joint_buffer, joint_vel=joint_buffer), + write_joint_stiffness_to_sim_index=lambda **_: gain_writes.append("stiffness"), + write_joint_damping_to_sim_index=lambda **_: gain_writes.append("damping"), + ) + wrapper = SimpleNamespace() + adapter = SimpleNamespace(joint_indices=wp.array([0], dtype=wp.int32), finalize=lambda _: None) + monkeypatch.setattr(actuator_control, "find_first_matching_prim", lambda _: None) + monkeypatch.setattr(PhysxActuatorWrapper, "create", lambda **_: wrapper) + monkeypatch.setattr(NewtonActuatorAdapter, "from_usd", lambda **_: adapter) + + native_groups = PhysxActuatorControl(articulation).prepare_native_actuators( + collection, + {"explicit": IdealPDActuatorCfg(joint_names_expr=["joint"], stiffness=None, damping=None)}, + ) + + assert native_groups == {"explicit"} + assert gain_writes == [] + + +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 = PhysxActuatorControl(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 == [] From eda3caba8dd5912b0485baa03204956f86c958bb Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 11:06:15 +0200 Subject: [PATCH 12/33] Cover PhysX deformables and surface grippers with local fixtures Replace the disabled deformable suite with one scene of locally authored surface and volume meshes that need no mesher or collision cooking, and check partial nodal writes, kinematic targets, and the surface target rejection. Author the surface gripper locally, keep the CI skip for the real open/close seam, and check the command and property payloads with a recording view. --- .../test/assets/test_deformable_object.py | 499 ++++++++---------- .../test/assets/test_surface_gripper.py | 309 ++++------- .../assets/unit/test_physx_surface_gripper.py | 57 ++ 3 files changed, 371 insertions(+), 494 deletions(-) create mode 100644 source/isaaclab_physx/test/assets/unit/test_physx_surface_gripper.py diff --git a/source/isaaclab_physx/test/assets/test_deformable_object.py b/source/isaaclab_physx/test/assets/test_deformable_object.py index bde56f791c8f..1f2404ccb2bd 100644 --- a/source/isaaclab_physx/test/assets/test_deformable_object.py +++ b/source/isaaclab_physx/test/assets/test_deformable_object.py @@ -3,337 +3,254 @@ # # 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 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, 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.") - - -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. - - 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. - - """ - 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), - ) - 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, - ) - # 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) - +_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)] -@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 +def _add_api_schemas(prim, schemas: list[str]) -> None: + """Set the applied API schemas of a prim explicitly.""" + schemas_op = Sdf.TokenListOp() + schemas_op.explicitItems = schemas + prim.SetMetadata("apiSchemas", schemas_op) - # 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) +def _author_deformable_body(body_prim, points: list[Gf.Vec3f], sim_api: str, material_cfg) -> None: + """Author the deformable body schemas, rest shape, and a bound physics material on a mesh prim.""" + _add_api_schemas( + body_prim, + ["OmniPhysicsDeformableBodyAPI", sim_api, "OmniPhysicsDeformablePoseAPI:default", "PhysicsCollisionAPI"], + ) + body_prim.CreateAttribute("deformablePose:default:omniphysics:points", Sdf.ValueTypeNames.Point3fArray).Set(points) + body_prim.CreateAttribute("deformablePose:default:omniphysics:purposes", Sdf.ValueTypeNames.TokenArray).Set( + ["bindPose"] + ) + 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", + ) - # 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(): """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): 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 - - 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) - - 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) - - 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) +@dataclass +class _DeformableScene: + """Volume and surface deformables that share one real PhysX lifecycle.""" - torch.testing.assert_close(cube_object.data.root_pos_w.torch, mean_nodal_pos_init, rtol=1e-4, atol=1e-4) + sim: sim_utils.SimulationContext + device: str + volume: DeformableObject + surface: DeformableObject + refcounts: dict[str, int] + 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) -@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) - - default_root_pos = cube_object.data.default_nodal_state_w.torch.mean(dim=1) - - cube_object.reset() +@pytest.fixture(scope="module", params=test_devices(DeviceScope.CUDA)) +def deformable_scene(request) -> _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) - 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]) +def test_deformable_initialization(deformable_scene: _DeformableScene): + """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) + + +def test_deformable_nodal_state_writes(deformable_scene: _DeformableScene): + """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) + torch.testing.assert_close(volume.data.root_pos_w.torch, nodal_pos.mean(dim=1)) + + +def test_volume_kinematic_targets_hold_selected_nodes(deformable_scene: _DeformableScene): + """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_surface_gripper.py b/source/isaaclab_physx/test/assets/test_surface_gripper.py index c2d02b226f56..ebfa1bfec7ef 100644 --- a/source/isaaclab_physx/test/assets/test_surface_gripper.py +++ b/source/isaaclab_physx/test/assets/test_surface_gripper.py @@ -3,13 +3,11 @@ # # 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() @@ -18,228 +16,133 @@ 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, - ), - }, - ) - - 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)), +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, + ) ) - 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 - ) - - 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]) -@pytest.mark.isaacsim_ci -def test_raise_error_if_not_cpu(sim, device, add_ground_plane) -> 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="cpu", gravity_enabled=False) as sim: + sim._app_control_on_stop_handle = None + surface_gripper = _author_surface_gripper() - with pytest.raises(Exception, match="only supported on CPU"): 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() + # step the simulation so the gripper reacts to the command + 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(device) -> None: + """Test that the SurfaceGripper raises an error if the device is not CPU.""" + 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 -if __name__ == "__main__": - pytest.main([__file__, "-v", "--maxfail=1"]) + with pytest.raises(Exception, match="only supported on CPU"): + sim.reset() diff --git a/source/isaaclab_physx/test/assets/unit/test_physx_surface_gripper.py b/source/isaaclab_physx/test/assets/unit/test_physx_surface_gripper.py new file mode 100644 index 000000000000..3007be316949 --- /dev/null +++ b/source/isaaclab_physx/test/assets/unit/test_physx_surface_gripper.py @@ -0,0 +1,57 @@ +# 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 PhysX surface-gripper command and property payloads.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import warp as wp +from isaaclab_physx.assets.surface_gripper.surface_gripper import SurfaceGripper + + +def _gripper() -> SurfaceGripper: + """Create 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(), + ) + return gripper + + +def test_command_filter_and_partial_property_update_use_literal_view_payloads() -> None: + """Submit only open/close commands and preserve property selector ordering.""" + gripper = _gripper() + 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, + ) + + 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], + } From 59d7a640b592ce78d16f596fbdf8c2f1b0a0debb Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 11:16:50 +0200 Subject: [PATCH 13/33] Author local fixtures for kitless Newton asset tests --- .../test/articulation_test_utils.py | 539 ++++++++++++++++++ 1 file changed, 539 insertions(+) create mode 100644 source/isaaclab_newton/test/articulation_test_utils.py diff --git a/source/isaaclab_newton/test/articulation_test_utils.py b/source/isaaclab_newton/test/articulation_test_utils.py new file mode 100644 index 000000000000..173740041c70 --- /dev/null +++ b/source/isaaclab_newton/test/articulation_test_utils.py @@ -0,0 +1,539 @@ +# 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 and controller tests. + +Test modules below ``source/isaaclab_newton/test`` import this module after putting this directory on +``sys.path``. The fixtures are small USD documents authored here and written to a per-process temporary +directory, so the scenes need neither Kit nor Nucleus and the repository carries no generated USD files. +""" + +from __future__ import annotations + +import atexit +import functools +import shutil +import tempfile +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager +from pathlib import Path +from typing import Any + +import numpy as np +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 + +DATA_DIR = Path(__file__).resolve().parent / "assets" / "data" +"""Directory holding the checked-in USD fixtures.""" + +NUM_ENVS = 2 +"""Environments per scene: enough to prove a partial write leaves the other environment untouched.""" + +ENV_SPACING = 40.0 +"""Distance between environment origins along x [m].""" + +Vec3 = tuple[float, float, float] + +## +# Fixture authoring. +## + + +def _body( + name: str, + *, + mass: float, + inertia: Vec3, + com: Vec3 | None = None, + translate: Vec3 | None = None, + prim_type: str = "Xform", + schemas: Sequence[str] = (), + body: str = "", +) -> str: + """Return a rigid body with explicit mass properties.""" + applied = ", ".join(f'"{schema}"' for schema in ("PhysicsRigidBodyAPI", "PhysicsMassAPI", *schemas)) + lines = [] + if translate is not None: + lines += [f"double3 xformOp:translate = {translate}", 'uniform token[] xformOpOrder = ["xformOp:translate"]'] + lines += [f"float physics:mass = {mass}", f"float3 physics:diagonalInertia = {inertia}"] + if com is not None: + lines.append(f"point3f physics:centerOfMass = {com}") + content = "\n".join(f" {line}" for line in lines) + body + return ( + f' def {prim_type} "{name}" (\n prepend apiSchemas = [{applied}]\n )\n {{\n{content}\n }}\n' + ) + + +def _collision_cube(size: float) -> str: + """Return a collider cube authored as a child of a rigid body.""" + return ( + f'\n\n def Cube "Collision" (\n prepend apiSchemas = ["PhysicsCollisionAPI"]\n )\n' + f" {{\n double size = {size}\n }}" + ) + + +def _joint( + root: str, + kind: str, + name: str, + body0: str | None, + body1: str, + *, + axis: str | None = None, + local_pos0: Vec3 | None = None, + local_pos1: Vec3 | None = None, + limits: tuple[float, float] | None = None, + drive: tuple[str, float, float] | None = None, +) -> str: + """Return a joint. + + ``body0=None`` attaches ``body1`` to the world, and ``drive`` is ``(kind, stiffness, damping)``. + """ + schemas = f' (\n prepend apiSchemas = ["PhysicsDriveAPI:{drive[0]}"]\n )' if drive else "" + lines = [] + if axis is not None: + lines.append(f'uniform token physics:axis = "{axis}"') + if body0 is not None: + lines.append(f"rel physics:body0 = ") + lines.append(f"rel physics:body1 = ") + if local_pos0 is not None: + lines.append(f"point3f physics:localPos0 = {local_pos0}") + if local_pos1 is not None: + lines.append(f"point3f physics:localPos1 = {local_pos1}") + if limits is not None: + lines += [f"float physics:lowerLimit = {limits[0]}", f"float physics:upperLimit = {limits[1]}"] + if drive is not None: + lines += [ + f"float drive:{drive[0]}:physics:stiffness = {drive[1]}", + f"float drive:{drive[0]}:physics:damping = {drive[2]}", + ] + content = "\n".join(f" {line}" for line in lines) + return f' def Physics{kind}Joint "{name}"{schemas}\n {{\n{content}\n }}\n' + + +def _document(root: str, prims: Sequence[str], *, root_schemas: Sequence[str] = ("PhysicsArticulationRootAPI",)) -> str: + """Return a USD document whose default prim is an ``Xform`` holding ``prims``.""" + header = f'#usda 1.0\n(\n defaultPrim = "{root}"\n metersPerUnit = 1\n upAxis = "Z"\n)\n\n' + applied = ", ".join(f'"{schema}"' for schema in root_schemas) + metadata = f" (\n prepend apiSchemas = [{applied}]\n)" if root_schemas else "" + return header + f'def Xform "{root}"{metadata}\n{{\n' + "\n".join(prims) + "}\n" + + +def _floating_two_link() -> str: + """Two colliding links joined about z; the articulation root sits on the ``Root`` body.""" + return _document( + "Robot", + [ + _body( + "Root", + mass=1.0, + inertia=(0.01, 0.02, 0.02), + com=(0.0, 0.0, 0.0), + schemas=("PhysicsArticulationRootAPI",), + body=_collision_cube(0.1), + ), + _body( + "Child", + mass=1.0, + inertia=(0.01, 0.02, 0.02), + com=(0.0, 0.0, 0.0), + translate=(0.5, 0.0, 0.0), + body=_collision_cube(0.1), + ), + _joint( + "Robot", + "Revolute", + "Joint", + "Root", + "Child", + axis="Z", + local_pos0=(0.25, 0.0, 0.0), + local_pos1=(-0.25, 0.0, 0.0), + limits=(-90.0, 90.0), + ), + ], + root_schemas=(), + ) + + +def _fixed_spatial_chain() -> str: + """A fixed-base chain of three slides and a three-axis wrist whose links carry center-of-mass offsets. + + Every joint has a USD drive; the wrist's angular gains are per degree. + """ + links = [ + _body("Root", mass=1.0, inertia=(0.01, 0.01, 0.01)), + _joint("Robot", "Fixed", "RootJoint", None, "Root"), + ] + link_specs = [ + (0.5, (0.02, 0.02, 0.02), (0.02, 0.0, 0.0), None), + (0.5, (0.02, 0.02, 0.02), (0.0, 0.02, 0.0), None), + (0.5, (0.02, 0.02, 0.02), (0.0, 0.0, 0.02), None), + (0.3, (0.01, 0.01, 0.01), (0.03, 0.0, 0.02), (0.0, 0.0, 0.1)), + (0.3, (0.01, 0.01, 0.01), (0.0, 0.03, 0.02), (0.0, 0.0, 0.2)), + (0.2, (0.01, 0.01, 0.01), (0.05, 0.02, 0.0), (0.0, 0.0, 0.3)), + ] + for index, (mass, inertia, com, translate) in enumerate(link_specs): + links.append(_body(f"Link_{index}", mass=mass, inertia=inertia, com=com, translate=translate)) + parents = ["Root", "Link_0", "Link_1", "Link_2", "Link_3", "Link_4"] + for index, axis in enumerate("XYZ"): + links.append( + _joint( + "Robot", + "Prismatic", + f"Joint_{index}", + parents[index], + f"Link_{index}", + axis=axis, + limits=(-0.2, 0.2), + drive=("linear", 200.0, 20.0), + ) + ) + for index, (axis, stiffness) in enumerate(zip("XYZ", WRIST_USD_STIFFNESS), start=3): + links.append( + _joint( + "Robot", + "Revolute", + f"Joint_{index}", + parents[index], + f"Link_{index}", + axis=axis, + local_pos0=(0.0, 0.0, 0.1), + local_pos1=(0.0, 0.0, 0.0), + limits=(-60.0, 60.0), + drive=("angular", stiffness, round(0.1 * stiffness, 6)), + ) + ) + return _document("Robot", links) + + +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.""" + + +def _revolute_pendulum() -> str: + """A fixed pivot and a unit arm turning about z.""" + return _document( + "Robot", + [ + _body("CenterPivot", mass=1.0, inertia=(0.01, 0.01, 0.01)), + _joint("Robot", "Fixed", "FixedJoint", None, "CenterPivot"), + _body("Arm", mass=1.0, inertia=(0.01, 0.01, 0.01), translate=(1.0, 0.0, 0.0)), + _joint( + "Robot", + "Revolute", + "RevoluteJoint", + "CenterPivot", + "Arm", + axis="Z", + local_pos0=(0.0, 0.0, 0.0), + local_pos1=(-1.0, 0.0, 0.0), + ), + ], + ) + + +def _fixed_tendon_hand() -> str: + """A fixed palm with two two-joint fingers; each fixed tendon ``XJ0`` sums ``XJ1`` and ``XJ2``. + + A MuJoCo position actuator drives each tendon. + """ + prims = [ + _body("Palm", mass=0.5, inertia=(0.001, 0.001, 0.001)), + _joint("Hand", "Fixed", "PalmJoint", None, "Palm"), + ] + for finger, y in (("FF", 0.02), ("MF", -0.02)): + prims += [ + _body( + f"{finger}_proximal", + mass=0.05, + inertia=(1e-5, 1e-5, 1e-5), + com=(0.025, 0.0, 0.0), + translate=(0.05, y, 0.0), + ), + _body( + f"{finger}_distal", mass=0.03, inertia=(1e-5, 1e-5, 1e-5), com=(0.02, 0.0, 0.0), translate=(0.1, y, 0.0) + ), + _joint( + "Hand", + "Revolute", + f"{finger}J2", + "Palm", + f"{finger}_proximal", + axis="Y", + local_pos0=(0.05, y, 0.0), + local_pos1=(0.0, 0.0, 0.0), + limits=(0.0, 90.0), + ), + _joint( + "Hand", + "Revolute", + f"{finger}J1", + f"{finger}_proximal", + f"{finger}_distal", + axis="Y", + local_pos0=(0.05, 0.0, 0.0), + local_pos1=(0.0, 0.0, 0.0), + limits=(0.0, 90.0), + ), + f' def MjcTendon "{finger}J0"\n {{\n uniform token mjc:type = "fixed"\n' + f" rel mjc:path = [, ]\n" + " uniform int[] mjc:path:indices = [0, 1]\n uniform double[] mjc:path:coef = [1, 1]\n }\n", + f' def MjcActuator "{finger}_tendon_actuator"\n {{\n uniform token mjc:biasType = "affine"\n' + " uniform double[] mjc:gainPrm = [0.05, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n" + " uniform double[] mjc:biasPrm = [0, -0.05, -0.005, 0, 0, 0, 0, 0, 0, 0]\n" + f" rel mjc:target = \n }}\n", + ] + return _document("Hand", prims) + + +def _floating_two_leg() -> str: + """A floating base with two three-joint legs named like ANYmal's (``HAA``, ``HFE``, ``KFE``).""" + prims = [_body("base", mass=10.0, inertia=(0.5, 0.5, 0.5), com=(0.0, 0.0, 0.0))] + joints = [] + for leg, x, y in (("LF", 0.3, 0.1), ("RH", -0.3, -0.1)): + prims += [ + _body(f"{leg}_HIP", mass=1.0, inertia=(0.05, 0.05, 0.05), com=(0.0, 0.02, 0.0), translate=(x, y, 0.0)), + _body(f"{leg}_THIGH", mass=1.0, inertia=(0.05, 0.05, 0.05), com=(0.0, 0.0, -0.1), translate=(x, y, -0.05)), + _body(f"{leg}_SHANK", mass=0.5, inertia=(0.05, 0.05, 0.05), com=(0.0, 0.0, -0.1), translate=(x, y, -0.3)), + ] + joints += [ + _joint( + "Robot", + "Revolute", + f"{leg}_HAA", + "base", + f"{leg}_HIP", + axis="X", + local_pos0=(x, y, 0.0), + local_pos1=(0.0, 0.0, 0.0), + limits=(-45.0, 45.0), + ), + _joint( + "Robot", + "Revolute", + f"{leg}_HFE", + f"{leg}_HIP", + f"{leg}_THIGH", + axis="Y", + local_pos0=(0.0, 0.0, -0.05), + local_pos1=(0.0, 0.0, 0.0), + limits=(-90.0, 90.0), + ), + _joint( + "Robot", + "Revolute", + f"{leg}_KFE", + f"{leg}_THIGH", + f"{leg}_SHANK", + axis="Y", + local_pos0=(0.0, 0.0, -0.25), + local_pos1=(0.0, 0.0, 0.0), + limits=(-160.0, 0.0), + ), + ] + return _document("Robot", prims + joints) + + +def _fixed_cartpole() -> str: + """A cart sliding on a fixed rail with a pole on a revolute joint.""" + return _document( + "Robot", + [ + _body("rail", mass=1.0, inertia=(0.01, 0.01, 0.01)), + _joint("Robot", "Fixed", "rail_joint", None, "rail"), + _body("cart", mass=1.0, inertia=(0.01, 0.01, 0.01)), + _body("pole", mass=0.2, inertia=(0.01, 0.01, 0.01), com=(0.0, 0.0, 0.5)), + _joint("Robot", "Prismatic", "slider_to_cart", "rail", "cart", axis="X", limits=(-2.0, 2.0)), + _joint("Robot", "Revolute", "cart_to_pole", "cart", "pole", axis="Y"), + ], + ) + + +def _rigid_body_with_articulation_root() -> str: + """A rigid body carrying an articulation root, with a second body fixed to it.""" + return """#usda 1.0 +( + defaultPrim = "Object" + metersPerUnit = 1 + upAxis = "Z" +) + +def Cube "Object" ( + prepend apiSchemas = ["PhysicsRigidBodyAPI", "PhysicsMassAPI", "PhysicsCollisionAPI", "PhysicsArticulationRootAPI"] +) +{ + double size = 0.1 + float physics:mass = 1 + + def Cube "Child" ( + prepend apiSchemas = ["PhysicsRigidBodyAPI", "PhysicsMassAPI", "PhysicsCollisionAPI"] + ) + { + double size = 0.1 + double3 xformOp:translate = (0, 0, 2) + uniform token[] xformOpOrder = ["xformOp:translate"] + float physics:mass = 1 + } + + def PhysicsFixedJoint "Joint" + { + rel physics:body0 = + rel physics:body1 = + } +} +""" + + +_FIXTURES: dict[str, Callable[[], str]] = { + "fixed_cartpole.usda": _fixed_cartpole, + "fixed_spatial_chain.usda": _fixed_spatial_chain, + "fixed_tendon_hand.usda": _fixed_tendon_hand, + "floating_two_leg.usda": _floating_two_leg, + "floating_two_link.usda": _floating_two_link, + "revolute_pendulum.usda": _revolute_pendulum, + "rigid_body_with_articulation_root.usda": _rigid_body_with_articulation_root, +} +"""Authored fixtures by file name.""" + + +@functools.cache +def _fixture_dir() -> Path: + """Return this process's directory for the authored fixtures, removed when the process exits.""" + directory = Path(tempfile.mkdtemp(prefix="isaaclab_newton_test_fixtures_")) + atexit.register(shutil.rmtree, directory, ignore_errors=True) + return directory + + +def local_usd(name: str, **kwargs: Any) -> sim_utils.UsdFileCfg: + """Return a spawner for one of the local USD fixtures. + + Args: + name: Authored fixture name, or a file name inside :data:`DATA_DIR`. + **kwargs: Additional :class:`~isaaclab.sim.UsdFileCfg` fields. + + Returns: + The spawner configuration. + """ + if name in _FIXTURES: + path = _fixture_dir() / name + if not path.exists(): + path.write_text(_FIXTURES[name](), encoding="utf-8") + else: + path = DATA_DIR / name + return sim_utils.UsdFileCfg(usd_path=str(path), **kwargs) + + +## +# Scenes. +## + + +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, Any], 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 = np.zeros((num_envs, 3), dtype=np.float32) + positions[:, 0] = ENV_SPACING * np.arange(num_envs) + # 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).expand_as(previous)) + SimulationManager.add_model_change(ModelFlags.MODEL_PROPERTIES) + try: + yield + finally: + world_gravity_view.copy_(previous) + SimulationManager.add_model_change(ModelFlags.MODEL_PROPERTIES) From 10b9844868ae395e43b8b1f036f3b82801d4010e Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 11:16:50 +0200 Subject: [PATCH 14/33] Rebuild Newton articulation tests on one composite local scene --- .../test/assets/test_articulation.py | 4569 ++++++----------- .../unit/test_articulation_accessors.py | 110 + 2 files changed, 1679 insertions(+), 3000 deletions(-) create mode 100644 source/isaaclab_newton/test/assets/unit/test_articulation_accessors.py diff --git a/source/isaaclab_newton/test/assets/test_articulation.py b/source/isaaclab_newton/test/assets/test_articulation.py index f0f4b5da6dbc..fb518e020f17 100644 --- a/source/isaaclab_newton/test/assets/test_articulation.py +++ b/source/isaaclab_newton/test/assets/test_articulation.py @@ -6,379 +6,75 @@ # 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. + +Configurations that fail initialization and the rebind scenario, which swaps the live Newton state, build +their own small scenes. +""" -HEADLESS = True +from isaaclab_newton.physics import NewtonCfg -launch_test_simulation() +from isaaclab.sim import SimulationCfg +from isaaclab.test.utils import launch_test_simulation + +launch_test_simulation(SimulationCfg(physics=NewtonCfg())) import sys -from copy import copy, deepcopy +from copy import copy +from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace +from unittest.mock import patch -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.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.selection import ArticulationView -from newton.solvers import SolverMuJoCo +from newton import ModelFlags, ShapeFlags 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 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.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 - -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_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 - +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, +) -def fix_reversed_joints(stage): - """Fix reversed joints on the USD stage. +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - 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 articulation_test_utils import ( # noqa: E402 + WRIST_USD_STIFFNESS, + env_origins, + local_usd, + newton_sim_cfg, + spawn_assets, + world_gravity, ) -_ANYMAL_C_ROOT_PRESERVING_REVERSED_BODY_NAMES = (_ANYMAL_C_BODY_NAMES[0], *reversed(_ANYMAL_C_BODY_NAMES[1:])) +pytestmark = [pytest.mark.integration, pytest.mark.kitless] + +_DEVICE = "cpu" +_GRAVITY = (0.0, 0.0, -9.81) _NEWTON_USER_ORDER_STATE_CACHES = ( "_joint_pos_user", @@ -387,639 +83,356 @@ 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. - - """ - # 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()) +_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. - return articulation, torch.as_tensor(translations, device=device) +Neither permutation is its own inverse, so a map applied in the wrong direction shows. +""" -# --------------------------------------------------------------------------- -# Franka task-space tracking helpers (shared between the OSC tests). -# --------------------------------------------------------------------------- +def _limit_actuators(actuator_type: type) -> dict: + """Return one limit-resolution group per branching-fixture joint. - -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_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_shoulder/panda_forearm - 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_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(clone(FRANKA_PANDA_HIGH_PD_CFG), prim_path="/World/Env_[^/]*/Robot") - if zero_actuator_pd: - cfg.actuators["panda_shoulder"].stiffness = 0.0 - cfg.actuators["panda_shoulder"].damping = 0.0 - cfg.actuators["panda_forearm"].stiffness = 0.0 - cfg.actuators["panda_forearm"].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 - - -def _compute_ee_vel_root(jacobian_b, joint_vel): - """Return the EE 6D velocity in the root frame as ``J · q_dot``. - - 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) - + if actuator_type is ImplicitActuatorCfg: + gains = {"stiffness": 10.0, "damping": 1.0} + actuator_limits = {"actuator_velocity_limit": 1e2} + else: + 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), + } -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) +def _island_cfgs() -> dict[str, ArticulationCfg]: + """Return the composite scene's islands, each at its own offset inside every environment.""" -@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 - 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}", - ) + def island(name: str, spawn, y: float, actuators: dict, **kwargs) -> ArticulationCfg: + init_state = kwargs.pop("init_state", ArticulationCfg.InitialStateCfg()) + init_state.pos = (0.0, y, 1.0) + return ArticulationCfg( + prim_path=f"/World/Env_[^/]*/{name}", spawn=spawn, init_state=init_state, actuators=actuators, **kwargs ) - 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}) - 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], + 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], + "passive": island( + "Passive", + local_usd("fixed_spatial_chain.usda"), + 4.0, + {"all": 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], + "pendulum": island( + "Pendulum", + local_usd("revolute_pendulum.usda"), + 6.0, + {"joint": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=0.0, damping=0.0)}, ), - ( - ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=0.0, damping=0.0), - JointTargetMode.EFFORT, - None, + "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, ), - ( - IdealPDActuatorCfg(joint_names_expr=[".*"], stiffness=10.0, damping=2.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", ), - ], -) -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 - + "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 + ) + }, + ), + "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), + ), + } -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.""" -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)] + sim: SimulationContext + articulations: dict[str, Articulation] + cfgs: dict[str, ArticulationCfg] + origins: torch.Tensor + def __getattr__(self, name: str) -> Articulation: + try: + return self.articulations[name] + except KeyError as error: + raise AttributeError(name) from error -@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 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. -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)] + 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_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)] +## +# Own scenes. These tests run before the composite scene exists: only one simulation context can be alive. +## -@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.""" +@pytest.mark.parametrize("invalid_field", ["joint_pos", "joint_vel", "articulation_root_prim_path"]) +def test_invalid_articulation_cfg_fails_initialization(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", - 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] - - -@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", - ) - 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 - - # 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 - - # 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 - - # 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() - - 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()) + 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"] - 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)) + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + with pytest.raises(error, match=match): + sim.reset() -def test_num_shapes_per_body_follows_public_body_order() -> None: - """Align Newton shape counts with the public body-name axis.""" - class _ShapeCountSurface: - backend_num_shapes_per_body = Articulation.backend_num_shapes_per_body - num_shapes_per_body = Articulation.num_shapes_per_body +def test_newton_rebind_refreshes_ordered_state_and_preserves_lab_owned_actuator_gains() -> None: + """Rebind public state to recreated Newton arrays, invalidate ordered caches, and keep Lab-owned gains. - articulation = _ShapeCountSurface() - articulation._num_shapes_per_body_backend = None - articulation._root_view = SimpleNamespace( - body_shapes=((), (object(), object()), (object(), object(), object())), - ) - articulation.body_ordering = SimpleNamespace( - user_to_backend_indices=(2, 0, 1), - ) + 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. - assert articulation.num_shapes_per_body == [3, 0, 2] - - -@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 + * 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. - 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, + 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 {}), ) - - 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, - ) - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) - - replicate(sim.get_clone_plan()) - sim.reset() - sim.step() - articulation.update(sim.cfg.dt) - - data = articulation.data - assert data.body_ordering is not None - root_body_idx = articulation.find_bodies("base")[0][0] - sim_timestamp = data._sim_timestamp - - 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) - - -@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, + for index, (actuation, ordering) in enumerate( + (actuation, ordering) for actuation in ("implicit", "explicit") for ordering in ("none", "reordered") ) - 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" - assert (articulation.data.joint_ordering is not None) is has_ordering - assert (articulation.data.body_ordering is not None) is has_ordering - - sim.step() - articulation.update(sim.cfg.dt) + } + # 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): + """Prime an implicit island's ordered state caches and return the check to run after the arrays are swapped.""" data = articulation.data + 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 @@ -1040,7 +453,7 @@ 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() - data.update(sim.cfg.dt) + data.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 @@ -1076,1191 +489,470 @@ 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 cache_name in _NEWTON_USER_ORDER_STATE_CACHES: + assert (getattr(data, cache_name) is not None) is has_ordering + + def check_rebind(new_state, new_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) - 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) + return check_rebind - 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) - - -@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) - - -@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. - """ - articulation_cfg = replace( - generate_articulation_cfg(articulation_type=articulation_type), - 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, - ) - 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() - - 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 - - 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.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. - - 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. - """ - 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 - - # 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) - - -@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`). - - The humanoid's articulation root sits on a rigid body, so this also checks floating-base initialization. - """ - 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() - - # 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) - - # -- 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 - - # 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,) - - # 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, - ) - 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) - - # 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) - - -@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. - - 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 - """ - 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 = [] - 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) - 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()) - prev_body_pos = body_pos.clone() - reported_max = torch.stack(reported_max).amax() - fin_diff_max = torch.stack(fin_diff_max).amax() - - # the commanded motion genuinely moves the bodies - 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 - - -@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. - - 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=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) +## +# Composite scene. +## - # 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"): - replicate(sim.get_clone_plan()) +@pytest.fixture(scope="module") +def composite_scene() -> _Scene: + """Initialize every articulation island once for this module.""" + 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)) -@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 +@pytest.fixture +def scene(composite_scene: _Scene) -> _Scene: + """Hand each test the composite scene with every island at rest. - Args: - sim: The simulation fixture - num_articulations: Number of articulations to test + Islands of one environment share a solver world, so a state one test leaves behind must not reach the next. """ - articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + for articulation in composite_scene.articulations.values(): + composite_scene.rest(articulation) + return composite_scene - # 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[..., 1] = 200.0 +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] - # 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 +def test_articulation_initialization_and_partial_state(scene: _Scene) -> None: + """Initialize every island and prove partial joint-state writes against live backend state. -@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. - - 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 + 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) - 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() - - 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)) -@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, + expected_fixed_base = { + "floating": False, + "fixed": True, + "passive": True, + "pendulum": True, + "ordered": False, + "tendon": True, + "root_fixed": True, + "root_freed": False, + "limits_implicit": False, + "limits_explicit": False, + "branching": False, } - 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) + 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 getattr(actuator, "is_implicit_model", False) == is_implicit_model_cfg + + floating = scene.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.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 ) - 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 - torch.testing.assert_close( - actuator.actuator_effort_limit, torch.full_like(newton_effort_limit, usd_effort_limit) - ) - + 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.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() -@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 - ) + 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) - # Play the simulator - replicate(sim.get_clone_plan()) - sim.reset() + # 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) - # Check that gains are loaded from the configuration - expected_stiffness = torch.full( - (articulation.num_instances, articulation.num_joints), 10.0, device=articulation.device + # 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_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())) - 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) - articulation.reset() - assert reset_env_ids == [None] - - instantaneous_composer = articulation.instantaneous_wrench_composer - permanent_composer = articulation.permanent_wrench_composer - - # Reset should zero external forces and torques - assert not instantaneous_composer.active - assert not permanent_composer.active - assert torch.count_nonzero(instantaneous_composer.out_force_b.torch) == 0 - assert torch.count_nonzero(instantaneous_composer.out_torque_b.torch) == 0 - assert torch.count_nonzero(permanent_composer.out_force_b.torch) == 0 - assert torch.count_nonzero(permanent_composer.out_torque_b.torch) == 0 + 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) - # A partial reset clears only the selected environment's wrenches - num_bodies = articulation.num_bodies - permanent_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), - ) - instantaneous_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), + # 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 ) - articulation.reset(env_ids=torch.tensor([0], device=device)) - assert instantaneous_composer.active - assert permanent_composer.active - assert torch.count_nonzero(instantaneous_composer.out_force_b.torch) == num_bodies * 3 - 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 - - -@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 - - # 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 - 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) - - # check they are set + # 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( - wp.to_torch(articulation.root_view.get_attribute("body_com", SimulationManager.get_model())), com + expected_com_quat.view(num_articulations, -1, 4), articulation.data.body_com_quat_w.torch ) - for i in range(50): - # perform step - sim.step() - # update buffers - articulation.update(sim.cfg.dt) + # 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.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) - # 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) + # 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) - # 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 + # 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) - # 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:]) +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. - # 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) + 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 = scene.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 inertial properties for the tests that drive this island + 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.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) - # 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) + # 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] * 2, device=_DEVICE) * 180.0 / torch.pi + 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.fixed + root_fixed = scene.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() - # 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) + # 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): + 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[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() - # 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:]) + # the commanded motion genuinely moves the bodies + 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("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) +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. - # Play sim - replicate(sim.get_clone_plan()) - sim.reset() + 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 = scene.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("base")[0][0] + root_idx = articulation.find_bodies("Root")[0][0] # 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) + 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 @@ -2281,432 +973,495 @@ def test_write_root_state_functions_data_consistency( assert not torch.allclose(updated_com_w, original_com_w) def random_root_state(num_envs: int) -> torch.Tensor: - state = torch.rand(num_envs, 13, device=device) + 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 - def write_root_state(state: torch.Tensor, env_ids: torch.Tensor | None = None): + for state_location in ("com", "link", "root"): + + 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 + scene.step("floating") + + # 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": - 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], + # 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) + + 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( + articulation.data.body_com_vel_w.torch[0], torch.zeros(2, 6, device=_DEVICE), atol=1e-6, rtol=0.0 ) - 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() + 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 = zeros[:, :1].clone() + positions[..., 1] = 0.1 + if is_global: + positions += articulation.data.body_com_pos_w.torch[1:, :1] + forces = zeros[:, :1].clone() + forces[..., 2] = force + 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 = 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) + articulation.reset() + assert reset_env_ids == [None] -@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): - """Joint limit and joint state writes update the joint buffers and refresh the body state without a step. + instantaneous_composer = articulation.instantaneous_wrench_composer + permanent_composer = articulation.permanent_wrench_composer - 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. Indexed joint limits that exclude a default joint position clamp it into the new limits - - Args: - sim: The simulation fixture - num_articulations: Number of articulations to test - device: The device to run the simulation on + # Reset should zero external forces and torques + assert not instantaneous_composer.active + assert not permanent_composer.active + assert torch.count_nonzero(instantaneous_composer.out_force_b.torch) == 0 + assert torch.count_nonzero(instantaneous_composer.out_torque_b.torch) == 0 + assert torch.count_nonzero(permanent_composer.out_force_b.torch) == 0 + assert torch.count_nonzero(permanent_composer.out_torque_b.torch) == 0 + + # A partial reset clears only the selected environment's wrenches + num_bodies = articulation.num_bodies + permanent_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), + ) + instantaneous_composer.add_forces_and_torques_index( + forces=torch.ones((num_articulations, num_bodies, 3), device=_DEVICE), + torques=torch.ones((num_articulations, num_bodies, 3), device=_DEVICE), + ) + articulation.reset(env_ids=torch.tensor([0], device=_DEVICE)) + assert instantaneous_composer.active + assert permanent_composer.active + assert torch.count_nonzero(instantaneous_composer.out_force_b.torch) == num_bodies * 3 + 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() + + +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. + + 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 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. """ - 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)]) + articulation = scene.branching + device = _DEVICE - # Play sim - replicate(sim.get_clone_plan()) - sim.reset() + # 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 - # Get current default joint pos - default_joint_pos = articulation.data.default_joint_pos.torch.clone() + # 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) - 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) + # 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 - # 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) + # 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 - # 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 + # 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, ) - 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) + 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() - from torch.distributions import Uniform + 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()) - 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) + 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)) - 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.set_masses_index(masses=initial_masses) + articulation.set_inertias_index(inertias=initial_inertias) - rand_joint_pos = pos_dist.sample() - rand_joint_vel = vel_dist.sample() - articulation.write_joint_position_to_sim_index(position=rand_joint_pos) - articulation.write_joint_velocity_to_sim_index(velocity=rand_joint_vel) - # make sure valued updated - body_link_pose_w = articulation.data.body_link_pose_w.torch - body_com_vel_w = articulation.data.body_com_vel_w.torch - original_body_states = torch.cat([original_body_link_pose_w, original_body_com_vel_w], dim=-1) - body_state_w = torch.cat([body_link_pose_w, body_com_vel_w], dim=-1) - assert torch.count_nonzero(original_body_states[:, 1:] != body_state_w[:, 1:]) > ( - len(original_body_states[:, 1:]) / 2 - ) - # 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:]) +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. - # 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), + 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. + + 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. + """ + articulation = scene.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) + + # 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() + + 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] ) - 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) + np.testing.assert_allclose(data._body_com_vel_w_user.numpy(), data._sim_bind_body_com_vel_w.numpy()[:, body_u2b]) + + articulation.update(scene.sim.cfg.dt) + root_body_idx = articulation.find_bodies("base")[0][0] + sim_timestamp = data._sim_timestamp - # 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.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) + 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) - # 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) + # 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) -@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. +@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. - 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 = generate_articulation_cfg(articulation_type) - articulation_cfg.actuators["panda_shoulder"].viscous_friction = 0.25 - articulation, _ = generate_articulation(articulation_cfg, num_articulations, device) - replicate(sim.get_clone_plan()) - sim.reset() - - shoulder_joint_ids = articulation.actuators["panda_shoulder"].joint_indices - expected_viscous_friction = torch.full((articulation.num_instances, 4), 0.25, device=device) - torch.testing.assert_close( - articulation.data.joint_viscous_friction_coeff.torch[:, shoulder_joint_ids], expected_viscous_friction - ) - torch.testing.assert_close( - wp.to_torch(articulation.root_view.get_attribute("joint_damping", SimulationManager.get_model()))[ - :, 0, shoulder_joint_ids - ], - expected_viscous_friction, - ) - - expected_pd_damping = torch.full_like(expected_viscous_friction, 4.0) - torch.testing.assert_close(articulation.data.joint_damping.torch[:, shoulder_joint_ids], expected_pd_damping) - torch.testing.assert_close( - wp.to_torch(articulation.root_view.get_attribute("joint_target_kd", SimulationManager.get_model()))[ - :, 0, shoulder_joint_ids - ], - expected_pd_damping, - ) + 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 - 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. """ - from unittest.mock import patch - - 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 - ) - - # Play the simulator - replicate(sim.get_clone_plan()) - sim.reset() + articulation = scene.floating + num_articulations = articulation.num_instances + device = _DEVICE - # 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) # 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() @@ -2720,18 +1475,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) @@ -2741,190 +1485,75 @@ def robot_shapes(env_index: int, selected_body_names: list[str] | None = None) - 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 - - -## -# 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_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, - ) + articulation = scene.tendon + scene.rest(articulation) + num_articulations = articulation.num_instances + assert articulation.num_fixed_tendons == 2 - 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" - ) + 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) + + # 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]) - franka_J = franka.data.body_link_jacobian_w.torch - anymal_J = anymal.data.body_link_jacobian_w.torch + 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() - # 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)}" - ) + 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") - sim.step() - franka.update(sim.cfg.dt) - anymal.update(sim.cfg.dt) + 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) - franka_M = franka.data.mass_matrix.torch - anymal_M = anymal.data.mass_matrix.torch + 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) - 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)) + # 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() - # 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" - ) - # 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" - - -@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 @@ -2936,139 +1565,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 = _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 = [island.num_joints + island.num_base_dofs for island 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 @@ -3079,8 +1692,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). @@ -3108,292 +1721,248 @@ 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_get_gravity_compensation_forces_static_equilibrium(scene: _Scene) -> None: + """Newton accuracy: ``τ_gc`` must hold a passive manipulator in static equilibrium. -@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. + 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 chain should not move. The island's actuators have zero gains, so the effort + target IS the joint torque applied and no PD spring-damper masks the gravity-compensation signal. Sign + errors, frame errors, and DoF-ordering errors all surface as joint drift. """ - 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. + articulation = scene.passive + scene.rest(articulation) + # a configuration that loads the prismatic and the revolute joints with gravity + q = articulation.data.default_joint_pos.torch.clone() + q[:, 3:] = torch.tensor([0.4, -0.3, 0.5], device=_DEVICE) + articulation.write_joint_position_to_sim_index(position=q) + articulation.update(scene.sim.cfg.dt) 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). + with world_gravity(_GRAVITY): 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) + assert tau_gc.abs().max() > 1.0, "the configuration carries no gravity load" + for _ in range(100): + # ``gravity_compensation_forces`` leads with ``num_base_dofs`` floating-base entries (0 on fixed-base); + # slice past them so the remaining tensor aligns with the actuated joints. + tau_gc = articulation.data.gravity_compensation_forces.torch[:, articulation.num_base_dofs :] + articulation.actuators.target_command.set_effort_index(value=tau_gc) + scene.step("passive") + final_q = articulation.data.joint_pos.torch.clone() - 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. + # A sign or frame bug in τ_gc produces drift of at least a degree per step on this configuration. 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." + f"max joint drift {drift:.5f} 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." ) + scene.rest(articulation) -@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)``. +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``. """ - robot, ee_frame_idx, ee_jacobi_idx, arm_joint_ids = _setup_franka_at_home_pose(sim, zero_actuator_pd=True) - - 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, - ) + articulation = scene.pendulum + scene.rest(articulation) + num_articulations = articulation.num_instances + num_bodies = articulation.num_bodies + 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) - - pos_history: list[float] = [] - rot_history: list[float] = [] - for _ in range(800): - 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) - 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=None, - ) + # 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") - 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_history.append(pos_error.norm(dim=-1).max().item()) - rot_history.append(rot_error.norm(dim=-1).max().item()) - - 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. - 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): - """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. + # check that the root is at the correct state - its default state as it is fixed base + 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 = (link_offset[0]) * joint_vel * torch.cos(joint_pos) + vz = torch.zeros(num_articulations, 1, 1, device=device) + 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 = (link_offset[0] + offset[0]) * torch.sin(joint_pos) + pz = torch.zeros(num_articulations, 1, 1, device=device) + pos_gt[:, arm_idx, :] = torch.cat([px, py, pz], dim=-1).squeeze(-2) + pos_gt += pivot_pos_w.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:]) + + # 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) + + +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`). """ - robot, ee_frame_idx, ee_jacobi_idx, arm_joint_ids = _setup_franka_at_home_pose( - sim, zero_actuator_pd=True, disable_gravity=False - ) + articulation = scene.floating + scene.rest(articulation) + num_articulations = articulation.num_instances + device = _DEVICE - 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, + # 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,) + + # 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, ) + with world_gravity((0.0, 0.0, 0.0)): + wp.to_torch(model_gravity_arr).copy_(new_gravity) + SimulationManager.add_model_change(ModelFlags.MODEL_PROPERTIES) - sim.step() - robot.update(sim.cfg.dt) - # 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) + # 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) - 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. - """ - a = sum(history[-200:-100]) / 100 - b = sum(history[-100:]) / 100 - 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" + # 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()) ) - return mean - - hist_off = run_phase(400) - osc.cfg.gravity_compensation = True - hist_on = run_phase(600) + 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) - 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)" +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. -if __name__ == "__main__": - pytest.main([__file__, "-v", "--maxfail=1"]) + 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. + """ + articulation = scene.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): + 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) diff --git a/source/isaaclab_newton/test/assets/unit/test_articulation_accessors.py b/source/isaaclab_newton/test/assets/unit/test_articulation_accessors.py new file mode 100644 index 000000000000..7e5e105f3e81 --- /dev/null +++ b/source/isaaclab_newton/test/assets/unit/test_articulation_accessors.py @@ -0,0 +1,110 @@ +# 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 + +"""Newton articulation accessors checked on a directly built model: lazy task-space allocation with graph capture, +and the public body order of the collision-shape counts. +""" + +from types import SimpleNamespace + +import newton +import numpy as np +import pytest +import warp as wp +from isaaclab_newton.assets import Articulation +from isaaclab_newton.assets.articulation.articulation_data import ArticulationData +from isaaclab_newton.physics import NewtonManager as SimulationManager +from newton import JointType, ModelBuilder +from newton.selection import ArticulationView + +from isaaclab.test.utils import test_devices + +pytestmark = pytest.mark.unit + + +def test_num_shapes_per_body_follows_public_body_order() -> None: + """Align Newton shape counts with the public body-name axis.""" + + class _ShapeCountSurface: + backend_num_shapes_per_body = Articulation.backend_num_shapes_per_body + num_shapes_per_body = Articulation.num_shapes_per_body + + articulation = _ShapeCountSurface() + articulation._num_shapes_per_body_backend = None + articulation._root_view = SimpleNamespace( + body_shapes=((), (object(), object()), (object(), object(), object())), + ) + articulation.body_ordering = SimpleNamespace( + user_to_backend_indices=(2, 0, 1), + ) + + assert articulation.num_shapes_per_body == [3, 0, 2] + + +@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 From 7ec1bc045ecaf37bcf1ebf37ffe5a61159e20163 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 11:16:50 +0200 Subject: [PATCH 15/33] Run Newton actuator equivalence on local composite scenes --- .../assets/test_newton_actuators_newton.py | 1255 ++++++++--------- .../unit/test_newton_actuator_adaptation.py | 295 ++++ 2 files changed, 869 insertions(+), 681 deletions(-) create mode 100644 source/isaaclab_newton/test/assets/unit/test_newton_actuator_adaptation.py 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..64286fb80a5d 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,54 @@ # # 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 +import sys +from dataclasses import dataclass, field +from pathlib import Path -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.physics import MJWarpSolverCfg from isaaclab_newton.physics import NewtonManager as SimulationManager -import isaaclab.sim as sim_utils from isaaclab.actuators import DCMotorCfg, DelayedPDActuatorCfg, IdealPDActuatorCfg +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 +58,41 @@ 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 +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from articulation_test_utils import NUM_ENVS, local_usd, newton_sim_cfg, spawn_assets # noqa: E402 + +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 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 +120,630 @@ ), } +# 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: + """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 + 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.""" + tolerances: dict[str, float] = field(default_factory=dict) + """Overrides of the equivalence oracle tolerances.""" + + +def _islands(mlp_path: str | None = None, lstm_path: str | None = None) -> dict[str, _Island]: + 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, tolerances={"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: + 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. - - 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. +) -> _Run: + """Spawn one island per actuator configuration, roll every island out, and record trajectories and telemetry. + + 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. - """ - 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`. + 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( + { + "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. """ - - __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) -> 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) -> _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. """ - - __test__ = True - actuators = MIXED_WITH_IMPLICIT_ACTUATORS + 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, **tolerances: float) -> None: + """Run the shared trajectory and telemetry oracles on one island's Lab and Newton rollouts.""" + oracle = type("_Oracle", (EquivalenceAssertionsMixin,), tolerances)() + 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() # --------------------------------------------------------------------------- -# Implicit + non-zero feedforward effort target +# Equivalence tests with different actuator types # --------------------------------------------------------------------------- -class TestImplicitWithFeedforwardEquivalence(_EquivalenceTestBase): - """Implicit-only actuators with a non-zero feedforward effort target. +@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. - 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. + 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. """ - - __test__ = True - actuators = IMPLICIT_ONLY_ACTUATORS - feedforward = 2.0 - torque_atol = 0.5 + _assert_equivalent(lab_run["results"][island], newton_run.results[island], **_islands()[island].tolerances) -# --------------------------------------------------------------------------- -# Multi-articulation Newton scene (regression test for class-attr clobber) -# --------------------------------------------------------------------------- - +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"]) -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 - 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 +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" - 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_newton_actuators_are_delayed_pd(newton_run: _Run) -> None: + """Command delays are authored as ``NewtonActuatorDelayAPI`` on a PD drive.""" + actuator_info = newton_run.results["delayed"]["actuator_info"] + assert actuator_info, "No Newton actuators were created" + for a in actuator_info: + assert a["has_delay"], "Delay not found on delayed PD actuator" + assert a["controller_type"] == "DrivePD" - return {"joint_pos_anymal": pos_anymal, "joint_pos_cartpole": pos_cartpole} - -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. - """ - - @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}", - ) +def test_newton_knee_actuators_are_delayed_position_clamped_pd(newton_run: _Run) -> None: + """RemotizedPD knees are authored as a delayed PD drive with position-based clamping.""" + kfe_acts = [ + a for a in newton_run.results["remotized"]["actuator_info"] if "ClampingPositionBased" in a["clamping_types"] + ] + assert len(kfe_acts) > 0, "No actuator with position-based clamping found" + for a in kfe_acts: + assert a["controller_type"] == "DrivePD" + assert a["has_delay"], "Delay not found on remotized KFE actuator" # --------------------------------------------------------------------------- -# Domain randomization via events.py — Newton backend +# Neural network actuator authoring: MLP and LSTM # --------------------------------------------------------------------------- -class TestRandomizeActuatorGainsViaEventsNewton(unittest.TestCase): - """End-to-end DR test for the Newton backend. +def test_mlp_has_dc_motor_clamping(newton_run: _Run) -> None: + """ActuatorNetMLPCfg is authored as a Newton neural MLP controller with DC motor clamping.""" + mlp_acts = [a for a in newton_run.results["neural"]["actuator_info"] if a["controller_type"] == "DriveNeuralMLP"] + assert len(mlp_acts) > 0, "No NeuralMLP controller found" + for a in mlp_acts: + assert "ClampingDCMotor" in a["clamping_types"] - 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``. - 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. - """ +def test_lstm_has_dc_motor_clamping(newton_run: _Run) -> None: + """ActuatorNetLSTMCfg is authored as a Newton neural LSTM controller with DC motor clamping.""" + lstm_acts = [a for a in newton_run.results["neural"]["actuator_info"] if a["controller_type"] == "DriveNeuralLSTM"] + assert len(lstm_acts) > 0, "No NeuralLSTM controller found" + for a in lstm_acts: + assert "ClampingDCMotor" in a["clamping_types"] - 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) - ) - - # 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] - ) - # 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] - ) +def test_positions_finite(newton_run: _Run) -> None: + """The neural actuators run on the Newton backend without producing non-finite joint positions.""" + for step_i, pos in enumerate(newton_run.results["neural"]["joint_pos"]): + assert torch.isfinite(pos).all(), f"Non-finite positions at step {step_i}" # --------------------------------------------------------------------------- -# 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 controllers of the articulation's Newton actuators. - 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``; the assertions read the controllers back via the public + ``read_group_parameter``. 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. """ + anymal = newton_run.articulations["ideal"] + cartpole = newton_run.articulations["cartpole"] + assert SimulationManager._adapter is not None + + 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 (6 free-root DOFs + the leg 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", + ) - 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)]), - ) - articulation = Articulation(cfg) - replicate(plan) - return articulation - - def _get_adapter(self, articulation): - return SimulationManager._adapter - - -# --------------------------------------------------------------------------- -# RemotizedPD actuator: PD + delay + position-based clamping lookup table -# --------------------------------------------------------------------------- - - -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 - - @classmethod - def setUpClass(cls): - cls.actuators = _remotized_pd_actuators() - super().setUpClass() + 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_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") + # The leg 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] + ) -class TestDecimationRemotizedPD(_DecimationMixin, TestRemotizedPDEquivalence): - """RemotizedPD — decimation=2 + CUDA graph.""" + # DR scoped to the floating-base leg 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]) # --------------------------------------------------------------------------- -# Neural network actuator authoring: MLP and LSTM +# Per-env reset: actuator state isolation # --------------------------------------------------------------------------- -class TestNeuralActuatorAuthoring(unittest.TestCase): - """Verify ActuatorNetMLPCfg and ActuatorNetLSTMCfg are authored as Newton neural - controllers with DC motor clamping and run on the Newton backend. - """ - - @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)) +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].""" + articulation = newton_run.articulations["delayed"] + adapter = SimulationManager._adapter + assert adapter is not None + # Find the delayed actuators (only they have 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 + ] + 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/unit/test_newton_actuator_adaptation.py b/source/isaaclab_newton/test/assets/unit/test_newton_actuator_adaptation.py new file mode 100644 index 000000000000..fbc11daab888 --- /dev/null +++ b/source/isaaclab_newton/test/assets/unit/test_newton_actuator_adaptation.py @@ -0,0 +1,295 @@ +# 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 + +"""Newton actuator adaptation: target modes resolved on the model builder, native-group selection, and telemetry. + +These run on a model builder and on Warp arrays directly; none of them steps a simulation. +""" + +from types import SimpleNamespace + +import numpy as np +import pytest +import warp as wp +from isaaclab_newton.assets import Articulation +from isaaclab_newton.assets.articulation.actuator_control import NewtonActuatorControl +from isaaclab_newton.physics import NewtonBuilderCfg, NewtonCfg +from isaaclab_newton.physics import NewtonManager as SimulationManager +from newton import JointTargetMode, JointType, ModelBuilder +from newton.solvers import SolverMuJoCo + +import isaaclab.sim as sim_utils +from isaaclab.actuators import IdealPDActuatorCfg, ImplicitActuator, ImplicitActuatorCfg +from isaaclab.actuators.newton.kernels import sync_torque_telemetry +from isaaclab.assets import ArticulationCfg + +pytestmark = pytest.mark.unit + + +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. + + 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)] + + +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)] + + +@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] + + +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)) From c3ed18b5466b2f758d460a28c47bfd3a27aaff34 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 11:16:50 +0200 Subject: [PATCH 16/33] Run Newton rigid-object tests on persistent local scenes --- .../test/assets/test_rigid_object.py | 888 ++++++++---------- .../assets/test_rigid_object_collection.py | 765 +++++++-------- 2 files changed, 762 insertions(+), 891 deletions(-) diff --git a/source/isaaclab_newton/test/assets/test_rigid_object.py b/source/isaaclab_newton/test/assets/test_rigid_object.py index 39a756a60b12..ef7b91fd22a8 100644 --- a/source/isaaclab_newton/test/assets/test_rigid_object.py +++ b/source/isaaclab_newton/test/assets/test_rigid_object.py @@ -6,33 +6,38 @@ # 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 coverage on one persistent two-environment scene per device. -launch_test_simulation() +The cubes are locally spawned shapes. 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 case build their own scenes. +""" + +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 dataclasses import dataclass +from pathlib import Path 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.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 -from isaaclab.cloner import CloneCfg, clone_plan_from_env_0, replicate +from isaaclab.assets import RigidObjectCfg 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 test_devices from isaaclab.utils.math import ( combine_frame_transforms, quat_apply_inverse, @@ -42,107 +47,138 @@ random_orientation, ) -NEWTON_SIM_CFG = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg(), - ), -) +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from articulation_test_utils import env_origins, local_usd, newton_sim_cfg, spawn_assets, world_gravity # noqa: E402 -def _newton_sim_context(device, gravity_enabled=True, dt=None, **kwargs): - """Helper to create a Newton simulation context with the correct device. +pytestmark = [pytest.mark.integration, pytest.mark.kitless] - 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. +_GRAVITY = (0.0, 0.0, -9.81) - """ - 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": - # 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( +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=spawn_cfg, + 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, 0.0, height)), ) - 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): +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() + assert cube_object.is_initialized + 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) + + 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) + + +## +# Shared scene. +## + + +@dataclass +class _Scene: + """Two-environment cubes that share one real Newton model and solver lifecycle.""" + + sim: SimulationContext + cube: RigidObject + origins: torch.Tensor + device: str + + def step(self, num_steps: int = 1) -> None: + for _ in range(num_steps): + self.cube.write_data_to_sim() + self.sim.step() + self.cube.update(self.sim.cfg.dt) + + def rest(self) -> None: + """Put the cubes at rest at their configured pose with the configured inertial properties.""" + 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() + + +@pytest.fixture(scope="module", params=test_devices()) +def shared_scene(request) -> _Scene: + """Initialize the cubes once per device for this module.""" + device = request.param + with build_simulation_context(sim_cfg=newton_sim_cfg(device)) as sim: + cube = spawn_assets({"cube": _cube_cfg()})["cube"] + sim.reset() + scene = _Scene(sim=sim, cube=cube, origins=env_origins(device), device=device) + initial_properties = ( + cube.data.body_mass.torch.clone(), + cube.data.body_com_pos_b.torch.clone(), + cube.data.body_inertia.torch.clone(), + ) + yield scene, initial_properties + + +@pytest.fixture +def scene(shared_scene) -> _Scene: + """Hand each test the shared scene at rest with the configured inertial properties.""" + scene, (masses, coms, inertias) = shared_scene + scene.cube.set_masses_index(masses=masses) + scene.cube.set_coms_index(coms=coms) + scene.cube.set_inertias_index(inertias=inertias) + scene.rest() + return scene + + +def test_external_force_on_single_body(scene: _Scene) -> None: """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 @@ -153,55 +189,40 @@ def test_external_force_on_single_body(num_cubes, device): 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()) - 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) - + cube_object = scene.cube + num_cubes = cube_object.num_instances + device = scene.device + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(cube_object) < 10 + + # 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(): + scene.rest() + + # 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 + + with world_gravity(_GRAVITY): # 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_b = torch.zeros(cube_object.num_instances, len(body_ids), 6, device=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] @@ -223,18 +244,18 @@ def simulate(): body_ids=body_ids, is_global=is_global, ) - simulate() + scene.step(5) # 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) + cube_object.data.root_pos_w.torch[0::2, 2], torch.ones(num_cubes // 2, device=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) # 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 = torch.zeros(cube_object.num_instances, len(body_ids), 6, device=device) + external_wrench_positions_b = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=device) external_wrench_b[0::2, :, 2] = 50.0 external_wrench_positions_b[0::2, :, 1] = 1.0 @@ -269,7 +290,7 @@ def simulate(): body_ids=body_ids, is_global=is_global, ) - simulate() + scene.step(5) # 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) @@ -277,165 +298,122 @@ def simulate(): assert torch.all(cube_object.data.root_pos_w.torch[1::2, 2] < 1.0) -@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): +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.""" - 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() - - # 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) + 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() + 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=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) - # Simulate physics - sim.step() - cube_object.update(sim.cfg.dt) + # Simulate physics + scene.step() - 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]) + 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]) -@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): +def test_rigid_body_set_mass(scene: _Scene) -> None: """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) - - # Play sim - replicate(sim.get_clone_plan()) - sim.reset() - - # 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 - # perform rendering - sim.step() - # update object - cube_object.update(sim.cfg.dt) + 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), + ) - masses_to_check = cube_object.data.body_mass.torch[env_ids][:, body_ids] + # 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) @@ -453,7 +431,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()): @@ -465,149 +443,113 @@ 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:]) - - -@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): + cube_object = scene.cube + num_cubes = cube_object.num_instances + device = scene.device + env_pos = scene.origins + torch.tensor([0.0, 0.0, 1.0], device=device) + + # 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) + + # 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)) + 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:]) + + +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() - - # Check if cube_object is initialized - assert cube_object.is_initialized + 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) - # change center of mass offset from link frame - offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_cubes, 1) + # 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)) - # 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) - - # 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( @@ -638,10 +580,12 @@ 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 already reflect the write on their own (regressions: + # body_com_pose_w returned the pre-write buffer 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 @@ -683,23 +627,20 @@ 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`` @@ -714,59 +655,52 @@ 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) - - 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" - ) - # 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) + 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_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" + ) + # 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) diff --git a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py index 8724ad8ac51d..a10e472056c4 100644 --- a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py @@ -6,28 +6,36 @@ # 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 +"""Kitless real-Newton rigid-object-collection coverage on one persistent scene per device. -launch_test_simulation() +Every environment holds three locally spawned cubes in the collection and an unrelated rigid body beside them, +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 +case build their own scenes. +""" + +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 dataclasses import dataclass +from pathlib import Path -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.assets import RigidObject, RigidObjectCollection 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.assets import RigidObjectCfg, RigidObjectCollectionCfg +from isaaclab.sim import SimulationContext, build_simulation_context +from isaaclab.test.utils import test_devices from isaaclab.utils.math import ( combine_frame_transforms, quat_apply_inverse, @@ -38,222 +46,217 @@ subtract_frame_transforms, ) -NEWTON_SIM_CFG = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg(), - ), -) +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from articulation_test_utils import env_origins, newton_sim_cfg, spawn_assets, world_gravity # noqa: E402 -def _newton_sim_context(device, gravity_enabled=True, dt=None, **kwargs): - """Helper to create a Newton simulation context with the correct device. +pytestmark = [pytest.mark.integration, pytest.mark.kitless] - 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, -) -> 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. - spawn_unrelated_sibling: Whether to spawn a rigid body outside the collection in each environment. - - Returns: - A tuple containing the rigid object collection representing the cubes and the origins of the cubes. +_GRAVITY = (0.0, 0.0, -9.81) +_NUM_CUBES = 3 - """ - 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 = {} - for i in range(num_cubes): - cube_object_cfg = RigidObjectCfg( - prim_path=f"/World/Env_[^/]*/Object_{i}", - 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_[^/]*/UnrelatedObject")) - 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/UnrelatedObject", - 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) +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 + 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(), + ) + + +def _collection_cfg(num_cubes: int, height: float = 1.0, rigid: bool = True) -> RigidObjectCollectionCfg: + """Return a collection of cubes spaced 3 m apart along y.""" + return RigidObjectCollectionCfg( + rigid_objects={ + f"cube_{i}": RigidObjectCfg( + prim_path=f"/World/Env_[^/]*/Object_{i}", + spawn=_cube_spawn_cfg(rigid), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 3.0 * i, height)), + ) + for i in range(num_cubes) + } + ) - return cube_object_collection, 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.parametrize(("num_envs", "num_cubes", "spawn_unrelated_sibling"), [(1, 1, False), (2, 3, True)]) -@pytest.mark.parametrize("device", test_devices()) -def test_initialization(num_envs, num_cubes, spawn_unrelated_sibling, device): - """Test initialization for prim with rigid body API at the provided prim path. - With an unrelated rigid body next to the cubes in each environment, the collection view must still - select only its configured rigid objects. - """ - with _newton_sim_context(device, 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, - spawn_unrelated_sibling=spawn_unrelated_sibling, - ) +def test_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"] # 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) + with pytest.raises(RuntimeError, match="Expected 1 prims at"): + sim.reset() @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(), +def test_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: + """Cube collections that share one real Newton model and solver lifecycle.""" + + sim: SimulationContext + collection: RigidObjectCollection + sibling: RigidObject + origins: torch.Tensor + device: str + + def step(self, num_steps: int = 1) -> None: + for _ in range(num_steps): + self.collection.write_data_to_sim() + self.sim.step() + self.collection.update(self.sim.cfg.dt) + + def rest(self) -> None: + """Put the cubes at rest at their configured poses.""" + 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) ) - 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) - } + self.collection.reset() + + +@pytest.fixture(scope="module", params=test_devices()) +def shared_scene(request): + """Initialize the collection and its unrelated sibling bodies once per device for this module.""" + device = request.param + sibling_cfg = RigidObjectCfg( + prim_path="/World/Env_[^/]*/UnrelatedObject", + spawn=_cube_spawn_cfg(), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, -3.0, 1.0)), + ) + with build_simulation_context(sim_cfg=newton_sim_cfg(device)) as sim: + assets = spawn_assets({"collection": _collection_cfg(_NUM_CUBES), "sibling": sibling_cfg}) + sim.reset() + collection = assets["collection"] + scene = _Scene( + sim=sim, collection=collection, sibling=assets["sibling"], origins=env_origins(device), device=device ) - clone_plan_from_env_0( - CloneCfg(clone_template="/World/Env_{}"), - cfg.rigid_objects.values(), - num_envs, - 1.0, - positions=origins, + initial_properties = ( + collection.data.body_mass.torch.clone(), + collection.data.body_com_pos_b.torch.clone(), + collection.data.body_inertia.torch.clone(), ) - object_collection = RigidObjectCollection(cfg) - replicate(sim.get_clone_plan()) - sim.reset() + yield scene, initial_properties - 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 +@pytest.fixture +def scene(shared_scene) -> _Scene: + """Hand each test the shared scene at rest with the configured inertial properties.""" + scene, (masses, coms, inertias) = shared_scene + scene.collection.set_masses_index(masses=masses) + scene.collection.set_coms_index(coms=coms) + scene.collection.set_inertias_index(inertias=inertias) + scene.rest() + return scene - # Play sim - replicate(sim.get_clone_plan()) - with pytest.raises(RuntimeError, match="Expected 1 prims at"): - sim.reset() +def test_initialization(scene: _Scene) -> None: + """Test initialization for prims with rigid body API at the provided prim paths. -@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): + With an unrelated rigid body next to the cubes 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.sibling.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 sibling. + 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_external_force_on_single_body(scene: _Scene) -> None: """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 @@ -261,38 +264,24 @@ def test_external_force_on_single_body(num_envs, num_cubes, device): 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() + object_collection = scene.collection + device = scene.device - # 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) + # find objects to apply the force + object_ids, object_names = object_collection.find_bodies(".*") + def reset_objects(): + scene.rest() + + # 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 + + with world_gravity(_GRAVITY): # 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_b = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=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] @@ -315,7 +304,7 @@ def simulate(): env_ids=None, is_global=is_global, ) - simulate() + scene.step(10) # First object should still be at the same Z position (1.0) torch.testing.assert_close( @@ -326,10 +315,8 @@ def simulate(): 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 = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=device) + external_wrench_positions_b = torch.zeros(object_collection.num_instances, len(object_ids), 3, device=device) external_wrench_b[:, 0::2, 2] = 50.0 external_wrench_positions_b[:, 0::2, 1] = 1.0 @@ -365,7 +352,7 @@ def simulate(): body_ids=object_ids, is_global=is_global, ) - simulate() + scene.step(10) # 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) @@ -373,32 +360,24 @@ def simulate(): 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): +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 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() + 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((0.0, 0.0, -9.81), device=device) - ) + 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. - for _ in range(2): - sim.step() - object_collection.update(sim.cfg.dt) + 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) @@ -422,140 +401,111 @@ def test_gravity_vec_w_tracks_model_gravity(num_envs, num_cubes, device): # 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) + 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) -@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): +def test_object_state_properties(scene: _Scene) -> None: """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], - ) + 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) + + # 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 _ 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) + 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) + # 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, - ) + # 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) + # 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:]) + # 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): +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. """ - 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() + 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) - # 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) - # 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)) - # 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) + # 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): - sim.step() - cube_object.update(sim.cfg.dt) + 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].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), + 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. @@ -589,8 +539,8 @@ def test_write_object_state(num_envs, num_cubes, device, state_location): 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), + 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), ) @@ -602,8 +552,8 @@ def test_write_object_state(num_envs, num_cubes, device, state_location): 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), + link_pose_w[..., :3].reshape(-1, 3), + link_pose_w[..., 3:].reshape(-1, 4), object_link_to_com_pos, object_link_to_com_quat, ) @@ -617,16 +567,13 @@ def test_write_object_state(num_envs, num_cubes, device, state_location): # 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) + 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.parametrize("device", test_devices()) -@pytest.mark.isaacsim_ci -def test_body_pose_write_marks_fk_reset_mask(device): +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 @@ -642,47 +589,37 @@ 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) + cube_object = scene.collection + scene.step() + + 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) From 8d797e39dfe215c3fe97254628ec5452669ffcc2 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 11:16:50 +0200 Subject: [PATCH 17/33] Cover Newton operational-space control on a local chain --- .../antoiner-asset-tests-v2-newton.skip | 0 .../test_articulation_ordering_kernels.py | 2 + .../test_newton_task_space_controllers.py | 252 ++++++++++++++++++ 3 files changed, 254 insertions(+) create mode 100644 source/isaaclab_newton/changelog.d/antoiner-asset-tests-v2-newton.skip create mode 100644 source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py diff --git a/source/isaaclab_newton/changelog.d/antoiner-asset-tests-v2-newton.skip b/source/isaaclab_newton/changelog.d/antoiner-asset-tests-v2-newton.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_newton/test/assets/test_articulation_ordering_kernels.py b/source/isaaclab_newton/test/assets/test_articulation_ordering_kernels.py index fcb4e7759dee..fc7fb6bd445e 100644 --- a/source/isaaclab_newton/test/assets/test_articulation_ordering_kernels.py +++ b/source/isaaclab_newton/test/assets/test_articulation_ordering_kernels.py @@ -10,6 +10,8 @@ import warp as wp from isaaclab_newton.assets.articulation import kernels as articulation_kernels +pytestmark = pytest.mark.unit + def _selector(values: list[int], dtype: type) -> wp.array: """Create a CPU Warp selector with the requested integer width.""" diff --git a/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py b/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py new file mode 100644 index 000000000000..f255a75efe2d --- /dev/null +++ b/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py @@ -0,0 +1,252 @@ +# 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 + +"""Operational-space control of a locally authored Newton chain through the articulation's dynamics accessors. + +The Isaac Lab :class:`~isaaclab.controllers.OperationalSpaceController` consumes +:attr:`~isaaclab.assets.BaseArticulationData.body_link_jacobian_w`, +:attr:`~isaaclab.assets.BaseArticulationData.mass_matrix`, and +:attr:`~isaaclab.assets.BaseArticulationData.gravity_compensation_forces` every step. These sentinels close the +loop on a six-DOF fixed chain whose links carry center-of-mass offsets, so a wrong Jacobian, mass matrix, gravity +force, or DoF ordering pushes the steady-state error well past the bounds. The chain's joints are unpowered, so the +controller's joint efforts are the only actuation. +""" + +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 dataclasses import dataclass +from pathlib import Path + +import pytest +import torch +from isaaclab_newton.assets import Articulation + +from isaaclab.actuators import ImplicitActuatorCfg +from isaaclab.assets import ArticulationCfg +from isaaclab.controllers import OperationalSpaceController, OperationalSpaceControllerCfg +from isaaclab.sim import SimulationContext, build_simulation_context +from isaaclab.utils.math import compute_pose_error, matrix_from_quat, quat_inv, subtract_frame_transforms + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from articulation_test_utils import NUM_ENVS, local_usd, newton_sim_cfg, spawn_assets, world_gravity # noqa: E402 + +pytestmark = [pytest.mark.integration, pytest.mark.kitless] + +_DEVICE = "cpu" + + +@dataclass +class _Chain: + """A passive six-DOF chain and the end-effector indices the controller reads.""" + + sim: SimulationContext + robot: Articulation + ee_frame_idx: int + ee_jacobi_idx: int + arm_joint_ids: list[int] + + def rest(self) -> None: + """Hold the default configuration at rest with no effort command.""" + default_joint_pos = self.robot.data.default_joint_pos.torch.clone() + self.robot.write_joint_state_to_sim_index( + position=default_joint_pos, velocity=torch.zeros_like(default_joint_pos) + ) + self.robot.actuators.target_command.set_effort_index(value=torch.zeros_like(default_joint_pos)) + self.robot.reset() + self.robot.write_data_to_sim() + self.sim.step() + self.robot.update(self.sim.cfg.dt) + + +@pytest.fixture(scope="module") +def chain() -> _Chain: + """Build the unpowered chain once, at a configuration away from the joint limits.""" + robot_cfg = ArticulationCfg( + prim_path="/World/Env_[^/]*/Robot", + spawn=local_usd("fixed_spatial_chain.usda"), + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.0, 0.0, 1.0), joint_pos={"Joint_3": 0.3, "Joint_4": -0.2, "Joint_5": 0.4} + ), + actuators={"arm": ImplicitActuatorCfg(joint_names_expr=["Joint_.*"], stiffness=0.0, damping=0.0)}, + ) + with build_simulation_context(sim_cfg=newton_sim_cfg(_DEVICE)) as sim: + robot = spawn_assets({"robot": robot_cfg})["robot"] + sim.reset() + assert robot.is_initialized and robot.is_fixed_base + ee_frame_idx = robot.find_bodies("Link_5")[0][0] + yield _Chain( + sim=sim, + robot=robot, + ee_frame_idx=ee_frame_idx, + # the fixed root has no Jacobian row + ee_jacobi_idx=ee_frame_idx - 1, + arm_joint_ids=robot.find_joints(["Joint_.*"])[0], + ) + + +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 + + +def _compute_ee_vel_root(jacobian_b, joint_vel): + """Return the EE 6D velocity in the root frame as ``J · q_dot``. + + Required to make OSC's ``kd * ee_vel_b`` damping term meaningful: zero EE velocity leaves the impedance + undamped. ``J`` correctness is pinned independently by the articulation's gravity-compensation identity test. + """ + return torch.bmm(jacobian_b, joint_vel.unsqueeze(-1)).squeeze(-1) + + +def _build_relative_pose_target(robot, ee_frame_idx, delta_xyz): + """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) + + +def _make_osc() -> OperationalSpaceController: + return 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=NUM_ENVS, + device=_DEVICE, + ) + + +def _run_osc(chain: _Chain, osc: OperationalSpaceController, target_pose_b, num_steps: int, *, gravity: bool): + """Close the OSC loop for ``num_steps`` steps; return the per-step max position and rotation errors.""" + robot = chain.robot + pos_history: list[float] = [] + rot_history: list[float] = [] + for _ in range(num_steps): + jacobian_b = _compute_jacobian_root_frame(robot, chain.ee_jacobi_idx, chain.arm_joint_ids) + mass_matrix = robot.data.mass_matrix.torch[:, chain.arm_joint_ids, :][:, :, chain.arm_joint_ids] + gravity_forces = robot.data.gravity_compensation_forces.torch[:, chain.arm_joint_ids] if gravity else None + ee_pos_b, ee_quat_b, _ = _compute_ee_pose_root(robot, chain.ee_frame_idx) + ee_pose_b = torch.cat([ee_pos_b, ee_quat_b], dim=-1) + joint_vel = robot.data.joint_vel.torch[:, chain.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_forces, + ) + robot.actuators.target_command.set_effort_index(value=joint_efforts, joint_ids=chain.arm_joint_ids) + robot.write_data_to_sim() + chain.sim.step() + robot.update(chain.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()) + return pos_history, rot_history + + +def test_osc_tracking_accuracy(chain: _Chain) -> None: + """OSC pose tracking sentinel for the Jacobian and mass-matrix bridge. + + 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. + """ + chain.rest() + target_pose_b = _build_relative_pose_target(chain.robot, chain.ee_frame_idx, (0.05, 0.0, 0.0)) + pos_history, rot_history = _run_osc(chain, _make_osc(), target_pose_b, 300, gravity=False) + + 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 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?" + + +def test_osc_gravity_compensation_precision(chain: _Chain) -> None: + """Two-phase EE hold: gravity sag without compensation, tight hold with it. + + 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. + """ + chain.rest() + osc = _make_osc() + # Hold the initial EE pose: phase-1 steady-state error is pure gravity sag. + target_pose_b = _build_relative_pose_target(chain.robot, chain.ee_frame_idx, (0.0, 0.0, 0.0)) + + 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 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 + 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 + + with world_gravity((0.0, 0.0, -9.81)): + hist_off, _ = _run_osc(chain, osc, target_pose_b, 300, gravity=True) + osc.cfg.gravity_compensation = True + hist_on, _ = _run_osc(chain, osc, target_pose_b, 300, gravity=True) + + pos_off = _stationary_tail_mean(hist_off, "phase-1 sag") + pos_on = _stationary_tail_mean(hist_on, "phase-2 hold") + + 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)" From 9b3baec7d082e87562251b19f2b0ad8b251017d6 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 14:50:49 +0200 Subject: [PATCH 18/33] Sync with the split asset-test series Bring the tracking branch to the union of the 24 series branches: restore OVPhysX keepers for articulation cfg validation, USD gains, explicit and invalid root paths, default joint clamping, rigid COM kinematics and fused collection wrench delivery, and rename the changelog fragments per part. --- ...=> antoiner-asset-tests-p01-contract.skip} | 0 ...iner-asset-tests-p02-wrench-composer.skip} | 0 ...set-tests-p03-physx-actuator-runtime.skip} | 0 ...ntoiner-asset-tests-p04-terminations.skip} | 0 ...oiner-asset-tests-p18-newton-fixtures.skip | 0 ...sts-p19-newton-task-space-controllers.skip | 0 ...iner-asset-tests-p20-newton-actuators.skip | 0 ...r-asset-tests-p21-newton-rigid-object.skip | 0 ...ts-p22-newton-rigid-object-collection.skip | 0 ...et-tests-p23-newton-articulation-core.skip | 0 ...s-p24-newton-articulation-regressions.skip | 0 ...> antoiner-asset-tests-p05-ov-manager.rst} | 0 ...-asset-tests-p06-ov-articulation-core.skip | 0 ...tests-p07-ov-articulation-regressions.skip | 0 ...oiner-asset-tests-p08-ov-rigid-object.skip | 0 ...-tests-p09-ov-rigid-object-collection.skip | 0 ...ntoiner-asset-tests-p10-ov-deformable.skip | 0 .../test/assets/test_articulation.py | 123 +++++++++++++++++- .../test/assets/test_rigid_object.py | 84 +++++++++++- .../assets/test_rigid_object_collection.py | 32 ++++- ...set-tests-p11-physx-articulation-core.skip | 0 ...ts-p12-physx-articulation-regressions.skip | 0 ...sset-tests-p13-physx-newton-actuators.skip | 0 ...er-asset-tests-p14-physx-rigid-object.skip | 0 ...sts-p15-physx-rigid-object-collection.skip | 0 ...iner-asset-tests-p16-physx-deformable.skip | 0 ...asset-tests-p17-physx-surface-gripper.skip | 0 27 files changed, 235 insertions(+), 4 deletions(-) rename source/isaaclab/changelog.d/{antoiner-asset-tests-v2-physx.skip => antoiner-asset-tests-p01-contract.skip} (100%) rename source/isaaclab/changelog.d/{antoiner-asset-tests-v2-shared.skip => antoiner-asset-tests-p02-wrench-composer.skip} (100%) rename source/{isaaclab_newton/changelog.d/antoiner-asset-tests-v2-newton.skip => isaaclab/changelog.d/antoiner-asset-tests-p03-physx-actuator-runtime.skip} (100%) rename source/{isaaclab_physx/changelog.d/antoiner-asset-tests-v2-physx.skip => isaaclab/changelog.d/antoiner-asset-tests-p04-terminations.skip} (100%) create mode 100644 source/isaaclab_newton/changelog.d/antoiner-asset-tests-p18-newton-fixtures.skip create mode 100644 source/isaaclab_newton/changelog.d/antoiner-asset-tests-p19-newton-task-space-controllers.skip create mode 100644 source/isaaclab_newton/changelog.d/antoiner-asset-tests-p20-newton-actuators.skip create mode 100644 source/isaaclab_newton/changelog.d/antoiner-asset-tests-p21-newton-rigid-object.skip create mode 100644 source/isaaclab_newton/changelog.d/antoiner-asset-tests-p22-newton-rigid-object-collection.skip create mode 100644 source/isaaclab_newton/changelog.d/antoiner-asset-tests-p23-newton-articulation-core.skip create mode 100644 source/isaaclab_newton/changelog.d/antoiner-asset-tests-p24-newton-articulation-regressions.skip rename source/isaaclab_ov/changelog.d/{antoiner-asset-tests-v2-ov.rst => antoiner-asset-tests-p05-ov-manager.rst} (100%) create mode 100644 source/isaaclab_ov/changelog.d/antoiner-asset-tests-p06-ov-articulation-core.skip create mode 100644 source/isaaclab_ov/changelog.d/antoiner-asset-tests-p07-ov-articulation-regressions.skip create mode 100644 source/isaaclab_ov/changelog.d/antoiner-asset-tests-p08-ov-rigid-object.skip create mode 100644 source/isaaclab_ov/changelog.d/antoiner-asset-tests-p09-ov-rigid-object-collection.skip create mode 100644 source/isaaclab_ov/changelog.d/antoiner-asset-tests-p10-ov-deformable.skip create mode 100644 source/isaaclab_physx/changelog.d/antoiner-asset-tests-p11-physx-articulation-core.skip create mode 100644 source/isaaclab_physx/changelog.d/antoiner-asset-tests-p12-physx-articulation-regressions.skip create mode 100644 source/isaaclab_physx/changelog.d/antoiner-asset-tests-p13-physx-newton-actuators.skip create mode 100644 source/isaaclab_physx/changelog.d/antoiner-asset-tests-p14-physx-rigid-object.skip create mode 100644 source/isaaclab_physx/changelog.d/antoiner-asset-tests-p15-physx-rigid-object-collection.skip create mode 100644 source/isaaclab_physx/changelog.d/antoiner-asset-tests-p16-physx-deformable.skip create mode 100644 source/isaaclab_physx/changelog.d/antoiner-asset-tests-p17-physx-surface-gripper.skip diff --git a/source/isaaclab/changelog.d/antoiner-asset-tests-v2-physx.skip b/source/isaaclab/changelog.d/antoiner-asset-tests-p01-contract.skip similarity index 100% rename from source/isaaclab/changelog.d/antoiner-asset-tests-v2-physx.skip rename to source/isaaclab/changelog.d/antoiner-asset-tests-p01-contract.skip diff --git a/source/isaaclab/changelog.d/antoiner-asset-tests-v2-shared.skip b/source/isaaclab/changelog.d/antoiner-asset-tests-p02-wrench-composer.skip similarity index 100% rename from source/isaaclab/changelog.d/antoiner-asset-tests-v2-shared.skip rename to source/isaaclab/changelog.d/antoiner-asset-tests-p02-wrench-composer.skip diff --git a/source/isaaclab_newton/changelog.d/antoiner-asset-tests-v2-newton.skip b/source/isaaclab/changelog.d/antoiner-asset-tests-p03-physx-actuator-runtime.skip similarity index 100% rename from source/isaaclab_newton/changelog.d/antoiner-asset-tests-v2-newton.skip rename to source/isaaclab/changelog.d/antoiner-asset-tests-p03-physx-actuator-runtime.skip diff --git a/source/isaaclab_physx/changelog.d/antoiner-asset-tests-v2-physx.skip b/source/isaaclab/changelog.d/antoiner-asset-tests-p04-terminations.skip similarity index 100% rename from source/isaaclab_physx/changelog.d/antoiner-asset-tests-v2-physx.skip rename to source/isaaclab/changelog.d/antoiner-asset-tests-p04-terminations.skip diff --git a/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p18-newton-fixtures.skip b/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p18-newton-fixtures.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p19-newton-task-space-controllers.skip b/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p19-newton-task-space-controllers.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p20-newton-actuators.skip b/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p20-newton-actuators.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p21-newton-rigid-object.skip b/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p21-newton-rigid-object.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p22-newton-rigid-object-collection.skip b/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p22-newton-rigid-object-collection.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p23-newton-articulation-core.skip b/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p23-newton-articulation-core.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p24-newton-articulation-regressions.skip b/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p24-newton-articulation-regressions.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_ov/changelog.d/antoiner-asset-tests-v2-ov.rst b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p05-ov-manager.rst similarity index 100% rename from source/isaaclab_ov/changelog.d/antoiner-asset-tests-v2-ov.rst rename to source/isaaclab_ov/changelog.d/antoiner-asset-tests-p05-ov-manager.rst diff --git a/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p06-ov-articulation-core.skip b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p06-ov-articulation-core.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p07-ov-articulation-regressions.skip b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p07-ov-articulation-regressions.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p08-ov-rigid-object.skip b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p08-ov-rigid-object.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p09-ov-rigid-object-collection.skip b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p09-ov-rigid-object-collection.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p10-ov-deformable.skip b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p10-ov-deformable.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_ov/test/assets/test_articulation.py b/source/isaaclab_ov/test/assets/test_articulation.py index 6a259c3981d3..39c513ad0ae7 100644 --- a/source/isaaclab_ov/test/assets/test_articulation.py +++ b/source/isaaclab_ov/test/assets/test_articulation.py @@ -15,11 +15,13 @@ 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. +staging of CPU-only bindings and the native Newton actuator path. Initialization failures need their own +scenes, so those tests run first and build them. """ from __future__ import annotations +import math from collections.abc import Iterator from dataclasses import dataclass from pathlib import Path @@ -85,6 +87,7 @@ class _IslandCfg: reversed_elbow: bool = False collision_shapes: bool = False tendons: bool = False + articulation_root_prim_path: str | None = None devices: DeviceScope = DeviceScope.ALL @@ -158,6 +161,16 @@ class _IslandCfg: 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 + ), + "pos_limits": _IslandCfg(_implicit(), devices=DeviceScope.CUDA), "com_order": _IslandCfg( _implicit(), num_envs=1, body_ordering=_ROOT_PRESERVING_REVERSED_BODY_NAMES, distinct_coms=True ), @@ -255,11 +268,15 @@ def _spawn_island(name: str, cfg: _IslandCfg, y_offset: float) -> Articulation: 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: @@ -329,6 +346,51 @@ def _backend_to_user(articulation: Articulation) -> list[int]: return list(articulation.joint_ordering.backend_to_user_indices) +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, + ) + ) + + +@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, init_state, match): + """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): + """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.""" @@ -658,6 +720,65 @@ def test_com_orientation_write_invalidates_static_inertia_cache_with_body_orderi ) +@pytest.mark.parametrize("scene", _CUDA_DEVICES, 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["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_DEVICES, 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["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("scene", _CUDA_DEVICES, indirect=True) +def test_joint_position_limits_clamp_default_joint_positions(scene: _ArticulationScene) -> None: + """Write partial joint position limits and clamp only the selected default joint positions into them.""" + articulation = scene["pos_limits"] + device = scene.device + default_before = articulation.data.default_joint_pos.torch.clone() + env_ids = torch.tensor([1], dtype=torch.int32, device=device) + joint_ids = torch.tensor([0, 2], dtype=torch.int32, device=device) + # The default positions are zero, so both new ranges exclude them. + limits = torch.tensor([[[0.2, 0.3], [0.25, 0.35]]], device=device) + articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) + + torch.testing.assert_close(articulation.data.joint_pos_limits.torch[env_ids][:, joint_ids], limits) + default_joint_pos = articulation.data.default_joint_pos.torch + selected = default_joint_pos[env_ids][:, joint_ids] + assert torch.all((selected >= limits[..., 0]) & (selected <= 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], default_before[unselected]) + + @pytest.mark.parametrize("scene", _CUDA_DEVICES, 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.""" diff --git a/source/isaaclab_ov/test/assets/test_rigid_object.py b/source/isaaclab_ov/test/assets/test_rigid_object.py index d58bcaa56c7e..e7cae9140158 100644 --- a/source/isaaclab_ov/test/assets/test_rigid_object.py +++ b/source/isaaclab_ov/test/assets/test_rigid_object.py @@ -34,6 +34,7 @@ from isaaclab.assets import RigidObjectCfg # 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.math import quat_apply_inverse, quat_mul # noqa: E402 pytestmark = pytest.mark.integration @@ -114,6 +115,7 @@ class _RigidScene: dynamic: RigidObject kinematic: RigidObject resettable: RigidObject + spinning: RigidObject @pytest.fixture(scope="module") @@ -124,8 +126,11 @@ def scene(request: pytest.FixtureRequest) -> Iterator[_RigidScene]: dynamic = _spawn_cubes("Dynamic", y_offset=0.0, disable_gravity=True) kinematic = _spawn_cubes("Kinematic", y_offset=2.0, kinematic_enabled=True) resettable = _spawn_cubes("Resettable", y_offset=4.0, disable_gravity=True) + spinning = _spawn_cubes("Spinning", y_offset=6.0, disable_gravity=True) sim.reset() - yield _RigidScene(sim=sim, device=device, dynamic=dynamic, kinematic=kinematic, resettable=resettable) + yield _RigidScene( + sim=sim, device=device, dynamic=dynamic, kinematic=kinematic, resettable=resettable, spinning=spinning + ) @pytest.mark.parametrize("scene", test_devices(), indirect=True) @@ -193,7 +198,8 @@ def test_initialization_with_kinematic_enabled(scene: _RigidScene) -> None: # SDP bindings must be ready on CPU and GPU before any asset pose read. provider = sim.get_scene_data_provider() - expected_paths = {f"/World/{name}/Env_{i}/Cube" for name in ("Dynamic", "Kinematic", "Resettable") for i in (0, 1)} + names = ("Dynamic", "Kinematic", "Resettable", "Spinning") + expected_paths = {f"/World/{name}/Env_{i}/Cube" for name in names for i in (0, 1)} assert provider.transform_count == len(expected_paths) assert set(provider.backend.transform_paths) == expected_paths @@ -235,3 +241,77 @@ def test_reset_clears_active_wrench_composers(scene: _RigidScene) -> None: 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("scene", test_devices(DeviceScope.CUDA), 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() + + # 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.unsqueeze(1) + cube_object.set_coms_index(coms=wp.from_torch(com.contiguous(), dtype=wp.transformf)) + torch.testing.assert_close(cube_object.data.body_com_pose_b.torch, com) + + # z spin velocity + spin_twist = torch.zeros(6, device=device) + spin_twist[5] = 2.0 + + for _ in range(10): + # spin the object around Z axis (com) + cube_object.write_root_velocity_to_sim_index(root_velocity=spin_twist.repeat(_NUM_CUBES, 1)) + sim.step() + 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:]) diff --git a/source/isaaclab_ov/test/assets/test_rigid_object_collection.py b/source/isaaclab_ov/test/assets/test_rigid_object_collection.py index 1acfe0d5b32a..35c6eafa90ba 100644 --- a/source/isaaclab_ov/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_ov/test/assets/test_rigid_object_collection.py @@ -94,6 +94,7 @@ class _CollectionScene: fused: RigidObjectCollection resettable: RigidObjectCollection stepped: RigidObjectCollection + wrenched: RigidObjectCollection @pytest.fixture(scope="module") @@ -104,8 +105,11 @@ def scene(request: pytest.FixtureRequest) -> Iterator[_CollectionScene]: fused = _spawn_collection("Fused", y_offset=0.0) resettable = _spawn_collection("Resettable", y_offset=4.0) stepped = _spawn_collection("Stepped", y_offset=8.0) + wrenched = _spawn_collection("Wrenched", y_offset=12.0) sim.reset() - yield _CollectionScene(sim=sim, device=device, fused=fused, resettable=resettable, stepped=stepped) + yield _CollectionScene( + sim=sim, device=device, fused=fused, resettable=resettable, stepped=stepped, wrenched=wrenched + ) @pytest.mark.parametrize("scene", test_devices(), indirect=True) @@ -207,6 +211,32 @@ def test_reset_clears_active_wrench_composers(scene: _CollectionScene) -> None: assert torch.count_nonzero(object_collection._permanent_wrench_composer.composed_torque.torch) == 0 +@pytest.mark.parametrize("scene", test_devices(DeviceScope.CUDA), indirect=True) +def test_wrench_reaches_only_the_selected_collection_body(scene: _CollectionScene) -> None: + """Deliver an external force through the fused body-major binding to the selected body only. + + 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) + collection.reset() + + 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) + + @pytest.mark.parametrize("scene", test_devices(DeviceScope.CUDA), indirect=True) @pytest.mark.xfail( strict=True, diff --git a/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p11-physx-articulation-core.skip b/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p11-physx-articulation-core.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p12-physx-articulation-regressions.skip b/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p12-physx-articulation-regressions.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p13-physx-newton-actuators.skip b/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p13-physx-newton-actuators.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p14-physx-rigid-object.skip b/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p14-physx-rigid-object.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p15-physx-rigid-object-collection.skip b/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p15-physx-rigid-object-collection.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p16-physx-deformable.skip b/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p16-physx-deformable.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p17-physx-surface-gripper.skip b/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p17-physx-surface-gripper.skip new file mode 100644 index 000000000000..e69de29bb2d1 From 7ee61847a62a8aef80bcc85cf7e960cc619e8fcc Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 15:22:47 +0200 Subject: [PATCH 19/33] Sync with the style-aligned asset-test series Bring the tracking branch to the union of the 24 series branches after aligning them with the coding and testing guidelines in docs/source/refs/contributing.rst. --- .../actuators/test_physx_actuator_runtime.py | 23 +- .../contract/_articulation_contract_cases.py | 77 ++--- .../contract/_articulation_contract_utils.py | 2 +- .../test/assets/contract/_contract_boot.py | 9 +- .../test/assets/contract/_mock_backends.py | 3 +- ..._rigid_object_collection_contract_cases.py | 15 +- .../contract/_rigid_object_contract_cases.py | 12 +- .../contract/_rigid_object_contract_utils.py | 2 +- .../isaaclab/test/assets/contract/conftest.py | 2 +- .../test/assets/contract/public_surface.py | 8 +- .../contract/test_asset_contract_api.py | 2 +- .../test/envs/mdp/test_terminations.py | 4 +- .../test/utils/test_wrench_composer.py | 37 ++- .../utils/test_wrench_composer_integration.py | 9 +- .../test/articulation_test_utils.py | 309 +++++++++--------- .../test/assets/test_articulation.py | 135 ++++---- .../assets/test_newton_actuators_newton.py | 16 +- .../test/assets/test_rigid_object.py | 44 ++- .../assets/test_rigid_object_collection.py | 40 ++- .../test_newton_task_space_controllers.py | 58 ++-- .../antoiner-asset-tests-p05-ov-manager.rst | 8 +- .../isaaclab_ov/physics/ovphysx_manager.py | 30 +- .../test/assets/test_articulation.py | 44 +-- .../test/assets/test_rigid_object.py | 32 +- .../assets/test_rigid_object_collection.py | 5 +- .../unit/test_ovphysx_actuator_control.py | 3 +- .../physics/test_ovphysx_manager_lifecycle.py | 7 +- .../test/assets/test_articulation.py | 152 +++++---- .../test/assets/test_deformable_object.py | 48 +-- .../assets/test_newton_actuators_physx.py | 4 - .../test/assets/test_rigid_object.py | 35 +- .../assets/test_rigid_object_collection.py | 25 +- .../test/assets/test_surface_gripper.py | 1 - .../assets/unit/test_physx_surface_gripper.py | 10 +- 34 files changed, 612 insertions(+), 599 deletions(-) diff --git a/source/isaaclab/test/actuators/test_physx_actuator_runtime.py b/source/isaaclab/test/actuators/test_physx_actuator_runtime.py index 6f9a144cbcde..419a65b363f4 100644 --- a/source/isaaclab/test/actuators/test_physx_actuator_runtime.py +++ b/source/isaaclab/test/actuators/test_physx_actuator_runtime.py @@ -19,18 +19,13 @@ def _runtime() -> PhysxActuatorRuntime: return PhysxActuatorRuntime(SimpleNamespace(device="cuda:0"), logger=Mock()) -def _cuda_device(*, is_capturing: bool) -> SimpleNamespace: - """Describe a CUDA device with the given capture state.""" - return SimpleNamespace(is_cuda=True, is_capturing=is_capturing) - - -def test_graph_capture_failure_restores_adapter_state_and_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch): +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.""" class _FailingCapture: capture_count = 0 - def __init__(self, *args, **kwargs): + def __init__(self, *args, **kwargs) -> None: self.capture_index = type(self).capture_count type(self).capture_count += 1 self.graph = object() @@ -40,7 +35,7 @@ def __enter__(self): raise RuntimeError("second capture unavailable") return self - def __exit__(self, exc_type, exc_value, traceback): + def __exit__(self, exc_type, exc_value, traceback) -> bool: return False state_a, state_b = object(), object() @@ -61,23 +56,23 @@ def _swap_adapter_state(*args, **kwargs) -> None: runtime._logger.warning.assert_called_once() -def test_compute_runs_eagerly_after_graph_capture_failure(monkeypatch: pytest.MonkeyPatch): +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.""" class _FailingCapture: - def __init__(self, *args, **kwargs): + def __init__(self, *args, **kwargs) -> None: pass def __enter__(self): raise RuntimeError("capture unavailable") - def __exit__(self, exc_type, exc_value, traceback): + def __exit__(self, exc_type, exc_value, traceback) -> bool: return False 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: _cuda_device(is_capturing=False)) + monkeypatch.setattr(wp, "get_device", lambda device: SimpleNamespace(is_cuda=True, is_capturing=False)) monkeypatch.setattr(wp, "ScopedCapture", _FailingCapture) monkeypatch.setattr(runtime, "_run_native_actuator_kernels", eager_compute) collection = SimpleNamespace() @@ -91,11 +86,11 @@ def __exit__(self, exc_type, exc_value, traceback): eager_compute.assert_called_with(collection, 0.01) -def test_stateful_actuator_rejects_outer_cuda_capture(monkeypatch: pytest.MonkeyPatch): +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: _cuda_device(is_capturing=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/contract/_articulation_contract_cases.py b/source/isaaclab/test/assets/contract/_articulation_contract_cases.py index d1a0a521f322..2ac86859e6bc 100644 --- a/source/isaaclab/test/assets/contract/_articulation_contract_cases.py +++ b/source/isaaclab/test/assets/contract/_articulation_contract_cases.py @@ -14,6 +14,7 @@ import math import warnings +from contextlib import AbstractContextManager, nullcontext from unittest.mock import MagicMock, patch import numpy as np @@ -175,33 +176,28 @@ def test_finder_returns_legacy_list_or_cached_proxy(self, backend, finder_name, assert finder(first_name) == ([0], [first_name]) +def _fixed_tendon_articulation(backend: str): + """Create a two-environment articulation with two fixed tendons and no spatial tendons.""" + 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 + + 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 - @contract_backend("fixed_tendon_target_scheduling") 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.""" - art = self._articulation(backend) + art = _fixed_tendon_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 @requires_backend("newton") def test_newton_reports_a_missing_tendon_actuator(self): """Without a MuJoCo tendon actuator the solver has no adapter, so commanding must say so.""" - art = self._articulation("newton") + art = _fixed_tendon_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)) @@ -210,7 +206,7 @@ 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.""" from isaaclab_newton.physics import NewtonManager - art = self._articulation("newton") + art = _fixed_tendon_articulation("newton") with pytest.raises(NotImplementedError, match="does not drive fixed tendons"): NewtonManager.create_fixed_tendon_control(art) @@ -997,9 +993,8 @@ def test_partial_joint_state_write_follows_int64_selectors_and_restarts_accelera 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. - caches = ("_body_link_vel_w",) if backend == "newton" else ("_body_link_pose_w", "_body_com_vel_w") - for cache in caches: - getattr(art.data, cache[1:]) + cached_properties = ("body_link_vel_w",) if backend == "newton" else ("body_link_pose_w", "body_com_vel_w") + caches = _prime_timestamped_properties(art.data, [(name, 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) @@ -1037,11 +1032,10 @@ def test_partial_joint_state_write_follows_int64_selectors_and_restarts_accelera 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 cache in caches: - assert getattr(art.data, cache).timestamp == art.data._sim_timestamp, cache + 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]) - for cache in caches: - assert getattr(art.data, cache).timestamp < art.data._sim_timestamp, cache + _assert_buffers_stale(art.data, caches) @_writes_backends def test_deprecated_joint_state_writer_matches_public_writers(self, backend): @@ -1314,14 +1308,18 @@ def test_body_writer_mask(self, backend, device, method_base, kwarg): 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 == "physx": - getter = {"root_pose": "get_root_transforms", "root_velocity": "get_root_velocities", "mass": "get_masses"} - values = wp.to_torch(getattr(raw_backend, getter[quantity])()) + getter = { + "root_pose": raw_backend.get_root_transforms, + "root_velocity": raw_backend.get_root_velocities, + "mass": raw_backend.get_masses, + } + values = wp.to_torch(getter[quantity]()) elif backend == "newton": if quantity == "mass": values = wp.to_torch(art.data._sim_bind_body_mass) else: - getter = {"root_pose": "get_root_transforms", "root_velocity": "get_root_velocities"}[quantity] - values = wp.to_torch(getattr(raw_backend, getter)(None)) + getter = {"root_pose": raw_backend.get_root_transforms, "root_velocity": raw_backend.get_root_velocities} + values = wp.to_torch(getter[quantity](None)) else: from isaaclab_ov import tensor_types as TT @@ -1470,15 +1468,13 @@ def test_articulation_tendon_contract( ] -def _push_tendon_properties(art, backend: str, kind: str) -> None: - """Write staged tendon properties to the backend so data read-backs see them on every backend.""" - if backend == "newton": - from isaaclab_newton.physics import NewtonManager +def _suppress_newton_model_changes(backend: str) -> AbstractContextManager: + """Return a context that drops Newton model-change notifications, which need a solver the mock lacks.""" + if backend != "newton": + return nullcontext() + from isaaclab_newton.physics import NewtonManager - with patch.object(NewtonManager, "add_model_change"): - getattr(art, f"write_{kind}_tendon_properties_to_sim_index")() - else: - getattr(art, f"write_{kind}_tendon_properties_to_sim_index")() + return patch.object(NewtonManager, "add_model_change") class TestArticulationWritersTendon: @@ -1544,7 +1540,8 @@ def test_tendon_writer( method(**{kwarg: _make_data_torch(subset_shape, device, wp_dtype)}, **subset) # warp, all envs + all tendons: after the push, the matching data property reads the written values back method(**{kwarg: _make_payload_warp(full_shape, device, wp_dtype)}) - _push_tendon_properties(art, backend, kind) + with _suppress_newton_model_changes(backend): + getattr(art, f"write_{kind}_tendon_properties_to_sim_index")() _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) @@ -1610,16 +1607,12 @@ def capture(name, values, *, indices=None, 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_kwargs = {"env_mask": mask} + with _suppress_newton_model_changes(backend): + write(**write_kwargs) if backend == "newton": - from isaaclab_newton.physics import NewtonManager - - with patch.object(NewtonManager, "add_model_change"): - write(**write_kwargs) 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])] ) - else: - write(**write_kwargs) 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/contract/_articulation_contract_utils.py b/source/isaaclab/test/assets/contract/_articulation_contract_utils.py index ccaf1fb211e1..c940a1e18b1b 100644 --- a/source/isaaclab/test/assets/contract/_articulation_contract_utils.py +++ b/source/isaaclab/test/assets/contract/_articulation_contract_utils.py @@ -116,7 +116,7 @@ def create_physx_articulation( # 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 (the PhysX manager is patched for this test above) + # The data reads gravity from the PhysX manager patched above. data = PhysXArticulationData(mock_view, device) object.__setattr__(articulation, "_data", data) diff --git a/source/isaaclab/test/assets/contract/_contract_boot.py b/source/isaaclab/test/assets/contract/_contract_boot.py index cdf33bfd3a54..6d62e38c8379 100644 --- a/source/isaaclab/test/assets/contract/_contract_boot.py +++ b/source/isaaclab/test/assets/contract/_contract_boot.py @@ -20,10 +20,10 @@ launch_test_simulation() else: - def _install_stub(module_name: str, *, is_package: bool = False) -> MagicMock: + def _install_stub(module_name: str, *, is_package: bool = False) -> None: """Install one import-boundary module double unless the real module is already loaded.""" if module_name in sys.modules: - return sys.modules[module_name] + return stub = MagicMock() stub.__spec__ = ModuleSpec(module_name, loader=None, is_package=is_package) if is_package: @@ -32,10 +32,9 @@ def _install_stub(module_name: str, *, is_package: bool = False) -> MagicMock: if "." in module_name: parent_name, attribute = module_name.rsplit(".", 1) setattr(sys.modules[parent_name], attribute, stub) - return stub - # Worktree installs include the real ``isaaclab_physx`` package but not Kit's Python runtime. Stub only that - # external boundary so the contracts still import the real PhysX asset and data classes and their fixture views. + # Without Kit's Python runtime, ``carb`` and ``usdrt`` are missing even when ``isaaclab_physx`` is installed. Stub + # only that runtime boundary so the contracts still import the real PhysX asset and data classes and their views. _install_stub("carb") _install_stub("usdrt", is_package=True) diff --git a/source/isaaclab/test/assets/contract/_mock_backends.py b/source/isaaclab/test/assets/contract/_mock_backends.py index 77c968ddfa7c..32c60dd1ae8a 100644 --- a/source/isaaclab/test/assets/contract/_mock_backends.py +++ b/source/isaaclab/test/assets/contract/_mock_backends.py @@ -5,6 +5,7 @@ """Backend-manager and tensor-view doubles shared by the mocked contract factories.""" +from collections.abc import Callable from unittest.mock import MagicMock import numpy as np @@ -44,7 +45,7 @@ def install_physx_recording_setters(mock_view, storage_by_method: dict[str, str] storage_by_method: Setter name to the view attribute that stores its values. """ - def make_setter(storage_name: str): + 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) diff --git a/source/isaaclab/test/assets/contract/_rigid_object_collection_contract_cases.py b/source/isaaclab/test/assets/contract/_rigid_object_collection_contract_cases.py index 210f2f43858f..20a956e79440 100644 --- a/source/isaaclab/test/assets/contract/_rigid_object_collection_contract_cases.py +++ b/source/isaaclab/test/assets/contract/_rigid_object_collection_contract_cases.py @@ -838,18 +838,23 @@ def test_aliases_match_canonical_values(self, backend): } -def _read_backend_cells(backend: str, raw_backend, quantity: str, num_instances: int, num_bodies: int): +def _read_backend_cells(backend: str, raw_backend, quantity: str, num_instances: int, num_bodies: int) -> torch.Tensor: """Read one collection quantity from backend storage in public ``(env, body, ...)`` order.""" if backend == "physx": - getter = {"body_pose": "get_transforms", "body_velocity": "get_velocities", "mass": "get_masses"}[quantity] + getter = { + "body_pose": raw_backend.get_transforms, + "body_velocity": raw_backend.get_velocities, + "mass": raw_backend.get_masses, + } + values = wp.to_torch(getter[quantity]()) # PhysX collection views are body-major. - values = wp.to_torch(getattr(raw_backend, getter)()).reshape(num_bodies, num_instances, -1).transpose(0, 1) + values = values.reshape(num_bodies, num_instances, -1).transpose(0, 1) elif backend == "newton": if quantity == "mass": values = wp.to_torch(raw_backend.get_attribute("body_mass", None)) else: - getter = {"body_pose": "get_root_transforms", "body_velocity": "get_root_velocities"}[quantity] - values = wp.to_torch(getattr(raw_backend, getter)(None)) + getter = {"body_pose": raw_backend.get_root_transforms, "body_velocity": raw_backend.get_root_velocities} + values = wp.to_torch(getter[quantity](None)) else: from isaaclab_ov import tensor_types as TT diff --git a/source/isaaclab/test/assets/contract/_rigid_object_contract_cases.py b/source/isaaclab/test/assets/contract/_rigid_object_contract_cases.py index 223a4224f039..92e877b1bbd6 100644 --- a/source/isaaclab/test/assets/contract/_rigid_object_contract_cases.py +++ b/source/isaaclab/test/assets/contract/_rigid_object_contract_cases.py @@ -791,14 +791,18 @@ def test_body_writer_mask(self, backend, device, method_base, kwarg): def _read_backend_rows(backend: str, raw_backend, quantity: str) -> torch.Tensor: """Read one rigid-object quantity from backend storage as one row per environment.""" if backend == "physx": - getter = {"root_pose": "get_transforms", "root_velocity": "get_velocities", "mass": "get_masses"}[quantity] - values = getattr(raw_backend, getter)() + getter = { + "root_pose": raw_backend.get_transforms, + "root_velocity": raw_backend.get_velocities, + "mass": raw_backend.get_masses, + } + values = getter[quantity]() elif backend == "newton": if quantity == "mass": values = raw_backend.get_attribute("body_mass", None) else: - getter = {"root_pose": "get_root_transforms", "root_velocity": "get_root_velocities"}[quantity] - values = getattr(raw_backend, getter)(None) + getter = {"root_pose": raw_backend.get_root_transforms, "root_velocity": raw_backend.get_root_velocities} + values = getter[quantity](None) else: from isaaclab_ov import tensor_types as TT diff --git a/source/isaaclab/test/assets/contract/_rigid_object_contract_utils.py b/source/isaaclab/test/assets/contract/_rigid_object_contract_utils.py index 836e361973a8..482ec3f67b03 100644 --- a/source/isaaclab/test/assets/contract/_rigid_object_contract_utils.py +++ b/source/isaaclab/test/assets/contract/_rigid_object_contract_utils.py @@ -69,7 +69,7 @@ def create_physx_rigid_object( object.__setattr__(rigid_object, "_root_view", mock_view) object.__setattr__(rigid_object, "_device", device) - # Create RigidObjectData instance (the PhysX manager is patched for this test above) + # The data reads gravity from the PhysX manager patched above. data = PhysXRigidObjectData(mock_view, device) object.__setattr__(rigid_object, "_data", data) diff --git a/source/isaaclab/test/assets/contract/conftest.py b/source/isaaclab/test/assets/contract/conftest.py index b8dc619426b1..11701c696da6 100644 --- a/source/isaaclab/test/assets/contract/conftest.py +++ b/source/isaaclab/test/assets/contract/conftest.py @@ -3,7 +3,7 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Fixtures shared by the focused asset contract gate.""" +"""Fixtures shared by the asset contract tests.""" from collections.abc import Iterator diff --git a/source/isaaclab/test/assets/contract/public_surface.py b/source/isaaclab/test/assets/contract/public_surface.py index acf4d0d3e773..04ac1c919bdf 100644 --- a/source/isaaclab/test/assets/contract/public_surface.py +++ b/source/isaaclab/test/assets/contract/public_surface.py @@ -50,7 +50,9 @@ class PublicMemberContract: ) -def _mapping(class_name: str, members: str, kind: ContractKind, reason: str | None = None): +def _mapping( + class_name: str, members: str, kind: ContractKind, reason: str | None = None +) -> dict[str, PublicMemberContract]: """Classify a whitespace-separated group of members declared by one class.""" return {f"{class_name}.{member}": PublicMemberContract(kind, reason) for member in members.split()} @@ -240,7 +242,9 @@ def _mapping(class_name: str, members: str, kind: ContractKind, reason: str | No } -def public_surface_mismatches(classes: tuple[type, ...], contracts: dict[str, PublicMemberContract]) -> dict[str, list]: +def public_surface_mismatches( + classes: tuple[type, ...], contracts: dict[str, PublicMemberContract] +) -> dict[str, list[str]]: """Diff the public members the classes declare against their contract classifications. Args: diff --git a/source/isaaclab/test/assets/contract/test_asset_contract_api.py b/source/isaaclab/test/assets/contract/test_asset_contract_api.py index 7eaf7ef12bba..a51c348f0d62 100644 --- a/source/isaaclab/test/assets/contract/test_asset_contract_api.py +++ b/source/isaaclab/test/assets/contract/test_asset_contract_api.py @@ -82,7 +82,7 @@ def _manager_bindings() -> dict[str, object]: def test_contract_factories_scope_manager_patches_to_one_test(reverse: bool) -> None: """Factories patch backend managers only inside a test scope and restore the production bindings after it.""" original = _manager_bindings() - # Collection alone must not patch a manager: the production bindings are in place before any factory runs. + # Importing the contract modules must not patch a manager: production bindings precede every factory call. if "physx" in available_backends(): assert isinstance(original["physx.get_physics_sim_view"], classmethod) if "newton" in available_backends(): diff --git a/source/isaaclab/test/envs/mdp/test_terminations.py b/source/isaaclab/test/envs/mdp/test_terminations.py index e885a2d029f8..681f613c6bde 100644 --- a/source/isaaclab/test/envs/mdp/test_terminations.py +++ b/source/isaaclab/test/envs/mdp/test_terminations.py @@ -13,9 +13,9 @@ from isaaclab.managers import SceneEntityCfg -def test_joint_effort_limit_terminates_only_environments_with_clipped_selected_joints(): +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], [10.0, 20.0, 30.0], [10.0, 20.0, 30.0], [10.0, 20.0, 30.0]]) + computed = torch.tensor([10.0, 20.0, 30.0]).repeat(4, 1) # Environment 0 clips the selected joint, environment 1 clips nothing, environment 2 clips the selected joint # downward, and environment 3 clips only joints that are not selected. applied = torch.tensor([[10.0, 19.0, 30.0], [10.0, 20.0, 30.0], [10.0, 17.0, 30.0], [9.0, 20.0, 29.0]]) diff --git a/source/isaaclab/test/utils/test_wrench_composer.py b/source/isaaclab/test/utils/test_wrench_composer.py index c688d4d2574f..eeaa0822a848 100644 --- a/source/isaaclab/test/utils/test_wrench_composer.py +++ b/source/isaaclab/test/utils/test_wrench_composer.py @@ -35,6 +35,8 @@ "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(): @@ -62,11 +64,13 @@ def __init__(self, com_pos_w: torch.Tensor, link_quat_w: torch.Tensor) -> None: self.set_pose(com_pos_w, link_quat_w) 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.contiguous(), dtype=wp.vec3f)) - self.body_link_quat_w = ProxyArray(wp.from_torch(link_quat_w.contiguous(), dtype=wp.quatf)) + 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)) -def _poses(com: dict | None = None, quat: dict | None = None, shape: tuple[int, int] = (2, 2)): +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) @@ -78,14 +82,14 @@ def _poses(com: dict | None = None, quat: dict | None = None, shape: tuple[int, def _make_composer( - com: dict | None = None, quat: dict | None = None, *, supports_world_at_com: bool = False + 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) -def _grid(cells: dict, shape: tuple[int, int] = (2, 2)) -> torch.Tensor: +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(): @@ -94,13 +98,13 @@ def _grid(cells: dict, shape: tuple[int, int] = (2, 2)) -> torch.Tensor: def _vectors(values: torch.Tensor) -> wp.array: - """Convert a ``(..., 3)`` tensor to a Warp vector array.""" - return wp.from_torch(values.contiguous(), dtype=wp.vec3f) + """Convert a contiguous ``(..., 3)`` tensor to a Warp vector array.""" + return wp.from_torch(values, dtype=wp.vec3f) def _mask(values: list[bool]) -> wp.array: - """Convert literal booleans to a Warp mask array with owned storage.""" - return wp.array(np.array(values, dtype=bool), dtype=wp.bool, device="cpu") + """Convert literal booleans to a CPU Warp mask array.""" + return wp.array(values, dtype=wp.bool, device="cpu") def _apply_to_cell(composer: WrenchComposer, method: str, cell: tuple[int, int], is_global: bool, **vectors) -> None: @@ -293,16 +297,15 @@ def test_index_dtype_combinations_preserve_selected_wrench_cells( def _fill_all_input_buffers(composer: WrenchComposer) -> None: """Give every cell a distinct contribution in each of the five input buffers.""" - ones = torch.ones((2, 2, 3)) - composer.add_forces_and_torques_index( - forces=_vectors(ones * torch.tensor([0.0, 0.0, 1.0])), - positions=_vectors(ones * torch.tensor([0.0, 1.0, 0.0])), - is_global=True, - ) - composer.add_forces_and_torques_index(forces=_vectors(ones * torch.tensor([1.0, 0.0, 0.0])), is_global=True) + + def every_cell(value: tuple[float, float, float]) -> wp.array: + return _vectors(torch.tensor(value).repeat(2, 2, 1)) + composer.add_forces_and_torques_index( - forces=_vectors(ones * torch.tensor([0.0, 1.0, 0.0])), torques=_vectors(ones * torch.tensor([0.0, 0.0, 1.0])) + 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))) # buffer -> per-cell value written by ``_fill_all_input_buffers`` diff --git a/source/isaaclab/test/utils/test_wrench_composer_integration.py b/source/isaaclab/test/utils/test_wrench_composer_integration.py index d0456e00a837..4c49ffd879ce 100644 --- a/source/isaaclab/test/utils/test_wrench_composer_integration.py +++ b/source/isaaclab/test/utils/test_wrench_composer_integration.py @@ -80,13 +80,16 @@ def test_composed_wrenches_reach_physx_in_the_frame_they_are_delivered(device: s ) 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=wp.from_torch(force.view(-1, 3).contiguous(), dtype=wp.float32), - torque_data=wp.from_torch(torch.zeros_like(force).view(-1, 3).contiguous(), dtype=wp.float32), - position_data=wp.from_torch(raw_position.view(-1, 3).contiguous(), dtype=wp.float32), + force_data=raw_force, + torque_data=raw_torque, + position_data=raw_position_data, indices=raw._ALL_INDICES, is_global=True, ) diff --git a/source/isaaclab_newton/test/articulation_test_utils.py b/source/isaaclab_newton/test/articulation_test_utils.py index 173740041c70..6255acc80832 100644 --- a/source/isaaclab_newton/test/articulation_test_utils.py +++ b/source/isaaclab_newton/test/articulation_test_utils.py @@ -21,7 +21,6 @@ from pathlib import Path from typing import Any -import numpy as np import torch import warp as wp from isaaclab_newton.assets import Articulation, RigidObject, RigidObjectCollection @@ -34,17 +33,140 @@ from isaaclab.cloner import CloneCfg, clone_plan_from_env_0, replicate from isaaclab.sim import SimulationCfg, SimulationContext -DATA_DIR = Path(__file__).resolve().parent / "assets" / "data" -"""Directory holding the checked-in USD fixtures.""" - NUM_ENVS = 2 """Environments per scene: enough to prove a partial write leaves the other environment untouched.""" -ENV_SPACING = 40.0 -"""Distance between environment origins along x [m].""" +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 / "assets" / "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: Authored fixture name, or the name of a checked-in file in ``test/assets/data``. + **kwargs: Additional :class:`~isaaclab.sim.UsdFileCfg` fields. + + Returns: + The spawner configuration. + """ + if name in _FIXTURES: + path = _fixture_dir() / name + if not path.exists(): + path.write_text(_FIXTURES[name](), encoding="utf-8") + else: + path = _DATA_DIR / name + return sim_utils.UsdFileCfg(usd_path=str(path), **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) + + ## # Fixture authoring. ## @@ -57,30 +179,26 @@ def _body( inertia: Vec3, com: Vec3 | None = None, translate: Vec3 | None = None, - prim_type: str = "Xform", - schemas: Sequence[str] = (), - body: str = "", + articulation_root: bool = False, + collider_size: float | None = None, ) -> str: - """Return a rigid body with explicit mass properties.""" - applied = ", ".join(f'"{schema}"' for schema in ("PhysicsRigidBodyAPI", "PhysicsMassAPI", *schemas)) + """Return a rigid body with explicit mass properties and, optionally, a child collider cube [m].""" + applied = '"PhysicsRigidBodyAPI", "PhysicsMassAPI"' + if articulation_root: + applied += ', "PhysicsArticulationRootAPI"' lines = [] if translate is not None: lines += [f"double3 xformOp:translate = {translate}", 'uniform token[] xformOpOrder = ["xformOp:translate"]'] lines += [f"float physics:mass = {mass}", f"float3 physics:diagonalInertia = {inertia}"] if com is not None: lines.append(f"point3f physics:centerOfMass = {com}") - content = "\n".join(f" {line}" for line in lines) + body - return ( - f' def {prim_type} "{name}" (\n prepend apiSchemas = [{applied}]\n )\n {{\n{content}\n }}\n' - ) - - -def _collision_cube(size: float) -> str: - """Return a collider cube authored as a child of a rigid body.""" - return ( - f'\n\n def Cube "Collision" (\n prepend apiSchemas = ["PhysicsCollisionAPI"]\n )\n' - f" {{\n double size = {size}\n }}" - ) + content = "\n".join(f" {line}" for line in lines) + if collider_size is not None: + content += ( + f'\n\n def Cube "Collision" (\n prepend apiSchemas = ["PhysicsCollisionAPI"]\n )\n' + f" {{\n double size = {collider_size}\n }}" + ) + return f' def Xform "{name}" (\n prepend apiSchemas = [{applied}]\n )\n {{\n{content}\n }}\n' def _joint( @@ -122,11 +240,13 @@ def _joint( return f' def Physics{kind}Joint "{name}"{schemas}\n {{\n{content}\n }}\n' -def _document(root: str, prims: Sequence[str], *, root_schemas: Sequence[str] = ("PhysicsArticulationRootAPI",)) -> str: - """Return a USD document whose default prim is an ``Xform`` holding ``prims``.""" +def _document(root: str, prims: Sequence[str], *, articulation_root: bool = True) -> str: + """Return a USD document whose default prim is an ``Xform`` holding ``prims``. + + ``articulation_root=False`` leaves the articulation root to one of the bodies. + """ header = f'#usda 1.0\n(\n defaultPrim = "{root}"\n metersPerUnit = 1\n upAxis = "Z"\n)\n\n' - applied = ", ".join(f'"{schema}"' for schema in root_schemas) - metadata = f" (\n prepend apiSchemas = [{applied}]\n)" if root_schemas else "" + metadata = ' (\n prepend apiSchemas = ["PhysicsArticulationRootAPI"]\n)' if articulation_root else "" return header + f'def Xform "{root}"{metadata}\n{{\n' + "\n".join(prims) + "}\n" @@ -140,8 +260,8 @@ def _floating_two_link() -> str: mass=1.0, inertia=(0.01, 0.02, 0.02), com=(0.0, 0.0, 0.0), - schemas=("PhysicsArticulationRootAPI",), - body=_collision_cube(0.1), + articulation_root=True, + collider_size=0.1, ), _body( "Child", @@ -149,7 +269,7 @@ def _floating_two_link() -> str: inertia=(0.01, 0.02, 0.02), com=(0.0, 0.0, 0.0), translate=(0.5, 0.0, 0.0), - body=_collision_cube(0.1), + collider_size=0.1, ), _joint( "Robot", @@ -163,7 +283,7 @@ def _floating_two_link() -> str: limits=(-90.0, 90.0), ), ], - root_schemas=(), + articulation_root=False, ) @@ -218,10 +338,6 @@ def _fixed_spatial_chain() -> str: return _document("Robot", links) -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.""" - - def _revolute_pendulum() -> str: """A fixed pivot and a unit arm turning about z.""" return _document( @@ -414,126 +530,3 @@ def _fixture_dir() -> Path: directory = Path(tempfile.mkdtemp(prefix="isaaclab_newton_test_fixtures_")) atexit.register(shutil.rmtree, directory, ignore_errors=True) return directory - - -def local_usd(name: str, **kwargs: Any) -> sim_utils.UsdFileCfg: - """Return a spawner for one of the local USD fixtures. - - Args: - name: Authored fixture name, or a file name inside :data:`DATA_DIR`. - **kwargs: Additional :class:`~isaaclab.sim.UsdFileCfg` fields. - - Returns: - The spawner configuration. - """ - if name in _FIXTURES: - path = _fixture_dir() / name - if not path.exists(): - path.write_text(_FIXTURES[name](), encoding="utf-8") - else: - path = DATA_DIR / name - return sim_utils.UsdFileCfg(usd_path=str(path), **kwargs) - - -## -# Scenes. -## - - -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, Any], 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 = np.zeros((num_envs, 3), dtype=np.float32) - positions[:, 0] = ENV_SPACING * np.arange(num_envs) - # 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).expand_as(previous)) - 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 fb518e020f17..3a5e223d68ec 100644 --- a/source/isaaclab_newton/test/assets/test_articulation.py +++ b/source/isaaclab_newton/test/assets/test_articulation.py @@ -27,6 +27,7 @@ launch_test_simulation(SimulationCfg(physics=NewtonCfg())) import sys +from collections.abc import Callable, Iterator from copy import copy from dataclasses import dataclass from pathlib import Path @@ -40,14 +41,14 @@ from isaaclab_newton.assets import Articulation from isaaclab_newton.physics import NewtonManager as SimulationManager from isaaclab_physx.sim.schemas import PhysxJointCfg -from newton import ModelFlags, ShapeFlags +from newton import Model, ModelFlags, ShapeFlags, State 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 -from isaaclab.actuators import IdealPDActuatorCfg, ImplicitActuatorCfg +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.envs.mdp.events import randomize_rigid_body_collider_offsets, randomize_rigid_body_material @@ -59,6 +60,7 @@ BRANCHING_PHYSX_BODY_NAMES, BRANCHING_PHYSX_JOINT_NAMES, ) +from isaaclab.utils import replace sys.path.insert(0, str(Path(__file__).resolve().parents[1])) @@ -94,7 +96,9 @@ """ -def _limit_actuators(actuator_type: type) -> dict: +def _limit_actuators( + actuator_type: type[ImplicitActuatorCfg] | type[IdealPDActuatorCfg], +) -> dict[str, ActuatorBaseCfg]: """Return one limit-resolution group per branching-fixture joint. Mirrors the single-joint implicit and explicit rows of the solver-clamp resolution contract: joint limits @@ -118,9 +122,16 @@ def _limit_actuators(actuator_type: type) -> dict: 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, y: float, actuators: dict, **kwargs) -> ArticulationCfg: - init_state = kwargs.pop("init_state", ArticulationCfg.InitialStateCfg()) - init_state.pos = (0.0, y, 1.0) + 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 ) @@ -221,12 +232,6 @@ class _Scene: cfgs: dict[str, ArticulationCfg] origins: torch.Tensor - def __getattr__(self, name: str) -> Articulation: - try: - return self.articulations[name] - except KeyError as error: - raise AttributeError(name) from error - 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): @@ -428,7 +433,7 @@ def _other_callback() -> None: SimulationManager.unregister_post_step_callback(_other_callback) -def _prime_ordered_state_rebind(articulation: Articulation, dt: float): +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 has_ordering = data.joint_ordering is not None @@ -492,7 +497,7 @@ def _prime_ordered_state_rebind(articulation: Articulation, dt: float): 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, new_model) -> None: + 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 = { @@ -596,7 +601,7 @@ def check_rebind(new_state, new_model) -> None: @pytest.fixture(scope="module") -def composite_scene() -> _Scene: +def composite_scene() -> Iterator[_Scene]: """Initialize every articulation island once for this module.""" sim_cfg = newton_sim_cfg(_DEVICE, use_newton_actuators=False, newton_contacts=True) with build_simulation_context(sim_cfg=sim_cfg) as sim: @@ -656,14 +661,14 @@ def test_articulation_initialization_and_partial_state(scene: _Scene) -> None: 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 getattr(actuator, "is_implicit_model", False) == is_implicit_model_cfg + assert actuator.is_implicit_model == is_implicit_model_cfg - floating = scene.floating + 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.tendon.num_fixed_tendons > 0 + 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( @@ -689,15 +694,19 @@ def test_articulation_initialization_and_partial_state(scene: _Scene) -> None: 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.ordered + 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() - 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 + 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) # Check new limits are in place and the in-limit defaults are preserved @@ -757,9 +766,13 @@ def test_articulation_initialization_and_partial_state(scene: _Scene) -> None: 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 + 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) # Check new limits are in place and the defaults are clamped into them @@ -784,7 +797,7 @@ def test_articulation_joint_and_body_properties_round_trip(scene: _Scene, monkey through both writer selectors, and actuator gains taken from the configuration or, when unset, from the USD drives. """ - articulation = scene.floating + 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) @@ -839,7 +852,7 @@ def record_model_change(change: ModelFlags) -> None: articulation.set_inertias_index(inertias=initial_inertia) # Passive viscous joint damping is distinct from actuator derivative gains. - articulation = scene.fixed + 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( @@ -878,7 +891,7 @@ def record_model_change(change: ModelFlags) -> None: 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] * 2, device=_DEVICE) * 180.0 / torch.pi + 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) @@ -890,8 +903,8 @@ def test_articulation_drive_and_dynamics(scene: _Scene) -> None: 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.fixed - root_fixed = scene.root_fixed + 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")} @@ -934,7 +947,7 @@ def test_floating_articulation_root_and_wrench_response(scene: _Scene, monkeypat 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 = scene.floating + articulation = scene.articulations["floating"] scene.rest(articulation) num_articulations = articulation.num_instances env_pos = scene.origins @@ -981,7 +994,7 @@ def random_root_state(num_envs: int) -> torch.Tensor: for state_location in ("com", "link", "root"): - def write_root_state(state: torch.Tensor, env_ids: torch.Tensor | None = None): + 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) @@ -1001,8 +1014,6 @@ def read_written_root_state() -> torch.Tensor: 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 scene.step("floating") # Prime the lazily-derived caches at the current sim timestamp. Without this they would @@ -1098,12 +1109,10 @@ def one_step_response() -> tuple[torch.Tensor, torch.Tensor]: force = 1.0 for is_global in (False, True): scene.rest(articulation) - positions = zeros[:, :1].clone() - positions[..., 1] = 0.1 + 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 = zeros[:, :1].clone() - forces[..., 2] = force + 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, @@ -1124,11 +1133,11 @@ def one_step_response() -> tuple[torch.Tensor, torch.Tensor]: 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 = next(iter(articulation.actuators.values())) + 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) @@ -1180,7 +1189,7 @@ def test_branching_fixture_physx_ordering_reorders_newton_to_bfs(scene: _Scene) 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.branching + articulation = scene.articulations["branching"] device = _DEVICE # Newton's native traversal is depth-first, so the live backend view already reflects MJWarp order. @@ -1248,7 +1257,7 @@ def test_newton_ordered_state_publishes_in_step_and_refreshes_same_timestamp_wri 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. """ - articulation = scene.ordered + articulation = scene.articulations["ordered"] scene.rest(articulation) data = articulation.data assert data.joint_ordering is not None @@ -1405,7 +1414,7 @@ def test_set_material_properties(scene: _Scene) -> None: 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 = scene.floating + articulation = scene.articulations["floating"] num_articulations = articulation.num_instances device = _DEVICE @@ -1495,7 +1504,7 @@ def test_hand_with_tendons_targets_only_given_envs_and_properties_reach_solver(s Written fixed tendon stiffness, damping, and position limits reach the MuJoCo solver through both the index and the mask setters and writers. """ - articulation = scene.tendon + articulation = scene.articulations["tendon"] scene.rest(articulation) num_articulations = articulation.num_instances assert articulation.num_fixed_tendons == 2 @@ -1594,7 +1603,7 @@ def test_get_gravity_compensation_forces_matches_jacobian_gravity(scene: _Scene, # 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 = [island.num_joints + island.num_base_dofs for island in scene.articulations.values()] + 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 @@ -1744,7 +1753,7 @@ def test_get_gravity_compensation_forces_static_equilibrium(scene: _Scene) -> No target IS the joint torque applied and no PD spring-damper masks the gravity-compensation signal. Sign errors, frame errors, and DoF-ordering errors all surface as joint drift. """ - articulation = scene.passive + articulation = scene.articulations["passive"] scene.rest(articulation) # a configuration that loads the prismatic and the revolute joints with gravity q = articulation.data.default_joint_pos.torch.clone() @@ -1780,7 +1789,7 @@ def test_body_root_state(scene: _Scene) -> None: 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.pendulum + articulation = scene.articulations["pendulum"] scene.rest(articulation) num_articulations = articulation.num_instances num_bodies = articulation.num_bodies @@ -1804,11 +1813,10 @@ def test_body_root_state(scene: _Scene) -> None: for _ in range(50): scene.step("pendulum") - # check that the root is at the correct state - its default state as it is fixed base + # 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) - # 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 @@ -1816,40 +1824,31 @@ def test_body_root_state(scene: _Scene) -> None: 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 + # 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:]) - # lin_vel arm + # 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) - vz = torch.zeros(num_articulations, 1, 1, device=device) - lin_vel_gt[:, arm_idx, :] = torch.cat([vx, vy, vz], dim=-1).squeeze(-2) - - # linear velocity of root link should be zero + 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) - # 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 + # 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) - pz = torch.zeros(num_articulations, 1, 1, device=device) - pos_gt[:, arm_idx, :] = torch.cat([px, py, pz], dim=-1).squeeze(-2) - pos_gt += pivot_pos_w.unsqueeze(-2).repeat(1, num_bodies, 1) + 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) - # orientation + # 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:]) @@ -1868,7 +1867,7 @@ def test_gravity_vec_w_tracks_model_gravity(scene: _Scene) -> None: gravity broadcast to every env, hiding per-env gravity randomization (e.g. :class:`~isaaclab.envs.mdp.randomize_physics_scene_gravity`). """ - articulation = scene.floating + articulation = scene.articulations["floating"] scene.rest(articulation) num_articulations = articulation.num_instances device = _DEVICE @@ -1919,7 +1918,7 @@ def test_body_q_consistent_after_root_write(scene: _Scene) -> None: 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. """ - articulation = scene.floating + articulation = scene.articulations["floating"] scene.rest(articulation) model = SimulationManager.get_model() jc_starts = model.joint_coord_world_start.numpy() @@ -1943,7 +1942,7 @@ def test_body_q_consistent_after_root_write(scene: _Scene) -> None: original_simulate = SimulationManager._simulate_physics_only.__func__ @classmethod # type: ignore[misc] - def _patched_simulate(cls): + 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) 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 64286fb80a5d..c6d6319952c7 100644 --- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py +++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py @@ -29,6 +29,7 @@ import functools import os import sys +from collections.abc import Iterator from dataclasses import dataclass, field from pathlib import Path @@ -39,7 +40,7 @@ from isaaclab_newton.physics import MJWarpSolverCfg from isaaclab_newton.physics import NewtonManager as SimulationManager -from isaaclab.actuators import DCMotorCfg, DelayedPDActuatorCfg, IdealPDActuatorCfg +from isaaclab.actuators import ActuatorBaseCfg, DCMotorCfg, DelayedPDActuatorCfg, IdealPDActuatorCfg 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 @@ -78,6 +79,7 @@ def _solver_cfg() -> MJWarpSolverCfg: + """Return the MJWarp solver configuration shared by the actuator scenes.""" return MJWarpSolverCfg( njmax=500, nconmax=500, @@ -140,7 +142,7 @@ def _solver_cfg() -> MJWarpSolverCfg: } -def _neural_actuators(mlp_path: str, lstm_path: str) -> dict: +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( @@ -175,7 +177,7 @@ def _neural_actuators(mlp_path: str, lstm_path: str) -> dict: class _Island: """One actuator configuration under test and the commands its rollout sends.""" - actuators: dict + actuators: dict[str, ActuatorBaseCfg] usd: str = "floating_two_leg.usda" joint_ordering: tuple[str, ...] | None = None feedforward: float | None = None @@ -191,6 +193,7 @@ class _Island: 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), @@ -230,6 +233,7 @@ class _Run: 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( @@ -401,7 +405,7 @@ def _record_lab_state_reset(run: _Run) -> dict[str, tuple[torch.Tensor, torch.Te @pytest.fixture(scope="module", params=test_devices(DeviceScope.CUDA)) -def device(request) -> str: +def device(request: pytest.FixtureRequest) -> str: """Simulation device of the actuator scenes.""" return request.param @@ -437,7 +441,7 @@ def decimated_runs(device: str) -> dict[str, dict]: @pytest.fixture(scope="module") -def newton_run(device: str, lab_run: dict, decimated_runs: dict) -> _Run: +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. Depends on the Lab-path and decimated runs so their scenes are built and closed first. @@ -719,7 +723,7 @@ def test_newton_state_reset_isolated_to_reset_env(newton_run: _Run) -> None: 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 + if 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" diff --git a/source/isaaclab_newton/test/assets/test_rigid_object.py b/source/isaaclab_newton/test/assets/test_rigid_object.py index ef7b91fd22a8..e12dcb92eda1 100644 --- a/source/isaaclab_newton/test/assets/test_rigid_object.py +++ b/source/isaaclab_newton/test/assets/test_rigid_object.py @@ -20,6 +20,7 @@ launch_test_simulation(SimulationCfg(physics=NewtonCfg())) import sys +from collections.abc import Iterator from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace @@ -133,15 +134,22 @@ class _Scene: cube: RigidObject origins: torch.Tensor device: str + masses: torch.Tensor + """Configured body masses [kg], restored before each test.""" + coms: torch.Tensor + """Configured center-of-mass positions in the body frames [m], restored before each test.""" + inertias: torch.Tensor + """Configured body inertias [kg·m^2], restored before each test.""" def step(self, num_steps: int = 1) -> None: + """Step ``num_steps`` times, writing the cube commands before and updating the cubes after each step.""" for _ in range(num_steps): self.cube.write_data_to_sim() self.sim.step() self.cube.update(self.sim.cfg.dt) def rest(self) -> None: - """Put the cubes at rest at their configured pose with the configured inertial properties.""" + """Put the cubes at rest at their 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) @@ -152,30 +160,31 @@ def rest(self) -> None: @pytest.fixture(scope="module", params=test_devices()) -def shared_scene(request) -> _Scene: +def shared_scene(request: pytest.FixtureRequest) -> Iterator[_Scene]: """Initialize the cubes once per device for this module.""" device = request.param with build_simulation_context(sim_cfg=newton_sim_cfg(device)) as sim: cube = spawn_assets({"cube": _cube_cfg()})["cube"] sim.reset() - scene = _Scene(sim=sim, cube=cube, origins=env_origins(device), device=device) - initial_properties = ( - cube.data.body_mass.torch.clone(), - cube.data.body_com_pos_b.torch.clone(), - cube.data.body_inertia.torch.clone(), + yield _Scene( + sim=sim, + cube=cube, + origins=env_origins(device), + device=device, + masses=cube.data.body_mass.torch.clone(), + coms=cube.data.body_com_pos_b.torch.clone(), + inertias=cube.data.body_inertia.torch.clone(), ) - yield scene, initial_properties @pytest.fixture -def scene(shared_scene) -> _Scene: +def scene(shared_scene: _Scene) -> _Scene: """Hand each test the shared scene at rest with the configured inertial properties.""" - scene, (masses, coms, inertias) = shared_scene - scene.cube.set_masses_index(masses=masses) - scene.cube.set_coms_index(coms=coms) - scene.cube.set_inertias_index(inertias=inertias) - scene.rest() - return scene + shared_scene.cube.set_masses_index(masses=shared_scene.masses) + shared_scene.cube.set_coms_index(coms=shared_scene.coms) + shared_scene.cube.set_inertias_index(inertias=shared_scene.inertias) + shared_scene.rest() + return shared_scene def test_external_force_on_single_body(scene: _Scene) -> None: @@ -580,9 +589,8 @@ def test_write_root_state(scene: _Scene) -> None: # 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 (regressions: - # body_com_pose_w returned the pre-write buffer after a link-frame pose write, and - # body_link_vel_w after a center-of-mass velocity 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() diff --git a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py index a10e472056c4..7626bc5e5aef 100644 --- a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py @@ -22,6 +22,7 @@ launch_test_simulation(SimulationCfg(physics=NewtonCfg())) import sys +from collections.abc import Iterator from dataclasses import dataclass from pathlib import Path @@ -131,8 +132,15 @@ class _Scene: sibling: RigidObject origins: torch.Tensor device: str + masses: torch.Tensor + """Configured body masses [kg], restored before each test.""" + coms: torch.Tensor + """Configured center-of-mass positions in the body frames [m], restored before each test.""" + inertias: torch.Tensor + """Configured body inertias [kg·m^2], restored before each test.""" def step(self, num_steps: int = 1) -> None: + """Step ``num_steps`` times, writing the collection's commands before and updating it after each step.""" for _ in range(num_steps): self.collection.write_data_to_sim() self.sim.step() @@ -150,7 +158,7 @@ def rest(self) -> None: @pytest.fixture(scope="module", params=test_devices()) -def shared_scene(request): +def shared_scene(request: pytest.FixtureRequest) -> Iterator[_Scene]: """Initialize the collection and its unrelated sibling bodies once per device for this module.""" device = request.param sibling_cfg = RigidObjectCfg( @@ -162,26 +170,26 @@ def shared_scene(request): assets = spawn_assets({"collection": _collection_cfg(_NUM_CUBES), "sibling": sibling_cfg}) sim.reset() collection = assets["collection"] - scene = _Scene( - sim=sim, collection=collection, sibling=assets["sibling"], origins=env_origins(device), device=device - ) - initial_properties = ( - collection.data.body_mass.torch.clone(), - collection.data.body_com_pos_b.torch.clone(), - collection.data.body_inertia.torch.clone(), + yield _Scene( + sim=sim, + collection=collection, + sibling=assets["sibling"], + origins=env_origins(device), + device=device, + masses=collection.data.body_mass.torch.clone(), + coms=collection.data.body_com_pos_b.torch.clone(), + inertias=collection.data.body_inertia.torch.clone(), ) - yield scene, initial_properties @pytest.fixture -def scene(shared_scene) -> _Scene: +def scene(shared_scene: _Scene) -> _Scene: """Hand each test the shared scene at rest with the configured inertial properties.""" - scene, (masses, coms, inertias) = shared_scene - scene.collection.set_masses_index(masses=masses) - scene.collection.set_coms_index(coms=coms) - scene.collection.set_inertias_index(inertias=inertias) - scene.rest() - return scene + shared_scene.collection.set_masses_index(masses=shared_scene.masses) + shared_scene.collection.set_coms_index(coms=shared_scene.coms) + shared_scene.collection.set_inertias_index(inertias=shared_scene.inertias) + shared_scene.rest() + return shared_scene def test_initialization(scene: _Scene) -> None: diff --git a/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py b/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py index f255a75efe2d..52ddbe6c7d34 100644 --- a/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py +++ b/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py @@ -22,6 +22,7 @@ launch_test_simulation(SimulationCfg(physics=NewtonCfg())) import sys +from collections.abc import Iterator from dataclasses import dataclass from pathlib import Path @@ -68,7 +69,7 @@ def rest(self) -> None: @pytest.fixture(scope="module") -def chain() -> _Chain: +def chain() -> Iterator[_Chain]: """Build the unpowered chain once, at a configuration away from the joint limits.""" robot_cfg = ArticulationCfg( prim_path="/World/Env_[^/]*/Robot", @@ -93,18 +94,15 @@ def chain() -> _Chain: ) -def _compute_ee_pose_root(robot, ee_frame_idx): - """Return ``(ee_pos_b, ee_quat_b, root_pose_w)`` in the root frame.""" +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 - 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 + return 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, ee_jacobi_idx, arm_joint_ids): - """Return the EE Jacobian sliced to ``arm_joint_ids`` and rotated to the root frame.""" +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 = matrix_from_quat(quat_inv(robot.data.root_pose_w.torch[:, 3:7])) jacobian[:, :3, :] = torch.bmm(base_rot_matrix, jacobian[:, :3, :]) @@ -112,29 +110,17 @@ def _compute_jacobian_root_frame(robot, ee_jacobi_idx, arm_joint_ids): return jacobian -def _compute_ee_vel_root(jacobian_b, joint_vel): - """Return the EE 6D velocity in the root frame as ``J · q_dot``. - - Required to make OSC's ``kd * ee_vel_b`` damping term meaningful: zero EE velocity leaves the impedance - undamped. ``J`` correctness is pinned independently by the articulation's gravity-compensation identity test. - """ - return torch.bmm(jacobian_b, joint_vel.unsqueeze(-1)).squeeze(-1) - - -def _build_relative_pose_target(robot, ee_frame_idx, delta_xyz): - """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) +def _build_relative_pose_target( + robot: Articulation, ee_frame_idx: int, delta_xyz: tuple[float, float, float] +) -> torch.Tensor: + """Return the current end-effector pose in the root frame offset by ``delta_xyz`` [m], keeping its 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(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) - - def _make_osc() -> OperationalSpaceController: + """Return a fixed-impedance absolute-pose controller with inertial decoupling and no gravity compensation.""" return OperationalSpaceController( OperationalSpaceControllerCfg( target_types=["pose_abs"], @@ -150,7 +136,9 @@ def _make_osc() -> OperationalSpaceController: ) -def _run_osc(chain: _Chain, osc: OperationalSpaceController, target_pose_b, num_steps: int, *, gravity: bool): +def _run_osc( + chain: _Chain, 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 = chain.robot pos_history: list[float] = [] @@ -159,10 +147,12 @@ def _run_osc(chain: _Chain, osc: OperationalSpaceController, target_pose_b, num_ jacobian_b = _compute_jacobian_root_frame(robot, chain.ee_jacobi_idx, chain.arm_joint_ids) mass_matrix = robot.data.mass_matrix.torch[:, chain.arm_joint_ids, :][:, :, chain.arm_joint_ids] gravity_forces = robot.data.gravity_compensation_forces.torch[:, chain.arm_joint_ids] if gravity else None - ee_pos_b, ee_quat_b, _ = _compute_ee_pose_root(robot, chain.ee_frame_idx) + ee_pos_b, ee_quat_b = _compute_ee_pose_root(robot, chain.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[:, chain.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( @@ -194,8 +184,8 @@ def test_osc_tracking_accuracy(chain: _Chain) -> None: target_pose_b = _build_relative_pose_target(chain.robot, chain.ee_frame_idx, (0.05, 0.0, 0.0)) pos_history, rot_history = _run_osc(chain, _make_osc(), target_pose_b, 300, gravity=False) - pos_mean = _tail_mean(pos_history) - rot_mean = _tail_mean(rot_history) + pos_mean = sum(pos_history[-200:]) / 200 + rot_mean = sum(rot_history[-200:]) / 200 # 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 @@ -225,7 +215,7 @@ def test_osc_gravity_compensation_precision(chain: _Chain) -> None: # Hold the initial EE pose: phase-1 steady-state error is pure gravity sag. target_pose_b = _build_relative_pose_target(chain.robot, chain.ee_frame_idx, (0.0, 0.0, 0.0)) - def _stationary_tail_mean(history, label): + def _stationary_tail_mean(history: list[float], label: str) -> float: """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 solver noise floor of the compensated hold, diff --git a/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p05-ov-manager.rst b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p05-ov-manager.rst index 56cb1423f640..a0e44b9aec14 100644 --- a/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p05-ov-manager.rst +++ b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p05-ov-manager.rst @@ -1,7 +1,7 @@ Fixed ^^^^^ -* Fixed :class:`~isaaclab_ov.physics.OvPhysxManager` refusing to start a CPU simulation after a CUDA one, or the - reverse, in the same process. CPU scenes now select CPU dynamics through their physics-scene attributes instead of - the OVPhysX process-wide CPU-only mode, which cannot be reverted. Deployments that must keep OVPhysX off CUDA can - set ``OVPHYSX_DISABLE_GPU``. +* Fixed :class:`~isaaclab_ov.physics.OvPhysxManager` refusing to start a CPU simulation after a CUDA simulation, or + the reverse, in the same process. CPU scenes selected CPU dynamics through their physics-scene attributes 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/isaaclab_ov/physics/ovphysx_manager.py b/source/isaaclab_ov/isaaclab_ov/physics/ovphysx_manager.py index e9a0b3b78443..e18f9a44f449 100644 --- a/source/isaaclab_ov/isaaclab_ov/physics/ovphysx_manager.py +++ b/source/isaaclab_ov/isaaclab_ov/physics/ovphysx_manager.py @@ -245,8 +245,8 @@ def __init__(self, cfg: OvPhysxBackendCfg): } if is_gpu: carbonite_overrides.update({"/physics/suppressReadback": True, "/physics/suppressFabricUpdate": True}) - # Hard CPU-only mode (``PhysX.set_cpu_mode(True)``) is sticky for the process, so CPU scenes - # select CPU dynamics through the scene attributes instead; see _configure_physx_scene_prim. + # CPU scenes select CPU dynamics through their scene attributes (see _configure_physx_scene_prim). + # 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 @@ -829,12 +829,11 @@ def _replay_pending_clones(cls, physx: Any, requires_full_stage: bool) -> None: 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 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 + 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: @@ -956,10 +955,10 @@ def _configure_physx_scene_prim(scene_prim, cfg, device: str) -> None: raw Sdf metadata manipulation instead of using the high-level USD API. The schema, scene-query-support, solver-determinism/accuracy, and dynamics/broadphase - attributes are applied regardless of device. GPU scenes use ``enableGPUDynamics=true`` and - the GPU broadphase. CPU scenes use ``enableGPUDynamics=false`` and the MBP broadphase; the - schema default enables GPU dynamics, so a CPU scene that omits them would silently simulate - on the GPU. The GPU buffer-capacity attributes are applied only when ``device == "gpu"``. + attributes are applied regardless of device. GPU scenes use GPU dynamics and the GPU + broadphase; CPU scenes use CPU dynamics and the MBP broadphase. The schema default enables + GPU dynamics, so CPU scenes must author it explicitly. The GPU buffer-capacity attributes + are applied only to GPU scenes. Args: scene_prim: The /World/PhysicsScene prim to configure. @@ -991,12 +990,13 @@ def _configure_physx_scene_prim(scene_prim, cfg, device: str) -> None: cfg.enable_external_forces_every_iteration ) - scene_prim.CreateAttribute("physxScene:enableGPUDynamics", Sdf.ValueTypeNames.Bool).Set(device == "gpu") + 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 device == "gpu" else "MBP" + "GPU" if use_gpu else "MBP" ) - if device == "gpu" and cfg is not None: + if use_gpu and cfg is not None: for attr, val in [ ("gpuMaxRigidContactCount", cfg.gpu_max_rigid_contact_count), ("gpuMaxRigidPatchCount", cfg.gpu_max_rigid_patch_count), diff --git a/source/isaaclab_ov/test/assets/test_articulation.py b/source/isaaclab_ov/test/assets/test_articulation.py index 39c513ad0ae7..afee01d99612 100644 --- a/source/isaaclab_ov/test/assets/test_articulation.py +++ b/source/isaaclab_ov/test/assets/test_articulation.py @@ -45,6 +45,8 @@ import isaaclab.utils.math as math_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.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 @@ -61,12 +63,12 @@ _ALL_DEVICES = test_devices() _CUDA_DEVICES = test_devices(DeviceScope.CUDA) _ROTATED_JOINT_NAMES = (*BRANCHING_PHYSX_JOINT_NAMES[1:], BRANCHING_PHYSX_JOINT_NAMES[0]) -# We set 80.0 for max force because the USD default is 10e10, which makes testing annoying. +# 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: +def _implicit(**kwargs) -> dict[str, ImplicitActuatorCfg]: return {"joints": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=_STIFFNESS, damping=_DAMPING, **kwargs)} @@ -339,13 +341,6 @@ def _read_binding_to_torch(articulation: Articulation, tensor_type: int, device: return wp.to_torch(articulation.root_view.get_attribute(tensor_type)).to(device) -def _backend_to_user(articulation: Articulation) -> list[int]: - """Return the public joint index of every backend joint column.""" - if articulation.joint_ordering is None: - return list(range(articulation.num_joints)) - return list(articulation.joint_ordering.backend_to_user_indices) - - def _spawn_own_articulation(**cfg) -> Articulation: """Spawn one branching articulation for a test that needs its own scene.""" return Articulation( @@ -367,7 +362,7 @@ def _spawn_own_articulation(**cfg) -> Articulation: ], ids=["position", "velocity"], ) -def test_out_of_range_default_joint_state(device, init_state, match): +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: @@ -381,7 +376,7 @@ def test_out_of_range_default_joint_state(device, init_state, match): @pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) -def test_setting_invalid_articulation_root_prim_path(device): +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: @@ -420,7 +415,7 @@ def test_articulation_initialization_and_partial_state(scene: _ArticulationScene ) 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 = _backend_to_user(articulation) + 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], @@ -529,8 +524,6 @@ def test_set_material_properties(scene: _ArticulationScene) -> None: :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 = scene["ordered"] num_articulations = articulation.num_instances device = scene.device @@ -580,9 +573,7 @@ def test_setting_velocity_and_effort_limits_write_to_solver(scene: _Articulation actuator_cfg = _ISLANDS[island].actuators["joints"] joint_drive_props = _ISLANDS[island].joint_drive_props - # 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 actuator_cfg.joint_velocity_limit is None: @@ -592,9 +583,7 @@ def test_setting_velocity_and_effort_limits_write_to_solver(scene: _Articulation expected_velocity_limit = torch.full_like(physx_vel_limit, limit) torch.testing.assert_close(physx_vel_limit, expected_velocity_limit) - # 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 actuator_cfg.joint_effort_limit is None: @@ -674,12 +663,12 @@ def test_com_orientation_write_invalidates_static_inertia_cache_with_body_orderi 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).contiguous(), dtype=wp.transformf), + 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), ) 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.contiguous())) + 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 @@ -702,7 +691,7 @@ def test_com_orientation_write_invalidates_static_inertia_cache_with_body_orderi # 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]) + articulation.set_coms_index(coms=wp.from_torch(coms, dtype=wp.transformf), body_ids=[public_body_id]) 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]) @@ -711,7 +700,7 @@ def test_com_orientation_write_invalidates_static_inertia_cache_with_body_orderi ) 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]) + 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( @@ -941,7 +930,7 @@ def assert_dynamics_match(jacobian: bool = True, gravity: bool = True) -> None: data.gravity_compensation_forces 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 @@ -1074,7 +1063,7 @@ def test_floating_articulation_root_state_dynamics_wrench_and_reset(scene: _Arti # 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) com[:, 0, :3] = torch.tensor([1.0, 0.0, 0.0], device=device) - articulation.set_coms_index(coms=wp.from_torch(com.contiguous(), dtype=wp.transformf)) + 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) @@ -1172,7 +1161,7 @@ def test_tendon_properties_and_position_targets(scene: _ArticulationScene) -> No + 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) + 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 ) @@ -1192,7 +1181,7 @@ def test_native_actuator_submits_real_effort(scene: _ArticulationScene) -> None: device = scene.device assert articulation._actuator_control.native_actuator_path_active assert articulation.newton_actuator_adapter is not None - backend_to_user = _backend_to_user(articulation) + 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() @@ -1226,9 +1215,8 @@ def test_native_actuator_submits_real_effort(scene: _ArticulationScene) -> None: @pytest.mark.parametrize("scene", _CUDA_DEVICES, 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 - from isaaclab.envs.mdp.events import randomize_actuator_gains # noqa: PLC0415 - from isaaclab.managers import EventTermCfg, SceneEntityCfg # noqa: PLC0415 class Env: def __init__(self, asset): diff --git a/source/isaaclab_ov/test/assets/test_rigid_object.py b/source/isaaclab_ov/test/assets/test_rigid_object.py index e7cae9140158..9cdd281b12bd 100644 --- a/source/isaaclab_ov/test/assets/test_rigid_object.py +++ b/source/isaaclab_ov/test/assets/test_rigid_object.py @@ -16,6 +16,7 @@ from __future__ import annotations from collections.abc import Iterator +from contextlib import AbstractContextManager from dataclasses import dataclass import pytest @@ -41,7 +42,7 @@ _NUM_CUBES = 2 -def _sim_context(device: str): +def _sim_context(device: str) -> AbstractContextManager[SimulationContext]: """Build a local OVPhysX context from an in-memory USD stage.""" return build_simulation_context(device=device, sim_cfg=SimulationCfg(physics=OvPhysxCfg(), device=device)) @@ -78,7 +79,7 @@ def _spawn_static_colliders() -> RigidObject: @pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) -def test_initialization_with_no_rigid_body(device): +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 _sim_context(device) as sim: # Keep the asset alive: only a live asset initializes, and fails, on reset. @@ -94,10 +95,10 @@ def test_initialization_with_no_rigid_body(device): raises=pytest.fail.Exception, reason="OVPhysX RigidObject has no ArticulationRootAPI guard; the rigid body initializes normally.", ) -def test_initialization_with_articulation_root(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 _sim_context(device) as sim: - # Generate cubes and mark each rigid body as an articulation root. Keep the asset alive for reset. + # 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 index in range(_NUM_CUBES): @@ -249,24 +250,22 @@ def test_body_root_state_properties(scene: _RigidScene) -> None: cube_object, sim, device = scene.spinning, scene.sim, scene.device env_pos = cube_object.data.root_link_pos_w.torch.clone() - # change center of mass offset from link frame + # 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.contiguous(), dtype=wp.transformf)) + 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) - # z spin velocity spin_twist = torch.zeros(6, device=device) spin_twist[5] = 2.0 for _ in range(10): - # spin the object around Z axis (com) + # 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) - # 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 @@ -276,12 +275,10 @@ def test_body_root_state_properties(scene: _RigidScene) -> None: 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) + # 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)) - # link position will be moving but should stay constant away from center of mass + # 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], @@ -293,17 +290,16 @@ def test_body_root_state_properties(scene: _RigidScene) -> None: ) torch.testing.assert_close(-offset, body_link_state_pos_rel_com.squeeze(-2)) - # orientation of com will be a constant rotation from link orientation + # 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:]) - # lin_vel will not match - # center of mass vel will be constant (i.e. spinning around com) + # 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]) - # link frame will be moving, and should account for the reported COM velocity and offset + # 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]) @@ -312,6 +308,6 @@ def test_body_root_state_properties(scene: _RigidScene) -> None: 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 + # 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:]) diff --git a/source/isaaclab_ov/test/assets/test_rigid_object_collection.py b/source/isaaclab_ov/test/assets/test_rigid_object_collection.py index 35c6eafa90ba..1ddf73b50868 100644 --- a/source/isaaclab_ov/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_ov/test/assets/test_rigid_object_collection.py @@ -16,6 +16,7 @@ from __future__ import annotations from collections.abc import Iterator +from contextlib import AbstractContextManager from dataclasses import dataclass import pytest @@ -38,7 +39,7 @@ _NUM_ENVS, _NUM_BODIES = 2, 3 -def _sim_context(device: str): +def _sim_context(device: str) -> AbstractContextManager[SimulationContext]: """Build a local OVPhysX context from an in-memory USD stage.""" return build_simulation_context( device=device, sim_cfg=SimulationCfg(physics=OvPhysxCfg(), device=device, gravity=(0.0, 0.0, 0.0)) @@ -71,7 +72,7 @@ def _spawn_collection(name: str, y_offset: float, spawn: sim_utils.SpawnerCfg | @pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) -def test_initialization_with_no_rigid_body(device): +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 _sim_context(device) as sim: # Keep the asset alive: only a live asset initializes, and fails, on reset. diff --git a/source/isaaclab_ov/test/assets/unit/test_ovphysx_actuator_control.py b/source/isaaclab_ov/test/assets/unit/test_ovphysx_actuator_control.py index 38a05f280ec1..0da4414e6197 100644 --- a/source/isaaclab_ov/test/assets/unit/test_ovphysx_actuator_control.py +++ b/source/isaaclab_ov/test/assets/unit/test_ovphysx_actuator_control.py @@ -17,6 +17,7 @@ pytest.importorskip("ovphysx.types", reason="ovphysx wheel not installed") +from isaaclab_ov.assets.articulation import actuator_control # noqa: E402 from isaaclab_ov.assets.articulation.actuator_control import OvPhysxActuatorControl # noqa: E402 from isaaclab.actuators import ImplicitActuatorCfg # noqa: E402 @@ -26,8 +27,6 @@ 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 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 2690c821d00f..24b9c1b0d6dd 100644 --- a/source/isaaclab_ov/test/physics/test_ovphysx_manager_lifecycle.py +++ b/source/isaaclab_ov/test/physics/test_ovphysx_manager_lifecycle.py @@ -20,6 +20,9 @@ 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): @@ -626,10 +629,6 @@ def test_default_cache_dir_rejects_a_directory_owned_by_another_user(manager_mod manager_module._prepare_default_cache_dir(str(target)) -_CPU_DEVICES = test_devices(DeviceScope.CPU) -_CUDA_DEVICES = test_devices(DeviceScope.CUDA) - - @pytest.mark.skipif(not (_CPU_DEVICES and _CUDA_DEVICES), reason="The device sequence requires a CPU and a CUDA device") @pytest.mark.parametrize("order", ["cpu-cuda-cpu", "cuda-cpu"]) def test_cpu_and_cuda_scenes_run_sequentially_in_one_process(order): diff --git a/source/isaaclab_physx/test/assets/test_articulation.py b/source/isaaclab_physx/test/assets/test_articulation.py index 2bb4c02d9073..e6d0e134aa78 100644 --- a/source/isaaclab_physx/test/assets/test_articulation.py +++ b/source/isaaclab_physx/test/assets/test_articulation.py @@ -28,6 +28,7 @@ import math import sys +from collections.abc import Iterator from dataclasses import dataclass from pathlib import Path @@ -163,14 +164,9 @@ def _author_branching_robot( leaf.CreateUpperLimitAttr(2.0) -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 _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]] * _NUM_ENVS, device=device) + 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]: @@ -179,7 +175,7 @@ def _yaw_quat(angle: float) -> tuple[float, float, float, float]: @pytest.fixture -def sim(request): +def sim(request) -> Iterator[SimulationContext]: """Create a function-scoped simulation context for tests that own their scene.""" device = request.getfixturevalue("device") gravity_enabled = request.getfixturevalue("gravity_enabled") if "gravity_enabled" in request.fixturenames else True @@ -222,7 +218,7 @@ def generate_articulation( ], ) @pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -def test_out_of_range_default_joint_state_raises(sim, device, init_state, message): +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. @@ -236,7 +232,7 @@ def test_out_of_range_default_joint_state_raises(sim, device, init_state, messag @pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -def test_setting_invalid_articulation_root_prim_path(sim, device): +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") @@ -250,7 +246,7 @@ def test_setting_invalid_articulation_root_prim_path(sim, device): @pytest.mark.parametrize("num_articulations", [2]) @pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -def test_fixed_tendon_position_target_writes_offset(sim, num_articulations, device): +def test_fixed_tendon_position_target_writes_offset(sim, num_articulations, device) -> None: """A tendon length target lands in the simulation as ``rest_length - target`` on the selected cells only. The index form commands every tendon of environment 0; the mask form commands tendon 0 of @@ -268,7 +264,7 @@ def test_fixed_tendon_position_target_writes_offset(sim, num_articulations, devi 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 + 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() @@ -302,7 +298,7 @@ def test_fixed_tendon_position_target_writes_offset(sim, num_articulations, devi @pytest.mark.parametrize("num_articulations", [1]) @pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) @pytest.mark.isaacsim_ci -def test_get_gravity_compensation_forces_static_equilibrium(sim, num_articulations, device): +def test_get_gravity_compensation_forces_static_equilibrium(sim, num_articulations, device) -> None: """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``. @@ -445,7 +441,7 @@ def _summarize_history(history, tail: int = 200): @pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) @pytest.mark.parametrize("gravity_enabled", [False]) @pytest.mark.isaacsim_ci -def test_franka_ik_tracking_accuracy(sim, device, gravity_enabled): +def test_franka_ik_tracking_accuracy(sim, device, gravity_enabled) -> None: """PhysX-side IK convergence sentinel — backend parity with the Newton test. Mirrors :func:`isaaclab_newton.test.assets.test_articulation.test_franka_ik_tracking_accuracy` @@ -526,7 +522,7 @@ class _ArticulationScene: @property def islands(self) -> dict[str, Articulation]: """All articulations keyed by island name.""" - return {name: getattr(self, name) for name in ("ordered", "floating", "reordered", "tendon")} + 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.""" @@ -553,7 +549,7 @@ def _spawn_island(name: str, y_offset: float, cfg: ArticulationCfg, **authoring) @pytest.fixture(scope="module", params=test_devices()) -def articulation_scene(request) -> _ArticulationScene: +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: @@ -604,7 +600,7 @@ def articulation_scene(request) -> _ArticulationScene: ) -def test_articulation_initialization_and_ordering(articulation_scene: _ArticulationScene): +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 @@ -615,7 +611,7 @@ def test_articulation_initialization_and_ordering(articulation_scene: _Articulat 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 getattr(actuator, "is_implicit_model", False) == is_implicit_model_cfg + assert actuator.is_implicit_model == is_implicit_model_cfg 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 @@ -652,27 +648,27 @@ def test_articulation_initialization_and_ordering(articulation_scene: _Articulat ): expected = _in_user_order(usd_values, ordered.joint_names, device) torch.testing.assert_close(public, expected) - torch.testing.assert_close(_to_device_tensor(raw, device), expected[:, joint_backend_to_user]) + 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 attribute, raw, values in ( - ("stiffness", floating.root_view.get_dof_stiffnesses(), (6.0, 4.0)), - ("damping", floating.root_view.get_dof_dampings(), (0.6, 0.4)), + 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(getattr(floating.actuators["joints"], attribute), expected) - torch.testing.assert_close(_to_device_tensor(raw, device), expected) + 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(_to_device_tensor(raw, device), public) + torch.testing.assert_close(wp.to_torch(raw).to(device), public) -def test_articulation_joint_state_writes_follow_ordering(articulation_scene: _ArticulationScene): +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 @@ -695,7 +691,7 @@ def test_articulation_joint_state_writes_follow_ordering(articulation_scene: _Ar (articulation.root_view.get_dof_velocities(), joint_vel), (articulation.root_view.get_dof_stiffnesses(), joint_stiffness), ): - torch.testing.assert_close(_to_device_tensor(raw, device), written[:, joint_backend_to_user]) + 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 @@ -723,7 +719,7 @@ def test_articulation_joint_state_writes_follow_ordering(articulation_scene: _Ar (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, _to_device_tensor(raw, device)[:, 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]) @@ -746,18 +742,18 @@ def test_articulation_joint_state_writes_follow_ordering(articulation_scene: _Ar 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( - _to_device_tensor(articulation.root_view.get_dof_positions(), device), + wp.to_torch(articulation.root_view.get_dof_positions()).to(device), expected_position[:, joint_backend_to_user], ) torch.testing.assert_close( - _to_device_tensor(articulation.root_view.get_dof_velocities(), device), + 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 test_articulation_joint_and_body_properties_round_trip(articulation_scene: _ArticulationScene): +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 @@ -766,17 +762,15 @@ def test_articulation_joint_and_body_properties_round_trip(articulation_scene: _ # buffer (identity order), and after the next update once the buffer is primed (non-identity order). cpu_env_ids = wp.array(list(range(_NUM_ENVS)), dtype=wp.int32, device="cpu") - def write_backend_mass_inertia(articulation: Articulation, delta_mass: float, delta_inertia: float): + 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.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 - ) + 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) backend_masses, backend_inertias = write_backend_mass_inertia(scene.tendon, 0.293, 0.023) @@ -862,11 +856,19 @@ def write_backend_mass_inertia(articulation: Articulation, delta_mass: float, de (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(_to_device_tensor(raw, device), public[:, body_backend_to_user]) + torch.testing.assert_close(wp.to_torch(raw).to(device), public[:, body_backend_to_user]) def _rotate_z(angle: torch.Tensor, vector: tuple[float, float, float] | torch.Tensor) -> torch.Tensor: - """Rotate vectors about the world z axis by per-environment angles [rad].""" + """Rotate vectors about the world z axis. + + 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( @@ -875,15 +877,34 @@ def _rotate_z(angle: torch.Tensor, vector: tuple[float, float, float] | torch.Te def _cross_z(rate: torch.Tensor, vector: torch.Tensor) -> torch.Tensor: - """Return ``(rate * z) x vector`` for per-environment angular rates about the world z axis.""" + """Return ``(rate * z) x vector``, the velocity induced by a rotation about the world z axis. + + 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) -def _expected_fixed_branching_state(articulation: Articulation, origins: torch.Tensor, com_b: torch.Tensor) -> dict: +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. 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. + + 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. + + 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))) @@ -913,20 +934,21 @@ def _expected_fixed_branching_state(articulation: Articulation, origins: torch.T return {key: torch.stack(values, dim=1) for key, values in expected.items()} -def _assert_jacobian_contract(articulation: Articulation, generalized_velocity: torch.Tensor): - """Check the link and center-of-mass Jacobians against body velocities and the mass-matrix contract.""" +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. + + 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. - body_velocities = { - "body_link_jacobian_w": data.body_link_lin_vel_w.torch.clone(), - "body_com_jacobian_w": data.body_com_lin_vel_w.torch.clone(), - } + 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_name, lin_vel in body_velocities.items(): - jacobian = getattr(data, jacobian_name).torch + 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) @@ -953,7 +975,7 @@ def _assert_jacobian_contract(articulation: Articulation, generalized_velocity: torch.testing.assert_close(generalized_energy, body_energy, atol=1e-5, rtol=1e-4) -def test_articulation_drive_and_dynamics(articulation_scene: _ArticulationScene): +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 @@ -981,11 +1003,11 @@ def test_articulation_drive_and_dynamics(articulation_scene: _ArticulationScene) expected_efforts = zeros.clone() expected_efforts[1, right_shoulder] = 10.0 torch.testing.assert_close( - _to_device_tensor(articulation.root_view.get_dof_position_targets(), device), + wp.to_torch(articulation.root_view.get_dof_position_targets()).to(device), expected_targets[:, joint_backend_to_user], ) torch.testing.assert_close( - _to_device_tensor(articulation.root_view.get_dof_actuation_forces(), device), + wp.to_torch(articulation.root_view.get_dof_actuation_forces()).to(device), expected_efforts[:, joint_backend_to_user], ) # The effort must reach the solver, not only its staging buffer. @@ -1072,7 +1094,7 @@ def _place_at_rest(articulation: Articulation, root_pose: torch.Tensor) -> None: articulation.reset() -def test_floating_articulation_root_writes_and_wrenches(articulation_scene: _ArticulationScene, monkeypatch): +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 @@ -1092,7 +1114,7 @@ def test_floating_articulation_root_writes_and_wrenches(articulation_scene: _Art 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(_to_device_tensor(floating.root_view.get_root_transforms(), device), 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 @@ -1106,8 +1128,8 @@ def test_floating_articulation_root_writes_and_wrenches(articulation_scene: _Art 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(_to_device_tensor(floating.root_view.get_coms(), device), backend_coms) - torch.testing.assert_close(_to_device_tensor(reordered.root_view.get_coms(), device), backend_coms) + 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) @@ -1120,14 +1142,12 @@ def test_floating_articulation_root_writes_and_wrenches(articulation_scene: _Art 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 = _to_device_tensor(reordered.root_view.get_root_transforms(), device).clone() + 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( - reordered_transforms, _to_device_tensor(floating.root_view.get_root_transforms(), device) - ) - torch.testing.assert_close( - _to_device_tensor(reordered.root_view.get_root_velocities(), device), - _to_device_tensor(floating.root_view.get_root_velocities(), 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. @@ -1220,7 +1240,7 @@ def test_floating_articulation_root_writes_and_wrenches(articulation_scene: _Art actuator_reset = actuator.reset reset_env_ids = [] - def record_actuator_reset(env_ids=None): + def record_actuator_reset(env_ids=None) -> None: reset_env_ids.append(env_ids) actuator_reset(env_ids) @@ -1244,7 +1264,7 @@ def record_actuator_reset(env_ids=None): assert torch.count_nonzero(composer.out_torque_b.torch) == 0 -def test_spatial_tendon_properties_round_trip(articulation_scene: _ArticulationScene): +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 @@ -1266,7 +1286,7 @@ def test_spatial_tendon_properties_round_trip(articulation_scene: _ArticulationS getattr(articulation, f"set_spatial_tendon_{name}_index")(**{name: value}) articulation.write_spatial_tendon_properties_to_sim_index() for name, getter in getters.items(): - torch.testing.assert_close(_to_device_tensor(getter(), device), values[name]) + torch.testing.assert_close(wp.to_torch(getter()).to(device), values[name]) # Partial writes change environment 1 in the data and in the solver and preserve environment 0. env_ids = torch.tensor([1], dtype=torch.int32, device=device) @@ -1279,4 +1299,4 @@ def test_spatial_tendon_properties_round_trip(articulation_scene: _ArticulationS torch.testing.assert_close(getattr(articulation.data, f"spatial_tendon_{name}").torch, values[name]) articulation.write_spatial_tendon_properties_to_sim_index(env_ids=env_ids) for name, getter in getters.items(): - torch.testing.assert_close(_to_device_tensor(getter(), device), values[name]) + torch.testing.assert_close(wp.to_torch(getter()).to(device), values[name]) diff --git a/source/isaaclab_physx/test/assets/test_deformable_object.py b/source/isaaclab_physx/test/assets/test_deformable_object.py index 1f2404ccb2bd..edda3ea457f8 100644 --- a/source/isaaclab_physx/test/assets/test_deformable_object.py +++ b/source/isaaclab_physx/test/assets/test_deformable_object.py @@ -19,6 +19,7 @@ launch_test_simulation() import sys +from collections.abc import Iterator from dataclasses import dataclass import pytest @@ -27,7 +28,7 @@ from isaaclab_physx.assets import DeformableObject from isaaclab_physx.sim import PhysxDeformableBodyMaterialCfg, PhysxSurfaceDeformableBodyMaterialCfg -from pxr import Gf, Sdf, UsdGeom, UsdShade +from pxr import Gf, Sdf, Usd, UsdGeom, UsdShade import isaaclab.sim as sim_utils from isaaclab.assets import DeformableObjectCfg @@ -42,19 +43,28 @@ _SURFACE_TRIANGLES = [(0, 1, 2), (0, 2, 3)] -def _add_api_schemas(prim, schemas: list[str]) -> None: - """Set the applied API schemas of a prim explicitly.""" - schemas_op = Sdf.TokenListOp() - schemas_op.explicitItems = schemas - prim.SetMetadata("apiSchemas", schemas_op) - - -def _author_deformable_body(body_prim, points: list[Gf.Vec3f], sim_api: str, material_cfg) -> None: - """Author the deformable body schemas, rest shape, and a bound physics material on a mesh prim.""" - _add_api_schemas( - body_prim, - ["OmniPhysicsDeformableBodyAPI", sim_api, "OmniPhysicsDeformablePoseAPI:default", "PhysicsCollisionAPI"], - ) +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: + 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. + """ + 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"] @@ -120,7 +130,7 @@ def _spawn_surface_deformables(root: str, y_offset: float) -> DeformableObject: @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", gravity_enabled=False) as sim: sim._app_control_on_stop_handle = None @@ -154,7 +164,7 @@ def step(self, num_steps: int = 1) -> None: @pytest.fixture(scope="module", params=test_devices(DeviceScope.CUDA)) -def deformable_scene(request) -> _DeformableScene: +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: @@ -166,7 +176,7 @@ def deformable_scene(request) -> _DeformableScene: yield _DeformableScene(sim=sim, device=device, volume=volume, surface=surface, refcounts=refcounts) -def test_deformable_initialization(deformable_scene: _DeformableScene): +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 ( @@ -193,7 +203,7 @@ def test_deformable_initialization(deformable_scene: _DeformableScene): scene.surface.write_nodal_kinematic_target_to_sim_index(dummy_targets) -def test_deformable_nodal_state_writes(deformable_scene: _DeformableScene): +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 @@ -231,7 +241,7 @@ def test_deformable_nodal_state_writes(deformable_scene: _DeformableScene): torch.testing.assert_close(volume.data.root_pos_w.torch, nodal_pos.mean(dim=1)) -def test_volume_kinematic_targets_hold_selected_nodes(deformable_scene: _DeformableScene): +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 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 0e4c0e482602..0f2710abe926 100644 --- a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py +++ b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py @@ -686,7 +686,3 @@ def tearDownClass(cls): 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 ee3d89174245..a3a81756c83c 100644 --- a/source/isaaclab_physx/test/assets/test_rigid_object.py +++ b/source/isaaclab_physx/test/assets/test_rigid_object.py @@ -20,8 +20,10 @@ 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 @@ -29,6 +31,8 @@ from isaaclab_physx.assets import RigidObject from isaaclab_physx.sim.schemas import PhysxRigidBodyCfg +import omni.kit.app +import omni.physx from pxr import UsdPhysics import isaaclab.sim as sim_utils @@ -81,14 +85,14 @@ def _yaw_quat(angle: float) -> tuple[float, float, float, float]: pytest.param( "articulation_root", marks=pytest.mark.xfail( - strict=True, reason="The PhysX rigid object no longer rejects an enabled articulation root." + strict=True, 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_initialization_rejects_invalid_rigid_body(device, api: Literal["none", "articulation_root"]): +def test_initialization_rejects_invalid_rigid_body(device, api: Literal["none", "articulation_root"]) -> None: """Initialization fails without a rigid body and when the rigid body is an articulation root.""" with build_simulation_context(device=device, auto_add_lighting=True) as sim: sim._app_control_on_stop_handle = None @@ -116,7 +120,7 @@ def test_initialization_rejects_invalid_rigid_body(device, api: Literal["none", @pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) @pytest.mark.parametrize("gravity_enabled", [True, False]) @pytest.mark.isaacsim_ci -def test_gravity_vec_w(device, gravity_enabled): +def test_gravity_vec_w(device, gravity_enabled) -> None: """Test that gravity vector direction is set correctly for the rigid object.""" with build_simulation_context(device=device, gravity_enabled=gravity_enabled) as sim: sim._app_control_on_stop_handle = None @@ -154,13 +158,8 @@ def test_gravity_vec_w(device, gravity_enabled): @pytest.mark.isaacsim_ci @pytest.mark.parametrize("device", test_devices()) -def test_warmup_loads_physics_once(device): +def test_warmup_loads_physics_once(device) -> None: """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 - 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 sim_utils.create_prim("/World/Table_0", "Xform", translation=(0.0, 0.0, 1.0)) @@ -226,7 +225,7 @@ def place_cubes_at_rest(self, root_pose: torch.Tensor) -> None: @pytest.fixture(scope="module", params=test_devices()) -def rigid_object_scene(request) -> _RigidObjectScene: +def rigid_object_scene(request) -> Iterator[_RigidObjectScene]: """Initialize the composite rigid-object scene once per device.""" device = request.param with build_simulation_context(device=device, gravity_enabled=True) as sim: @@ -246,7 +245,7 @@ def rigid_object_scene(request) -> _RigidObjectScene: ) -def test_rigid_object_initialization(rigid_object_scene: _RigidObjectScene): +def test_rigid_object_initialization(rigid_object_scene: _RigidObjectScene) -> None: """Initialize local cubes with the expected buffers; kinematic cubes hold their default pose under gravity.""" scene = rigid_object_scene for name, rigid_object in (("cubes", scene.cubes), ("kinematic", scene.kinematic)): @@ -269,7 +268,7 @@ def test_rigid_object_initialization(rigid_object_scene: _RigidObjectScene): torch.testing.assert_close(kinematic.data.root_com_vel_w.torch, kinematic.data.default_root_vel.torch) -def test_rigid_object_inertial_properties(rigid_object_scene: _RigidObjectScene): +def test_rigid_object_inertial_properties(rigid_object_scene: _RigidObjectScene) -> None: """Mass, center-of-mass, and inertia writes reach only the selected PhysX entries and survive a step.""" scene = rigid_object_scene device = scene.device @@ -309,12 +308,14 @@ def test_rigid_object_inertial_properties(rigid_object_scene: _RigidObjectScene) torch.testing.assert_close(wp.to_torch(raw).to(device).reshape(expected[name].shape), expected[name]) -def test_rigid_object_root_state_writes(rigid_object_scene: _RigidObjectScene): +def test_rigid_object_root_state_writes(rigid_object_scene: _RigidObjectScene) -> None: """Root writes round-trip through the written frame and keep the other frame consistent with the offset.""" scene = rigid_object_scene device = scene.device cube_object = scene.cubes - rest_pose = torch.cat((scene.origins["cubes"], torch.tensor([[0.0, 0.0, 0.0, 1.0]] * _NUM_ENVS, device=device)), -1) + 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) @@ -412,7 +413,7 @@ def test_rigid_object_root_state_writes(rigid_object_scene: _RigidObjectScene): torch.testing.assert_close(root_com_vel_w[:, 3:], root_link_vel_w[:, 3:]) -def test_rigid_object_wrench_delivery_and_reset(rigid_object_scene: _RigidObjectScene): +def test_rigid_object_wrench_delivery_and_reset(rigid_object_scene: _RigidObjectScene) -> None: """External wrenches act in the frame they are given in on the selected root; reset clears them.""" scene = rigid_object_scene device = scene.device @@ -463,7 +464,9 @@ def test_rigid_object_wrench_delivery_and_reset(rigid_object_scene: _RigidObject # 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]] * 2, device=device))[:, 1] + 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 ) diff --git a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py index 86cfb0972106..794e2f2a98ab 100644 --- a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py @@ -20,6 +20,7 @@ import math import sys +from collections.abc import Iterator from dataclasses import dataclass import pytest @@ -97,7 +98,7 @@ def _yaw_quat(angle: float) -> tuple[float, float, float, float]: @pytest.fixture -def sim(request): +def sim(request) -> Iterator[SimulationContext]: """Create a function-scoped simulation context for tests that own their scene.""" device = request.getfixturevalue("device") gravity_enabled = request.getfixturevalue("gravity_enabled") if "gravity_enabled" in request.fixturenames else True @@ -113,7 +114,7 @@ def sim(request): @pytest.mark.parametrize("num_cubes", [2]) @pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -def test_initialization_with_no_rigid_body(sim, num_cubes, device): +def test_initialization_with_no_rigid_body(sim, num_cubes, device) -> None: """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) @@ -127,7 +128,7 @@ def test_initialization_with_no_rigid_body(sim, num_cubes, device): @pytest.mark.parametrize("device", test_devices()) @pytest.mark.parametrize("gravity_enabled", [False]) -def test_subset_write_reaches_selected_view_entry(sim, device, gravity_enabled): +def test_subset_write_reaches_selected_view_entry(sim, device, gravity_enabled) -> None: """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) @@ -151,7 +152,7 @@ def test_subset_write_reaches_selected_view_entry(sim, device, gravity_enabled): @pytest.mark.parametrize("device", test_devices()) -def test_inertial_property_subset_writes_reach_selected_view_entries(sim, device): +def test_inertial_property_subset_writes_reach_selected_view_entries(sim, device) -> None: """COM and inertia writes to a subset of (env, body) cells must reach only those PhysX view entries.""" num_envs, num_cubes = 2, 3 object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes) @@ -190,22 +191,19 @@ def read_view(values: wp.array, data_dim: int) -> torch.Tensor: @pytest.mark.parametrize("num_cubes", [2]) @pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) @pytest.mark.parametrize("gravity_enabled", [True, False]) -def test_gravity_vec_w(sim, num_envs, num_cubes, device, gravity_enabled): +def test_gravity_vec_w(sim, num_envs, num_cubes, device, gravity_enabled) -> None: """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) - # 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) @@ -215,7 +213,6 @@ def test_gravity_vec_w(sim, num_envs, num_cubes, device, gravity_enabled): 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) @@ -262,7 +259,7 @@ def place_cubes_at_rest(self, body_poses: torch.Tensor) -> None: @pytest.fixture(scope="module", params=test_devices()) -def collection_scene(request) -> _CollectionScene: +def collection_scene(request) -> Iterator[_CollectionScene]: """Initialize the composite collection scene once per device.""" device = request.param with build_simulation_context(device=device, gravity_enabled=True) as sim: @@ -293,7 +290,7 @@ def _rest_poses(scene: _CollectionScene, yaw: float = 0.0) -> torch.Tensor: return poses -def test_collection_initialization(collection_scene: _CollectionScene): +def test_collection_initialization(collection_scene: _CollectionScene) -> None: """Initialize local collections, including a single-cube one; kinematic cubes hold their pose under gravity.""" scene = collection_scene for name, collection, num_envs, num_cubes in ( @@ -320,7 +317,7 @@ def test_collection_initialization(collection_scene: _CollectionScene): torch.testing.assert_close(kinematic.data.body_link_vel_w.torch, default_body_vel) -def test_collection_state_writes(collection_scene: _CollectionScene): +def test_collection_state_writes(collection_scene: _CollectionScene) -> None: """Body state writes reach the selected body-major view entries in the frame they are given in.""" scene = collection_scene device = scene.device @@ -428,7 +425,7 @@ def test_collection_state_writes(collection_scene: _CollectionScene): torch.testing.assert_close(object_com_vel_w[..., 3:], object_link_vel_w[..., 3:]) -def test_collection_masses_reach_selected_view_entries(collection_scene: _CollectionScene): +def test_collection_masses_reach_selected_view_entries(collection_scene: _CollectionScene) -> None: """Mass writes to a non-sorted subset reach only the selected body-major view entries.""" scene = collection_scene device = scene.device @@ -446,7 +443,7 @@ def test_collection_masses_reach_selected_view_entries(collection_scene: _Collec collection.set_masses_index(masses=torch.ones((_NUM_ENVS, _NUM_CUBES), device=device)) -def test_collection_wrench_delivery_and_reset(collection_scene: _CollectionScene): +def test_collection_wrench_delivery_and_reset(collection_scene: _CollectionScene) -> None: """External wrenches act on the selected bodies in the frame they are given in; reset clears them.""" scene = collection_scene device = scene.device diff --git a/source/isaaclab_physx/test/assets/test_surface_gripper.py b/source/isaaclab_physx/test/assets/test_surface_gripper.py index ebfa1bfec7ef..b135510573b7 100644 --- a/source/isaaclab_physx/test/assets/test_surface_gripper.py +++ b/source/isaaclab_physx/test/assets/test_surface_gripper.py @@ -116,7 +116,6 @@ def test_close_and_open_command() -> None: # 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() - # step the simulation so the gripper reacts to the command for _ in range(3): sim.step() surface_gripper.update(sim.cfg.dt) diff --git a/source/isaaclab_physx/test/assets/unit/test_physx_surface_gripper.py b/source/isaaclab_physx/test/assets/unit/test_physx_surface_gripper.py index 3007be316949..4a20299a672d 100644 --- a/source/isaaclab_physx/test/assets/unit/test_physx_surface_gripper.py +++ b/source/isaaclab_physx/test/assets/unit/test_physx_surface_gripper.py @@ -13,8 +13,9 @@ from isaaclab_physx.assets.surface_gripper.surface_gripper import SurfaceGripper -def _gripper() -> SurfaceGripper: - """Create a three-environment CPU gripper around a recording view.""" +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 @@ -28,12 +29,7 @@ def _gripper() -> SurfaceGripper: apply_gripper_action=Mock(), set_surface_gripper_properties=Mock(), ) - return gripper - -def test_command_filter_and_partial_property_update_use_literal_view_payloads() -> None: - """Submit only open/close commands and preserve property selector ordering.""" - gripper = _gripper() gripper.set_grippers_command_index(wp.array([0.5, 0.0, -0.5], dtype=wp.float32, device="cpu")) gripper.write_data_to_sim() From 4fc854cad46b45b6122f629dc854ffe182f6ce2e Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 16:57:03 +0200 Subject: [PATCH 20/33] Sync with the reviewed asset-test series Bring the tracking branch to the union of the series branches after the independent review round. --- docs/source/concepts/ovphysx.rst | 4 +- pyproject.toml | 2 +- ...sset-tests-p13-physx-newton-actuators.skip | 0 .../test/actuators/test_newton_adapter.py | 36 ++- .../actuators/test_physx_actuator_runtime.py | 5 +- .../contract/_articulation_contract_cases.py | 50 ++-- .../test/assets/contract/_mock_backends.py | 9 +- ..._rigid_object_collection_contract_cases.py | 3 +- .../test/assets/contract/capabilities.py | 5 +- .../test/envs/mdp/test_terminations.py | 8 +- .../test/utils/test_wrench_composer.py | 4 +- .../utils/test_wrench_composer_integration.py | 14 +- .../test/articulation_test_utils.py | 2 + .../test/assets/test_articulation.py | 224 ++++++++++++------ .../assets/test_newton_actuators_newton.py | 40 +++- .../test/assets/test_rigid_object.py | 11 +- .../assets/test_rigid_object_collection.py | 14 +- .../test_newton_task_space_controllers.py | 24 +- .../antoiner-asset-tests-p05-ov-manager.rst | 9 +- .../isaaclab_ov/physics/ovphysx_manager.py | 89 ++++--- .../test/assets/test_articulation.py | 13 +- .../test/assets/test_rigid_object.py | 10 +- .../assets/test_rigid_object_collection.py | 13 +- .../physics/test_ovphysx_manager_lifecycle.py | 20 +- .../test/sensors/test_contact_sensor.py | 39 --- .../test/sensors/test_frame_transformer.py | 33 --- source/isaaclab_ov/test/sensors/test_imu.py | 28 --- .../test/sensors/test_joint_wrench_sensor.py | 34 --- source/isaaclab_ov/test/sensors/test_pva.py | 28 --- .../test/sim/test_views_xform_prim_ovphysx.py | 2 - ...ndomize_rigid_body_collider_offsets_mdp.py | 24 +- .../test_randomize_rigid_body_material_mdp.py | 25 +- ...oiner-fix-physx-collection-com-inertia.rst | 15 +- .../test/assets/test_articulation.py | 212 +++++++++++------ .../test/assets/test_deformable_object.py | 8 +- .../test/assets/test_rigid_object.py | 71 ++++-- .../assets/test_rigid_object_collection.py | 50 ++-- 37 files changed, 660 insertions(+), 518 deletions(-) create mode 100644 source/isaaclab/changelog.d/antoiner-asset-tests-p13-physx-newton-actuators.skip diff --git a/docs/source/concepts/ovphysx.rst b/docs/source/concepts/ovphysx.rst index c955be7704aa..7223c3c30e74 100644 --- a/docs/source/concepts/ovphysx.rst +++ b/docs/source/concepts/ovphysx.rst @@ -123,13 +123,13 @@ 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 .. tab-item:: isaaclab.sh / isaaclab.bat .. code-block:: bash - ./isaaclab.sh -p -m pytest source/isaaclab_ov/test/assets/test_rigid_object.py::test_initialization -k cpu + ./isaaclab.sh -p -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 10ef9a8d5506..6a6cb6c78472 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -362,7 +362,7 @@ 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)", + "device_split: re-invoke this file in a separate process per device (CPU, then GPU) in CI", "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-p13-physx-newton-actuators.skip b/source/isaaclab/changelog.d/antoiner-asset-tests-p13-physx-newton-actuators.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab/test/actuators/test_newton_adapter.py b/source/isaaclab/test/actuators/test_newton_adapter.py index 52841965c6ce..1cd1e319b92d 100644 --- a/source/isaaclab/test/actuators/test_newton_adapter.py +++ b/source/isaaclab/test/actuators/test_newton_adapter.py @@ -54,7 +54,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") @@ -65,7 +66,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", @@ -178,6 +183,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))], diff --git a/source/isaaclab/test/actuators/test_physx_actuator_runtime.py b/source/isaaclab/test/actuators/test_physx_actuator_runtime.py index 419a65b363f4..3bbdf0c35224 100644 --- a/source/isaaclab/test/actuators/test_physx_actuator_runtime.py +++ b/source/isaaclab/test/actuators/test_physx_actuator_runtime.py @@ -60,8 +60,10 @@ def test_compute_runs_eagerly_after_graph_capture_failure(monkeypatch: pytest.Mo """A graphable adapter whose capture fails still computes the current command eagerly.""" class _FailingCapture: + capture_count = 0 + def __init__(self, *args, **kwargs) -> None: - pass + type(self).capture_count += 1 def __enter__(self): raise RuntimeError("capture unavailable") @@ -81,6 +83,7 @@ def __exit__(self, exc_type, exc_value, traceback) -> bool: runtime.compute(collection, 0.01) # The failed capture is not retried, and every step falls back to the eager kernels. + assert _FailingCapture.capture_count == 1 assert runtime.native_actuator_graphs == () assert eager_compute.call_count == 2 eager_compute.assert_called_with(collection, 0.01) diff --git a/source/isaaclab/test/assets/contract/_articulation_contract_cases.py b/source/isaaclab/test/assets/contract/_articulation_contract_cases.py index 2ac86859e6bc..fda0c57eac11 100644 --- a/source/isaaclab/test/assets/contract/_articulation_contract_cases.py +++ b/source/isaaclab/test/assets/contract/_articulation_contract_cases.py @@ -47,7 +47,6 @@ def _check_proxy_array(arr, *, expected_shape: tuple, expected_dtype: type, name _api_backends = contract_backend("api") _data_backends = contract_backend("data") _writes_backends = contract_backend("writes") -_index_resolution_backends = contract_backend("index_resolution") _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 @@ -80,7 +79,7 @@ def test_resolve_env_ids_handles_tensor_view_shape(self, backend, device): assert resolved.data_ptr() == cached[selection].data_ptr() assert resolved.stride() == cached[selection].stride() - @_index_resolution_backends + @_api_backends def test_resolve_joint_ids_handles_tensor_view_shape(self, backend): art, _ = get_articulation(backend, num_joints=4, device="cpu") @@ -91,7 +90,7 @@ def test_resolve_joint_ids_handles_tensor_view_shape(self, backend): assert resolved_full.shape[0] == 4 assert resolved_view.shape[0] == 2 - @_index_resolution_backends + @_api_backends def test_resolve_body_ids_handles_tensor_view_shape(self, backend): art, _ = get_articulation(backend, num_bodies=4, device="cpu") @@ -1489,7 +1488,6 @@ class TestArticulationWritersTendon: ) def test_tendon_writer( self, - request, backend, device, selection, @@ -1504,16 +1502,6 @@ def test_tendon_writer( num_spatial_tendons, ): require_backend_capability(backend, _TENDON_FAMILIES[_TENDON_PROPERTY_FAMILY[getter]][1]) - if backend == "newton" and selection == "mask": - # Product bug: Newton's fixed-tendon mask setters slice full data without validating its shape, so an - # oversized input is silently truncated. Every assertion before the negative shape checks still runs. - request.applymarker( - pytest.mark.xfail( - raises=pytest.fail.Exception, - strict=True, - reason="Newton fixed-tendon mask setters accept data with extra environments", - ) - ) num_tendons = num_fixed_tendons if kind == "fixed" else num_spatial_tendons art, _ = get_articulation( backend, num_instances, num_joints, num_bodies, num_fixed_tendons, num_spatial_tendons, device @@ -1551,12 +1539,36 @@ def test_tendon_writer( 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 + + # Shape validation runs before any kernel launch, so one device covers it. + @_tendon_backends + @pytest.mark.parametrize("array_kind", ["torch", "warp"]) + @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_rejects_extra_environments( + self, request, backend, selection, array_kind, kind, method_base, kwarg, getter + ): + require_backend_capability(backend, _TENDON_FAMILIES[_TENDON_PROPERTY_FAMILY[getter]][1]) + if backend == "newton" and selection == "mask": + # Product bug: Newton's fixed-tendon mask setters slice full data without validating its shape, so an + # oversized input is silently truncated. + request.applymarker( + pytest.mark.xfail( + raises=pytest.fail.Exception, + strict=True, + reason="Newton fixed-tendon mask setters accept data with extra environments", + ) + ) + num_instances, num_joints, num_bodies, num_fixed_tendons, num_spatial_tendons = _TENDON_DIMS[1] + art, _ = get_articulation( + backend, num_instances, num_joints, num_bodies, num_fixed_tendons, num_spatial_tendons, "cpu" + ) + full_shape = (num_instances, num_fixed_tendons if kind == "fixed" else num_spatial_tendons) + wp_dtype = wp.vec2f if getter == "fixed_tendon_pos_limits" else wp.float32 + make_bad_data = _make_bad_data_torch if array_kind == "torch" else _make_bad_data_warp + with pytest.raises((AssertionError, RuntimeError)): - method(**{kwarg: _make_bad_data_warp(full_shape, device, wp_dtype)}) + getattr(art, f"{method_base}_{selection}")(**{kwarg: make_bad_data(full_shape, "cpu", wp_dtype)}) @_tendon_backends @_devices diff --git a/source/isaaclab/test/assets/contract/_mock_backends.py b/source/isaaclab/test/assets/contract/_mock_backends.py index 32c60dd1ae8a..1d9b14853d28 100644 --- a/source/isaaclab/test/assets/contract/_mock_backends.py +++ b/source/isaaclab/test/assets/contract/_mock_backends.py @@ -37,7 +37,7 @@ def install_physx_recording_setters(mock_view, storage_by_method: dict[str, str] """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 accept both layouts + 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: @@ -56,11 +56,10 @@ def setter(values: wp.array, indices: wp.array | None = None) -> None: 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) - if values_np.size == stored_np.size: - stored_np[rows] = values_np.reshape(stored_np.shape)[rows] - else: - stored_np[rows] = values_np.reshape(len(rows), *stored_np.shape[1:]) + 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)) diff --git a/source/isaaclab/test/assets/contract/_rigid_object_collection_contract_cases.py b/source/isaaclab/test/assets/contract/_rigid_object_collection_contract_cases.py index 20a956e79440..fd4ca35d083b 100644 --- a/source/isaaclab/test/assets/contract/_rigid_object_collection_contract_cases.py +++ b/source/isaaclab/test/assets/contract/_rigid_object_collection_contract_cases.py @@ -48,7 +48,6 @@ def _check_proxy_array(arr, *, expected_shape: tuple, expected_dtype: type, name _api_backends = contract_backend("api") _data_backends = contract_backend("data") _writes_backends = contract_backend("writes") -_index_resolution_backends = contract_backend("index_resolution") _devices = pytest.mark.parametrize("device", test_devices(DeviceScope.CPU_AND_DEFAULT_CUDA)) @@ -206,7 +205,7 @@ def test_resolve_env_ids_handles_tensor_view_shape(self, backend, device): assert resolved.data_ptr() == cached[selection].data_ptr() assert resolved.stride() == cached[selection].stride() - @_index_resolution_backends + @_api_backends def test_resolve_body_ids_handles_tensor_view_shape(self, backend): obj, _ = get_rigid_object_collection(backend, num_bodies=4, device="cpu") diff --git a/source/isaaclab/test/assets/contract/capabilities.py b/source/isaaclab/test/assets/contract/capabilities.py index 8c5dbaf2b07d..a8d6948ce036 100644 --- a/source/isaaclab/test/assets/contract/capabilities.py +++ b/source/isaaclab/test/assets/contract/capabilities.py @@ -50,12 +50,12 @@ class BackendStatus: BackendDeclaration( name="physx", required_modules=("carb", "isaaclab_physx"), - capabilities=_SHARED_CAPABILITIES | _PHYSX_FAMILY_CAPABILITIES | {"index_resolution"}, + capabilities=_SHARED_CAPABILITIES | _PHYSX_FAMILY_CAPABILITIES, ), BackendDeclaration( name="newton", required_modules=("isaaclab_newton",), - capabilities=_SHARED_CAPABILITIES | {"index_resolution"}, + capabilities=_SHARED_CAPABILITIES, unsupported={ "fixed_tendon_extended": "Newton does not implement fixed-tendon limit stiffness, rest length, or offset", "spatial_tendons": "Newton does not support spatial tendons", @@ -69,7 +69,6 @@ class BackendStatus: name="ovphysx", required_modules=("ovphysx", "isaaclab_ov"), capabilities=_SHARED_CAPABILITIES | _PHYSX_FAMILY_CAPABILITIES, - unsupported={"index_resolution": "OVPhysX does not expose the joint and body index-resolution helpers"}, # The mocked bindings allocate pinned host staging buffers even for CPU tensors. requires_cuda_runtime=True, ), diff --git a/source/isaaclab/test/envs/mdp/test_terminations.py b/source/isaaclab/test/envs/mdp/test_terminations.py index 681f613c6bde..0c89b5d712e1 100644 --- a/source/isaaclab/test/envs/mdp/test_terminations.py +++ b/source/isaaclab/test/envs/mdp/test_terminations.py @@ -16,9 +16,9 @@ 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) - # Environment 0 clips the selected joint, environment 1 clips nothing, environment 2 clips the selected joint - # downward, and environment 3 clips only joints that are not selected. - applied = torch.tensor([[10.0, 19.0, 30.0], [10.0, 20.0, 30.0], [10.0, 17.0, 30.0], [9.0, 20.0, 29.0]]) + # 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) @@ -26,6 +26,6 @@ def test_joint_effort_limit_terminates_only_environments_with_clipped_selected_j ) env = SimpleNamespace(scene={"robot": robot}) - result = joint_effort_out_of_limit(env, SceneEntityCfg("robot", joint_ids=[1])) + 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 eeaa0822a848..de7dce62bace 100644 --- a/source/isaaclab/test/utils/test_wrench_composer.py +++ b/source/isaaclab/test/utils/test_wrench_composer.py @@ -482,8 +482,8 @@ def test_composed_wrench_emits_deprecation_warning(deprecated_name: str, name: s 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) + # The deprecated property aliases the new output property. + assert result is getattr(composer, name) _assert_vectors(result, _grid({(1, 0): (1.0, 2.0, 3.0)})) diff --git a/source/isaaclab/test/utils/test_wrench_composer_integration.py b/source/isaaclab/test/utils/test_wrench_composer_integration.py index 4c49ffd879ce..4095291fbbee 100644 --- a/source/isaaclab/test/utils/test_wrench_composer_integration.py +++ b/source/isaaclab/test/utils/test_wrench_composer_integration.py @@ -19,11 +19,12 @@ 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.test.utils import test_devices pytestmark = pytest.mark.integration @@ -37,9 +38,12 @@ def _make_cube(name: str, position: tuple[float, float, float], rotation: tuple[ """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.RigidBodyPropertiesCfg(disable_gravity=True, linear_damping=0.0, angular_damping=0.0), - mass_props=sim_utils.MassPropertiesCfg(mass=1.0), - collision_props=sim_utils.CollisionPropertiesCfg(), + 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), ) cfg = RigidObjectCfg( prim_path=f"/World/{name}", spawn=spawn, init_state=RigidObjectCfg.InitialStateCfg(pos=position, rot=rotation) @@ -47,7 +51,7 @@ def _make_cube(name: str, position: tuple[float, float, float], rotation: tuple[ return RigidObject(cfg=cfg) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +@pytest.mark.parametrize("device", test_devices()) 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( diff --git a/source/isaaclab_newton/test/articulation_test_utils.py b/source/isaaclab_newton/test/articulation_test_utils.py index 6255acc80832..bc6729aa9eba 100644 --- a/source/isaaclab_newton/test/articulation_test_utils.py +++ b/source/isaaclab_newton/test/articulation_test_utils.py @@ -507,6 +507,8 @@ def PhysicsFixedJoint "Joint" { rel physics:body0 = rel physics:body1 = + point3f physics:localPos0 = (0, 0, 2) + point3f physics:localPos1 = (0, 0, 0) } } """ diff --git a/source/isaaclab_newton/test/assets/test_articulation.py b/source/isaaclab_newton/test/assets/test_articulation.py index 3a5e223d68ec..09c23243c02c 100644 --- a/source/isaaclab_newton/test/assets/test_articulation.py +++ b/source/isaaclab_newton/test/assets/test_articulation.py @@ -15,14 +15,18 @@ 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. -Configurations that fail initialization and the rebind scenario, which swaps the live Newton state, build -their own small scenes. +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. + +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 launch_test_simulation +from isaaclab.test.utils import DeviceScope, launch_test_simulation, test_devices launch_test_simulation(SimulationCfg(physics=NewtonCfg())) @@ -51,6 +55,7 @@ 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.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 SimulationContext, build_simulation_context @@ -75,7 +80,6 @@ pytestmark = [pytest.mark.integration, pytest.mark.kitless] -_DEVICE = "cpu" _GRAVITY = (0.0, 0.0, -9.81) _NEWTON_USER_ORDER_STATE_CACHES = ( @@ -231,6 +235,7 @@ class _Scene: articulations: dict[str, Articulation] cfgs: dict[str, ArticulationCfg] origins: torch.Tensor + device: str def step(self, *names: str, num_steps: int = 1) -> None: """Write the named islands' commands, step every island, and update the named islands.""" @@ -274,8 +279,9 @@ def rest(self, articulation: Articulation) -> None: ## +@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(invalid_field: str) -> None: +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", @@ -291,7 +297,7 @@ def test_invalid_articulation_cfg_fails_initialization(invalid_field: str) -> No 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: + 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. @@ -301,7 +307,8 @@ def test_invalid_articulation_cfg_fails_initialization(invalid_field: str) -> No sim.reset() -def test_newton_rebind_refreshes_ordered_state_and_preserves_lab_owned_actuator_gains() -> None: +@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 @@ -343,7 +350,7 @@ def test_newton_rebind_refreshes_ordered_state_and_preserves_lab_owned_actuator_ 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: + with build_simulation_context(sim_cfg=newton_sim_cfg(device, use_newton_actuators=False)) as sim: articulations = spawn_assets(cfgs) sim.reset() sim.step() @@ -595,20 +602,65 @@ def check_rebind(new_state: State, new_model: Model) -> None: return check_rebind +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +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. + + 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 = 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=_LEG_PUBLIC_JOINT_NAMES, + body_ordering=_LEG_PUBLIC_BODY_NAMES, + ) + 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] + ) + + ## # Composite scene. ## -@pytest.fixture(scope="module") -def composite_scene() -> Iterator[_Scene]: +@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.""" - sim_cfg = newton_sim_cfg(_DEVICE, use_newton_actuators=False, newton_contacts=True) + 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)) + yield _Scene(sim=sim, articulations=articulations, cfgs=cfgs, origins=env_origins(device), device=device) @pytest.fixture @@ -635,6 +687,7 @@ def test_articulation_initialization_and_partial_state(scene: _Scene) -> None: state refresh after joint writes without a step, and joint position limit writes that keep or clamp the default joint positions. """ + device = scene.device expected_fixed_base = { "floating": False, "fixed": True, @@ -679,12 +732,12 @@ def test_articulation_initialization_and_partial_state(scene: _Scene) -> 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) + 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) + 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 ) @@ -702,8 +755,8 @@ def test_articulation_initialization_and_partial_state(scene: _Scene) -> None: 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), + -5.0 - torch.rand(num_articulations, articulation.num_joints, device=device), + 5.0 + torch.rand(num_articulations, articulation.num_joints, device=device), ), dim=-1, ) @@ -714,10 +767,10 @@ def test_articulation_initialization_and_partial_state(scene: _Scene) -> None: torch.testing.assert_close(articulation.data.default_joint_pos.torch, default_joint_pos) # 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) + 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( @@ -763,13 +816,13 @@ def test_articulation_initialization_and_partial_state(scene: _Scene) -> None: ) # Set new joint limits with indexing that invalidate the selected default joint positions - env_ids = torch.arange(1, device=_DEVICE, dtype=torch.int32) + 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), + -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, ) @@ -797,10 +850,11 @@ def test_articulation_joint_and_body_properties_round_trip(scene: _Scene, monkey through both writer selectors, and actuator gains taken from the configuration or, when unset, from the USD drives. """ + 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) + 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 @@ -810,7 +864,7 @@ def record_model_change(change: ModelFlags) -> None: 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) + 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 ) @@ -822,7 +876,7 @@ def record_model_change(change: ModelFlags) -> None: notifications.clear() initial_mass = articulation.data.body_mass.torch.clone() - masses = torch.tensor([[3.0]], device=_DEVICE) + 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]) @@ -831,7 +885,7 @@ def record_model_change(change: ModelFlags) -> None: notifications.clear() initial_com = articulation.data.body_com_pos_b.torch.clone() - coms = torch.tensor([[[0.05, -0.02, 0.01]]], device=_DEVICE) + 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]) @@ -839,14 +893,15 @@ def record_model_change(change: ModelFlags) -> None: 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) + 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 inertial properties for the tests that drive this island + # 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) @@ -854,7 +909,7 @@ def record_model_change(change: ModelFlags) -> None: # 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) + 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 ) @@ -870,13 +925,13 @@ def record_model_change(change: ModelFlags) -> None: (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) + 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) + 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) @@ -889,9 +944,9 @@ def record_model_change(change: ModelFlags) -> None: # 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(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) @@ -947,6 +1002,7 @@ def test_floating_articulation_root_and_wrench_response(scene: _Scene, monkeypat 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. """ + device = scene.device articulation = scene.articulations["floating"] scene.rest(articulation) num_articulations = articulation.num_instances @@ -964,8 +1020,8 @@ def test_floating_articulation_root_and_wrench_response(scene: _Scene, monkeypat # 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) + 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 @@ -986,7 +1042,7 @@ def test_floating_articulation_root_and_wrench_response(scene: _Scene, monkeypat assert not torch.allclose(updated_com_w, original_com_w) def random_root_state(num_envs: int) -> torch.Tensor: - state = torch.rand(num_envs, 13, device=_DEVICE) + 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) @@ -1067,7 +1123,7 @@ def read_written_root_state() -> torch.Tensor: # 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)) + 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) @@ -1080,14 +1136,14 @@ def read_written_root_state() -> torch.Tensor: 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) + 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( - articulation.data.body_com_vel_w.torch[0], torch.zeros(2, 6, device=_DEVICE), atol=1e-6, rtol=0.0 + articulation.data.body_com_vel_w.torch[0], torch.zeros(2, 6, device=device), atol=1e-6, rtol=0.0 ) 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 @@ -1096,23 +1152,23 @@ def one_step_response() -> tuple[torch.Tensor, torch.Tensor]: # 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) + 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) + 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) + 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) + 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) + 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, @@ -1122,7 +1178,7 @@ def one_step_response() -> tuple[torch.Tensor, torch.Tensor]: torques=zeros[:, :1], positions=positions, env_ids=env_1, - body_ids=torch.tensor([root_idx], dtype=torch.int32, device=_DEVICE), + body_ids=torch.tensor([root_idx], dtype=torch.int32, device=device), is_global=is_global, ) lin_vel, ang_vel = one_step_response() @@ -1130,7 +1186,7 @@ def one_step_response() -> tuple[torch.Tensor, torch.Tensor]: 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) + 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"] @@ -1159,14 +1215,14 @@ def record_actuator_reset(env_ids: torch.Tensor | None = None) -> None: # A partial reset clears only the selected environment's wrenches num_bodies = articulation.num_bodies permanent_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), + forces=torch.ones((num_articulations, num_bodies, 3), device=device), + torques=torch.ones((num_articulations, num_bodies, 3), device=device), ) instantaneous_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), + forces=torch.ones((num_articulations, num_bodies, 3), device=device), + torques=torch.ones((num_articulations, num_bodies, 3), device=device), ) - articulation.reset(env_ids=torch.tensor([0], device=_DEVICE)) + articulation.reset(env_ids=torch.tensor([0], device=device)) assert instantaneous_composer.active assert permanent_composer.active assert torch.count_nonzero(instantaneous_composer.out_force_b.torch) == num_bodies * 3 @@ -1190,7 +1246,7 @@ def test_branching_fixture_physx_ordering_reorders_newton_to_bfs(scene: _Scene) selected inertial-property writes keep Newton's inverse arrays current under the body ordering. """ articulation = scene.articulations["branching"] - device = _DEVICE + device = scene.device # 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 @@ -1257,6 +1313,7 @@ def test_newton_ordered_state_publishes_in_step_and_refreshes_same_timestamp_wri 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. """ + device = scene.device articulation = scene.articulations["ordered"] scene.rest(articulation) data = articulation.data @@ -1299,7 +1356,7 @@ def test_newton_ordered_state_publishes_in_step_and_refreshes_same_timestamp_wri # 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) + 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) @@ -1316,6 +1373,7 @@ def test_write_data_to_sim_writes_joint_targets_in_backend_order(scene: _Scene, 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. """ + device = scene.device articulation = scene.articulations[island] scene.rest(articulation) assert (articulation.data.joint_ordering is not None) is (island == "ordered") @@ -1323,7 +1381,7 @@ def test_write_data_to_sim_writes_joint_targets_in_backend_order(scene: _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) + 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() @@ -1416,7 +1474,7 @@ def test_set_material_properties(scene: _Scene) -> None: """ articulation = scene.articulations["floating"] num_articulations = articulation.num_instances - device = _DEVICE + device = scene.device # Resolve this articulation's shapes per environment and body from the flat Newton model. model = SimulationManager.get_model() @@ -1438,6 +1496,8 @@ def robot_shapes(env_index: int, selected_body_names: list[str] | None = None) - 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. for body_subset, friction, restitution_value in ((True, 0.55, 0.15), (False, 0.65, 0.25)): @@ -1493,6 +1553,13 @@ 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]) + # 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) + 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. @@ -1504,12 +1571,13 @@ def test_hand_with_tendons_targets_only_given_envs_and_properties_reach_solver(s Written fixed tendon stiffness, damping, and position limits reach the MuJoCo solver through both the index and the mask setters and writers. """ + device = scene.device articulation = scene.articulations["tendon"] scene.rest(articulation) num_articulations = articulation.num_instances assert articulation.num_fixed_tendons == 2 - target = torch.full((1, articulation.num_fixed_tendons), 1.0, dtype=torch.float32, device=_DEVICE) + 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) @@ -1526,13 +1594,13 @@ def test_hand_with_tendons_targets_only_given_envs_and_properties_reach_solver(s articulation.set_fixed_tendon_position_target_index(target=torch.zeros_like(target), env_ids=[0]) shape = (num_articulations, articulation.num_fixed_tendons) - limits = torch.tensor([-0.1, 0.2], device=_DEVICE).expand(*shape, 2) + 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() - 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_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") @@ -1561,6 +1629,7 @@ def test_hand_with_tendons_targets_only_given_envs_and_properties_reach_solver(s articulation.write_fixed_tendon_properties_to_sim_index() +@pytest.mark.isaacsim_ci @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. @@ -1597,7 +1666,7 @@ def test_get_gravity_compensation_forces_matches_jacobian_gravity(scene: _Scene, """ articulation = scene.articulations[island] num_articulations = articulation.num_instances - device = _DEVICE + device = scene.device scene.rest(articulation) # Sanity: the islands have different DoF counts, so a regression to model-wide sizing would manifest as wrong @@ -1745,6 +1814,7 @@ def test_get_gravity_compensation_forces_matches_jacobian_gravity(scene: _Scene, scene.rest(articulation) +@pytest.mark.isaacsim_ci def test_get_gravity_compensation_forces_static_equilibrium(scene: _Scene) -> None: """Newton accuracy: ``τ_gc`` must hold a passive manipulator in static equilibrium. @@ -1753,11 +1823,12 @@ def test_get_gravity_compensation_forces_static_equilibrium(scene: _Scene) -> No target IS the joint torque applied and no PD spring-damper masks the gravity-compensation signal. Sign errors, frame errors, and DoF-ordering errors all surface as joint drift. """ + device = scene.device articulation = scene.articulations["passive"] scene.rest(articulation) # a configuration that loads the prismatic and the revolute joints with gravity q = articulation.data.default_joint_pos.torch.clone() - q[:, 3:] = torch.tensor([0.4, -0.3, 0.5], device=_DEVICE) + q[:, 3:] = torch.tensor([0.4, -0.3, 0.5], device=device) articulation.write_joint_position_to_sim_index(position=q) articulation.update(scene.sim.cfg.dt) init_q = articulation.data.joint_pos.torch.clone() @@ -1793,7 +1864,7 @@ def test_body_root_state(scene: _Scene) -> None: scene.rest(articulation) num_articulations = articulation.num_instances num_bodies = articulation.num_bodies - device = _DEVICE + device = scene.device # Resolve body indices by name (ordering may differ across physics backends) root_idx = articulation.body_names.index("CenterPivot") @@ -1860,6 +1931,7 @@ def test_body_root_state(scene: _Scene) -> None: scene.rest(articulation) +@pytest.mark.isaacsim_ci 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``. @@ -1870,7 +1942,7 @@ def test_gravity_vec_w_tracks_model_gravity(scene: _Scene) -> None: articulation = scene.articulations["floating"] scene.rest(articulation) num_articulations = articulation.num_instances - device = _DEVICE + device = scene.device # GRAVITY_VEC_W must share storage with Newton's per-env gravity array. model = SimulationManager.get_model() @@ -1879,15 +1951,18 @@ def test_gravity_vec_w_tracks_model_gravity(scene: _Scene) -> None: assert articulation.data.GRAVITY_VEC_W.warp.ptr == model_gravity_arr.ptr assert articulation.data.GRAVITY_VEC_W.shape == (num_articulations,) - # Mutate model.gravity per-env in place, as randomize_physics_scene_gravity does. + # 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)): - wp.to_torch(model_gravity_arr).copy_(new_gravity) - SimulationManager.add_model_change(ModelFlags.MODEL_PROPERTIES) + 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) @@ -1918,6 +1993,7 @@ def test_body_q_consistent_after_root_write(scene: _Scene) -> None: 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. """ + device = scene.device articulation = scene.articulations["floating"] scene.rest(articulation) model = SimulationManager.get_model() @@ -1934,7 +2010,7 @@ def test_body_q_consistent_after_root_write(scene: _Scene) -> None: 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), + env_ids=torch.tensor([0], device=device, dtype=torch.int32), ) # Patch _simulate_physics_only to capture body_q before collide runs 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 c6d6319952c7..b1ec09f4e457 100644 --- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py +++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py @@ -30,7 +30,7 @@ import os import sys from collections.abc import Iterator -from dataclasses import dataclass, field +from dataclasses import dataclass from pathlib import Path import pytest @@ -188,8 +188,8 @@ class _Island: """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.""" - tolerances: dict[str, float] = field(default_factory=dict) - """Overrides of the equivalence oracle tolerances.""" + 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]: @@ -200,7 +200,7 @@ def _islands(mlp_path: str | None = None, lstm_path: str | None = None) -> dict[ "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, tolerances={"torque_atol": 0.5}), + "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 @@ -457,9 +457,18 @@ def newton_run(device: str, lab_run: dict, decimated_runs: dict) -> Iterator[_Ru os.unlink(lstm_path) -def _assert_equivalent(lab_result: dict, newton_result: dict, **tolerances: float) -> None: - """Run the shared trajectory and telemetry oracles on one island's Lab and Newton rollouts.""" - oracle = type("_Oracle", (EquivalenceAssertionsMixin,), tolerances)() +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. + + 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]. + """ + oracle = EquivalenceAssertionsMixin() + oracle.torque_atol = torque_atol oracle.lab_result = lab_result oracle.newton_result = newton_result oracle.test_joint_positions_match() @@ -484,7 +493,9 @@ def test_newton_actuators_match_lab_actuators(lab_run: dict, newton_run: _Run, i 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], **_islands()[island].tolerances) + _assert_equivalent( + lab_run["results"][island], newton_run.results[island], torque_atol=_islands()[island].torque_atol + ) def test_decimated_remotized_pd_matches_lab_actuators(decimated_runs: dict) -> None: @@ -715,15 +726,22 @@ def test_randomize_actuator_gains_reaches_newton_controllers(newton_run: _Run) - 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].""" + """Newton: ``num_pushes`` zeroes for env 0's DOFs only after reset of [0]. + + 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. + """ articulation = newton_run.articulations["delayed"] adapter = SimulationManager._adapter assert adapter is not None - # Find the delayed actuators (only they have delay_state). + 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 st is not None and st.delay_state is not None + 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" diff --git a/source/isaaclab_newton/test/assets/test_rigid_object.py b/source/isaaclab_newton/test/assets/test_rigid_object.py index e12dcb92eda1..7c3b8e9ddcc5 100644 --- a/source/isaaclab_newton/test/assets/test_rigid_object.py +++ b/source/isaaclab_newton/test/assets/test_rigid_object.py @@ -9,7 +9,8 @@ """Kitless real-Newton rigid-object coverage on one persistent two-environment scene per device. The cubes are locally spawned shapes. 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 case build their own scenes. +its own duration. Configurations that fail initialization and the single-instance case 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 @@ -76,6 +77,7 @@ def _cube_cfg(height: float = 1.0) -> RigidObjectCfg: ## +@pytest.mark.isaacsim_ci @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. @@ -187,6 +189,7 @@ def scene(shared_scene: _Scene) -> _Scene: return shared_scene +@pytest.mark.isaacsim_ci def test_external_force_on_single_body(scene: _Scene) -> None: """Test initialization and external forces on the base of the object. @@ -307,6 +310,7 @@ def reset_cubes(): assert torch.all(cube_object.data.root_pos_w.torch[1::2, 2] < 1.0) +@pytest.mark.isaacsim_ci 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 @@ -351,6 +355,7 @@ def test_rigid_body_set_material_properties(scene: _Scene) -> None: np.testing.assert_array_equal(restitution[shapes], original_restitution[shapes]) +@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 @@ -406,6 +411,7 @@ def test_rigid_body_set_mass(scene: _Scene) -> None: torch.testing.assert_close(masses, cube_object.data.body_mass.torch[env_ids][:, body_ids]) +@pytest.mark.isaacsim_ci 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``. @@ -458,6 +464,7 @@ def test_gravity_vec_w_tracks_model_gravity(scene: _Scene) -> None: torch.testing.assert_close(cube_object.data.projected_gravity_b.torch, expected, atol=1e-5, rtol=1e-5) +@pytest.mark.isaacsim_ci 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.""" cube_object = scene.cube @@ -536,6 +543,7 @@ def test_body_root_state_properties(scene: _Scene) -> None: torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) +@pytest.mark.isaacsim_ci 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. @@ -648,6 +656,7 @@ def test_write_root_state(scene: _Scene) -> None: torch.testing.assert_close(cube_object.data.root_com_vel_w.torch, written_com_vel_w, rtol=1e-1, atol=1e-1) +@pytest.mark.isaacsim_ci 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. diff --git a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py index 7626bc5e5aef..56aa0918a4d5 100644 --- a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py @@ -11,7 +11,8 @@ Every environment holds three locally spawned cubes in the collection and an unrelated rigid body beside them, 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 -case build their own scenes. +case 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 @@ -147,7 +148,7 @@ def step(self, num_steps: int = 1) -> None: self.collection.update(self.sim.cfg.dt) def rest(self) -> None: - """Put the cubes at rest at their configured poses.""" + """Put the cubes and the sibling body at rest at their configured poses.""" 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) @@ -155,6 +156,13 @@ def rest(self) -> None: body_velocities=torch.zeros_like(self.collection.data.default_body_vel.torch) ) self.collection.reset() + sibling_pose = self.sibling.data.default_root_pose.torch.clone() + sibling_pose[:, :3] += self.origins + self.sibling.write_root_pose_to_sim_index(root_pose=sibling_pose) + self.sibling.write_root_velocity_to_sim_index( + root_velocity=torch.zeros_like(self.sibling.data.default_root_vel.torch) + ) + self.sibling.reset() @pytest.fixture(scope="module", params=test_devices()) @@ -368,6 +376,7 @@ def reset_objects(): assert torch.all(object_collection.data.body_link_pos_w.torch[:, 1::2, 2] < 1.0) +@pytest.mark.isaacsim_ci 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``. @@ -581,6 +590,7 @@ def test_write_object_state(scene: _Scene) -> None: 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. diff --git a/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py b/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py index 52ddbe6c7d34..c397712f1257 100644 --- a/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py +++ b/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py @@ -34,6 +34,7 @@ from isaaclab.assets import ArticulationCfg from isaaclab.controllers import OperationalSpaceController, OperationalSpaceControllerCfg from isaaclab.sim import SimulationContext, build_simulation_context +from isaaclab.test.utils import DeviceScope, test_devices from isaaclab.utils.math import compute_pose_error, matrix_from_quat, quat_inv, subtract_frame_transforms sys.path.insert(0, str(Path(__file__).resolve().parents[1])) @@ -42,14 +43,13 @@ pytestmark = [pytest.mark.integration, pytest.mark.kitless] -_DEVICE = "cpu" - @dataclass class _Chain: """A passive six-DOF chain and the end-effector indices the controller reads.""" sim: SimulationContext + device: str robot: Articulation ee_frame_idx: int ee_jacobi_idx: int @@ -68,8 +68,8 @@ def rest(self) -> None: self.robot.update(self.sim.cfg.dt) -@pytest.fixture(scope="module") -def chain() -> Iterator[_Chain]: +@pytest.fixture(scope="module", params=test_devices(DeviceScope.CPU)) +def chain(request: pytest.FixtureRequest) -> Iterator[_Chain]: """Build the unpowered chain once, at a configuration away from the joint limits.""" robot_cfg = ArticulationCfg( prim_path="/World/Env_[^/]*/Robot", @@ -79,13 +79,15 @@ def chain() -> Iterator[_Chain]: ), actuators={"arm": ImplicitActuatorCfg(joint_names_expr=["Joint_.*"], stiffness=0.0, damping=0.0)}, ) - with build_simulation_context(sim_cfg=newton_sim_cfg(_DEVICE)) as sim: + device = request.param + with build_simulation_context(sim_cfg=newton_sim_cfg(device)) as sim: robot = spawn_assets({"robot": robot_cfg})["robot"] sim.reset() assert robot.is_initialized and robot.is_fixed_base ee_frame_idx = robot.find_bodies("Link_5")[0][0] yield _Chain( sim=sim, + device=device, robot=robot, ee_frame_idx=ee_frame_idx, # the fixed root has no Jacobian row @@ -115,11 +117,11 @@ def _build_relative_pose_target( ) -> torch.Tensor: """Return the current end-effector pose in the root frame offset by ``delta_xyz`` [m], keeping its 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(delta_xyz, device=_DEVICE, dtype=initial_ee_pos_b.dtype) + target_pos_b = initial_ee_pos_b + initial_ee_pos_b.new_tensor(delta_xyz) return torch.cat([target_pos_b, initial_ee_quat_b], dim=-1) -def _make_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( @@ -132,7 +134,7 @@ def _make_osc() -> OperationalSpaceController: motion_damping_ratio_task=1.0, ), num_envs=NUM_ENVS, - device=_DEVICE, + device=device, ) @@ -173,6 +175,7 @@ def _run_osc( return pos_history, rot_history +@pytest.mark.isaacsim_ci def test_osc_tracking_accuracy(chain: _Chain) -> None: """OSC pose tracking sentinel for the Jacobian and mass-matrix bridge. @@ -182,7 +185,7 @@ def test_osc_tracking_accuracy(chain: _Chain) -> None: """ chain.rest() target_pose_b = _build_relative_pose_target(chain.robot, chain.ee_frame_idx, (0.05, 0.0, 0.0)) - pos_history, rot_history = _run_osc(chain, _make_osc(), target_pose_b, 300, gravity=False) + pos_history, rot_history = _run_osc(chain, _make_osc(chain.device), target_pose_b, 300, gravity=False) pos_mean = sum(pos_history[-200:]) / 200 rot_mean = sum(rot_history[-200:]) / 200 @@ -195,6 +198,7 @@ def test_osc_tracking_accuracy(chain: _Chain) -> None: assert rot_mean < 5e-2, f"OSC rot_mean {rot_mean:.5f} > 0.05 rad — bridge regression?" +@pytest.mark.isaacsim_ci def test_osc_gravity_compensation_precision(chain: _Chain) -> None: """Two-phase EE hold: gravity sag without compensation, tight hold with it. @@ -211,7 +215,7 @@ def test_osc_gravity_compensation_precision(chain: _Chain) -> None: phase 2. Both phases reach a true steady state, enforced by tail-half stationarity guards. """ chain.rest() - osc = _make_osc() + osc = _make_osc(chain.device) # Hold the initial EE pose: phase-1 steady-state error is pure gravity sag. target_pose_b = _build_relative_pose_target(chain.robot, chain.ee_frame_idx, (0.0, 0.0, 0.0)) diff --git a/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p05-ov-manager.rst b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p05-ov-manager.rst index a0e44b9aec14..39a071784d95 100644 --- a/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p05-ov-manager.rst +++ b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p05-ov-manager.rst @@ -2,6 +2,11 @@ Fixed ^^^^^ * Fixed :class:`~isaaclab_ov.physics.OvPhysxManager` refusing to start a CPU simulation after a CUDA simulation, or - the reverse, in the same process. CPU scenes selected CPU dynamics through their physics-scene attributes instead of - the OVPhysX process-wide CPU-only mode, which cannot be reverted. To keep OVPhysX off CUDA entirely, set + the reverse, in the same process. + +Changed +^^^^^^^ + +* 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/isaaclab_ov/physics/ovphysx_manager.py b/source/isaaclab_ov/isaaclab_ov/physics/ovphysx_manager.py index e18f9a44f449..7183572d2f0c 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 @@ -85,6 +85,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__) @@ -245,7 +269,7 @@ def __init__(self, cfg: OvPhysxBackendCfg): } if is_gpu: carbonite_overrides.update({"/physics/suppressReadback": True, "/physics/suppressFabricUpdate": True}) - # CPU scenes select CPU dynamics through their scene attributes (see _configure_physx_scene_prim). + # 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( @@ -851,10 +875,9 @@ def _warmup_and_load(cls) -> None: ovphysx_device = "gpu" if "cuda" in PhysicsManager._device else "cpu" scene_prim = sim.stage.GetPrimAtPath(sim.cfg.physics_prim_path) - if scene_prim.IsValid(): - if cls._active_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._active_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) # Flatten the current USD stage to USDA text so OVStage can populate it # without an intermediate file. @@ -945,20 +968,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, solver-determinism/accuracy, and dynamics/broadphase - attributes are applied regardless of device. GPU scenes use GPU dynamics and the GPU - broadphase; CPU scenes use CPU dynamics and the MBP broadphase. The schema default enables - GPU dynamics, so CPU scenes must author it explicitly. The GPU buffer-capacity attributes - are applied only to GPU scenes. + 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. @@ -966,13 +1005,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) # Propagate scene query support from SimulationCfg so omni.physx creates # the scene with the correct query mode. OvPhysxCfg does not carry this field. @@ -990,13 +1023,7 @@ def _configure_physx_scene_prim(scene_prim, cfg, device: str) -> None: cfg.enable_external_forces_every_iteration ) - 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" - ) - - if use_gpu and cfg is not None: + 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), diff --git a/source/isaaclab_ov/test/assets/test_articulation.py b/source/isaaclab_ov/test/assets/test_articulation.py index afee01d99612..1cfeb8162e74 100644 --- a/source/isaaclab_ov/test/assets/test_articulation.py +++ b/source/isaaclab_ov/test/assets/test_articulation.py @@ -9,14 +9,17 @@ """Real OVPhysX articulation coverage on one module-scoped composite scene per device. Each device builds one scene of isolated articulation islands from the local branching fixture, -resets it once, and keeps it alive for every test in the module. Each island belongs to one test -unless noted, so the tests do not depend on each other's order. Partial writes target the last -environment and prove the other environments keep their real backend state. +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. Initialization failures need their own -scenes, so those tests run first and build them. +staging of CPU-only bindings and the native Newton actuator path. + +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 diff --git a/source/isaaclab_ov/test/assets/test_rigid_object.py b/source/isaaclab_ov/test/assets/test_rigid_object.py index 9cdd281b12bd..1973391af9c2 100644 --- a/source/isaaclab_ov/test/assets/test_rigid_object.py +++ b/source/isaaclab_ov/test/assets/test_rigid_object.py @@ -9,8 +9,12 @@ """Real OVPhysX rigid-object coverage on one module-scoped scene per device. Each device builds one scene of locally authored cuboid pairs, resets it once, and keeps it alive for -every test that uses it. Each pair belongs to one test, so the tests do not depend on each other's -order. Initialization failures need their own scenes and therefore build them before any shared scene. +every test that uses it. Each pair belongs to one test, so those tests do not depend on each other's +order. + +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 @@ -150,6 +154,8 @@ def test_rigid_object_real_ovphysx_seams(scene: _RigidScene) -> None: 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]) diff --git a/source/isaaclab_ov/test/assets/test_rigid_object_collection.py b/source/isaaclab_ov/test/assets/test_rigid_object_collection.py index 1ddf73b50868..b5a5c4429b5f 100644 --- a/source/isaaclab_ov/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_ov/test/assets/test_rigid_object_collection.py @@ -9,8 +9,12 @@ """Real OVPhysX rigid-object-collection coverage on one module-scoped scene per device. Each device builds one scene of locally authored ``N=2, B=3`` collections, resets it once, and keeps it -alive for every test that uses it. Each collection belongs to one test, so the tests do not depend on each -other's order. Initialization failures need their own scenes and therefore build them before any shared scene. +alive for every test that uses it. Each collection belongs to one test, so those tests do not depend on each +other's order. + +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 @@ -134,6 +138,10 @@ def test_rigid_object_collection_real_ovphysx_seams(scene: _CollectionScene) -> # 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) @@ -241,6 +249,7 @@ def test_wrench_reaches_only_the_selected_collection_body(scene: _CollectionScen @pytest.mark.parametrize("scene", test_devices(DeviceScope.CUDA), 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." 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 24b9c1b0d6dd..67c1a532b907 100644 --- a/source/isaaclab_ov/test/physics/test_ovphysx_manager_lifecycle.py +++ b/source/isaaclab_ov/test/physics/test_ovphysx_manager_lifecycle.py @@ -562,20 +562,26 @@ def test_scene_external_forces_every_iteration(monkeypatch, manager_module, devi @pytest.mark.parametrize(("device", "gpu_dynamics", "broadphase"), [("cpu", False, "MBP"), ("gpu", True, "GPU")]) -def test_scene_authors_device_dynamics_and_broadphase(monkeypatch, manager_module, device, gpu_dynamics, broadphase): - """Each scene selects its own dynamics device, since CPU and GPU scenes share one process.""" +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 Usd + from pxr import Sdf, Usd, UsdPhysics from isaaclab.physics import PhysicsManager monkeypatch.setattr(PhysicsManager, "_sim", None) stage = Usd.Stage.CreateInMemory() - prim = stage.DefinePrim("/World/PhysicsScene", "PhysicsScene") - manager_module.OvPhysxManager._configure_physx_scene_prim(prim, OvPhysxCfg(), device) - assert prim.GetAttribute("physxScene:enableGPUDynamics").Get() is gpu_dynamics - assert prim.GetAttribute("physxScene:broadphaseType").Get() == broadphase + 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): diff --git a/source/isaaclab_ov/test/sensors/test_contact_sensor.py b/source/isaaclab_ov/test/sensors/test_contact_sensor.py index 7d61cc909aa1..28a6f93aa981 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 4f520bd382b0..afa92a20db4c 100644 --- a/source/isaaclab_ov/test/sensors/test_frame_transformer.py +++ b/source/isaaclab_ov/test/sensors/test_frame_transformer.py @@ -53,39 +53,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 32061ee27a19..a6289b9bc297 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 @@ -11,8 +11,7 @@ friction/restitution through the asset's ``OvPhysxView``. The ``cfg`` / ``env`` / ``asset_cfg`` inputs are stubbed: the term only reads ``cfg.params``, ``env.scene[...]`` 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 @@ -36,28 +35,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.rst b/source/isaaclab_physx/changelog.d/antoiner-fix-physx-collection-com-inertia.rst index c985ac498742..531bfb955601 100644 --- a/source/isaaclab_physx/changelog.d/antoiner-fix-physx-collection-com-inertia.rst +++ b/source/isaaclab_physx/changelog.d/antoiner-fix-physx-collection-com-inertia.rst @@ -1,10 +1,13 @@ Fixed ^^^^^ -* Fixed :meth:`~isaaclab_physx.assets.RigidObjectCollection.set_coms_index` raising a ``KeyError`` because the +* 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` leaving the unselected entries of - :attr:`~isaaclab_physx.assets.RigidObjectCollectionData.body_com_pose_b` zeroed when it was called for a subset of - bodies or environments before the property was first read. -* Fixed :meth:`~isaaclab_physx.assets.RigidObjectCollection.set_inertias_index` failing in the PhysX backend because - the inertias were passed to the PhysX tensor API with the wrong layout. +* 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/test/assets/test_articulation.py b/source/isaaclab_physx/test/assets/test_articulation.py index e6d0e134aa78..a586d2e2d919 100644 --- a/source/isaaclab_physx/test/assets/test_articulation.py +++ b/source/isaaclab_physx/test/assets/test_articulation.py @@ -84,6 +84,16 @@ _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 _branching_cfg(prim_path: str = "/World/Robot", **kwargs) -> ArticulationCfg: @@ -245,7 +255,7 @@ def test_setting_invalid_articulation_root_prim_path(sim, device) -> None: @pytest.mark.parametrize("num_articulations", [2]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +@pytest.mark.parametrize("device", test_devices()) def test_fixed_tendon_position_target_writes_offset(sim, num_articulations, device) -> None: """A tendon length target lands in the simulation as ``rest_length - target`` on the selected cells only. @@ -518,6 +528,10 @@ class _ArticulationScene: """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]: @@ -564,22 +578,10 @@ def articulation_scene(request) -> Iterator[_ArticulationScene]: ), {"fixed_base": True}, ), - "floating": ( - _branching_cfg( - 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, - ) - }, - ), - {"root_on_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"), ), @@ -591,11 +593,31 @@ def articulation_scene(request) -> Iterator[_ArticulationScene]: 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, ) @@ -759,7 +781,9 @@ def test_articulation_joint_and_body_properties_round_trip(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 (identity order), and after the next update once the buffer is primed (non-identity order). + # 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 write_backend_mass_inertia( @@ -773,16 +797,16 @@ def write_backend_mass_inertia( 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) - backend_masses, backend_inertias = write_backend_mass_inertia(scene.tendon, 0.293, 0.023) - torch.testing.assert_close(scene.tendon.data.body_mass.torch, backend_masses) - torch.testing.assert_close(scene.tendon.data.body_inertia.torch, backend_inertias) - 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) - articulation.data.body_mass.torch # prime the lazy buffers - articulation.data.body_inertia.torch + # 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]) @@ -792,12 +816,22 @@ def write_backend_mass_inertia( 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), } - expected_frictions = {name: getattr(articulation.data, name).torch.clone() for name in frictions} + 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"], @@ -808,7 +842,7 @@ def write_backend_mass_inertia( 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(getattr(articulation.data, name).torch, expected_frictions[name]) + 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) @@ -842,15 +876,18 @@ def write_backend_mass_inertia( inertias = articulation.data.body_inertia.torch[env_ids][:, body_ids].clone() inertias[0, 0, 0] *= 1.2 inertias[0, 1, 4] *= 1.3 - for setter, kwarg, value, name in ( - (articulation.set_masses_index, "masses", masses, "body_mass"), - (articulation.set_coms_index, "coms", coms, "body_com_pose_b"), - (articulation.set_inertias_index, "inertias", inertias, "body_inertia"), - ): - expected = getattr(articulation.data, name).torch.clone() - expected[env_ids[:, None], body_ids] = value - setter(**{kwarg: value, "env_ids": env_ids, "body_ids": body_ids}) - torch.testing.assert_close(getattr(articulation.data, name).torch, expected) + 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()), @@ -858,6 +895,12 @@ def write_backend_mass_inertia( ): 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. @@ -975,6 +1018,7 @@ def _assert_jacobian_contract(articulation: Articulation, generalized_velocity: torch.testing.assert_close(generalized_energy, body_energy, atol=1e-5, rtol=1e-4) +@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 @@ -1061,7 +1105,10 @@ def test_articulation_drive_and_dynamics(articulation_scene: _ArticulationScene) 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. + # 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) @@ -1187,11 +1234,11 @@ def test_floating_articulation_root_writes_and_wrenches(articulation_scene: _Art # 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]]]}, "root_com_lin_vel_w"), + ({"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]]]}, - "root_ang_vel_b", + lambda: floating.data.root_ang_vel_b, ), ): _place_at_rest(floating, rest_pose) @@ -1205,7 +1252,7 @@ def test_floating_articulation_root_writes_and_wrenches(articulation_scene: _Art body_ids=base, env_ids=[0], is_global=True, **global_wrench ) scene.step() - response_value = getattr(floating.data, response).torch + 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. @@ -1219,20 +1266,28 @@ def test_floating_articulation_root_writes_and_wrenches(articulation_scene: _Art 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, also when the public body - # order differs from the PhysX link order. - _place_at_rest(reordered, reordered_rest_pose) - tips = reordered.find_bodies(["left_tip", "right_tip"], preserve_order=True)[0] - reordered.permanent_wrench_composer.set_forces_and_torques_index( - forces=torch.tensor([[[0.0, 4.0, 0.0], [0.0, -4.0, 0.0]]], device=device), body_ids=tips, env_ids=[1] - ) + # 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], torch.tensor([[0.0, 1.0, 0.0], [0.0, -1.0, 0.0]], device=device) - ) + 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 + ) # A partial reset clears the selected environment; a full reset resets every actuator environment and # clears all external forces and torques. @@ -1264,39 +1319,58 @@ def record_actuator_reset(env_ids=None) -> None: assert torch.count_nonzero(composer.out_torque_b.torch) == 0 +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) + + 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 assert articulation.num_spatial_tendons == 1 + data = articulation.data root_view = articulation.root_view - getters = { - "stiffness": root_view.get_spatial_tendon_stiffnesses, - "limit_stiffness": root_view.get_spatial_tendon_limit_stiffnesses, - "damping": root_view.get_spatial_tendon_dampings, - "offset": root_view.get_spatial_tendon_offsets, + # 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), } # 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 - for name, value in values.items(): - getattr(articulation, f"set_spatial_tendon_{name}_index")(**{name: value}) + _set_spatial_tendon_properties(articulation, values) articulation.write_spatial_tendon_properties_to_sim_index() - for name, getter in getters.items(): - torch.testing.assert_close(wp.to_torch(getter()).to(device), values[name]) - - # Partial writes change environment 1 in the data and in the solver and preserve environment 0. - env_ids = torch.tensor([1], dtype=torch.int32, device=device) - partial = {"stiffness": 12.0, "limit_stiffness": 3.0, "damping": 1.5, "offset": 0.1} - for name, value in partial.items(): - getattr(articulation, f"set_spatial_tendon_{name}_index")( - **{name: torch.tensor([[value]], device=device), "env_ids": env_ids} - ) - values[name][1] = value - torch.testing.assert_close(getattr(articulation.data, f"spatial_tendon_{name}").torch, values[name]) - articulation.write_spatial_tendon_properties_to_sim_index(env_ids=env_ids) - for name, getter in getters.items(): - torch.testing.assert_close(wp.to_torch(getter()).to(device), values[name]) + 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]) diff --git a/source/isaaclab_physx/test/assets/test_deformable_object.py b/source/isaaclab_physx/test/assets/test_deformable_object.py index edda3ea457f8..cfd389fa99ef 100644 --- a/source/isaaclab_physx/test/assets/test_deformable_object.py +++ b/source/isaaclab_physx/test/assets/test_deformable_object.py @@ -139,7 +139,7 @@ def test_initialization_on_device_cpu() -> None: # Check that the framework doesn't hold excessive strong references. assert sys.getrefcount(deformable) < 10 - with pytest.raises(RuntimeError): + with pytest.raises(RuntimeError, match="Failed to create deformable body at"): sim.reset() @@ -176,6 +176,7 @@ def deformable_scene(request) -> Iterator[_DeformableScene]: yield _DeformableScene(sim=sim, device=device, volume=volume, surface=surface, refcounts=refcounts) +@pytest.mark.isaacsim_ci def test_deformable_initialization(deformable_scene: _DeformableScene) -> None: """Resolve the volume and surface views, their materials, and their nodal buffers.""" scene = deformable_scene @@ -203,6 +204,7 @@ def test_deformable_initialization(deformable_scene: _DeformableScene) -> None: scene.surface.write_nodal_kinematic_target_to_sim_index(dummy_targets) +@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 @@ -238,9 +240,11 @@ def test_deformable_nodal_state_writes(deformable_scene: _DeformableScene) -> No nodal_state[..., :3] = nodal_pos nodal_state[..., 3:] = 0.0 volume.write_nodal_state_to_sim_index(nodal_state) - torch.testing.assert_close(volume.data.root_pos_w.torch, nodal_pos.mean(dim=1)) + 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) +@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 diff --git a/source/isaaclab_physx/test/assets/test_rigid_object.py b/source/isaaclab_physx/test/assets/test_rigid_object.py index a3a81756c83c..af3057c2e342 100644 --- a/source/isaaclab_physx/test/assets/test_rigid_object.py +++ b/source/isaaclab_physx/test/assets/test_rigid_object.py @@ -14,7 +14,7 @@ # ignore private usage of variables warning # pyright: reportPrivateUsage=none -from isaaclab.test.utils import DeviceScope, launch_test_simulation, test_devices +from isaaclab.test.utils import launch_test_simulation, test_devices launch_test_simulation() @@ -29,7 +29,7 @@ import torch import warp as wp from isaaclab_physx.assets import RigidObject -from isaaclab_physx.sim.schemas import PhysxRigidBodyCfg +from isaaclab_physx.sim.schemas import PhysxArticulationCfg, PhysxRigidBodyCfg import omni.kit.app import omni.physx @@ -38,6 +38,7 @@ import isaaclab.sim as sim_utils from isaaclab.assets import RigidObjectCfg 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 @@ -85,12 +86,14 @@ def _yaw_quat(angle: float) -> tuple[float, float, float, float]: pytest.param( "articulation_root", marks=pytest.mark.xfail( - strict=True, reason="The PhysX rigid object does not reject an enabled articulation root." + 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.parametrize("device", test_devices()) @pytest.mark.isaacsim_ci def test_initialization_rejects_invalid_rigid_body(device, api: Literal["none", "articulation_root"]) -> None: """Initialization fails without a rigid body and when the rigid body is an articulation root.""" @@ -106,9 +109,17 @@ def test_initialization_rejects_invalid_rigid_body(device, api: Literal["none", cube_object = RigidObject(cfg=cfg) else: cube_object = RigidObject(cfg=_cube_cfg("/World/Env_[^/]*/Object")) - # The articulation root above the rigid body turns the body into an articulation link. + # 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): - UsdPhysics.ArticulationRootAPI.Apply(sim.stage.GetPrimAtPath(f"/World/Env_{index}")) + prim = sim.stage.GetPrimAtPath(f"/World/Env_{index}/Object") + assert prim.HasAPI(UsdPhysics.ArticulationRootAPI) + assert prim.GetAttribute("physxArticulation:articulationEnabled").Get() # Check that the framework doesn't hold excessive strong references. assert sys.getrefcount(cube_object) < 10 @@ -117,7 +128,7 @@ def test_initialization_rejects_invalid_rigid_body(device, api: Literal["none", sim.reset() -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +@pytest.mark.parametrize("device", test_devices()) @pytest.mark.parametrize("gravity_enabled", [True, False]) @pytest.mark.isaacsim_ci def test_gravity_vec_w(device, gravity_enabled) -> None: @@ -245,6 +256,7 @@ def rigid_object_scene(request) -> Iterator[_RigidObjectScene]: ) +@pytest.mark.isaacsim_ci def test_rigid_object_initialization(rigid_object_scene: _RigidObjectScene) -> None: """Initialize local cubes with the expected buffers; kinematic cubes hold their default pose under gravity.""" scene = rigid_object_scene @@ -268,6 +280,7 @@ def test_rigid_object_initialization(rigid_object_scene: _RigidObjectScene) -> N torch.testing.assert_close(kinematic.data.root_com_vel_w.torch, kinematic.data.default_root_vel.torch) +@pytest.mark.isaacsim_ci def test_rigid_object_inertial_properties(rigid_object_scene: _RigidObjectScene) -> None: """Mass, center-of-mass, and inertia writes reach only the selected PhysX entries and survive a step.""" scene = rigid_object_scene @@ -299,15 +312,16 @@ def test_rigid_object_inertial_properties(rigid_object_scene: _RigidObjectScene) 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, raw in ( - ("body_mass", cube_object.root_view.get_masses()), - ("body_com_pose_b", cube_object.root_view.get_coms().view(wp.float32)), - ("body_inertia", cube_object.root_view.get_inertias()), + 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(getattr(cube_object.data, name).torch, expected[name]) + torch.testing.assert_close(data.torch, expected[name]) torch.testing.assert_close(wp.to_torch(raw).to(device).reshape(expected[name].shape), expected[name]) +@pytest.mark.isaacsim_ci def test_rigid_object_root_state_writes(rigid_object_scene: _RigidObjectScene) -> None: """Root writes round-trip through the written frame and keep the other frame consistent with the offset.""" scene = rigid_object_scene @@ -322,6 +336,7 @@ def test_rigid_object_root_state_writes(rigid_object_scene: _RigidObjectScene) - # 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) @@ -334,6 +349,13 @@ def test_rigid_object_root_state_writes(rigid_object_scene: _RigidObjectScene) - 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) @@ -379,6 +401,8 @@ def test_rigid_object_root_state_writes(rigid_object_scene: _RigidObjectScene) - 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) @@ -413,6 +437,7 @@ def test_rigid_object_root_state_writes(rigid_object_scene: _RigidObjectScene) - torch.testing.assert_close(root_com_vel_w[:, 3:], root_link_vel_w[:, 3:]) +@pytest.mark.isaacsim_ci def test_rigid_object_wrench_delivery_and_reset(rigid_object_scene: _RigidObjectScene) -> None: """External wrenches act in the frame they are given in on the selected root; reset clears them.""" scene = rigid_object_scene @@ -422,14 +447,14 @@ def test_rigid_object_wrench_delivery_and_reset(rigid_object_scene: _RigidObject (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 environment 0 so that their responses are comparable. - for setter, kwarg, name in ( - (cube_object.set_masses_index, "masses", "body_mass"), - (cube_object.set_coms_index, "coms", "body_com_pose_b"), - (cube_object.set_inertias_index, "inertias", "body_inertia"), - ): - value = getattr(cube_object.data, name).torch[:1] - setter(**{kwarg: value.expand(_NUM_ENVS, *value.shape[1:]).contiguous()}) + # 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) @@ -443,11 +468,11 @@ def test_rigid_object_wrench_delivery_and_reset(rigid_object_scene: _RigidObject # 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]]]}, "root_com_lin_vel_w"), + ({"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]]]}, - "root_com_ang_vel_b", + lambda: cube_object.data.root_com_ang_vel_b, ), ): scene.place_cubes_at_rest(rest_pose) @@ -459,7 +484,7 @@ def test_rigid_object_wrench_delivery_and_reset(rigid_object_scene: _RigidObject 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 = getattr(cube_object.data, response).torch + 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) diff --git a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py index 794e2f2a98ab..97d10830f887 100644 --- a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py @@ -14,7 +14,7 @@ # ignore private usage of variables warning # pyright: reportPrivateUsage=none -from isaaclab.test.utils import DeviceScope, launch_test_simulation, test_devices +from isaaclab.test.utils import launch_test_simulation, test_devices launch_test_simulation() @@ -113,7 +113,7 @@ def sim(request) -> Iterator[SimulationContext]: @pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +@pytest.mark.parametrize("device", test_devices()) def test_initialization_with_no_rigid_body(sim, num_cubes, device) -> None: """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) @@ -189,7 +189,7 @@ def read_view(values: wp.array, data_dim: int) -> torch.Tensor: @pytest.mark.parametrize("num_envs", [3]) @pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +@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) -> None: """Test that gravity vector direction is set correctly for the rigid object.""" @@ -374,6 +374,14 @@ def test_collection_state_writes(collection_scene: _CollectionScene) -> None: 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) @@ -394,6 +402,8 @@ def test_collection_state_writes(collection_scene: _CollectionScene) -> None: 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_cubes_at_rest(_rest_poses(scene)) @@ -449,23 +459,29 @@ def test_collection_wrench_delivery_and_reset(collection_scene: _CollectionScene device = scene.device collection = scene.cubes rest_poses = _rest_poses(scene, yaw=0.5 * math.pi) - # Give every body the same inertial properties so that their responses are comparable. - for setter, kwarg, name in ( - (collection.set_coms_index, "coms", "body_com_pose_b"), - (collection.set_inertias_index, "inertias", "body_inertia"), - ): - value = getattr(collection.data, name).torch[:1, :1] - setter(**{kwarg: value.expand(_NUM_ENVS, _NUM_CUBES, *value.shape[2:]).contiguous()}) + # 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 in ( - ({"forces": [[6.0, 0.0, 0.0]]}, {"forces": [[0.0, 6.0, 0.0]]}, "body_com_lin_vel_w"), + 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]]}, - "body_com_ang_vel_b", + lambda: collection.data.body_com_ang_vel_b, + False, ), ): scene.place_cubes_at_rest(rest_poses) @@ -478,11 +494,15 @@ def test_collection_wrench_delivery_and_reset(collection_scene: _CollectionScene body_ids=object_ids[0::2], env_ids=[env_id], is_global=is_global, **wrench ) scene.step() - response_value = getattr(collection.data, response).torch + 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 response == "body_com_lin_vel_w": + 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( From 0a0ded20cc198169477eea845ae9a998b0c9a8a2 Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Fri, 2 Oct 2026 14:11:00 -0700 Subject: [PATCH 21/33] Simplify asset tests and remove custom test scaffolding --- .../changelog.d/asset-tests-readability.skip | 0 .../contract/_articulation_contract_cases.py | 1630 ----------- .../_articulation_ordering_contract_cases.py | 2405 ----------------- .../assets/contract/_manager_patch_scope.py | 44 - ..._rigid_object_collection_contract_cases.py | 913 ------- .../contract/_rigid_object_contract_cases.py | 841 ------ ...tract_utils.py => articulation_factory.py} | 253 +- .../test/assets/contract/capabilities.py | 7 +- .../isaaclab/test/assets/contract/conftest.py | 23 +- .../{_contract_boot.py => contract_boot.py} | 0 .../{_mock_backends.py => mock_backends.py} | 28 +- ....py => rigid_object_collection_factory.py} | 138 +- ...tract_utils.py => rigid_object_factory.py} | 135 +- .../contract/test_articulation_contract.py | 1613 +++++++++++ .../test_articulation_ordering_contract.py | 2367 ++++++++++++++++ .../contract/test_asset_contract_api.py | 36 +- .../contract/test_asset_contract_data.py | 30 - .../contract/test_asset_contract_writes.py | 39 - .../test_rigid_object_collection_contract.py | 867 ++++++ .../contract/test_rigid_object_contract.py | 790 ++++++ .../changelog.d/asset-tests-readability.skip | 0 .../test/articulation_test_utils.py | 389 +-- .../test/assets/data/fixed_cartpole.usda | 3 + .../test/assets/data/fixed_spatial_chain.usda | 3 + .../test/assets/data/fixed_tendon_hand.usda | 3 + .../test/assets/data/floating_two_leg.usda | 3 + .../test/assets/data/floating_two_link.usda | 3 + .../test/assets/data/revolute_pendulum.usda | 3 + .../rigid_body_with_articulation_root.usda | 3 + .../assets/test_newton_actuators_newton.py | 48 +- .../test_newton_task_space_controllers.py | 15 +- .../changelog.d/asset-tests-readability.skip | 0 .../test/assets/test_articulation.py | 39 +- .../test/assets/test_rigid_object.py | 12 +- .../assets/test_rigid_object_collection.py | 16 +- .../changelog.d/asset-tests-readability.skip | 0 .../test/assets/test_articulation.py | 18 +- .../assets/test_rigid_object_collection.py | 20 +- 38 files changed, 6035 insertions(+), 6702 deletions(-) create mode 100644 source/isaaclab/changelog.d/asset-tests-readability.skip delete mode 100644 source/isaaclab/test/assets/contract/_articulation_contract_cases.py delete mode 100644 source/isaaclab/test/assets/contract/_articulation_ordering_contract_cases.py delete mode 100644 source/isaaclab/test/assets/contract/_manager_patch_scope.py delete mode 100644 source/isaaclab/test/assets/contract/_rigid_object_collection_contract_cases.py delete mode 100644 source/isaaclab/test/assets/contract/_rigid_object_contract_cases.py rename source/isaaclab/test/assets/contract/{_articulation_contract_utils.py => articulation_factory.py} (60%) rename source/isaaclab/test/assets/contract/{_contract_boot.py => contract_boot.py} (100%) rename source/isaaclab/test/assets/contract/{_mock_backends.py => mock_backends.py} (69%) rename source/isaaclab/test/assets/contract/{_rigid_object_collection_contract_utils.py => rigid_object_collection_factory.py} (63%) rename source/isaaclab/test/assets/contract/{_rigid_object_contract_utils.py => rigid_object_factory.py} (58%) create mode 100644 source/isaaclab/test/assets/contract/test_articulation_contract.py create mode 100644 source/isaaclab/test/assets/contract/test_articulation_ordering_contract.py delete mode 100644 source/isaaclab/test/assets/contract/test_asset_contract_data.py delete mode 100644 source/isaaclab/test/assets/contract/test_asset_contract_writes.py create mode 100644 source/isaaclab/test/assets/contract/test_rigid_object_collection_contract.py create mode 100644 source/isaaclab/test/assets/contract/test_rigid_object_contract.py create mode 100644 source/isaaclab_newton/changelog.d/asset-tests-readability.skip create mode 100644 source/isaaclab_newton/test/assets/data/fixed_cartpole.usda create mode 100644 source/isaaclab_newton/test/assets/data/fixed_spatial_chain.usda create mode 100644 source/isaaclab_newton/test/assets/data/fixed_tendon_hand.usda create mode 100644 source/isaaclab_newton/test/assets/data/floating_two_leg.usda create mode 100644 source/isaaclab_newton/test/assets/data/floating_two_link.usda create mode 100644 source/isaaclab_newton/test/assets/data/revolute_pendulum.usda create mode 100644 source/isaaclab_newton/test/assets/data/rigid_body_with_articulation_root.usda create mode 100644 source/isaaclab_ov/changelog.d/asset-tests-readability.skip create mode 100644 source/isaaclab_physx/changelog.d/asset-tests-readability.skip diff --git a/source/isaaclab/changelog.d/asset-tests-readability.skip b/source/isaaclab/changelog.d/asset-tests-readability.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab/test/assets/contract/_articulation_contract_cases.py b/source/isaaclab/test/assets/contract/_articulation_contract_cases.py deleted file mode 100644 index fda0c57eac11..000000000000 --- a/source/isaaclab/test/assets/contract/_articulation_contract_cases.py +++ /dev/null @@ -1,1630 +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 articulation contract cases, collected by the ``test_asset_contract_*`` entry modules. - -Checks that every articulation backend provides the data and writer behavior the base articulation class advertises. -The backends run on mocked views, so these cases need neither Isaac Sim nor a GPU simulation. -""" - -import math -import warnings -from contextlib import AbstractContextManager, nullcontext -from unittest.mock import MagicMock, patch - -import numpy as np -import pytest -import torch -import warp as wp - -from isaaclab.test.utils import DeviceScope, test_devices - -from ._articulation_contract_utils import get_articulation -from ._articulation_ordering_contract_cases import _read_backend_joint_state -from .capabilities import contract_backend, require_backend_capability, requires_backend - -# --------------------------------------------------------------------------- -# 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. -_api_backends = contract_backend("api") -_data_backends = contract_backend("data") -_writes_backends = contract_backend("writes") -_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 - - -# --------------------------------------------------------------------------- -# Tests: Index resolution helpers -# --------------------------------------------------------------------------- - - -class TestArticulationIndexResolution: - """Test backend-specific index resolution helpers.""" - - @_api_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() - - @_api_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 - - @_api_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.""" - - @_api_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.""" - - @_api_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 finder_name == "find_spatial_tendons": - require_backend_capability(backend, "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]) - - -def _fixed_tendon_articulation(backend: str): - """Create a two-environment articulation with two fixed tendons and no spatial tendons.""" - 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 - - -class TestFixedTendonTargetScheduling: - """Commanding a tendon target is what schedules the write, not merely having tendons.""" - - @contract_backend("fixed_tendon_target_scheduling") - 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.""" - art = _fixed_tendon_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 - - @requires_backend("newton") - def test_newton_reports_a_missing_tendon_actuator(self): - """Without a MuJoCo tendon actuator the solver has no adapter, so commanding must say so.""" - art = _fixed_tendon_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)) - - @requires_backend("newton") - 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.""" - from isaaclab_newton.physics import NewtonManager - - art = _fixed_tendon_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.""" - - @_api_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") - - @_api_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 - - @_api_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.""" - - @_data_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() - - @requires_backend("physx") - 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 - - @requires_backend("physx") - 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 - - @contract_backend("mocked_dynamics") - @_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 - - @_data_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] - - @_data_backends - def test_deprecated_command_aliases_return_actuator_collection_buffers(self, backend): - """The deprecated data command and torque properties warn and alias the actuator collection's buffers.""" - art, _ = get_articulation(backend, 2, 4, 5, device="cpu") - actuators = art.actuators - 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 - - -# --------------------------------------------------------------------------- -# 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.""" - - @_data_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.""" - - @_writes_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) - - @_writes_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) - - @_writes_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 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 -- - - @_writes_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)) - - @_writes_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 -- - - @_writes_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)) - - @_writes_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)) - - -@_writes_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.""" - - @_writes_backends - def test_partial_joint_state_write_follows_int64_selectors_and_restarts_acceleration(self, backend): - """A fused joint-state write honors unsorted int64 selectors and restarts the finite-difference baseline.""" - art, raw_backend = get_articulation(backend, 2, 3, 2, device="cpu") - 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, [(name, 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) - - @_writes_backends - def test_deprecated_joint_state_writer_matches_public_writers(self, backend): - """The deprecated combined joint-state writer warns and writes what the public writers would.""" - art, raw_backend = get_articulation(backend, 2, 3, 2, device="cpu") - 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) - - @_writes_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)}) - - @_writes_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.""" - - @_writes_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)}) - - @_writes_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)}) - - -# --------------------------------------------------------------------------- -# Tests: partial root and body writes reach only the selected backend rows -# --------------------------------------------------------------------------- - -# quantity -> (writer base name, writer keyword, trailing size). Joint partial writes are covered, with and without -# a joint ordering, by ``TestArticulationOrderingWriteParity``. -_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 == "physx": - getter = { - "root_pose": raw_backend.get_root_transforms, - "root_velocity": raw_backend.get_root_velocities, - "mass": raw_backend.get_masses, - } - values = wp.to_torch(getter[quantity]()) - elif backend == "newton": - if quantity == "mass": - values = wp.to_torch(art.data._sim_bind_body_mass) - else: - getter = {"root_pose": raw_backend.get_root_transforms, "root_velocity": raw_backend.get_root_velocities} - values = wp.to_torch(getter[quantity](None)) - else: - 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) - return values.reshape(art.num_instances, -1).cpu().clone() - - -class TestArticulationPartialWriteRows: - """Test that partial root and body writers reach only the selected rows of backend storage.""" - - @_writes_backends - @pytest.mark.parametrize("selection", ["index", "mask"]) - @pytest.mark.parametrize("quantity", _PARTIAL_WRITES) - def test_partial_write_preserves_unselected_backend_rows(self, backend, selection, quantity): - num_instances, num_bodies = 2, 4 - art, raw_backend = get_articulation(backend, num_instances, 3, num_bodies, device="cpu") - writer, kwarg, trailing = _PARTIAL_WRITES[quantity] - width = num_bodies if quantity == "mass" else trailing - # Literal, per-element distinct payload; poses keep an identity rotation. - values = 100.0 * torch.arange(1, num_instances + 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 - kwargs = {} - if selection == "index": - kwargs[kwarg] = values[env : env + 1, body : body + 1] if quantity == "mass" else values[env : env + 1] - kwargs["env_ids"] = torch.tensor([env], dtype=torch.int32) - if quantity == "mass": - kwargs["body_ids"] = [body] - else: - kwargs[kwarg] = values - kwargs["env_mask"] = _make_item_mask(num_instances, [env], "cpu") - if quantity == "mass": - kwargs["body_mask"] = _make_item_mask(num_bodies, [body], "cpu") - - getattr(art, f"{writer}_{selection}")(**kwargs) - - after = _read_backend_rows(backend, art, raw_backend, quantity) - expected = before.clone() - if quantity == "mass": - expected[env, body] = values[env, body] - else: - expected[env] = values[env] - torch.testing.assert_close(after, expected, rtol=0.0, atol=0.0) - - -# --------------------------------------------------------------------------- -# Tendon tests — counts, names, data, writers -# --------------------------------------------------------------------------- - -# Every backend supports the core fixed-tendon properties. Newton declares the extended fixed-tendon properties and -# spatial tendons unsupported, so those rows stay visible as reasoned capability skips. -_tendon_backends = contract_backend("fixed_tendons") - -# The first row covers the zero-spatial-tendon edge; the second has distinct counts on every axis. -_TENDON_DIMS = [ - (1, 2, 2, 1, 0), # fixed only - (2, 6, 7, 3, 2), # both types -] -_TENDON_DIM_NAMES = ("num_instances", "num_joints", "num_bodies", "num_fixed_tendons", "num_spatial_tendons") -_tendon_dims = pytest.mark.parametrize(_TENDON_DIM_NAMES, _TENDON_DIMS) - -# tendon property family -> (tendon kind, backend capability) -_TENDON_FAMILIES = { - "fixed": ("fixed", "fixed_tendons"), - "fixed_extended": ("fixed", "fixed_tendon_extended"), - "spatial": ("spatial", "spatial_tendons"), -} - -# (property, family, dtype) -_TENDON_DATA_PROPERTIES = [ - ("fixed_tendon_stiffness", "fixed", wp.float32), - ("fixed_tendon_damping", "fixed", wp.float32), - ("fixed_tendon_pos_limits", "fixed", wp.vec2f), - ("fixed_tendon_limit_stiffness", "fixed_extended", wp.float32), - ("fixed_tendon_rest_length", "fixed_extended", wp.float32), - ("fixed_tendon_offset", "fixed_extended", wp.float32), - ("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), -] -_TENDON_PROPERTY_FAMILY = {name: family for name, family, _ in _TENDON_DATA_PROPERTIES} - - -class TestArticulationTendons: - """Test tendon counts, names, and data properties.""" - - @_tendon_backends - @_tendon_dims - @_devices - @pytest.mark.parametrize("family", _TENDON_FAMILIES) - def test_articulation_tendon_contract( - self, backend, num_instances, num_joints, num_bodies, num_fixed_tendons, num_spatial_tendons, device, family - ): - kind, capability = _TENDON_FAMILIES[family] - require_backend_capability(backend, capability) - art, _ = get_articulation( - backend, num_instances, num_joints, num_bodies, num_fixed_tendons, num_spatial_tendons, device - ) - num_tendons = num_fixed_tendons if kind == "fixed" else num_spatial_tendons - assert getattr(art, f"num_{kind}_tendons") == num_tendons - names = getattr(art, f"{kind}_tendon_names") - assert isinstance(names, list) - assert len(names) == num_tendons - assert all(isinstance(n, str) for n in names) - - art.data.update(dt=0.01) - for name, property_family, dtype in _TENDON_DATA_PROPERTIES: - if property_family != family or num_tendons == 0: - continue - expected_shape = (num_instances, num_tendons) - 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] -# Writer cases pair each method with the dimension rows that configure at least one tendon of its kind. -_TENDON_WRITER_CASES = [ - pytest.param(*method, *dims, id=f"{method[1]}-{'-'.join(map(str, dims))}") - for method in _TENDON_METHODS - for dims in _TENDON_DIMS - if dims[3 if method[0] == "fixed" else 4] > 0 -] - - -def _suppress_newton_model_changes(backend: str) -> AbstractContextManager: - """Return a context that drops Newton model-change notifications, which need a solver the mock lacks.""" - if backend != "newton": - return nullcontext() - from isaaclab_newton.physics import NewtonManager - - return patch.object(NewtonManager, "add_model_change") - - -class TestArticulationWritersTendon: - """Test tendon writers/setters with all input combinations.""" - - @_tendon_backends - @_devices - @pytest.mark.parametrize("selection", ["index", "mask"]) - @pytest.mark.parametrize( - ("kind", "method_base", "kwarg", "getter", *_TENDON_DIM_NAMES), - _TENDON_WRITER_CASES, - ) - def test_tendon_writer( - self, - backend, - device, - selection, - kind, - method_base, - kwarg, - getter, - num_instances, - num_joints, - num_bodies, - num_fixed_tendons, - num_spatial_tendons, - ): - require_backend_capability(backend, _TENDON_FAMILIES[_TENDON_PROPERTY_FAMILY[getter]][1]) - num_tendons = num_fixed_tendons if kind == "fixed" else num_spatial_tendons - 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: after the push, the matching data property reads the written values back - method(**{kwarg: _make_payload_warp(full_shape, device, wp_dtype)}) - with _suppress_newton_model_changes(backend): - getattr(art, f"write_{kind}_tendon_properties_to_sim_index")() - _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}) - - # Shape validation runs before any kernel launch, so one device covers it. - @_tendon_backends - @pytest.mark.parametrize("array_kind", ["torch", "warp"]) - @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_rejects_extra_environments( - self, request, backend, selection, array_kind, kind, method_base, kwarg, getter - ): - require_backend_capability(backend, _TENDON_FAMILIES[_TENDON_PROPERTY_FAMILY[getter]][1]) - if backend == "newton" and selection == "mask": - # Product bug: Newton's fixed-tendon mask setters slice full data without validating its shape, so an - # oversized input is silently truncated. - request.applymarker( - pytest.mark.xfail( - raises=pytest.fail.Exception, - strict=True, - reason="Newton fixed-tendon mask setters accept data with extra environments", - ) - ) - num_instances, num_joints, num_bodies, num_fixed_tendons, num_spatial_tendons = _TENDON_DIMS[1] - art, _ = get_articulation( - backend, num_instances, num_joints, num_bodies, num_fixed_tendons, num_spatial_tendons, "cpu" - ) - full_shape = (num_instances, num_fixed_tendons if kind == "fixed" else num_spatial_tendons) - wp_dtype = wp.vec2f if getter == "fixed_tendon_pos_limits" else wp.float32 - make_bad_data = _make_bad_data_torch if array_kind == "torch" else _make_bad_data_warp - - with pytest.raises((AssertionError, RuntimeError)): - getattr(art, f"{method_base}_{selection}")(**{kwarg: make_bad_data(full_shape, "cpu", 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.""" - require_backend_capability(backend, _TENDON_FAMILIES[kind][1]) - 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)) - 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() - if selection == "index": - env_ids = None if len(selected_envs) == num_instances else torch.tensor(selected_envs, device=device) - write_kwargs = {"env_ids": env_ids} - 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_kwargs = {"env_mask": mask} - with _suppress_newton_model_changes(backend): - write(**write_kwargs) - 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 diff --git a/source/isaaclab/test/assets/contract/_articulation_ordering_contract_cases.py b/source/isaaclab/test/assets/contract/_articulation_ordering_contract_cases.py deleted file mode 100644 index ca836b6ad9d3..000000000000 --- a/source/isaaclab/test/assets/contract/_articulation_ordering_contract_cases.py +++ /dev/null @@ -1,2405 +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 joint- and body-ordering contract cases.""" - -from unittest.mock import MagicMock, patch - -import numpy as np -import pytest -import torch -import warp as wp -from _pytest.mark.structures import ParameterSet - -from isaaclab.utils import replace -from isaaclab.utils.buffers import TimestampedBuffer - -from ._articulation_contract_utils import get_articulation -from .capabilities import backend_parameters, contract_backend, requires_backend - - -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 an ordering parameter row that carries the backend's declared skip reason.""" - marks = [*kwargs.pop("marks", ()), *backend_parameters("ordering", names=(backend,))[0].marks] - return pytest.param(backend, *values, marks=marks, **kwargs) - - -_requires_physx = requires_backend("physx") -_requires_ovphysx = requires_backend("ovphysx") -_requires_newton = requires_backend("newton") -_all_backends = contract_backend("ordering") -_physx_ovphysx_backends = contract_backend("ordering", names=("physx", "ovphysx")) -_dynamics_ordering_backends = contract_backend("ordering", names=("physx", "newton")) -_non_mock_backends = _all_backends - - -class TestArticulationOrderingBodyState: - """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 TestArticulationOrderingJointState: - """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 TestArticulationOrderingOperations: - """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 TestArticulationOrderingJointWriters: - """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 TestArticulationOrderingBodyWriters: - """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/contract/_manager_patch_scope.py b/source/isaaclab/test/assets/contract/_manager_patch_scope.py deleted file mode 100644 index 243991ce51ca..000000000000 --- a/source/isaaclab/test/assets/contract/_manager_patch_scope.py +++ /dev/null @@ -1,44 +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 - -"""Test-local backend-manager patching for mocked asset factories.""" - -import inspect -from collections.abc import Iterator -from contextlib import contextmanager - -_MISSING = object() -_patch_scopes: list[dict[tuple[object, str], object]] = [] - - -@contextmanager -def contract_manager_patch_scope() -> Iterator[None]: - """Restore backend manager attributes changed by factories in this scope.""" - patches: dict[tuple[object, str], object] = {} - _patch_scopes.append(patches) - try: - yield - finally: - for (owner, name), original in reversed(patches.items()): - if original is _MISSING: - delattr(owner, name) - else: - setattr(owner, name, original) - popped = _patch_scopes.pop() - assert popped is patches - - -def patch_contract_manager(owner: object, name: str, replacement: object) -> None: - """Patch a manager binding until the current contract test finishes.""" - if not _patch_scopes: - raise RuntimeError("contract factories must run inside contract_manager_patch_scope()") - patches = _patch_scopes[-1] - key = (owner, name) - if key not in patches: - try: - patches[key] = inspect.getattr_static(owner, name) - except AttributeError: - patches[key] = _MISSING - setattr(owner, name, replacement) diff --git a/source/isaaclab/test/assets/contract/_rigid_object_collection_contract_cases.py b/source/isaaclab/test/assets/contract/_rigid_object_collection_contract_cases.py deleted file mode 100644 index fd4ca35d083b..000000000000 --- a/source/isaaclab/test/assets/contract/_rigid_object_collection_contract_cases.py +++ /dev/null @@ -1,913 +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 rigid-object-collection contract cases, collected by the ``test_asset_contract_*`` entry modules. - -Checks that every collection backend provides the data and writer behavior the base collection class advertises. -The backends run on mocked views, so these cases need neither Isaac Sim nor a GPU simulation. -""" - -import math -from unittest.mock import patch - -import numpy as np -import pytest -import torch -import warp as wp - -from isaaclab.test.utils import DeviceScope, test_devices - -from ._rigid_object_collection_contract_utils import get_rigid_object_collection -from .capabilities import contract_backend - -# 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. -_api_backends = contract_backend("api") -_data_backends = contract_backend("data") -_writes_backends = contract_backend("writes") -_devices = pytest.mark.parametrize("device", test_devices(DeviceScope.CPU_AND_DEFAULT_CUDA)) - - -# --------------------------------------------------------------------------- -# 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.""" - - @_api_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() - - @_api_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.""" - - @_data_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) - - @_data_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.""" - - @_api_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.""" - - @_api_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] - - @_api_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.""" - - @_data_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: - @_writes_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) - - @_writes_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) - - @_writes_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 -- - - @_writes_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 -- - - @_writes_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 -- - - @_writes_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 -- - - @_writes_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.""" - - @_writes_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)}) - - @_writes_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.""" - - @_data_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 - - -# --------------------------------------------------------------------------- -# Tests: partial writes reach only the selected backend cells -# --------------------------------------------------------------------------- - -# quantity -> (writer base name, writer keyword, trailing size) -_PARTIAL_WRITES = { - "body_pose": ("write_body_link_pose_to_sim", "body_poses", 7), - "body_velocity": ("write_body_com_velocity_to_sim", "body_velocities", 6), - "mass": ("set_masses", "masses", 1), -} - - -def _read_backend_cells(backend: str, raw_backend, quantity: str, num_instances: int, num_bodies: int) -> torch.Tensor: - """Read one collection quantity from backend storage in public ``(env, body, ...)`` order.""" - if backend == "physx": - getter = { - "body_pose": raw_backend.get_transforms, - "body_velocity": raw_backend.get_velocities, - "mass": raw_backend.get_masses, - } - values = wp.to_torch(getter[quantity]()) - # 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_backend.get_attribute("body_mass", None)) - else: - getter = {"body_pose": raw_backend.get_root_transforms, "body_velocity": raw_backend.get_root_velocities} - values = wp.to_torch(getter[quantity](None)) - else: - from isaaclab_ov import tensor_types as TT - - binding = {"body_pose": TT.LINK_POSE, "body_velocity": TT.LINK_VELOCITY, "mass": TT.BODY_MASS}[quantity] - values = torch.as_tensor(raw_backend.bindings[binding]._data) - return values.reshape(num_instances, num_bodies, -1).cpu().clone() - - -class TestCollectionPartialWriteCells: - """Test that partial writers reach only the selected cells of backend storage.""" - - @_writes_backends - @pytest.mark.parametrize("selection", ["index", "mask"]) - @pytest.mark.parametrize("quantity", _PARTIAL_WRITES) - def test_partial_write_preserves_unselected_backend_cells(self, request, backend, selection, quantity): - if backend == "ovphysx" 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", - ) - ) - num_instances, num_bodies = _NUM_INSTANCES, _NUM_BODIES - obj, raw_backend = get_rigid_object_collection(backend, num_instances, num_bodies, "cpu") - writer, kwarg, trailing = _PARTIAL_WRITES[quantity] - # Literal, per-cell distinct payload; poses keep an identity rotation. - values = 100.0 * torch.arange(1, num_instances * num_bodies + 1, dtype=torch.float32).reshape( - num_instances, num_bodies, 1 - ) + torch.arange(trailing, dtype=torch.float32) - if quantity == "body_pose": - values[..., 3:6] = 0.0 - values[..., 6] = 1.0 - if quantity == "mass": - values = values[..., 0] - before = _read_backend_cells(backend, raw_backend, quantity, num_instances, num_bodies) - # Select the second environment and the last body so identity or transposed routing cannot pass. - env, body = 1, 2 - if selection == "index": - getattr(obj, f"{writer}_index")( - **{kwarg: values[env : env + 1, body : body + 1]}, - env_ids=torch.tensor([env], dtype=torch.int32), - body_ids=_make_body_ids("cpu", [body]), - ) - else: - getattr(obj, f"{writer}_mask")( - **{kwarg: values}, - env_mask=_make_item_mask(num_instances, [env], "cpu"), - body_mask=_make_item_mask(num_bodies, [body], "cpu"), - ) - - after = _read_backend_cells(backend, raw_backend, quantity, num_instances, num_bodies) - expected = before.clone() - expected[env, body] = values[env, body].reshape(-1) - torch.testing.assert_close(after, expected, rtol=0.0, atol=0.0) diff --git a/source/isaaclab/test/assets/contract/_rigid_object_contract_cases.py b/source/isaaclab/test/assets/contract/_rigid_object_contract_cases.py deleted file mode 100644 index 92e877b1bbd6..000000000000 --- a/source/isaaclab/test/assets/contract/_rigid_object_contract_cases.py +++ /dev/null @@ -1,841 +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 rigid-object contract cases, collected by the ``test_asset_contract_*`` entry modules. - -Checks that every rigid-object backend provides the data and writer behavior the base rigid-object class advertises. -The backends run on mocked views, so these cases need neither Isaac Sim nor a GPU simulation. -""" - -import math -from unittest.mock import MagicMock, patch - -import numpy as np -import pytest -import torch -import warp as wp - -from isaaclab.test.utils import DeviceScope, test_devices - -from ._rigid_object_contract_utils import get_rigid_object -from .capabilities import contract_backend - -# --------------------------------------------------------------------------- -# 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. -_api_backends = contract_backend("api") -_data_backends = contract_backend("data") -_writes_backends = contract_backend("writes") -_devices = pytest.mark.parametrize("device", test_devices(DeviceScope.CPU_AND_DEFAULT_CUDA)) -_NUM_INSTANCES = 2 - - -# --------------------------------------------------------------------------- -# Tests: Index resolution helpers -# --------------------------------------------------------------------------- - - -class TestRigidObjectIndexResolution: - """Test backend-specific index resolution helpers.""" - - @_api_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.""" - - @_api_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.""" - - @_api_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.""" - - @_data_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.""" - - @_data_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: - @_writes_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) - - @_writes_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 - - @_writes_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 -- - - @_writes_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)) - - @_writes_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 -- - - @_writes_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)) - - @_writes_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]: - """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 TestRigidObjectWritersBody: - """Test body property writers/setters with all input combinations.""" - - @_writes_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) - - @_writes_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)}) - 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)}) - - @_writes_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)}) - 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: partial writes reach only the selected backend rows -# --------------------------------------------------------------------------- - -# quantity -> (writer base name, writer keyword, literal per-environment payload) -_PARTIAL_WRITES = { - "root_pose": ( - "write_root_link_pose_to_sim", - "root_pose", - [[31.0, 32.0, 33.0, 0.0, 0.0, 0.0, 1.0], [41.0, 42.0, 43.0, 0.0, 0.0, 0.0, 1.0]], - ), - "root_velocity": ( - "write_root_com_velocity_to_sim", - "root_velocity", - [[21.0, 22.0, 23.0, 24.0, 25.0, 26.0], [71.0, 72.0, 73.0, 74.0, 75.0, 76.0]], - ), - "mass": ("set_masses", "masses", [[51.0], [61.0]]), -} - - -def _read_backend_rows(backend: str, raw_backend, quantity: str) -> torch.Tensor: - """Read one rigid-object quantity from backend storage as one row per environment.""" - if backend == "physx": - getter = { - "root_pose": raw_backend.get_transforms, - "root_velocity": raw_backend.get_velocities, - "mass": raw_backend.get_masses, - } - values = getter[quantity]() - elif backend == "newton": - if quantity == "mass": - values = raw_backend.get_attribute("body_mass", None) - else: - getter = {"root_pose": raw_backend.get_root_transforms, "root_velocity": raw_backend.get_root_velocities} - values = getter[quantity](None) - else: - from isaaclab_ov import tensor_types as TT - - binding = {"root_pose": TT.RIGID_BODY_POSE, "root_velocity": TT.RIGID_BODY_VELOCITY, "mass": TT.RIGID_BODY_MASS} - values = raw_backend.bindings[binding[quantity]]._data - values = wp.to_torch(values) if isinstance(values, wp.array) else torch.as_tensor(values) - return values.reshape(2, -1).cpu().clone() - - -class TestRigidObjectPartialWriteRows: - """Test that partial writers reach only the selected rows of backend storage.""" - - @_writes_backends - @pytest.mark.parametrize("selection", ["index", "mask"]) - @pytest.mark.parametrize("quantity", _PARTIAL_WRITES) - def test_partial_write_preserves_unselected_backend_rows(self, backend, selection, quantity): - obj, raw_backend = get_rigid_object(backend, num_instances=2, device="cpu") - writer, kwarg, payload = _PARTIAL_WRITES[quantity] - values = torch.tensor(payload, dtype=torch.float32) - before = _read_backend_rows(backend, raw_backend, quantity) - # Select the second environment so a writer that ignores the selection cannot pass by writing row 0. - selected, unselected = 1, 0 - if selection == "index": - kwargs = {kwarg: values[selected : selected + 1], "env_ids": torch.tensor([selected], dtype=torch.int32)} - if quantity == "mass": - kwargs["body_ids"] = [0] - else: - kwargs = {kwarg: values, "env_mask": _make_item_mask(2, [selected], "cpu")} - if quantity == "mass": - kwargs["body_mask"] = _make_item_mask(1, [0], "cpu") - - getattr(obj, f"{writer}_{selection}")(**kwargs) - - after = _read_backend_rows(backend, raw_backend, quantity) - torch.testing.assert_close(after[selected], values[selected], rtol=0.0, atol=0.0) - torch.testing.assert_close(after[unselected], before[unselected], rtol=0.0, atol=0.0) diff --git a/source/isaaclab/test/assets/contract/_articulation_contract_utils.py b/source/isaaclab/test/assets/contract/articulation_factory.py similarity index 60% rename from source/isaaclab/test/assets/contract/_articulation_contract_utils.py rename to source/isaaclab/test/assets/contract/articulation_factory.py index c940a1e18b1b..fa4e2a7d6eb1 100644 --- a/source/isaaclab/test/assets/contract/_articulation_contract_utils.py +++ b/source/isaaclab/test/assets/contract/articulation_factory.py @@ -11,14 +11,14 @@ from unittest.mock import MagicMock import numpy as np +import pytest import warp as wp from isaaclab.assets.articulation.articulation_cfg import ArticulationCfg from isaaclab.utils.wrench_composer import WrenchComposer -from ._manager_patch_scope import patch_contract_manager -from ._mock_backends import install_physx_recording_setters, patch_ovphysx_manager, patch_physx_manager from .capabilities import available_backends +from .mock_backends import install_physx_recording_setters, patch_ovphysx_manager, patch_physx_manager BACKENDS = available_backends() @@ -69,9 +69,11 @@ def create_physx_articulation( is_fixed_base: bool = False, joint_ordering: tuple[str, ...] | None = None, body_ordering: tuple[str, ...] | None = None, + *, + monkeypatch: pytest.MonkeyPatch, ): """Create a test Articulation instance with mocked dependencies.""" - patch_physx_manager() + patch_physx_manager(monkeypatch=monkeypatch) 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)] @@ -86,8 +88,8 @@ def create_physx_articulation( joint_ordering=joint_ordering, body_ordering=body_ordering, ) - object.__setattr__(articulation, "_sim_cfg", None) - object.__setattr__(articulation, "_fixed_tendon_target_dirty", False) + articulation._sim_cfg = None + articulation._fixed_tendon_target_dirty = False # Create PhysX mock view mock_view = PhysXMockArticulationViewWarp( @@ -108,103 +110,85 @@ def create_physx_articulation( 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) + mock_view._shared_metatype = mock_metatype - object.__setattr__(articulation, "_root_view", mock_view) - object.__setattr__(articulation, "_device", device) + articulation._root_view = mock_view + articulation._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. # The data reads gravity from the PhysX manager patched above. data = PhysXArticulationData(mock_view, device) - object.__setattr__(articulation, "_data", data) + 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) + articulation._fixed_tendon_names = fixed_tendon_names + 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) + articulation._instantaneous_wrench_composer = mock_inst_wrench + 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) + articulation._initialize_handle = None + articulation._invalidate_initialize_handle = None + articulation._prim_deletion_handle = None + 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) - ) + articulation._ALL_INDICES = wp.array(np.arange(num_instances, dtype=np.int32), device=device) + articulation._ALL_BODY_INDICES = wp.array(np.arange(num_bodies, dtype=np.int32), device=device) + 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), - ) + articulation._ALL_FIXED_TENDON_INDICES = wp.array(np.arange(num_fixed_tendons, dtype=np.int32), device=device) + 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) - ) + articulation._ALL_INDICES_WP = wp.array(np.arange(num_instances, dtype=np.int32), device=device) + 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._joint_pos_target_backend = None + articulation._joint_vel_target_backend = None + articulation._joint_effort_target_backend = None + articulation._body_wrench_force_backend = None + 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) + articulation._root_link_pose_w_f32 = None + articulation._root_com_vel_w_f32 = None + articulation._root_link_vel_w_f32 = None + articulation._inst_wrench_force_f32 = None + articulation._inst_wrench_torque_f32 = None + articulation._perm_wrench_force_f32 = None + 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", {}) + articulation._sim_env_ids = wp.empty(N, dtype=wp.int32, device=device) + 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=wp.is_cuda_available()) - ) - 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._cpu_env_ids_all = cpu_env_ids + articulation._cpu_env_ids = wp.empty(N, dtype=wp.int32, device="cpu", pinned=wp.is_cuda_available()) + articulation._cpu_env_ids_views = {} + articulation._cpu_joint_stiffness = wp.zeros((N, J), dtype=wp.float32, device="cpu") + articulation._cpu_joint_damping = wp.zeros((N, J), dtype=wp.float32, device="cpu") + articulation._cpu_joint_pos_limits = wp.zeros((N, J, 2), dtype=wp.float32, device="cpu") + articulation._cpu_joint_vel_limits = wp.zeros((N, J), dtype=wp.float32, device="cpu") + articulation._cpu_joint_effort_limits = wp.zeros((N, J), dtype=wp.float32, device="cpu") + articulation._cpu_joint_armature = wp.zeros((N, J), dtype=wp.float32, device="cpu") + articulation._cpu_joint_friction_props = wp.zeros((N, J, 3), dtype=wp.float32, device="cpu") + articulation._cpu_body_mass = wp.zeros((N, B), dtype=wp.float32, device="cpu") + articulation._cpu_body_coms = wp.zeros((N, B, 7), dtype=wp.float32, device="cpu") + articulation._cpu_body_inertia = wp.zeros((N, B, 9), dtype=wp.float32, device="cpu") articulation._process_actuators_cfg() @@ -221,9 +205,11 @@ def create_ovphysx_articulation( is_fixed_base: bool = False, joint_ordering: tuple[str, ...] | None = None, body_ordering: tuple[str, ...] | None = None, + *, + monkeypatch: pytest.MonkeyPatch, ): """Create a test OvPhysX Articulation instance with mocked tensor bindings.""" - patch_ovphysx_manager() + patch_ovphysx_manager(monkeypatch=monkeypatch) joint_names = [f"joint_{i}" for i in range(num_joints)] body_names = [f"body_{i}" for i in range(num_bodies)] @@ -236,8 +222,8 @@ def create_ovphysx_articulation( joint_ordering=joint_ordering, body_ordering=body_ordering, ) - object.__setattr__(articulation, "_sim_cfg", None) - object.__setattr__(articulation, "_fixed_tendon_target_dirty", False) + articulation._sim_cfg = None + articulation._fixed_tendon_target_dirty = False # Create mock binding set mock_bindings = MockOvPhysxBindingSet( @@ -255,20 +241,20 @@ def create_ovphysx_articulation( 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) + articulation._device = device + articulation._ovphysx = MagicMock() + articulation._root_view = mock_bindings.view + articulation._bindings = mock_bindings.bindings + articulation._num_instances = num_instances + articulation._num_joints = num_joints + articulation._num_bodies = num_bodies + articulation._is_fixed_base = is_fixed_base + articulation._joint_names = joint_names + articulation._body_names = body_names + articulation._fixed_tendon_names = fixed_tendon_names + articulation._spatial_tendon_names = spatial_tendon_names + articulation._num_fixed_tendons = num_fixed_tendons + articulation._num_spatial_tendons = num_spatial_tendons # Create ArticulationData; counts come from the view, names are set after. data = OvPhysxArticulationData(mock_bindings.view, device) @@ -277,7 +263,7 @@ def create_ovphysx_articulation( 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) + articulation._data = data # Allocate the articulation-side index/mask caches and wrench buffer that # _initialize_impl would normally populate. Wrench composers created here @@ -288,24 +274,20 @@ def create_ovphysx_articulation( # 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) + articulation._instantaneous_wrench_composer = mock_inst_wrench + 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._initialize_handle = None + articulation._invalidate_initialize_handle = None + articulation._prim_deletion_handle = None + 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 - ) + 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, mock_bindings @@ -320,6 +302,8 @@ def create_newton_articulation( is_fixed_base: bool = False, joint_ordering: tuple[str, ...] | None = None, body_ordering: tuple[str, ...] | None = None, + *, + monkeypatch: pytest.MonkeyPatch, ): """Create a test Newton Articulation instance with mocked dependencies.""" import isaaclab_newton.assets.articulation.articulation_data as newton_data_module @@ -380,7 +364,7 @@ def create_newton_articulation( mock_manager.get_control.return_value = mock_control # Patch SimulationManager in the Newton data module until the test finishes. - patch_contract_manager(newton_data_module, "SimulationManager", mock_manager) + monkeypatch.setattr(newton_data_module, "SimulationManager", mock_manager, raising=False) data = NewtonArticulationData(mock_view, device) # Create Articulation shell (bypass __init__) @@ -393,64 +377,52 @@ def create_newton_articulation( joint_ordering=joint_ordering, body_ordering=body_ordering, ) - object.__setattr__(articulation, "_sim_cfg", None) - object.__setattr__(articulation, "_fixed_tendon_target_dirty", False) + articulation._sim_cfg = None + 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._root_view = mock_view + articulation._device = device + articulation._data = data + articulation._test_simulation_manager = mock_manager # the solver builds this adapter; the shell has no model, so it stays absent - object.__setattr__(articulation, "_fixed_tendon_control", None) + 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", []) + articulation._fixed_tendon_names = fixed_tendon_names + 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) + articulation._instantaneous_wrench_composer = mock_inst_wrench + 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) + articulation._initialize_handle = None + articulation._invalidate_initialize_handle = None + articulation._prim_deletion_handle = None + 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) - ) + articulation._ALL_INDICES = wp.array(np.arange(num_instances, dtype=np.int32), device=device) + articulation._ALL_BODY_INDICES = wp.array(np.arange(num_bodies, dtype=np.int32), device=device) + 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._ALL_ENV_MASK = wp.ones((num_instances,), dtype=wp.bool, device=device) + articulation._ALL_BODY_MASK = wp.ones((num_bodies,), dtype=wp.bool, device=device) + 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._ALL_FIXED_TENDON_INDICES = wp.array(np.arange(num_fixed_tendons, dtype=np.int32), device=device) + articulation._ALL_FIXED_TENDON_MASK = wp.ones((num_fixed_tendons,), dtype=wp.bool, device=device) + articulation._ALL_SPATIAL_TENDON_INDICES = wp.array(np.array([], dtype=np.int32), device=device) + articulation._ALL_SPATIAL_TENDON_MASK = wp.ones((0,), dtype=wp.bool, device=device) articulation._process_actuators_cfg() @@ -468,6 +440,8 @@ def get_articulation( is_fixed_base: bool = False, joint_ordering: tuple[str, ...] | None = None, body_ordering: tuple[str, ...] | None = None, + *, + monkeypatch: pytest.MonkeyPatch, ): if backend == "physx": return create_physx_articulation( @@ -480,6 +454,7 @@ def get_articulation( is_fixed_base=is_fixed_base, joint_ordering=joint_ordering, body_ordering=body_ordering, + monkeypatch=monkeypatch, ) elif backend == "ovphysx": return create_ovphysx_articulation( @@ -492,6 +467,7 @@ def get_articulation( is_fixed_base=is_fixed_base, joint_ordering=joint_ordering, body_ordering=body_ordering, + monkeypatch=monkeypatch, ) elif backend == "newton": return create_newton_articulation( @@ -504,6 +480,7 @@ def get_articulation( is_fixed_base=is_fixed_base, joint_ordering=joint_ordering, body_ordering=body_ordering, + monkeypatch=monkeypatch, ) else: raise ValueError(f"Invalid backend: {backend}") diff --git a/source/isaaclab/test/assets/contract/capabilities.py b/source/isaaclab/test/assets/contract/capabilities.py index a8d6948ce036..a26a48e1baf5 100644 --- a/source/isaaclab/test/assets/contract/capabilities.py +++ b/source/isaaclab/test/assets/contract/capabilities.py @@ -18,7 +18,7 @@ from _pytest.mark.structures import ParameterSet # Install the kitless import stubs before resolving the declared backend dependencies. -from . import _contract_boot # noqa: F401 +from . import contract_boot # noqa: F401 @dataclass(frozen=True) @@ -128,11 +128,6 @@ def backend_parameters(capability: str, *, names: tuple[str, ...] | None = None) return parameters -def contract_backend(capability: str, *, names: tuple[str, ...] | None = None) -> pytest.MarkDecorator: - """Parametrize ``backend`` over the declared backends for one capability.""" - return pytest.mark.parametrize("backend", backend_parameters(capability, names=names)) - - def requires_backend(name: str) -> pytest.MarkDecorator: """Skip a backend-specific test with the declared reason when that backend is unavailable.""" reason = _STATUS_BY_NAME[name].reason diff --git a/source/isaaclab/test/assets/contract/conftest.py b/source/isaaclab/test/assets/contract/conftest.py index 11701c696da6..5b1048be3538 100644 --- a/source/isaaclab/test/assets/contract/conftest.py +++ b/source/isaaclab/test/assets/contract/conftest.py @@ -3,17 +3,22 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Fixtures shared by the asset contract tests.""" - -from collections.abc import Iterator +"""Backend and device parameters shared by the asset contracts.""" import pytest -from ._manager_patch_scope import contract_manager_patch_scope +from isaaclab.test.utils import DeviceScope, test_devices + +from .capabilities import backend_parameters + + +@pytest.fixture(params=backend_parameters("api")) +def backend(request: pytest.FixtureRequest) -> str: + """Run each shared contract on every declared backend, with explicit unavailable skips.""" + return request.param -@pytest.fixture(autouse=True) -def _restore_contract_manager_patches() -> Iterator[None]: - """Limit factory manager substitutions to one test invocation.""" - with contract_manager_patch_scope(): - yield +@pytest.fixture(params=test_devices(DeviceScope.CPU_AND_DEFAULT_CUDA)) +def device(request: pytest.FixtureRequest) -> str: + """Exercise CPU buffers and CUDA staging without repeating CPU-only bookkeeping.""" + return request.param diff --git a/source/isaaclab/test/assets/contract/_contract_boot.py b/source/isaaclab/test/assets/contract/contract_boot.py similarity index 100% rename from source/isaaclab/test/assets/contract/_contract_boot.py rename to source/isaaclab/test/assets/contract/contract_boot.py diff --git a/source/isaaclab/test/assets/contract/_mock_backends.py b/source/isaaclab/test/assets/contract/mock_backends.py similarity index 69% rename from source/isaaclab/test/assets/contract/_mock_backends.py rename to source/isaaclab/test/assets/contract/mock_backends.py index 1d9b14853d28..c9fe9dbcf5f0 100644 --- a/source/isaaclab/test/assets/contract/_mock_backends.py +++ b/source/isaaclab/test/assets/contract/mock_backends.py @@ -9,28 +9,27 @@ from unittest.mock import MagicMock import numpy as np +import pytest import warp as wp -from ._manager_patch_scope import patch_contract_manager - -def patch_physx_manager() -> None: +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) - patch_contract_manager(PhysxManager, "get_physics_sim_view", MagicMock(return_value=physics_sim_view)) + 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. - patch_contract_manager(PhysxManager, "_scene_data_backend", PhysxSceneDataBackend()) + monkeypatch.setattr(PhysxManager, "_scene_data_backend", PhysxSceneDataBackend(), raising=False) -def patch_ovphysx_manager() -> None: +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 - patch_contract_manager(OvPhysxManager, "_scene_data_backend", OvPhysxSceneDataBackend()) + monkeypatch.setattr(OvPhysxManager, "_scene_data_backend", OvPhysxSceneDataBackend(), raising=False) def install_physx_recording_setters(mock_view, storage_by_method: dict[str, str]) -> None: @@ -68,3 +67,18 @@ def setter(values: wp.array, indices: wp.array | None = None) -> None: 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 + 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}") diff --git a/source/isaaclab/test/assets/contract/_rigid_object_collection_contract_utils.py b/source/isaaclab/test/assets/contract/rigid_object_collection_factory.py similarity index 63% rename from source/isaaclab/test/assets/contract/_rigid_object_collection_contract_utils.py rename to source/isaaclab/test/assets/contract/rigid_object_collection_factory.py index 371b3e574d25..4f9442c80162 100644 --- a/source/isaaclab/test/assets/contract/_rigid_object_collection_contract_utils.py +++ b/source/isaaclab/test/assets/contract/rigid_object_collection_factory.py @@ -11,15 +11,15 @@ from unittest.mock import MagicMock import numpy as np +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 isaaclab.utils.wrench_composer import WrenchComposer -from ._manager_patch_scope import patch_contract_manager -from ._mock_backends import install_physx_recording_setters, patch_ovphysx_manager, patch_physx_manager from .capabilities import available_backends +from .mock_backends import install_physx_recording_setters, patch_ovphysx_manager, patch_physx_manager BACKENDS = available_backends() @@ -60,12 +60,10 @@ def create_physx_rigid_object_collection( - num_instances: int = 2, - num_bodies: int = 3, - device: str = "cuda:0", + num_instances: int = 2, num_bodies: int = 3, device: str = "cuda:0", *, monkeypatch: pytest.MonkeyPatch ): """Create a test RigidObjectCollection instance with mocked dependencies.""" - patch_physx_manager() + patch_physx_manager(monkeypatch=monkeypatch) collection = object.__new__(PhysXRigidObjectCollection) rigid_objects = {f"object_{i}": RigidObjectCfg(prim_path=f"/World/Object_{i}") for i in range(num_bodies)} @@ -79,54 +77,48 @@ def create_physx_rigid_object_collection( mock_view.set_random_mock_data() install_physx_recording_setters(mock_view, _PHYSX_RIGID_BODY_STORAGE) - 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)]) + collection._root_view = mock_view + collection._device = device + collection._num_bodies = num_bodies + collection._num_instances = num_instances + 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) + 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) + collection._instantaneous_wrench_composer = mock_inst_wrench + 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) + collection._initialize_handle = None + collection._invalidate_initialize_handle = None + collection._prim_deletion_handle = None + 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)) + collection._ALL_ENV_INDICES = wp.array(np.arange(num_instances, dtype=np.int32), device=device) + 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", {}) + collection._ALL_VIEW_INDICES = all_view_ids + collection._sim_view_ids = wp.empty(num_view_ids, dtype=wp.int32, device=device) + collection._sim_view_ids_views = {} cpu_all_view_ids = wp.empty(num_view_ids, dtype=wp.int32, device="cpu", pinned=wp.is_cuda_available()) 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=wp.is_cuda_available()) - ) - object.__setattr__(collection, "_cpu_view_ids_views", {}) + collection._cpu_all_view_ids = cpu_all_view_ids + collection._cpu_view_ids = wp.empty(num_view_ids, dtype=wp.int32, device="cpu", pinned=wp.is_cuda_available()) + 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", + num_instances: int = 2, num_bodies: int = 3, device: str = "cuda:0", *, monkeypatch: pytest.MonkeyPatch ): """Create a test Newton RigidObjectCollection instance with mocked dependencies.""" import isaaclab_newton.assets.rigid_object_collection.rigid_object_collection as newton_coll_module @@ -161,8 +153,8 @@ def create_newton_rigid_object_collection( mock_manager.get_control.return_value = mock_control # Patch SimulationManager in both data and collection modules until the test finishes. - patch_contract_manager(newton_data_module, "SimulationManager", mock_manager) - patch_contract_manager(newton_coll_module, "SimulationManager", mock_manager) + monkeypatch.setattr(newton_data_module, "SimulationManager", mock_manager, raising=False) + monkeypatch.setattr(newton_coll_module, "SimulationManager", mock_manager, raising=False) data = NewtonRigidObjectCollectionData(mock_view, num_bodies, device) # Create collection shell (bypass __init__) @@ -171,44 +163,40 @@ def create_newton_rigid_object_collection( 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) + collection._root_view = mock_view + collection._device = device + collection._num_bodies = num_bodies + collection._num_instances = num_instances + collection._body_names_list = body_names + 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) + collection._instantaneous_wrench_composer = mock_inst_wrench + 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) + collection._initialize_handle = None + collection._invalidate_initialize_handle = None + collection._prim_deletion_handle = None + 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)) + collection._ALL_ENV_INDICES = wp.array(np.arange(num_instances, dtype=np.int32), device=device) + collection._ALL_BODY_INDICES = wp.array(np.arange(num_bodies, dtype=np.int32), device=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, mock_view def create_ovphysx_rigid_object_collection( - num_instances: int = 2, - num_bodies: int = 3, - device: str = "cuda:0", + num_instances: int = 2, num_bodies: int = 3, device: str = "cuda:0", *, monkeypatch: pytest.MonkeyPatch ): """Create a test OVPhysX RigidObjectCollection instance with mocked tensor bindings.""" - patch_ovphysx_manager() + patch_ovphysx_manager(monkeypatch=monkeypatch) body_names = [f"object_{i}" for i in range(num_bodies)] collection = object.__new__(OvPhysxRigidObjectCollection) @@ -226,20 +214,20 @@ def create_ovphysx_rigid_object_collection( ) 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) + collection._device = device + collection._ovphysx = MagicMock() + collection._root_view = mock_bindings.view + collection._bindings = mock_bindings.bindings + collection._num_instances = num_instances + collection._num_bodies = num_bodies + 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) + collection._data = data # Allocate the buffers that RigidObjectCollection normally allocates in _initialize_impl. collection._create_buffers() @@ -247,14 +235,14 @@ def create_ovphysx_rigid_object_collection( # 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) + collection._instantaneous_wrench_composer = mock_inst_wrench + 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) + collection._initialize_handle = None + collection._invalidate_initialize_handle = None + collection._prim_deletion_handle = None + collection._debug_vis_handle = None return collection, mock_bindings @@ -264,12 +252,14 @@ def get_rigid_object_collection( num_instances: int = 2, num_bodies: int = 3, device: str = "cuda:0", + *, + monkeypatch: pytest.MonkeyPatch, ): if backend == "physx": - return create_physx_rigid_object_collection(num_instances, num_bodies, device) + return create_physx_rigid_object_collection(num_instances, num_bodies, device, monkeypatch=monkeypatch) elif backend == "ovphysx": - return create_ovphysx_rigid_object_collection(num_instances, num_bodies, device) + return create_ovphysx_rigid_object_collection(num_instances, num_bodies, device, monkeypatch=monkeypatch) elif backend == "newton": - return create_newton_rigid_object_collection(num_instances, num_bodies, device) + return create_newton_rigid_object_collection(num_instances, num_bodies, device, monkeypatch=monkeypatch) else: raise ValueError(f"Invalid backend: {backend}") diff --git a/source/isaaclab/test/assets/contract/_rigid_object_contract_utils.py b/source/isaaclab/test/assets/contract/rigid_object_factory.py similarity index 58% rename from source/isaaclab/test/assets/contract/_rigid_object_contract_utils.py rename to source/isaaclab/test/assets/contract/rigid_object_factory.py index 482ec3f67b03..7caad02b17a0 100644 --- a/source/isaaclab/test/assets/contract/_rigid_object_contract_utils.py +++ b/source/isaaclab/test/assets/contract/rigid_object_factory.py @@ -11,14 +11,14 @@ from unittest.mock import MagicMock import numpy as np +import pytest import warp as wp from isaaclab.assets.rigid_object.rigid_object_cfg import RigidObjectCfg from isaaclab.utils.wrench_composer import WrenchComposer -from ._manager_patch_scope import patch_contract_manager -from ._mock_backends import install_physx_recording_setters, patch_ovphysx_manager, patch_physx_manager from .capabilities import available_backends +from .mock_backends import install_physx_recording_setters, patch_ovphysx_manager, patch_physx_manager BACKENDS = available_backends() @@ -46,12 +46,9 @@ } -def create_physx_rigid_object( - num_instances: int = 2, - device: str = "cuda:0", -): +def create_physx_rigid_object(num_instances: int = 2, device: str = "cuda:0", *, monkeypatch: pytest.MonkeyPatch): """Create a test RigidObject instance with mocked dependencies.""" - patch_physx_manager() + patch_physx_manager(monkeypatch=monkeypatch) body_names = ["body_0"] rigid_object = object.__new__(PhysXRigidObject) @@ -66,12 +63,12 @@ def create_physx_rigid_object( mock_view.set_random_mock_data() install_physx_recording_setters(mock_view, _PHYSX_RIGID_BODY_STORAGE) - object.__setattr__(rigid_object, "_root_view", mock_view) - object.__setattr__(rigid_object, "_device", device) + rigid_object._root_view = mock_view + rigid_object._device = device # The data reads gravity from the PhysX manager patched above. data = PhysXRigidObjectData(mock_view, device) - object.__setattr__(rigid_object, "_data", data) + rigid_object._data = data # Set body names on data data.body_names = body_names @@ -79,44 +76,39 @@ def create_physx_rigid_object( # 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) + rigid_object._instantaneous_wrench_composer = mock_inst_wrench + 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) + rigid_object._initialize_handle = None + rigid_object._invalidate_initialize_handle = None + rigid_object._prim_deletion_handle = None + 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)) + rigid_object._ALL_INDICES = wp.array(np.arange(num_instances, dtype=np.int32), device=device) + 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) + rigid_object._root_link_pose_w_f32 = None + 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", {}) + rigid_object._sim_env_ids = wp.empty(N, dtype=wp.int32, device=device) + 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=wp.is_cuda_available()) - ) - 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")) + rigid_object._cpu_env_ids_all = cpu_env_ids + rigid_object._cpu_env_ids = wp.empty(N, dtype=wp.int32, device="cpu", pinned=wp.is_cuda_available()) + rigid_object._cpu_env_ids_views = {} + rigid_object._cpu_body_mass = wp.zeros((N, B), dtype=wp.float32, device="cpu") + rigid_object._cpu_body_coms = wp.zeros((N, B, 7), dtype=wp.float32, device="cpu") + 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", -): +def create_newton_rigid_object(num_instances: int = 2, device: str = "cuda:0", *, monkeypatch: pytest.MonkeyPatch): """Create a test Newton RigidObject instance with mocked dependencies.""" import isaaclab_newton.assets.rigid_object.rigid_object_data as newton_data_module @@ -152,7 +144,7 @@ def create_newton_rigid_object( mock_manager.get_control.return_value = mock_control # Patch SimulationManager in the Newton data module until the test finishes. - patch_contract_manager(newton_data_module, "SimulationManager", mock_manager) + monkeypatch.setattr(newton_data_module, "SimulationManager", mock_manager, raising=False) data = NewtonRigidObjectData(mock_view, device) # Create RigidObject shell (bypass __init__) @@ -160,39 +152,36 @@ def create_newton_rigid_object( 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) + rigid_object._root_view = mock_view + rigid_object._device = device + 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) + rigid_object._instantaneous_wrench_composer = mock_inst_wrench + 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) + rigid_object._initialize_handle = None + rigid_object._invalidate_initialize_handle = None + rigid_object._prim_deletion_handle = None + 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)) + rigid_object._ALL_INDICES = wp.array(np.arange(num_instances, dtype=np.int32), device=device) + 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)) + rigid_object._ALL_ENV_MASK = wp.ones((num_instances,), dtype=wp.bool, device=device) + 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", -): +def create_ovphysx_rigid_object(num_instances: int = 2, device: str = "cuda:0", *, monkeypatch: pytest.MonkeyPatch): """Create a test OvPhysX RigidObject instance with mocked tensor bindings.""" - patch_ovphysx_manager() + patch_ovphysx_manager(monkeypatch=monkeypatch) body_names = ["base_link"] obj = object.__new__(OvPhysxRigidObject) @@ -209,20 +198,20 @@ def create_ovphysx_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) + obj._device = device + obj._ovphysx = MagicMock() + obj._root_view = mock_bindings.view + obj._bindings = mock_bindings.bindings + obj._num_instances = num_instances + obj._num_bodies = 1 + 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) + obj._data = data # Build the buffers RigidObject normally allocates in _initialize_impl # (_ALL_INDICES, _ALL_*_MASK, pinned CPU staging buffers, wrench buf). @@ -233,28 +222,24 @@ def create_ovphysx_rigid_object( # 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) + obj._instantaneous_wrench_composer = mock_inst_wrench + 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) + obj._initialize_handle = None + obj._invalidate_initialize_handle = None + obj._prim_deletion_handle = None + obj._debug_vis_handle = None return obj, mock_bindings -def get_rigid_object( - backend: str, - num_instances: int = 2, - device: str = "cuda:0", -): +def get_rigid_object(backend: str, num_instances: int = 2, device: str = "cuda:0", *, monkeypatch: pytest.MonkeyPatch): if backend == "physx": - return create_physx_rigid_object(num_instances, device) + return create_physx_rigid_object(num_instances, device, monkeypatch=monkeypatch) elif backend == "ovphysx": - return create_ovphysx_rigid_object(num_instances, device) + return create_ovphysx_rigid_object(num_instances, device, monkeypatch=monkeypatch) elif backend == "newton": - return create_newton_rigid_object(num_instances, device) + return create_newton_rigid_object(num_instances, device, monkeypatch=monkeypatch) else: raise ValueError(f"Invalid backend: {backend}") 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..38ef1ad62c61 --- /dev/null +++ b/source/isaaclab/test/assets/contract/test_articulation_contract.py @@ -0,0 +1,1613 @@ +# 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 and writer behavior the base articulation class advertises. +The backends run on mocked views, so these cases need neither Isaac Sim nor a GPU simulation. +""" + +import math +import warnings +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +import torch +import warp as wp + +from .articulation_factory import get_articulation +from .capabilities import backend_parameters, require_backend_capability, requires_backend +from .mock_backends import read_backend_joint_state + +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. +# One fixture with distinct instance, joint, and body counts so any swapped axis shows up. +_NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES = 2, 6, 7 + + +# --------------------------------------------------------------------------- +# Tests: Index resolution helpers +# --------------------------------------------------------------------------- + + +def test_resolve_env_ids_handles_tensor_view_shape(monkeypatch, backend, device): + art, _ = get_articulation(backend, num_instances=4, device=device, monkeypatch=monkeypatch) + + 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() + + +def test_resolve_joint_ids_handles_tensor_view_shape(monkeypatch, backend): + art, _ = get_articulation(backend, num_joints=4, device="cpu", monkeypatch=monkeypatch) + + 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 + + +def test_resolve_body_ids_handles_tensor_view_shape(monkeypatch, backend): + art, _ = get_articulation(backend, num_bodies=4, device="cpu", monkeypatch=monkeypatch) + + 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 +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("is_fixed_base", [False, True], ids=["floating", "fixed"]) +def test_articulation_counts_names_and_finders(monkeypatch, 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, + monkeypatch=monkeypatch, + ) + + 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) + + +@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(monkeypatch, 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", + monkeypatch=monkeypatch, + ) + if finder_name == "find_spatial_tendons": + require_backend_capability(backend, "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]) + + +def _fixed_tendon_articulation(backend: str, *, monkeypatch: pytest.MonkeyPatch): + """Create a two-environment articulation with two fixed tendons and no spatial tendons.""" + art, _ = get_articulation( + backend, + num_instances=2, + num_joints=2, + num_bodies=2, + num_fixed_tendons=2, + num_spatial_tendons=0, + device="cpu", + monkeypatch=monkeypatch, + ) + return art + + +@pytest.mark.parametrize("backend", backend_parameters("fixed_tendon_target_scheduling")) +def test_commanding_a_target_schedules_the_write(monkeypatch, backend): + """The target setter marks the write pending; a static offset write keeps its own contract.""" + art = _fixed_tendon_articulation(backend, monkeypatch=monkeypatch) + art.set_fixed_tendon_position_target_index(target=torch.zeros((2, 2), dtype=torch.float32)) + assert art._fixed_tendon_target_dirty is True + + +@requires_backend("newton") +def test_newton_reports_a_missing_tendon_actuator(monkeypatch): + """Without a MuJoCo tendon actuator the solver has no adapter, so commanding must say so.""" + art = _fixed_tendon_articulation("newton", monkeypatch=monkeypatch) + with pytest.raises(RuntimeError, match="no MuJoCo tendon actuator"): + art.set_fixed_tendon_position_target_index(target=torch.zeros((2, 2), dtype=torch.float32)) + + +@requires_backend("newton") +def test_a_solver_without_tendon_transmission_refuses_to_build_an_adapter(monkeypatch): + """A solver with no tendon transmission says so, rather than returning nothing.""" + from isaaclab_newton.physics import NewtonManager + + art = _fixed_tendon_articulation("newton", monkeypatch=monkeypatch) + with pytest.raises(NotImplementedError, match="does not drive fixed tendons"): + NewtonManager.create_fixed_tendon_control(art) + + +# --------------------------------------------------------------------------- +# Tests: resolve_matching_names caching behavior +# --------------------------------------------------------------------------- + + +def test_unmatched_regex_raises(monkeypatch, backend): + """ValueError from resolve_matching_names propagates correctly.""" + art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device="cpu", monkeypatch=monkeypatch) + with pytest.raises(ValueError): + art.find_bodies("nonexistent_body_xyz") + with pytest.raises(ValueError): + art.find_joints("nonexistent_joint_xyz") + + +def test_mutating_result_does_not_corrupt_cache(monkeypatch, backend): + """Mutating returned lists must not affect future cached results.""" + art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device="cpu", monkeypatch=monkeypatch) + + 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 + + +def test_find_with_preserve_order(monkeypatch, 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", monkeypatch=monkeypatch) + 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), +] + + +def test_articulation_data_property_contract(monkeypatch, backend, device): + art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device, monkeypatch=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) + 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() + + +@requires_backend("physx") +def test_physx_set_coms_index_updates_body_com_pose_b_cache(monkeypatch): + art, view = get_articulation( + "physx", num_instances=2, num_joints=3, num_bodies=4, device="cpu", monkeypatch=monkeypatch + ) + + 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 + + +@requires_backend("physx") +def test_physx_joint_position_write_preserves_body_com_pose_b_cache(monkeypatch): + art, view = get_articulation( + "physx", num_instances=2, num_joints=3, num_bodies=4, device="cpu", monkeypatch=monkeypatch + ) + + 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", backend_parameters("mocked_dynamics")) +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, _ = get_articulation( + backend, num_instances=2, num_joints=3, num_bodies=4, device=device, monkeypatch=monkeypatch + ) + 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 + + +def test_actuator_compatibility_projections_are_stable(monkeypatch, backend): + num_instances, num_joints = 2, 4 + art, _ = get_articulation(backend, num_instances, num_joints, 5, device="cpu", monkeypatch=monkeypatch) + 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] + + +def test_deprecated_command_aliases_return_actuator_collection_buffers(monkeypatch, backend): + """The deprecated data command and torque properties warn and alias the actuator collection's buffers.""" + art, _ = get_articulation(backend, 2, 4, 5, device="cpu", monkeypatch=monkeypatch) + actuators = art.actuators + 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 + + +# --------------------------------------------------------------------------- +# 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"), +] + + +def test_aliases_match_canonical_values(monkeypatch, 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", monkeypatch=monkeypatch) + 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", +} + + +@pytest.mark.parametrize( + "body_ordering", + [None, ("body_0", "body_3", "body_2", "body_1")], +) +def test_pose_write_invalidates_pose_dependent_caches(monkeypatch, backend, body_ordering): + art, _ = get_articulation( + backend, + num_instances=2, + num_joints=3, + num_bodies=4, + device="cpu", + body_ordering=body_ordering, + monkeypatch=monkeypatch, + ) + 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) + + +@pytest.mark.parametrize( + "body_ordering", + [None, ("body_0", "body_3", "body_2", "body_1")], +) +def test_velocity_write_invalidates_body_frame_caches(monkeypatch, backend, body_ordering): + art, _ = get_articulation( + backend, + num_instances=2, + num_joints=3, + num_bodies=4, + device="cpu", + body_ordering=body_ordering, + monkeypatch=monkeypatch, + ) + 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) + + +@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(monkeypatch, backend, setter_kind, body_ordering): + art, _ = get_articulation( + backend, + num_instances=2, + num_joints=3, + num_bodies=4, + device="cpu", + body_ordering=body_ordering, + monkeypatch=monkeypatch, + ) + 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 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 -- + + +@pytest.mark.parametrize("method_suffix", _ROOT_POSE_METHODS) +def test_write_root_pose_to_sim_index(monkeypatch, backend, device, method_suffix): + art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device, monkeypatch=monkeypatch) + 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)) + + +@pytest.mark.parametrize("method_suffix", _ROOT_VEL_METHODS) +def test_write_root_velocity_to_sim_index(monkeypatch, backend, device, method_suffix): + art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device, monkeypatch=monkeypatch) + 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 -- + + +@pytest.mark.parametrize("method_suffix", _ROOT_POSE_METHODS) +def test_write_root_pose_to_sim_mask(monkeypatch, backend, device, method_suffix): + art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device, monkeypatch=monkeypatch) + 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)) + + +@pytest.mark.parametrize("method_suffix", _ROOT_VEL_METHODS) +def test_write_root_velocity_to_sim_mask(monkeypatch, backend, device, method_suffix): + art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device, monkeypatch=monkeypatch) + 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)) + + +def test_deprecated_joint_friction_writers(monkeypatch, 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", monkeypatch=monkeypatch) + 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 + + +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 = get_articulation(backend, 2, 3, 2, device="cpu", monkeypatch=monkeypatch) + 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, [(name, 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_joint_state_writer_matches_public_writers(monkeypatch, backend): + """The deprecated combined joint-state writer warns and writes what the public writers would.""" + art, raw_backend = get_articulation(backend, 2, 3, 2, device="cpu", monkeypatch=monkeypatch) + 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) + + +@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(monkeypatch, backend, device, method_base, kwarg, wp_dtype, accepts_float, getter): + art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device, monkeypatch=monkeypatch) + 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)}) + + +@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(monkeypatch, backend, device, method_base, kwarg, wp_dtype, accepts_float, getter): + art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device, monkeypatch=monkeypatch) + 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) + + +@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(monkeypatch, backend, device, method_base, kwarg): + art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device, monkeypatch=monkeypatch) + 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)}) + + +@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(monkeypatch, backend, device, method_base, kwarg): + art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device, monkeypatch=monkeypatch) + 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)}) + + +# --------------------------------------------------------------------------- +# Tests: partial root and body writes reach only the selected backend rows +# --------------------------------------------------------------------------- + +# quantity -> (writer base name, writer keyword, trailing size). Joint partial writes are covered, with and without +# a joint ordering, by ``TestArticulationOrderingWriteParity``. +_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 == "physx": + getter = { + "root_pose": raw_backend.get_root_transforms, + "root_velocity": raw_backend.get_root_velocities, + "mass": raw_backend.get_masses, + } + values = wp.to_torch(getter[quantity]()) + elif backend == "newton": + if quantity == "mass": + values = wp.to_torch(art.data._sim_bind_body_mass) + else: + getter = {"root_pose": raw_backend.get_root_transforms, "root_velocity": raw_backend.get_root_velocities} + values = wp.to_torch(getter[quantity](None)) + else: + 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) + 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): + num_instances, num_bodies = 2, 4 + art, raw_backend = get_articulation(backend, num_instances, 3, num_bodies, device="cpu", monkeypatch=monkeypatch) + writer, kwarg, trailing = _PARTIAL_WRITES[quantity] + width = num_bodies if quantity == "mass" else trailing + # Literal, per-element distinct payload; poses keep an identity rotation. + values = 100.0 * torch.arange(1, num_instances + 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 + kwargs = {} + if selection == "index": + kwargs[kwarg] = values[env : env + 1, body : body + 1] if quantity == "mass" else values[env : env + 1] + kwargs["env_ids"] = torch.tensor([env], dtype=torch.int32) + if quantity == "mass": + kwargs["body_ids"] = [body] + else: + kwargs[kwarg] = values + kwargs["env_mask"] = _make_item_mask(num_instances, [env], "cpu") + if quantity == "mass": + kwargs["body_mask"] = _make_item_mask(num_bodies, [body], "cpu") + + getattr(art, f"{writer}_{selection}")(**kwargs) + + after = _read_backend_rows(backend, art, raw_backend, quantity) + expected = before.clone() + if quantity == "mass": + expected[env, body] = values[env, body] + else: + expected[env] = values[env] + torch.testing.assert_close(after, expected, rtol=0.0, atol=0.0) + + +# --------------------------------------------------------------------------- +# Tendon tests — counts, names, data, writers +# --------------------------------------------------------------------------- + +# Every backend supports the core fixed-tendon properties. Newton declares the extended fixed-tendon properties and +# spatial tendons unsupported, so those rows stay visible as reasoned capability skips. + +# The first row covers the zero-spatial-tendon edge; the second has distinct counts on every axis. +_TENDON_DIMS = [ + (1, 2, 2, 1, 0), # fixed only + (2, 6, 7, 3, 2), # both types +] +_TENDON_DIM_NAMES = ("num_instances", "num_joints", "num_bodies", "num_fixed_tendons", "num_spatial_tendons") + +# tendon property family -> (tendon kind, backend capability) +_TENDON_FAMILIES = { + "fixed": ("fixed", "fixed_tendons"), + "fixed_extended": ("fixed", "fixed_tendon_extended"), + "spatial": ("spatial", "spatial_tendons"), +} + +# (property, family, dtype) +_TENDON_DATA_PROPERTIES = [ + ("fixed_tendon_stiffness", "fixed", wp.float32), + ("fixed_tendon_damping", "fixed", wp.float32), + ("fixed_tendon_pos_limits", "fixed", wp.vec2f), + ("fixed_tendon_limit_stiffness", "fixed_extended", wp.float32), + ("fixed_tendon_rest_length", "fixed_extended", wp.float32), + ("fixed_tendon_offset", "fixed_extended", wp.float32), + ("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), +] +_TENDON_PROPERTY_FAMILY = {name: family for name, family, _ in _TENDON_DATA_PROPERTIES} + + +@pytest.mark.parametrize(_TENDON_DIM_NAMES, _TENDON_DIMS) +@pytest.mark.parametrize("family", _TENDON_FAMILIES) +def test_articulation_tendon_contract( + monkeypatch, + backend, + num_instances, + num_joints, + num_bodies, + num_fixed_tendons, + num_spatial_tendons, + device, + family, +): + kind, capability = _TENDON_FAMILIES[family] + require_backend_capability(backend, capability) + art, _ = get_articulation( + backend, + num_instances, + num_joints, + num_bodies, + num_fixed_tendons, + num_spatial_tendons, + device, + monkeypatch=monkeypatch, + ) + num_tendons = num_fixed_tendons if kind == "fixed" else num_spatial_tendons + assert getattr(art, f"num_{kind}_tendons") == num_tendons + names = getattr(art, f"{kind}_tendon_names") + assert isinstance(names, list) + assert len(names) == num_tendons + assert all(isinstance(n, str) for n in names) + + art.data.update(dt=0.01) + for name, property_family, dtype in _TENDON_DATA_PROPERTIES: + if property_family != family or num_tendons == 0: + continue + expected_shape = (num_instances, num_tendons) + 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] +# Writer cases pair each method with the dimension rows that configure at least one tendon of its kind. +_TENDON_WRITER_CASES = [ + pytest.param(*method, *dims, id=f"{method[1]}-{'-'.join(map(str, dims))}") + for method in _TENDON_METHODS + for dims in _TENDON_DIMS + if dims[3 if method[0] == "fixed" else 4] > 0 +] + + +@pytest.mark.parametrize("selection", ["index", "mask"]) +@pytest.mark.parametrize( + ("kind", "method_base", "kwarg", "getter", *_TENDON_DIM_NAMES), + _TENDON_WRITER_CASES, +) +def test_tendon_writer( + monkeypatch, + backend, + device, + selection, + kind, + method_base, + kwarg, + getter, + num_instances, + num_joints, + num_bodies, + num_fixed_tendons, + num_spatial_tendons, +): + require_backend_capability(backend, _TENDON_FAMILIES[_TENDON_PROPERTY_FAMILY[getter]][1]) + if backend == "newton": + from isaaclab_newton.physics import NewtonManager + + # The mocked asset has no solver to receive model-change notifications. + monkeypatch.setattr(NewtonManager, "add_model_change", MagicMock()) + + num_tendons = num_fixed_tendons if kind == "fixed" else num_spatial_tendons + art, _ = get_articulation( + backend, + num_instances, + num_joints, + num_bodies, + num_fixed_tendons, + num_spatial_tendons, + device, + monkeypatch=monkeypatch, + ) + 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: after the push, the matching data property reads the written values back + method(**{kwarg: _make_payload_warp(full_shape, device, wp_dtype)}) + getattr(art, f"write_{kind}_tendon_properties_to_sim_index")() + _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}) + + +# Shape validation runs before any kernel launch, so one device covers it. + + +@pytest.mark.parametrize("array_kind", ["torch", "warp"]) +@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_rejects_extra_environments( + monkeypatch, request, backend, selection, array_kind, kind, method_base, kwarg, getter +): + require_backend_capability(backend, _TENDON_FAMILIES[_TENDON_PROPERTY_FAMILY[getter]][1]) + if backend == "newton" and selection == "mask": + # Product bug: Newton's fixed-tendon mask setters slice full data without validating its shape, so an + # oversized input is silently truncated. + request.applymarker( + pytest.mark.xfail( + raises=pytest.fail.Exception, + strict=True, + reason="Newton fixed-tendon mask setters accept data with extra environments", + ) + ) + num_instances, num_joints, num_bodies, num_fixed_tendons, num_spatial_tendons = _TENDON_DIMS[1] + art, _ = get_articulation( + backend, + num_instances, + num_joints, + num_bodies, + num_fixed_tendons, + num_spatial_tendons, + "cpu", + monkeypatch=monkeypatch, + ) + full_shape = (num_instances, num_fixed_tendons if kind == "fixed" else num_spatial_tendons) + wp_dtype = wp.vec2f if getter == "fixed_tendon_pos_limits" else wp.float32 + make_bad_data = _make_bad_data_torch if array_kind == "torch" else _make_bad_data_warp + + with pytest.raises((AssertionError, RuntimeError)): + getattr(art, f"{method_base}_{selection}")(**{kwarg: make_bad_data(full_shape, "cpu", wp_dtype)}) + + +@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": + from isaaclab_newton.physics import NewtonManager + + # The mocked asset has no solver to receive model-change notifications. + monkeypatch.setattr(NewtonManager, "add_model_change", MagicMock()) + + require_backend_capability(backend, _TENDON_FAMILIES[kind][1]) + num_instances = 3 + art, raw_backend = get_articulation(backend, num_instances, 2, 2, 2, 2, device, monkeypatch=monkeypatch) + 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() + if selection == "index": + env_ids = None if len(selected_envs) == num_instances else torch.tensor(selected_envs, device=device) + write_kwargs = {"env_ids": env_ids} + 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_kwargs = {"env_mask": mask} + write(**write_kwargs) + 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 diff --git a/source/isaaclab/test/assets/contract/test_articulation_ordering_contract.py b/source/isaaclab/test/assets/contract/test_articulation_ordering_contract.py new file mode 100644 index 000000000000..4c955bbc5fa0 --- /dev/null +++ b/source/isaaclab/test/assets/contract/test_articulation_ordering_contract.py @@ -0,0 +1,2367 @@ +# 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 joint- and body-ordering contract cases.""" + +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +import torch +import warp as wp +from _pytest.mark.structures import ParameterSet + +from isaaclab.utils import replace +from isaaclab.utils.buffers import TimestampedBuffer + +from .articulation_factory import get_articulation +from .capabilities import backend_parameters, requires_backend +from .mock_backends import read_backend_joint_state + +pytestmark = pytest.mark.integration + + +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 _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, *, monkeypatch: pytest.MonkeyPatch +) -> 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, + monkeypatch=monkeypatch, + ) + 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, *, monkeypatch: pytest.MonkeyPatch) -> 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), + monkeypatch=monkeypatch, + ) + 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, *, monkeypatch: pytest.MonkeyPatch +) -> 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), + monkeypatch=monkeypatch, + ) + 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 an ordering parameter row that carries the backend's declared skip reason.""" + marks = [*kwargs.pop("marks", ()), *backend_parameters("ordering", names=(backend,))[0].marks] + return pytest.param(backend, *values, marks=marks, **kwargs) + + +def test_reversed_body_ordering_reorders_public_body_quantities(monkeypatch, backend): + num_instances = 2 + num_joints = 1 + num_bodies = 3 + device = "cpu" + identity_art, identity_raw = get_articulation( + backend, num_instances, num_joints, num_bodies, device=device, monkeypatch=monkeypatch + ) + ordered_art, ordered_raw = get_articulation( + backend, num_instances, num_joints, num_bodies, device=device, monkeypatch=monkeypatch + ) + 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) + + +@pytest.mark.parametrize("ordering_mode", ["reversed", "cyclic"]) +def test_body_ordering_reorders_public_body_properties(monkeypatch, backend, ordering_mode): + """Expose mass and every inertia component under the matching public body name.""" + num_instances = 2 + num_joints = 1 + num_bodies = 3 + device = "cpu" + 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, + monkeypatch=monkeypatch, + ) + 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_backend("ovphysx") +def test_ovphysx_reversed_body_ordering_rereads_backend_shadows_after_reset(monkeypatch): + """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", monkeypatch=monkeypatch + ) + 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_backend("ovphysx") +def test_ovphysx_reversed_body_ordering_rereads_all_velocity_shadows_after_reset(monkeypatch): + """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"), monkeypatch=monkeypatch + ) + 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_backend("ovphysx") +def test_ovphysx_ordered_root_and_body_velocity_share_one_backend_read(monkeypatch) -> 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"), monkeypatch=monkeypatch + ) + 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 + + +@pytest.mark.parametrize("backend", backend_parameters("ordering", names=("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 = 2 + num_joints = 3 + num_bodies = 4 + device = "cpu" + identity_art, identity_raw = get_articulation( + backend, + num_instances, + num_joints, + num_bodies, + device=device, + is_fixed_base=is_fixed_base, + monkeypatch=monkeypatch, + ) + ordered_art, ordered_raw = get_articulation( + backend, + num_instances, + num_joints, + num_bodies, + device=device, + is_fixed_base=is_fixed_base, + monkeypatch=monkeypatch, + ) + 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 + + +_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": (), +} + + +@pytest.mark.parametrize("ordering_mode", ["none", "identity", "reversed"]) +def test_backends_allocate_shadows_only_for_nonidentity_ordering(monkeypatch, 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), + monkeypatch=monkeypatch, + ) + + 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, _DATA_SHADOWS[backend]), + (art, _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) + + +@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_backend("ovphysx") +def test_ovphysx_partial_com_write_preserves_backend_rows( + monkeypatch, 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), + monkeypatch=monkeypatch, + ) + 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_backend("newton") +def test_newton_partial_com_write_routes_to_backend_bodies(monkeypatch, 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), + monkeypatch=monkeypatch, + ) + 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]) + + +@pytest.mark.parametrize("backend", backend_parameters("ordering", names=("physx", "ovphysx"))) +@pytest.mark.parametrize("ordering_mode", ["none", "reversed"]) +def test_duplicate_body_ids_preserve_omitted_backend_com(monkeypatch, 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), + monkeypatch=monkeypatch, + ) + 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]) + + +@pytest.mark.parametrize("backend", backend_parameters("ordering", names=("physx", "ovphysx"))) +@pytest.mark.parametrize("ordering_mode", ["none", "reversed"]) +def test_duplicate_env_ids_leave_global_com_cache_invalid(monkeypatch, 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), + monkeypatch=monkeypatch, + ) + 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 + + +@pytest.mark.parametrize("backend", backend_parameters("ordering", names=("physx", "ovphysx"))) +@pytest.mark.parametrize("ordering_mode", ["none", "reversed"]) +def test_static_com_cache_does_not_follow_sim_timestamp(monkeypatch, 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), + monkeypatch=monkeypatch, + ) + 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 +# --------------------------------------------------------------------------- + + +@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( + monkeypatch, 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, + monkeypatch=monkeypatch, + ) + 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, + monkeypatch=monkeypatch, + ) + # 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) + + +# 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( + monkeypatch, backend: str, ordering_mode: str, selection: str, operation: str, coverage: str +) -> None: + _exercise_partial_joint_write(backend, ordering_mode, selection, operation, coverage, monkeypatch=monkeypatch) + + +@pytest.mark.parametrize( + "operation, expected_position_reads, expected_velocity_reads", + [("position", 1, 0), ("velocity", 0, 1), ("state", 1, 1)], +) +@requires_backend("physx") +def test_physx_mask_partial_joint_write_reads_each_staging_axis_once( + monkeypatch, 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), + monkeypatch=monkeypatch, + ) + 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) + + +@pytest.mark.parametrize("backend", backend_parameters("ordering", names=("physx", "ovphysx"))) +@pytest.mark.parametrize("operation", ["position", "velocity", "state"]) +def test_duplicate_full_length_joint_selector_preserves_newer_backend_rows( + monkeypatch, backend: str, operation: str +) -> None: + _exercise_duplicate_full_length_joint_write(backend, operation, monkeypatch=monkeypatch) + + +@pytest.mark.parametrize("selector", ["default_mask", "explicit_full_index"]) +@pytest.mark.parametrize("operation", ["position", "velocity", "state"]) +@requires_backend("physx") +def test_physx_joint_selector_provenance_controls_staging_read(monkeypatch, 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), + monkeypatch=monkeypatch, + ) + 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. +@pytest.mark.parametrize("backend", backend_parameters("ordering", names=("physx", "ovphysx"))) +@pytest.mark.parametrize("operation", ["position", "velocity", "state"]) +def test_stale_partial_joint_write_preserves_newer_backend_rows(monkeypatch, backend: str, operation: str) -> None: + _exercise_stale_partial_joint_write(backend, operation, monkeypatch=monkeypatch) + + +@requires_backend("ovphysx") +def test_ovphysx_partial_effort_target_write_preserves_unselected_backend_rows(monkeypatch) -> 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), + monkeypatch=monkeypatch, + ) + 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) + + +@requires_backend("ovphysx") +def test_ovphysx_joint_acceleration_differences_public_order_velocities(monkeypatch): + """Finite-difference OVPhysX joint velocity entirely in public joint order.""" + art, raw_backend = get_articulation("ovphysx", 2, 3, 2, device="cpu", monkeypatch=monkeypatch) + 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_backend("ovphysx") +def test_ovphysx_ordered_joint_state_is_cached_per_sim_timestamp(monkeypatch): + """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"), monkeypatch=monkeypatch + ) + 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 + + +def test_joint_ordering_reorders_public_joint_properties(monkeypatch, backend): + """Expose every backend joint property under the matching public joint name.""" + num_instances = 2 + num_joints = 3 + num_bodies = 2 + device = "cpu" + 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, + monkeypatch=monkeypatch, + ) + 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 +# --------------------------------------------------------------------------- + + +def test_legacy_position_target_accepts_joint_slice(monkeypatch, backend: str) -> None: + art, _ = get_articulation( + backend, num_instances=2, num_joints=4, num_bodies=2, device="cpu", monkeypatch=monkeypatch + ) + 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) + + +@pytest.mark.parametrize("ordering_mode", ["none", "reversed", "cyclic"]) +@pytest.mark.parametrize("is_fixed_base", [False, True], ids=["floating", "fixed"]) +def test_external_wrenches_are_written_in_backend_body_order(monkeypatch, backend, ordering_mode, is_fixed_base): + """Write public-order body wrenches to each backend in backend body order.""" + device = "cpu" + 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, + monkeypatch=monkeypatch, + ) + _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_backend("ovphysx") +def test_ovphysx_configured_defaults_use_public_joint_names(monkeypatch): + """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"), monkeypatch=monkeypatch + ) + 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_backend("ovphysx") +@pytest.mark.parametrize("ordering_mode", ["reversed", "cyclic"]) +def test_ovphysx_implicit_targets_are_written_in_backend_order(monkeypatch, 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, + monkeypatch=monkeypatch, + ) + 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_backend("ovphysx") +@pytest.mark.parametrize("selector_kind", ["torch", "warp"]) +def test_ovphysx_int64_effort_target_selector_reaches_binding( + 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", monkeypatch=monkeypatch) + 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, raising=False) + 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_backend("physx") +@pytest.mark.parametrize("ordering_mode", ["reversed", "cyclic"]) +def test_physx_newton_actuator_forces_are_written_in_backend_order(monkeypatch, 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, + monkeypatch=monkeypatch, + ) + 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_backend("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) + 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() + + +@pytest.mark.parametrize("ordering_mode", ["reversed", "cyclic"]) +@pytest.mark.parametrize("selection", ["index", "mask"]) +def test_reversed_joint_ordering_routes_property_writes_to_backend( + monkeypatch, 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, + monkeypatch=monkeypatch, + ) + 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], + ) + + +@pytest.mark.parametrize("ordering_mode", ["reversed", "cyclic"]) +@pytest.mark.parametrize("selection", ["index", "mask"]) +def test_body_ordering_routes_property_writes_to_backend(monkeypatch, 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, + monkeypatch=monkeypatch, + ) + 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/contract/test_asset_contract_api.py b/source/isaaclab/test/assets/contract/test_asset_contract_api.py index a51c348f0d62..bbac11d7379e 100644 --- a/source/isaaclab/test/assets/contract/test_asset_contract_api.py +++ b/source/isaaclab/test/assets/contract/test_asset_contract_api.py @@ -10,29 +10,11 @@ import pytest -from ._articulation_contract_cases import ( # noqa: F401 - TestArticulationFinderReturnModes, - TestArticulationIndexResolution, - TestArticulationProperties, - TestFixedTendonTargetScheduling, - TestResolveMatchingNamesCache, -) -from ._articulation_contract_utils import get_articulation -from ._manager_patch_scope import contract_manager_patch_scope -from ._rigid_object_collection_contract_cases import ( # noqa: F401 - TestCollectionFinderReturnModes, - TestCollectionIndexResolution, - TestCollectionProperties, -) -from ._rigid_object_collection_contract_utils import get_rigid_object_collection -from ._rigid_object_contract_cases import ( # noqa: F401 - TestRigidObjectFinderReturnModes, - TestRigidObjectIndexResolution, - TestRigidObjectProperties, -) -from ._rigid_object_contract_utils import get_rigid_object +from .articulation_factory import get_articulation from .capabilities import BACKEND_STATUSES, available_backends from .public_surface import BASE_SURFACE_CLASSES, PUBLIC_SURFACE_CONTRACTS, public_surface_mismatches +from .rigid_object_collection_factory import get_rigid_object_collection +from .rigid_object_factory import get_rigid_object pytestmark = pytest.mark.integration @@ -91,16 +73,20 @@ def test_contract_factories_scope_manager_patches_to_one_test(reverse: bool) -> assert all(manager is NewtonManager for name, manager in original.items() if name.startswith("newton.")) backends = available_backends()[::-1] if reverse else available_backends() - with contract_manager_patch_scope(): + patches = pytest.MonkeyPatch() + try: for backend in backends: - rigid_object, _ = get_rigid_object(backend, device="cpu") - collection, _ = get_rigid_object_collection(backend, device="cpu") - articulation, _ = get_articulation(backend, device="cpu") + rigid_object, _ = get_rigid_object(backend, device="cpu", monkeypatch=patches) + collection, _ = get_rigid_object_collection(backend, device="cpu", monkeypatch=patches) + articulation, _ = get_articulation(backend, device="cpu", monkeypatch=patches) # Later factories must not break the patched managers the earlier families read through. assert rigid_object.data.root_link_pose_w.shape == (2,) assert collection.data.body_link_pose_w.shape == (2, 3) assert articulation.data.root_link_pose_w.shape == (2,) + finally: + patches.undo() + restored = _manager_bindings() assert restored.keys() == original.keys() assert all(restored[name] is binding for name, binding in original.items()) diff --git a/source/isaaclab/test/assets/contract/test_asset_contract_data.py b/source/isaaclab/test/assets/contract/test_asset_contract_data.py deleted file mode 100644 index 717011671607..000000000000 --- a/source/isaaclab/test/assets/contract/test_asset_contract_data.py +++ /dev/null @@ -1,30 +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 - -"""Shared asset data contract: property shapes and dtypes, aliases, lazy buffers, and ordered reads.""" - -import pytest - -from ._articulation_contract_cases import ( # noqa: F401 - TestArticulationDataAliases, - TestArticulationDataProperties, - TestArticulationTendons, -) -from ._articulation_ordering_contract_cases import ( # noqa: F401 - TestArticulationOrderingAllocation, - TestArticulationOrderingBodyState, - TestArticulationOrderingJointState, -) -from ._rigid_object_collection_contract_cases import ( # noqa: F401 - TestCollectionDataAliases, - TestCollectionDataProperties, - TestCollectionViewReshape, -) -from ._rigid_object_contract_cases import ( # noqa: F401 - TestRigidObjectDataAliases, - TestRigidObjectDataProperties, -) - -pytestmark = pytest.mark.integration diff --git a/source/isaaclab/test/assets/contract/test_asset_contract_writes.py b/source/isaaclab/test/assets/contract/test_asset_contract_writes.py deleted file mode 100644 index 4805968d3179..000000000000 --- a/source/isaaclab/test/assets/contract/test_asset_contract_writes.py +++ /dev/null @@ -1,39 +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 - -"""Shared asset write contract: writers, read-backs, partial selections, cache invalidation, and ordered writes.""" - -import pytest - -from ._articulation_contract_cases import ( # noqa: F401 - TestArticulationPartialWriteRows, - TestArticulationWritersBody, - TestArticulationWritersJoint, - TestArticulationWritersRoot, - TestArticulationWritersTendon, - test_deprecated_joint_friction_writers, -) -from ._articulation_ordering_contract_cases import ( # noqa: F401 - TestArticulationOrderingBodyWriters, - TestArticulationOrderingComWrites, - TestArticulationOrderingJointWriters, - TestArticulationOrderingOperations, - TestArticulationOrderingRootWriteParity, - TestArticulationOrderingWriteParity, -) -from ._rigid_object_collection_contract_cases import ( # noqa: F401 - TestCollectionCacheInvalidation, - TestCollectionPartialWriteCells, - TestCollectionWritersBody, - TestCollectionWritersPose, -) -from ._rigid_object_contract_cases import ( # noqa: F401 - TestRigidObjectCacheInvalidation, - TestRigidObjectPartialWriteRows, - TestRigidObjectWritersBody, - TestRigidObjectWritersRoot, -) - -pytestmark = pytest.mark.integration diff --git a/source/isaaclab/test/assets/contract/test_rigid_object_collection_contract.py b/source/isaaclab/test/assets/contract/test_rigid_object_collection_contract.py new file mode 100644 index 000000000000..762ce672fb87 --- /dev/null +++ b/source/isaaclab/test/assets/contract/test_rigid_object_collection_contract.py @@ -0,0 +1,867 @@ +# 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-collection contracts across the production backends. + +Checks that every collection backend provides the data and writer behavior the base collection class advertises. +The backends run on mocked views, so these cases need neither Isaac Sim nor a 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_factory import get_rigid_object_collection + +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. + + +# --------------------------------------------------------------------------- +# 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 +# --------------------------------------------------------------------------- + + +def test_resolve_env_ids_handles_tensor_view_shape(monkeypatch, backend, device): + obj, _ = get_rigid_object_collection(backend, num_instances=4, device=device, monkeypatch=monkeypatch) + + 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() + + +def test_resolve_body_ids_handles_tensor_view_shape(monkeypatch, backend): + obj, _ = get_rigid_object_collection(backend, num_bodies=4, device="cpu", monkeypatch=monkeypatch) + + 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 +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_reshape_data_to_view_3d_accepts_torch_tensor(monkeypatch, 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, monkeypatch=monkeypatch + ) + 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) + + +@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(monkeypatch, 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", monkeypatch=monkeypatch + ) + 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 +# --------------------------------------------------------------------------- + + +def test_collection_counts_and_names(monkeypatch, 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", monkeypatch=monkeypatch) + + 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) + + +def test_find_bodies_returns_legacy_tensor_or_cached_proxy(monkeypatch, backend): + collection, _ = get_rigid_object_collection( + backend, num_instances=2, num_bodies=3, device="cpu", monkeypatch=monkeypatch + ) + + 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] + + +def test_find_objects_forwards_return_mode_with_alias_warning(monkeypatch, backend): + collection, _ = get_rigid_object_collection( + backend, num_instances=2, num_bodies=3, device="cpu", monkeypatch=monkeypatch + ) + + 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), +] + + +def test_collection_data_property_contract(monkeypatch, backend, device): + obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device, monkeypatch=monkeypatch) + 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", +} + + +def test_pose_write_invalidates_pose_dependent_caches(monkeypatch, backend): + obj, _ = get_rigid_object_collection(backend, num_instances=2, num_bodies=3, device="cpu", monkeypatch=monkeypatch) + 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) + + +def test_velocity_write_invalidates_body_frame_caches(monkeypatch, backend): + obj, _ = get_rigid_object_collection(backend, num_instances=2, num_bodies=3, device="cpu", monkeypatch=monkeypatch) + 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) + + +@pytest.mark.parametrize("setter_kind", ["index", "mask"]) +def test_set_coms_invalidates_same_timestamp_dependents(monkeypatch, backend, setter_kind): + obj, _ = get_rigid_object_collection(backend, num_instances=2, num_bodies=3, device="cpu", monkeypatch=monkeypatch) + 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) + + +# -- index variants for pose -- + + +@pytest.mark.parametrize("method_suffix", _BODY_POSE_METHODS) +def test_write_body_pose_to_sim_index(monkeypatch, backend, device, method_suffix): + obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device, monkeypatch=monkeypatch) + 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 -- + + +@pytest.mark.parametrize("method_suffix", _BODY_VEL_METHODS) +def test_write_body_velocity_to_sim_index(monkeypatch, backend, device, method_suffix): + obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device, monkeypatch=monkeypatch) + 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 -- + + +@pytest.mark.parametrize("method_suffix", _BODY_POSE_METHODS) +def test_write_body_pose_to_sim_mask(monkeypatch, backend, device, method_suffix): + obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device, monkeypatch=monkeypatch) + 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 -- + + +@pytest.mark.parametrize("method_suffix", _BODY_VEL_METHODS) +def test_write_body_velocity_to_sim_mask(monkeypatch, backend, device, method_suffix): + obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device, monkeypatch=monkeypatch) + 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) + + +@pytest.mark.parametrize("method_base, kwarg", _BODY_METHODS, ids=[m[0] for m in _BODY_METHODS]) +def test_body_writer_index(monkeypatch, backend, device, method_base, kwarg): + obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device, monkeypatch=monkeypatch) + 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)}) + + +@pytest.mark.parametrize("method_base, kwarg", _BODY_METHODS, ids=[m[0] for m in _BODY_METHODS]) +def test_body_writer_mask(monkeypatch, backend, device, method_base, kwarg): + obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device, monkeypatch=monkeypatch) + 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"), +] + + +def test_aliases_match_canonical_values(monkeypatch, 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", monkeypatch=monkeypatch) + 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 + + +# --------------------------------------------------------------------------- +# Tests: partial writes reach only the selected backend cells +# --------------------------------------------------------------------------- + +# quantity -> (writer base name, writer keyword, trailing size) +_PARTIAL_WRITES = { + "body_pose": ("write_body_link_pose_to_sim", "body_poses", 7), + "body_velocity": ("write_body_com_velocity_to_sim", "body_velocities", 6), + "mass": ("set_masses", "masses", 1), +} + + +def _read_backend_cells(backend: str, raw_backend, quantity: str, num_instances: int, num_bodies: int) -> torch.Tensor: + """Read one collection quantity from backend storage in public ``(env, body, ...)`` order.""" + if backend == "physx": + getter = { + "body_pose": raw_backend.get_transforms, + "body_velocity": raw_backend.get_velocities, + "mass": raw_backend.get_masses, + } + values = wp.to_torch(getter[quantity]()) + # 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_backend.get_attribute("body_mass", None)) + else: + getter = {"body_pose": raw_backend.get_root_transforms, "body_velocity": raw_backend.get_root_velocities} + values = wp.to_torch(getter[quantity](None)) + else: + from isaaclab_ov import tensor_types as TT + + binding = {"body_pose": TT.LINK_POSE, "body_velocity": TT.LINK_VELOCITY, "mass": TT.BODY_MASS}[quantity] + values = torch.as_tensor(raw_backend.bindings[binding]._data) + return values.reshape(num_instances, num_bodies, -1).cpu().clone() + + +@pytest.mark.parametrize("selection", ["index", "mask"]) +@pytest.mark.parametrize("quantity", _PARTIAL_WRITES) +def test_partial_write_preserves_unselected_backend_cells(monkeypatch, request, backend, selection, quantity): + if backend == "ovphysx" 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", + ) + ) + num_instances, num_bodies = _NUM_INSTANCES, _NUM_BODIES + obj, raw_backend = get_rigid_object_collection(backend, num_instances, num_bodies, "cpu", monkeypatch=monkeypatch) + writer, kwarg, trailing = _PARTIAL_WRITES[quantity] + # Literal, per-cell distinct payload; poses keep an identity rotation. + values = 100.0 * torch.arange(1, num_instances * num_bodies + 1, dtype=torch.float32).reshape( + num_instances, num_bodies, 1 + ) + torch.arange(trailing, dtype=torch.float32) + if quantity == "body_pose": + values[..., 3:6] = 0.0 + values[..., 6] = 1.0 + if quantity == "mass": + values = values[..., 0] + before = _read_backend_cells(backend, raw_backend, quantity, num_instances, num_bodies) + # Select the second environment and the last body so identity or transposed routing cannot pass. + env, body = 1, 2 + if selection == "index": + getattr(obj, f"{writer}_index")( + **{kwarg: values[env : env + 1, body : body + 1]}, + env_ids=torch.tensor([env], dtype=torch.int32), + body_ids=_make_body_ids("cpu", [body]), + ) + else: + getattr(obj, f"{writer}_mask")( + **{kwarg: values}, + env_mask=_make_item_mask(num_instances, [env], "cpu"), + body_mask=_make_item_mask(num_bodies, [body], "cpu"), + ) + + after = _read_backend_cells(backend, raw_backend, quantity, num_instances, num_bodies) + expected = before.clone() + expected[env, body] = values[env, body].reshape(-1) + torch.testing.assert_close(after, expected, rtol=0.0, atol=0.0) 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..ad1f5a7d8562 --- /dev/null +++ b/source/isaaclab/test/assets/contract/test_rigid_object_contract.py @@ -0,0 +1,790 @@ +# 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 contracts across the production backends. + +Checks that every rigid-object backend provides the data and writer behavior the base rigid-object class advertises. +The backends run on mocked views, so these cases need neither Isaac Sim nor a 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_factory import get_rigid_object + +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. +_NUM_INSTANCES = 2 + + +# --------------------------------------------------------------------------- +# Tests: Index resolution helpers +# --------------------------------------------------------------------------- + + +def test_resolve_env_ids_handles_tensor_view_shape(monkeypatch, backend, device): + obj, _ = get_rigid_object(backend, num_instances=4, device=device, monkeypatch=monkeypatch) + + 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 +# --------------------------------------------------------------------------- + + +def test_rigid_object_counts_and_names(monkeypatch, backend): + from isaaclab.assets.rigid_object.base_rigid_object_data import BaseRigidObjectData + + obj, _ = get_rigid_object(backend, _NUM_INSTANCES, device="cpu", monkeypatch=monkeypatch) + + 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) + + +def test_find_bodies_returns_legacy_list_or_cached_proxy(monkeypatch, backend): + obj, _ = get_rigid_object(backend, num_instances=2, device="cpu", monkeypatch=monkeypatch) + + 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), +] + + +def test_rigid_object_data_property_contract(monkeypatch, backend, device): + obj, _ = get_rigid_object(backend, _NUM_INSTANCES, device, monkeypatch=monkeypatch) + 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"), +] + + +def test_aliases_match_canonical_values(monkeypatch, backend): + # Random mock state makes link and COM quantities differ, so a retargeted alias fails. + obj, _ = get_rigid_object(backend, _NUM_INSTANCES, device="cpu", monkeypatch=monkeypatch) + 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", +} + + +def test_pose_write_invalidates_pose_dependent_caches(monkeypatch, backend): + obj, _ = get_rigid_object(backend, num_instances=2, device="cpu", monkeypatch=monkeypatch) + 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) + + +def test_velocity_write_invalidates_body_frame_caches(monkeypatch, backend): + obj, _ = get_rigid_object(backend, num_instances=2, device="cpu", monkeypatch=monkeypatch) + 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 + + +@pytest.mark.parametrize("setter_kind", ["index", "mask"]) +def test_set_coms_invalidates_same_timestamp_dependents(monkeypatch, backend, setter_kind): + obj, _ = get_rigid_object(backend, num_instances=2, device="cpu", monkeypatch=monkeypatch) + 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) + + +# -- index variants -- + + +@pytest.mark.parametrize("method_suffix", _ROOT_POSE_METHODS) +def test_write_root_pose_to_sim_index(monkeypatch, backend, device, method_suffix): + num_instances = _NUM_INSTANCES + obj, _ = get_rigid_object(backend, num_instances, device, monkeypatch=monkeypatch) + 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)) + + +@pytest.mark.parametrize("method_suffix", _ROOT_VEL_METHODS) +def test_write_root_velocity_to_sim_index(monkeypatch, backend, device, method_suffix): + num_instances = _NUM_INSTANCES + obj, _ = get_rigid_object(backend, num_instances, device, monkeypatch=monkeypatch) + 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 -- + + +@pytest.mark.parametrize("method_suffix", _ROOT_POSE_METHODS) +def test_write_root_pose_to_sim_mask(monkeypatch, backend, device, method_suffix): + num_instances = _NUM_INSTANCES + obj, _ = get_rigid_object(backend, num_instances, device, monkeypatch=monkeypatch) + 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)) + + +@pytest.mark.parametrize("method_suffix", _ROOT_VEL_METHODS) +def test_write_root_velocity_to_sim_mask(monkeypatch, backend, device, method_suffix): + num_instances = _NUM_INSTANCES + obj, _ = get_rigid_object(backend, num_instances, device, monkeypatch=monkeypatch) + 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]: + """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) + + +def test_external_wrench_frames(monkeypatch, 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, monkeypatch=monkeypatch) + # 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) + + +@pytest.mark.parametrize("method_base, kwarg", _BODY_METHODS, ids=[m[0] for m in _BODY_METHODS]) +def test_body_writer_index(monkeypatch, backend, device, method_base, kwarg): + num_instances, num_bodies = _NUM_INSTANCES, 1 + obj, _ = get_rigid_object(backend, num_instances, device, monkeypatch=monkeypatch) + 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)}) + 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)}) + + +@pytest.mark.parametrize("method_base, kwarg", _BODY_METHODS, ids=[m[0] for m in _BODY_METHODS]) +def test_body_writer_mask(monkeypatch, backend, device, method_base, kwarg): + num_instances, num_bodies = _NUM_INSTANCES, 1 + obj, _ = get_rigid_object(backend, num_instances, device, monkeypatch=monkeypatch) + 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: partial writes reach only the selected backend rows +# --------------------------------------------------------------------------- + +# quantity -> (writer base name, writer keyword, literal per-environment payload) +_PARTIAL_WRITES = { + "root_pose": ( + "write_root_link_pose_to_sim", + "root_pose", + [[31.0, 32.0, 33.0, 0.0, 0.0, 0.0, 1.0], [41.0, 42.0, 43.0, 0.0, 0.0, 0.0, 1.0]], + ), + "root_velocity": ( + "write_root_com_velocity_to_sim", + "root_velocity", + [[21.0, 22.0, 23.0, 24.0, 25.0, 26.0], [71.0, 72.0, 73.0, 74.0, 75.0, 76.0]], + ), + "mass": ("set_masses", "masses", [[51.0], [61.0]]), +} + + +def _read_backend_rows(backend: str, raw_backend, quantity: str) -> torch.Tensor: + """Read one rigid-object quantity from backend storage as one row per environment.""" + if backend == "physx": + getter = { + "root_pose": raw_backend.get_transforms, + "root_velocity": raw_backend.get_velocities, + "mass": raw_backend.get_masses, + } + values = getter[quantity]() + elif backend == "newton": + if quantity == "mass": + values = raw_backend.get_attribute("body_mass", None) + else: + getter = {"root_pose": raw_backend.get_root_transforms, "root_velocity": raw_backend.get_root_velocities} + values = getter[quantity](None) + else: + from isaaclab_ov import tensor_types as TT + + binding = {"root_pose": TT.RIGID_BODY_POSE, "root_velocity": TT.RIGID_BODY_VELOCITY, "mass": TT.RIGID_BODY_MASS} + values = raw_backend.bindings[binding[quantity]]._data + values = wp.to_torch(values) if isinstance(values, wp.array) else torch.as_tensor(values) + return values.reshape(2, -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): + obj, raw_backend = get_rigid_object(backend, num_instances=2, device="cpu", monkeypatch=monkeypatch) + writer, kwarg, payload = _PARTIAL_WRITES[quantity] + values = torch.tensor(payload, dtype=torch.float32) + before = _read_backend_rows(backend, raw_backend, quantity) + # Select the second environment so a writer that ignores the selection cannot pass by writing row 0. + selected, unselected = 1, 0 + if selection == "index": + kwargs = {kwarg: values[selected : selected + 1], "env_ids": torch.tensor([selected], dtype=torch.int32)} + if quantity == "mass": + kwargs["body_ids"] = [0] + else: + kwargs = {kwarg: values, "env_mask": _make_item_mask(2, [selected], "cpu")} + if quantity == "mass": + kwargs["body_mask"] = _make_item_mask(1, [0], "cpu") + + getattr(obj, f"{writer}_{selection}")(**kwargs) + + after = _read_backend_rows(backend, raw_backend, quantity) + torch.testing.assert_close(after[selected], values[selected], rtol=0.0, atol=0.0) + torch.testing.assert_close(after[unselected], before[unselected], rtol=0.0, atol=0.0) diff --git a/source/isaaclab_newton/changelog.d/asset-tests-readability.skip b/source/isaaclab_newton/changelog.d/asset-tests-readability.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_newton/test/articulation_test_utils.py b/source/isaaclab_newton/test/articulation_test_utils.py index bc6729aa9eba..b95e46621f04 100644 --- a/source/isaaclab_newton/test/articulation_test_utils.py +++ b/source/isaaclab_newton/test/articulation_test_utils.py @@ -5,18 +5,14 @@ """Locally authored scenes shared by the kitless Newton asset and controller tests. -Test modules below ``source/isaaclab_newton/test`` import this module after putting this directory on -``sys.path``. The fixtures are small USD documents authored here and written to a per-process temporary -directory, so the scenes need neither Kit nor Nucleus and the repository carries no generated USD files. +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 -import atexit -import functools -import shutil -import tempfile -from collections.abc import Callable, Iterator, Mapping, Sequence +from collections.abc import Iterator, Mapping from contextlib import contextmanager from pathlib import Path from typing import Any @@ -52,19 +48,13 @@ def local_usd(name: str, **kwargs: Any) -> sim_utils.UsdFileCfg: """Return a spawner for one of the local USD fixtures. Args: - name: Authored fixture name, or the name of a checked-in file in ``test/assets/data``. + name: File name in ``test/assets/data``. **kwargs: Additional :class:`~isaaclab.sim.UsdFileCfg` fields. Returns: The spawner configuration. """ - if name in _FIXTURES: - path = _fixture_dir() / name - if not path.exists(): - path.write_text(_FIXTURES[name](), encoding="utf-8") - else: - path = _DATA_DIR / name - return sim_utils.UsdFileCfg(usd_path=str(path), **kwargs) + return sim_utils.UsdFileCfg(usd_path=str(_DATA_DIR / name), **kwargs) def newton_sim_cfg( @@ -165,370 +155,3 @@ def world_gravity(gravity: Vec3) -> Iterator[None]: finally: world_gravity_view.copy_(previous) SimulationManager.add_model_change(ModelFlags.MODEL_PROPERTIES) - - -## -# Fixture authoring. -## - - -def _body( - name: str, - *, - mass: float, - inertia: Vec3, - com: Vec3 | None = None, - translate: Vec3 | None = None, - articulation_root: bool = False, - collider_size: float | None = None, -) -> str: - """Return a rigid body with explicit mass properties and, optionally, a child collider cube [m].""" - applied = '"PhysicsRigidBodyAPI", "PhysicsMassAPI"' - if articulation_root: - applied += ', "PhysicsArticulationRootAPI"' - lines = [] - if translate is not None: - lines += [f"double3 xformOp:translate = {translate}", 'uniform token[] xformOpOrder = ["xformOp:translate"]'] - lines += [f"float physics:mass = {mass}", f"float3 physics:diagonalInertia = {inertia}"] - if com is not None: - lines.append(f"point3f physics:centerOfMass = {com}") - content = "\n".join(f" {line}" for line in lines) - if collider_size is not None: - content += ( - f'\n\n def Cube "Collision" (\n prepend apiSchemas = ["PhysicsCollisionAPI"]\n )\n' - f" {{\n double size = {collider_size}\n }}" - ) - return f' def Xform "{name}" (\n prepend apiSchemas = [{applied}]\n )\n {{\n{content}\n }}\n' - - -def _joint( - root: str, - kind: str, - name: str, - body0: str | None, - body1: str, - *, - axis: str | None = None, - local_pos0: Vec3 | None = None, - local_pos1: Vec3 | None = None, - limits: tuple[float, float] | None = None, - drive: tuple[str, float, float] | None = None, -) -> str: - """Return a joint. - - ``body0=None`` attaches ``body1`` to the world, and ``drive`` is ``(kind, stiffness, damping)``. - """ - schemas = f' (\n prepend apiSchemas = ["PhysicsDriveAPI:{drive[0]}"]\n )' if drive else "" - lines = [] - if axis is not None: - lines.append(f'uniform token physics:axis = "{axis}"') - if body0 is not None: - lines.append(f"rel physics:body0 = ") - lines.append(f"rel physics:body1 = ") - if local_pos0 is not None: - lines.append(f"point3f physics:localPos0 = {local_pos0}") - if local_pos1 is not None: - lines.append(f"point3f physics:localPos1 = {local_pos1}") - if limits is not None: - lines += [f"float physics:lowerLimit = {limits[0]}", f"float physics:upperLimit = {limits[1]}"] - if drive is not None: - lines += [ - f"float drive:{drive[0]}:physics:stiffness = {drive[1]}", - f"float drive:{drive[0]}:physics:damping = {drive[2]}", - ] - content = "\n".join(f" {line}" for line in lines) - return f' def Physics{kind}Joint "{name}"{schemas}\n {{\n{content}\n }}\n' - - -def _document(root: str, prims: Sequence[str], *, articulation_root: bool = True) -> str: - """Return a USD document whose default prim is an ``Xform`` holding ``prims``. - - ``articulation_root=False`` leaves the articulation root to one of the bodies. - """ - header = f'#usda 1.0\n(\n defaultPrim = "{root}"\n metersPerUnit = 1\n upAxis = "Z"\n)\n\n' - metadata = ' (\n prepend apiSchemas = ["PhysicsArticulationRootAPI"]\n)' if articulation_root else "" - return header + f'def Xform "{root}"{metadata}\n{{\n' + "\n".join(prims) + "}\n" - - -def _floating_two_link() -> str: - """Two colliding links joined about z; the articulation root sits on the ``Root`` body.""" - return _document( - "Robot", - [ - _body( - "Root", - mass=1.0, - inertia=(0.01, 0.02, 0.02), - com=(0.0, 0.0, 0.0), - articulation_root=True, - collider_size=0.1, - ), - _body( - "Child", - mass=1.0, - inertia=(0.01, 0.02, 0.02), - com=(0.0, 0.0, 0.0), - translate=(0.5, 0.0, 0.0), - collider_size=0.1, - ), - _joint( - "Robot", - "Revolute", - "Joint", - "Root", - "Child", - axis="Z", - local_pos0=(0.25, 0.0, 0.0), - local_pos1=(-0.25, 0.0, 0.0), - limits=(-90.0, 90.0), - ), - ], - articulation_root=False, - ) - - -def _fixed_spatial_chain() -> str: - """A fixed-base chain of three slides and a three-axis wrist whose links carry center-of-mass offsets. - - Every joint has a USD drive; the wrist's angular gains are per degree. - """ - links = [ - _body("Root", mass=1.0, inertia=(0.01, 0.01, 0.01)), - _joint("Robot", "Fixed", "RootJoint", None, "Root"), - ] - link_specs = [ - (0.5, (0.02, 0.02, 0.02), (0.02, 0.0, 0.0), None), - (0.5, (0.02, 0.02, 0.02), (0.0, 0.02, 0.0), None), - (0.5, (0.02, 0.02, 0.02), (0.0, 0.0, 0.02), None), - (0.3, (0.01, 0.01, 0.01), (0.03, 0.0, 0.02), (0.0, 0.0, 0.1)), - (0.3, (0.01, 0.01, 0.01), (0.0, 0.03, 0.02), (0.0, 0.0, 0.2)), - (0.2, (0.01, 0.01, 0.01), (0.05, 0.02, 0.0), (0.0, 0.0, 0.3)), - ] - for index, (mass, inertia, com, translate) in enumerate(link_specs): - links.append(_body(f"Link_{index}", mass=mass, inertia=inertia, com=com, translate=translate)) - parents = ["Root", "Link_0", "Link_1", "Link_2", "Link_3", "Link_4"] - for index, axis in enumerate("XYZ"): - links.append( - _joint( - "Robot", - "Prismatic", - f"Joint_{index}", - parents[index], - f"Link_{index}", - axis=axis, - limits=(-0.2, 0.2), - drive=("linear", 200.0, 20.0), - ) - ) - for index, (axis, stiffness) in enumerate(zip("XYZ", WRIST_USD_STIFFNESS), start=3): - links.append( - _joint( - "Robot", - "Revolute", - f"Joint_{index}", - parents[index], - f"Link_{index}", - axis=axis, - local_pos0=(0.0, 0.0, 0.1), - local_pos1=(0.0, 0.0, 0.0), - limits=(-60.0, 60.0), - drive=("angular", stiffness, round(0.1 * stiffness, 6)), - ) - ) - return _document("Robot", links) - - -def _revolute_pendulum() -> str: - """A fixed pivot and a unit arm turning about z.""" - return _document( - "Robot", - [ - _body("CenterPivot", mass=1.0, inertia=(0.01, 0.01, 0.01)), - _joint("Robot", "Fixed", "FixedJoint", None, "CenterPivot"), - _body("Arm", mass=1.0, inertia=(0.01, 0.01, 0.01), translate=(1.0, 0.0, 0.0)), - _joint( - "Robot", - "Revolute", - "RevoluteJoint", - "CenterPivot", - "Arm", - axis="Z", - local_pos0=(0.0, 0.0, 0.0), - local_pos1=(-1.0, 0.0, 0.0), - ), - ], - ) - - -def _fixed_tendon_hand() -> str: - """A fixed palm with two two-joint fingers; each fixed tendon ``XJ0`` sums ``XJ1`` and ``XJ2``. - - A MuJoCo position actuator drives each tendon. - """ - prims = [ - _body("Palm", mass=0.5, inertia=(0.001, 0.001, 0.001)), - _joint("Hand", "Fixed", "PalmJoint", None, "Palm"), - ] - for finger, y in (("FF", 0.02), ("MF", -0.02)): - prims += [ - _body( - f"{finger}_proximal", - mass=0.05, - inertia=(1e-5, 1e-5, 1e-5), - com=(0.025, 0.0, 0.0), - translate=(0.05, y, 0.0), - ), - _body( - f"{finger}_distal", mass=0.03, inertia=(1e-5, 1e-5, 1e-5), com=(0.02, 0.0, 0.0), translate=(0.1, y, 0.0) - ), - _joint( - "Hand", - "Revolute", - f"{finger}J2", - "Palm", - f"{finger}_proximal", - axis="Y", - local_pos0=(0.05, y, 0.0), - local_pos1=(0.0, 0.0, 0.0), - limits=(0.0, 90.0), - ), - _joint( - "Hand", - "Revolute", - f"{finger}J1", - f"{finger}_proximal", - f"{finger}_distal", - axis="Y", - local_pos0=(0.05, 0.0, 0.0), - local_pos1=(0.0, 0.0, 0.0), - limits=(0.0, 90.0), - ), - f' def MjcTendon "{finger}J0"\n {{\n uniform token mjc:type = "fixed"\n' - f" rel mjc:path = [, ]\n" - " uniform int[] mjc:path:indices = [0, 1]\n uniform double[] mjc:path:coef = [1, 1]\n }\n", - f' def MjcActuator "{finger}_tendon_actuator"\n {{\n uniform token mjc:biasType = "affine"\n' - " uniform double[] mjc:gainPrm = [0.05, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n" - " uniform double[] mjc:biasPrm = [0, -0.05, -0.005, 0, 0, 0, 0, 0, 0, 0]\n" - f" rel mjc:target = \n }}\n", - ] - return _document("Hand", prims) - - -def _floating_two_leg() -> str: - """A floating base with two three-joint legs named like ANYmal's (``HAA``, ``HFE``, ``KFE``).""" - prims = [_body("base", mass=10.0, inertia=(0.5, 0.5, 0.5), com=(0.0, 0.0, 0.0))] - joints = [] - for leg, x, y in (("LF", 0.3, 0.1), ("RH", -0.3, -0.1)): - prims += [ - _body(f"{leg}_HIP", mass=1.0, inertia=(0.05, 0.05, 0.05), com=(0.0, 0.02, 0.0), translate=(x, y, 0.0)), - _body(f"{leg}_THIGH", mass=1.0, inertia=(0.05, 0.05, 0.05), com=(0.0, 0.0, -0.1), translate=(x, y, -0.05)), - _body(f"{leg}_SHANK", mass=0.5, inertia=(0.05, 0.05, 0.05), com=(0.0, 0.0, -0.1), translate=(x, y, -0.3)), - ] - joints += [ - _joint( - "Robot", - "Revolute", - f"{leg}_HAA", - "base", - f"{leg}_HIP", - axis="X", - local_pos0=(x, y, 0.0), - local_pos1=(0.0, 0.0, 0.0), - limits=(-45.0, 45.0), - ), - _joint( - "Robot", - "Revolute", - f"{leg}_HFE", - f"{leg}_HIP", - f"{leg}_THIGH", - axis="Y", - local_pos0=(0.0, 0.0, -0.05), - local_pos1=(0.0, 0.0, 0.0), - limits=(-90.0, 90.0), - ), - _joint( - "Robot", - "Revolute", - f"{leg}_KFE", - f"{leg}_THIGH", - f"{leg}_SHANK", - axis="Y", - local_pos0=(0.0, 0.0, -0.25), - local_pos1=(0.0, 0.0, 0.0), - limits=(-160.0, 0.0), - ), - ] - return _document("Robot", prims + joints) - - -def _fixed_cartpole() -> str: - """A cart sliding on a fixed rail with a pole on a revolute joint.""" - return _document( - "Robot", - [ - _body("rail", mass=1.0, inertia=(0.01, 0.01, 0.01)), - _joint("Robot", "Fixed", "rail_joint", None, "rail"), - _body("cart", mass=1.0, inertia=(0.01, 0.01, 0.01)), - _body("pole", mass=0.2, inertia=(0.01, 0.01, 0.01), com=(0.0, 0.0, 0.5)), - _joint("Robot", "Prismatic", "slider_to_cart", "rail", "cart", axis="X", limits=(-2.0, 2.0)), - _joint("Robot", "Revolute", "cart_to_pole", "cart", "pole", axis="Y"), - ], - ) - - -def _rigid_body_with_articulation_root() -> str: - """A rigid body carrying an articulation root, with a second body fixed to it.""" - return """#usda 1.0 -( - defaultPrim = "Object" - metersPerUnit = 1 - upAxis = "Z" -) - -def Cube "Object" ( - prepend apiSchemas = ["PhysicsRigidBodyAPI", "PhysicsMassAPI", "PhysicsCollisionAPI", "PhysicsArticulationRootAPI"] -) -{ - double size = 0.1 - float physics:mass = 1 - - def Cube "Child" ( - prepend apiSchemas = ["PhysicsRigidBodyAPI", "PhysicsMassAPI", "PhysicsCollisionAPI"] - ) - { - double size = 0.1 - double3 xformOp:translate = (0, 0, 2) - uniform token[] xformOpOrder = ["xformOp:translate"] - float physics:mass = 1 - } - - def PhysicsFixedJoint "Joint" - { - rel physics:body0 = - rel physics:body1 = - point3f physics:localPos0 = (0, 0, 2) - point3f physics:localPos1 = (0, 0, 0) - } -} -""" - - -_FIXTURES: dict[str, Callable[[], str]] = { - "fixed_cartpole.usda": _fixed_cartpole, - "fixed_spatial_chain.usda": _fixed_spatial_chain, - "fixed_tendon_hand.usda": _fixed_tendon_hand, - "floating_two_leg.usda": _floating_two_leg, - "floating_two_link.usda": _floating_two_link, - "revolute_pendulum.usda": _revolute_pendulum, - "rigid_body_with_articulation_root.usda": _rigid_body_with_articulation_root, -} -"""Authored fixtures by file name.""" - - -@functools.cache -def _fixture_dir() -> Path: - """Return this process's directory for the authored fixtures, removed when the process exits.""" - directory = Path(tempfile.mkdtemp(prefix="isaaclab_newton_test_fixtures_")) - atexit.register(shutil.rmtree, directory, ignore_errors=True) - return directory 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/test_newton_actuators_newton.py b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py index b1ec09f4e457..9820fc174911 100644 --- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py +++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py @@ -26,7 +26,6 @@ launch_test_simulation(SimulationCfg(physics=NewtonCfg())) -import functools import os import sys from collections.abc import Iterator @@ -651,10 +650,8 @@ def test_randomize_actuator_gains_reaches_newton_controllers(newton_run: _Run) - cartpole = newton_run.articulations["cartpole"] assert SimulationManager._adapter is not None - 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() + anymal_stiffness_before = read_group_parameter(anymal.actuators, "legs", "controller", "kp").clone() + anymal_damping_before = read_group_parameter(anymal.actuators, "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 (6 free-root DOFs + the leg joints): a wrong stride corrupts every @@ -662,8 +659,8 @@ def test_randomize_actuator_gains_reaches_newton_controllers(newton_run: _Run) - 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() + cartpole_stiffness_before = read_group_parameter(cartpole.actuators, "all_joints", "controller", "kp").clone() + cartpole_damping_before = read_group_parameter(cartpole.actuators, "all_joints", "controller", "kd").clone() env = MockEnv({"anymal": anymal, "cartpole": cartpole}, NUM_ENVS, anymal.device) term, asset_cfg = build_dr_term(env, "cartpole") @@ -681,23 +678,31 @@ def test_randomize_actuator_gains_reaches_newton_controllers(newton_run: _Run) - 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) + read_group_parameter(cartpole.actuators, "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) + read_group_parameter(cartpole.actuators, "all_joints", "controller", "kd")[0], + torch.full((n_cp,), 5.0, device=anymal.device), ) # The leg 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) + torch.testing.assert_close( + read_group_parameter(anymal.actuators, "legs", "controller", "kp"), anymal_stiffness_before + ) + torch.testing.assert_close( + read_group_parameter(anymal.actuators, "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] + read_group_parameter(cartpole.actuators, "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] + read_group_parameter(cartpole.actuators, "all_joints", "controller", "kd")[env_idx], + cartpole_damping_before[env_idx], ) # DR scoped to the floating-base leg updates only its selected env. @@ -712,12 +717,21 @@ def test_randomize_actuator_gains_reaches_newton_controllers(newton_run: _Run) - distribution="uniform", ) torch.testing.assert_close( - anymal_read("legs", "controller", "kp")[0], torch.full((n,), 100.0, device=anymal.device) + read_group_parameter(anymal.actuators, "legs", "controller", "kp")[0], + torch.full((n,), 100.0, device=anymal.device), + ) + torch.testing.assert_close( + read_group_parameter(anymal.actuators, "legs", "controller", "kd")[0], + torch.full((n,), 5.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]) + torch.testing.assert_close( + read_group_parameter(anymal.actuators, "legs", "controller", "kp")[env_idx], + anymal_stiffness_before[env_idx], + ) + torch.testing.assert_close( + read_group_parameter(anymal.actuators, "legs", "controller", "kd")[env_idx], anymal_damping_before[env_idx] + ) # --------------------------------------------------------------------------- diff --git a/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py b/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py index c397712f1257..d6d25c43f9f3 100644 --- a/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py +++ b/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py @@ -112,15 +112,6 @@ def _compute_jacobian_root_frame(robot: Articulation, ee_jacobi_idx: int, arm_jo return jacobian -def _build_relative_pose_target( - robot: Articulation, ee_frame_idx: int, delta_xyz: tuple[float, float, float] -) -> torch.Tensor: - """Return the current end-effector pose in the root frame offset by ``delta_xyz`` [m], keeping its orientation.""" - initial_ee_pos_b, initial_ee_quat_b = _compute_ee_pose_root(robot, ee_frame_idx) - target_pos_b = initial_ee_pos_b + initial_ee_pos_b.new_tensor(delta_xyz) - return torch.cat([target_pos_b, initial_ee_quat_b], dim=-1) - - def _make_osc(device: str) -> OperationalSpaceController: """Return a fixed-impedance absolute-pose controller with inertial decoupling and no gravity compensation.""" return OperationalSpaceController( @@ -184,7 +175,8 @@ def test_osc_tracking_accuracy(chain: _Chain) -> None: ``inertial_dynamics_decoupling=True`` exercises ``mass_matrix`` and the COM-referenced J → M_b → J product. """ chain.rest() - target_pose_b = _build_relative_pose_target(chain.robot, chain.ee_frame_idx, (0.05, 0.0, 0.0)) + ee_pos_b, ee_quat_b = _compute_ee_pose_root(chain.robot, chain.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(chain, _make_osc(chain.device), target_pose_b, 300, gravity=False) pos_mean = sum(pos_history[-200:]) / 200 @@ -217,7 +209,8 @@ def test_osc_gravity_compensation_precision(chain: _Chain) -> None: chain.rest() osc = _make_osc(chain.device) # Hold the initial EE pose: phase-1 steady-state error is pure gravity sag. - target_pose_b = _build_relative_pose_target(chain.robot, chain.ee_frame_idx, (0.0, 0.0, 0.0)) + ee_pos_b, ee_quat_b = _compute_ee_pose_root(chain.robot, chain.ee_frame_idx) + target_pose_b = torch.cat([ee_pos_b, ee_quat_b], dim=-1) def _stationary_tail_mean(history: list[float], label: str) -> float: """Mean of the last 200 samples, asserting the two tail halves agree within 25%. diff --git a/source/isaaclab_ov/changelog.d/asset-tests-readability.skip b/source/isaaclab_ov/changelog.d/asset-tests-readability.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_ov/test/assets/test_articulation.py b/source/isaaclab_ov/test/assets/test_articulation.py index 1cfeb8162e74..f2cbf996e403 100644 --- a/source/isaaclab_ov/test/assets/test_articulation.py +++ b/source/isaaclab_ov/test/assets/test_articulation.py @@ -199,9 +199,6 @@ class _ArticulationScene: native_actuators: bool islands: dict[str, Articulation] - def __getitem__(self, name: str) -> Articulation: - return self.islands[name] - def step(self, articulation: Articulation, count: int = 1) -> None: """Step the shared scene while writing and refreshing one island.""" for _ in range(count): @@ -392,7 +389,7 @@ def test_setting_invalid_articulation_root_prim_path(device: str) -> None: @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["ordered"] + articulation = scene.islands["ordered"] device = scene.device assert articulation.is_initialized assert articulation.is_fixed_base @@ -443,7 +440,7 @@ def test_articulation_initialization_and_partial_state(scene: _ArticulationScene @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["ordered"] + 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) @@ -527,7 +524,7 @@ def test_set_material_properties(scene: _ArticulationScene) -> None: :class:`~isaaclab_ov.sim.views.OvPhysxView`. The binding is CPU-native, so the buffer lives in host memory. """ - articulation = scene["ordered"] + articulation = scene.islands["ordered"] num_articulations = articulation.num_instances device = scene.device sim = scene.sim @@ -571,7 +568,7 @@ def test_setting_velocity_and_effort_limits_write_to_solver(scene: _Articulation USD-authored defaults when unset) land in the native solver buffers and match ``data.joint_vel_limits`` and ``data.joint_effort_limits``. """ - articulation = scene[island] + articulation = scene.islands[island] device = scene.device actuator_cfg = _ISLANDS[island].actuators["joints"] joint_drive_props = _ISLANDS[island].joint_drive_props @@ -604,7 +601,7 @@ def test_effort_binding_excludes_implicit_pd(scene: _ArticulationScene) -> None: 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["mixed"] + articulation = scene.islands["mixed"] stiffness, effort_limit = 20.0, 400.0 joint_names = articulation.joint_names shoulder_ids = [joint_names.index(name) for name in ("left_shoulder", "right_shoulder")] @@ -648,7 +645,7 @@ def test_com_orientation_write_invalidates_static_inertia_cache_with_body_orderi 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["com_order"] + articulation = scene.islands["com_order"] device = scene.device assert articulation.body_ordering is not None @@ -715,7 +712,7 @@ def test_com_orientation_write_invalidates_static_inertia_cache_with_body_orderi @pytest.mark.parametrize("scene", _CUDA_DEVICES, 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["usd_gains"] + 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)) @@ -729,7 +726,7 @@ def test_loading_gains_from_usd(scene: _ArticulationScene) -> None: @pytest.mark.parametrize("scene", _CUDA_DEVICES, 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["non_root"] + articulation = scene.islands["non_root"] num_articulations = articulation.num_instances assert articulation.is_initialized assert not articulation.is_fixed_base @@ -753,7 +750,7 @@ def test_explicit_articulation_root_on_a_floating_base_body(scene: _Articulation @pytest.mark.parametrize("scene", _CUDA_DEVICES, indirect=True) def test_joint_position_limits_clamp_default_joint_positions(scene: _ArticulationScene) -> None: """Write partial joint position limits and clamp only the selected default joint positions into them.""" - articulation = scene["pos_limits"] + articulation = scene.islands["pos_limits"] device = scene.device default_before = articulation.data.default_joint_pos.torch.clone() env_ids = torch.tensor([1], dtype=torch.int32, device=device) @@ -774,7 +771,7 @@ def test_joint_position_limits_clamp_default_joint_positions(scene: _Articulatio @pytest.mark.parametrize("scene", _CUDA_DEVICES, 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["drive"] + articulation = scene.islands["drive"] device = scene.device ordering = articulation.joint_ordering assert ordering is not None @@ -840,7 +837,7 @@ def test_computed_dynamics_follow_public_order_and_model_writes(scene: _Articula 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 = scene["dynamics"] + articulation = scene.islands["dynamics"] device = scene.device num_envs = articulation.num_instances @@ -965,7 +962,7 @@ def assert_dynamics_match(jacobian: bool = True, gravity: bool = True) -> None: @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[island] + articulation = scene.islands[island] device = scene.device assert (articulation.joint_ordering is not None) == (island == "reversed_ordered") @@ -1004,7 +1001,7 @@ def test_reversed_joint_dynamics_use_public_joint_basis(scene: _ArticulationScen @pytest.mark.parametrize("scene", _CUDA_DEVICES, 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["floating"] + articulation = scene.islands["floating"] device = scene.device num_articulations = articulation.num_instances num_bodies = articulation.num_bodies @@ -1127,7 +1124,7 @@ def test_tendon_properties_and_position_targets(scene: _ArticulationScene) -> No every tendon of environment 0; the mask form commands tendon 0 of environment 1. Every other cell must keep its initial offset. """ - articulation = scene["tendon"] + articulation = scene.islands["tendon"] device = scene.device num_articulations = articulation.num_instances assert articulation.is_initialized @@ -1180,7 +1177,7 @@ def test_tendon_properties_and_position_targets(scene: _ArticulationScene) -> No @pytest.mark.parametrize("scene", _CUDA_DEVICES, 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["native"] + articulation = scene.islands["native"] device = scene.device assert articulation._actuator_control.native_actuator_path_active assert articulation.newton_actuator_adapter is not None @@ -1232,7 +1229,7 @@ def __getitem__(self, name): assert name == "robot" return self._asset - articulation = scene["delayed"] + articulation = scene.islands["delayed"] device = scene.device num_joints = articulation.num_joints scene.step(articulation, count=3) @@ -1273,7 +1270,7 @@ def test_root_link_vel_w_refreshes_fk_before_body_com_vel_w_read(scene: _Articul 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. """ - articulation = scene["ordered"] + articulation = scene.islands["ordered"] articulation.update(scene.sim.cfg.dt) # Prime the derived buffers before the write so their TimestampedBuffers are populated; otherwise @@ -1315,7 +1312,7 @@ def test_cpu_only_property_writes_wait_for_pinned_host_staging(scene: _Articulat 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 = scene["staging"] + articulation = scene.islands["staging"] device = scene.device num_envs, num_joints = articulation.num_instances, articulation.num_joints diff --git a/source/isaaclab_ov/test/assets/test_rigid_object.py b/source/isaaclab_ov/test/assets/test_rigid_object.py index 1973391af9c2..22e9e69a57a9 100644 --- a/source/isaaclab_ov/test/assets/test_rigid_object.py +++ b/source/isaaclab_ov/test/assets/test_rigid_object.py @@ -20,7 +20,6 @@ from __future__ import annotations from collections.abc import Iterator -from contextlib import AbstractContextManager from dataclasses import dataclass import pytest @@ -46,11 +45,6 @@ _NUM_CUBES = 2 -def _sim_context(device: str) -> AbstractContextManager[SimulationContext]: - """Build a local OVPhysX context from an in-memory USD stage.""" - return build_simulation_context(device=device, sim_cfg=SimulationCfg(physics=OvPhysxCfg(), device=device)) - - def _spawn_cubes(name: str, y_offset: float, **rigid_props) -> RigidObject: """Author one pair of local cuboids.""" for index in range(_NUM_CUBES): @@ -85,7 +79,7 @@ def _spawn_static_colliders() -> RigidObject: @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 _sim_context(device) as sim: + 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"): @@ -101,7 +95,7 @@ def test_initialization_with_no_rigid_body(device: str) -> None: ) 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 _sim_context(device) as sim: + 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() @@ -127,7 +121,7 @@ class _RigidScene: def scene(request: pytest.FixtureRequest) -> Iterator[_RigidScene]: """Initialize every cuboid pair for one device once for this module.""" device = request.param - with _sim_context(device) as sim: + 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) resettable = _spawn_cubes("Resettable", y_offset=4.0, disable_gravity=True) diff --git a/source/isaaclab_ov/test/assets/test_rigid_object_collection.py b/source/isaaclab_ov/test/assets/test_rigid_object_collection.py index b5a5c4429b5f..3ed2a3f16b24 100644 --- a/source/isaaclab_ov/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_ov/test/assets/test_rigid_object_collection.py @@ -20,7 +20,6 @@ from __future__ import annotations from collections.abc import Iterator -from contextlib import AbstractContextManager from dataclasses import dataclass import pytest @@ -43,13 +42,6 @@ _NUM_ENVS, _NUM_BODIES = 2, 3 -def _sim_context(device: str) -> AbstractContextManager[SimulationContext]: - """Build a local OVPhysX context from an in-memory USD stage.""" - return build_simulation_context( - device=device, sim_cfg=SimulationCfg(physics=OvPhysxCfg(), device=device, gravity=(0.0, 0.0, 0.0)) - ) - - 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: @@ -78,7 +70,9 @@ def _spawn_collection(name: str, y_offset: float, spawn: sim_utils.SpawnerCfg | @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 _sim_context(device) as sim: + 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", @@ -106,7 +100,9 @@ class _CollectionScene: def scene(request: pytest.FixtureRequest) -> Iterator[_CollectionScene]: """Initialize every collection for one device once for this module.""" device = request.param - with _sim_context(device) as sim: + with build_simulation_context( + device=device, sim_cfg=SimulationCfg(physics=OvPhysxCfg(), device=device, gravity=(0.0, 0.0, 0.0)) + ) as sim: fused = _spawn_collection("Fused", y_offset=0.0) resettable = _spawn_collection("Resettable", y_offset=4.0) stepped = _spawn_collection("Stepped", y_offset=8.0) diff --git a/source/isaaclab_physx/changelog.d/asset-tests-readability.skip b/source/isaaclab_physx/changelog.d/asset-tests-readability.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_physx/test/assets/test_articulation.py b/source/isaaclab_physx/test/assets/test_articulation.py index a586d2e2d919..e85e18505ac7 100644 --- a/source/isaaclab_physx/test/assets/test_articulation.py +++ b/source/isaaclab_physx/test/assets/test_articulation.py @@ -185,10 +185,14 @@ def _yaw_quat(angle: float) -> tuple[float, float, float, float]: @pytest.fixture -def sim(request) -> Iterator[SimulationContext]: +def gravity_enabled() -> bool: + """Enable gravity unless a test explicitly parametrizes it otherwise.""" + return True + + +@pytest.fixture +def sim(device: str, gravity_enabled: bool) -> Iterator[SimulationContext]: """Create a function-scoped simulation context for tests that own their scene.""" - device = request.getfixturevalue("device") - gravity_enabled = request.getfixturevalue("gravity_enabled") if "gravity_enabled" in request.fixturenames else True 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 @@ -254,15 +258,15 @@ def test_setting_invalid_articulation_root_prim_path(sim, device) -> None: sim.reset() -@pytest.mark.parametrize("num_articulations", [2]) @pytest.mark.parametrize("device", test_devices()) -def test_fixed_tendon_position_target_writes_offset(sim, num_articulations, device) -> None: +def test_fixed_tendon_position_target_writes_offset(sim, device) -> None: """A tendon length target lands in the simulation as ``rest_length - target`` on the selected cells only. 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. The Shadow Hand is the shipped asset whose native PhysX fixed tendons these writers command. """ + num_articulations = 2 articulation_cfg = SHADOW_HAND_PHYSX_CFG articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device) @@ -305,10 +309,9 @@ def test_fixed_tendon_position_target_writes_offset(sim, num_articulations, devi torch.testing.assert_close(articulation.data.fixed_tendon_offset.torch, expected) -@pytest.mark.parametrize("num_articulations", [1]) @pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) @pytest.mark.isaacsim_ci -def test_get_gravity_compensation_forces_static_equilibrium(sim, num_articulations, device) -> None: +def test_get_gravity_compensation_forces_static_equilibrium(sim, device) -> None: """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``. @@ -322,6 +325,7 @@ def test_get_gravity_compensation_forces_static_equilibrium(sim, num_articulatio Newton-side variant of the same name lives in ``isaaclab_newton/test/assets/test_articulation.py`` (backend parity). """ + num_articulations = 1 # 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 diff --git a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py index 97d10830f887..1a1022929003 100644 --- a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py @@ -98,10 +98,14 @@ def _yaw_quat(angle: float) -> tuple[float, float, float, float]: @pytest.fixture -def sim(request) -> Iterator[SimulationContext]: +def gravity_enabled() -> bool: + """Enable gravity unless a test explicitly parametrizes it otherwise.""" + return True + + +@pytest.fixture +def sim(device: str, gravity_enabled: bool) -> Iterator[SimulationContext]: """Create a function-scoped simulation context for tests that own their scene.""" - device = request.getfixturevalue("device") - gravity_enabled = request.getfixturevalue("gravity_enabled") if "gravity_enabled" in request.fixturenames else True 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 @@ -112,10 +116,10 @@ def sim(request) -> Iterator[SimulationContext]: ## -@pytest.mark.parametrize("num_cubes", [2]) @pytest.mark.parametrize("device", test_devices()) -def test_initialization_with_no_rigid_body(sim, num_cubes, device) -> None: +def test_initialization_with_no_rigid_body(sim, device) -> None: """Test that initialization fails when no rigid body is found at the provided prim path.""" + num_cubes = 2 object_collection, _ = generate_cubes_scene(num_cubes=num_cubes, has_api=False) # Check that the framework doesn't hold excessive strong references. @@ -187,12 +191,12 @@ def read_view(values: wp.array, data_dim: int) -> torch.Tensor: torch.testing.assert_close(object_collection.data.body_inertia.torch.cpu(), expected_inertias) -@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) -> None: +def test_gravity_vec_w(sim, device, gravity_enabled) -> None: """Test that gravity vector direction is set correctly for the rigid object.""" + num_envs = 3 + num_cubes = 2 object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes) gravity_dir = (0.0, 0.0, -1.0) if gravity_enabled else (0.0, 0.0, 0.0) From 406fa72eb1efc1c54b2a3dde419f9aba5ddbb280 Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Fri, 2 Oct 2026 15:20:41 -0700 Subject: [PATCH 22/33] Replace contract capability framework and merge asset contracts --- .../assets/contract/articulation_factory.py | 65 +- .../{mock_backends.py => backends.py} | 92 +- .../test/assets/contract/capabilities.py | 146 - .../isaaclab/test/assets/contract/conftest.py | 8 +- .../test/assets/contract/contract_boot.py | 58 - .../test/assets/contract/public_surface.py | 267 -- .../rigid_object_collection_factory.py | 265 -- .../assets/contract/rigid_object_factory.py | 404 +-- .../contract/test_articulation_contract.py | 1810 +++-------- .../test_articulation_ordering_contract.py | 2884 +++++------------ .../contract/test_asset_contract_api.py | 101 - .../test_rigid_object_collection_contract.py | 867 ----- .../contract/test_rigid_object_contract.py | 1021 ++---- 13 files changed, 1915 insertions(+), 6073 deletions(-) rename source/isaaclab/test/assets/contract/{mock_backends.py => backends.py} (52%) delete mode 100644 source/isaaclab/test/assets/contract/capabilities.py delete mode 100644 source/isaaclab/test/assets/contract/contract_boot.py delete mode 100644 source/isaaclab/test/assets/contract/public_surface.py delete mode 100644 source/isaaclab/test/assets/contract/rigid_object_collection_factory.py delete mode 100644 source/isaaclab/test/assets/contract/test_asset_contract_api.py delete mode 100644 source/isaaclab/test/assets/contract/test_rigid_object_collection_contract.py diff --git a/source/isaaclab/test/assets/contract/articulation_factory.py b/source/isaaclab/test/assets/contract/articulation_factory.py index 9adda86b06b8..07a4dd7e67b6 100644 --- a/source/isaaclab/test/assets/contract/articulation_factory.py +++ b/source/isaaclab/test/assets/contract/articulation_factory.py @@ -17,22 +17,19 @@ from isaaclab.assets.articulation.articulation_cfg import ArticulationCfg from isaaclab.utils.wrench_composer import WrenchComposer -from .capabilities import available_backends -from .mock_backends import install_physx_recording_setters, patch_ovphysx_manager, patch_physx_manager +from .backends import AVAILABLE, install_physx_recording_setters, patch_ovphysx_manager, patch_physx_manager -BACKENDS = available_backends() - -if "physx" in BACKENDS: +if "physx" in AVAILABLE: from isaaclab_physx.assets.articulation.articulation import Articulation as PhysXArticulation from isaaclab_physx.assets.articulation.articulation_data import ArticulationData as PhysXArticulationData from isaaclab_physx.test.fixtures.views import MockArticulationViewWarp as PhysXMockArticulationViewWarp -if "newton" in BACKENDS: +if "newton" in AVAILABLE: 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 -if "ovphysx" in BACKENDS: +if "ovphysx" in AVAILABLE: 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 @@ -439,50 +436,12 @@ def get_articulation( 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, + **kwargs, ): - 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, - monkeypatch=monkeypatch, - ) - 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, - monkeypatch=monkeypatch, - ) - 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, - monkeypatch=monkeypatch, - ) - else: - raise ValueError(f"Invalid backend: {backend}") + """Create a mocked articulation of one backend and return it with its backend view or bindings.""" + factory = { + "physx": create_physx_articulation, + "ovphysx": create_ovphysx_articulation, + "newton": create_newton_articulation, + }[backend] + return factory(num_instances, num_joints, num_bodies, num_fixed_tendons, num_spatial_tendons, device, **kwargs) diff --git a/source/isaaclab/test/assets/contract/mock_backends.py b/source/isaaclab/test/assets/contract/backends.py similarity index 52% rename from source/isaaclab/test/assets/contract/mock_backends.py rename to source/isaaclab/test/assets/contract/backends.py index c9fe9dbcf5f0..a2650715070a 100644 --- a/source/isaaclab/test/assets/contract/mock_backends.py +++ b/source/isaaclab/test/assets/contract/backends.py @@ -3,17 +3,97 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Backend-manager and tensor-view doubles shared by the mocked contract factories.""" +"""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 - -def patch_physx_manager(*, monkeypatch: pytest.MonkeyPatch) -> None: +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 @@ -25,7 +105,7 @@ def patch_physx_manager(*, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(PhysxManager, "_scene_data_backend", PhysxSceneDataBackend(), raising=False) -def patch_ovphysx_manager(*, monkeypatch: pytest.MonkeyPatch) -> None: +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 @@ -79,6 +159,4 @@ def read_backend_joint_state(backend: str, art, raw_backend) -> tuple[np.ndarray 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}") + return art.data._sim_bind_joint_pos.numpy().copy(), art.data._sim_bind_joint_vel.numpy().copy() diff --git a/source/isaaclab/test/assets/contract/capabilities.py b/source/isaaclab/test/assets/contract/capabilities.py deleted file mode 100644 index a26a48e1baf5..000000000000 --- a/source/isaaclab/test/assets/contract/capabilities.py +++ /dev/null @@ -1,146 +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 - -"""Explicit backend capability declarations for the shared asset contracts. - -Every contract test is parametrized over all declared backends. A backend that is unavailable in this process, or -that declares a capability unsupported, still collects as a skipped case with the declared reason, so a missing -backend or an unsupported branch never silently disappears from the matrix. -""" - -from dataclasses import dataclass, field -from importlib.util import find_spec - -import pytest -import warp as wp -from _pytest.mark.structures import ParameterSet - -# Install the kitless import stubs before resolving the declared backend dependencies. -from . import contract_boot # noqa: F401 - - -@dataclass(frozen=True) -class BackendDeclaration: - """Dependencies and supported contract capabilities of one backend.""" - - name: str - required_modules: tuple[str, ...] - capabilities: frozenset[str] - unsupported: dict[str, str] = field(default_factory=dict) - requires_cuda_runtime: bool = False - - -@dataclass(frozen=True) -class BackendStatus: - """Whether one declared backend can run the contracts in this process.""" - - declaration: BackendDeclaration - available: bool - reason: str | None - - -_SHARED_CAPABILITIES = frozenset({"api", "data", "writes", "ordering", "fixed_tendons"}) -_PHYSX_FAMILY_CAPABILITIES = frozenset( - {"fixed_tendon_extended", "spatial_tendons", "fixed_tendon_target_scheduling", "mocked_dynamics"} -) - -BACKEND_DECLARATIONS = ( - BackendDeclaration( - name="physx", - required_modules=("carb", "isaaclab_physx"), - capabilities=_SHARED_CAPABILITIES | _PHYSX_FAMILY_CAPABILITIES, - ), - BackendDeclaration( - name="newton", - required_modules=("isaaclab_newton",), - capabilities=_SHARED_CAPABILITIES, - unsupported={ - "fixed_tendon_extended": "Newton does not implement fixed-tendon limit stiffness, rest length, or offset", - "spatial_tendons": "Newton does not support spatial tendons", - "fixed_tendon_target_scheduling": ( - "Newton commands tendon targets through a MuJoCo tendon actuator the mocked contract does not build" - ), - "mocked_dynamics": "Newton computes task-space dynamics from a solver model that the mocked contract lacks", - }, - ), - BackendDeclaration( - name="ovphysx", - required_modules=("ovphysx", "isaaclab_ov"), - capabilities=_SHARED_CAPABILITIES | _PHYSX_FAMILY_CAPABILITIES, - # The mocked bindings allocate pinned host staging buffers even for CPU tensors. - requires_cuda_runtime=True, - ), -) - - -def _module_available(module_name: str) -> bool: - """Return whether an explicitly declared dependency can be resolved.""" - try: - return find_spec(module_name) is not None - except ModuleNotFoundError: - return False - - -def _evaluate_backend(declaration: BackendDeclaration, cuda_available: bool) -> BackendStatus: - """Evaluate one backend declaration against the current process.""" - for module_name in declaration.required_modules: - if not _module_available(module_name): - return BackendStatus(declaration, available=False, reason=f"missing required module: {module_name}") - if declaration.requires_cuda_runtime and not cuda_available: - return BackendStatus(declaration, available=False, reason="CUDA runtime unavailable") - return BackendStatus(declaration, available=True, reason=None) - - -BACKEND_STATUSES = tuple(_evaluate_backend(declaration, wp.is_cuda_available()) for declaration in BACKEND_DECLARATIONS) -_STATUS_BY_NAME = {status.declaration.name: status for status in BACKEND_STATUSES} - - -def _unsupported_reason(declaration: BackendDeclaration, capability: str) -> str | None: - """Return why a backend does not run a capability, or None when it does.""" - if capability in declaration.capabilities: - return None - reason = declaration.unsupported.get(capability) - if reason is None: - raise ValueError(f"{declaration.name} has no declaration for capability {capability!r}") - return reason - - -def backend_parameters(capability: str, *, names: tuple[str, ...] | None = None) -> list[ParameterSet]: - """Return one pytest parameter per declared backend, skipped with its reason when it cannot run. - - Args: - capability: Contract capability the test exercises. - names: Optional subset of backend names for a backend-specific branch. Defaults to every backend. - - Returns: - Backend parameters in declaration order. - """ - parameters = [] - for status in BACKEND_STATUSES: - declaration = status.declaration - if names is not None and declaration.name not in names: - continue - reason = status.reason or _unsupported_reason(declaration, capability) - marks = () if reason is None else pytest.mark.skip(reason=reason) - parameters.append(pytest.param(declaration.name, marks=marks, id=declaration.name)) - return parameters - - -def requires_backend(name: str) -> pytest.MarkDecorator: - """Skip a backend-specific test with the declared reason when that backend is unavailable.""" - reason = _STATUS_BY_NAME[name].reason - return pytest.mark.skipif(reason is not None, reason=f"{name}: {reason}") - - -def require_backend_capability(backend: str, capability: str) -> None: - """Skip the current case with the declared reason when a backend lacks a capability.""" - reason = _unsupported_reason(_STATUS_BY_NAME[backend].declaration, capability) - if reason is not None: - pytest.skip(reason) - - -def available_backends() -> list[str]: - """Return the declared backends that can run in this process.""" - return [status.declaration.name for status in BACKEND_STATUSES if status.available] diff --git a/source/isaaclab/test/assets/contract/conftest.py b/source/isaaclab/test/assets/contract/conftest.py index 5b1048be3538..1e7989bb5bb0 100644 --- a/source/isaaclab/test/assets/contract/conftest.py +++ b/source/isaaclab/test/assets/contract/conftest.py @@ -9,16 +9,16 @@ from isaaclab.test.utils import DeviceScope, test_devices -from .capabilities import backend_parameters +from .backends import backends -@pytest.fixture(params=backend_parameters("api")) +@pytest.fixture(params=backends()) def backend(request: pytest.FixtureRequest) -> str: - """Run each shared contract on every declared backend, with explicit unavailable skips.""" + """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 CUDA staging without repeating CPU-only bookkeeping.""" + """Exercise CPU buffers and CUDA staging.""" return request.param diff --git a/source/isaaclab/test/assets/contract/contract_boot.py b/source/isaaclab/test/assets/contract/contract_boot.py deleted file mode 100644 index 6d62e38c8379..000000000000 --- a/source/isaaclab/test/assets/contract/contract_boot.py +++ /dev/null @@ -1,58 +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 - -"""Shared Kit and kitless bootstrap for the asset contract tests.""" - -import os -import sys -from importlib.machinery import ModuleSpec -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: - - def _install_stub(module_name: str, *, is_package: bool = False) -> None: - """Install one import-boundary module double unless the real module is already loaded.""" - if module_name in sys.modules: - return - stub = MagicMock() - stub.__spec__ = ModuleSpec(module_name, loader=None, is_package=is_package) - if is_package: - stub.__path__ = [] - sys.modules[module_name] = stub - if "." in module_name: - parent_name, attribute = module_name.rsplit(".", 1) - setattr(sys.modules[parent_name], attribute, stub) - - # Without Kit's Python runtime, ``carb`` and ``usdrt`` are missing even when ``isaaclab_physx`` is installed. Stub - # only that runtime boundary so the contracts still import the real PhysX asset and data classes and their views. - _install_stub("carb") - _install_stub("usdrt", is_package=True) - - # ``omni`` is a real namespace package in kitless runs. Install missing submodules in both ``sys.modules`` and the - # namespace attributes. - import omni as _omni # noqa: F401 - - for _module_name, _is_package in ( - ("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), - ): - _install_stub(_module_name, is_package=_is_package) - sys.modules["omni.kit.app"].get_app.return_value = None diff --git a/source/isaaclab/test/assets/contract/public_surface.py b/source/isaaclab/test/assets/contract/public_surface.py deleted file mode 100644 index 04ac1c919bdf..000000000000 --- a/source/isaaclab/test/assets/contract/public_surface.py +++ /dev/null @@ -1,267 +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 - -"""Public base-class surface classified by the shared contract that owns it. - -:data:`PUBLIC_SURFACE_CONTRACTS` is the single inventory of the public members declared by the asset base classes. -``test_base_public_surface_has_an_explicit_contract_classification`` diffs it against the classes, so adding a public -member forces its contract classification and removing one is reported as a stale entry. -""" - -from dataclasses import dataclass -from enum import StrEnum - -from isaaclab.assets.articulation.base_articulation import BaseArticulation -from isaaclab.assets.articulation.base_articulation_data import BaseArticulationData -from isaaclab.assets.asset_base import AssetBase -from isaaclab.assets.rigid_object.base_rigid_object import BaseRigidObject -from isaaclab.assets.rigid_object.base_rigid_object_data import BaseRigidObjectData -from isaaclab.assets.rigid_object_collection.base_rigid_object_collection import BaseRigidObjectCollection -from isaaclab.assets.rigid_object_collection.base_rigid_object_collection_data import BaseRigidObjectCollectionData - - -class ContractKind(StrEnum): - """Contract family that owns a public member.""" - - API = "api" - DATA = "data" - WRITE = "write" - OUT_OF_SCOPE = "out_of_scope" - - -@dataclass(frozen=True) -class PublicMemberContract: - """Classify one public member, with the reason for an out-of-scope exclusion.""" - - kind: ContractKind - reason: str | None = None - - -BASE_SURFACE_CLASSES = ( - AssetBase, - BaseRigidObject, - BaseRigidObjectData, - BaseRigidObjectCollection, - BaseRigidObjectCollectionData, - BaseArticulation, - BaseArticulationData, -) - - -def _mapping( - class_name: str, members: str, kind: ContractKind, reason: str | None = None -) -> dict[str, PublicMemberContract]: - """Classify a whitespace-separated group of members declared by one class.""" - return {f"{class_name}.{member}": PublicMemberContract(kind, reason) for member in members.split()} - - -PUBLIC_SURFACE_CONTRACTS = { - **_mapping("AssetBase", "data device is_initialized num_instances", ContractKind.API), - **_mapping( - "AssetBase", - "assert_shape_and_dtype assert_shape_and_dtype_mask reset update write_data_to_sim", - ContractKind.WRITE, - ), - **_mapping( - "AssetBase", - "has_debug_vis_implementation set_debug_vis", - ContractKind.OUT_OF_SCOPE, - reason="debug visualization requires a live USD stage and Kit application", - ), - **_mapping( - "BaseRigidObject", - """ - body_names data find_bodies instantaneous_wrench_composer num_bodies num_instances permanent_wrench_composer - root_view - """, - ContractKind.API, - ), - **_mapping( - "BaseRigidObject", - """ - reset set_coms set_coms_index set_coms_mask set_external_force_and_torque set_inertias set_inertias_index - set_inertias_mask set_masses set_masses_index set_masses_mask update write_data_to_sim - write_root_com_pose_to_sim write_root_com_pose_to_sim_index write_root_com_pose_to_sim_mask - write_root_com_state_to_sim write_root_com_velocity_to_sim write_root_com_velocity_to_sim_index - write_root_com_velocity_to_sim_mask write_root_link_pose_to_sim write_root_link_pose_to_sim_index - write_root_link_pose_to_sim_mask write_root_link_state_to_sim write_root_link_velocity_to_sim - write_root_link_velocity_to_sim_index write_root_link_velocity_to_sim_mask write_root_pose_to_sim - write_root_pose_to_sim_index write_root_pose_to_sim_mask write_root_state_to_sim write_root_velocity_to_sim - write_root_velocity_to_sim_index write_root_velocity_to_sim_mask - """, - ContractKind.WRITE, - ), - **_mapping( - "BaseRigidObjectData", - """ - body_acc_w body_ang_acc_w body_ang_vel_w body_com_acc_w body_com_ang_acc_w body_com_ang_vel_w - body_com_lin_acc_w body_com_lin_vel_w body_com_pos_b body_com_pos_w body_com_pose_b body_com_pose_w - body_com_quat_b body_com_quat_w body_com_state_w body_com_vel_w body_inertia body_lin_acc_w body_lin_vel_w - body_link_ang_vel_w body_link_lin_vel_w body_link_pos_w body_link_pose_w body_link_quat_w body_link_state_w - body_link_vel_w body_mass body_names body_pos_w body_pose_w body_quat_w body_state_w body_vel_w com_pos_b - com_quat_b default_inertia default_mass default_root_pose default_root_state default_root_vel heading_w - projected_gravity_b root_ang_vel_b root_ang_vel_w root_com_ang_vel_b root_com_ang_vel_w root_com_lin_vel_b - root_com_lin_vel_w root_com_pos_w root_com_pose_w root_com_quat_w root_com_state_w root_com_vel_w - root_lin_vel_b root_lin_vel_w root_link_ang_vel_b root_link_ang_vel_w root_link_lin_vel_b root_link_lin_vel_w - root_link_pos_w root_link_pose_w root_link_quat_w root_link_state_w root_link_vel_w root_pos_w root_pose_w - root_quat_w root_state_w root_vel_w update - """, - ContractKind.DATA, - ), - **_mapping( - "BaseRigidObjectCollection", - """ - body_names data find_bodies find_objects instantaneous_wrench_composer num_bodies num_instances num_objects - object_names permanent_wrench_composer root_view - """, - ContractKind.API, - ), - **_mapping( - "BaseRigidObjectCollection", - """ - reset set_coms set_coms_index set_coms_mask set_external_force_and_torque set_inertias set_inertias_index - set_inertias_mask set_masses set_masses_index set_masses_mask update write_body_com_pose_to_sim - write_body_com_pose_to_sim_index write_body_com_pose_to_sim_mask write_body_com_state_to_sim - write_body_com_velocity_to_sim write_body_com_velocity_to_sim_index write_body_com_velocity_to_sim_mask - write_body_link_pose_to_sim write_body_link_pose_to_sim_index write_body_link_pose_to_sim_mask - write_body_link_state_to_sim write_body_link_velocity_to_sim write_body_link_velocity_to_sim_index - write_body_link_velocity_to_sim_mask write_body_pose_to_sim write_body_pose_to_sim_index - write_body_pose_to_sim_mask write_body_state_to_sim write_body_velocity_to_sim write_body_velocity_to_sim_index - write_body_velocity_to_sim_mask write_data_to_sim write_object_com_pose_to_sim write_object_com_state_to_sim - write_object_com_velocity_to_sim write_object_link_pose_to_sim write_object_link_state_to_sim - write_object_link_velocity_to_sim write_object_pose_to_sim write_object_state_to_sim - write_object_velocity_to_sim - """, - ContractKind.WRITE, - ), - **_mapping( - "BaseRigidObjectCollectionData", - """ - body_acc_w body_ang_acc_w body_ang_vel_w body_com_acc_w body_com_ang_acc_w body_com_ang_vel_b - body_com_ang_vel_w body_com_lin_acc_w body_com_lin_vel_b body_com_lin_vel_w body_com_pos_b body_com_pos_w - body_com_pose_b body_com_pose_w body_com_quat_b body_com_quat_w body_com_state_w body_com_vel_w body_inertia - body_lin_acc_w body_lin_vel_w body_link_ang_vel_b body_link_ang_vel_w body_link_lin_vel_b body_link_lin_vel_w - body_link_pos_w body_link_pose_w body_link_quat_w body_link_state_w body_link_vel_w body_mass body_names - body_pos_w body_pose_w body_quat_w body_state_w body_vel_w com_pos_b com_quat_b default_body_pose - default_body_state default_body_vel default_inertia default_mass default_object_pose default_object_state - default_object_vel heading_w object_acc_w object_ang_acc_w object_ang_vel_b object_ang_vel_w object_com_acc_w - object_com_ang_acc_w object_com_ang_vel_b object_com_ang_vel_w object_com_lin_acc_w object_com_lin_vel_b - object_com_lin_vel_w object_com_pos_b object_com_pos_w object_com_pose_b object_com_pose_w object_com_quat_b - object_com_quat_w object_com_state_w object_com_vel_w object_lin_acc_w object_lin_vel_b object_lin_vel_w - object_link_ang_vel_b object_link_ang_vel_w object_link_lin_vel_b object_link_lin_vel_w object_link_pos_w - object_link_pose_w object_link_quat_w object_link_state_w object_link_vel_w object_pos_w object_pose_w - object_quat_w object_state_w object_vel_w projected_gravity_b update - """, - ContractKind.DATA, - ), - **_mapping( - "BaseArticulation", - """ - backend_body_names backend_joint_names body_names body_ordering data find_bodies find_fixed_tendons find_joints - find_spatial_tendons fixed_tendon_names instantaneous_wrench_composer is_fixed_base joint_names joint_ordering - map_body_ids_to_backend map_joint_ids_to_backend num_base_dofs num_bodies num_fixed_tendons num_instances - num_joints num_spatial_tendons permanent_wrench_composer root_view spatial_tendon_names - """, - ContractKind.API, - ), - **_mapping( - "BaseArticulation", - """ - reset set_coms set_coms_index set_coms_mask set_external_force_and_torque set_fixed_tendon_damping - set_fixed_tendon_damping_index set_fixed_tendon_damping_mask set_fixed_tendon_limit - set_fixed_tendon_limit_stiffness set_fixed_tendon_limit_stiffness_index set_fixed_tendon_limit_stiffness_mask - set_fixed_tendon_offset set_fixed_tendon_offset_index set_fixed_tendon_offset_mask - set_fixed_tendon_position_limit set_fixed_tendon_position_limit_index set_fixed_tendon_position_limit_mask - set_fixed_tendon_position_target_index set_fixed_tendon_position_target_mask set_fixed_tendon_rest_length - set_fixed_tendon_rest_length_index set_fixed_tendon_rest_length_mask set_fixed_tendon_stiffness - set_fixed_tendon_stiffness_index set_fixed_tendon_stiffness_mask set_inertias set_inertias_index - set_inertias_mask set_joint_effort_target set_joint_effort_target_index set_joint_effort_target_mask - set_joint_position_target set_joint_position_target_index set_joint_position_target_mask - set_joint_velocity_target set_joint_velocity_target_index set_joint_velocity_target_mask set_masses - set_masses_index set_masses_mask set_spatial_tendon_damping set_spatial_tendon_damping_index - set_spatial_tendon_damping_mask set_spatial_tendon_limit_stiffness set_spatial_tendon_limit_stiffness_index - set_spatial_tendon_limit_stiffness_mask set_spatial_tendon_offset set_spatial_tendon_offset_index - set_spatial_tendon_offset_mask set_spatial_tendon_stiffness set_spatial_tendon_stiffness_index - set_spatial_tendon_stiffness_mask update write_data_to_sim write_fixed_tendon_properties_to_sim - write_fixed_tendon_properties_to_sim_index write_fixed_tendon_properties_to_sim_mask - write_joint_armature_to_sim write_joint_armature_to_sim_index write_joint_armature_to_sim_mask - write_joint_damping_to_sim write_joint_damping_to_sim_index write_joint_damping_to_sim_mask - write_joint_effort_limit_to_sim write_joint_effort_limit_to_sim_index write_joint_effort_limit_to_sim_mask - write_joint_friction_coefficient_to_sim write_joint_friction_coefficient_to_sim_index - write_joint_friction_coefficient_to_sim_mask write_joint_friction_to_sim write_joint_limits_to_sim - write_joint_position_limit_to_sim write_joint_position_limit_to_sim_index - write_joint_position_limit_to_sim_mask write_joint_position_to_sim write_joint_position_to_sim_index - write_joint_position_to_sim_mask write_joint_state_to_sim write_joint_stiffness_to_sim - write_joint_stiffness_to_sim_index write_joint_stiffness_to_sim_mask write_joint_velocity_limit_to_sim - write_joint_velocity_limit_to_sim_index write_joint_velocity_limit_to_sim_mask write_joint_velocity_to_sim - write_joint_velocity_to_sim_index write_joint_velocity_to_sim_mask write_root_com_pose_to_sim - write_root_com_pose_to_sim_index write_root_com_pose_to_sim_mask write_root_com_state_to_sim - write_root_com_velocity_to_sim write_root_com_velocity_to_sim_index write_root_com_velocity_to_sim_mask - write_root_link_pose_to_sim write_root_link_pose_to_sim_index write_root_link_pose_to_sim_mask - write_root_link_state_to_sim write_root_link_velocity_to_sim write_root_link_velocity_to_sim_index - write_root_link_velocity_to_sim_mask write_root_pose_to_sim write_root_pose_to_sim_index - write_root_pose_to_sim_mask write_root_state_to_sim write_root_velocity_to_sim write_root_velocity_to_sim_index - write_root_velocity_to_sim_mask write_spatial_tendon_properties_to_sim - write_spatial_tendon_properties_to_sim_index write_spatial_tendon_properties_to_sim_mask - """, - ContractKind.WRITE, - ), - **_mapping( - "BaseArticulationData", - """ - applied_torque bind_actuator_collection body_acc_w body_ang_acc_w body_ang_vel_w body_com_acc_w - body_com_ang_acc_w body_com_ang_vel_w body_com_jacobian_w body_com_lin_acc_w body_com_lin_vel_w body_com_pos_b - body_com_pos_w body_com_pose_b body_com_pose_w body_com_quat_b body_com_quat_w body_com_state_w body_com_vel_w - body_inertia body_lin_acc_w body_lin_vel_w body_link_ang_vel_w body_link_jacobian_w body_link_lin_vel_w - body_link_pos_w body_link_pose_w body_link_quat_w body_link_state_w body_link_vel_w body_mass body_names - body_ordering body_pos_w body_pose_w body_quat_w body_state_w body_vel_w com_pos_b com_quat_b computed_torque - default_fixed_tendon_damping default_fixed_tendon_limit default_fixed_tendon_limit_stiffness - default_fixed_tendon_offset default_fixed_tendon_pos_limits default_fixed_tendon_rest_length - default_fixed_tendon_stiffness default_inertia default_joint_armature default_joint_damping - default_joint_friction default_joint_friction_coeff default_joint_limits default_joint_pos - default_joint_pos_limits default_joint_stiffness default_joint_vel default_joint_viscous_friction_coeff - default_mass default_root_pose default_root_state default_root_vel default_spatial_tendon_damping - default_spatial_tendon_limit_stiffness default_spatial_tendon_offset default_spatial_tendon_stiffness - fixed_tendon_damping fixed_tendon_limit fixed_tendon_limit_stiffness fixed_tendon_names fixed_tendon_offset - fixed_tendon_pos_limits fixed_tendon_rest_length fixed_tendon_stiffness gravity_compensation_forces - has_body_ordering has_joint_ordering heading_w joint_acc joint_armature joint_damping joint_effort_limits - joint_effort_target joint_friction joint_friction_coeff joint_limits joint_names joint_ordering joint_pos - joint_pos_limits joint_pos_target joint_stiffness joint_vel joint_vel_limits joint_vel_target - joint_velocity_limits mass_matrix projected_gravity_b root_ang_vel_b root_ang_vel_w root_com_ang_vel_b - root_com_ang_vel_w root_com_lin_vel_b root_com_lin_vel_w root_com_pos_w root_com_pose_w root_com_quat_w - root_com_state_w root_com_vel_w root_lin_vel_b root_lin_vel_w root_link_ang_vel_b root_link_ang_vel_w - root_link_lin_vel_b root_link_lin_vel_w root_link_pos_w root_link_pose_w root_link_quat_w root_link_state_w - root_link_vel_w root_pos_w root_pose_w root_quat_w root_state_w root_vel_w soft_joint_pos_limits - soft_joint_vel_limits spatial_tendon_damping spatial_tendon_limit_stiffness spatial_tendon_names - spatial_tendon_offset spatial_tendon_stiffness update - """, - ContractKind.DATA, - ), -} - - -def public_surface_mismatches( - classes: tuple[type, ...], contracts: dict[str, PublicMemberContract] -) -> dict[str, list[str]]: - """Diff the public members the classes declare against their contract classifications. - - Args: - classes: Base classes whose own public members form the audited surface. - contracts: Qualified member name to contract classification. - - Returns: - Sorted ``missing`` (declared but unclassified), ``stale`` (classified but no longer declared), and - ``unreasoned`` (out-of-scope without a reason) member names; all empty when the inventory is exact. - """ - declared = {f"{cls.__name__}.{name}" for cls in classes for name in vars(cls) if not name.startswith("_")} - return { - "missing": sorted(declared - contracts.keys()), - "stale": sorted(contracts.keys() - declared), - "unreasoned": sorted( - name - for name, contract in contracts.items() - if contract.kind is ContractKind.OUT_OF_SCOPE and not contract.reason - ), - } diff --git a/source/isaaclab/test/assets/contract/rigid_object_collection_factory.py b/source/isaaclab/test/assets/contract/rigid_object_collection_factory.py deleted file mode 100644 index 4f9442c80162..000000000000 --- a/source/isaaclab/test/assets/contract/rigid_object_collection_factory.py +++ /dev/null @@ -1,265 +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 the contract tests.""" - -from unittest.mock import MagicMock - -import numpy as np -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 isaaclab.utils.wrench_composer import WrenchComposer - -from .capabilities import available_backends -from .mock_backends import install_physx_recording_setters, patch_ovphysx_manager, patch_physx_manager - -BACKENDS = available_backends() - -if "physx" in BACKENDS: - 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.test.fixtures.views import MockRigidBodyViewWarp as PhysXMockRigidBodyViewWarp - -if "newton" in BACKENDS: - 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 - -if "ovphysx" in BACKENDS: - 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 - -_PHYSX_RIGID_BODY_STORAGE = { - "set_transforms": "_transforms", - "set_velocities": "_velocities", - "set_masses": "_masses", - "set_coms": "_coms", - "set_inertias": "_inertias", -} - - -def create_physx_rigid_object_collection( - num_instances: int = 2, num_bodies: int = 3, device: str = "cuda:0", *, monkeypatch: pytest.MonkeyPatch -): - """Create a test RigidObjectCollection instance with mocked dependencies.""" - patch_physx_manager(monkeypatch=monkeypatch) - 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() - install_physx_recording_setters(mock_view, _PHYSX_RIGID_BODY_STORAGE) - - collection._root_view = mock_view - collection._device = device - collection._num_bodies = num_bodies - collection._num_instances = num_instances - collection._body_names_list = [f"object_{i}" for i in range(num_bodies)] - - # Create RigidObjectCollectionData instance - data = PhysXRigidObjectCollectionData(mock_view, num_bodies, device) - 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) - collection._instantaneous_wrench_composer = mock_inst_wrench - collection._permanent_wrench_composer = mock_perm_wrench - - # Prevent __del__ / _clear_callbacks from raising AttributeError - collection._initialize_handle = None - collection._invalidate_initialize_handle = None - collection._prim_deletion_handle = None - collection._debug_vis_handle = None - - # Set up index arrays - collection._ALL_ENV_INDICES = wp.array(np.arange(num_instances, dtype=np.int32), device=device) - 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) - collection._ALL_VIEW_INDICES = all_view_ids - collection._sim_view_ids = wp.empty(num_view_ids, dtype=wp.int32, device=device) - collection._sim_view_ids_views = {} - cpu_all_view_ids = wp.empty(num_view_ids, dtype=wp.int32, device="cpu", pinned=wp.is_cuda_available()) - wp.copy(cpu_all_view_ids, all_view_ids) - collection._cpu_all_view_ids = cpu_all_view_ids - collection._cpu_view_ids = wp.empty(num_view_ids, dtype=wp.int32, device="cpu", pinned=wp.is_cuda_available()) - 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", *, monkeypatch: pytest.MonkeyPatch -): - """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 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 until the test finishes. - monkeypatch.setattr(newton_data_module, "SimulationManager", mock_manager, raising=False) - monkeypatch.setattr(newton_coll_module, "SimulationManager", mock_manager, raising=False) - data = NewtonRigidObjectCollectionData(mock_view, num_bodies, device) - - # 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) - - collection._root_view = mock_view - collection._device = device - collection._num_bodies = num_bodies - collection._num_instances = num_instances - collection._body_names_list = body_names - collection._data = data - data.body_names = body_names - - # Wrench composers (Newton-specific) - mock_inst_wrench = WrenchComposer(collection) - mock_perm_wrench = WrenchComposer(collection) - collection._instantaneous_wrench_composer = mock_inst_wrench - collection._permanent_wrench_composer = mock_perm_wrench - - # Prevent __del__ / _clear_callbacks from raising AttributeError - collection._initialize_handle = None - collection._invalidate_initialize_handle = None - collection._prim_deletion_handle = None - collection._debug_vis_handle = None - - # Index arrays (warp) - collection._ALL_ENV_INDICES = wp.array(np.arange(num_instances, dtype=np.int32), device=device) - collection._ALL_BODY_INDICES = wp.array(np.arange(num_bodies, dtype=np.int32), device=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, mock_view - - -def create_ovphysx_rigid_object_collection( - num_instances: int = 2, num_bodies: int = 3, device: str = "cuda:0", *, monkeypatch: pytest.MonkeyPatch -): - """Create a test OVPhysX RigidObjectCollection instance with mocked tensor bindings.""" - patch_ovphysx_manager(monkeypatch=monkeypatch) - 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() - - collection._device = device - collection._ovphysx = MagicMock() - collection._root_view = mock_bindings.view - collection._bindings = mock_bindings.bindings - collection._num_instances = num_instances - collection._num_bodies = num_bodies - 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 - 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) - collection._instantaneous_wrench_composer = mock_inst_wrench - collection._permanent_wrench_composer = mock_perm_wrench - - # Prevent __del__ / _clear_callbacks from raising - collection._initialize_handle = None - collection._invalidate_initialize_handle = None - collection._prim_deletion_handle = None - 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", - *, - monkeypatch: pytest.MonkeyPatch, -): - if backend == "physx": - return create_physx_rigid_object_collection(num_instances, num_bodies, device, monkeypatch=monkeypatch) - elif backend == "ovphysx": - return create_ovphysx_rigid_object_collection(num_instances, num_bodies, device, monkeypatch=monkeypatch) - elif backend == "newton": - return create_newton_rigid_object_collection(num_instances, num_bodies, device, monkeypatch=monkeypatch) - else: - raise ValueError(f"Invalid backend: {backend}") diff --git a/source/isaaclab/test/assets/contract/rigid_object_factory.py b/source/isaaclab/test/assets/contract/rigid_object_factory.py index 7caad02b17a0..3b22f9c78424 100644 --- a/source/isaaclab/test/assets/contract/rigid_object_factory.py +++ b/source/isaaclab/test/assets/contract/rigid_object_factory.py @@ -6,7 +6,11 @@ # ignore private usage of variables warning # pyright: reportPrivateUsage=none -"""Shared mocked rigid-object backend factories for the contract tests.""" +"""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 @@ -15,29 +19,27 @@ 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 -from .capabilities import available_backends -from .mock_backends import install_physx_recording_setters, patch_ovphysx_manager, patch_physx_manager - -BACKENDS = available_backends() +from .backends import AVAILABLE, install_physx_recording_setters, patch_ovphysx_manager, patch_physx_manager -if "physx" in BACKENDS: - 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.test.fixtures.views import MockRigidBodyViewWarp as PhysXMockRigidBodyViewWarp +if "physx" in AVAILABLE: + import isaaclab_physx.assets as physx_assets + from isaaclab_physx.test.fixtures.views import MockRigidBodyViewWarp -if "newton" in BACKENDS: - 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 +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 BACKENDS: - 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 +if "ovphysx" in AVAILABLE: + import isaaclab_ov.assets as ovphysx_assets from isaaclab_ov.test.fixtures.views import MockOvPhysxBindingSet -_PHYSX_RIGID_BODY_STORAGE = { +_PHYSX_STORAGE = { "set_transforms": "_transforms", "set_velocities": "_velocities", "set_masses": "_masses", @@ -46,200 +48,214 @@ } -def create_physx_rigid_object(num_instances: int = 2, device: str = "cuda:0", *, monkeypatch: pytest.MonkeyPatch): - """Create a test RigidObject instance with mocked dependencies.""" - patch_physx_manager(monkeypatch=monkeypatch) - body_names = ["body_0"] +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 - rigid_object = object.__new__(PhysXRigidObject) - rigid_object.cfg = RigidObjectCfg(prim_path="/World/Object") +def _indices(count: int, device: str) -> wp.array: + return wp.array(np.arange(count, dtype=np.int32), device=device) - # Create PhysX mock view - mock_view = PhysXMockRigidBodyViewWarp( - count=num_instances, - device=device, - ) - mock_view.set_random_mock_data() - install_physx_recording_setters(mock_view, _PHYSX_RIGID_BODY_STORAGE) - - rigid_object._root_view = mock_view - rigid_object._device = device - - # The data reads gravity from the PhysX manager patched above. - data = PhysXRigidObjectData(mock_view, device) - 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) - rigid_object._instantaneous_wrench_composer = mock_inst_wrench - rigid_object._permanent_wrench_composer = mock_perm_wrench - - # Prevent __del__ / _clear_callbacks from raising AttributeError - rigid_object._initialize_handle = None - rigid_object._invalidate_initialize_handle = None - rigid_object._prim_deletion_handle = None - rigid_object._debug_vis_handle = None - - # Set up index arrays (warp arrays for rigid object) - rigid_object._ALL_INDICES = wp.array(np.arange(num_instances, dtype=np.int32), device=device) - rigid_object._ALL_BODY_INDICES = wp.array(np.array([0], dtype=np.int32), device=device) - - # Cached .view(wp.float32) wrappers - rigid_object._root_link_pose_w_f32 = None - 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 - rigid_object._sim_env_ids = wp.empty(N, dtype=wp.int32, device=device) - rigid_object._sim_env_ids_views = {} - cpu_env_ids = wp.array(np.arange(N, dtype=np.int32), device="cpu") - rigid_object._cpu_env_ids_all = cpu_env_ids - rigid_object._cpu_env_ids = wp.empty(N, dtype=wp.int32, device="cpu", pinned=wp.is_cuda_available()) - rigid_object._cpu_env_ids_views = {} - rigid_object._cpu_body_mass = wp.zeros((N, B), dtype=wp.float32, device="cpu") - rigid_object._cpu_body_coms = wp.zeros((N, B, 7), dtype=wp.float32, device="cpu") - 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", *, monkeypatch: pytest.MonkeyPatch): - """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( + +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 + + +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_names, + body_names=["body_0"], ) - mock_view.set_random_mock_data() - - # 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 until the test finishes. - monkeypatch.setattr(newton_data_module, "SimulationManager", mock_manager, raising=False) - data = NewtonRigidObjectData(mock_view, device) - - # Create RigidObject shell (bypass __init__) - rigid_object = object.__new__(NewtonRigidObject) - - rigid_object.cfg = RigidObjectCfg(prim_path="/World/Object") - - rigid_object._root_view = mock_view - rigid_object._device = device - rigid_object._data = data - - # Wrench composers - mock_inst_wrench = WrenchComposer(rigid_object) - mock_perm_wrench = WrenchComposer(rigid_object) - rigid_object._instantaneous_wrench_composer = mock_inst_wrench - rigid_object._permanent_wrench_composer = mock_perm_wrench - - # Prevent __del__ / _clear_callbacks from raising AttributeError - rigid_object._initialize_handle = None - rigid_object._invalidate_initialize_handle = None - rigid_object._prim_deletion_handle = None - rigid_object._debug_vis_handle = None + view.set_random_mock_data() + monkeypatch.setattr(newton_rigid_object_data, "SimulationManager", _newton_manager(num_instances, device)) - # Newton uses wp.array for indices - rigid_object._ALL_INDICES = wp.array(np.arange(num_instances, dtype=np.int32), device=device) - rigid_object._ALL_BODY_INDICES = wp.array(np.array([0], dtype=np.int32), device=device) - - # Newton uses wp.bool masks - rigid_object._ALL_ENV_MASK = wp.ones((num_instances,), dtype=wp.bool, device=device) - rigid_object._ALL_BODY_MASK = wp.ones((1,), dtype=wp.bool, device=device) - - return rigid_object, mock_view + 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 create_ovphysx_rigid_object(num_instances: int = 2, device: str = "cuda:0", *, monkeypatch: pytest.MonkeyPatch): - """Create a test OvPhysX RigidObject instance with mocked tensor bindings.""" - patch_ovphysx_manager(monkeypatch=monkeypatch) - body_names = ["base_link"] - obj = object.__new__(OvPhysxRigidObject) +def _collection_cfg(body_names: list[str]) -> RigidObjectCollectionCfg: + return RigidObjectCollectionCfg( + rigid_objects={name: RigidObjectCfg(prim_path=f"/World/{name}") for name in body_names} + ) - obj.cfg = RigidObjectCfg(prim_path="/World/object") - # Create mock binding set - mock_bindings = MockOvPhysxBindingSet( +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=1, + num_bodies=num_bodies, body_names=body_names, - asset_kind="rigid_object", + asset_kind="articulation", ) - mock_bindings.set_random_data() - - obj._device = device - obj._ovphysx = MagicMock() - obj._root_view = mock_bindings.view - obj._bindings = mock_bindings.bindings - obj._num_instances = num_instances - obj._num_bodies = 1 - 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 - 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) - obj._instantaneous_wrench_composer = mock_inst_wrench - obj._permanent_wrench_composer = mock_perm_wrench - - # Prevent __del__ / _clear_callbacks from raising - obj._initialize_handle = None - obj._invalidate_initialize_handle = None - obj._prim_deletion_handle = None - obj._debug_vis_handle = None - - return obj, mock_bindings - - -def get_rigid_object(backend: str, num_instances: int = 2, device: str = "cuda:0", *, monkeypatch: pytest.MonkeyPatch): - if backend == "physx": - return create_physx_rigid_object(num_instances, device, monkeypatch=monkeypatch) - elif backend == "ovphysx": - return create_ovphysx_rigid_object(num_instances, device, monkeypatch=monkeypatch) - elif backend == "newton": - return create_newton_rigid_object(num_instances, device, monkeypatch=monkeypatch) - else: - raise ValueError(f"Invalid backend: {backend}") + 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 index 38ef1ad62c61..3ae0e19e691b 100644 --- a/source/isaaclab/test/assets/contract/test_articulation_contract.py +++ b/source/isaaclab/test/assets/contract/test_articulation_contract.py @@ -9,339 +9,283 @@ """Shared articulation contracts across the production backends. Checks that every articulation backend provides the data and writer behavior the base articulation class advertises. -The backends run on mocked views, so these cases need neither Isaac Sim nor a GPU simulation. +The backends run on mocked views, so these cases need neither Isaac Sim nor a GPU simulation. Pure bookkeeping runs on +CPU only; getters and writers also run on CUDA because PhysX stages through CPU-pinned buffers there. """ import math import warnings -from unittest.mock import MagicMock, patch +from operator import attrgetter +from unittest.mock import MagicMock import numpy as np import pytest import torch import warp as wp +from isaaclab.assets.articulation.base_articulation_data import BaseArticulationData +from isaaclab.utils.warp import ProxyArray + from .articulation_factory import get_articulation -from .capabilities import backend_parameters, require_backend_capability, requires_backend -from .mock_backends import read_backend_joint_state +from .backends import backends, read_backend_joint_state, requires pytestmark = pytest.mark.integration -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - +# 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 _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}" +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)) -# 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. -# One fixture with distinct instance, joint, and body counts so any swapped axis shows up. -_NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES = 2, 6, 7 +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) -# --------------------------------------------------------------------------- -# Tests: Index resolution helpers -# --------------------------------------------------------------------------- - +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. -def test_resolve_env_ids_handles_tensor_view_shape(monkeypatch, backend, device): - art, _ = get_articulation(backend, num_instances=4, device=device, monkeypatch=monkeypatch) + 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 - 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() +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 test_resolve_joint_ids_handles_tensor_view_shape(monkeypatch, backend): - art, _ = get_articulation(backend, num_joints=4, device="cpu", monkeypatch=monkeypatch) - 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]) +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 - assert resolved_full.shape[0] == 4 - assert resolved_view.shape[0] == 2 +def _assert_buffers_stale(data, buffers) -> None: + for name, buffer in buffers: + assert buffer.timestamp < data._sim_timestamp, name -def test_resolve_body_ids_handles_tensor_view_shape(monkeypatch, backend): - art, _ = get_articulation(backend, num_bodies=4, device="cpu", monkeypatch=monkeypatch) - 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]) +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 - assert resolved_full.shape[0] == 4 - assert resolved_view.shape[0] == 2 + monkeypatch.setattr(NewtonManager, "add_model_change", MagicMock()) # --------------------------------------------------------------------------- -# Tests: Articulation properties and finders +# Counts, names, finders, and selectors # --------------------------------------------------------------------------- @pytest.mark.parametrize("is_fixed_base", [False, True], ids=["floating", "fixed"]) -def test_articulation_counts_names_and_finders(monkeypatch, 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, - monkeypatch=monkeypatch, - ) +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 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) + 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 - -@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(monkeypatch, 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", - monkeypatch=monkeypatch, - ) - if finder_name == "find_spatial_tendons": - require_backend_capability(backend, "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]) - - -def _fixed_tendon_articulation(backend: str, *, monkeypatch: pytest.MonkeyPatch): - """Create a two-environment articulation with two fixed tendons and no spatial tendons.""" - art, _ = get_articulation( - backend, - num_instances=2, - num_joints=2, - num_bodies=2, - num_fixed_tendons=2, - num_spatial_tendons=0, - device="cpu", - monkeypatch=monkeypatch, - ) - return art + 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() -@pytest.mark.parametrize("backend", backend_parameters("fixed_tendon_target_scheduling")) -def test_commanding_a_target_schedules_the_write(monkeypatch, backend): - """The target setter marks the write pending; a static offset write keeps its own contract.""" - art = _fixed_tendon_articulation(backend, monkeypatch=monkeypatch) - art.set_fixed_tendon_position_target_index(target=torch.zeros((2, 2), dtype=torch.float32)) - assert art._fixed_tendon_target_dirty is True +# --------------------------------------------------------------------------- +# Fixed-tendon targets +# --------------------------------------------------------------------------- -@requires_backend("newton") -def test_newton_reports_a_missing_tendon_actuator(monkeypatch): - """Without a MuJoCo tendon actuator the solver has no adapter, so commanding must say so.""" - art = _fixed_tendon_articulation("newton", monkeypatch=monkeypatch) - with pytest.raises(RuntimeError, match="no MuJoCo tendon actuator"): - art.set_fixed_tendon_position_target_index(target=torch.zeros((2, 2), dtype=torch.float32)) +@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_backend("newton") -def test_a_solver_without_tendon_transmission_refuses_to_build_an_adapter(monkeypatch): - """A solver with no tendon transmission says so, rather than returning nothing.""" +@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 = _fixed_tendon_articulation("newton", monkeypatch=monkeypatch) + 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) # --------------------------------------------------------------------------- -# Tests: resolve_matching_names caching behavior +# Data properties # --------------------------------------------------------------------------- -def test_unmatched_regex_raises(monkeypatch, backend): - """ValueError from resolve_matching_names propagates correctly.""" - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device="cpu", monkeypatch=monkeypatch) - with pytest.raises(ValueError): - art.find_bodies("nonexistent_body_xyz") - with pytest.raises(ValueError): - art.find_joints("nonexistent_joint_xyz") - - -def test_mutating_result_does_not_corrupt_cache(monkeypatch, backend): - """Mutating returned lists must not affect future cached results.""" - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device="cpu", monkeypatch=monkeypatch) - - 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 - - -def test_find_with_preserve_order(monkeypatch, 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", monkeypatch=monkeypatch) - 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), +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_articulation_data_property_contract(monkeypatch, backend, device): - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device, monkeypatch=monkeypatch) +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) - storage = wp.empty((2 * binding.shape[0], *binding.shape[1:]), dtype=binding.dtype, device=device) - view = storage[::2] + 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": - 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: + # 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) - _check_proxy_array(value, expected_shape=shapes[shape_kind], expected_dtype=dtype, name=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"))): @@ -352,63 +296,43 @@ def test_articulation_data_property_contract(monkeypatch, backend, device): 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 -@requires_backend("physx") -def test_physx_set_coms_index_updates_body_com_pose_b_cache(monkeypatch): - art, view = get_articulation( - "physx", num_instances=2, num_joints=3, num_bodies=4, device="cpu", monkeypatch=monkeypatch - ) - - 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 - - -@requires_backend("physx") -def test_physx_joint_position_write_preserves_body_com_pose_b_cache(monkeypatch): - art, view = get_articulation( - "physx", num_instances=2, num_joints=3, num_bodies=4, device="cpu", monkeypatch=monkeypatch - ) - - 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 +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)) - 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 + 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) - assert num_get_coms_calls == 1 + 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", backend_parameters("mocked_dynamics")) +@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, _ = get_articulation( - backend, num_instances=2, num_joints=3, num_bodies=4, device=device, monkeypatch=monkeypatch - ) + 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): @@ -424,554 +348,278 @@ def test_dynamics_buffers_allocate_on_read_and_reuse_across_refreshes(monkeypatc assert art.data.body_link_jacobian_w is wrapper -def test_actuator_compatibility_projections_are_stable(monkeypatch, backend): - num_instances, num_joints = 2, 4 - art, _ = get_articulation(backend, num_instances, num_joints, 5, device="cpu", monkeypatch=monkeypatch) - 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] - - -def test_deprecated_command_aliases_return_actuator_collection_buffers(monkeypatch, backend): - """The deprecated data command and torque properties warn and alias the actuator collection's buffers.""" - art, _ = get_articulation(backend, 2, 4, 5, device="cpu", monkeypatch=monkeypatch) - actuators = art.actuators - 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 - - # --------------------------------------------------------------------------- -# Tests: Alias/shorthand properties +# Cache invalidation # --------------------------------------------------------------------------- -_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"), -] +_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"] -def test_aliases_match_canonical_values(monkeypatch, 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", monkeypatch=monkeypatch) +@_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) - 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 -# --------------------------------------------------------------------------- + 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") -# 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 -} + 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) -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) +@_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: - 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.""" + 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": - 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) - + 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) -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) + 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) -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) +@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 -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) + assert view.get_coms.call_count == 0 # --------------------------------------------------------------------------- -# Tests: Root writers — torch/warp × index/mask × all/subset × negative +# Writers # --------------------------------------------------------------------------- -# 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", +# (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( - "body_ordering", - [None, ("body_0", "body_3", "body_2", "body_1")], -) -def test_pose_write_invalidates_pose_dependent_caches(monkeypatch, backend, body_ordering): - art, _ = get_articulation( - backend, - num_instances=2, - num_joints=3, - num_bodies=4, - device="cpu", - body_ordering=body_ordering, - monkeypatch=monkeypatch, - ) - 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) - - -@pytest.mark.parametrize( - "body_ordering", - [None, ("body_0", "body_3", "body_2", "body_1")], -) -def test_velocity_write_invalidates_body_frame_caches(monkeypatch, backend, body_ordering): - art, _ = get_articulation( - backend, - num_instances=2, - num_joints=3, - num_bodies=4, - device="cpu", - body_ordering=body_ordering, - monkeypatch=monkeypatch, - ) - 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) - - -@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(monkeypatch, backend, setter_kind, body_ordering): - art, _ = get_articulation( - backend, - num_instances=2, - num_joints=3, - num_bodies=4, - device="cpu", - body_ordering=body_ordering, - monkeypatch=monkeypatch, - ) +@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) - 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 isaaclab_newton.physics import NewtonManager - - with patch.object(NewtonManager, "add_model_change"): - set_coms() + 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: - set_coms() - _assert_buffers_stale(art.data, buffers) - - -# -- index variants -- - - -@pytest.mark.parametrize("method_suffix", _ROOT_POSE_METHODS) -def test_write_root_pose_to_sim_index(monkeypatch, backend, device, method_suffix): - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device, monkeypatch=monkeypatch) - 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)) + 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") -@pytest.mark.parametrize("method_suffix", _ROOT_VEL_METHODS) -def test_write_root_velocity_to_sim_index(monkeypatch, backend, device, method_suffix): - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device, monkeypatch=monkeypatch) - 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 + 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(root_velocity=_make_bad_data_torch((num_instances,), device, wp.spatial_vectorf)) - # negative: bad warp shape + method(**{kwarg: extra_env}) with pytest.raises((AssertionError, RuntimeError)): - method(root_velocity=_make_bad_data_warp((num_instances,), device, wp.spatial_vectorf)) + method(**{kwarg: wp.from_torch(extra_env.contiguous(), dtype=dtype)}) -# -- mask variants -- - - -@pytest.mark.parametrize("method_suffix", _ROOT_POSE_METHODS) -def test_write_root_pose_to_sim_mask(monkeypatch, backend, device, method_suffix): - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device, monkeypatch=monkeypatch) - 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 +@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)): - method(root_pose=_make_bad_data_warp((num_instances,), device, wp.transformf)) + art.set_fixed_tendon_stiffness_mask(stiffness=torch.ones((N + 1, FT))) -@pytest.mark.parametrize("method_suffix", _ROOT_VEL_METHODS) -def test_write_root_velocity_to_sim_mask(monkeypatch, backend, device, method_suffix): - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device, monkeypatch=monkeypatch) - 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)) +# 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 test_deprecated_joint_friction_writers(monkeypatch, 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", monkeypatch=monkeypatch) - 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) +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() -# --------------------------------------------------------------------------- -# 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"), -] +@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) -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 + 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 = get_articulation(backend, 2, 3, 2, device="cpu", monkeypatch=monkeypatch) + 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, [(name, f"_{name}") for name in cached_properties]) + 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) @@ -1015,9 +663,9 @@ def test_partial_joint_state_write_follows_int64_selectors_and_restarts_accelera _assert_buffers_stale(art.data, caches) -def test_deprecated_joint_state_writer_matches_public_writers(monkeypatch, backend): - """The deprecated combined joint-state writer warns and writes what the public writers would.""" - art, raw_backend = get_articulation(backend, 2, 3, 2, device="cpu", monkeypatch=monkeypatch) +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() @@ -1026,7 +674,6 @@ def test_deprecated_joint_state_writer_matches_public_writers(monkeypatch, backe 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) @@ -1035,536 +682,29 @@ def test_deprecated_joint_state_writer_matches_public_writers(monkeypatch, backe 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) -@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(monkeypatch, backend, device, method_base, kwarg, wp_dtype, accepts_float, getter): - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device, monkeypatch=monkeypatch) - 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)}) - - -@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(monkeypatch, backend, device, method_base, kwarg, wp_dtype, accepts_float, getter): - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device, monkeypatch=monkeypatch) - 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) - - -@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(monkeypatch, backend, device, method_base, kwarg): - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device, monkeypatch=monkeypatch) - 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)}) - - -@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(monkeypatch, backend, device, method_base, kwarg): - art, _ = get_articulation(backend, _NUM_INSTANCES, _NUM_JOINTS, _NUM_BODIES, device=device, monkeypatch=monkeypatch) - 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)}) - - -# --------------------------------------------------------------------------- -# Tests: partial root and body writes reach only the selected backend rows -# --------------------------------------------------------------------------- - -# quantity -> (writer base name, writer keyword, trailing size). Joint partial writes are covered, with and without -# a joint ordering, by ``TestArticulationOrderingWriteParity``. -_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 == "physx": - getter = { - "root_pose": raw_backend.get_root_transforms, - "root_velocity": raw_backend.get_root_velocities, - "mass": raw_backend.get_masses, - } - values = wp.to_torch(getter[quantity]()) - elif backend == "newton": - if quantity == "mass": - values = wp.to_torch(art.data._sim_bind_body_mass) - else: - getter = {"root_pose": raw_backend.get_root_transforms, "root_velocity": raw_backend.get_root_velocities} - values = wp.to_torch(getter[quantity](None)) - else: - 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) - 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): - num_instances, num_bodies = 2, 4 - art, raw_backend = get_articulation(backend, num_instances, 3, num_bodies, device="cpu", monkeypatch=monkeypatch) - writer, kwarg, trailing = _PARTIAL_WRITES[quantity] - width = num_bodies if quantity == "mass" else trailing - # Literal, per-element distinct payload; poses keep an identity rotation. - values = 100.0 * torch.arange(1, num_instances + 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 - kwargs = {} - if selection == "index": - kwargs[kwarg] = values[env : env + 1, body : body + 1] if quantity == "mass" else values[env : env + 1] - kwargs["env_ids"] = torch.tensor([env], dtype=torch.int32) - if quantity == "mass": - kwargs["body_ids"] = [body] - else: - kwargs[kwarg] = values - kwargs["env_mask"] = _make_item_mask(num_instances, [env], "cpu") - if quantity == "mass": - kwargs["body_mask"] = _make_item_mask(num_bodies, [body], "cpu") - - getattr(art, f"{writer}_{selection}")(**kwargs) - - after = _read_backend_rows(backend, art, raw_backend, quantity) - expected = before.clone() - if quantity == "mass": - expected[env, body] = values[env, body] - else: - expected[env] = values[env] - torch.testing.assert_close(after, expected, rtol=0.0, atol=0.0) - - -# --------------------------------------------------------------------------- -# Tendon tests — counts, names, data, writers -# --------------------------------------------------------------------------- - -# Every backend supports the core fixed-tendon properties. Newton declares the extended fixed-tendon properties and -# spatial tendons unsupported, so those rows stay visible as reasoned capability skips. - -# The first row covers the zero-spatial-tendon edge; the second has distinct counts on every axis. -_TENDON_DIMS = [ - (1, 2, 2, 1, 0), # fixed only - (2, 6, 7, 3, 2), # both types -] -_TENDON_DIM_NAMES = ("num_instances", "num_joints", "num_bodies", "num_fixed_tendons", "num_spatial_tendons") - -# tendon property family -> (tendon kind, backend capability) -_TENDON_FAMILIES = { - "fixed": ("fixed", "fixed_tendons"), - "fixed_extended": ("fixed", "fixed_tendon_extended"), - "spatial": ("spatial", "spatial_tendons"), -} - -# (property, family, dtype) -_TENDON_DATA_PROPERTIES = [ - ("fixed_tendon_stiffness", "fixed", wp.float32), - ("fixed_tendon_damping", "fixed", wp.float32), - ("fixed_tendon_pos_limits", "fixed", wp.vec2f), - ("fixed_tendon_limit_stiffness", "fixed_extended", wp.float32), - ("fixed_tendon_rest_length", "fixed_extended", wp.float32), - ("fixed_tendon_offset", "fixed_extended", wp.float32), - ("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), -] -_TENDON_PROPERTY_FAMILY = {name: family for name, family, _ in _TENDON_DATA_PROPERTIES} - - -@pytest.mark.parametrize(_TENDON_DIM_NAMES, _TENDON_DIMS) -@pytest.mark.parametrize("family", _TENDON_FAMILIES) -def test_articulation_tendon_contract( - monkeypatch, - backend, - num_instances, - num_joints, - num_bodies, - num_fixed_tendons, - num_spatial_tendons, - device, - family, -): - kind, capability = _TENDON_FAMILIES[family] - require_backend_capability(backend, capability) - art, _ = get_articulation( - backend, - num_instances, - num_joints, - num_bodies, - num_fixed_tendons, - num_spatial_tendons, - device, - monkeypatch=monkeypatch, - ) - num_tendons = num_fixed_tendons if kind == "fixed" else num_spatial_tendons - assert getattr(art, f"num_{kind}_tendons") == num_tendons - names = getattr(art, f"{kind}_tendon_names") - assert isinstance(names, list) - assert len(names) == num_tendons - assert all(isinstance(n, str) for n in names) - - art.data.update(dt=0.01) - for name, property_family, dtype in _TENDON_DATA_PROPERTIES: - if property_family != family or num_tendons == 0: - continue - expected_shape = (num_instances, num_tendons) - 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] -# Writer cases pair each method with the dimension rows that configure at least one tendon of its kind. -_TENDON_WRITER_CASES = [ - pytest.param(*method, *dims, id=f"{method[1]}-{'-'.join(map(str, dims))}") - for method in _TENDON_METHODS - for dims in _TENDON_DIMS - if dims[3 if method[0] == "fixed" else 4] > 0 -] - - -@pytest.mark.parametrize("selection", ["index", "mask"]) -@pytest.mark.parametrize( - ("kind", "method_base", "kwarg", "getter", *_TENDON_DIM_NAMES), - _TENDON_WRITER_CASES, -) -def test_tendon_writer( - monkeypatch, - backend, - device, - selection, - kind, - method_base, - kwarg, - getter, - num_instances, - num_joints, - num_bodies, - num_fixed_tendons, - num_spatial_tendons, -): - require_backend_capability(backend, _TENDON_FAMILIES[_TENDON_PROPERTY_FAMILY[getter]][1]) - if backend == "newton": - from isaaclab_newton.physics import NewtonManager - - # The mocked asset has no solver to receive model-change notifications. - monkeypatch.setattr(NewtonManager, "add_model_change", MagicMock()) - - num_tendons = num_fixed_tendons if kind == "fixed" else num_spatial_tendons - art, _ = get_articulation( - backend, - num_instances, - num_joints, - num_bodies, - num_fixed_tendons, - num_spatial_tendons, - device, - monkeypatch=monkeypatch, - ) - 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: after the push, the matching data property reads the written values back - method(**{kwarg: _make_payload_warp(full_shape, device, wp_dtype)}) - getattr(art, f"write_{kind}_tendon_properties_to_sim_index")() - _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}) - - -# Shape validation runs before any kernel launch, so one device covers it. - - -@pytest.mark.parametrize("array_kind", ["torch", "warp"]) -@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_rejects_extra_environments( - monkeypatch, request, backend, selection, array_kind, kind, method_base, kwarg, getter -): - require_backend_capability(backend, _TENDON_FAMILIES[_TENDON_PROPERTY_FAMILY[getter]][1]) - if backend == "newton" and selection == "mask": - # Product bug: Newton's fixed-tendon mask setters slice full data without validating its shape, so an - # oversized input is silently truncated. - request.applymarker( - pytest.mark.xfail( - raises=pytest.fail.Exception, - strict=True, - reason="Newton fixed-tendon mask setters accept data with extra environments", - ) - ) - num_instances, num_joints, num_bodies, num_fixed_tendons, num_spatial_tendons = _TENDON_DIMS[1] - art, _ = get_articulation( - backend, - num_instances, - num_joints, - num_bodies, - num_fixed_tendons, - num_spatial_tendons, - "cpu", - monkeypatch=monkeypatch, - ) - full_shape = (num_instances, num_fixed_tendons if kind == "fixed" else num_spatial_tendons) - wp_dtype = wp.vec2f if getter == "fixed_tendon_pos_limits" else wp.float32 - make_bad_data = _make_bad_data_torch if array_kind == "torch" else _make_bad_data_warp - - with pytest.raises((AssertionError, RuntimeError)): - getattr(art, f"{method_base}_{selection}")(**{kwarg: make_bad_data(full_shape, "cpu", wp_dtype)}) + 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": - from isaaclab_newton.physics import NewtonManager - - # The mocked asset has no solver to receive model-change notifications. - monkeypatch.setattr(NewtonManager, "add_model_change", MagicMock()) - - require_backend_capability(backend, _TENDON_FAMILIES[kind][1]) + 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 = get_articulation(backend, num_instances, 2, 2, 2, 2, device, monkeypatch=monkeypatch) + 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": @@ -1597,15 +737,11 @@ def capture(name, values, *, indices=None, mask=None): # 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": - env_ids = None if len(selected_envs) == num_instances else torch.tensor(selected_envs, device=device) - write_kwargs = {"env_ids": env_ids} + write(env_ids=torch.tensor(selected_envs, device=device) if partial else None) 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_kwargs = {"env_mask": mask} - write(**write_kwargs) + 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])]) diff --git a/source/isaaclab/test/assets/contract/test_articulation_ordering_contract.py b/source/isaaclab/test/assets/contract/test_articulation_ordering_contract.py index 4c955bbc5fa0..02f76ce76107 100644 --- a/source/isaaclab/test/assets/contract/test_articulation_ordering_contract.py +++ b/source/isaaclab/test/assets/contract/test_articulation_ordering_contract.py @@ -5,7 +5,12 @@ # pyright: reportPrivateUsage=none -"""Mocked cross-backend articulation joint- and body-ordering contract cases.""" +"""Cross-backend articulation joint- and body-ordering contracts on mocked backends. + +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. +""" from unittest.mock import MagicMock, patch @@ -13,651 +18,133 @@ import pytest import torch import warp as wp -from _pytest.mark.structures import ParameterSet from isaaclab.utils import replace -from isaaclab.utils.buffers import TimestampedBuffer from .articulation_factory import get_articulation -from .capabilities import backend_parameters, requires_backend -from .mock_backends import read_backend_joint_state +from .backends import backends, read_backend_joint_state, requires pytestmark = pytest.mark.integration -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) +@pytest.fixture(autouse=True) +def _cpu_warp_device(): + """Every ordering case runs on CPU; allocate unqualified Warp arrays there too.""" + with wp.ScopedDevice("cpu"): + yield - 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 +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- -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 _names(prefix: str, count: int) -> tuple[str, ...]: + return tuple(f"{prefix}_{index}" for index in range(count)) -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: +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.""" - if ordering is None: - return np.arange(item_count, dtype=np.int64) - return np.asarray(ordering.user_to_backend_indices, dtype=np.int64) + return np.arange(count) if ordering is None else 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: +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), 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 + return np.arange(len(backend_names)) + return np.asarray([public_names.index(name) for name in backend_names]) - 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 _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, *, monkeypatch: pytest.MonkeyPatch -) -> 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, - monkeypatch=monkeypatch, - ) - 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, *, monkeypatch: pytest.MonkeyPatch) -> 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), - monkeypatch=monkeypatch, +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), ) - 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) - + # 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 _exercise_duplicate_full_length_joint_write( - backend: str, operation: str, *, monkeypatch: pytest.MonkeyPatch -) -> 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), - monkeypatch=monkeypatch, - ) - 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) +def _mask(total: int, selected, device: str) -> wp.array: + mask = np.zeros(total, dtype=bool) + mask[list(selected)] = True + return wp.array(mask, dtype=wp.bool, device=device) - 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 _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 _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 _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 _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 _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_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.""" +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.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._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 - 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) + 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) @@ -668,1336 +155,742 @@ def _set_body_ordering_backend_data( 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}") + art.data._previous_body_com_vel.zero_() -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 _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 _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.""" +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": - 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 + raw_backend._coms = wp.array(coms, dtype=wp.float32, device="cpu") 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) - + from isaaclab_ov import tensor_types as TT -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}") + raw_backend.bindings[TT.BODY_COM_POSE]._data = coms.copy() + _invalidate_com_caches(art) + return coms -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.""" +def _read_backend_coms(backend: str, art, raw_backend) -> np.ndarray: if backend == "physx": - return ( - raw_backend._root_transforms.numpy().copy(), - art.data.root_com_vel_w.warp.numpy().copy(), - ) - if backend == "ovphysx": + coms = raw_backend.get_coms() + else: 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}") - + coms = art.root_view.get_attribute(TT.BODY_COM_POSE) + return coms.numpy().reshape(art.num_instances, art.num_bodies, 7).copy() -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 _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 _get_backend_joint_property_tensors(backend: str, art, raw_backend) -> dict[str, torch.Tensor]: - """Return backend-order joint property values for ordering parity checks.""" +def _backend_properties(backend: str, art, raw_backend) -> dict[str, torch.Tensor]: + """Return backend-order joint and body properties.""" 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], + 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(), } - if backend == "ovphysx": + elif 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], + 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, } - if backend == "newton": + 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": _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), + "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), } - raise AssertionError(f"Unsupported backend for joint-property ordering test: {backend}") + 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 _get_backend_body_property_tensors(backend: str, art, raw_backend) -> dict[str, torch.Tensor]: - """Return backend-order body property values for ordering parity checks.""" +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": - 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 an ordering parameter row that carries the backend's declared skip reason.""" - marks = [*kwargs.pop("marks", ()), *backend_parameters("ordering", names=(backend,))[0].marks] - return pytest.param(backend, *values, marks=marks, **kwargs) - - -def test_reversed_body_ordering_reorders_public_body_quantities(monkeypatch, backend): - num_instances = 2 - num_joints = 1 - num_bodies = 3 - device = "cpu" - identity_art, identity_raw = get_articulation( - backend, num_instances, num_joints, num_bodies, device=device, monkeypatch=monkeypatch - ) - ordered_art, ordered_raw = get_articulation( - backend, num_instances, num_joints, num_bodies, device=device, monkeypatch=monkeypatch - ) - 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) - - -@pytest.mark.parametrize("ordering_mode", ["reversed", "cyclic"]) -def test_body_ordering_reorders_public_body_properties(monkeypatch, backend, ordering_mode): - """Expose mass and every inertia component under the matching public body name.""" - num_instances = 2 - num_joints = 1 - num_bodies = 3 - device = "cpu" - 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, - monkeypatch=monkeypatch, - ) - 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_backend("ovphysx") -def test_ovphysx_reversed_body_ordering_rereads_backend_shadows_after_reset(monkeypatch): - """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", monkeypatch=monkeypatch - ) - 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) + 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() - _ = _clone_proxy_tensor(art.data.body_link_pose_w) - _ = _clone_proxy_tensor(art.data.body_com_vel_w) +def _set_ovphysx_joint_state(raw_backend, position: np.ndarray, velocity: np.ndarray) -> None: 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) - + raw_backend.bindings[TT.DOF_POSITION]._data = position.copy() + raw_backend.bindings[TT.DOF_VELOCITY]._data = velocity.copy() -@requires_backend("ovphysx") -def test_ovphysx_reversed_body_ordering_rereads_all_velocity_shadows_after_reset(monkeypatch): - """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"), monkeypatch=monkeypatch - ) - 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() +_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 - 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_backend("ovphysx") -def test_ovphysx_ordered_root_and_body_velocity_share_one_backend_read(monkeypatch) -> None: - """Read the shared backend link-velocity binding once per timestamp.""" +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( - "ovphysx", 2, 1, 3, device="cpu", body_ordering=("body_2", "body_1", "body_0"), monkeypatch=monkeypatch - ) - 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 - - -@pytest.mark.parametrize("backend", backend_parameters("ordering", names=("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 = 2 - num_joints = 3 - num_bodies = 4 - device = "cpu" - identity_art, identity_raw = get_articulation( - backend, - num_instances, - num_joints, - num_bodies, - device=device, - is_fixed_base=is_fixed_base, - monkeypatch=monkeypatch, + backend, 2, 3, 2, device="cpu", joint_ordering=_ordering(mode, _names("joint", 3)), monkeypatch=monkeypatch ) - ordered_art, ordered_raw = get_articulation( - backend, - num_instances, - num_joints, - num_bodies, - device=device, - is_fixed_base=is_fixed_base, - monkeypatch=monkeypatch, - ) - 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) + _seed_joint_state(backend, art, raw_backend, _INITIAL_POSITION, _INITIAL_VELOCITY) + return art, raw_backend - 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] +# --------------------------------------------------------------------------- +# Ordering resolution and public reads +# --------------------------------------------------------------------------- - _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 - - -_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": (), -} - - -@pytest.mark.parametrize("ordering_mode", ["none", "identity", "reversed"]) -def test_backends_allocate_shadows_only_for_nonidentity_ordering(monkeypatch, backend: str, ordering_mode: str) -> None: - num_instances = 2 - num_joints = 3 - num_bodies = 4 +@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, - num_instances, - num_joints, - num_bodies, + 2, + 3, + 4, device="cpu", - joint_ordering=_joint_ordering_for_mode(ordering_mode, num_joints), - body_ordering=_body_ordering_for_mode(ordering_mode, num_bodies), + joint_ordering=joint_ordering, + body_ordering=body_ordering, monkeypatch=monkeypatch, ) - 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, _DATA_SHADOWS[backend]), - (art, _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) - - -@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_backend("ovphysx") -def test_ovphysx_partial_com_write_preserves_backend_rows( - monkeypatch, 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), - monkeypatch=monkeypatch, - ) - 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_backend("newton") -def test_newton_partial_com_write_routes_to_backend_bodies(monkeypatch, 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), - monkeypatch=monkeypatch, - ) - 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)] + 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( - 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]) - - -@pytest.mark.parametrize("backend", backend_parameters("ordering", names=("physx", "ovphysx"))) -@pytest.mark.parametrize("ordering_mode", ["none", "reversed"]) -def test_duplicate_body_ids_preserve_omitted_backend_com(monkeypatch, 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), - monkeypatch=monkeypatch, - ) - 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), + art.joint_ordering.backend_to_user_indices, _backend_to_user(joint_ordering, joint_names) ) - 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]) - - -@pytest.mark.parametrize("backend", backend_parameters("ordering", names=("physx", "ovphysx"))) -@pytest.mark.parametrize("ordering_mode", ["none", "reversed"]) -def test_duplicate_env_ids_leave_global_com_cache_invalid(monkeypatch, 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), - monkeypatch=monkeypatch, + art.body_ordering.backend_to_user_indices, _backend_to_user(body_ordering, body_names) ) - 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 - - -@pytest.mark.parametrize("backend", backend_parameters("ordering", names=("physx", "ovphysx"))) -@pytest.mark.parametrize("ordering_mode", ["none", "reversed"]) -def test_static_com_cache_does_not_follow_sim_timestamp(monkeypatch, 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), - monkeypatch=monkeypatch, - ) - 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 +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) -# --------------------------------------------------------------------------- -# Tests: ArticulationData joint state and properties -# --------------------------------------------------------------------------- + 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("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( - monkeypatch, 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) +@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_instances, - num_joints=1, - num_bodies=num_bodies, + num_instances, + num_joints, + num_bodies, device="cpu", - is_fixed_base=False, - body_ordering=backend_body_names, + 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_instances, - num_joints=1, - num_bodies=num_bodies, - device="cpu", - is_fixed_base=False, - body_ordering=ordered_body_names, - monkeypatch=monkeypatch, - ) - # 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) - - -# 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( - monkeypatch, backend: str, ordering_mode: str, selection: str, operation: str, coverage: str -) -> None: - _exercise_partial_joint_write(backend, ordering_mode, selection, operation, coverage, monkeypatch=monkeypatch) - - -@pytest.mark.parametrize( - "operation, expected_position_reads, expected_velocity_reads", - [("position", 1, 0), ("velocity", 0, 1), ("state", 1, 1)], -) -@requires_backend("physx") -def test_physx_mask_partial_joint_write_reads_each_staging_axis_once( - monkeypatch, 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, + backend, + num_instances, + num_joints, + num_bodies, device="cpu", - joint_ordering=_joint_ordering_for_mode("reversed", 3), + is_fixed_base=is_fixed_base, monkeypatch=monkeypatch, ) - 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, - ) + 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 - 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) + 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) -@pytest.mark.parametrize("backend", backend_parameters("ordering", names=("physx", "ovphysx"))) -@pytest.mark.parametrize("operation", ["position", "velocity", "state"]) -def test_duplicate_full_length_joint_selector_preserves_newer_backend_rows( - monkeypatch, backend: str, operation: str -) -> None: - _exercise_duplicate_full_length_joint_write(backend, operation, monkeypatch=monkeypatch) +@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 -@pytest.mark.parametrize("selector", ["default_mask", "explicit_full_index"]) -@pytest.mark.parametrize("operation", ["position", "velocity", "state"]) -@requires_backend("physx") -def test_physx_joint_selector_provenance_controls_staging_read(monkeypatch, operation: str, selector: str) -> None: art, raw_backend = get_articulation( - "physx", - num_instances=2, - num_joints=3, - num_bodies=2, + "ovphysx", + 2, + 3, + 3, device="cpu", - joint_ordering=_joint_ordering_for_mode("reversed", 3), + joint_ordering=_ordering("cyclic", _names("joint", 3)), + body_ordering=_ordering("cyclic", _names("body", 3)), monkeypatch=monkeypatch, ) - 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")) - + 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() -# 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. -@pytest.mark.parametrize("backend", backend_parameters("ordering", names=("physx", "ovphysx"))) -@pytest.mark.parametrize("operation", ["position", "velocity", "state"]) -def test_stale_partial_joint_write_preserves_newer_backend_rows(monkeypatch, backend: str, operation: str) -> None: - _exercise_stale_partial_joint_write(backend, operation, monkeypatch=monkeypatch) + _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_backend("ovphysx") -def test_ovphysx_partial_effort_target_write_preserves_unselected_backend_rows(monkeypatch) -> 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. - """ +@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 - num_instances = 2 - num_joints = 3 art, raw_backend = get_articulation( "ovphysx", - num_instances=num_instances, - num_joints=num_joints, - num_bodies=2, + 2, + 3, + 3, device="cpu", - joint_ordering=_joint_ordering_for_mode("reversed", num_joints), + joint_ordering=("joint_2", "joint_1", "joint_0"), + body_ordering=("body_2", "body_1", "body_0"), monkeypatch=monkeypatch, ) - 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) + _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) - expected_targets = raw_targets.copy() - expected_targets[:, 1] = new_target.cpu().numpy()[:, 0] - expected_backend = expected_targets[:, user_to_backend] + 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() - pushed_after_partial_write = raw_backend.bindings[TT.DOF_ACTUATION_FORCE]._data.copy() - np.testing.assert_array_equal(pushed_after_partial_write, expected_backend) + 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_backend("ovphysx") +@requires("ovphysx") def test_ovphysx_joint_acceleration_differences_public_order_velocities(monkeypatch): - """Finite-difference OVPhysX joint velocity entirely in public joint order.""" - art, raw_backend = get_articulation("ovphysx", 2, 3, 2, device="cpu", monkeypatch=monkeypatch) - user_to_backend = _install_test_joint_ordering(art) - from isaaclab_ov import tensor_types as TT - + 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) - raw_backend.bindings[TT.DOF_VELOCITY]._data = first + _set_ovphysx_joint_state(raw_backend, _INITIAL_POSITION, first) art.data.update(0.1) art.data.joint_acc.torch.clone() - raw_backend.bindings[TT.DOF_VELOCITY]._data = second + _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), - ) + torch.testing.assert_close(art.data.joint_acc.torch, torch.from_numpy((second - first)[:, user_to_backend] / 0.1)) -@requires_backend("ovphysx") -def test_ovphysx_ordered_joint_state_is_cached_per_sim_timestamp(monkeypatch): - """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"), monkeypatch=monkeypatch - ) - 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 +# --------------------------------------------------------------------------- +# Property writes +# --------------------------------------------------------------------------- - for _ in range(2): - art.data.joint_pos.torch.clone() - art.data.joint_vel.torch.clone() +# (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 - 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 +@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 -def test_joint_ordering_reorders_public_joint_properties(monkeypatch, backend): - """Expose every backend joint property under the matching public joint name.""" - num_instances = 2 - num_joints = 3 - num_bodies = 2 - device = "cpu" - joint_ordering = _joint_ordering_for_mode("cyclic", num_joints) + monkeypatch.setattr(NewtonManager, "add_model_change", MagicMock()) art, raw_backend = get_articulation( backend, num_instances, - num_joints, - num_bodies, - device=device, - joint_ordering=joint_ordering, + count, + count, + device="cpu", + joint_ordering=_ordering(mode, _names("joint", count)), + body_ordering=_ordering(mode, _names("body", count)), monkeypatch=monkeypatch, ) - 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, - } + 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)] ) - if backend in ("physx", "ovphysx"): - public_properties["dynamic_friction"] = art.data.joint_dynamic_friction_coeff + 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 +# --------------------------------------------------------------------------- - for property_name, public_property in public_properties.items(): - _assert_proxy_close(public_property, backend_properties[property_name][:, user_to_backend]) +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 - # 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]) - ) + 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() -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) +@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) # --------------------------------------------------------------------------- -# Tests: Articulation operations +# Joint state writes # --------------------------------------------------------------------------- -def test_legacy_position_target_accepts_joint_slice(monkeypatch, backend: str) -> None: - art, _ = get_articulation( - backend, num_instances=2, num_joints=4, num_bodies=2, device="cpu", monkeypatch=monkeypatch - ) - target = torch.tensor([[11.0, 12.0], [21.0, 22.0]], dtype=torch.float32) +# 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]) - 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) +@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)) -@pytest.mark.parametrize("ordering_mode", ["none", "reversed", "cyclic"]) -@pytest.mark.parametrize("is_fixed_base", [False, True], ids=["floating", "fixed"]) -def test_external_wrenches_are_written_in_backend_body_order(monkeypatch, backend, ordering_mode, is_fixed_base): - """Write public-order body wrenches to each backend in backend body order.""" - device = "cpu" - 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]) + # 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: - body_ordering = _body_ordering_for_mode(ordering_mode, num_bodies) + 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, - num_joints, + 1, num_bodies, - device=device, + device="cpu", is_fixed_base=is_fixed_base, body_ordering=body_ordering, monkeypatch=monkeypatch, ) - _set_rotated_body_poses(backend, art, raw_backend) + # 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): - 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() + 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 = _expected_backend_to_user(body_ordering, backend_body_names) + 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, device=device), - torques=wp.array(torques, dtype=wp.vec3f, device=device), - is_global=is_global, + 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() @@ -2005,363 +898,166 @@ def capture_wrench(*, force_data, torque_data, position_data, indices, is_global 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 - 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]) + 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) -@requires_backend("ovphysx") -def test_ovphysx_configured_defaults_use_public_joint_names(monkeypatch): - """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"), monkeypatch=monkeypatch - ) - 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.write_data_to_sim() - art._resolve_joint_values(patterns, art.data._default_joint_pos) + 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]) - 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") +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. -@requires_backend("ovphysx") -@pytest.mark.parametrize("ordering_mode", ["reversed", "cyclic"]) -def test_ovphysx_implicit_targets_are_written_in_backend_order(monkeypatch, ordering_mode: str) -> None: - """Write implicit position and velocity targets under their matching backend joint names.""" + 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 - 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, - monkeypatch=monkeypatch, - ) - 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) + 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() - - 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], + 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_backend("ovphysx") +@requires("ovphysx") @pytest.mark.parametrize("selector_kind", ["torch", "warp"]) -def test_ovphysx_int64_effort_target_selector_reaches_binding( - selector_kind: str, monkeypatch: pytest.MonkeyPatch -) -> None: - """Convert indexed effort-target environment selectors for the OVPhysX binding.""" +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() - captured_indices = None set_attribute = art._root_view.set_attribute + captured = [] 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 + 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 = 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) + 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=target, env_ids=env_ids, joint_ids=joint_ids) + art.set_joint_effort_target_index(target=torch.tensor([[7.0]]), env_ids=env_ids, joint_ids=torch.tensor([2])) - expected[1] = 0.0 - expected[1, 2] = 7.0 - assert captured_indices is not None - np.testing.assert_array_equal(captured_indices.numpy(), [1]) + 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_backend("physx") -@pytest.mark.parametrize("ordering_mode", ["reversed", "cyclic"]) -def test_physx_newton_actuator_forces_are_written_in_backend_order(monkeypatch, 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) +@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", - num_instances, - num_joints, - num_bodies, + 2, + 4, + 2, device="cpu", - joint_ordering=joint_ordering, + joint_ordering=_ordering("cyclic", _names("joint", 4)), monkeypatch=monkeypatch, ) - 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) + 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)) - wrapper = MagicMock() - wrapper.joint_f_2d = user_forces - object.__setattr__(art, "_physx_actuator_wrapper", wrapper) + 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) - 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 + written = [] + raw_backend.set_dof_actuation_forces = lambda values, indices: written.append(values.numpy().copy()) 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]) + np.testing.assert_allclose(written[-1], forces[:, _backend_to_user(tuple(art.joint_names), _names("joint", 4))]) -@requires_backend("physx") +@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) - 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") - ) + 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() - - -@pytest.mark.parametrize("ordering_mode", ["reversed", "cyclic"]) -@pytest.mark.parametrize("selection", ["index", "mask"]) -def test_reversed_joint_ordering_routes_property_writes_to_backend( - monkeypatch, 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, - monkeypatch=monkeypatch, - ) - 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], - ) - - -@pytest.mark.parametrize("ordering_mode", ["reversed", "cyclic"]) -@pytest.mark.parametrize("selection", ["index", "mask"]) -def test_body_ordering_routes_property_writes_to_backend(monkeypatch, 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, - monkeypatch=monkeypatch, - ) - 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/contract/test_asset_contract_api.py b/source/isaaclab/test/assets/contract/test_asset_contract_api.py deleted file mode 100644 index bbac11d7379e..000000000000 --- a/source/isaaclab/test/assets/contract/test_asset_contract_api.py +++ /dev/null @@ -1,101 +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 - -"""Shared asset API contract: counts, names, finders, index resolution, and the classified public surface.""" - -import inspect -from importlib.util import find_spec - -import pytest - -from .articulation_factory import get_articulation -from .capabilities import BACKEND_STATUSES, available_backends -from .public_surface import BASE_SURFACE_CLASSES, PUBLIC_SURFACE_CONTRACTS, public_surface_mismatches -from .rigid_object_collection_factory import get_rigid_object_collection -from .rigid_object_factory import get_rigid_object - -pytestmark = pytest.mark.integration - -_MISSING = object() - - -def test_installed_backend_packages_run_the_contract() -> None: - """An installed backend package must not silently drop out of the contract matrix.""" - unavailable = { - status.declaration.name: status.reason - for status in BACKEND_STATUSES - if not status.available - and status.reason != "CUDA runtime unavailable" - and find_spec(status.declaration.required_modules[-1]) is not None - } - - assert unavailable == {} - - -def _manager_bindings() -> dict[str, object]: - """Return the backend-manager bindings that the mocked factories replace during a test.""" - bindings = {} - backends = available_backends() - if "physx" in backends: - from isaaclab_physx.physics import PhysxManager - - for name in ("get_physics_sim_view", "_scene_data_backend"): - bindings[f"physx.{name}"] = inspect.getattr_static(PhysxManager, name, _MISSING) - if "ovphysx" in backends: - from isaaclab_ov.physics.ovphysx_manager import OvPhysxManager - - bindings["ovphysx._scene_data_backend"] = inspect.getattr_static( - OvPhysxManager, "_scene_data_backend", _MISSING - ) - if "newton" in backends: - import isaaclab_newton.assets.articulation.articulation_data as articulation_data - import isaaclab_newton.assets.rigid_object.rigid_object_data as rigid_object_data - import isaaclab_newton.assets.rigid_object_collection.rigid_object_collection as collection - import isaaclab_newton.assets.rigid_object_collection.rigid_object_collection_data as collection_data - - for module in (articulation_data, rigid_object_data, collection, collection_data): - bindings[f"newton.{module.__name__}"] = module.SimulationManager - return bindings - - -@pytest.mark.parametrize("reverse", [False, True], ids=["forward", "reverse"]) -def test_contract_factories_scope_manager_patches_to_one_test(reverse: bool) -> None: - """Factories patch backend managers only inside a test scope and restore the production bindings after it.""" - original = _manager_bindings() - # Importing the contract modules must not patch a manager: production bindings precede every factory call. - if "physx" in available_backends(): - assert isinstance(original["physx.get_physics_sim_view"], classmethod) - if "newton" in available_backends(): - from isaaclab_newton.physics import NewtonManager - - assert all(manager is NewtonManager for name, manager in original.items() if name.startswith("newton.")) - - backends = available_backends()[::-1] if reverse else available_backends() - patches = pytest.MonkeyPatch() - try: - for backend in backends: - rigid_object, _ = get_rigid_object(backend, device="cpu", monkeypatch=patches) - collection, _ = get_rigid_object_collection(backend, device="cpu", monkeypatch=patches) - articulation, _ = get_articulation(backend, device="cpu", monkeypatch=patches) - # Later factories must not break the patched managers the earlier families read through. - assert rigid_object.data.root_link_pose_w.shape == (2,) - assert collection.data.body_link_pose_w.shape == (2, 3) - assert articulation.data.root_link_pose_w.shape == (2,) - - finally: - patches.undo() - - restored = _manager_bindings() - assert restored.keys() == original.keys() - assert all(restored[name] is binding for name, binding in original.items()) - - -def test_base_public_surface_has_an_explicit_contract_classification() -> None: - """Every public member of the asset base classes is classified, and every classified member still exists.""" - assert public_surface_mismatches(BASE_SURFACE_CLASSES, PUBLIC_SURFACE_CONTRACTS) == { - "missing": [], - "stale": [], - "unreasoned": [], - } diff --git a/source/isaaclab/test/assets/contract/test_rigid_object_collection_contract.py b/source/isaaclab/test/assets/contract/test_rigid_object_collection_contract.py deleted file mode 100644 index 762ce672fb87..000000000000 --- a/source/isaaclab/test/assets/contract/test_rigid_object_collection_contract.py +++ /dev/null @@ -1,867 +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 rigid-object-collection contracts across the production backends. - -Checks that every collection backend provides the data and writer behavior the base collection class advertises. -The backends run on mocked views, so these cases need neither Isaac Sim nor a 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_factory import get_rigid_object_collection - -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. - - -# --------------------------------------------------------------------------- -# 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 -# --------------------------------------------------------------------------- - - -def test_resolve_env_ids_handles_tensor_view_shape(monkeypatch, backend, device): - obj, _ = get_rigid_object_collection(backend, num_instances=4, device=device, monkeypatch=monkeypatch) - - 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() - - -def test_resolve_body_ids_handles_tensor_view_shape(monkeypatch, backend): - obj, _ = get_rigid_object_collection(backend, num_bodies=4, device="cpu", monkeypatch=monkeypatch) - - 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 -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) -def test_reshape_data_to_view_3d_accepts_torch_tensor(monkeypatch, 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, monkeypatch=monkeypatch - ) - 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) - - -@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(monkeypatch, 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", monkeypatch=monkeypatch - ) - 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 -# --------------------------------------------------------------------------- - - -def test_collection_counts_and_names(monkeypatch, 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", monkeypatch=monkeypatch) - - 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) - - -def test_find_bodies_returns_legacy_tensor_or_cached_proxy(monkeypatch, backend): - collection, _ = get_rigid_object_collection( - backend, num_instances=2, num_bodies=3, device="cpu", monkeypatch=monkeypatch - ) - - 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] - - -def test_find_objects_forwards_return_mode_with_alias_warning(monkeypatch, backend): - collection, _ = get_rigid_object_collection( - backend, num_instances=2, num_bodies=3, device="cpu", monkeypatch=monkeypatch - ) - - 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), -] - - -def test_collection_data_property_contract(monkeypatch, backend, device): - obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device, monkeypatch=monkeypatch) - 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", -} - - -def test_pose_write_invalidates_pose_dependent_caches(monkeypatch, backend): - obj, _ = get_rigid_object_collection(backend, num_instances=2, num_bodies=3, device="cpu", monkeypatch=monkeypatch) - 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) - - -def test_velocity_write_invalidates_body_frame_caches(monkeypatch, backend): - obj, _ = get_rigid_object_collection(backend, num_instances=2, num_bodies=3, device="cpu", monkeypatch=monkeypatch) - 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) - - -@pytest.mark.parametrize("setter_kind", ["index", "mask"]) -def test_set_coms_invalidates_same_timestamp_dependents(monkeypatch, backend, setter_kind): - obj, _ = get_rigid_object_collection(backend, num_instances=2, num_bodies=3, device="cpu", monkeypatch=monkeypatch) - 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) - - -# -- index variants for pose -- - - -@pytest.mark.parametrize("method_suffix", _BODY_POSE_METHODS) -def test_write_body_pose_to_sim_index(monkeypatch, backend, device, method_suffix): - obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device, monkeypatch=monkeypatch) - 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 -- - - -@pytest.mark.parametrize("method_suffix", _BODY_VEL_METHODS) -def test_write_body_velocity_to_sim_index(monkeypatch, backend, device, method_suffix): - obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device, monkeypatch=monkeypatch) - 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 -- - - -@pytest.mark.parametrize("method_suffix", _BODY_POSE_METHODS) -def test_write_body_pose_to_sim_mask(monkeypatch, backend, device, method_suffix): - obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device, monkeypatch=monkeypatch) - 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 -- - - -@pytest.mark.parametrize("method_suffix", _BODY_VEL_METHODS) -def test_write_body_velocity_to_sim_mask(monkeypatch, backend, device, method_suffix): - obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device, monkeypatch=monkeypatch) - 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) - - -@pytest.mark.parametrize("method_base, kwarg", _BODY_METHODS, ids=[m[0] for m in _BODY_METHODS]) -def test_body_writer_index(monkeypatch, backend, device, method_base, kwarg): - obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device, monkeypatch=monkeypatch) - 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)}) - - -@pytest.mark.parametrize("method_base, kwarg", _BODY_METHODS, ids=[m[0] for m in _BODY_METHODS]) -def test_body_writer_mask(monkeypatch, backend, device, method_base, kwarg): - obj, _ = get_rigid_object_collection(backend, _NUM_INSTANCES, _NUM_BODIES, device, monkeypatch=monkeypatch) - 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"), -] - - -def test_aliases_match_canonical_values(monkeypatch, 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", monkeypatch=monkeypatch) - 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 - - -# --------------------------------------------------------------------------- -# Tests: partial writes reach only the selected backend cells -# --------------------------------------------------------------------------- - -# quantity -> (writer base name, writer keyword, trailing size) -_PARTIAL_WRITES = { - "body_pose": ("write_body_link_pose_to_sim", "body_poses", 7), - "body_velocity": ("write_body_com_velocity_to_sim", "body_velocities", 6), - "mass": ("set_masses", "masses", 1), -} - - -def _read_backend_cells(backend: str, raw_backend, quantity: str, num_instances: int, num_bodies: int) -> torch.Tensor: - """Read one collection quantity from backend storage in public ``(env, body, ...)`` order.""" - if backend == "physx": - getter = { - "body_pose": raw_backend.get_transforms, - "body_velocity": raw_backend.get_velocities, - "mass": raw_backend.get_masses, - } - values = wp.to_torch(getter[quantity]()) - # 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_backend.get_attribute("body_mass", None)) - else: - getter = {"body_pose": raw_backend.get_root_transforms, "body_velocity": raw_backend.get_root_velocities} - values = wp.to_torch(getter[quantity](None)) - else: - from isaaclab_ov import tensor_types as TT - - binding = {"body_pose": TT.LINK_POSE, "body_velocity": TT.LINK_VELOCITY, "mass": TT.BODY_MASS}[quantity] - values = torch.as_tensor(raw_backend.bindings[binding]._data) - return values.reshape(num_instances, num_bodies, -1).cpu().clone() - - -@pytest.mark.parametrize("selection", ["index", "mask"]) -@pytest.mark.parametrize("quantity", _PARTIAL_WRITES) -def test_partial_write_preserves_unselected_backend_cells(monkeypatch, request, backend, selection, quantity): - if backend == "ovphysx" 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", - ) - ) - num_instances, num_bodies = _NUM_INSTANCES, _NUM_BODIES - obj, raw_backend = get_rigid_object_collection(backend, num_instances, num_bodies, "cpu", monkeypatch=monkeypatch) - writer, kwarg, trailing = _PARTIAL_WRITES[quantity] - # Literal, per-cell distinct payload; poses keep an identity rotation. - values = 100.0 * torch.arange(1, num_instances * num_bodies + 1, dtype=torch.float32).reshape( - num_instances, num_bodies, 1 - ) + torch.arange(trailing, dtype=torch.float32) - if quantity == "body_pose": - values[..., 3:6] = 0.0 - values[..., 6] = 1.0 - if quantity == "mass": - values = values[..., 0] - before = _read_backend_cells(backend, raw_backend, quantity, num_instances, num_bodies) - # Select the second environment and the last body so identity or transposed routing cannot pass. - env, body = 1, 2 - if selection == "index": - getattr(obj, f"{writer}_index")( - **{kwarg: values[env : env + 1, body : body + 1]}, - env_ids=torch.tensor([env], dtype=torch.int32), - body_ids=_make_body_ids("cpu", [body]), - ) - else: - getattr(obj, f"{writer}_mask")( - **{kwarg: values}, - env_mask=_make_item_mask(num_instances, [env], "cpu"), - body_mask=_make_item_mask(num_bodies, [body], "cpu"), - ) - - after = _read_backend_cells(backend, raw_backend, quantity, num_instances, num_bodies) - expected = before.clone() - expected[env, body] = values[env, body].reshape(-1) - torch.testing.assert_close(after, expected, rtol=0.0, atol=0.0) diff --git a/source/isaaclab/test/assets/contract/test_rigid_object_contract.py b/source/isaaclab/test/assets/contract/test_rigid_object_contract.py index ad1f5a7d8562..07922735d759 100644 --- a/source/isaaclab/test/assets/contract/test_rigid_object_contract.py +++ b/source/isaaclab/test/assets/contract/test_rigid_object_contract.py @@ -6,13 +6,14 @@ # ignore private usage of variables warning # pyright: reportPrivateUsage=none -"""Shared rigid-object contracts across the production backends. +"""Shared rigid-object and rigid-object-collection contracts across the production backends. -Checks that every rigid-object backend provides the data and writer behavior the base rigid-object class advertises. -The backends run on mocked views, so these cases need neither Isaac Sim nor a GPU simulation. +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. """ -import math from unittest.mock import MagicMock, patch import numpy as np @@ -20,596 +21,382 @@ import torch import warp as wp -from .rigid_object_factory import get_rigid_object - -pytestmark = pytest.mark.integration - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- +from isaaclab.assets import BaseRigidObjectCollectionData, BaseRigidObjectData +from isaaclab.utils.warp import ProxyArray +from .rigid_object_factory import get_rigid_object, get_rigid_object_collection -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 +pytestmark = pytest.mark.integration - 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}" +# Distinct instance and body counts make swapped axes visible. +_NUM_INSTANCES, _NUM_BODIES = 2, 3 +_ASSETS = ["object", "collection"] -# 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. -_NUM_INSTANCES = 2 +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 -# --------------------------------------------------------------------------- -# Tests: Index resolution helpers -# --------------------------------------------------------------------------- - +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 test_resolve_env_ids_handles_tensor_view_shape(monkeypatch, backend, device): - obj, _ = get_rigid_object(backend, num_instances=4, device=device, monkeypatch=monkeypatch) - 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]) +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``. - 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() + 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 # --------------------------------------------------------------------------- -# Tests: RigidObject properties and finders +# Bookkeeping, finders, and data properties # --------------------------------------------------------------------------- -def test_rigid_object_counts_and_names(monkeypatch, backend): - from isaaclab.assets.rigid_object.base_rigid_object_data import BaseRigidObjectData - - obj, _ = get_rigid_object(backend, _NUM_INSTANCES, device="cpu", monkeypatch=monkeypatch) - - 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) - +@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) -def test_find_bodies_returns_legacy_list_or_cached_proxy(monkeypatch, backend): - obj, _ = get_rigid_object(backend, num_instances=2, device="cpu", monkeypatch=monkeypatch) + 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) 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 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 - assert str(proxy.device) == obj.device + 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() - first_body = obj.body_names[0] - assert obj.find_bodies(first_body) == ([0], [first_body]) +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 -# --------------------------------------------------------------------------- -# 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), -] - - -def test_rigid_object_data_property_contract(monkeypatch, backend, device): - obj, _ = get_rigid_object(backend, _NUM_INSTANCES, device, monkeypatch=monkeypatch) - obj.data.update(dt=0.01) - if backend == "newton": +@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(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() + 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 -# --------------------------------------------------------------------------- -# 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"), -] - - -def test_aliases_match_canonical_values(monkeypatch, backend): # Random mock state makes link and COM quantities differ, so a retargeted alias fails. - obj, _ = get_rigid_object(backend, _NUM_INSTANCES, device="cpu", monkeypatch=monkeypatch) - 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 + 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 - -# --------------------------------------------------------------------------- -# 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) + # 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: - 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) + 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) -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: + for name, buffer in buffers.items(): 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 +# Writers # --------------------------------------------------------------------------- -# 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", +_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 test_pose_write_invalidates_pose_dependent_caches(monkeypatch, backend): - obj, _ = get_rigid_object(backend, num_instances=2, device="cpu", monkeypatch=monkeypatch) - 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) - - -def test_velocity_write_invalidates_body_frame_caches(monkeypatch, backend): - obj, _ = get_rigid_object(backend, num_instances=2, device="cpu", monkeypatch=monkeypatch) - 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 +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("setter_kind", ["index", "mask"]) -def test_set_coms_invalidates_same_timestamp_dependents(monkeypatch, backend, setter_kind): - obj, _ = get_rigid_object(backend, num_instances=2, device="cpu", monkeypatch=monkeypatch) +@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) - 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) + 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: - obj.set_coms_mask(coms=coms) - - if backend == "newton": - from isaaclab_newton.physics import NewtonManager - - with patch.object(NewtonManager, "add_model_change"): - set_coms() + values = wp.to_torch({"pose": raw.get_root_transforms, "velocity": raw.get_root_velocities}[quantity](None)) else: - set_coms() - _assert_buffers_stale(obj.data, buffers) - - -# -- index variants -- - - -@pytest.mark.parametrize("method_suffix", _ROOT_POSE_METHODS) -def test_write_root_pose_to_sim_index(monkeypatch, backend, device, method_suffix): - num_instances = _NUM_INSTANCES - obj, _ = get_rigid_object(backend, num_instances, device, monkeypatch=monkeypatch) - 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)) - - -@pytest.mark.parametrize("method_suffix", _ROOT_VEL_METHODS) -def test_write_root_velocity_to_sim_index(monkeypatch, backend, device, method_suffix): - num_instances = _NUM_INSTANCES - obj, _ = get_rigid_object(backend, num_instances, device, monkeypatch=monkeypatch) - 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)) - + from isaaclab_ov import tensor_types as TT -# -- mask variants -- + 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("method_suffix", _ROOT_POSE_METHODS) -def test_write_root_pose_to_sim_mask(monkeypatch, backend, device, method_suffix): - num_instances = _NUM_INSTANCES - obj, _ = get_rigid_object(backend, num_instances, device, monkeypatch=monkeypatch) - 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)) - - -@pytest.mark.parametrize("method_suffix", _ROOT_VEL_METHODS) -def test_write_root_velocity_to_sim_mask(monkeypatch, backend, device, method_suffix): - num_instances = _NUM_INSTANCES - obj, _ = get_rigid_object(backend, num_instances, device, monkeypatch=monkeypatch) - 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)) +@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) # --------------------------------------------------------------------------- -# Tests: Body writers — torch/warp × index/mask × all/subset × negative +# Family-specific behavior # --------------------------------------------------------------------------- -_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) - -def test_external_wrench_frames(monkeypatch, backend): +def test_rigid_object_external_wrench_frames(monkeypatch, 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, monkeypatch=monkeypatch) - # Seed a known 90-degree rotation about Z so the local-frame expectation is independent of production. + 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]], - device=device, + [[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 FK would publish. + # 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, device=device).reshape(2, 1, 3) + forces = torch.arange(1.0, 7.0).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() @@ -620,11 +407,10 @@ def test_external_wrench_frames(monkeypatch, backend): expected_force, expected_torque = forces, torques if backend == "physx": - call = raw_backend.apply_forces_and_torques_at_position.call_args.kwargs + 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 = call["force_data"].numpy().reshape(2, 1, 3) - actual_torque = call["torque_data"].numpy().reshape(2, 1, 3) + 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). @@ -635,156 +421,31 @@ def test_external_wrench_frames(monkeypatch, backend): else: from isaaclab_ov import tensor_types as TT - packed = raw_backend.bindings[TT.RIGID_BODY_WRENCH]._data.reshape(2, 1, 9) + 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, 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) + 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("method_base, kwarg", _BODY_METHODS, ids=[m[0] for m in _BODY_METHODS]) -def test_body_writer_index(monkeypatch, backend, device, method_base, kwarg): - num_instances, num_bodies = _NUM_INSTANCES, 1 - obj, _ = get_rigid_object(backend, num_instances, device, monkeypatch=monkeypatch) - 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)}) - 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)}) - - -@pytest.mark.parametrize("method_base, kwarg", _BODY_METHODS, ids=[m[0] for m in _BODY_METHODS]) -def test_body_writer_mask(monkeypatch, backend, device, method_base, kwarg): - num_instances, num_bodies = _NUM_INSTANCES, 1 - obj, _ = get_rigid_object(backend, num_instances, device, monkeypatch=monkeypatch) - 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), - } +@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 ) - # 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: partial writes reach only the selected backend rows -# --------------------------------------------------------------------------- - -# quantity -> (writer base name, writer keyword, literal per-environment payload) -_PARTIAL_WRITES = { - "root_pose": ( - "write_root_link_pose_to_sim", - "root_pose", - [[31.0, 32.0, 33.0, 0.0, 0.0, 0.0, 1.0], [41.0, 42.0, 43.0, 0.0, 0.0, 0.0, 1.0]], - ), - "root_velocity": ( - "write_root_com_velocity_to_sim", - "root_velocity", - [[21.0, 22.0, 23.0, 24.0, 25.0, 26.0], [71.0, 72.0, 73.0, 74.0, 75.0, 76.0]], - ), - "mass": ("set_masses", "masses", [[51.0], [61.0]]), -} - - -def _read_backend_rows(backend: str, raw_backend, quantity: str) -> torch.Tensor: - """Read one rigid-object quantity from backend storage as one row per environment.""" - if backend == "physx": - getter = { - "root_pose": raw_backend.get_transforms, - "root_velocity": raw_backend.get_velocities, - "mass": raw_backend.get_masses, - } - values = getter[quantity]() - elif backend == "newton": - if quantity == "mass": - values = raw_backend.get_attribute("body_mass", None) - else: - getter = {"root_pose": raw_backend.get_root_transforms, "root_velocity": raw_backend.get_root_velocities} - values = getter[quantity](None) - else: - from isaaclab_ov import tensor_types as TT - - binding = {"root_pose": TT.RIGID_BODY_POSE, "root_velocity": TT.RIGID_BODY_VELOCITY, "mass": TT.RIGID_BODY_MASS} - values = raw_backend.bindings[binding[quantity]]._data - values = wp.to_torch(values) if isinstance(values, wp.array) else torch.as_tensor(values) - return values.reshape(2, -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): - obj, raw_backend = get_rigid_object(backend, num_instances=2, device="cpu", monkeypatch=monkeypatch) - writer, kwarg, payload = _PARTIAL_WRITES[quantity] - values = torch.tensor(payload, dtype=torch.float32) - before = _read_backend_rows(backend, raw_backend, quantity) - # Select the second environment so a writer that ignores the selection cannot pass by writing row 0. - selected, unselected = 1, 0 - if selection == "index": - kwargs = {kwarg: values[selected : selected + 1], "env_ids": torch.tensor([selected], dtype=torch.int32)} - if quantity == "mass": - kwargs["body_ids"] = [0] - else: - kwargs = {kwarg: values, "env_mask": _make_item_mask(2, [selected], "cpu")} - if quantity == "mass": - kwargs["body_mask"] = _make_item_mask(1, [0], "cpu") - - getattr(obj, f"{writer}_{selection}")(**kwargs) - - after = _read_backend_rows(backend, raw_backend, quantity) - torch.testing.assert_close(after[selected], values[selected], rtol=0.0, atol=0.0) - torch.testing.assert_close(after[unselected], before[unselected], rtol=0.0, atol=0.0) + 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()) From b72ab09f49fb3934503ec7ea7b18c94770b9839d Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Fri, 2 Oct 2026 15:20:41 -0700 Subject: [PATCH 23/33] Trim Newton asset and controller tests --- .../test/assets/test_articulation.py | 46 ----- .../assets/test_newton_actuators_newton.py | 191 ++++++------------ .../test/assets/test_rigid_object.py | 189 ++++------------- .../assets/test_rigid_object_collection.py | 161 ++++----------- .../unit/test_articulation_accessors.py | 14 +- .../test_newton_task_space_controllers.py | 16 +- 6 files changed, 154 insertions(+), 463 deletions(-) diff --git a/source/isaaclab_newton/test/assets/test_articulation.py b/source/isaaclab_newton/test/assets/test_articulation.py index 87fd7fde3f0e..74d4af973ffb 100644 --- a/source/isaaclab_newton/test/assets/test_articulation.py +++ b/source/isaaclab_newton/test/assets/test_articulation.py @@ -161,12 +161,6 @@ def island( "wrist": ImplicitActuatorCfg(joint_names_expr=["Joint_[3-5]"], stiffness=None, damping=None), }, ), - "passive": island( - "Passive", - local_usd("fixed_spatial_chain.usda"), - 4.0, - {"all": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=0.0, damping=0.0)}, - ), "pendulum": island( "Pendulum", local_usd("revolute_pendulum.usda"), @@ -705,7 +699,6 @@ def test_articulation_initialization_and_partial_state(scene: _Scene) -> None: expected_fixed_base = { "floating": False, "fixed": True, - "passive": True, "pendulum": True, "ordered": False, "tendon": True, @@ -1828,45 +1821,6 @@ def test_get_gravity_compensation_forces_matches_jacobian_gravity(scene: _Scene, scene.rest(articulation) -@pytest.mark.isaacsim_ci -def test_get_gravity_compensation_forces_static_equilibrium(scene: _Scene) -> None: - """Newton accuracy: ``τ_gc`` must hold a passive 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 chain should not move. The island's actuators have zero gains, so the effort - target IS the joint torque applied and no PD spring-damper masks the gravity-compensation signal. Sign - errors, frame errors, and DoF-ordering errors all surface as joint drift. - """ - device = scene.device - articulation = scene.articulations["passive"] - scene.rest(articulation) - # a configuration that loads the prismatic and the revolute joints with gravity - q = articulation.data.default_joint_pos.torch.clone() - q[:, 3:] = torch.tensor([0.4, -0.3, 0.5], device=device) - articulation.write_joint_position_to_sim_index(position=q) - articulation.update(scene.sim.cfg.dt) - init_q = articulation.data.joint_pos.torch.clone() - - with world_gravity(_GRAVITY): - tau_gc = articulation.data.gravity_compensation_forces.torch[:, articulation.num_base_dofs :] - assert tau_gc.abs().max() > 1.0, "the configuration carries no gravity load" - for _ in range(100): - # ``gravity_compensation_forces`` leads with ``num_base_dofs`` floating-base entries (0 on fixed-base); - # slice past them so the remaining tensor aligns with the actuated joints. - tau_gc = articulation.data.gravity_compensation_forces.torch[:, articulation.num_base_dofs :] - articulation.actuators.target_command.set_effort_index(value=tau_gc) - scene.step("passive") - final_q = articulation.data.joint_pos.torch.clone() - - drift = (final_q - init_q).abs().max() - # A sign or frame bug in τ_gc produces drift of at least a degree per step on this configuration. - assert drift < 5e-3, ( - f"max joint drift {drift:.5f} 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." - ) - 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. 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 9820fc174911..683e6746f7a8 100644 --- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py +++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py @@ -365,6 +365,7 @@ def _run( 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, @@ -511,51 +512,29 @@ def test_dc_motor_clamp_binds(lab_run: dict, newton_run: _Run) -> None: ), "the DC-motor clamp never bound, so the equivalence cannot detect missing clamping" -def test_newton_actuators_are_delayed_pd(newton_run: _Run) -> None: - """Command delays are authored as ``NewtonActuatorDelayAPI`` on a PD drive.""" - actuator_info = newton_run.results["delayed"]["actuator_info"] - assert actuator_info, "No Newton actuators were created" - for a in actuator_info: - assert a["has_delay"], "Delay not found on delayed PD actuator" - assert a["controller_type"] == "DrivePD" - - -def test_newton_knee_actuators_are_delayed_position_clamped_pd(newton_run: _Run) -> None: - """RemotizedPD knees are authored as a delayed PD drive with position-based clamping.""" - kfe_acts = [ - a for a in newton_run.results["remotized"]["actuator_info"] if "ClampingPositionBased" in a["clamping_types"] - ] - assert len(kfe_acts) > 0, "No actuator with position-based clamping found" - for a in kfe_acts: - assert a["controller_type"] == "DrivePD" - assert a["has_delay"], "Delay not found on remotized KFE actuator" - - -# --------------------------------------------------------------------------- -# Neural network actuator authoring: MLP and LSTM -# --------------------------------------------------------------------------- - - -def test_mlp_has_dc_motor_clamping(newton_run: _Run) -> None: - """ActuatorNetMLPCfg is authored as a Newton neural MLP controller with DC motor clamping.""" - mlp_acts = [a for a in newton_run.results["neural"]["actuator_info"] if a["controller_type"] == "DriveNeuralMLP"] - assert len(mlp_acts) > 0, "No NeuralMLP controller found" - for a in mlp_acts: - assert "ClampingDCMotor" in a["clamping_types"] - - -def test_lstm_has_dc_motor_clamping(newton_run: _Run) -> None: - """ActuatorNetLSTMCfg is authored as a Newton neural LSTM controller with DC motor clamping.""" - lstm_acts = [a for a in newton_run.results["neural"]["actuator_info"] if a["controller_type"] == "DriveNeuralLSTM"] - assert len(lstm_acts) > 0, "No NeuralLSTM controller found" - for a in lstm_acts: - assert "ClampingDCMotor" in a["clamping_types"] - - -def test_positions_finite(newton_run: _Run) -> None: - """The neural actuators run on the Newton backend without producing non-finite joint positions.""" - for step_i, pos in enumerate(newton_run.results["neural"]["joint_pos"]): - assert torch.isfinite(pos).all(), f"Non-finite positions at step {step_i}" +@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"]) # --------------------------------------------------------------------------- @@ -640,98 +619,46 @@ def gather_stiffness() -> torch.Tensor: def test_randomize_actuator_gains_reaches_newton_controllers(newton_run: _Run) -> None: - """``randomize_actuator_gains`` writes kp/kd into the controllers of the articulation's Newton actuators. + """``randomize_actuator_gains`` writes kp/kd into the selected environment of one articulation's controllers. - The event writes through ``write_group_parameter``; the assertions read the controllers back via the public - ``read_group_parameter``. 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 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``. """ - anymal = newton_run.articulations["ideal"] - cartpole = newton_run.articulations["cartpole"] + groups = { + "legs": (newton_run.articulations["ideal"], "legs"), + "cartpole": (newton_run.articulations["cartpole"], "all_joints"), + } + legs = groups["legs"][0] assert SimulationManager._adapter is not None - anymal_stiffness_before = read_group_parameter(anymal.actuators, "legs", "controller", "kp").clone() - anymal_damping_before = read_group_parameter(anymal.actuators, "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 (6 free-root DOFs + the leg 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 = read_group_parameter(cartpole.actuators, "all_joints", "controller", "kp").clone() - cartpole_damping_before = read_group_parameter(cartpole.actuators, "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( - read_group_parameter(cartpole.actuators, "all_joints", "controller", "kp")[0], - torch.full((n_cp,), 100.0, device=anymal.device), - ) - torch.testing.assert_close( - read_group_parameter(cartpole.actuators, "all_joints", "controller", "kd")[0], - torch.full((n_cp,), 5.0, device=anymal.device), - ) - - # The leg is untouched (DR was scoped to cartpole). - torch.testing.assert_close( - read_group_parameter(anymal.actuators, "legs", "controller", "kp"), anymal_stiffness_before - ) - torch.testing.assert_close( - read_group_parameter(anymal.actuators, "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( - read_group_parameter(cartpole.actuators, "all_joints", "controller", "kp")[env_idx], - cartpole_stiffness_before[env_idx], - ) - torch.testing.assert_close( - read_group_parameter(cartpole.actuators, "all_joints", "controller", "kd")[env_idx], - cartpole_damping_before[env_idx], - ) - - # DR scoped to the floating-base leg 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( - read_group_parameter(anymal.actuators, "legs", "controller", "kp")[0], - torch.full((n,), 100.0, device=anymal.device), - ) - torch.testing.assert_close( - read_group_parameter(anymal.actuators, "legs", "controller", "kd")[0], - torch.full((n,), 5.0, device=anymal.device), - ) - for env_idx in range(1, NUM_ENVS): - torch.testing.assert_close( - read_group_parameter(anymal.actuators, "legs", "controller", "kp")[env_idx], - anymal_stiffness_before[env_idx], - ) - torch.testing.assert_close( - read_group_parameter(anymal.actuators, "legs", "controller", "kd")[env_idx], anymal_damping_before[env_idx] + 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", ) + # only environment 0 of the selected articulation changes + for other in groups: + torch.testing.assert_close(gains(other), expected[other]) # --------------------------------------------------------------------------- diff --git a/source/isaaclab_newton/test/assets/test_rigid_object.py b/source/isaaclab_newton/test/assets/test_rigid_object.py index c1353fa34385..ffb69b1f2026 100644 --- a/source/isaaclab_newton/test/assets/test_rigid_object.py +++ b/source/isaaclab_newton/test/assets/test_rigid_object.py @@ -6,7 +6,7 @@ # ignore private usage of variables warning # pyright: reportPrivateUsage=none -"""Kitless real-Newton rigid-object coverage on one persistent two-environment scene per device. +"""Kitless real-Newton rigid-object coverage on one persistent two-environment CPU scene. The cubes are locally spawned shapes. 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 case build their own scenes. Only @@ -39,7 +39,7 @@ from isaaclab.envs.mdp.events import randomize_rigid_body_material from isaaclab.managers import EventTermCfg, SceneEntityCfg from isaaclab.sim import SimulationContext, build_simulation_context -from isaaclab.test.utils import test_devices +from isaaclab.test.utils import DeviceScope, test_devices from isaaclab.utils.math import ( combine_frame_transforms, quat_apply_inverse, @@ -100,7 +100,7 @@ def test_initialization_rejects_non_rigid_body_prims(api: str) -> None: sim.reset() -@pytest.mark.parametrize("device", test_devices()) +@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: @@ -161,9 +161,9 @@ def rest(self) -> None: self.cube.reset() -@pytest.fixture(scope="module", params=test_devices()) +@pytest.fixture(scope="module", params=test_devices(DeviceScope.CPU)) def shared_scene(request: pytest.FixtureRequest) -> Iterator[_Scene]: - """Initialize the cubes once per device for this module.""" + """Initialize the cubes once for this module.""" device = request.param with build_simulation_context(sim_cfg=newton_sim_cfg(device)) as sim: cube = spawn_assets({"cube": _cube_cfg()})["cube"] @@ -191,123 +191,37 @@ def scene(shared_scene: _Scene) -> _Scene: @pytest.mark.isaacsim_ci def test_external_force_on_single_body(scene: _Scene) -> None: - """Test initialization and external forces on the base of the object. + """A permanent wrench reaches the solver in the local and the global frame, and a reset clears it. - 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. + 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_object = scene.cube - num_cubes = cube_object.num_instances - device = scene.device - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(cube_object) < 10 - - # 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(): - scene.rest() - - # 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 + 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): - # Sample a force equal to the weight of the object - external_wrench_b = torch.zeros(cube_object.num_instances, len(body_ids), 6, device=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] - - for i in range(2): - reset_cubes() - - 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, - ) - scene.step(5) - - # 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=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) - - # 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=device) - external_wrench_positions_b = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=device) - external_wrench_b[0::2, :, 2] = 50.0 - external_wrench_positions_b[0::2, :, 1] = 1.0 - - 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, - ) - scene.step(5) - - # 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) + 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 @@ -485,12 +399,12 @@ def test_body_root_state_properties(scene: _Scene) -> None: # 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 about z, fast enough for the link-frame velocity to discriminate spin_twist = torch.zeros(6, device=device) - spin_twist[5] = 0.5 * torch.randn(1, device=device).clamp(-1.0, 1.0) + spin_twist[5] = 2.0 # Simulate physics - for _ in range(100): + 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() @@ -684,45 +598,28 @@ def _fk_reset_mask_dirty() -> bool: # Step once so that _sim_timestamp > 0 and caches are primed. scene.step() - for writer in ("link_index", "link_mask", "com_index", "com_mask"): + 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 = wp.to_torch(cube_object.data.body_link_pose_w).clone().view(num_cubes, 7) + 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() - # 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) + # 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 = 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. + 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" ) - # 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)) diff --git a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py index fe96326b9314..bb1f955ce3fc 100644 --- a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py @@ -6,7 +6,7 @@ # ignore private usage of variables warning # pyright: reportPrivateUsage=none -"""Kitless real-Newton rigid-object-collection coverage on one persistent scene per device. +"""Kitless real-Newton rigid-object-collection coverage on one persistent CPU scene. Every environment holds three locally spawned cubes in the collection and an unrelated rigid body beside them, so the collection view must select exactly its configured bodies. Scene gravity is off; a test that needs gravity @@ -37,7 +37,7 @@ import isaaclab.sim as sim_utils from isaaclab.assets import RigidObjectCfg, RigidObjectCollectionCfg from isaaclab.sim import SimulationContext, build_simulation_context -from isaaclab.test.utils import test_devices +from isaaclab.test.utils import DeviceScope, test_devices from isaaclab.utils.math import ( combine_frame_transforms, quat_apply_inverse, @@ -106,7 +106,7 @@ def test_initialization_with_no_rigid_body() -> None: sim.reset() -@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) def test_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: @@ -169,9 +169,9 @@ def rest(self) -> None: self.sibling.reset() -@pytest.fixture(scope="module", params=test_devices()) +@pytest.fixture(scope="module", params=test_devices(DeviceScope.CPU)) def shared_scene(request: pytest.FixtureRequest) -> Iterator[_Scene]: - """Initialize the collection and its unrelated sibling bodies once per device for this module.""" + """Initialize the collection and its unrelated sibling bodies once for this module.""" device = request.param sibling_cfg = RigidObjectCfg( prim_path="/World/Env_[^/]*/Object_Target", @@ -284,107 +284,37 @@ def test_set_body_inertial_properties_updates_inverses(scene: _Scene) -> None: def test_external_force_on_single_body(scene: _Scene) -> None: - """Test application of external force on the base of the object. + """A permanent wrench reaches the solver in the local and the global frame, and a reset clears it. - 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. + 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. """ - object_collection = scene.collection - device = scene.device - - # find objects to apply the force - object_ids, object_names = object_collection.find_bodies(".*") - - def reset_objects(): - scene.rest() - - # 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 + 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): - # Sample a force equal to the weight of the object - external_wrench_b = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=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, - ) - scene.step(10) - - # 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=device) - external_wrench_positions_b = torch.zeros(object_collection.num_instances, len(object_ids), 3, device=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, - ) - scene.step(10) - - # 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) + 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 @@ -457,9 +387,9 @@ def test_object_state_properties(scene: _Scene) -> None: # 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 about z, fast enough for the link-frame velocity to discriminate spin_twist = torch.zeros(6, device=device) - spin_twist[5] = torch.randn(1, device=device) + spin_twist[5] = 2.0 # initial spawn point init_com = cube_object.data.body_com_pose_w.torch[..., :3] @@ -627,33 +557,22 @@ def _fk_reset_mask_dirty() -> bool: cube_object = scene.collection scene.step() - for writer in ("link_index", "link_mask", "com_index", "com_mask"): + 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 = 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) + 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 = wp.to_torch(cube_object.data.body_link_pose_w) + 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" ) diff --git a/source/isaaclab_newton/test/assets/unit/test_articulation_accessors.py b/source/isaaclab_newton/test/assets/unit/test_articulation_accessors.py index 7e5e105f3e81..609a5fa1bf23 100644 --- a/source/isaaclab_newton/test/assets/unit/test_articulation_accessors.py +++ b/source/isaaclab_newton/test/assets/unit/test_articulation_accessors.py @@ -22,7 +22,7 @@ from newton import JointType, ModelBuilder from newton.selection import ArticulationView -from isaaclab.test.utils import test_devices +from isaaclab.test.utils import DeviceScope, test_devices pytestmark = pytest.mark.unit @@ -46,7 +46,7 @@ class _ShapeCountSurface: assert articulation.num_shapes_per_body == [3, 0, 2] -@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) @pytest.mark.parametrize( "first_property", ["body_com_jacobian_w", "body_link_jacobian_w", "mass_matrix", "gravity_compensation_forces"], @@ -82,10 +82,7 @@ def test_task_space_allocation_and_capture(monkeypatch, device, first_property): 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: + with wp.ScopedCapture(device=device) as capture: output = getattr(data, first_property) required = {first_property} if first_property == "body_link_jacobian_w": @@ -100,10 +97,7 @@ def test_task_space_allocation_and_capture(monkeypatch, device, first_property): 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 + 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() diff --git a/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py b/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py index d6d25c43f9f3..28621c0ea1c5 100644 --- a/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py +++ b/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py @@ -177,10 +177,10 @@ def test_osc_tracking_accuracy(chain: _Chain) -> None: chain.rest() ee_pos_b, ee_quat_b = _compute_ee_pose_root(chain.robot, chain.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(chain, _make_osc(chain.device), target_pose_b, 300, gravity=False) + pos_history, rot_history = _run_osc(chain, _make_osc(chain.device), target_pose_b, 150, gravity=False) - pos_mean = sum(pos_history[-200:]) / 200 - rot_mean = sum(rot_history[-200:]) / 200 + 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 @@ -213,13 +213,13 @@ def test_osc_gravity_compensation_precision(chain: _Chain) -> None: target_pose_b = torch.cat([ee_pos_b, ee_quat_b], dim=-1) def _stationary_tail_mean(history: list[float], label: str) -> float: - """Mean of the last 200 samples, asserting the two tail halves agree within 25%. + """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" @@ -227,9 +227,9 @@ def _stationary_tail_mean(history: list[float], label: str) -> float: return mean with world_gravity((0.0, 0.0, -9.81)): - hist_off, _ = _run_osc(chain, osc, target_pose_b, 300, gravity=True) + hist_off, _ = _run_osc(chain, osc, target_pose_b, 200, gravity=True) osc.cfg.gravity_compensation = True - hist_on, _ = _run_osc(chain, osc, target_pose_b, 300, gravity=True) + hist_on, _ = _run_osc(chain, 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") From a5aecd6b6d478280582f4979f9ec7fa0f82db1bc Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Fri, 2 Oct 2026 15:20:41 -0700 Subject: [PATCH 24/33] Trim PhysX asset tests --- .../test/assets/test_articulation.py | 395 +++----- .../assets/test_newton_actuators_physx.py | 895 ++++++------------ .../test/assets/test_rigid_object.py | 177 ++-- .../assets/test_rigid_object_collection.py | 176 ++-- .../assets/unit/test_physx_articulation.py | 108 --- 5 files changed, 546 insertions(+), 1205 deletions(-) diff --git a/source/isaaclab_physx/test/assets/test_articulation.py b/source/isaaclab_physx/test/assets/test_articulation.py index 98cee3031a34..2e025bd791b2 100644 --- a/source/isaaclab_physx/test/assets/test_articulation.py +++ b/source/isaaclab_physx/test/assets/test_articulation.py @@ -22,7 +22,7 @@ BRANCHING_PHYSX_BODY_NAMES, BRANCHING_PHYSX_JOINT_NAMES, ) -from isaaclab.utils import clone, replace +from isaaclab.utils import replace launch_test_simulation() @@ -43,22 +43,8 @@ import isaaclab.sim as sim_utils from isaaclab.actuators import ImplicitActuatorCfg from isaaclab.assets import ArticulationCfg, get_articulation_name_ordering -from isaaclab.controllers import DifferentialIKController, DifferentialIKControllerCfg from isaaclab.sim import SimulationContext, build_simulation_context -from isaaclab.utils.math import ( - combine_frame_transforms, - compute_pose_error, - matrix_from_quat, - quat_apply, - quat_inv, - subtract_frame_transforms, -) - -## -# Pre-defined configs -## -from isaaclab_assets import FRANKA_PANDA_CFG, FRANKA_PANDA_HIGH_PD_CFG # isort:skip -from isaaclab_assets.robots.shadow_hand import SHADOW_HAND_PHYSX_CFG # isort:skip +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 @@ -110,6 +96,7 @@ def _author_branching_robot( 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. @@ -120,6 +107,7 @@ def _author_branching_robot( 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. """ @@ -173,6 +161,17 @@ def _author_branching_robot( 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: @@ -186,15 +185,11 @@ def _yaw_quat(angle: float) -> tuple[float, float, float, float]: @pytest.fixture -def gravity_enabled() -> bool: - """Enable gravity unless a test explicitly parametrizes it otherwise.""" - return True - - -@pytest.fixture -def sim(device: str, gravity_enabled: bool) -> Iterator[SimulationContext]: +def sim(device: str) -> Iterator[SimulationContext]: """Create a function-scoped simulation context for tests that own their scene.""" - with build_simulation_context(device=device, auto_add_lighting=True, gravity_enabled=gravity_enabled) as sim: + # 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 @@ -204,27 +199,6 @@ def sim(device: str, gravity_enabled: bool) -> Iterator[SimulationContext]: ## -def generate_articulation( - articulation_cfg: ArticulationCfg, num_articulations: int, device: str -) -> tuple[Articulation, torch.Tensor]: - """Spawn ``num_articulations`` copies of an articulation 2.5 m apart along x. - - 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. - """ - translations = torch.zeros(num_articulations, 3, device=device) - translations[:, 0] = torch.arange(num_articulations) * 2.5 - 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 - - @pytest.mark.parametrize( ("init_state", "message"), [ @@ -260,254 +234,44 @@ def test_setting_invalid_articulation_root_prim_path(sim, device) -> None: @pytest.mark.parametrize("device", test_devices()) -def test_fixed_tendon_position_target_writes_offset(sim, device) -> None: - """A tendon length target lands in the simulation as ``rest_length - target`` on the selected cells only. - - 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. The Shadow Hand is the shipped asset whose - native PhysX fixed tendons these writers command. - """ - num_articulations = 2 - articulation_cfg = SHADOW_HAND_PHYSX_CFG - 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 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) - +def test_gravity_compensation_holds_static_equilibrium(sim, device) -> None: + """Gravity compensation efforts hold the arms still and follow mass writes. -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.isaacsim_ci -def test_get_gravity_compensation_forces_static_equilibrium(sim, device) -> None: - """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). + 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. """ - num_articulations = 1 - # 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( - FRANKA_PANDA_CFG, actuators={"all": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=0.0, damping=0.0)} + 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))), + ) ) - # 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) + _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 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. + 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)) + 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) + 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) - 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." - ) - - -# --------------------------------------------------------------------------- -# Franka task-space tracking helpers for the IK test. -# Mirrors the helpers in ``isaaclab_newton/test/assets/test_articulation.py``. -# --------------------------------------------------------------------------- - - -def _setup_franka_at_home_pose(sim): - """Build a Franka articulation at its configured home pose. - - 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. - - Returns: - Tuple of ``(robot, ee_frame_idx, ee_jacobi_idx, arm_joint_ids)``. - """ - 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) - 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_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 - - -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 - - -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 _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) - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@pytest.mark.parametrize("gravity_enabled", [False]) -@pytest.mark.isaacsim_ci -def test_franka_ik_tracking_accuracy(sim, device, gravity_enabled) -> None: - """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?" + # A mass write refreshes efforts already read in this step. GPU PhysX applies mass writes on the next step + # (see ``test_inertial_writes_refresh_the_mass_matrix``). + if device == "cpu": + primed = articulation.data.gravity_compensation_forces.torch.clone() + tips = articulation.find_bodies(".*_tip")[0] + articulation.set_masses_index(masses=2.0 * articulation.data.body_mass.torch[:, tips], body_ids=tips) + assert not torch.allclose(articulation.data.gravity_compensation_forces.torch, primed, atol=1e-3) ## @@ -528,7 +292,7 @@ class _ArticulationScene: 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 and a spatial tendon.""" + """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] @@ -592,7 +356,10 @@ def articulation_scene(request) -> Iterator[_ArticulationScene]: ), {"root_on_base": True}, ), - "tendon": (_branching_cfg(), {"fixed_base": True, "reversed_left_elbow": True, "spatial_tendon": 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) @@ -639,6 +406,8 @@ def test_articulation_initialization_and_ordering(articulation_scene: _Articulat 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 @@ -1294,8 +1063,8 @@ def test_floating_articulation_root_writes_and_wrenches(articulation_scene: _Art reordered.data.body_link_lin_vel_w.torch[:, floating_order], floating.data.body_link_lin_vel_w.torch ) - # A partial reset clears the selected environment; a full reset resets every actuator environment and - # clears all external forces and torques. + # 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 = [] @@ -1310,6 +1079,7 @@ def record_actuator_reset(env_ids=None) -> None: 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): @@ -1381,6 +1151,55 @@ def test_spatial_tendon_properties_round_trip(articulation_scene: _ArticulationS 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. + """ + device = articulation_scene.device + articulation = articulation_scene.tendon + assert articulation.num_fixed_tendons == 1 + rest_length = articulation.data.fixed_tendon_rest_length.torch.clone() + 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() + 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) + + 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) + + +def test_inertial_writes_refresh_the_mass_matrix(articulation_scene: _ArticulationScene, request) -> None: + """Mass and inertia writes refresh a mass matrix that was already read in the same step.""" + if articulation_scene.device.startswith("cuda"): + request.applymarker( + pytest.mark.xfail(strict=True, reason="GPU PhysX applies mass and inertia writes on the next 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 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 0f2710abe926..26203c59d9d2 100644 --- a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py +++ b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py @@ -3,24 +3,26 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Newton-native actuator seams on the PhysX backend, driven on ANYmal-C (floating-base quadruped). - -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. 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, per-environment resets, gain randomization, and network actuators. +"""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 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, and network actuators. """ -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 +import pytest import torch import warp as wp from isaaclab_physx.assets import Articulation @@ -28,16 +30,16 @@ import isaaclab.sim as sim_utils from isaaclab.actuators import IdealPDActuatorCfg +from isaaclab.actuators.actuator_net_cfg import ActuatorNetLSTMCfg, ActuatorNetMLPCfg from isaaclab.actuators.newton import read_group_parameter -from isaaclab.sim import SimulationCfg, build_simulation_context +from isaaclab.assets import ArticulationCfg +from isaaclab.sim import SimulationCfg, SimulationContext, build_simulation_context 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, make_dummy_lstm_checkpoint, @@ -45,16 +47,11 @@ ) from isaaclab.test.utils.articulation_ordering import assert_articulation_ordering_trace_matches -from isaaclab_assets import ANYMAL_C_CFG - -# --------------------------------------------------------------------------- -# Constants -# --------------------------------------------------------------------------- +from isaaclab_assets import ANYMAL_C_CFG, CARTPOLE_CFG 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", @@ -69,620 +66,320 @@ "RF_KFE", "RH_KFE", ) +# 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), +} + + +_ROWS = {"delayed": 0, "mixed": 1, "ideal": 2, "reversed": 2, "cartpole": 3, "implicit": 4, "mlp": 5, "lstm": 6} +"""Row of each robot along y. Robots compared across rollouts share a row: ground contacts amplify the floating-point +differences between placements at different world positions.""" +_LAB_ROBOTS = ("delayed", "mixed", "ideal", "cartpole") +"""Robots of the Lab-actuator rollout, which the Newton-actuator rollout reproduces.""" +_NEWTON_ROBOTS = (*_LAB_ROBOTS, "implicit", "mlp", "lstm") +"""Robots of the Newton-actuator rollout.""" -# --------------------------------------------------------------------------- -# Simulation runner -# --------------------------------------------------------------------------- +def _robot_cfgs(names: tuple[str, ...], networks: dict[str, str]) -> dict[str, ArticulationCfg]: + """Return the configurations of the named robots. + Args: + names: Robot names. ``reversed`` is the ``ideal`` robot with a reversed 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 + ) + anymal_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, + }, + } + cfgs = {} + for name in names: + if name == "cartpole": + cfgs[name] = replace(CARTPOLE_CFG, actuators=CARTPOLE_EXPLICIT_ACTUATORS) + elif name == "reversed": + joint_ordering = tuple(reversed(_ANYMAL_C_PHYSX_JOINT_NAMES)) + cfgs[name] = replace(ANYMAL_C_CFG, actuators=IDEAL_PD_ACTUATORS, joint_ordering=joint_ordering) + else: + cfgs[name] = replace(ANYMAL_C_CFG, actuators=anymal_actuators[name]) + return cfgs -def _run_simulation( - actuators: dict, - use_newton_actuators: bool, - *, - num_steps: int = NUM_STEPS, - joint_ordering: tuple[str, ...] | None = None, - permutation_sensitive_commands: bool = False, -) -> dict: - """Run ANYmal-C on PhysX and return recorded trajectories + telemetry. - Always records public joint state/telemetry and the Newton adapter outputs. +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: - 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. - joint_ordering: Optional explicit public joint-name order. - permutation_sensitive_commands: Whether to command distinct position, velocity, and effort values by - physical joint name. + sim: Simulation context the robots are spawned in. + cfgs: Robot configurations keyed by name. Returns: - Recorded joint-name metadata, commands, public trajectories and torque telemetry, and adapter effort traces. + The articulations and their recorded traces, both keyed by robot name. """ - 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 + 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( + replace( + cfg, + prim_path=f"/World/Env_[^/]*/{name}", + init_state=replace(cfg.init_state, pos=(0.0, 3.0 * _ROWS[name], cfg.init_state.pos[2])), + ) ) - articulation = Articulation(art_cfg) - sim.reset() - assert articulation.is_initialized + 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) - installed_ordering = articulation.joint_ordering - joint_ordering_state = ( - None - if installed_ordering is None + 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": 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 - - 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 = [] - for _ in range(num_steps): + "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() - 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()) + # 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: - 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_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() + 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 -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 -# --------------------------------------------------------------------------- - +@dataclass +class _Rollouts: + """Recorded traces of the three rollouts and the live robots of the Newton-actuator rollout.""" -class _EquivalenceTestBase(EquivalenceAssertionsMixin, unittest.TestCase): - """Base for Lab-vs-Newton equivalence tests on the PhysX backend. + lab: dict[str, dict] + """Traces of the Lab-actuator rollout.""" + newton: dict[str, dict] + """Traces of the Newton-actuator rollout.""" + reversed: dict + """Trace of the ``reversed`` robot, recorded alone with Newton-native actuators.""" + robots: dict[str, Articulation] + """Robots of the Newton-actuator rollout, whose simulation stays alive for the tests.""" - 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 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 +def _simulation(device: str, use_newton_actuators: bool): + """Return a simulation context with gravity and a ground plane on the PhysX backend.""" 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"]) + return build_simulation_context(device=device, gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg) + + +@pytest.fixture(scope="module", params=test_devices(DeviceScope.DEFAULT_CUDA)) +def rollouts(request) -> Iterator[_Rollouts]: + """Record the Lab-actuator and reversed-ordering rollouts, then keep the Newton-actuator rollout alive.""" + device = request.param + networks = {"mlp": make_dummy_mlp_checkpoint(), "lstm": make_dummy_lstm_checkpoint()} + try: + traces = {} + for name, use_newton_actuators, names in (("lab", False, _LAB_ROBOTS), ("reversed", True, ("reversed",))): + with _simulation(device, use_newton_actuators) as sim: + sim._app_control_on_stop_handle = None + traces[name] = _rollout(sim, _robot_cfgs(names, 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(traces["lab"], newton_traces, traces["reversed"]["reversed"], 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=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", + newton, lab, atol=atol, rtol=rtol, msg=f"{robot} {field} diverged at step {step}" ) - 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 -# --------------------------------------------------------------------------- +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 -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 test_reversed_joint_ordering_matches_identity_ordering(rollouts: _Rollouts) -> None: + """A reversed public joint order commands and observes the same physical joints as the backend order.""" + assert_articulation_ordering_trace_matches( + rollouts.newton["ideal"], rollouts.reversed, tuple(reversed(_ANYMAL_C_PHYSX_JOINT_NAMES)) + ) - 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 - - -# --------------------------------------------------------------------------- -# 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, - ) +def test_reversed_joint_ordering_uses_current_joint_state(rollouts: _Rollouts) -> None: + """Regression: under a reversed joint order, each step's effort follows the IdealPD law on that step's state.""" + trace = rollouts.reversed + 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) == 4 + 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", + ) - @classmethod - def tearDownClass(cls): - os.unlink(cls.lstm_path) + 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 ANYmal controllers alone.""" + anymal, cartpole = rollouts.robots["ideal"], rollouts.robots["cartpole"] + assert anymal.newton_actuator_adapter is not None and cartpole.newton_actuator_adapter is not None + assert anymal.newton_actuator_adapter is not cartpole.newton_actuator_adapter + groups = {"anymal": (anymal, "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({"anymal": anymal, "cartpole": cartpole}, NUM_ENVS, anymal.device) + term, asset_cfg = build_dr_term(env, "cartpole") + + term( + env, + env_ids=torch.tensor([0], device=anymal.device), + asset_cfg=asset_cfg, + stiffness_distribution_params=(100.0, 100.0), + damping_distribution_params=(5.0, 5.0), + operation="abs", + distribution="uniform", + ) - def test_network_drives_haa_joints(self): - _assert_network_drives_haa_joints(self, self.result) + 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)) diff --git a/source/isaaclab_physx/test/assets/test_rigid_object.py b/source/isaaclab_physx/test/assets/test_rigid_object.py index af3057c2e342..75887a65edfb 100644 --- a/source/isaaclab_physx/test/assets/test_rigid_object.py +++ b/source/isaaclab_physx/test/assets/test_rigid_object.py @@ -14,7 +14,7 @@ # 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() @@ -74,6 +74,16 @@ def _yaw_quat(angle: float) -> tuple[float, float, float, float]: return (0.0, 0.0, math.sin(0.5 * angle), math.cos(0.5 * angle)) +@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 + + ## # Tests that own their simulation context. Keep them above the composite scene. ## @@ -93,110 +103,67 @@ def _yaw_quat(angle: float) -> tuple[float, float, float, float]: ), ], ) -@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) @pytest.mark.isaacsim_ci -def test_initialization_rejects_invalid_rigid_body(device, api: Literal["none", "articulation_root"]) -> None: +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.""" - with build_simulation_context(device=device, auto_add_lighting=True) as sim: - sim._app_control_on_stop_handle = None - _spawn_envs("/World") - if api == "none": - # Without rigid body properties the cube is a static collider. - cfg = RigidObjectCfg( - prim_path="/World/Env_[^/]*/Object", - spawn=sim_utils.CuboidCfg(size=(0.1, 0.1, 0.1), collision_props=sim_utils.UsdPhysicsCollisionCfg()), - ) - cube_object = RigidObject(cfg=cfg) - 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() - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(cube_object) < 10 - - with pytest.raises(RuntimeError): - sim.reset() - - -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("gravity_enabled", [True, False]) -@pytest.mark.isaacsim_ci -def test_gravity_vec_w(device, gravity_enabled) -> None: - """Test that gravity vector direction is set correctly for the rigid object.""" - with build_simulation_context(device=device, gravity_enabled=gravity_enabled) as sim: - sim._app_control_on_stop_handle = None - _spawn_envs("/World") + _spawn_envs("/World") + if api == "none": + # Without rigid body properties the cube is a static collider. + cfg = RigidObjectCfg( + prim_path="/World/Env_[^/]*/Object", + spawn=sim_utils.CuboidCfg(size=(0.1, 0.1, 0.1), collision_props=sim_utils.UsdPhysicsCollisionCfg()), + ) + cube_object = RigidObject(cfg=cfg) + 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() - # Obtain gravity direction - if gravity_enabled: - gravity_dir = (0.0, 0.0, -1.0) - else: - gravity_dir = (0.0, 0.0, 0.0) + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(cube_object) < 10 - # Play sim + with pytest.raises(RuntimeError): 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_ENVS, 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.isaacsim_ci @pytest.mark.parametrize("device", test_devices()) -def test_warmup_loads_physics_once(device) -> None: +def test_warmup_loads_physics_once(sim, device) -> None: """Attach on GPU or force-load on CPU, without destroying and rebuilding native objects.""" - 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 - 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() - - 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]) + 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() + + 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]) ## @@ -214,6 +181,8 @@ class _RigidObjectScene: """Dynamic cubes that ignore gravity.""" kinematic: RigidObject """Kinematic cubes.""" + falling: RigidObject + """Dynamic cubes under gravity.""" origins: dict[str, torch.Tensor] """Environment origins keyed by object name.""" refcounts: dict[str, int] @@ -222,10 +191,10 @@ class _RigidObjectScene: def step(self, num_steps: int = 1) -> None: """Write, step, and update every object.""" for _ in range(num_steps): - for rigid_object in (self.cubes, self.kinematic): + for rigid_object in (self.cubes, self.kinematic, self.falling): rigid_object.write_data_to_sim() self.sim.step() - for rigid_object in (self.cubes, self.kinematic): + for rigid_object in (self.cubes, self.kinematic, self.falling): rigid_object.update(self.sim.cfg.dt) def place_cubes_at_rest(self, root_pose: torch.Tensor) -> None: @@ -241,9 +210,14 @@ def rigid_object_scene(request) -> Iterator[_RigidObjectScene]: device = request.param with build_simulation_context(device=device, gravity_enabled=True) as sim: sim._app_control_on_stop_handle = None - origins = {"cubes": _spawn_envs("/World/Cubes"), "kinematic": _spawn_envs("/World/Kinematic", 3.0)} + origins = { + "cubes": _spawn_envs("/World/Cubes"), + "kinematic": _spawn_envs("/World/Kinematic", 3.0), + "falling": _spawn_envs("/World/Falling", 6.0), + } 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")) refcounts = {"cubes": sys.getrefcount(cubes), "kinematic": sys.getrefcount(kinematic)} sim.reset() yield _RigidObjectScene( @@ -251,6 +225,7 @@ def rigid_object_scene(request) -> Iterator[_RigidObjectScene]: device=device, cubes=cubes, kinematic=kinematic, + falling=falling, origins={name: value.to(device) for name, value in origins.items()}, refcounts=refcounts, ) @@ -258,7 +233,8 @@ def rigid_object_scene(request) -> Iterator[_RigidObjectScene]: @pytest.mark.isaacsim_ci def test_rigid_object_initialization(rigid_object_scene: _RigidObjectScene) -> None: - """Initialize local cubes with the expected buffers; kinematic cubes hold their default pose under gravity.""" + """Initialize local cubes with the expected buffers; under gravity, kinematic cubes hold their default pose and + dynamic cubes accelerate downward on every step.""" scene = rigid_object_scene for name, rigid_object in (("cubes", scene.cubes), ("kinematic", scene.kinematic)): # Check that the framework doesn't hold excessive strong references. @@ -270,10 +246,15 @@ def test_rigid_object_initialization(rigid_object_scene: _RigidObjectScene) -> N 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) + ) kinematic = scene.kinematic + gravity_acceleration = torch.tensor([[[0.0, 0.0, -9.81, 0.0, 0.0, 0.0]]] * _NUM_ENVS, device=scene.device) for _ in range(2): scene.step() + torch.testing.assert_close(scene.falling.data.body_acc_w.torch, gravity_acceleration) 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) diff --git a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py index 1a1022929003..2925f229df63 100644 --- a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py @@ -14,7 +14,7 @@ # 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() @@ -98,15 +98,11 @@ def _yaw_quat(angle: float) -> tuple[float, float, float, float]: @pytest.fixture -def gravity_enabled() -> bool: - """Enable gravity unless a test explicitly parametrizes it otherwise.""" - return True - - -@pytest.fixture -def sim(device: str, gravity_enabled: bool) -> Iterator[SimulationContext]: +def sim(device: str) -> Iterator[SimulationContext]: """Create a function-scoped simulation context for tests that own their scene.""" - with build_simulation_context(device=device, auto_add_lighting=True, gravity_enabled=gravity_enabled) as sim: + # 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 @@ -116,7 +112,7 @@ def sim(device: str, gravity_enabled: bool) -> Iterator[SimulationContext]: ## -@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) def test_initialization_with_no_rigid_body(sim, device) -> None: """Test that initialization fails when no rigid body is found at the provided prim path.""" num_cubes = 2 @@ -130,94 +126,21 @@ def test_initialization_with_no_rigid_body(sim, device) -> None: sim.reset() -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("gravity_enabled", [False]) -def test_subset_write_reaches_selected_view_entry(sim, device, gravity_enabled) -> None: - """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) - - # 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("device", test_devices()) -def test_inertial_property_subset_writes_reach_selected_view_entries(sim, device) -> None: - """COM and inertia writes to a subset of (env, body) cells must reach only those PhysX view entries.""" - num_envs, num_cubes = 2, 3 - object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes) - sim.reset() - - def read_view(values: wp.array, data_dim: int) -> torch.Tensor: - # The view is body-major, (num_cubes * num_envs, data_dim); return it as (num_envs, num_cubes, data_dim) - return wp.to_torch(values).reshape(num_cubes, num_envs, data_dim).transpose(0, 1).cpu() - - # Write in non-sorted env order so that a wrong env/body to view index mapping cannot pass - env_ids = [1, 0] - body_ids = [2] - - initial_coms = read_view(object_collection.root_view.get_coms().view(wp.float32), 7) - coms = initial_coms[env_ids][:, body_ids].clone() - coms[0, 0, :3] = torch.tensor([0.02, 0.03, 0.04]) - coms[1, 0, :3] = torch.tensor([-0.01, 0.01, 0.02]) - object_collection.set_coms_index(coms=coms.to(device), env_ids=env_ids, body_ids=body_ids) - expected_coms = initial_coms.clone() - expected_coms[env_ids, 2] = coms[:, 0] - torch.testing.assert_close(read_view(object_collection.root_view.get_coms().view(wp.float32), 7), expected_coms) - torch.testing.assert_close(object_collection.data.body_com_pose_b.torch.cpu(), expected_coms) - - initial_inertias = read_view(object_collection.root_view.get_inertias(), 9) - inertias = initial_inertias[env_ids][:, body_ids].clone() - inertias[0, 0, [0, 4, 8]] *= 1.5 - inertias[1, 0, [0, 4, 8]] *= 2.0 - object_collection.set_inertias_index(inertias=inertias.to(device), env_ids=env_ids, body_ids=body_ids) - expected_inertias = initial_inertias.clone() - expected_inertias[env_ids, 2] = inertias[:, 0] - torch.testing.assert_close(read_view(object_collection.root_view.get_inertias(), 9), expected_inertias) - torch.testing.assert_close(object_collection.data.body_inertia.torch.cpu(), expected_inertias) - - -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("gravity_enabled", [True, False]) -def test_gravity_vec_w(sim, device, gravity_enabled) -> None: - """Test that gravity vector direction is set correctly for the rigid object.""" - num_envs = 3 - num_cubes = 2 - object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes) - - gravity_dir = (0.0, 0.0, -1.0) if gravity_enabled else (0.0, 0.0, 0.0) - - sim.reset() - - 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] - - 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 - - torch.testing.assert_close(object_collection.data.body_com_acc_w.torch, gravity) +@pytest.mark.xfail( + strict=True, raises=AssertionError, reason="The first CPU simulation after a failed initialization reads zero poses" +) +def test_failed_initialization_does_not_leak_into_the_next_cpu_simulation() -> None: + """A kinematic cube reports its spawn pose in the CPU simulation that follows a failed initialization.""" + with build_simulation_context(device="cpu") as sim: + sim._app_control_on_stop_handle = None + invalid_collection, _ = generate_cubes_scene(num_cubes=2, has_api=False) # noqa: F841 # keep it alive + with pytest.raises(RuntimeError): + sim.reset() + with build_simulation_context(device="cpu") as sim: + sim._app_control_on_stop_handle = None + kinematic, origins = generate_cubes_scene(kinematic_enabled=True) + sim.reset() + torch.testing.assert_close(kinematic.data.body_link_pos_w.torch[:, 0], origins + torch.tensor([0.0, 0.0, 1.0])) ## @@ -238,6 +161,8 @@ class _CollectionScene: """Two environments of three dynamic cubes that ignore gravity.""" kinematic: RigidObjectCollection """One environment of one kinematic cube.""" + falling: RigidObjectCollection + """Two environments of two dynamic cubes under gravity.""" origins: dict[str, torch.Tensor] """Environment origins keyed by collection name.""" refcounts: dict[str, int] @@ -246,10 +171,10 @@ class _CollectionScene: def step(self, num_steps: int = 1) -> None: """Write, step, and update every collection.""" for _ in range(num_steps): - for collection in (self.cubes, self.kinematic): + for collection in (self.cubes, self.kinematic, self.falling): collection.write_data_to_sim() self.sim.step() - for collection in (self.cubes, self.kinematic): + for collection in (self.cubes, self.kinematic, self.falling): collection.update(self.sim.cfg.dt) def place_cubes_at_rest(self, body_poses: torch.Tensor) -> None: @@ -274,6 +199,7 @@ def collection_scene(request) -> Iterator[_CollectionScene]: kinematic, kinematic_origins = generate_cubes_scene( num_envs=1, num_cubes=1, kinematic_enabled=True, root="/World/Kinematic", y_offset=12.0 ) + falling, _ = generate_cubes_scene(num_envs=_NUM_ENVS, num_cubes=2, root="/World/Falling", y_offset=18.0) refcounts = {"cubes": sys.getrefcount(cubes), "kinematic": sys.getrefcount(kinematic)} sim.reset() yield _CollectionScene( @@ -281,6 +207,7 @@ def collection_scene(request) -> Iterator[_CollectionScene]: device=device, cubes=cubes, kinematic=kinematic, + falling=falling, origins={"cubes": cube_origins.to(device), "kinematic": kinematic_origins.to(device)}, refcounts=refcounts, ) @@ -295,7 +222,8 @@ def _rest_poses(scene: _CollectionScene, yaw: float = 0.0) -> torch.Tensor: def test_collection_initialization(collection_scene: _CollectionScene) -> None: - """Initialize local collections, including a single-cube one; kinematic cubes hold their pose under gravity.""" + """Initialize local collections, including a single-cube one; under gravity, kinematic cubes hold their pose and + dynamic cubes accelerate downward on every step.""" scene = collection_scene for name, collection, num_envs, num_cubes in ( ("cubes", scene.cubes, _NUM_ENVS, _NUM_CUBES), @@ -310,10 +238,15 @@ def test_collection_initialization(collection_scene: _CollectionScene) -> None: 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) + ) kinematic = scene.kinematic + gravity_acceleration = torch.tensor([0.0, 0.0, -9.81, 0.0, 0.0, 0.0], device=scene.device).repeat(_NUM_ENVS, 2, 1) for _ in range(2): scene.step() + torch.testing.assert_close(scene.falling.data.body_com_acc_w.torch, gravity_acceleration) default_body_pose = kinematic.data.default_body_pose.torch.clone() default_body_vel = kinematic.data.default_body_vel.torch.clone() default_body_pose[..., :3] += scene.origins["kinematic"].unsqueeze(1) @@ -439,22 +372,41 @@ def test_collection_state_writes(collection_scene: _CollectionScene) -> None: torch.testing.assert_close(object_com_vel_w[..., 3:], object_link_vel_w[..., 3:]) -def test_collection_masses_reach_selected_view_entries(collection_scene: _CollectionScene) -> None: - """Mass writes to a non-sorted subset reach only the selected body-major view entries.""" +def test_collection_inertial_properties_reach_selected_view_entries(collection_scene: _CollectionScene) -> None: + """Mass, center-of-mass, and inertia writes to a non-sorted subset reach only the selected body-major entries.""" scene = collection_scene device = scene.device collection = scene.cubes env_ids = torch.tensor([1, 0], dtype=torch.int32, device=device) body_ids = torch.tensor([1], dtype=torch.int32, device=device) - expected = collection.data.body_mass.torch.clone() - masses = torch.tensor([[5.0], [7.0]], device=device) - expected[env_ids[:, None], body_ids] = masses - collection.set_masses_index(masses=masses, env_ids=env_ids, body_ids=body_ids) - torch.testing.assert_close(collection.data.body_mass.torch, expected) - raw_mass = wp.to_torch(collection.root_view.get_masses()).to(device).reshape(_NUM_CUBES, _NUM_ENVS).T - torch.testing.assert_close(raw_mass, expected) - # Restore uniform masses for the wrench checks. - collection.set_masses_index(masses=torch.ones((_NUM_ENVS, _NUM_CUBES), 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(collection_scene: _CollectionScene) -> None: diff --git a/source/isaaclab_physx/test/assets/unit/test_physx_articulation.py b/source/isaaclab_physx/test/assets/unit/test_physx_articulation.py index 93d2368b68ba..023f038f49e4 100644 --- a/source/isaaclab_physx/test/assets/unit/test_physx_articulation.py +++ b/source/isaaclab_physx/test/assets/unit/test_physx_articulation.py @@ -5,11 +5,6 @@ """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 @@ -24,109 +19,6 @@ def _selector(values: list[int], dtype: type) -> wp.array: 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: From 7cd30713df293afe04ba2423b5e7c522909a8cbc Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Fri, 2 Oct 2026 15:20:41 -0700 Subject: [PATCH 25/33] Trim OVPhysX asset, lifecycle, and unit tests --- .../test/assets/test_articulation.py | 95 +++++---- .../test/assets/test_rigid_object.py | 63 +++--- .../assets/test_rigid_object_collection.py | 50 ++--- .../assets/unit/test_articulation_helpers.py | 122 ------------ .../assets/unit/test_articulation_kernels.py | 52 ----- .../test/assets/unit/test_deformable_views.py | 122 ------------ .../unit/test_ovphysx_actuator_control.py | 70 ------- .../assets/unit/test_ovphysx_articulation.py | 187 +++++++++++++++++- .../unit/test_ovphysx_deformable_object.py | 53 ++++- .../physics/test_ovphysx_manager_lifecycle.py | 6 +- 10 files changed, 310 insertions(+), 510 deletions(-) delete mode 100644 source/isaaclab_ov/test/assets/unit/test_articulation_helpers.py delete mode 100644 source/isaaclab_ov/test/assets/unit/test_articulation_kernels.py delete mode 100644 source/isaaclab_ov/test/assets/unit/test_deformable_views.py delete mode 100644 source/isaaclab_ov/test/assets/unit/test_ovphysx_actuator_control.py diff --git a/source/isaaclab_ov/test/assets/test_articulation.py b/source/isaaclab_ov/test/assets/test_articulation.py index 785058530825..5e10eab1927e 100644 --- a/source/isaaclab_ov/test/assets/test_articulation.py +++ b/source/isaaclab_ov/test/assets/test_articulation.py @@ -180,7 +180,6 @@ class _IslandCfg: "non_root": _IslandCfg( _implicit(), fixed_base=False, articulation_root_prim_path="/base", devices=DeviceScope.CUDA ), - "pos_limits": _IslandCfg(_implicit(), devices=DeviceScope.CUDA), "com_order": _IslandCfg( _implicit(), num_envs=1, body_ordering=_ROOT_PRESERVING_REVERSED_BODY_NAMES, distinct_coms=True ), @@ -381,9 +380,15 @@ def test_mimic_finger_follows_commanded_finger(device: str) -> None: 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"), + [ + (device, variant) + for variant in ("geometry", "joint_type", "d6_translation", "disabled_joint") + for device in _ALL_DEVICES + ] + # 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.""" @@ -874,27 +879,6 @@ def test_explicit_articulation_root_on_a_floating_base_body(scene: _Articulation assert torch.isfinite(articulation.data.root_link_pose_w.torch).all() -@pytest.mark.parametrize("scene", _CUDA_DEVICES, indirect=True) -def test_joint_position_limits_clamp_default_joint_positions(scene: _ArticulationScene) -> None: - """Write partial joint position limits and clamp only the selected default joint positions into them.""" - articulation = scene.islands["pos_limits"] - device = scene.device - default_before = articulation.data.default_joint_pos.torch.clone() - env_ids = torch.tensor([1], dtype=torch.int32, device=device) - joint_ids = torch.tensor([0, 2], dtype=torch.int32, device=device) - # The default positions are zero, so both new ranges exclude them. - limits = torch.tensor([[[0.2, 0.3], [0.25, 0.35]]], device=device) - articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) - - torch.testing.assert_close(articulation.data.joint_pos_limits.torch[env_ids][:, joint_ids], limits) - default_joint_pos = articulation.data.default_joint_pos.torch - selected = default_joint_pos[env_ids][:, joint_ids] - assert torch.all((selected >= limits[..., 0]) & (selected <= 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], default_before[unselected]) - - @pytest.mark.parametrize("scene", _CUDA_DEVICES, 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.""" @@ -1161,31 +1145,21 @@ def test_floating_articulation_root_state_dynamics_wrench_and_reset(scene: _Arti raw_gravity[:, generalized_user_to_backend], ) - # Reset should zero external forces and torques - articulation.reset() - 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 - - articulation.permanent_wrench_composer.set_forces_and_torques_index( - forces=torch.ones((num_articulations, num_bodies, 3), device=device), - torques=torch.ones((num_articulations, num_bodies, 3), device=device), - ) - articulation.instantaneous_wrench_composer.add_forces_and_torques_index( - forces=torch.ones((num_articulations, num_bodies, 3), device=device), - torques=torch.ones((num_articulations, num_bodies, 3), device=device), - ) + # 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)) - 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 + 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 # 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) @@ -1483,15 +1457,34 @@ def occupy_device_stream(iterations: int = 500_000): @pytest.mark.parametrize("scene", _ALL_DEVICES, indirect=True) -def test_joint_position_limit_clamping_respects_logging(scene: _ArticulationScene, caplog) -> None: - """Logging level controls reporting without changing clamping or reusing a stale violation count.""" +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() - limits = torch.zeros_like(original_limits) - limits[..., 1] = 0.5 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() diff --git a/source/isaaclab_ov/test/assets/test_rigid_object.py b/source/isaaclab_ov/test/assets/test_rigid_object.py index 771007636a2d..3cd7835d2f67 100644 --- a/source/isaaclab_ov/test/assets/test_rigid_object.py +++ b/source/isaaclab_ov/test/assets/test_rigid_object.py @@ -104,10 +104,15 @@ class HeterogeneousRigidSceneCfg(InteractiveSceneCfg): @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.""" +@pytest.mark.parametrize( + ("support_cloning", "filter_collisions"), + [("native", False), ("usd_and_native", True), ("usd", False), ("usd_nested", True)], +) +def test_heterogeneous_clone_contacts(device, support_cloning, filter_collisions): + """Sources and clones from different variants contact their own support. + + 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: @@ -222,7 +227,6 @@ class _RigidScene: device: str dynamic: RigidObject kinematic: RigidObject - resettable: RigidObject spinning: RigidObject @@ -233,17 +237,14 @@ def scene(request: pytest.FixtureRequest) -> Iterator[_RigidScene]: 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) - resettable = _spawn_cubes("Resettable", y_offset=4.0, disable_gravity=True) - spinning = _spawn_cubes("Spinning", y_offset=6.0, disable_gravity=True) + spinning = _spawn_cubes("Spinning", y_offset=4.0, disable_gravity=True) sim.reset() - yield _RigidScene( - sim=sim, device=device, dynamic=dynamic, kinematic=kinematic, resettable=resettable, spinning=spinning - ) + yield _RigidScene(sim=sim, device=device, dynamic=dynamic, kinematic=kinematic, spinning=spinning) @pytest.mark.parametrize("scene", test_devices(), indirect=True) def test_rigid_object_real_ovphysx_seams(scene: _RigidScene) -> None: - """Prove partial state, inertial properties, and one real wrench delivery.""" + """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 @@ -300,6 +301,16 @@ def test_rigid_object_real_ovphysx_seams(scene: _RigidScene) -> None: 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: @@ -308,7 +319,7 @@ def test_initialization_with_kinematic_enabled(scene: _RigidScene) -> None: # SDP bindings must be ready on CPU and GPU before any asset pose read. provider = sim.get_scene_data_provider() - names = ("Dynamic", "Kinematic", "Resettable", "Spinning") + names = ("Dynamic", "Kinematic", "Spinning") expected_paths = {f"/World/{name}/Env_{i}/Cube" for name in names for i in (0, 1)} assert provider.transform_count == len(expected_paths) assert set(provider.backend.transform_paths) == expected_paths @@ -325,34 +336,6 @@ def test_initialization_with_kinematic_enabled(scene: _RigidScene) -> None: torch.testing.assert_close(cube_object.data.root_com_vel_w.torch, default_root_vel) -@pytest.mark.parametrize("scene", test_devices(DeviceScope.CUDA), indirect=True) -def test_reset_clears_active_wrench_composers(scene: _RigidScene) -> None: - """Test that resetting the rigid object clears both active wrench composers.""" - cube_object, device = scene.resettable, scene.device - - # 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 - - 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("scene", test_devices(DeviceScope.CUDA), 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.""" diff --git a/source/isaaclab_ov/test/assets/test_rigid_object_collection.py b/source/isaaclab_ov/test/assets/test_rigid_object_collection.py index 3ed2a3f16b24..fc0d3ae0671a 100644 --- a/source/isaaclab_ov/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_ov/test/assets/test_rigid_object_collection.py @@ -91,7 +91,6 @@ class _CollectionScene: sim: SimulationContext device: str fused: RigidObjectCollection - resettable: RigidObjectCollection stepped: RigidObjectCollection wrenched: RigidObjectCollection @@ -104,13 +103,10 @@ def scene(request: pytest.FixtureRequest) -> Iterator[_CollectionScene]: device=device, sim_cfg=SimulationCfg(physics=OvPhysxCfg(), device=device, gravity=(0.0, 0.0, 0.0)) ) as sim: fused = _spawn_collection("Fused", y_offset=0.0) - resettable = _spawn_collection("Resettable", y_offset=4.0) - stepped = _spawn_collection("Stepped", y_offset=8.0) - wrenched = _spawn_collection("Wrenched", y_offset=12.0) + stepped = _spawn_collection("Stepped", y_offset=4.0) + wrenched = _spawn_collection("Wrenched", y_offset=8.0) sim.reset() - yield _CollectionScene( - sim=sim, device=device, fused=fused, resettable=resettable, stepped=stepped, wrenched=wrenched - ) + yield _CollectionScene(sim=sim, device=device, fused=fused, stepped=stepped, wrenched=wrenched) @pytest.mark.parametrize("scene", test_devices(), indirect=True) @@ -188,37 +184,9 @@ def test_rigid_object_collection_real_ovphysx_seams(scene: _CollectionScene) -> torch.testing.assert_close(raw_materials.reshape(3, 2, 3).transpose(0, 1), materials) -@pytest.mark.parametrize("scene", test_devices(DeviceScope.CUDA), indirect=True) -def test_reset_clears_active_wrench_composers(scene: _CollectionScene) -> None: - """Test that resetting the collection clears both active wrench composers.""" - object_collection, device = scene.resettable, scene.device - - # 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_BODIES, 3), device=device), - torques=torch.ones((_NUM_ENVS, _NUM_BODIES, 3), device=device), - ) - object_collection.instantaneous_wrench_composer.add_forces_and_torques_index( - forces=torch.ones((_NUM_ENVS, _NUM_BODIES, 3), device=device), - torques=torch.ones((_NUM_ENVS, _NUM_BODIES, 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("scene", test_devices(DeviceScope.CUDA), indirect=True) def test_wrench_reaches_only_the_selected_collection_body(scene: _CollectionScene) -> None: - """Deliver an external force through the fused body-major binding to the selected body only. + """Deliver an external force through the fused body-major binding to the selected body only, then reset it. Environment 1, body 1 has different flat indices in body-major and environment-major layouts. """ @@ -233,7 +201,6 @@ def test_wrench_reaches_only_the_selected_collection_body(scene: _CollectionScen collection.write_data_to_sim() sim.step() collection.update(sim.cfg.dt) - collection.reset() velocity = collection.data.body_com_vel_w.torch assert velocity[1, 1, 0] > velocity_before[1, 1, 0] + 1e-3 @@ -241,6 +208,15 @@ def test_wrench_reaches_only_the_selected_collection_body(scene: _CollectionScen 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", test_devices(DeviceScope.CUDA), indirect=True) @pytest.mark.xfail( diff --git a/source/isaaclab_ov/test/assets/unit/test_articulation_helpers.py b/source/isaaclab_ov/test/assets/unit/test_articulation_helpers.py deleted file mode 100644 index b2fb56274e71..000000000000 --- a/source/isaaclab_ov/test/assets/unit/test_articulation_helpers.py +++ /dev/null @@ -1,122 +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 - -pytestmark = pytest.mark.unit - - -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/unit/test_articulation_kernels.py b/source/isaaclab_ov/test/assets/unit/test_articulation_kernels.py deleted file mode 100644 index 3cb1b7abd9e5..000000000000 --- a/source/isaaclab_ov/test/assets/unit/test_articulation_kernels.py +++ /dev/null @@ -1,52 +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 - -pytestmark = pytest.mark.unit - - -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/unit/test_deformable_views.py b/source/isaaclab_ov/test/assets/unit/test_deformable_views.py deleted file mode 100644 index 6211cf9d9d60..000000000000 --- a/source/isaaclab_ov/test/assets/unit/test_deformable_views.py +++ /dev/null @@ -1,122 +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 - -pytestmark = pytest.mark.unit - - -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/unit/test_ovphysx_actuator_control.py b/source/isaaclab_ov/test/assets/unit/test_ovphysx_actuator_control.py deleted file mode 100644 index 0da4414e6197..000000000000 --- a/source/isaaclab_ov/test/assets/unit/test_ovphysx_actuator_control.py +++ /dev/null @@ -1,70 +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 - -"""OVPhysX actuator-control unit tests.""" - -from __future__ import annotations - -import importlib -from types import SimpleNamespace - -import pytest - -pytest.importorskip("ovphysx.types", reason="ovphysx wheel not installed") - -from isaaclab_ov.assets.articulation import actuator_control # noqa: E402 -from isaaclab_ov.assets.articulation.actuator_control import OvPhysxActuatorControl # noqa: E402 - -from isaaclab.actuators import ImplicitActuatorCfg # noqa: E402 - -pytestmark = pytest.mark.unit - - -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)) diff --git a/source/isaaclab_ov/test/assets/unit/test_ovphysx_articulation.py b/source/isaaclab_ov/test/assets/unit/test_ovphysx_articulation.py index c15cd99c69e1..aeb11ebf3178 100644 --- a/source/isaaclab_ov/test/assets/unit/test_ovphysx_articulation.py +++ b/source/isaaclab_ov/test/assets/unit/test_ovphysx_articulation.py @@ -6,25 +6,32 @@ # ignore private usage of variables warning # pyright: reportPrivateUsage=none -"""OVPhysX articulation data-cache and joint-direction unit tests.""" +"""OVPhysX articulation unit tests: data caches, joint directions, tendon scoping, kernels, and actuator control.""" from __future__ import annotations +import importlib from types import SimpleNamespace from unittest.mock import Mock +import numpy as np import pytest import torch import warp as wp -from pxr import Usd, UsdGeom, UsdPhysics +from pxr import Sdf, Usd, UsdGeom, UsdPhysics 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 OvPhysxManager # noqa: E402 +from isaaclab_ov.test.fixtures.views import MockOvPhysxBindingSet # noqa: E402 +from isaaclab.actuators import ImplicitActuatorCfg # noqa: E402 from isaaclab.utils.warp.launch_cache import _WarpLaunchCache # noqa: E402 pytestmark = pytest.mark.unit @@ -146,3 +153,177 @@ def test_joint_dof_sign_resolution_traverses_instance_proxies(): ) assert Articulation._resolve_joint_dof_signs(articulation, stage) == (-1,) + + +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"] + + +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]]) + + +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)) diff --git a/source/isaaclab_ov/test/assets/unit/test_ovphysx_deformable_object.py b/source/isaaclab_ov/test/assets/unit/test_ovphysx_deformable_object.py index fc1e2441eff6..9ae5882c8daf 100644 --- a/source/isaaclab_ov/test/assets/unit/test_ovphysx_deformable_object.py +++ b/source/isaaclab_ov/test/assets/unit/test_ovphysx_deformable_object.py @@ -3,11 +3,10 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Unit tests for the kitless OVPhysX deformable object asset.""" +"""Unit tests for the kitless OVPhysX deformable object asset and its body view.""" from __future__ import annotations -import importlib from types import SimpleNamespace import pytest @@ -17,11 +16,12 @@ 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 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.assets.deformable_object.views import OvPhysxDeformableBodyView # noqa: E402 from isaaclab_ov.physics.ovphysx_manager import OvPhysxManager, OvPhysxSceneDataBackend # noqa: E402 pytestmark = pytest.mark.unit @@ -289,12 +289,11 @@ def test_full_overwrite_stale_cache_does_not_read_simulator() -> None: 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): +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, trailing_dimension), device=asset.device) + 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]) @@ -342,7 +341,43 @@ def test_surface_debug_visualization_uses_below_ground_sentinel(): 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") +class _MixedTopologyBinding: + """Padded simulation connectivity with five nodes in one body and four in another.""" - assert assets_module.DeformableObject is DeformableObject + _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")) + + +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, + ) 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 47d94d4a17ce..10fba55987d9 100644 --- a/source/isaaclab_ov/test/physics/test_ovphysx_manager_lifecycle.py +++ b/source/isaaclab_ov/test/physics/test_ovphysx_manager_lifecycle.py @@ -647,12 +647,10 @@ def test_default_cache_dir_rejects_a_directory_owned_by_another_user(manager_mod @pytest.mark.skipif(not (_CPU_DEVICES and _CUDA_DEVICES), reason="The device sequence requires a CPU and a CUDA device") -@pytest.mark.parametrize("order", ["cpu-cuda-cpu", "cuda-cpu"]) -def test_cpu_and_cuda_scenes_run_sequentially_in_one_process(order): +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] - devices = (cpu, cuda, cpu) if order == "cpu-cuda-cpu" else (cuda, cpu) - completed, output = _run_child(_device_sequence_script(devices)) + completed, output = _run_child(_device_sequence_script((cpu, cuda, cpu))) assert completed.returncode == 0, output[-8000:] assert "DEVICE_SEQUENCE_OK" in output, output[-8000:] From 3346171af0595d37527b846ad459e799ea2f45fc Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Fri, 2 Oct 2026 15:20:41 -0700 Subject: [PATCH 26/33] Trim WrenchComposer and actuator runtime tests --- .../test/utils/actuator_equivalence.py | 157 +----------------- .../actuators/test_physx_actuator_runtime.py | 43 ++--- .../test/utils/test_wrench_composer.py | 82 +++------ 3 files changed, 43 insertions(+), 239 deletions(-) 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_physx_actuator_runtime.py b/source/isaaclab/test/actuators/test_physx_actuator_runtime.py index 3bbdf0c35224..ec77abf5357b 100644 --- a/source/isaaclab/test/actuators/test_physx_actuator_runtime.py +++ b/source/isaaclab/test/actuators/test_physx_actuator_runtime.py @@ -19,25 +19,30 @@ def _runtime() -> PhysxActuatorRuntime: return PhysxActuatorRuntime(SimpleNamespace(device="cuda:0"), logger=Mock()) -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.""" +def _failing_capture(fail_at: int) -> type: + """Return a ``wp.ScopedCapture`` double whose capture number ``fail_at`` (0-based) raises on entry.""" - class _FailingCapture: - capture_count = 0 + class _Capture: + count = 0 def __init__(self, *args, **kwargs) -> None: - self.capture_index = type(self).capture_count - type(self).capture_count += 1 + self.index = type(self).count + type(self).count += 1 self.graph = object() def __enter__(self): - if self.capture_index == 1: - raise RuntimeError("second capture unavailable") + if self.index >= fail_at: + raise RuntimeError("capture unavailable") return self - def __exit__(self, exc_type, exc_value, traceback) -> bool: + 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) @@ -45,7 +50,7 @@ def __exit__(self, exc_type, exc_value, traceback) -> bool: 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", _FailingCapture) + 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) @@ -58,24 +63,12 @@ def _swap_adapter_state(*args, **kwargs) -> None: 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.""" - - class _FailingCapture: - capture_count = 0 - - def __init__(self, *args, **kwargs) -> None: - type(self).capture_count += 1 - - def __enter__(self): - raise RuntimeError("capture unavailable") - - def __exit__(self, exc_type, exc_value, traceback) -> bool: - return False - + 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", _FailingCapture) + monkeypatch.setattr(wp, "ScopedCapture", capture) monkeypatch.setattr(runtime, "_run_native_actuator_kernels", eager_compute) collection = SimpleNamespace() @@ -83,7 +76,7 @@ def __exit__(self, exc_type, exc_value, traceback) -> bool: runtime.compute(collection, 0.01) # The failed capture is not retried, and every step falls back to the eager kernels. - assert _FailingCapture.capture_count == 1 + assert capture.count == 1 assert runtime.native_actuator_graphs == () assert eager_compute.call_count == 2 eager_compute.assert_called_with(collection, 0.01) diff --git a/source/isaaclab/test/utils/test_wrench_composer.py b/source/isaaclab/test/utils/test_wrench_composer.py index de7dce62bace..10daf0d06612 100644 --- a/source/isaaclab/test/utils/test_wrench_composer.py +++ b/source/isaaclab/test/utils/test_wrench_composer.py @@ -10,10 +10,7 @@ ``(-x, -y, z)``. Real-solver delivery is covered by ``test_wrench_composer_integration.py``. """ -import ast -import inspect from types import SimpleNamespace -from typing import get_type_hints from unittest.mock import patch import numpy as np @@ -39,19 +36,6 @@ _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] - - # --------------------------------------------------------------------------- # Literal fixture helpers # --------------------------------------------------------------------------- @@ -252,11 +236,8 @@ def test_index_and_mask_selection_change_only_selected_cells() -> None: _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.""" asset = SimpleNamespace(num_instances=3, num_bodies=3, device="cpu", data=_AssetData(*_poses(shape=(3, 3)))) composer = WrenchComposer(asset) @@ -332,8 +313,7 @@ def test_set_clears_every_buffer_of_targeted_environments_only(method: str) -> N _assert_vectors(getattr(composer, name), expected) -@pytest.mark.parametrize("source_is_global", [False, True], ids=["local_source", "global_source"]) -@pytest.mark.parametrize("selector", ["list", "slice", "mask"]) +@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) @@ -435,34 +415,21 @@ def test_raw_buffer_merge_accumulates_all_five_buffers_and_ignores_inactive_sour # --------------------------------------------------------------------------- -@pytest.mark.parametrize("input_name, output_name", [("forces", "out_force_b"), ("torques", "out_torque_b")]) -def test_output_property_triggers_lazy_composition(input_name: str, output_name: str): - """Test that reading out_force_b/out_torque_b without explicit compose_to_body_frame returns correct results.""" +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, **{input_name: (2.0, 0.0, 0.0)}) - - # Do NOT call compose_to_body_frame -- rely on lazy composition + _apply_to_cell(composer, "add_forces_and_torques_index", (0, 1), True, forces=(2.0, 0.0, 0.0)) assert composer._dirty - _assert_vectors(getattr(composer, output_name), _grid({(0, 1): (0.0, -2.0, 0.0)})) + _assert_vectors(composer.out_force_b, _grid({(0, 1): (0.0, -2.0, 0.0)})) assert not composer._dirty - -def test_lazy_composition_reflects_later_adds(): - """Test that an add after a lazy composition is reflected by the next output read.""" - composer = _make_composer() - ones = torch.ones((2, 2, 3)) - composer.add_forces_and_torques_index(forces=_vectors(ones)) - # Reading out_force_b composes lazily - _assert_vectors(composer.out_force_b, ones) - - # Another add must be reflected by the next read of either output - composer.add_forces_and_torques_index(forces=_vectors(ones), torques=_vectors(ones)) - _assert_vectors(composer.out_torque_b, ones) - _assert_vectors(composer.out_force_b, 2.0 * ones) + _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() - _assert_vectors(composer.out_force_b, 2.0 * ones) - _assert_vectors(composer.out_torque_b, ones) + _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)})) # --------------------------------------------------------------------------- @@ -470,21 +437,20 @@ def test_lazy_composition_reflects_later_adds(): # --------------------------------------------------------------------------- -@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(deprecated_name: str, name: str, input_name: str): - """Test that accessing composed_force/composed_torque emits a DeprecationWarning and aliases the output.""" +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, **{input_name: (1.0, 2.0, 3.0)}) - - with pytest.warns(DeprecationWarning, match=f"{deprecated_name}.*is deprecated"): - result = getattr(composer, deprecated_name) + _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) + ) - # The deprecated property aliases the new output property. - assert result is getattr(composer, name) - _assert_vectors(result, _grid({(1, 0): (1.0, 2.0, 3.0)})) + 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})) def test_deprecated_writers_warn_and_keep_add_and_set_behavior() -> None: From 364734367eb564badff0d01f18b49da4d8e1880c Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Fri, 2 Oct 2026 15:20:41 -0700 Subject: [PATCH 27/33] Remove device_split runner support --- .github/actions/run-package-tests/action.yml | 2 +- .github/actions/run-tests/action.yml | 2 +- .github/workflows/tools-tests.yml | 4 +- pyproject.toml | 1 - tools/_device_split.py | 65 ----------- tools/conftest.py | 21 +--- tools/test_device_split.py | 108 ------------------- 7 files changed, 5 insertions(+), 198 deletions(-) delete mode 100644 tools/_device_split.py delete mode 100644 tools/test_device_split.py 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/pyproject.toml b/pyproject.toml index 93dcd5a2ffbf..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 in a separate process per device (CPU, then GPU) in CI", "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/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 From 61d6e4b8c19c46425a3561e4836c2aada76135c4 Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Fri, 2 Oct 2026 15:20:41 -0700 Subject: [PATCH 28/33] Consolidate asset-test changelog fragments --- .../antoiner-asset-tests-p03-physx-actuator-runtime.skip | 0 .../changelog.d/antoiner-asset-tests-p04-terminations.skip | 0 .../antoiner-asset-tests-p13-physx-newton-actuators.skip | 0 ...tests-p01-contract.skip => antoiner-asset-tests-redesign.skip} | 0 source/isaaclab/changelog.d/asset-tests-readability.skip | 0 .../changelog.d/antoiner-asset-tests-p18-newton-fixtures.skip | 0 .../antoiner-asset-tests-p19-newton-task-space-controllers.skip | 0 .../changelog.d/antoiner-asset-tests-p20-newton-actuators.skip | 0 .../changelog.d/antoiner-asset-tests-p21-newton-rigid-object.skip | 0 .../antoiner-asset-tests-p22-newton-rigid-object-collection.skip | 0 .../antoiner-asset-tests-p23-newton-articulation-core.skip | 0 .../antoiner-asset-tests-p24-newton-articulation-regressions.skip | 0 .../changelog.d/antoiner-asset-tests-redesign.skip} | 0 source/isaaclab_newton/changelog.d/asset-tests-readability.skip | 0 .../antoiner-asset-tests-p06-ov-articulation-core.skip | 0 .../antoiner-asset-tests-p07-ov-articulation-regressions.skip | 0 .../changelog.d/antoiner-asset-tests-p08-ov-rigid-object.skip | 0 .../antoiner-asset-tests-p09-ov-rigid-object-collection.skip | 0 .../changelog.d/antoiner-asset-tests-p10-ov-deformable.skip | 0 ...ager.changed.rst => antoiner-asset-tests-redesign.changed.rst} | 0 ...-manager.fixed.rst => antoiner-asset-tests-redesign.fixed.rst} | 0 source/isaaclab_ov/changelog.d/asset-tests-readability.skip | 0 .../antoiner-asset-tests-p11-physx-articulation-core.skip | 0 .../antoiner-asset-tests-p12-physx-articulation-regressions.skip | 0 .../antoiner-asset-tests-p13-physx-newton-actuators.skip | 0 .../changelog.d/antoiner-asset-tests-p14-physx-rigid-object.skip | 0 .../antoiner-asset-tests-p15-physx-rigid-object-collection.skip | 0 .../changelog.d/antoiner-asset-tests-p16-physx-deformable.skip | 0 .../antoiner-asset-tests-p17-physx-surface-gripper.skip | 0 source/isaaclab_physx/changelog.d/asset-tests-readability.skip | 0 30 files changed, 0 insertions(+), 0 deletions(-) delete mode 100644 source/isaaclab/changelog.d/antoiner-asset-tests-p03-physx-actuator-runtime.skip delete mode 100644 source/isaaclab/changelog.d/antoiner-asset-tests-p04-terminations.skip delete mode 100644 source/isaaclab/changelog.d/antoiner-asset-tests-p13-physx-newton-actuators.skip rename source/isaaclab/changelog.d/{antoiner-asset-tests-p01-contract.skip => antoiner-asset-tests-redesign.skip} (100%) delete mode 100644 source/isaaclab/changelog.d/asset-tests-readability.skip delete mode 100644 source/isaaclab_newton/changelog.d/antoiner-asset-tests-p18-newton-fixtures.skip delete mode 100644 source/isaaclab_newton/changelog.d/antoiner-asset-tests-p19-newton-task-space-controllers.skip delete mode 100644 source/isaaclab_newton/changelog.d/antoiner-asset-tests-p20-newton-actuators.skip delete mode 100644 source/isaaclab_newton/changelog.d/antoiner-asset-tests-p21-newton-rigid-object.skip delete mode 100644 source/isaaclab_newton/changelog.d/antoiner-asset-tests-p22-newton-rigid-object-collection.skip delete mode 100644 source/isaaclab_newton/changelog.d/antoiner-asset-tests-p23-newton-articulation-core.skip delete mode 100644 source/isaaclab_newton/changelog.d/antoiner-asset-tests-p24-newton-articulation-regressions.skip rename source/{isaaclab/changelog.d/antoiner-asset-tests-p02-wrench-composer.skip => isaaclab_newton/changelog.d/antoiner-asset-tests-redesign.skip} (100%) delete mode 100644 source/isaaclab_newton/changelog.d/asset-tests-readability.skip delete mode 100644 source/isaaclab_ov/changelog.d/antoiner-asset-tests-p06-ov-articulation-core.skip delete mode 100644 source/isaaclab_ov/changelog.d/antoiner-asset-tests-p07-ov-articulation-regressions.skip delete mode 100644 source/isaaclab_ov/changelog.d/antoiner-asset-tests-p08-ov-rigid-object.skip delete mode 100644 source/isaaclab_ov/changelog.d/antoiner-asset-tests-p09-ov-rigid-object-collection.skip delete mode 100644 source/isaaclab_ov/changelog.d/antoiner-asset-tests-p10-ov-deformable.skip rename source/isaaclab_ov/changelog.d/{antoiner-asset-tests-p05-ov-manager.changed.rst => antoiner-asset-tests-redesign.changed.rst} (100%) rename source/isaaclab_ov/changelog.d/{antoiner-asset-tests-p05-ov-manager.fixed.rst => antoiner-asset-tests-redesign.fixed.rst} (100%) delete mode 100644 source/isaaclab_ov/changelog.d/asset-tests-readability.skip delete mode 100644 source/isaaclab_physx/changelog.d/antoiner-asset-tests-p11-physx-articulation-core.skip delete mode 100644 source/isaaclab_physx/changelog.d/antoiner-asset-tests-p12-physx-articulation-regressions.skip delete mode 100644 source/isaaclab_physx/changelog.d/antoiner-asset-tests-p13-physx-newton-actuators.skip delete mode 100644 source/isaaclab_physx/changelog.d/antoiner-asset-tests-p14-physx-rigid-object.skip delete mode 100644 source/isaaclab_physx/changelog.d/antoiner-asset-tests-p15-physx-rigid-object-collection.skip delete mode 100644 source/isaaclab_physx/changelog.d/antoiner-asset-tests-p16-physx-deformable.skip delete mode 100644 source/isaaclab_physx/changelog.d/antoiner-asset-tests-p17-physx-surface-gripper.skip delete mode 100644 source/isaaclab_physx/changelog.d/asset-tests-readability.skip diff --git a/source/isaaclab/changelog.d/antoiner-asset-tests-p03-physx-actuator-runtime.skip b/source/isaaclab/changelog.d/antoiner-asset-tests-p03-physx-actuator-runtime.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab/changelog.d/antoiner-asset-tests-p04-terminations.skip b/source/isaaclab/changelog.d/antoiner-asset-tests-p04-terminations.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab/changelog.d/antoiner-asset-tests-p13-physx-newton-actuators.skip b/source/isaaclab/changelog.d/antoiner-asset-tests-p13-physx-newton-actuators.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab/changelog.d/antoiner-asset-tests-p01-contract.skip b/source/isaaclab/changelog.d/antoiner-asset-tests-redesign.skip similarity index 100% rename from source/isaaclab/changelog.d/antoiner-asset-tests-p01-contract.skip rename to source/isaaclab/changelog.d/antoiner-asset-tests-redesign.skip diff --git a/source/isaaclab/changelog.d/asset-tests-readability.skip b/source/isaaclab/changelog.d/asset-tests-readability.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p18-newton-fixtures.skip b/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p18-newton-fixtures.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p19-newton-task-space-controllers.skip b/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p19-newton-task-space-controllers.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p20-newton-actuators.skip b/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p20-newton-actuators.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p21-newton-rigid-object.skip b/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p21-newton-rigid-object.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p22-newton-rigid-object-collection.skip b/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p22-newton-rigid-object-collection.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p23-newton-articulation-core.skip b/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p23-newton-articulation-core.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p24-newton-articulation-regressions.skip b/source/isaaclab_newton/changelog.d/antoiner-asset-tests-p24-newton-articulation-regressions.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab/changelog.d/antoiner-asset-tests-p02-wrench-composer.skip b/source/isaaclab_newton/changelog.d/antoiner-asset-tests-redesign.skip similarity index 100% rename from source/isaaclab/changelog.d/antoiner-asset-tests-p02-wrench-composer.skip rename to source/isaaclab_newton/changelog.d/antoiner-asset-tests-redesign.skip diff --git a/source/isaaclab_newton/changelog.d/asset-tests-readability.skip b/source/isaaclab_newton/changelog.d/asset-tests-readability.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p06-ov-articulation-core.skip b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p06-ov-articulation-core.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p07-ov-articulation-regressions.skip b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p07-ov-articulation-regressions.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p08-ov-rigid-object.skip b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p08-ov-rigid-object.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p09-ov-rigid-object-collection.skip b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p09-ov-rigid-object-collection.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p10-ov-deformable.skip b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p10-ov-deformable.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p05-ov-manager.changed.rst b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-redesign.changed.rst similarity index 100% rename from source/isaaclab_ov/changelog.d/antoiner-asset-tests-p05-ov-manager.changed.rst rename to source/isaaclab_ov/changelog.d/antoiner-asset-tests-redesign.changed.rst diff --git a/source/isaaclab_ov/changelog.d/antoiner-asset-tests-p05-ov-manager.fixed.rst b/source/isaaclab_ov/changelog.d/antoiner-asset-tests-redesign.fixed.rst similarity index 100% rename from source/isaaclab_ov/changelog.d/antoiner-asset-tests-p05-ov-manager.fixed.rst rename to source/isaaclab_ov/changelog.d/antoiner-asset-tests-redesign.fixed.rst diff --git a/source/isaaclab_ov/changelog.d/asset-tests-readability.skip b/source/isaaclab_ov/changelog.d/asset-tests-readability.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p11-physx-articulation-core.skip b/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p11-physx-articulation-core.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p12-physx-articulation-regressions.skip b/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p12-physx-articulation-regressions.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p13-physx-newton-actuators.skip b/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p13-physx-newton-actuators.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p14-physx-rigid-object.skip b/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p14-physx-rigid-object.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p15-physx-rigid-object-collection.skip b/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p15-physx-rigid-object-collection.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p16-physx-deformable.skip b/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p16-physx-deformable.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p17-physx-surface-gripper.skip b/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p17-physx-surface-gripper.skip deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/source/isaaclab_physx/changelog.d/asset-tests-readability.skip b/source/isaaclab_physx/changelog.d/asset-tests-readability.skip deleted file mode 100644 index e69de29bb2d1..000000000000 From c5c43d5aaadda5ff9e8f5c8c5e85bbacc9d0e1c9 Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Fri, 2 Oct 2026 16:34:53 -0700 Subject: [PATCH 29/33] Merge contract and actuator runtime test files --- .../test/actuators/test_newton_adapter.py | 88 +- .../actuators/test_physx_actuator_runtime.py | 92 -- .../assets/contract/articulation_factory.py | 486 +++----- .../isaaclab/test/assets/contract/backends.py | 27 + .../isaaclab/test/assets/contract/conftest.py | 23 +- .../assets/contract/rigid_object_factory.py | 72 +- .../contract/test_articulation_contract.py | 1038 +++++++++++++++- .../test_articulation_ordering_contract.py | 1063 ----------------- 8 files changed, 1351 insertions(+), 1538 deletions(-) delete mode 100644 source/isaaclab/test/actuators/test_physx_actuator_runtime.py delete mode 100644 source/isaaclab/test/assets/contract/test_articulation_ordering_contract.py diff --git a/source/isaaclab/test/actuators/test_newton_adapter.py b/source/isaaclab/test/actuators/test_newton_adapter.py index e295d5d4d420..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, @@ -345,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/actuators/test_physx_actuator_runtime.py b/source/isaaclab/test/actuators/test_physx_actuator_runtime.py deleted file mode 100644 index ec77abf5357b..000000000000 --- a/source/isaaclab/test/actuators/test_physx_actuator_runtime.py +++ /dev/null @@ -1,92 +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 CUDA-graph handling of the host-side PhysX actuator runtime.""" - -from types import SimpleNamespace -from unittest.mock import Mock - -import pytest -import warp as wp - -from isaaclab.actuators.newton.physx_runtime import PhysxActuatorRuntime - - -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/contract/articulation_factory.py b/source/isaaclab/test/assets/contract/articulation_factory.py index 07a4dd7e67b6..0a5ed0b754b1 100644 --- a/source/isaaclab/test/assets/contract/articulation_factory.py +++ b/source/isaaclab/test/assets/contract/articulation_factory.py @@ -6,7 +6,12 @@ # ignore private usage of variables warning # pyright: reportPrivateUsage=none -"""Shared mocked articulation backend factories for the contract tests.""" +"""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 @@ -15,26 +20,32 @@ import warp as wp from isaaclab.assets.articulation.articulation_cfg import ArticulationCfg -from isaaclab.utils.wrench_composer import WrenchComposer -from .backends import AVAILABLE, install_physx_recording_setters, patch_ovphysx_manager, patch_physx_manager +from .backends import ( + AVAILABLE, + finish_shell, + indices, + install_physx_recording_setters, + newton_manager, + patch_ovphysx_manager, + patch_physx_manager, +) if "physx" in AVAILABLE: - from isaaclab_physx.assets.articulation.articulation import Articulation as PhysXArticulation - from isaaclab_physx.assets.articulation.articulation_data import ArticulationData as PhysXArticulationData - from isaaclab_physx.test.fixtures.views import MockArticulationViewWarp as PhysXMockArticulationViewWarp + import isaaclab_physx.assets as physx_assets + from isaaclab_physx.test.fixtures.views import MockArticulationViewWarp if "newton" in AVAILABLE: - 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 + 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: - from isaaclab_ov.assets.articulation.articulation import Articulation as OvPhysxArticulation - from isaaclab_ov.assets.articulation.articulation_data import ArticulationData as OvPhysxArticulationData + import isaaclab_ov.assets as ovphysx_assets + from isaaclab_ov import tensor_types as TT from isaaclab_ov.test.fixtures.views import MockOvPhysxBindingSet -_PHYSX_ARTICULATION_STORAGE = { +_PHYSX_STORAGE = { "set_root_transforms": "_root_transforms", "set_root_velocities": "_root_velocities", "set_dof_positions": "_dof_positions", @@ -56,28 +67,13 @@ } -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, - *, - monkeypatch: pytest.MonkeyPatch, -): - """Create a test Articulation instance with mocked dependencies.""" - patch_physx_manager(monkeypatch=monkeypatch) - 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)] +def _names(prefix: str, count: int) -> list[str]: + return [f"{prefix}_{i}" for i in range(count)] - articulation = object.__new__(PhysXArticulation) +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, @@ -87,345 +83,170 @@ def create_physx_articulation( ) articulation._sim_cfg = None articulation._fixed_tendon_target_dirty = False + return articulation - # 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() - install_physx_recording_setters(mock_view, _PHYSX_ARTICULATION_STORAGE) - - # 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 - mock_view._shared_metatype = mock_metatype - - articulation._root_view = mock_view - 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. - - # The data reads gravity from the PhysX manager patched above. - data = PhysXArticulationData(mock_view, device) - articulation._data = data - # Set tendon names on articulation and data - articulation._fixed_tendon_names = fixed_tendon_names - articulation._spatial_tendon_names = spatial_tendon_names - data.fixed_tendon_names = fixed_tendon_names - data.spatial_tendon_names = spatial_tendon_names +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 - # 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) - articulation._instantaneous_wrench_composer = mock_inst_wrench - articulation._permanent_wrench_composer = mock_perm_wrench - # Prevent __del__ / _clear_callbacks from raising AttributeError - articulation._initialize_handle = None - articulation._invalidate_initialize_handle = None - articulation._prim_deletion_handle = None - articulation._debug_vis_handle = None - - # Set up other required attributes - articulation._ALL_INDICES = wp.array(np.arange(num_instances, dtype=np.int32), device=device) - articulation._ALL_BODY_INDICES = wp.array(np.arange(num_bodies, dtype=np.int32), device=device) - articulation._ALL_JOINT_INDICES = wp.array(np.arange(num_joints, dtype=np.int32), device=device) - - # Tendon index arrays - articulation._ALL_FIXED_TENDON_INDICES = wp.array(np.arange(num_fixed_tendons, dtype=np.int32), device=device) - 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 - articulation._ALL_INDICES_WP = wp.array(np.arange(num_instances, dtype=np.int32), device=device) - articulation._ALL_BODY_INDICES_WP = wp.array(np.arange(num_bodies, dtype=np.int32), device=device) +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._joint_pos_target_backend = None - articulation._joint_vel_target_backend = None - articulation._joint_effort_target_backend = None - articulation._body_wrench_force_backend = None + 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 - articulation._root_link_pose_w_f32 = None - articulation._root_com_vel_w_f32 = None - articulation._root_link_vel_w_f32 = None - articulation._inst_wrench_force_f32 = None - articulation._inst_wrench_torque_f32 = None - articulation._perm_wrench_force_f32 = None - articulation._perm_wrench_torque_f32 = None - - # Pre-allocated pinned CPU buffers for PhysX TensorAPI writes - N, J, B = num_instances, num_joints, num_bodies + # 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 = {} - cpu_env_ids = wp.array(np.arange(N, dtype=np.int32), device="cpu") - articulation._cpu_env_ids_all = cpu_env_ids + 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 = {} - articulation._cpu_joint_stiffness = wp.zeros((N, J), dtype=wp.float32, device="cpu") - articulation._cpu_joint_damping = wp.zeros((N, J), dtype=wp.float32, device="cpu") - articulation._cpu_joint_pos_limits = wp.zeros((N, J, 2), dtype=wp.float32, device="cpu") - articulation._cpu_joint_vel_limits = wp.zeros((N, J), dtype=wp.float32, device="cpu") - articulation._cpu_joint_effort_limits = wp.zeros((N, J), dtype=wp.float32, device="cpu") - articulation._cpu_joint_armature = wp.zeros((N, J), dtype=wp.float32, device="cpu") - articulation._cpu_joint_friction_props = wp.zeros((N, J, 3), dtype=wp.float32, device="cpu") - articulation._cpu_body_mass = wp.zeros((N, B), dtype=wp.float32, device="cpu") - articulation._cpu_body_coms = wp.zeros((N, B, 7), dtype=wp.float32, device="cpu") - articulation._cpu_body_inertia = wp.zeros((N, B, 9), dtype=wp.float32, device="cpu") - + 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 - 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, - *, - monkeypatch: pytest.MonkeyPatch, -): - """Create a test OvPhysX Articulation instance with mocked tensor bindings.""" - patch_ovphysx_manager(monkeypatch=monkeypatch) - 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, - ) - articulation._sim_cfg = None - 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, +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=num_fixed_tendons, - num_spatial_tendons=num_spatial_tendons, + num_fixed_tendons=FT, + num_spatial_tendons=ST, ) - 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)] + bindings.set_random_data() + articulation = _shell(ovphysx_assets.Articulation, joint_ordering, body_ordering) articulation._device = device articulation._ovphysx = MagicMock() - articulation._root_view = mock_bindings.view - articulation._bindings = mock_bindings.bindings - articulation._num_instances = num_instances - articulation._num_joints = num_joints - articulation._num_bodies = num_bodies + 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._fixed_tendon_names = fixed_tendon_names - articulation._spatial_tendon_names = spatial_tendon_names - articulation._num_fixed_tendons = num_fixed_tendons - 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 - 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._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() - - # Wrench composers - mock_inst_wrench = WrenchComposer(articulation, supports_world_at_com=True) - mock_perm_wrench = WrenchComposer(articulation, supports_world_at_com=True) - articulation._instantaneous_wrench_composer = mock_inst_wrench - articulation._permanent_wrench_composer = mock_perm_wrench - # Prevent __del__ / _clear_callbacks from raising - articulation._initialize_handle = None - articulation._invalidate_initialize_handle = None - articulation._prim_deletion_handle = None - articulation._debug_vis_handle = None + finish_shell(articulation) articulation._process_actuators_cfg() - - from isaaclab_ov import tensor_types as TT - 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, mock_bindings + return articulation, 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, - *, - monkeypatch: pytest.MonkeyPatch, -): - """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, - num_tendons=num_fixed_tendons, +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=joint_names, - body_names=body_names, + joint_names=_names("joint", J), + body_names=_names("body", B), tendon_names=fixed_tendon_names, ) - mock_view.set_random_mock_data() + view.set_random_mock_data() # MuJoCo tendon properties the data binds when the articulation has fixed tendons, with distinct values. - tendon_shape = (num_instances, 1, num_fixed_tendons) - tendon_values = np.arange(1, num_instances * num_fixed_tendons + 1, dtype=np.float32).reshape(tendon_shape) + 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) - for name, values, dtype in ( - ("mujoco.tendon_stiffness", tendon_values, wp.float32), - ("mujoco.tendon_damping", tendon_values + 0.5, wp.float32), - ("mujoco.tendon_range", tendon_range, wp.vec2f), - ): - mock_view._attributes[name] = wp.array(values, dtype=dtype, device=device) - - # 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 until the test finishes. - monkeypatch.setattr(newton_data_module, "SimulationManager", mock_manager, raising=False) - data = NewtonArticulationData(mock_view, device) - - # 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, - ) - articulation._sim_cfg = None - articulation._fixed_tendon_target_dirty = False - - articulation._root_view = mock_view + 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 = data - articulation._test_simulation_manager = mock_manager - # the solver builds this adapter; the shell has no model, so it stays absent + 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 - - # Newton supports fixed tendons but not spatial tendons. - articulation._fixed_tendon_names = fixed_tendon_names - 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) - articulation._instantaneous_wrench_composer = mock_inst_wrench - articulation._permanent_wrench_composer = mock_perm_wrench - - # Prevent __del__ / _clear_callbacks from raising AttributeError - articulation._initialize_handle = None - articulation._invalidate_initialize_handle = None - articulation._prim_deletion_handle = None - articulation._debug_vis_handle = None - - # Newton uses wp.array for indices (not torch) - articulation._ALL_INDICES = wp.array(np.arange(num_instances, dtype=np.int32), device=device) - articulation._ALL_BODY_INDICES = wp.array(np.arange(num_bodies, dtype=np.int32), device=device) - articulation._ALL_JOINT_INDICES = wp.array(np.arange(num_joints, dtype=np.int32), device=device) - - # Newton uses wp.bool masks - articulation._ALL_ENV_MASK = wp.ones((num_instances,), dtype=wp.bool, device=device) - articulation._ALL_BODY_MASK = wp.ones((num_bodies,), dtype=wp.bool, device=device) - articulation._ALL_JOINT_MASK = wp.ones((num_joints,), dtype=wp.bool, device=device) - + _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() - - # Tendon arrays - articulation._ALL_FIXED_TENDON_INDICES = wp.array(np.arange(num_fixed_tendons, dtype=np.int32), device=device) - articulation._ALL_FIXED_TENDON_MASK = wp.ones((num_fixed_tendons,), dtype=wp.bool, device=device) - articulation._ALL_SPATIAL_TENDON_INDICES = wp.array(np.array([], dtype=np.int32), device=device) + 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, mock_view + return articulation, view def get_articulation( @@ -436,12 +257,13 @@ def get_articulation( num_fixed_tendons: int = 0, num_spatial_tendons: int = 0, device: str = "cuda:0", - **kwargs, + 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": create_physx_articulation, - "ovphysx": create_ovphysx_articulation, - "newton": create_newton_articulation, - }[backend] - return factory(num_instances, num_joints, num_bodies, num_fixed_tendons, num_spatial_tendons, device, **kwargs) + 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 index a2650715070a..bbd5bdfda8de 100644 --- a/source/isaaclab/test/assets/contract/backends.py +++ b/source/isaaclab/test/assets/contract/backends.py @@ -20,6 +20,8 @@ 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 @@ -160,3 +162,28 @@ def read_backend_joint_state(backend: str, art, raw_backend) -> tuple[np.ndarray 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 index 1e7989bb5bb0..16e2a0e7acce 100644 --- a/source/isaaclab/test/assets/contract/conftest.py +++ b/source/isaaclab/test/assets/contract/conftest.py @@ -6,6 +6,7 @@ """Backend and device parameters shared by the asset contracts.""" import pytest +import warp as wp from isaaclab.test.utils import DeviceScope, test_devices @@ -20,5 +21,25 @@ def backend(request: pytest.FixtureRequest) -> str: @pytest.fixture(params=test_devices(DeviceScope.CPU_AND_DEFAULT_CUDA)) def device(request: pytest.FixtureRequest) -> str: - """Exercise CPU buffers and CUDA staging.""" + """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 index 3b22f9c78424..b5a5530327f0 100644 --- a/source/isaaclab/test/assets/contract/rigid_object_factory.py +++ b/source/isaaclab/test/assets/contract/rigid_object_factory.py @@ -14,15 +14,21 @@ from unittest.mock import MagicMock -import numpy as np 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 isaaclab.utils.wrench_composer import WrenchComposer -from .backends import AVAILABLE, install_physx_recording_setters, patch_ovphysx_manager, patch_physx_manager +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 @@ -48,30 +54,6 @@ } -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 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 - - def _physx_rigid_object(num_instances: int, device: str, monkeypatch: pytest.MonkeyPatch): patch_physx_manager(monkeypatch) view = MockRigidBodyViewWarp(count=num_instances, device=device) @@ -84,16 +66,16 @@ def _physx_rigid_object(num_instances: int, device: str, monkeypatch: pytest.Mon 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) + 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_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") @@ -113,16 +95,16 @@ def _newton_rigid_object(num_instances: int, device: str, monkeypatch: pytest.Mo body_names=["body_0"], ) view.set_random_mock_data() - monkeypatch.setattr(newton_rigid_object_data, "SimulationManager", _newton_manager(num_instances, device)) + 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) + 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 @@ -149,7 +131,7 @@ def _ovphysx_rigid_object(num_instances: int, device: str, monkeypatch: pytest.M obj._data.num_bodies = 1 obj._data._is_primed = True obj._create_buffers() - _finish_shell(obj) + finish_shell(obj) return obj, bindings @@ -177,10 +159,10 @@ def _physx_collection(num_instances: int, num_bodies: int, device: str, monkeypa 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) + 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 = {} @@ -195,7 +177,7 @@ def _newton_collection(num_instances: int, num_bodies: int, device: str, monkeyp 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) + manager = newton_manager(num_instances, device) monkeypatch.setattr(newton_collection_data, "SimulationManager", manager) monkeypatch.setattr(newton_collection, "SimulationManager", manager) @@ -208,9 +190,9 @@ def _newton_collection(num_instances: int, num_bodies: int, device: str, monkeyp 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) + 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 @@ -243,7 +225,7 @@ def _ovphysx_collection(num_instances: int, num_bodies: int, device: str, monkey collection._data.num_bodies = num_bodies collection._data._is_primed = True collection._create_buffers() - _finish_shell(collection) + finish_shell(collection) return collection, bindings diff --git a/source/isaaclab/test/assets/contract/test_articulation_contract.py b/source/isaaclab/test/assets/contract/test_articulation_contract.py index 3ae0e19e691b..83f141148791 100644 --- a/source/isaaclab/test/assets/contract/test_articulation_contract.py +++ b/source/isaaclab/test/assets/contract/test_articulation_contract.py @@ -8,15 +8,16 @@ """Shared articulation contracts across the production backends. -Checks that every articulation backend provides the data and writer 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 runs on -CPU only; getters and writers also run on CUDA because PhysX stages through CPU-pinned buffers there. +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 +from unittest.mock import MagicMock, patch import numpy as np import pytest @@ -24,6 +25,7 @@ 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 @@ -747,3 +749,1031 @@ def capture(name, values, *, indices=None, mask=None): 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_articulation_ordering_contract.py b/source/isaaclab/test/assets/contract/test_articulation_ordering_contract.py deleted file mode 100644 index 02f76ce76107..000000000000 --- a/source/isaaclab/test/assets/contract/test_articulation_ordering_contract.py +++ /dev/null @@ -1,1063 +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 - -"""Cross-backend articulation joint- and body-ordering contracts on mocked backends. - -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. -""" - -from unittest.mock import MagicMock, patch - -import numpy as np -import pytest -import torch -import warp as wp - -from isaaclab.utils import replace - -from .articulation_factory import get_articulation -from .backends import backends, read_backend_joint_state, requires - -pytestmark = pytest.mark.integration - - -@pytest.fixture(autouse=True) -def _cpu_warp_device(): - """Every ordering case runs on CPU; allocate unqualified Warp arrays there too.""" - with wp.ScopedDevice("cpu"): - yield - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -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 _mask(total: int, selected, device: str) -> wp.array: - mask = np.zeros(total, dtype=bool) - mask[list(selected)] = True - return wp.array(mask, dtype=wp.bool, device=device) - - -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() From 14dc49da9efef4ec3019452f9f8a43d45aa3b51a Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Fri, 2 Oct 2026 16:34:53 -0700 Subject: [PATCH 30/33] Consolidate Newton asset tests into three files --- .../newton_test_utils.py} | 4 +- .../test/assets/test_articulation.py | 407 +++++++++++- .../test_articulation_ordering_kernels.py | 151 ----- .../assets/test_newton_actuators_newton.py | 296 ++++++++- .../test/assets/test_rigid_object.py | 526 +++++++++++++++- .../assets/test_rigid_object_collection.py | 579 ------------------ .../unit/test_articulation_accessors.py | 104 ---- .../unit/test_newton_actuator_adaptation.py | 295 --------- .../test_newton_task_space_controllers.py | 239 -------- 9 files changed, 1177 insertions(+), 1424 deletions(-) rename source/isaaclab_newton/test/{articulation_test_utils.py => assets/newton_test_utils.py} (97%) delete mode 100644 source/isaaclab_newton/test/assets/test_articulation_ordering_kernels.py delete mode 100644 source/isaaclab_newton/test/assets/test_rigid_object_collection.py delete mode 100644 source/isaaclab_newton/test/assets/unit/test_articulation_accessors.py delete mode 100644 source/isaaclab_newton/test/assets/unit/test_newton_actuator_adaptation.py delete mode 100644 source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py diff --git a/source/isaaclab_newton/test/articulation_test_utils.py b/source/isaaclab_newton/test/assets/newton_test_utils.py similarity index 97% rename from source/isaaclab_newton/test/articulation_test_utils.py rename to source/isaaclab_newton/test/assets/newton_test_utils.py index b95e46621f04..b98d6b67e715 100644 --- a/source/isaaclab_newton/test/articulation_test_utils.py +++ b/source/isaaclab_newton/test/assets/newton_test_utils.py @@ -3,7 +3,7 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Locally authored scenes shared by the kitless Newton asset and controller tests. +"""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 @@ -37,7 +37,7 @@ Vec3 = tuple[float, float, float] -_DATA_DIR = Path(__file__).resolve().parent / "assets" / "data" +_DATA_DIR = Path(__file__).resolve().parent / "data" """Directory holding the checked-in USD fixtures.""" _ENV_SPACING = 40.0 diff --git a/source/isaaclab_newton/test/assets/test_articulation.py b/source/isaaclab_newton/test/assets/test_articulation.py index 74d4af973ffb..89efea0f137a 100644 --- a/source/isaaclab_newton/test/assets/test_articulation.py +++ b/source/isaaclab_newton/test/assets/test_articulation.py @@ -35,18 +35,31 @@ from collections.abc import Callable, Iterator from copy import copy from dataclasses import dataclass -from pathlib import Path 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 import kernels as articulation_kernels +from isaaclab_newton.assets.articulation.articulation_data import ArticulationData from isaaclab_newton.physics import NewtonManager as SimulationManager from isaaclab_physx.sim.schemas import PhysxJointCfg -from newton import Model, ModelFlags, ShapeFlags, State +from newton import JointType, Model, ModelBuilder, ModelFlags, ShapeFlags, State +from newton.selection import ArticulationView +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 @@ -56,6 +69,7 @@ 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.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 @@ -68,17 +82,6 @@ ) from isaaclab.utils import replace -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -from articulation_test_utils import ( # noqa: E402 - WRIST_USD_STIFFNESS, - env_origins, - local_usd, - newton_sim_cfg, - spawn_assets, - world_gravity, -) - pytestmark = [pytest.mark.integration, pytest.mark.kitless] _GRAVITY = (0.0, 0.0, -9.81) @@ -219,6 +222,14 @@ def island( 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}), + ), } @@ -269,6 +280,214 @@ def rest(self, articulation: Articulation) -> None: articulation.write_data_to_sim() +## +# Kernels and model-level accessors. These build no simulation context. +## + + +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.""" + + 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_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", + ) + + 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: + """Align Newton shape counts with the public body-name axis.""" + + class _ShapeCountSurface: + backend_num_shapes_per_body = Articulation.backend_num_shapes_per_body + num_shapes_per_body = Articulation.num_shapes_per_body + + articulation = _ShapeCountSurface() + articulation._num_shapes_per_body_backend = None + articulation._root_view = SimpleNamespace( + body_shapes=((), (object(), object()), (object(), object(), object())), + ) + articulation.body_ordering = SimpleNamespace( + user_to_backend_indices=(2, 0, 1), + ) + + assert articulation.num_shapes_per_body == [3, 0, 2] + + +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +@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 + 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") + + 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 + + ## # Own scenes. These tests run before the composite scene exists: only one simulation context can be alive. ## @@ -707,6 +926,7 @@ def test_articulation_initialization_and_partial_state(scene: _Scene) -> None: "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. @@ -2034,3 +2254,164 @@ def test_joint_position_limit_clamping_respects_logging(scene: _Scene, caplog) - 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 + + +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", + 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=NUM_ENVS, + device=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(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_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 = 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( + jacobian_b=jacobian_b, + current_ee_pose_b=ee_pose_b, + current_ee_vel_b=ee_vel_b, + mass_matrix=mass_matrix, + gravity=gravity_forces, + ) + robot.actuators.target_command.set_effort_index(value=joint_efforts, joint_ids=arm_joint_ids) + scene.step("osc") + + 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. + + 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.isaacsim_ci +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_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, _, _ = _osc_chain(scene) + osc = _make_osc(scene.device) + # Hold the initial EE pose: phase-1 steady-state error is pure gravity sag. + 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 _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[-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 + + 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") + + 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)" 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 fc7fb6bd445e..000000000000 --- a/source/isaaclab_newton/test/assets/test_articulation_ordering_kernels.py +++ /dev/null @@ -1,151 +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 - -pytestmark = pytest.mark.unit - - -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 683e6746f7a8..025371088a75 100644 --- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py +++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py @@ -27,22 +27,35 @@ launch_test_simulation(SimulationCfg(physics=NewtonCfg())) import os -import sys from collections.abc import Iterator from dataclasses import dataclass -from pathlib import Path +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 +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 isaaclab.actuators import ActuatorBaseCfg, DCMotorCfg, DelayedPDActuatorCfg, IdealPDActuatorCfg +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 ( + 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 ArticulationCfg from isaaclab.sim import SimulationContext, build_simulation_context from isaaclab.test.utils import DeviceScope, test_devices @@ -62,10 +75,6 @@ from isaaclab_assets.robots.spot import joint_parameter_lookup as SPOT_KNEE_LOOKUP -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -from articulation_test_utils import NUM_ENVS, local_usd, newton_sim_cfg, spawn_assets # noqa: E402 - pytestmark = [pytest.mark.integration, pytest.mark.kitless] # --------------------------------------------------------------------------- @@ -477,6 +486,275 @@ def _assert_equivalent( oracle.test_computed_effort_match() +# --------------------------------------------------------------------------- +# Adaptation on a model builder and on Warp arrays; these step no simulation and run before the scenes exist +# --------------------------------------------------------------------------- + + +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. + + 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)] + + +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)] + + +@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] + + +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)) + + # --------------------------------------------------------------------------- # Equivalence tests with different actuator types # --------------------------------------------------------------------------- diff --git a/source/isaaclab_newton/test/assets/test_rigid_object.py b/source/isaaclab_newton/test/assets/test_rigid_object.py index ffb69b1f2026..fac10ee7bdfa 100644 --- a/source/isaaclab_newton/test/assets/test_rigid_object.py +++ b/source/isaaclab_newton/test/assets/test_rigid_object.py @@ -6,11 +6,13 @@ # ignore private usage of variables warning # pyright: reportPrivateUsage=none -"""Kitless real-Newton rigid-object coverage on one persistent two-environment CPU scene. +"""Kitless real-Newton rigid-object and rigid-object-collection coverage on one persistent two-environment CPU scene. -The cubes are locally spawned shapes. 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 case 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. +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 @@ -23,18 +25,19 @@ import sys from collections.abc import Iterator from dataclasses import dataclass -from pathlib import Path from types import SimpleNamespace import numpy as np import pytest import torch import warp as wp -from isaaclab_newton.assets import RigidObject +from isaaclab_newton.assets import RigidObject, RigidObjectCollection from isaaclab_newton.physics import NewtonManager as SimulationManager +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 RigidObjectCfg +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 @@ -47,12 +50,9 @@ quat_mul, quat_rotate, random_orientation, + subtract_frame_transforms, ) -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -from articulation_test_utils import env_origins, local_usd, newton_sim_cfg, spawn_assets, world_gravity # noqa: E402 - pytestmark = [pytest.mark.integration, pytest.mark.kitless] _GRAVITY = (0.0, 0.0, -9.81) @@ -68,7 +68,41 @@ def _cube_cfg(height: float = 1.0) -> RigidObjectCfg: mass_props=sim_utils.MassCfg(mass=1.0), collision_props=sim_utils.UsdPhysicsCollisionCfg(), ), - init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, height)), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, -3.0, height)), + ) + + +_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 + 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(), + ) + + +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) + } ) @@ -123,6 +157,35 @@ def test_single_instance_initialization_and_root_state_write(device: str) -> Non torch.testing.assert_close(cube_object.data.body_link_pose_w.torch.squeeze(1), target_pose) +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"] + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(object_collection) < 10 + + with pytest.raises(RuntimeError, match="Expected 1 prims at"): + sim.reset() + + +@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. ## @@ -130,28 +193,28 @@ def test_single_instance_initialization_and_root_state_write(device: str) -> Non @dataclass class _Scene: - """Two-environment cubes that share one real Newton model and solver lifecycle.""" + """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 - masses: torch.Tensor - """Configured body masses [kg], restored before each test.""" - coms: torch.Tensor - """Configured center-of-mass positions in the body frames [m], restored before each test.""" - inertias: torch.Tensor - """Configured body inertias [kg·m^2], restored before each test.""" + 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 cube commands before and updating the cubes after each step.""" + """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 the cubes at rest at their configured pose.""" + """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) @@ -159,32 +222,56 @@ def rest(self) -> None: 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 once for this module.""" + """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: - cube = spawn_assets({"cube": _cube_cfg()})["cube"] + 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=cube, + cube=assets["cube"], + collection=assets["collection"], origins=env_origins(device), device=device, - masses=cube.data.body_mass.torch.clone(), - coms=cube.data.body_com_pos_b.torch.clone(), - inertias=cube.data.body_inertia.torch.clone(), + 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.fixture def scene(shared_scene: _Scene) -> _Scene: """Hand each test the shared scene at rest with the configured inertial properties.""" - shared_scene.cube.set_masses_index(masses=shared_scene.masses) - shared_scene.cube.set_coms_index(coms=shared_scene.coms) - shared_scene.cube.set_inertias_index(inertias=shared_scene.inertias) + 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 @@ -235,8 +322,10 @@ def test_rigid_body_set_material_properties(scene: _Scene) -> None: 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(body_world == index))) for index in range(num_cubes) + 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() @@ -390,7 +479,7 @@ def test_body_root_state_properties(scene: _Scene) -> None: cube_object = scene.cube num_cubes = cube_object.num_instances device = scene.device - env_pos = scene.origins + torch.tensor([0.0, 0.0, 1.0], device=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) @@ -624,3 +713,376 @@ def _fk_reset_mask_dirty() -> bool: # 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:]) + + +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, + ) + 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 bb1f955ce3fc..000000000000 --- a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py +++ /dev/null @@ -1,579 +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 - -"""Kitless real-Newton rigid-object-collection coverage on one persistent CPU scene. - -Every environment holds three locally spawned cubes in the collection and an unrelated rigid body beside them, -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 -case 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 pathlib import Path - -import pytest -import torch -import warp as wp -from isaaclab_newton.assets import RigidObject, RigidObjectCollection -from isaaclab_newton.physics import NewtonManager as SimulationManager -from newton import ModelFlags - -import isaaclab.sim as sim_utils -from isaaclab.assets import RigidObjectCfg, RigidObjectCollectionCfg -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, - quat_inv, - quat_mul, - quat_rotate, - random_orientation, - subtract_frame_transforms, -) - -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -from articulation_test_utils import env_origins, newton_sim_cfg, spawn_assets, world_gravity # noqa: E402 - -pytestmark = [pytest.mark.integration, pytest.mark.kitless] - -_GRAVITY = (0.0, 0.0, -9.81) -_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 - 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(), - ) - - -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) - } - ) - - -## -# Own scenes. These tests run before the shared scene exists: only one simulation context can be alive. -## - - -def test_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"] - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(object_collection) < 10 - - with pytest.raises(RuntimeError, match="Expected 1 prims at"): - sim.reset() - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) -def test_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: - """Cube collections that share one real Newton model and solver lifecycle.""" - - sim: SimulationContext - collection: RigidObjectCollection - sibling: RigidObject - origins: torch.Tensor - device: str - masses: torch.Tensor - """Configured body masses [kg], restored before each test.""" - coms: torch.Tensor - """Configured center-of-mass positions in the body frames [m], restored before each test.""" - inertias: torch.Tensor - """Configured body inertias [kg·m^2], restored before each test.""" - - def step(self, num_steps: int = 1) -> None: - """Step ``num_steps`` times, writing the collection's commands before and updating it after each step.""" - for _ in range(num_steps): - self.collection.write_data_to_sim() - self.sim.step() - self.collection.update(self.sim.cfg.dt) - - def rest(self) -> None: - """Put the cubes and the sibling body at rest at their configured poses.""" - 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() - sibling_pose = self.sibling.data.default_root_pose.torch.clone() - sibling_pose[:, :3] += self.origins - self.sibling.write_root_pose_to_sim_index(root_pose=sibling_pose) - self.sibling.write_root_velocity_to_sim_index( - root_velocity=torch.zeros_like(self.sibling.data.default_root_vel.torch) - ) - self.sibling.reset() - - -@pytest.fixture(scope="module", params=test_devices(DeviceScope.CPU)) -def shared_scene(request: pytest.FixtureRequest) -> Iterator[_Scene]: - """Initialize the collection and its unrelated sibling bodies once for this module.""" - device = request.param - sibling_cfg = RigidObjectCfg( - prim_path="/World/Env_[^/]*/Object_Target", - spawn=_cube_spawn_cfg(), - init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, -3.0, 1.0)), - ) - with build_simulation_context(sim_cfg=newton_sim_cfg(device)) as sim: - assets = spawn_assets( - { - "collection": _collection_cfg( - _NUM_CUBES, object_paths=("Object_0", "RobotLeft/Body", "RobotRight/Body") - ), - "sibling": sibling_cfg, - } - ) - sim.reset() - collection = assets["collection"] - yield _Scene( - sim=sim, - collection=collection, - sibling=assets["sibling"], - origins=env_origins(device), - device=device, - masses=collection.data.body_mass.torch.clone(), - coms=collection.data.body_com_pos_b.torch.clone(), - inertias=collection.data.body_inertia.torch.clone(), - ) - - -@pytest.fixture -def scene(shared_scene: _Scene) -> _Scene: - """Hand each test the shared scene at rest with the configured inertial properties.""" - shared_scene.collection.set_masses_index(masses=shared_scene.masses) - shared_scene.collection.set_coms_index(coms=shared_scene.coms) - shared_scene.collection.set_inertias_index(inertias=shared_scene.inertias) - shared_scene.rest() - return shared_scene - - -def test_initialization(scene: _Scene) -> None: - """Test initialization for prims with rigid body API at the provided prim paths. - - With an unrelated rigid body next to the cubes 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.sibling.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 sibling. - 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_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_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:]) - - -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, - ) - 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/unit/test_articulation_accessors.py b/source/isaaclab_newton/test/assets/unit/test_articulation_accessors.py deleted file mode 100644 index 609a5fa1bf23..000000000000 --- a/source/isaaclab_newton/test/assets/unit/test_articulation_accessors.py +++ /dev/null @@ -1,104 +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 - -"""Newton articulation accessors checked on a directly built model: lazy task-space allocation with graph capture, -and the public body order of the collision-shape counts. -""" - -from types import SimpleNamespace - -import newton -import numpy as np -import pytest -import warp as wp -from isaaclab_newton.assets import Articulation -from isaaclab_newton.assets.articulation.articulation_data import ArticulationData -from isaaclab_newton.physics import NewtonManager as SimulationManager -from newton import JointType, ModelBuilder -from newton.selection import ArticulationView - -from isaaclab.test.utils import DeviceScope, test_devices - -pytestmark = pytest.mark.unit - - -def test_num_shapes_per_body_follows_public_body_order() -> None: - """Align Newton shape counts with the public body-name axis.""" - - class _ShapeCountSurface: - backend_num_shapes_per_body = Articulation.backend_num_shapes_per_body - num_shapes_per_body = Articulation.num_shapes_per_body - - articulation = _ShapeCountSurface() - articulation._num_shapes_per_body_backend = None - articulation._root_view = SimpleNamespace( - body_shapes=((), (object(), object()), (object(), object(), object())), - ) - articulation.body_ordering = SimpleNamespace( - user_to_backend_indices=(2, 0, 1), - ) - - assert articulation.num_shapes_per_body == [3, 0, 2] - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -@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 - 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") - - 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 diff --git a/source/isaaclab_newton/test/assets/unit/test_newton_actuator_adaptation.py b/source/isaaclab_newton/test/assets/unit/test_newton_actuator_adaptation.py deleted file mode 100644 index fbc11daab888..000000000000 --- a/source/isaaclab_newton/test/assets/unit/test_newton_actuator_adaptation.py +++ /dev/null @@ -1,295 +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 - -"""Newton actuator adaptation: target modes resolved on the model builder, native-group selection, and telemetry. - -These run on a model builder and on Warp arrays directly; none of them steps a simulation. -""" - -from types import SimpleNamespace - -import numpy as np -import pytest -import warp as wp -from isaaclab_newton.assets import Articulation -from isaaclab_newton.assets.articulation.actuator_control import NewtonActuatorControl -from isaaclab_newton.physics import NewtonBuilderCfg, NewtonCfg -from isaaclab_newton.physics import NewtonManager as SimulationManager -from newton import JointTargetMode, JointType, ModelBuilder -from newton.solvers import SolverMuJoCo - -import isaaclab.sim as sim_utils -from isaaclab.actuators import IdealPDActuatorCfg, ImplicitActuator, ImplicitActuatorCfg -from isaaclab.actuators.newton.kernels import sync_torque_telemetry -from isaaclab.assets import ArticulationCfg - -pytestmark = pytest.mark.unit - - -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. - - 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)] - - -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)] - - -@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] - - -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)) diff --git a/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py b/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py deleted file mode 100644 index 28621c0ea1c5..000000000000 --- a/source/isaaclab_newton/test/controllers/test_newton_task_space_controllers.py +++ /dev/null @@ -1,239 +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 - -"""Operational-space control of a locally authored Newton chain through the articulation's dynamics accessors. - -The Isaac Lab :class:`~isaaclab.controllers.OperationalSpaceController` consumes -:attr:`~isaaclab.assets.BaseArticulationData.body_link_jacobian_w`, -:attr:`~isaaclab.assets.BaseArticulationData.mass_matrix`, and -:attr:`~isaaclab.assets.BaseArticulationData.gravity_compensation_forces` every step. These sentinels close the -loop on a six-DOF fixed chain whose links carry center-of-mass offsets, so a wrong Jacobian, mass matrix, gravity -force, or DoF ordering pushes the steady-state error well past the bounds. The chain's joints are unpowered, so the -controller's joint efforts are the only actuation. -""" - -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 pathlib import Path - -import pytest -import torch -from isaaclab_newton.assets import Articulation - -from isaaclab.actuators import ImplicitActuatorCfg -from isaaclab.assets import ArticulationCfg -from isaaclab.controllers import OperationalSpaceController, OperationalSpaceControllerCfg -from isaaclab.sim import SimulationContext, build_simulation_context -from isaaclab.test.utils import DeviceScope, test_devices -from isaaclab.utils.math import compute_pose_error, matrix_from_quat, quat_inv, subtract_frame_transforms - -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -from articulation_test_utils import NUM_ENVS, local_usd, newton_sim_cfg, spawn_assets, world_gravity # noqa: E402 - -pytestmark = [pytest.mark.integration, pytest.mark.kitless] - - -@dataclass -class _Chain: - """A passive six-DOF chain and the end-effector indices the controller reads.""" - - sim: SimulationContext - device: str - robot: Articulation - ee_frame_idx: int - ee_jacobi_idx: int - arm_joint_ids: list[int] - - def rest(self) -> None: - """Hold the default configuration at rest with no effort command.""" - default_joint_pos = self.robot.data.default_joint_pos.torch.clone() - self.robot.write_joint_state_to_sim_index( - position=default_joint_pos, velocity=torch.zeros_like(default_joint_pos) - ) - self.robot.actuators.target_command.set_effort_index(value=torch.zeros_like(default_joint_pos)) - self.robot.reset() - self.robot.write_data_to_sim() - self.sim.step() - self.robot.update(self.sim.cfg.dt) - - -@pytest.fixture(scope="module", params=test_devices(DeviceScope.CPU)) -def chain(request: pytest.FixtureRequest) -> Iterator[_Chain]: - """Build the unpowered chain once, at a configuration away from the joint limits.""" - robot_cfg = ArticulationCfg( - prim_path="/World/Env_[^/]*/Robot", - spawn=local_usd("fixed_spatial_chain.usda"), - init_state=ArticulationCfg.InitialStateCfg( - pos=(0.0, 0.0, 1.0), joint_pos={"Joint_3": 0.3, "Joint_4": -0.2, "Joint_5": 0.4} - ), - actuators={"arm": ImplicitActuatorCfg(joint_names_expr=["Joint_.*"], stiffness=0.0, damping=0.0)}, - ) - device = request.param - with build_simulation_context(sim_cfg=newton_sim_cfg(device)) as sim: - robot = spawn_assets({"robot": robot_cfg})["robot"] - sim.reset() - assert robot.is_initialized and robot.is_fixed_base - ee_frame_idx = robot.find_bodies("Link_5")[0][0] - yield _Chain( - sim=sim, - device=device, - robot=robot, - ee_frame_idx=ee_frame_idx, - # the fixed root has no Jacobian row - ee_jacobi_idx=ee_frame_idx - 1, - arm_joint_ids=robot.find_joints(["Joint_.*"])[0], - ) - - -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 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 = 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 - - -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", - 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=NUM_ENVS, - device=device, - ) - - -def _run_osc( - chain: _Chain, 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 = chain.robot - pos_history: list[float] = [] - rot_history: list[float] = [] - for _ in range(num_steps): - jacobian_b = _compute_jacobian_root_frame(robot, chain.ee_jacobi_idx, chain.arm_joint_ids) - mass_matrix = robot.data.mass_matrix.torch[:, chain.arm_joint_ids, :][:, :, chain.arm_joint_ids] - gravity_forces = robot.data.gravity_compensation_forces.torch[:, chain.arm_joint_ids] if gravity else None - ee_pos_b, ee_quat_b = _compute_ee_pose_root(robot, chain.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[:, chain.arm_joint_ids] - 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( - jacobian_b=jacobian_b, - current_ee_pose_b=ee_pose_b, - current_ee_vel_b=ee_vel_b, - mass_matrix=mass_matrix, - gravity=gravity_forces, - ) - robot.actuators.target_command.set_effort_index(value=joint_efforts, joint_ids=chain.arm_joint_ids) - robot.write_data_to_sim() - chain.sim.step() - robot.update(chain.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()) - return pos_history, rot_history - - -@pytest.mark.isaacsim_ci -def test_osc_tracking_accuracy(chain: _Chain) -> None: - """OSC pose tracking sentinel for the Jacobian and mass-matrix bridge. - - 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. - """ - chain.rest() - ee_pos_b, ee_quat_b = _compute_ee_pose_root(chain.robot, chain.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(chain, _make_osc(chain.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.isaacsim_ci -def test_osc_gravity_compensation_precision(chain: _Chain) -> None: - """Two-phase EE hold: gravity sag without compensation, tight hold with it. - - 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. - """ - chain.rest() - osc = _make_osc(chain.device) - # Hold the initial EE pose: phase-1 steady-state error is pure gravity sag. - ee_pos_b, ee_quat_b = _compute_ee_pose_root(chain.robot, chain.ee_frame_idx) - target_pose_b = torch.cat([ee_pos_b, ee_quat_b], dim=-1) - - 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[-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 - - with world_gravity((0.0, 0.0, -9.81)): - hist_off, _ = _run_osc(chain, osc, target_pose_b, 200, gravity=True) - osc.cfg.gravity_compensation = True - hist_on, _ = _run_osc(chain, 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") - - 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)" From 9800f71329045c55cabb7d36bc2060dd5e916f1e Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Fri, 2 Oct 2026 16:34:53 -0700 Subject: [PATCH 31/33] Consolidate PhysX asset tests and use local actuator fixtures --- .../test/assets/test_articulation.py | 79 ++- .../assets/test_newton_actuators_physx.py | 200 +++++--- .../test/assets/test_rigid_object.py | 462 ++++++++++++++--- .../assets/test_rigid_object_collection.py | 485 ------------------ .../test/assets/test_surface_gripper.py | 49 ++ .../unit/test_physx_actuator_control.py | 73 --- .../assets/unit/test_physx_articulation.py | 52 -- .../assets/unit/test_physx_surface_gripper.py | 53 -- 8 files changed, 642 insertions(+), 811 deletions(-) delete mode 100644 source/isaaclab_physx/test/assets/test_rigid_object_collection.py delete mode 100644 source/isaaclab_physx/test/assets/unit/test_physx_actuator_control.py delete mode 100644 source/isaaclab_physx/test/assets/unit/test_physx_articulation.py delete mode 100644 source/isaaclab_physx/test/assets/unit/test_physx_surface_gripper.py diff --git a/source/isaaclab_physx/test/assets/test_articulation.py b/source/isaaclab_physx/test/assets/test_articulation.py index 2e025bd791b2..5cb2a55b479f 100644 --- a/source/isaaclab_physx/test/assets/test_articulation.py +++ b/source/isaaclab_physx/test/assets/test_articulation.py @@ -33,10 +33,12 @@ 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.assets.articulation.kernels import write_joint_state_data, write_joint_state_data_kernel from pxr import Gf, PhysxSchema, UsdGeom, UsdPhysics @@ -233,13 +235,14 @@ def test_setting_invalid_articulation_root_prim_path(sim, device) -> None: sim.reset() -@pytest.mark.parametrize("device", test_devices()) +@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 and follow mass writes. + """Gravity compensation efforts hold the arms still; gravity compensation and the mass matrix follow mass writes. 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. + 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). """ articulation = Articulation( _branching_cfg( @@ -265,13 +268,17 @@ def test_gravity_compensation_holds_static_equilibrium(sim, device) -> None: articulation.update(sim.cfg.dt) torch.testing.assert_close(articulation.data.joint_pos.torch, joint_pos, atol=5e-3, rtol=0.0) - # A mass write refreshes efforts already read in this step. GPU PhysX applies mass writes on the next step - # (see ``test_inertial_writes_refresh_the_mass_matrix``). - if device == "cpu": - primed = articulation.data.gravity_compensation_forces.torch.clone() - tips = articulation.find_bodies(".*_tip")[0] - articulation.set_masses_index(masses=2.0 * articulation.data.body_mass.torch[:, tips], body_ids=tips) - assert not torch.allclose(articulation.data.gravity_compensation_forces.torch, primed, atol=1e-3) + # 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" ## @@ -331,7 +338,7 @@ def _spawn_island(name: str, y_offset: float, cfg: ArticulationCfg, **authoring) return articulation, torch.tensor(origins) -@pytest.fixture(scope="module", params=test_devices()) +@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 @@ -1179,12 +1186,9 @@ def test_fixed_tendon_position_target_writes_offset(articulation_scene: _Articul torch.testing.assert_close(wp.to_torch(articulation.root_view.get_fixed_tendon_offsets()).to(device), offset) -def test_inertial_writes_refresh_the_mass_matrix(articulation_scene: _ArticulationScene, request) -> None: +@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.""" - if articulation_scene.device.startswith("cuda"): - request.applymarker( - pytest.mark.xfail(strict=True, reason="GPU PhysX applies mass and inertia writes on the next step") - ) articulation = articulation_scene.ordered data = articulation.data writes = ( @@ -1223,3 +1227,46 @@ def test_joint_position_limit_clamping_respects_logging(articulation_scene: _Art finally: articulation.write_joint_position_limit_to_sim_index(limits=original_limits) articulation.data.default_joint_pos.torch.copy_(original_defaults) + + +## +# Kernels, without a simulation. +## + + +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", + ) + + 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_newton_actuators_physx.py b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py index 26203c59d9d2..2895693cab1d 100644 --- a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py +++ b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py @@ -7,10 +7,11 @@ 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 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, and network actuators. +: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 DeviceScope, launch_test_simulation, test_devices @@ -21,17 +22,21 @@ import os 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 +from isaaclab.actuators import IdealPDActuatorCfg, ImplicitActuatorCfg from isaaclab.actuators.actuator_net_cfg import ActuatorNetLSTMCfg, ActuatorNetMLPCfg -from isaaclab.actuators.newton import read_group_parameter +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 ( @@ -47,25 +52,14 @@ ) from isaaclab.test.utils.articulation_ordering import assert_articulation_ordering_trace_matches -from isaaclab_assets import ANYMAL_C_CFG, CARTPOLE_CFG - NUM_ENVS = 2 NUM_STEPS = 10 DT = 1.0 / 120.0 -_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), @@ -74,13 +68,9 @@ "applied_effort": (1e-3, 1e-3), } - -_ROWS = {"delayed": 0, "mixed": 1, "ideal": 2, "reversed": 2, "cartpole": 3, "implicit": 4, "mlp": 5, "lstm": 6} -"""Row of each robot along y. Robots compared across rollouts share a row: ground contacts amplify the floating-point -differences between placements at different world positions.""" _LAB_ROBOTS = ("delayed", "mixed", "ideal", "cartpole") """Robots of the Lab-actuator rollout, which the Newton-actuator rollout reproduces.""" -_NEWTON_ROBOTS = (*_LAB_ROBOTS, "implicit", "mlp", "lstm") +_NEWTON_ROBOTS = (*_LAB_ROBOTS, "permuted", "implicit", "mlp", "lstm") """Robots of the Newton-actuator rollout.""" @@ -88,14 +78,14 @@ def _robot_cfgs(names: tuple[str, ...], networks: dict[str, str]) -> dict[str, A """Return the configurations of the named robots. Args: - names: Robot names. ``reversed`` is the ``ideal`` robot with a reversed public joint order. + 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 ) - anymal_actuators = { + leg_actuators = { "delayed": DELAYED_PD_ACTUATORS, "mixed": MIXED_WITH_IMPLICIT_ACTUATORS, "ideal": IDEAL_PD_ACTUATORS, @@ -120,15 +110,27 @@ def _robot_cfgs(names: tuple[str, ...], networks: dict[str, str]) -> dict[str, A "pd_legs": pd_legs, }, } + leg_actuators["permuted"] = IDEAL_PD_ACTUATORS cfgs = {} - for name in names: + for row, name in enumerate(names): + init_state = ArticulationCfg.InitialStateCfg(pos=(0.0, 3.0 * row, 1.0)) if name == "cartpole": - cfgs[name] = replace(CARTPOLE_CFG, actuators=CARTPOLE_EXPLICIT_ACTUATORS) - elif name == "reversed": - joint_ordering = tuple(reversed(_ANYMAL_C_PHYSX_JOINT_NAMES)) - cfgs[name] = replace(ANYMAL_C_CFG, actuators=IDEAL_PD_ACTUATORS, joint_ordering=joint_ordering) + usd, actuators = "fixed_cartpole.usda", CARTPOLE_EXPLICIT_ACTUATORS else: - cfgs[name] = replace(ANYMAL_C_CFG, actuators=anymal_actuators[name]) + 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 @@ -144,16 +146,7 @@ def _rollout(sim: SimulationContext, cfgs: dict[str, ArticulationCfg]) -> tuple[ """ 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( - replace( - cfg, - prim_path=f"/World/Env_[^/]*/{name}", - init_state=replace(cfg.init_state, pos=(0.0, 3.0 * _ROWS[name], cfg.init_state.pos[2])), - ) - ) - for name, cfg in cfgs.items() - } + articulations = {name: Articulation(cfg) for name, cfg in cfgs.items()} sim.reset() traces = {} @@ -216,39 +209,35 @@ def _rollout(sim: SimulationContext, cfgs: dict[str, ArticulationCfg]) -> tuple[ @dataclass class _Rollouts: - """Recorded traces of the three rollouts and the live robots of the Newton-actuator rollout.""" + """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.""" - reversed: dict - """Trace of the ``reversed`` robot, recorded alone with Newton-native actuators.""" 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 and a ground plane on the PhysX backend.""" + """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, add_ground_plane=True, sim_cfg=sim_cfg) + 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 and reversed-ordering rollouts, then keep the Newton-actuator rollout alive.""" + """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: - traces = {} - for name, use_newton_actuators, names in (("lab", False, _LAB_ROBOTS), ("reversed", True, ("reversed",))): - with _simulation(device, use_newton_actuators) as sim: - sim._app_control_on_stop_handle = None - traces[name] = _rollout(sim, _robot_cfgs(names, networks))[1] + 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(traces["lab"], newton_traces, traces["reversed"]["reversed"], robots) + yield _Rollouts(lab_traces, newton_traces, robots) finally: for path in networks.values(): os.unlink(path) @@ -262,7 +251,11 @@ def test_newton_actuators_match_lab_actuators(rollouts: _Rollouts, robot: str) - zip(rollouts.lab[robot][field], rollouts.newton[robot][field], strict=True) ): torch.testing.assert_close( - newton, lab, atol=atol, rtol=rtol, msg=f"{robot} {field} diverged at step {step}" + newton, + lab, + atol=atol, + rtol=rtol, + msg=lambda detail: f"{robot} {field} diverged at step {step}: {detail}", ) @@ -271,16 +264,16 @@ def test_stateful_actuators_capture_ping_pong_graphs(rollouts: _Rollouts) -> Non assert len(rollouts.robots["delayed"]._actuator_control._native_actuator_graphs) == 2 -def test_reversed_joint_ordering_matches_identity_ordering(rollouts: _Rollouts) -> None: - """A reversed public joint order commands and observes the same physical joints as the backend order.""" +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.reversed, tuple(reversed(_ANYMAL_C_PHYSX_JOINT_NAMES)) + rollouts.newton["ideal"], rollouts.newton["permuted"], _PERMUTED_JOINT_NAMES ) -def test_reversed_joint_ordering_uses_current_joint_state(rollouts: _Rollouts) -> None: - """Regression: under a reversed joint order, each step's effort follows the IdealPD law on that step's state.""" - trace = rollouts.reversed +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) @@ -295,7 +288,7 @@ def test_network_actuators_drive_haa_joints(rollouts: _Rollouts, robot: str) -> """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) == 4 + 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}" @@ -352,11 +345,11 @@ def test_gain_randomization_reuses_payload_for_implicit_storage(rollouts: _Rollo def test_gain_randomization_reaches_only_the_selected_articulation(rollouts: _Rollouts) -> None: - """Each PhysX articulation owns its adapter, so randomizing the cartpole leaves the ANYmal controllers alone.""" - anymal, cartpole = rollouts.robots["ideal"], rollouts.robots["cartpole"] - assert anymal.newton_actuator_adapter is not None and cartpole.newton_actuator_adapter is not None - assert anymal.newton_actuator_adapter is not cartpole.newton_actuator_adapter - groups = {"anymal": (anymal, "legs"), "cartpole": (cartpole, "all_joints")} + """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 { @@ -366,12 +359,12 @@ def read_gains() -> dict[tuple[str, str], torch.Tensor]: } expected = read_gains() - env = MockEnv({"anymal": anymal, "cartpole": cartpole}, NUM_ENVS, anymal.device) + 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=anymal.device), + 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), @@ -383,3 +376,66 @@ def read_gains() -> dict[tuple[str, str], torch.Tensor]: 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) + ) + gain_writes = [] + articulation = SimpleNamespace( + _sim_cfg=SimpleNamespace(use_newton_actuators=True), + cfg=SimpleNamespace(prim_path="/World/Robot"), + joint_names=["joint"], + num_instances=1, + num_joints=1, + device="cpu", + _data=SimpleNamespace(joint_pos=joint_buffer, joint_vel=joint_buffer), + write_joint_stiffness_to_sim_index=lambda **_: gain_writes.append("stiffness"), + write_joint_damping_to_sim_index=lambda **_: gain_writes.append("damping"), + ) + wrapper = SimpleNamespace() + adapter = SimpleNamespace(joint_indices=wp.array([0], dtype=wp.int32), finalize=lambda _: None) + monkeypatch.setattr(actuator_control, "find_first_matching_prim", lambda _: None) + monkeypatch.setattr(PhysxActuatorWrapper, "create", lambda **_: wrapper) + monkeypatch.setattr(NewtonActuatorAdapter, "from_usd", lambda **_: adapter) + + native_groups = PhysxActuatorControl(articulation).prepare_native_actuators( + collection, + {"explicit": IdealPDActuatorCfg(joint_names_expr=["joint"], stiffness=None, damping=None)}, + ) + + assert native_groups == {"explicit"} + assert gain_writes == [] + + +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 = PhysxActuatorControl(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 == [] diff --git a/source/isaaclab_physx/test/assets/test_rigid_object.py b/source/isaaclab_physx/test/assets/test_rigid_object.py index 75887a65edfb..8a0d87c96c6e 100644 --- a/source/isaaclab_physx/test/assets/test_rigid_object.py +++ b/source/isaaclab_physx/test/assets/test_rigid_object.py @@ -3,12 +3,14 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Real PhysX rigid-object coverage. - -Most checks run against one module-scoped scene of locally authored cubes. Each rigid object 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 are defined first: pytest runs them before the -composite scene is created, and a new simulation context would replace the composite stage. +"""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 @@ -28,7 +30,7 @@ import pytest import torch import warp as wp -from isaaclab_physx.assets import RigidObject +from isaaclab_physx.assets import RigidObject, RigidObjectCollection from isaaclab_physx.sim.schemas import PhysxArticulationCfg, PhysxRigidBodyCfg import omni.kit.app @@ -36,27 +38,38 @@ from pxr import UsdPhysics import isaaclab.sim as sim_utils -from isaaclab.assets import RigidObjectCfg +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(), + ) + +_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.""" -def _cube_cfg(prim_path: str, *, kinematic: bool = False, disable_gravity: bool = False) -> RigidObjectCfg: - """Create a local 1 kg collision cube one meter above its environment origin.""" + +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=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(), - ), + spawn=_cube_spawn(**spawn_kwargs), init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 1.0)), ) @@ -69,6 +82,28 @@ def _spawn_envs(root: str, y_offset: float = 0.0) -> torch.Tensor: return origins +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. + + 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 + + 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)) @@ -89,6 +124,29 @@ def sim(device: str) -> Iterator[SimulationContext]: ## +@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. + + 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 + invalid_collection, _ = _collection("/World", 1, 2, _STATIC_COLLIDER) + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(invalid_collection) < 10 + with pytest.raises(RuntimeError): + sim.reset() + with build_simulation_context(device="cpu") as sim: + sim._app_control_on_stop_handle = None + kinematic, origins = _collection("/World", 1, 1, _cube_spawn(kinematic=True)) + sim.reset() + 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", [ @@ -110,11 +168,7 @@ def test_initialization_rejects_invalid_rigid_body(sim, device, api: Literal["no _spawn_envs("/World") if api == "none": # Without rigid body properties the cube is a static collider. - cfg = RigidObjectCfg( - prim_path="/World/Env_[^/]*/Object", - spawn=sim_utils.CuboidCfg(size=(0.1, 0.1, 0.1), collision_props=sim_utils.UsdPhysicsCollisionCfg()), - ) - cube_object = RigidObject(cfg=cfg) + 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. @@ -172,8 +226,8 @@ def test_warmup_loads_physics_once(sim, device) -> None: @dataclass -class _RigidObjectScene: - """Rigid objects that share one real PhysX lifecycle under gravity.""" +class _Scene: + """Rigid objects and collections that share one real PhysX lifecycle under gravity.""" sim: SimulationContext device: str @@ -183,19 +237,36 @@ class _RigidObjectScene: """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 object name.""" + """Environment origins keyed by asset name.""" refcounts: dict[str, int] - """Reference count of each object right after construction.""" + """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, + ) def step(self, num_steps: int = 1) -> None: - """Write, step, and update every object.""" + """Write, step, and update every asset.""" for _ in range(num_steps): - for rigid_object in (self.cubes, self.kinematic, self.falling): - rigid_object.write_data_to_sim() + for asset in self.assets: + asset.write_data_to_sim() self.sim.step() - for rigid_object in (self.cubes, self.kinematic, self.falling): - rigid_object.update(self.sim.cfg.dt) + 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.""" @@ -203,10 +274,27 @@ def place_cubes_at_rest(self, root_pose: torch.Tensor) -> None: 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.fixture(scope="module", params=test_devices()) -def rigid_object_scene(request) -> Iterator[_RigidObjectScene]: - """Initialize the composite rigid-object scene once per device.""" +@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 @@ -215,31 +303,39 @@ def rigid_object_scene(request) -> Iterator[_RigidObjectScene]: "kinematic": _spawn_envs("/World/Kinematic", 3.0), "falling": _spawn_envs("/World/Falling", 6.0), } - 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")) - refcounts = {"cubes": sys.getrefcount(cubes), "kinematic": sys.getrefcount(kinematic)} + 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() - yield _RigidObjectScene( + yield _Scene( sim=sim, device=device, - cubes=cubes, - kinematic=kinematic, - falling=falling, origins={name: value.to(device) for name, value in origins.items()}, refcounts=refcounts, + **assets, ) @pytest.mark.isaacsim_ci -def test_rigid_object_initialization(rigid_object_scene: _RigidObjectScene) -> None: - """Initialize local cubes with the expected buffers; under gravity, kinematic cubes hold their default pose and - dynamic cubes accelerate downward on every step.""" - scene = rigid_object_scene - for name, rigid_object in (("cubes", scene.cubes), ("kinematic", scene.kinematic)): +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 - assert rigid_object.is_initialized + 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) @@ -249,22 +345,43 @@ def test_rigid_object_initialization(rigid_object_scene: _RigidObjectScene) -> N 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) + ) - kinematic = scene.kinematic - gravity_acceleration = torch.tensor([[[0.0, 0.0, -9.81, 0.0, 0.0, 0.0]]] * _NUM_ENVS, device=scene.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) + 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 + ) @pytest.mark.isaacsim_ci -def test_rigid_object_inertial_properties(rigid_object_scene: _RigidObjectScene) -> None: +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.""" - scene = rigid_object_scene device = scene.device cube_object = scene.cubes @@ -303,9 +420,8 @@ def test_rigid_object_inertial_properties(rigid_object_scene: _RigidObjectScene) @pytest.mark.isaacsim_ci -def test_rigid_object_root_state_writes(rigid_object_scene: _RigidObjectScene) -> None: +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.""" - scene = rigid_object_scene device = scene.device cube_object = scene.cubes rest_pose = torch.cat( @@ -419,9 +535,8 @@ def test_rigid_object_root_state_writes(rigid_object_scene: _RigidObjectScene) - @pytest.mark.isaacsim_ci -def test_rigid_object_wrench_delivery_and_reset(rigid_object_scene: _RigidObjectScene) -> None: +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.""" - scene = rigid_object_scene device = scene.device cube_object = scene.cubes rest_pose = torch.cat( @@ -492,3 +607,230 @@ def test_rigid_object_wrench_delivery_and_reset(rigid_object_scene: _RigidObject 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) + + # 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] + ) + 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 2925f229df63..000000000000 --- a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py +++ /dev/null @@ -1,485 +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 - -"""Real PhysX rigid-object-collection coverage. - -Most checks run against one module-scoped scene of locally authored cubes. PhysX stores collection bodies in -body-major view order, so 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 DeviceScope, launch_test_simulation, test_devices - -launch_test_simulation() - -import math -import sys -from collections.abc import Iterator -from dataclasses import dataclass - -import pytest -import torch -import warp as wp -from isaaclab_physx.assets import RigidObjectCollection -from isaaclab_physx.sim.schemas import PhysxRigidBodyCfg - -import isaaclab.sim as sim_utils -from isaaclab.assets import RigidObjectCfg, RigidObjectCollectionCfg -from isaaclab.sim import SimulationContext, build_simulation_context -from isaaclab.utils.math import combine_frame_transforms, quat_apply, quat_apply_inverse, quat_mul - - -def generate_cubes_scene( - num_envs: int = 1, - num_cubes: int = 1, - height=1.0, - has_api: bool = True, - kinematic_enabled: bool = False, - disable_gravity: bool = False, - root: str = "/World", - y_offset: float = 0.0, -) -> tuple[RigidObjectCollection, torch.Tensor]: - """Generate a collection of local 1 kg cubes, 3 m apart along y within each environment. - - 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. - disable_gravity: Whether the cubes ignore gravity. - root: Prim path under which the environments are created. - y_offset: Offset of the environments along y [m]. - - Returns: - A tuple containing the rigid object collection and the origins of the environments. - """ - origins = torch.tensor([(i * 3.0, y_offset, height) for i in range(num_envs)]) - # Create Top-level Xforms, one for each cube - for i, origin in enumerate(origins.tolist()): - sim_utils.create_prim(f"{root}/Table_{i}", "Xform", translation=origin) - - # Resolve spawn configuration - if has_api: - spawn_cfg = sim_utils.CuboidCfg( - size=(0.2, 0.2, 0.2), - rigid_props=[ - sim_utils.UsdPhysicsRigidBodyCfg(kinematic_enabled=kinematic_enabled), - PhysxRigidBodyCfg(disable_gravity=disable_gravity), - ], - mass_props=sim_utils.MassCfg(mass=1.0), - collision_props=sim_utils.UsdPhysicsCollisionCfg(), - ) - 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_config_dict[f"cube_{i}"] = RigidObjectCfg( - prim_path=f"{root}/Table_[^/]*/Object_{i}", - spawn=spawn_cfg, - init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 3 * i, height)), - ) - # create the rigid object collection - return RigidObjectCollection(cfg=RigidObjectCollectionCfg(rigid_objects=cube_config_dict)), origins - - -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.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 - - -## -# Tests that own their simulation context. Keep them above the composite scene. -## - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) -def test_initialization_with_no_rigid_body(sim, device) -> None: - """Test that initialization fails when no rigid body is found at the provided prim path.""" - num_cubes = 2 - object_collection, _ = generate_cubes_scene(num_cubes=num_cubes, has_api=False) - - # 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.xfail( - strict=True, raises=AssertionError, reason="The first CPU simulation after a failed initialization reads zero poses" -) -def test_failed_initialization_does_not_leak_into_the_next_cpu_simulation() -> None: - """A kinematic cube reports its spawn pose in the CPU simulation that follows a failed initialization.""" - with build_simulation_context(device="cpu") as sim: - sim._app_control_on_stop_handle = None - invalid_collection, _ = generate_cubes_scene(num_cubes=2, has_api=False) # noqa: F841 # keep it alive - with pytest.raises(RuntimeError): - sim.reset() - with build_simulation_context(device="cpu") as sim: - sim._app_control_on_stop_handle = None - kinematic, origins = generate_cubes_scene(kinematic_enabled=True) - sim.reset() - torch.testing.assert_close(kinematic.data.body_link_pos_w.torch[:, 0], origins + torch.tensor([0.0, 0.0, 1.0])) - - -## -# Composite scene shared by the remaining tests. -## - -_NUM_ENVS = 2 -_NUM_CUBES = 3 - - -@dataclass -class _CollectionScene: - """Rigid-object collections that share one real PhysX lifecycle under gravity.""" - - sim: SimulationContext - device: str - cubes: RigidObjectCollection - """Two environments of three dynamic cubes that ignore gravity.""" - kinematic: RigidObjectCollection - """One environment of one kinematic cube.""" - falling: RigidObjectCollection - """Two environments of two dynamic cubes under gravity.""" - origins: dict[str, torch.Tensor] - """Environment origins keyed by collection name.""" - refcounts: dict[str, int] - """Reference count of each collection right after construction.""" - - def step(self, num_steps: int = 1) -> None: - """Write, step, and update every collection.""" - for _ in range(num_steps): - for collection in (self.cubes, self.kinematic, self.falling): - collection.write_data_to_sim() - self.sim.step() - for collection in (self.cubes, self.kinematic, self.falling): - collection.update(self.sim.cfg.dt) - - def place_cubes_at_rest(self, body_poses: torch.Tensor) -> None: - """Teleport the dynamic cubes to ``body_poses`` at rest and clear their external wrenches.""" - self.cubes.write_body_link_pose_to_sim_index(body_poses=body_poses) - self.cubes.write_body_com_velocity_to_sim_index( - body_velocities=torch.zeros((_NUM_ENVS, _NUM_CUBES, 6), device=self.device) - ) - self.cubes.permanent_wrench_composer.reset() - self.cubes.instantaneous_wrench_composer.reset() - - -@pytest.fixture(scope="module", params=test_devices()) -def collection_scene(request) -> Iterator[_CollectionScene]: - """Initialize the composite collection scene once per device.""" - device = request.param - with build_simulation_context(device=device, gravity_enabled=True) as sim: - sim._app_control_on_stop_handle = None - cubes, cube_origins = generate_cubes_scene( - num_envs=_NUM_ENVS, num_cubes=_NUM_CUBES, disable_gravity=True, root="/World/Cubes" - ) - kinematic, kinematic_origins = generate_cubes_scene( - num_envs=1, num_cubes=1, kinematic_enabled=True, root="/World/Kinematic", y_offset=12.0 - ) - falling, _ = generate_cubes_scene(num_envs=_NUM_ENVS, num_cubes=2, root="/World/Falling", y_offset=18.0) - refcounts = {"cubes": sys.getrefcount(cubes), "kinematic": sys.getrefcount(kinematic)} - sim.reset() - yield _CollectionScene( - sim=sim, - device=device, - cubes=cubes, - kinematic=kinematic, - falling=falling, - origins={"cubes": cube_origins.to(device), "kinematic": kinematic_origins.to(device)}, - refcounts=refcounts, - ) - - -def _rest_poses(scene: _CollectionScene, yaw: float = 0.0) -> torch.Tensor: - """Return the default body poses of the dynamic cubes at their environment origins with a given yaw.""" - poses = scene.cubes.data.default_body_pose.torch.clone() - poses[..., :2] += scene.origins["cubes"].unsqueeze(1)[..., :2] - poses[..., 3:] = torch.tensor(_yaw_quat(yaw), device=scene.device) - return poses - - -def test_collection_initialization(collection_scene: _CollectionScene) -> None: - """Initialize local collections, including a single-cube one; under gravity, kinematic cubes hold their pose and - dynamic cubes accelerate downward on every step.""" - scene = collection_scene - for name, collection, num_envs, num_cubes in ( - ("cubes", scene.cubes, _NUM_ENVS, _NUM_CUBES), - ("kinematic", scene.kinematic, 1, 1), - ): - # Check that the framework doesn't hold excessive strong references. - assert scene.refcounts[name] < 10 - assert collection.is_initialized - 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) - ) - - kinematic = scene.kinematic - gravity_acceleration = torch.tensor([0.0, 0.0, -9.81, 0.0, 0.0, 0.0], device=scene.device).repeat(_NUM_ENVS, 2, 1) - for _ in range(2): - scene.step() - torch.testing.assert_close(scene.falling.data.body_com_acc_w.torch, gravity_acceleration) - default_body_pose = kinematic.data.default_body_pose.torch.clone() - default_body_vel = kinematic.data.default_body_vel.torch.clone() - default_body_pose[..., :3] += scene.origins["kinematic"].unsqueeze(1) - torch.testing.assert_close(kinematic.data.body_link_pose_w.torch, default_body_pose) - torch.testing.assert_close(kinematic.data.body_link_vel_w.torch, default_body_vel) - - -def test_collection_state_writes(collection_scene: _CollectionScene) -> None: - """Body state writes reach the selected body-major view entries in the frame they are given in.""" - scene = collection_scene - device = scene.device - collection = scene.cubes - scene.place_cubes_at_rest(_rest_poses(scene)) - - # 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_cubes_at_rest(_rest_poses(scene)) - 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) - - # 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_cubes_at_rest(_rest_poses(scene)) - 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] - ) - 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(collection_scene: _CollectionScene) -> None: - """Mass, center-of-mass, and inertia writes to a non-sorted subset reach only the selected body-major entries.""" - scene = collection_scene - device = scene.device - collection = scene.cubes - 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(collection_scene: _CollectionScene) -> None: - """External wrenches act on the selected bodies in the frame they are given in; reset clears them.""" - scene = collection_scene - device = scene.device - collection = scene.cubes - rest_poses = _rest_poses(scene, yaw=0.5 * math.pi) - # 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_cubes_at_rest(rest_poses) - 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_surface_gripper.py b/source/isaaclab_physx/test/assets/test_surface_gripper.py index b135510573b7..729282e55e5a 100644 --- a/source/isaaclab_physx/test/assets/test_surface_gripper.py +++ b/source/isaaclab_physx/test/assets/test_surface_gripper.py @@ -11,6 +11,10 @@ launch_test_simulation() +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np import pytest import torch import warp as wp @@ -145,3 +149,48 @@ def test_raise_error_if_not_cpu(device) -> None: 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(), + ) + + 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, + ) + + 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/source/isaaclab_physx/test/assets/unit/test_physx_actuator_control.py b/source/isaaclab_physx/test/assets/unit/test_physx_actuator_control.py deleted file mode 100644 index e83e7577af38..000000000000 --- a/source/isaaclab_physx/test/assets/unit/test_physx_actuator_control.py +++ /dev/null @@ -1,73 +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 PhysX Newton-actuator path preparation.""" - -from types import SimpleNamespace - -import warp as wp -from isaaclab_physx.assets.articulation import actuator_control -from isaaclab_physx.assets.articulation.actuator_control import PhysxActuatorControl - -from isaaclab.actuators import IdealPDActuatorCfg, ImplicitActuatorCfg -from isaaclab.actuators.newton import NewtonActuatorAdapter, PhysxActuatorWrapper - - -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) - ) - gain_writes = [] - articulation = SimpleNamespace( - _sim_cfg=SimpleNamespace(use_newton_actuators=True), - cfg=SimpleNamespace(prim_path="/World/Robot"), - joint_names=["joint"], - num_instances=1, - num_joints=1, - device="cpu", - _data=SimpleNamespace(joint_pos=joint_buffer, joint_vel=joint_buffer), - write_joint_stiffness_to_sim_index=lambda **_: gain_writes.append("stiffness"), - write_joint_damping_to_sim_index=lambda **_: gain_writes.append("damping"), - ) - wrapper = SimpleNamespace() - adapter = SimpleNamespace(joint_indices=wp.array([0], dtype=wp.int32), finalize=lambda _: None) - monkeypatch.setattr(actuator_control, "find_first_matching_prim", lambda _: None) - monkeypatch.setattr(PhysxActuatorWrapper, "create", lambda **_: wrapper) - monkeypatch.setattr(NewtonActuatorAdapter, "from_usd", lambda **_: adapter) - - native_groups = PhysxActuatorControl(articulation).prepare_native_actuators( - collection, - {"explicit": IdealPDActuatorCfg(joint_names_expr=["joint"], stiffness=None, damping=None)}, - ) - - assert native_groups == {"explicit"} - assert gain_writes == [] - - -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 = PhysxActuatorControl(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 == [] diff --git a/source/isaaclab_physx/test/assets/unit/test_physx_articulation.py b/source/isaaclab_physx/test/assets/unit/test_physx_articulation.py deleted file mode 100644 index 023f038f49e4..000000000000 --- a/source/isaaclab_physx/test/assets/unit/test_physx_articulation.py +++ /dev/null @@ -1,52 +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 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") - - -@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/unit/test_physx_surface_gripper.py b/source/isaaclab_physx/test/assets/unit/test_physx_surface_gripper.py deleted file mode 100644 index 4a20299a672d..000000000000 --- a/source/isaaclab_physx/test/assets/unit/test_physx_surface_gripper.py +++ /dev/null @@ -1,53 +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 PhysX surface-gripper command and property payloads.""" - -from types import SimpleNamespace -from unittest.mock import Mock - -import numpy as np -import warp as wp -from isaaclab_physx.assets.surface_gripper.surface_gripper import SurfaceGripper - - -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(), - ) - - 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, - ) - - 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], - } From bd3f55501c61981fae9fc923c4cfa23775f49401 Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Fri, 2 Oct 2026 16:34:54 -0700 Subject: [PATCH 32/33] Consolidate OVPhysX asset tests and share scenes across devices --- .../test/assets/test_articulation.py | 354 ++++++++++++++- .../test/assets/test_deformable_object.py | 417 +++++++++++++++++- .../test/assets/test_rigid_object.py | 320 +++++++++++--- .../assets/test_rigid_object_collection.py | 259 ----------- .../isaaclab_ov/test/assets/unit/__init__.py | 4 - .../assets/unit/test_ovphysx_articulation.py | 329 -------------- .../unit/test_ovphysx_deformable_object.py | 383 ---------------- 7 files changed, 998 insertions(+), 1068 deletions(-) delete mode 100644 source/isaaclab_ov/test/assets/test_rigid_object_collection.py delete mode 100644 source/isaaclab_ov/test/assets/unit/__init__.py delete mode 100644 source/isaaclab_ov/test/assets/unit/test_ovphysx_articulation.py delete mode 100644 source/isaaclab_ov/test/assets/unit/test_ovphysx_deformable_object.py diff --git a/source/isaaclab_ov/test/assets/test_articulation.py b/source/isaaclab_ov/test/assets/test_articulation.py index 5e10eab1927e..aa1afb4b8737 100644 --- a/source/isaaclab_ov/test/assets/test_articulation.py +++ b/source/isaaclab_ov/test/assets/test_articulation.py @@ -17,6 +17,7 @@ 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. @@ -24,6 +25,7 @@ from __future__ import annotations +import importlib import logging import math from collections.abc import Iterator @@ -32,6 +34,7 @@ from types import SimpleNamespace from unittest.mock import Mock +import numpy as np import pytest import torch import warp as wp @@ -41,8 +44,12 @@ 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 @@ -61,6 +68,7 @@ BRANCHING_PHYSX_JOINT_NAMES, ) from isaaclab.utils import replace # noqa: E402 +from isaaclab.utils.warp.launch_cache import _WarpLaunchCache # noqa: E402 from isaaclab_assets import FRANKA_PANDA_CFG # noqa: E402 @@ -70,6 +78,12 @@ _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)) @@ -357,9 +371,315 @@ def _spawn_own_articulation(**cfg) -> Articulation: ) -@pytest.mark.parametrize("device", test_devices()) +## +# 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(): + """Ordering installation and simulation invalidation should discard recorded reads.""" + + class MinimalData(ArticulationData): + def __dir__(self): + return [] + + class Buffer: + timestamp = 1.0 + + data = MinimalData.__new__(MinimalData) + read_launch_cache = Mock() + data._read_launch_cache = read_launch_cache + data._configure_ordering_buffers = lambda: None + data._make_jacobian_body_user_to_backend = lambda: object() + data.joint_ordering = None + data._body_com_jacobian_w = Buffer() + data._mass_matrix = Buffer() + data._gravity_compensation_forces = Buffer() + + data._apply_ordering_maps_after_resolve() + + read_launch_cache.clear.assert_called_once_with() + assert data._body_com_jacobian_w.timestamp == -1.0 + assert data._mass_matrix.timestamp == -1.0 + assert data._gravity_compensation_forces.timestamp == -1.0 + + data._is_primed = True + data._sim_timestamp = 1.0 + data._invalidate_initialize_callback(None) + + assert read_launch_cache.clear.call_count == 2 + assert data._is_primed is False + assert data._sim_timestamp == 0.0 + + +def test_static_property_reads_are_not_invalidated_by_simulation_steps(): + """Joint properties and body mass/inertia should be read once per invalidation, not once per step. + + On OVPhysX these are blocking CPU-only binding reads whose cost scales with the number of + environments, so re-reading them every physics step made per-step consumers (such as the + native actuator telemetry sync) host-bound. State buffers must still refresh every step. + """ + + class Buffer: + def __init__(self, shape): + self.data = wp.zeros(shape, dtype=wp.float32, device="cpu") + self.timestamp = -1.0 + + data = ArticulationData.__new__(ArticulationData) + data.device = "cpu" + data.num_instances = 1 + data.num_joints = 2 + data._sim_timestamp = 1.0 + data.body_ordering = None + data._get_binding = lambda tensor_type: object() + reads: list[int] = [] + data._binding_read = lambda tensor_type, dst: reads.append(tensor_type) + + # Joint properties: one read across several steps, one more after explicit invalidation + # (simulation reinitialization). + data.joint_ordering = None + stiffness = Buffer((1, 2)) + for _ in range(3): + data._sim_timestamp += 1.0 + data._read_joint_property_binding(TT.DOF_STIFFNESS, stiffness, None) + assert reads.count(TT.DOF_STIFFNESS) == 1 + stiffness.timestamp = -1.0 + data._read_joint_property_binding(TT.DOF_STIFFNESS, stiffness, None) + assert reads.count(TT.DOF_STIFFNESS) == 2 + + # Body properties behave the same; body state buffers still refresh every step. + mass, link_pose = Buffer((1, 2)), Buffer((1, 2)) + for _ in range(3): + data._sim_timestamp += 1.0 + data._refresh_reordered_body_buffer(mass, None, TT.BODY_MASS, static=True) + data._refresh_reordered_body_buffer(link_pose, None, TT.LINK_POSE) + assert reads.count(TT.BODY_MASS) == 1 + assert reads.count(TT.LINK_POSE) == 3 + mass.timestamp = -1.0 + data._refresh_reordered_body_buffer(mass, None, TT.BODY_MASS, static=True) + assert reads.count(TT.BODY_MASS) == 2 + + # Under a non-identity joint ordering the property is gathered once, then served from cache. + data.joint_ordering = SimpleNamespace(user_to_backend=wp.array([1, 0], dtype=wp.int32, device="cpu")) + data._read_launch_cache = _WarpLaunchCache("cpu") + user_buffer, backend_buffer = Buffer((1, 2)), Buffer((1, 2)) + data._binding_read = lambda tensor_type, dst: (reads.append(tensor_type), dst.assign([[1.0, 2.0]])) + for _ in range(3): + data._sim_timestamp += 1.0 + data._read_joint_property_binding(TT.DOF_DAMPING, user_buffer, backend_buffer) + assert reads.count(TT.DOF_DAMPING) == 1 + torch.testing.assert_close(wp.to_torch(user_buffer.data), torch.tensor([[2.0, 1.0]])) + + +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,) + + +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"] + + +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]]) + + +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.""" + """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")) @@ -370,7 +690,7 @@ def test_mimic_finger_follows_commanded_finger(device: str) -> None: 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(100): + for _ in range(50): articulation.write_data_to_sim() sim.step() articulation.update(sim.cfg.dt) @@ -382,11 +702,13 @@ def test_mimic_finger_follows_commanded_finger(device: str) -> None: @pytest.mark.parametrize( ("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") - for device in _ALL_DEVICES ] + + [(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")], ) @@ -841,7 +1163,7 @@ def test_com_orientation_write_invalidates_static_inertia_cache_with_body_orderi ) -@pytest.mark.parametrize("scene", _CUDA_DEVICES, indirect=True) +@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"] @@ -855,7 +1177,7 @@ def test_loading_gains_from_usd(scene: _ArticulationScene) -> None: torch.testing.assert_close(_read_binding_to_torch(articulation, TT.DOF_DAMPING, scene.device), expected_damping) -@pytest.mark.parametrize("scene", _CUDA_DEVICES, indirect=True) +@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"] @@ -879,7 +1201,7 @@ def test_explicit_articulation_root_on_a_floating_base_body(scene: _Articulation assert torch.isfinite(articulation.data.root_link_pose_w.torch).all() -@pytest.mark.parametrize("scene", _CUDA_DEVICES, indirect=True) +@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"] @@ -940,7 +1262,7 @@ def test_implicit_drive_targets_submit_feedforward_effort_and_track(scene: _Arti torch.testing.assert_close(articulation.data.joint_pos.torch, expected_position, atol=0.1, rtol=0.0) -@pytest.mark.parametrize("scene", _CUDA_DEVICES, indirect=True) +@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. @@ -1069,7 +1391,7 @@ def assert_dynamics_match(jacobian: bool = True, gravity: bool = True) -> None: assert_dynamics_match(jacobian=False, gravity=False) -@pytest.mark.parametrize("scene", _CUDA_DEVICES, indirect=True) +@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.""" @@ -1109,7 +1431,7 @@ def test_reversed_joint_dynamics_use_public_joint_basis(scene: _ArticulationScen ) -@pytest.mark.parametrize("scene", _CUDA_DEVICES, indirect=True) +@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"] @@ -1217,7 +1539,7 @@ def test_floating_articulation_root_state_dynamics_wrench_and_reset(scene: _Arti articulation.reset() -@pytest.mark.parametrize("scene", _CUDA_DEVICES, indirect=True) +@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. @@ -1275,7 +1597,7 @@ def test_tendon_properties_and_position_targets(scene: _ArticulationScene) -> No torch.testing.assert_close(articulation.data.fixed_tendon_offset.torch, expected) -@pytest.mark.parametrize("scene", _CUDA_DEVICES, indirect=True) +@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"] @@ -1313,7 +1635,7 @@ def test_native_actuator_submits_real_effort(scene: _ArticulationScene) -> None: ) -@pytest.mark.parametrize("scene", _CUDA_DEVICES, indirect=True) +@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. @@ -1364,7 +1686,7 @@ def __getitem__(self, name): torch.testing.assert_close(damping, torch.tensor([[3.0], [1.0]], device=device).repeat(1, num_joints)) -@pytest.mark.parametrize("scene", _CUDA_DEVICES, indirect=True) +@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. @@ -1403,7 +1725,7 @@ def _occupy_stream_kernel(iterations: int, out: wp.array(dtype=wp.float32)): out[0] = total -@pytest.mark.parametrize("scene", _CUDA_DEVICES, indirect=True) +@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. diff --git a/source/isaaclab_ov/test/assets/test_deformable_object.py b/source/isaaclab_ov/test/assets/test_deformable_object.py index 5b207f4fe459..3e46afa80da8 100644 --- a/source/isaaclab_ov/test/assets/test_deformable_object.py +++ b/source/isaaclab_ov/test/assets/test_deformable_object.py @@ -6,12 +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 re import sys +from types import SimpleNamespace import ovphysx.types # noqa: F401 import pytest @@ -19,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 @@ -127,25 +133,394 @@ def _assert_rest_positions_match_authored( ) +## +# 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() @@ -154,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, @@ -195,12 +567,13 @@ def test_initialization_on_device_cpu(): @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) @@ -215,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) @@ -233,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() @@ -245,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_rigid_object.py b/source/isaaclab_ov/test/assets/test_rigid_object.py index 3cd7835d2f67..f0d2deb35bbf 100644 --- a/source/isaaclab_ov/test/assets/test_rigid_object.py +++ b/source/isaaclab_ov/test/assets/test_rigid_object.py @@ -6,11 +6,12 @@ # ignore private usage of variables warning # pyright: reportPrivateUsage=none -"""Real OVPhysX rigid-object coverage on one module-scoped scene per device. +"""Real OVPhysX rigid-object and rigid-object-collection coverage on one module-scoped scene per device. -Each device builds one scene of locally authored cuboid pairs, resets it once, and keeps it alive for -every test that uses it. Each pair belongs to one test, so those tests do not depend on each other's -order. +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. 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 @@ -31,12 +32,12 @@ 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, SimulationContext, build_simulation_context # noqa: E402 from isaaclab.test.utils import DeviceScope, test_devices # noqa: E402 @@ -46,6 +47,7 @@ pytestmark = pytest.mark.integration _NUM_CUBES = 2 +_NUM_ENVS, _NUM_BODIES = 2, 3 def _spawn_cubes(name: str, y_offset: float, **rigid_props) -> RigidObject: @@ -79,6 +81,31 @@ def _spawn_static_colliders() -> RigidObject: ) +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 class HeterogeneousRigidSceneCfg(InteractiveSceneCfg): """Two object variants dropping onto an independently cloned support.""" @@ -99,17 +126,31 @@ 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()) +_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( - ("support_cloning", "filter_collisions"), - [("native", False), ("usd_and_native", True), ("usd", False), ("usd_nested", True)], + ("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(device, support_cloning, filter_collisions): - """Sources and clones from different variants contact their own support. +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. """ @@ -130,64 +171,53 @@ def test_heterogeneous_clone_contacts(device, support_cloning, filter_collisions 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 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. - 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 build_simulation_context( - device=device, sim_cfg=SimulationCfg(physics=OvPhysxCfg(), device=device, gravity=(0.0, 0.0, 0.0)) - ) 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]) + + +@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() + for _ in range(60): + 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(DeviceScope.CPU)) @@ -219,27 +249,51 @@ def test_initialization_with_articulation_root(device: str) -> None: sim.reset() +@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()), + ) + with pytest.raises(RuntimeError, match="Expected 1 prims"): + sim.reset() + assert not object_collection.is_initialized + + @dataclass class _RigidScene: - """Cuboid pairs that share one real OVPhysX lifecycle.""" + """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 for one device once for this module.""" + """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=sim, device=device, dynamic=dynamic, kinematic=kinematic, spinning=spinning) + yield _RigidScene(sim, device, dynamic, kinematic, spinning, fused, stepped, wrenched) @pytest.mark.parametrize("scene", test_devices(), indirect=True) @@ -321,6 +375,12 @@ def test_initialization_with_kinematic_enabled(scene: _RigidScene) -> None: 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 @@ -336,7 +396,7 @@ def test_initialization_with_kinematic_enabled(scene: _RigidScene) -> None: torch.testing.assert_close(cube_object.data.root_com_vel_w.torch, default_root_vel) -@pytest.mark.parametrize("scene", test_devices(DeviceScope.CUDA), indirect=True) +@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 @@ -403,3 +463,153 @@ def test_body_root_state_properties(scene: _RigidScene) -> None: # 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) + + +@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. + + 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 fc0d3ae0671a..000000000000 --- a/source/isaaclab_ov/test/assets/test_rigid_object_collection.py +++ /dev/null @@ -1,259 +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 OVPhysX rigid-object-collection coverage on one module-scoped scene per device. - -Each device builds one scene of locally authored ``N=2, B=3`` collections, resets it once, and keeps it -alive for every test that uses it. Each collection belongs to one test, so those tests do not depend on each -other's order. - -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 - -from collections.abc import Iterator -from dataclasses import dataclass - -import pytest -import torch -import warp as wp - -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, SimulationContext, build_simulation_context # noqa: E402 -from isaaclab.test.utils import DeviceScope, test_devices # noqa: E402 - -pytestmark = pytest.mark.integration - -_NUM_ENVS, _NUM_BODIES = 2, 3 - - -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")) - } - ) - ) - - -@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, 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()), - ) - with pytest.raises(RuntimeError, match="Expected 1 prims"): - sim.reset() - assert not object_collection.is_initialized - - -@dataclass -class _CollectionScene: - """Collections that share one real OVPhysX lifecycle.""" - - sim: SimulationContext - device: str - fused: RigidObjectCollection - stepped: RigidObjectCollection - wrenched: RigidObjectCollection - - -@pytest.fixture(scope="module") -def scene(request: pytest.FixtureRequest) -> Iterator[_CollectionScene]: - """Initialize every collection for one device once for this module.""" - device = request.param - with build_simulation_context( - device=device, sim_cfg=SimulationCfg(physics=OvPhysxCfg(), device=device, gravity=(0.0, 0.0, 0.0)) - ) as sim: - fused = _spawn_collection("Fused", y_offset=0.0) - stepped = _spawn_collection("Stepped", y_offset=4.0) - wrenched = _spawn_collection("Wrenched", y_offset=8.0) - sim.reset() - yield _CollectionScene(sim=sim, device=device, fused=fused, stepped=stepped, wrenched=wrenched) - - -@pytest.mark.parametrize("scene", test_devices(), indirect=True) -def test_rigid_object_collection_real_ovphysx_seams(scene: _CollectionScene) -> 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) - - -@pytest.mark.parametrize("scene", test_devices(DeviceScope.CUDA), indirect=True) -def test_wrench_reaches_only_the_selected_collection_body(scene: _CollectionScene) -> None: - """Deliver an external force through the fused body-major binding to the selected body only, then reset it. - - 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", test_devices(DeviceScope.CUDA), 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: _CollectionScene) -> 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/unit/__init__.py b/source/isaaclab_ov/test/assets/unit/__init__.py deleted file mode 100644 index 460a30569089..000000000000 --- a/source/isaaclab_ov/test/assets/unit/__init__.py +++ /dev/null @@ -1,4 +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 diff --git a/source/isaaclab_ov/test/assets/unit/test_ovphysx_articulation.py b/source/isaaclab_ov/test/assets/unit/test_ovphysx_articulation.py deleted file mode 100644 index aeb11ebf3178..000000000000 --- a/source/isaaclab_ov/test/assets/unit/test_ovphysx_articulation.py +++ /dev/null @@ -1,329 +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 - -"""OVPhysX articulation unit tests: data caches, joint directions, tendon scoping, kernels, and actuator control.""" - -from __future__ import annotations - -import importlib -from types import SimpleNamespace -from unittest.mock import Mock - -import numpy as np -import pytest -import torch -import warp as wp - -from pxr import Sdf, Usd, UsdGeom, UsdPhysics - -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, 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 OvPhysxManager # noqa: E402 -from isaaclab_ov.test.fixtures.views import MockOvPhysxBindingSet # noqa: E402 - -from isaaclab.actuators import ImplicitActuatorCfg # noqa: E402 -from isaaclab.utils.warp.launch_cache import _WarpLaunchCache # noqa: E402 - -pytestmark = pytest.mark.unit - - -def test_cached_read_launches_reset_on_ordering_and_invalidation(): - """Ordering installation and simulation invalidation should discard recorded reads.""" - - class MinimalData(ArticulationData): - def __dir__(self): - return [] - - class Buffer: - timestamp = 1.0 - - data = MinimalData.__new__(MinimalData) - read_launch_cache = Mock() - data._read_launch_cache = read_launch_cache - data._configure_ordering_buffers = lambda: None - data._make_jacobian_body_user_to_backend = lambda: object() - data.joint_ordering = None - data._body_com_jacobian_w = Buffer() - data._mass_matrix = Buffer() - data._gravity_compensation_forces = Buffer() - - data._apply_ordering_maps_after_resolve() - - read_launch_cache.clear.assert_called_once_with() - assert data._body_com_jacobian_w.timestamp == -1.0 - assert data._mass_matrix.timestamp == -1.0 - assert data._gravity_compensation_forces.timestamp == -1.0 - - data._is_primed = True - data._sim_timestamp = 1.0 - data._invalidate_initialize_callback(None) - - assert read_launch_cache.clear.call_count == 2 - assert data._is_primed is False - assert data._sim_timestamp == 0.0 - - -def test_static_property_reads_are_not_invalidated_by_simulation_steps(): - """Joint properties and body mass/inertia should be read once per invalidation, not once per step. - - On OVPhysX these are blocking CPU-only binding reads whose cost scales with the number of - environments, so re-reading them every physics step made per-step consumers (such as the - native actuator telemetry sync) host-bound. State buffers must still refresh every step. - """ - - class Buffer: - def __init__(self, shape): - self.data = wp.zeros(shape, dtype=wp.float32, device="cpu") - self.timestamp = -1.0 - - data = ArticulationData.__new__(ArticulationData) - data.device = "cpu" - data.num_instances = 1 - data.num_joints = 2 - data._sim_timestamp = 1.0 - data.body_ordering = None - data._get_binding = lambda tensor_type: object() - reads: list[int] = [] - data._binding_read = lambda tensor_type, dst: reads.append(tensor_type) - - # Joint properties: one read across several steps, one more after explicit invalidation - # (simulation reinitialization). - data.joint_ordering = None - stiffness = Buffer((1, 2)) - for _ in range(3): - data._sim_timestamp += 1.0 - data._read_joint_property_binding(TT.DOF_STIFFNESS, stiffness, None) - assert reads.count(TT.DOF_STIFFNESS) == 1 - stiffness.timestamp = -1.0 - data._read_joint_property_binding(TT.DOF_STIFFNESS, stiffness, None) - assert reads.count(TT.DOF_STIFFNESS) == 2 - - # Body properties behave the same; body state buffers still refresh every step. - mass, link_pose = Buffer((1, 2)), Buffer((1, 2)) - for _ in range(3): - data._sim_timestamp += 1.0 - data._refresh_reordered_body_buffer(mass, None, TT.BODY_MASS, static=True) - data._refresh_reordered_body_buffer(link_pose, None, TT.LINK_POSE) - assert reads.count(TT.BODY_MASS) == 1 - assert reads.count(TT.LINK_POSE) == 3 - mass.timestamp = -1.0 - data._refresh_reordered_body_buffer(mass, None, TT.BODY_MASS, static=True) - assert reads.count(TT.BODY_MASS) == 2 - - # Under a non-identity joint ordering the property is gathered once, then served from cache. - data.joint_ordering = SimpleNamespace(user_to_backend=wp.array([1, 0], dtype=wp.int32, device="cpu")) - data._read_launch_cache = _WarpLaunchCache("cpu") - user_buffer, backend_buffer = Buffer((1, 2)), Buffer((1, 2)) - data._binding_read = lambda tensor_type, dst: (reads.append(tensor_type), dst.assign([[1.0, 2.0]])) - for _ in range(3): - data._sim_timestamp += 1.0 - data._read_joint_property_binding(TT.DOF_DAMPING, user_buffer, backend_buffer) - assert reads.count(TT.DOF_DAMPING) == 1 - torch.testing.assert_close(wp.to_torch(user_buffer.data), torch.tensor([[2.0, 1.0]])) - - -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,) - - -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"] - - -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]]) - - -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)) diff --git a/source/isaaclab_ov/test/assets/unit/test_ovphysx_deformable_object.py b/source/isaaclab_ov/test/assets/unit/test_ovphysx_deformable_object.py deleted file mode 100644 index 9ae5882c8daf..000000000000 --- a/source/isaaclab_ov/test/assets/unit/test_ovphysx_deformable_object.py +++ /dev/null @@ -1,383 +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 and its body view.""" - -from __future__ import annotations - -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 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.assets.deformable_object.views import OvPhysxDeformableBodyView # noqa: E402 -from isaaclab_ov.physics.ovphysx_manager import OvPhysxManager, OvPhysxSceneDataBackend # noqa: E402 - -pytestmark = pytest.mark.unit - - -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) - - -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 - - -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]])) - - -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")) - - -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, - ) From 5e13ef9b88c0fef5736ff02ed9625bc5065dbd6b Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Fri, 2 Oct 2026 16:34:54 -0700 Subject: [PATCH 33/33] Drop stale rigid-object-collection test timeout --- tools/test_settings.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tools/test_settings.py b/tools/test_settings.py index a5996d57ee6a..8ab74ce71154 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,