From 556bd814f2382e84610ed7d90659a4daa605c058 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 09:49:17 +0200 Subject: [PATCH 1/8] 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 06f7d93b5b7212c09b8e171eda87095c6f6d73c1 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 13:49:29 +0200 Subject: [PATCH 2/8] Run PhysX rigid-object tests on one composite scene Share one module-scoped scene of locally authored cubes instead of a Nucleus scene per test. Partial state, inertial, and wrench writes select environment 1 and read the PhysX view back; a center-of-mass offset checks every root-frame writer and the spin kinematics. Failure modes, the gravity vector, and physics warmup keep their own scenes. The 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. --- ...er-asset-tests-p14-physx-rigid-object.skip | 0 .../test/assets/test_rigid_object.py | 1082 ++++++----------- 2 files changed, 370 insertions(+), 712 deletions(-) create mode 100644 source/isaaclab_physx/changelog.d/antoiner-asset-tests-p14-physx-rigid-object.skip 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/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 From 4c897f37c22685db426abc21c5501853099ce5b3 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 13:49:29 +0200 Subject: [PATCH 3/8] 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 cddf9dd4d71fb3c290715e6dd4ee5f78531cc82a Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 15:04:29 +0200 Subject: [PATCH 4/8] Tighten style of the composite PhysX rigid-object tests Import the mock and PhysX interfaces at module level, allocate repeated rows directly, state the articulation-root xfail as the current contract, and annotate the tests and the scene fixture. --- .../test/assets/test_rigid_object.py | 35 ++++++++++--------- 1 file changed, 19 insertions(+), 16 deletions(-) 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 ) From 200e176a2ada42f0a812c89a1a921c4c77cbeff4 Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 15:04:54 +0200 Subject: [PATCH 5/8] 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 6/8] 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 a35418cd74501bdf0d3364ab5ae0b776484c433a Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 16:30:09 +0200 Subject: [PATCH 7/8] Pin the articulation root gap and tighten rigid object keepers The strict xfail now authors an enabled articulation root on the rigid body itself and only accepts the missing-rejection failure, so it turns red once the check is restored. Every test carries the Isaac Sim CI marker, the invalid-body and gravity tests run on every device, the writer passes step in between, the wrench node sets its own masses, and the partial velocity write is read back from PhysX. --- .../test/assets/test_rigid_object.py | 71 +++++++++++++------ 1 file changed, 48 insertions(+), 23 deletions(-) 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) From 197372f52cf7213a4dbe75755a238214fac9c98f Mon Sep 17 00:00:00 2001 From: AntoineRichard Date: Tue, 29 Sep 2026 16:32:31 +0200 Subject: [PATCH 8/8] 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(