From 556bd814f2382e84610ed7d90659a4daa605c058 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 09:49:17 +0200 Subject: [PATCH 1/5] 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 4c897f37c22685db426abc21c5501853099ce5b3 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 13:49:29 +0200 Subject: [PATCH 2/5] Run PhysX rigid-object-collection tests on one composite scene Share one module-scoped scene of locally authored cubes instead of a Nucleus scene per test. Partial state, mass, and wrench writes select non-sorted body-major subsets and read the PhysX view back; a center-of-mass offset checks every body-frame writer and the spin kinematics. Failure modes, the gravity vector, and the subset-write regressions keep their own scenes. --- ...sts-p15-physx-rigid-object-collection.skip | 0 .../assets/test_rigid_object_collection.py | 866 +++++++----------- 2 files changed, 346 insertions(+), 520 deletions(-) create mode 100644 source/isaaclab_physx/changelog.d/antoiner-asset-tests-p15-physx-rigid-object-collection.skip 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/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 200e176a2ada42f0a812c89a1a921c4c77cbeff4 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 15:04:54 +0200 Subject: [PATCH 3/5] Tighten style of the composite PhysX collection tests Drop step-narrating comments from the gravity test and annotate the tests and fixtures. --- .../assets/test_rigid_object_collection.py | 25 ++++++++----------- 1 file changed, 11 insertions(+), 14 deletions(-) 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 From a8300a319382ead1b7972bdbd994b08c67c84487 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 15:59:50 +0200 Subject: [PATCH 4/5] Mention the mask setters in the collection fix changelog --- .../antoiner-fix-physx-collection-com-inertia.rst | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) 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. From 197372f52cf7213a4dbe75755a238214fac9c98f Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 16:32:31 +0200 Subject: [PATCH 5/5] Check permanent wrench persistence in the collection wrench node A second step must double the velocity from a permanent force, which the deleted single-body force test covered. The writer passes step in between, the link velocity write is read back from PhysX, and the invalid-body and gravity tests run on every device again. --- .../assets/test_rigid_object_collection.py | 50 +++++++++++++------ 1 file changed, 35 insertions(+), 15 deletions(-) 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(