diff --git a/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p14-physx-rigid-object.skip b/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p14-physx-rigid-object.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p15-physx-rigid-object-collection.skip b/source/isaaclab_physx/changelog.d/antoiner-asset-tests-p15-physx-rigid-object-collection.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_physx/changelog.d/antoiner-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..531bfb955601 --- /dev/null +++ b/source/isaaclab_physx/changelog.d/antoiner-fix-physx-collection-com-inertia.rst @@ -0,0 +1,13 @@ +Fixed +^^^^^ + +* Fixed :meth:`~isaaclab_physx.assets.RigidObjectCollection.set_coms_index` and + :meth:`~isaaclab_physx.assets.RigidObjectCollection.set_coms_mask` raising a ``KeyError`` because the + center-of-mass poses were passed to the PhysX tensor API as ``wp.transformf`` instead of ``float32``. +* Fixed :meth:`~isaaclab_physx.assets.RigidObjectCollection.set_coms_index` and + :meth:`~isaaclab_physx.assets.RigidObjectCollection.set_coms_mask` leaving the unselected entries of + :attr:`~isaaclab_physx.assets.RigidObjectCollectionData.body_com_pose_b` zeroed when they were called for a + subset of bodies or environments before the property was first read. +* Fixed :meth:`~isaaclab_physx.assets.RigidObjectCollection.set_inertias_index` and + :meth:`~isaaclab_physx.assets.RigidObjectCollection.set_inertias_mask` failing in the PhysX backend because the + inertias were passed to the PhysX tensor API with the wrong layout. diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py index 341d60b62bbb..fc5f1245589d 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py @@ -1032,13 +1032,13 @@ def set_coms_index( ], device=self.device, ) - # Invalidate the cached buffer - self.data._body_com_pose_b.timestamp = self.data._sim_timestamp + # Invalidate the cached buffers. The next read refreshes from the view, so that a partial write does not + # expose entries of the buffer that were never read back from the simulation. self.data._reset_body_com_pose_b_dependents() # Set into simulation, note that when updating "model" properties with PhysX we need to do it on CPU. # Convert from instance order (num_instances, num_bodies, 7) to view order (num_bodies*num_instances, 7) for # PhysX. - com_view_order = self.reshape_data_to_view_2d(self.data._body_com_pose_b.data, device="cpu") # (B*I, 7) + com_view_order = self.reshape_data_to_view_2d(self.data._body_com_pose_b.data, device="cpu").view(wp.float32) view_ids = self._env_body_ids_to_view_ids(env_ids, body_ids, device="cpu") self.root_view.set_coms(com_view_order, indices=view_ids) @@ -1128,8 +1128,9 @@ def set_inertias_index( device=self.device, ) # Set into simulation, note that when updating "model" properties with PhysX we need to do it on CPU. - # Convert from instance order (num_instances, num_bodies) to view order for PhysX. - inertia_view_order = self.reshape_data_to_view_2d(self.data._body_inertia, device="cpu") + # Convert from instance order (num_instances, num_bodies, 9) to view order (num_bodies*num_instances, 9) for + # PhysX. + inertia_view_order = self.reshape_data_to_view_3d(self.data._body_inertia, 9, device="cpu") view_ids = self._env_body_ids_to_view_ids(env_ids, body_ids, device="cpu") self.root_view.set_inertias(inertia_view_order, indices=view_ids) diff --git a/source/isaaclab_physx/test/assets/test_rigid_object.py b/source/isaaclab_physx/test/assets/test_rigid_object.py index ac84e269f1b6..af3057c2e342 100644 --- a/source/isaaclab_physx/test/assets/test_rigid_object.py +++ b/source/isaaclab_physx/test/assets/test_rigid_object.py @@ -3,6 +3,14 @@ # # SPDX-License-Identifier: BSD-3-Clause +"""Real PhysX rigid-object coverage. + +Most checks run against one module-scoped scene of locally authored cubes. Each rigid object holds two +environments so that partial writes can target environment 1 and prove that environment 0 is preserved in the +real PhysX state. Tests that need their own simulation context are defined first: pytest runs them before the +composite scene is created, and a new simulation context would replace the composite stage. +""" + # ignore private usage of variables warning # pyright: reportPrivateUsage=none @@ -10,570 +18,125 @@ launch_test_simulation() +import math import sys +from collections.abc import Iterator +from dataclasses import dataclass from typing import Literal +from unittest.mock import MagicMock, patch import pytest import torch import warp as wp -from flaky import flaky from isaaclab_physx.assets import RigidObject +from isaaclab_physx.sim.schemas import PhysxArticulationCfg, PhysxRigidBodyCfg + +import omni.kit.app +import omni.physx +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.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 + + +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() - - # Check if object is initialized - assert cube_object.is_initialized - assert len(cube_object.body_names) == 1 - - # Check buffers that exists and have correct shapes - assert cube_object.data.root_pos_w.torch.shape == (num_cubes, 3) - assert cube_object.data.root_quat_w.torch.shape == (num_cubes, 4) - # Simulate physics - for _ in range(2): - # perform rendering - sim.step() - # update object - cube_object.update(sim.cfg.dt) - # check that the object is kinematic - default_root_pose = cube_object.data.default_root_pose.torch.clone() - default_root_vel = cube_object.data.default_root_vel.torch.clone() - default_root_pose[:, :3] += origins - torch.testing.assert_close(cube_object.data.root_link_pose_w.torch, default_root_pose) - torch.testing.assert_close(cube_object.data.root_com_vel_w.torch, default_root_vel) +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 -@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) +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)) - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(cube_object) < 10 - # Play sim - with pytest.raises(RuntimeError): - sim.reset() +## +# Tests that own their simulation context. Keep them above the composite scene. +## -@pytest.mark.parametrize("num_cubes", [2]) +@pytest.mark.parametrize( + "api", + [ + "none", + pytest.param( + "articulation_root", + marks=pytest.mark.xfail( + strict=True, + raises=pytest.fail.Exception, + reason="The PhysX rigid object does not reject an enabled articulation root.", + ), + ), + ], +) @pytest.mark.parametrize("device", test_devices()) @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"]) -> 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 - # 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")) + # An enabled articulation root on the rigid body turns the body into an articulation link. + assert apply_articulation_root_properties( + "/World/Env_[^/]*/Object", + [PhysxArticulationCfg(articulation_enabled=True)], + stage=sim.stage, + create_if_missing=True, + ) + for index in range(_NUM_ENVS): + prim = sim.stage.GetPrimAtPath(f"/World/Env_{index}/Object") + assert prim.HasAPI(UsdPhysics.ArticulationRootAPI) + assert prim.GetAttribute("physxArticulation:articulationEnabled").Get() # Check that the framework doesn't hold excessive strong references. assert sys.getrefcount(cube_object) < 10 - # 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("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) -> 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 - # 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 +160,229 @@ 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: +@pytest.mark.parametrize("device", test_devices()) +def test_warmup_loads_physics_once(device) -> None: + """Attach on GPU or force-load on CPU, without destroying and rebuilding native objects.""" + with build_simulation_context(device=device, add_ground_plane=True, dt=0.01, auto_add_lighting=True) as sim: sim._app_control_on_stop_handle = None - # 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) + sim_utils.create_prim("/World/Table_0", "Xform", translation=(0.0, 0.0, 1.0)) + cube = RigidObject(cfg=_cube_cfg("/World/Table_[^/]*/Object")) - com = wp.to_torch(cube_object.root_view.get_coms()) - com[..., :3] = offset.to("cpu") - cube_object.set_coms_index(coms=wp.from_torch(com.unsqueeze(1).to(device).contiguous(), dtype=wp.transformf)) - - # check center of mass has been set - torch.testing.assert_close(wp.to_torch(cube_object.root_view.get_coms()), com) - - # 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) -> 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: 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)) +@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 + 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) - # 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) +@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 + device = scene.device + cube_object = scene.cubes + + # Full-data mass writes reach the view and survive a simulation step. + original_masses = wp.to_torch(cube_object.root_view.get_masses()).clone() + assert original_masses.shape == (_NUM_ENVS, 1) + masses = original_masses + torch.tensor([[4.5], [6.5]]) + cube_object.set_masses_index(masses=masses.to(device)) + torch.testing.assert_close(cube_object.data.body_mass.torch, masses.to(device)) + torch.testing.assert_close(wp.to_torch(cube_object.root_view.get_masses()), masses) + scene.step() + torch.testing.assert_close(cube_object.data.body_mass.torch, masses.to(device)) + torch.testing.assert_close(wp.to_torch(cube_object.root_view.get_masses()), masses) + + # Partial writes with an int64 selector change environment 1 only. + env_ids = torch.tensor([1], dtype=torch.int64, device=device) + body_ids = torch.tensor([0], dtype=torch.int32, device=device) + expected = { + "body_mass": cube_object.data.body_mass.torch.clone(), + "body_com_pose_b": cube_object.data.body_com_pose_b.torch.clone(), + "body_inertia": cube_object.data.body_inertia.torch.clone(), + } + expected["body_mass"][1, 0] = 3.0 + expected["body_com_pose_b"][1, 0, :3] = torch.tensor([0.03, -0.02, 0.01], device=device) + expected["body_inertia"][1, 0, [0, 4, 8]] *= torch.tensor([1.2, 1.3, 1.4], device=device) + cube_object.set_masses_index(masses=expected["body_mass"][1:], env_ids=env_ids, body_ids=body_ids) + cube_object.set_coms_index(coms=expected["body_com_pose_b"][1:], env_ids=env_ids, body_ids=body_ids) + cube_object.set_inertias_index(inertias=expected["body_inertia"][1:], env_ids=env_ids, body_ids=body_ids) + for name, data, raw in ( + ("body_mass", cube_object.data.body_mass, cube_object.root_view.get_masses()), + ("body_com_pose_b", cube_object.data.body_com_pose_b, cube_object.root_view.get_coms().view(wp.float32)), + ("body_inertia", cube_object.data.body_inertia, cube_object.root_view.get_inertias()), + ): + torch.testing.assert_close(data.torch, expected[name]) + torch.testing.assert_close(wp.to_torch(raw).to(device).reshape(expected[name].shape), expected[name]) - env_idx = env_idx.to(device) - for i in range(10): - # perform step - sim.step() - # update buffers - cube_object.update(sim.cfg.dt) - env_ids = None if i % 2 == 0 else env_idx +@pytest.mark.isaacsim_ci +def test_rigid_object_root_state_writes(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], device=device).repeat(_NUM_ENVS, 1)), -1 + ) + rest_pose[:, 2] += 1.0 + scene.place_cubes_at_rest(rest_pose) + + # A partial pose and velocity write reaches only the selected root. + initial_pose = cube_object.data.root_link_pose_w.torch.clone() + initial_velocity = cube_object.data.root_link_vel_w.torch.clone() + initial_com_velocity = cube_object.data.root_com_vel_w.torch.clone() + target_pose = initial_pose[1:].clone() + target_pose[:, :3] += torch.tensor([0.25, -0.1, 0.3], device=device) + target_pose[:, 3:] = torch.tensor(_yaw_quat(0.4), device=device) + target_velocity = torch.tensor([[0.0, 0.2, 0.0, 0.0, 0.0, 0.1]], device=device) + cube_object.write_root_link_pose_to_sim_index(root_pose=target_pose, env_ids=[1]) + cube_object.write_root_link_velocity_to_sim_index(root_velocity=target_velocity, env_ids=[1]) + torch.testing.assert_close(cube_object.data.root_link_pose_w.torch[1:], target_pose) + torch.testing.assert_close(cube_object.data.root_link_vel_w.torch[1:], target_velocity) + torch.testing.assert_close(cube_object.data.root_link_pose_w.torch[:1], initial_pose[:1]) + torch.testing.assert_close(cube_object.data.root_link_vel_w.torch[:1], initial_velocity[:1]) + raw_pose = wp.to_torch(cube_object.root_view.get_transforms()).to(device) + torch.testing.assert_close(raw_pose, torch.cat((initial_pose[:1], target_pose))) + # PhysX holds the center-of-mass velocity: the link velocity plus the angular velocity crossed with the + # world-frame center-of-mass offset. + com_offset_w = quat_apply(target_pose[:, 3:], cube_object.data.body_com_pos_b.torch[1:, 0]) + target_com_velocity = target_velocity.clone() + target_com_velocity[:, :3] += torch.linalg.cross(target_velocity[:, 3:], com_offset_w) + raw_velocity = wp.to_torch(cube_object.root_view.get_velocities()).to(device) + torch.testing.assert_close(raw_velocity, torch.cat((initial_com_velocity[:1], target_com_velocity))) + + # With a center-of-mass offset, every writer round-trips through its own frame. + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(_NUM_ENVS, 1) + coms = cube_object.data.body_com_pose_b.torch.clone() + coms[:, 0, :3] = offset + coms[:, 0, 3:] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) + cube_object.set_coms_index(coms=coms) + # random pose and velocity so frame conversions see a non-trivial rotation + rand_state = torch.rand(_NUM_ENVS, 13, device=device) + rand_state[..., :3] += rest_pose[:, :3] + rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) + env_idx = torch.arange(_NUM_ENVS, dtype=torch.int32, device=device) + for state_location in ("com", "link", "root"): + for env_ids in (None, env_idx): if state_location == "com": cube_object.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_ids) cube_object.write_root_com_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_ids) - elif state_location == "link": - cube_object.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_ids) - cube_object.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_ids) - elif state_location == "root": - cube_object.write_root_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_ids) - cube_object.write_root_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_ids) - - if state_location == "com": torch.testing.assert_close(rand_state[..., :7], cube_object.data.root_com_pose_w.torch) torch.testing.assert_close(rand_state[..., 7:], cube_object.data.root_com_vel_w.torch) elif state_location == "link": + cube_object.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_ids) + cube_object.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_ids) torch.testing.assert_close(rand_state[..., :7], cube_object.data.root_link_pose_w.torch) torch.testing.assert_close(rand_state[..., 7:], cube_object.data.root_link_vel_w.torch) - elif state_location == "root": + else: + cube_object.write_root_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_ids) + cube_object.write_root_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_ids) torch.testing.assert_close(rand_state[..., :7], cube_object.data.root_link_pose_w.torch) torch.testing.assert_close(rand_state[..., 7:], cube_object.data.root_com_vel_w.torch) @@ -789,39 +401,113 @@ def test_write_root_state(num_cubes, device, with_offset, state_location): torch.testing.assert_close( cube_object.data.root_com_vel_w.torch[:, 3:], cube_object.data.root_link_vel_w.torch[:, 3:] ) + # Move away from the written state so that the next pass writes it again. + scene.step() + + # Spinning about the center of mass keeps it in place while the link frame orbits it. + scene.place_cubes_at_rest(rest_pose) + spin_twist = torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0, 1.3], device=device).repeat(_NUM_ENVS, 1) + initial_com_pos = cube_object.data.root_com_pose_w.torch[:, :3].clone() + for _ in range(10): + cube_object.write_root_com_velocity_to_sim_index(root_velocity=spin_twist) + scene.step() + root_link_pose_w = cube_object.data.root_link_pose_w.torch + root_link_vel_w = cube_object.data.root_link_vel_w.torch + root_com_pose_w = cube_object.data.root_com_pose_w.torch + root_com_vel_w = cube_object.data.root_com_vel_w.torch + body_link_pose_w = cube_object.data.body_link_pose_w.torch + body_com_pose_w = cube_object.data.body_com_pose_w.torch + # center of mass position will be constant (i.e. spinning around com) + torch.testing.assert_close(root_com_pose_w[:, :3], initial_com_pos) + torch.testing.assert_close(body_com_pose_w[:, 0, :3], initial_com_pos) + # link position will be moving but should stay constant away from center of mass + torch.testing.assert_close( + -offset, quat_apply_inverse(root_link_pose_w[:, 3:], root_link_pose_w[:, :3] - root_com_pose_w[:, :3]) + ) + # orientation of com will be a constant rotation from link orientation + com_quat_w = quat_mul(body_link_pose_w[..., 3:], cube_object.data.body_com_quat_b.torch) + torch.testing.assert_close(com_quat_w, body_com_pose_w[..., 3:]) + torch.testing.assert_close(com_quat_w[:, 0], root_com_pose_w[:, 3:]) + # center of mass is at rest while the link frame moves with angular velocity cross offset + torch.testing.assert_close(torch.zeros_like(root_com_vel_w[:, :3]), root_com_vel_w[:, :3]) + lin_vel_rel_root = quat_apply_inverse(root_link_pose_w[:, 3:], root_link_vel_w[:, :3]) + lin_vel_rel_gt = torch.linalg.cross(spin_twist[:, 3:], -offset) + torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_root, atol=1e-4, rtol=1e-4) + # ang_vel will always match + torch.testing.assert_close(root_com_vel_w[:, 3:], root_link_vel_w[:, 3:]) @pytest.mark.isaacsim_ci -@pytest.mark.parametrize("device", test_devices()) -def test_warmup_loads_physics_once(device): - """Attach on GPU or force-load on CPU, without destroying and rebuilding native objects.""" - from unittest.mock import MagicMock, patch - - import omni.kit.app - import omni.physx - - 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() +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 + cube_object = scene.cubes + rest_pose = torch.cat( + (scene.origins["cubes"], torch.tensor(_yaw_quat(0.5 * math.pi), device=device).expand(_NUM_ENVS, 4)), -1 + ) + rest_pose[:, 2] += 1.0 + # Give both environments the inertial properties of a uniform 1 kg cube so that their responses are comparable. + cube_object.set_masses_index(masses=torch.ones((_NUM_ENVS, 1), device=device)) + coms = torch.zeros((_NUM_ENVS, 1, 7), device=device) + coms[..., 6] = 1.0 + cube_object.set_coms_index(coms=coms) + inertias = torch.zeros((_NUM_ENVS, 1, 9), device=device) + inertias[..., [0, 4, 8]] = 1.0 * (0.2**2 + 0.2**2) / 12.0 + cube_object.set_inertias_index(inertias=inertias) + + # A body-frame force on environment 1 accelerates only that cube, along the rotated force direction. + scene.place_cubes_at_rest(rest_pose) + cube_object.permanent_wrench_composer.set_forces_and_torques_index( + forces=torch.tensor([[[6.0, 0.0, 0.0]]], device=device), env_ids=[1] + ) + scene.step() + velocity = cube_object.data.root_com_lin_vel_w.torch + assert velocity[1, 1] > 1e-2, velocity + torch.testing.assert_close(velocity[0], torch.zeros(3, device=device), atol=1e-5, rtol=0.0) + + # A world-frame force, and a force at a position, match their body-frame counterparts. + for local_wrench, global_wrench, response in ( + ({"forces": [[[6.0, 0.0, 0.0]]]}, {"forces": [[[0.0, 6.0, 0.0]]]}, lambda: cube_object.data.root_com_lin_vel_w), + ( + {"forces": [[[0.0, 0.0, 6.0]]], "positions": [[[0.0, 0.1, 0.0]]]}, + {"forces": [[[0.0, 0.0, 6.0]]], "positions": [[[-0.1, 0.0, 0.0]]]}, + lambda: cube_object.data.root_com_ang_vel_b, + ), + ): + scene.place_cubes_at_rest(rest_pose) + local_wrench = {key: torch.tensor(value, device=device) for key, value in local_wrench.items()} + global_wrench = {key: torch.tensor(value, device=device) for key, value in global_wrench.items()} + if "positions" in global_wrench: + # World-frame positions add the rotated body-frame lever arm to the center of mass. + global_wrench["positions"] = global_wrench["positions"] + cube_object.data.body_com_pos_w.torch[:1] + cube_object.permanent_wrench_composer.set_forces_and_torques_index(env_ids=[1], **local_wrench) + cube_object.permanent_wrench_composer.set_forces_and_torques_index(env_ids=[0], is_global=True, **global_wrench) + scene.step() + response_value = response().torch + torch.testing.assert_close(response_value[0], response_value[1], atol=1e-4, rtol=1e-3) + # An upward force 0.1 m along the body y-axis rolls the cube about its x-axis. + assert torch.all(cube_object.data.root_com_ang_vel_b.torch[:, 0] > 0.1) + assert torch.all( + quat_apply( + cube_object.data.root_link_quat_w.torch, torch.tensor([1.0, 0.0, 0.0], device=device).repeat(_NUM_ENVS, 1) + )[:, 1] + > 0.99 + ) - 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]) + # A partial reset clears only the selected environment; a full reset deactivates the composers. + composers = (cube_object.instantaneous_wrench_composer, cube_object.permanent_wrench_composer) + ones = torch.ones((_NUM_ENVS, 1, 3), device=device) + cube_object.permanent_wrench_composer.set_forces_and_torques_index(forces=ones, torques=ones) + cube_object.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones) + cube_object.reset(env_ids=torch.tensor([0], device=device)) + for composer in composers: + assert composer.active + for buffer in (composer.out_force_b.torch, composer.out_torque_b.torch): + assert torch.count_nonzero(buffer[0]) == 0 + assert torch.count_nonzero(buffer[1:]) == buffer[1:].numel() + cube_object.reset() + for composer in composers: + assert not composer.active + assert torch.count_nonzero(composer.out_force_b.torch) == 0 + assert torch.count_nonzero(composer.out_torque_b.torch) == 0 diff --git a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py index 2e0af9a4be8b..97d10830f887 100644 --- a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py @@ -3,6 +3,14 @@ # # 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 @@ -10,24 +18,21 @@ launch_test_simulation() +import math import sys +from collections.abc import Iterator +from dataclasses import dataclass import pytest import torch import warp as wp from isaaclab_physx.assets import RigidObjectCollection +from isaaclab_physx.sim.schemas import PhysxRigidBodyCfg import isaaclab.sim as sim_utils from isaaclab.assets import RigidObjectCfg, RigidObjectCollectionCfg -from isaaclab.sim import 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 +41,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 +53,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.""" +def sim(request) -> Iterator[SimulationContext]: + """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)]) +## +# Tests that own their simulation context. Keep them above the composite scene. +## + + +@pytest.mark.parametrize("num_cubes", [2]) @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) +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) # 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): +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, device=device) + object_collection, _ = generate_cubes_scene(num_envs=2, num_cubes=3) # Play sim sim.reset() @@ -152,458 +151,371 @@ 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("num_envs", [2]) -@pytest.mark.parametrize("num_cubes", [3]) @pytest.mark.parametrize("device", test_devices()) -def test_initialization_with_kinematic_enabled(sim, num_envs, num_cubes, device): - """Test that initialization for prim with kinematic flag enabled.""" - object_collection, origins = generate_cubes_scene( - num_envs=num_envs, num_cubes=num_cubes, kinematic_enabled=True, device=device - ) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(object_collection) < 10 - - # Play sim +def test_inertial_property_subset_writes_reach_selected_view_entries(sim, device) -> None: + """COM and inertia writes to a subset of (env, body) cells must reach only those PhysX view entries.""" + num_envs, num_cubes = 2, 3 + object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes) sim.reset() - # Check if object is initialized - assert object_collection.is_initialized - assert len(object_collection.body_names) == num_cubes - - # Check buffers that exist and have correct shapes - assert object_collection.data.body_link_pos_w.torch.shape == (num_envs, num_cubes, 3) - assert object_collection.data.body_link_quat_w.torch.shape == (num_envs, num_cubes, 4) - - # Simulate physics - for _ in range(2): - sim.step() - object_collection.update(sim.cfg.dt) - # check that the object is kinematic - default_body_pose = object_collection.data.default_body_pose.torch.clone() - default_body_vel = object_collection.data.default_body_vel.torch.clone() - default_body_pose[..., :3] += origins.unsqueeze(1) - torch.testing.assert_close(object_collection.data.body_link_pose_w.torch, default_body_pose) - torch.testing.assert_close(object_collection.data.body_link_vel_w.torch, default_body_vel) - - -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -def test_initialization_with_no_rigid_body(sim, num_cubes, device): - """Test that initialization fails when no rigid body is found at the provided prim path.""" - object_collection, _ = generate_cubes_scene(num_cubes=num_cubes, has_api=False, device=device) - - # Check that the framework doesn't hold excessive strong references. - assert sys.getrefcount(object_collection) < 10 - - # Play sim - with pytest.raises(RuntimeError): - sim.reset() - - -@pytest.mark.parametrize("num_envs", [2]) -@pytest.mark.parametrize("num_cubes", [4]) -@pytest.mark.parametrize("device", test_devices()) -def test_external_force_on_single_body(sim, num_envs, num_cubes, device): - """Test application of external force on the base of the object.""" - object_collection, origins = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) - sim.reset() - - # find objects to apply the force - object_ids, object_names = object_collection.find_bodies(".*") - - # Sample a force equal to the weight of the object - external_wrench_b = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=sim.device) - # Every 2nd cube should have a force applied to it - external_wrench_b[:, 0::2, 2] = 9.81 * object_collection.data.body_mass.torch[:, 0::2] - - for i in range(5): - # reset object state - body_pose = object_collection.data.default_body_pose.torch.clone() - body_vel = object_collection.data.default_body_vel.torch.clone() - # need to shift the position of the cubes otherwise they will be on top of each other - body_pose[..., :2] += origins.unsqueeze(1)[..., :2] - object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) - object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) - # reset object - object_collection.reset() - - is_global = False - if i % 2 == 0: - positions = object_collection.data.body_link_pos_w.torch[:, object_ids, :3] - is_global = True - else: - positions = None - - # apply force - object_collection.permanent_wrench_composer.set_forces_and_torques_index( - forces=external_wrench_b[..., :3], - torques=external_wrench_b[..., 3:], - positions=positions, - body_ids=object_ids, - env_ids=None, - is_global=is_global, - ) - for _ in range(10): - # write data to sim - object_collection.write_data_to_sim() - # step sim - sim.step() - # update object collection - object_collection.update(sim.cfg.dt) - - # First object should still be at the same Z position (1.0) - torch.testing.assert_close( - object_collection.data.body_link_pos_w.torch[:, 0::2, 2], - torch.ones_like(object_collection.data.body_link_pos_w.torch[:, 0::2, 2]), - ) - # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(object_collection.data.body_link_pos_w.torch[:, 1::2, 2] < 1.0) - - -@pytest.mark.parametrize("num_envs", [2]) -@pytest.mark.parametrize("num_cubes", [4]) -@pytest.mark.parametrize("device", test_devices()) -def test_external_force_on_single_body_at_position(sim, num_envs, num_cubes, device): - """Test application of external force on the base of the object at a specific position. - - In this test, we apply a force equal to the weight of an object on the base of - one of the objects at 1m in the Y direction, we check that the object rotates around it's X axis. - For the other object, we do not apply any force and check that it falls down. - """ - object_collection, origins = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) - sim.reset() - - # find objects to apply the force - object_ids, object_names = object_collection.find_bodies(".*") - - # Sample a force equal to the weight of the object - external_wrench_b = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=sim.device) - external_wrench_positions_b = torch.zeros(object_collection.num_instances, len(object_ids), 3, device=sim.device) - # Every 2nd cube should have a force applied to it - external_wrench_b[:, 0::2, 2] = 500.0 - external_wrench_positions_b[:, 0::2, 1] = 1.0 - - # Desired force and torque - for i in range(5): - # reset object state - body_pose = object_collection.data.default_body_pose.torch.clone() - body_vel = object_collection.data.default_body_vel.torch.clone() - # need to shift the position of the cubes otherwise they will be on top of each other - body_pose[..., :2] += origins.unsqueeze(1)[..., :2] - object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) - object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) - # reset object - object_collection.reset() - - is_global = False - if i % 2 == 0: - body_com_pos_w = object_collection.data.body_link_pos_w.torch[:, object_ids, :3] - external_wrench_positions_b[..., 0] = 0.0 - external_wrench_positions_b[..., 1] = 1.0 - external_wrench_positions_b[..., 2] = 0.0 - external_wrench_positions_b += body_com_pos_w - is_global = True - else: - external_wrench_positions_b[..., 0] = 0.0 - external_wrench_positions_b[..., 1] = 1.0 - external_wrench_positions_b[..., 2] = 0.0 - - # apply force - object_collection.permanent_wrench_composer.set_forces_and_torques_index( - forces=external_wrench_b[..., :3], - torques=external_wrench_b[..., 3:], - positions=external_wrench_positions_b, - body_ids=object_ids, - env_ids=None, - is_global=is_global, - ) - object_collection.permanent_wrench_composer.add_forces_and_torques_index( - forces=external_wrench_b[..., :3], - torques=external_wrench_b[..., 3:], - positions=external_wrench_positions_b, - body_ids=object_ids, - is_global=is_global, - ) - - for _ in range(10): - # write data to sim - object_collection.write_data_to_sim() - # step sim - sim.step() - # update object collection - object_collection.update(sim.cfg.dt) - - # First object should be rotating around it's X axis - assert torch.all(object_collection.data.body_com_ang_vel_b.torch[:, 0::2, 0] > 0.1) - # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(object_collection.data.body_link_pos_w.torch[:, 1::2, 2] < 1.0) + 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", [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() +@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.""" + object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes) - state_types = ["body_link_pos_w", "body_link_quat_w", "body_com_lin_vel_w", "body_com_ang_vel_w"] - - # Set each state type individually as they are dependent on each other - for state_type_to_randomize in state_types: - state_dict = { - "body_link_pos_w": torch.zeros_like(object_collection.data.body_link_pos_w.torch, device=sim.device), - "body_link_quat_w": default_orientation(num=num_cubes * num_envs, device=sim.device).view( - num_envs, num_cubes, 4 - ), - "body_com_lin_vel_w": torch.zeros_like(object_collection.data.body_com_lin_vel_w.torch, device=sim.device), - "body_com_ang_vel_w": torch.zeros_like(object_collection.data.body_com_ang_vel_w.torch, device=sim.device), - } - - for _ in range(5): - # reset object - object_collection.reset() - - # Set random state - if state_type_to_randomize == "body_link_quat_w": - state_dict[state_type_to_randomize] = random_orientation( - num=num_cubes * num_envs, device=sim.device - ).view(num_envs, num_cubes, 4) - else: - state_dict[state_type_to_randomize] = torch.randn(num_envs, num_cubes, 3, device=sim.device) - # make sure objects do not overlap - if state_type_to_randomize == "body_link_pos_w": - state_dict[state_type_to_randomize][..., :2] += origins.unsqueeze(1)[..., :2] - - # perform simulation - for _ in range(5): - body_pose = torch.cat( - [state_dict["body_link_pos_w"], state_dict["body_link_quat_w"]], - dim=-1, - ) - body_vel = torch.cat( - [state_dict["body_com_lin_vel_w"], state_dict["body_com_ang_vel_w"]], - dim=-1, - ) - # reset object state - object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) - object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) - sim.step() - - # assert that set object quantities are equal to the ones set in the state_dict - for key, expected_value in state_dict.items(): - value = getattr(object_collection.data, key).torch - torch.testing.assert_close(value, expected_value, rtol=1e-5, atol=1e-5) - - object_collection.update(sim.cfg.dt) - - -@pytest.mark.parametrize("num_envs", [4]) -@pytest.mark.parametrize("num_cubes", [2]) -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("with_offset", [True, False]) -@pytest.mark.parametrize("gravity_enabled", [False]) -def test_object_state_properties(sim, num_envs, num_cubes, device, with_offset, gravity_enabled): - """Test the object_com_state_w and object_link_state_w properties.""" - cube_object, env_pos = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, height=0.0, device=device) - view_ids = torch.tensor([x for x in range(num_cubes * num_envs)], dtype=torch.int32) + gravity_dir = (0.0, 0.0, -1.0) if gravity_enabled else (0.0, 0.0, 0.0) 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] + gravity_vec = object_collection.data.GRAVITY_VEC_W.torch + assert gravity_vec[0, 0, 0] == gravity_dir[0] + assert gravity_vec[0, 0, 1] == gravity_dir[1] + assert gravity_vec[0, 0, 2] == gravity_dir[2] - for 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)) + for _ in range(2): sim.step() - cube_object.update(sim.cfg.dt) - - # get state properties - object_link_pose_w = cube_object.data.body_link_pose_w.torch - object_link_vel_w = cube_object.data.body_link_vel_w.torch - object_com_pose_w = cube_object.data.body_com_pose_w.torch - object_com_vel_w = cube_object.data.body_com_vel_w.torch - - # if offset is [0,0,0] all object_state_%_w will match and all body_%_w will match - if not with_offset: - torch.testing.assert_close(object_link_pose_w, object_com_pose_w) - torch.testing.assert_close(object_com_vel_w, object_link_vel_w) - else: - # cubes are spinning around center of mass - # position will not match - # center of mass position will be constant (i.e. spinning around com) - torch.testing.assert_close(init_com, object_com_pose_w[..., :3]) - - # link position will be moving but should stay constant away from center of mass - object_link_state_pos_rel_com = quat_apply_inverse( - object_link_pose_w[..., 3:], - object_link_pose_w[..., :3] - object_com_pose_w[..., :3], - ) - - torch.testing.assert_close(-offset, object_link_state_pos_rel_com) - - # orientation of com will be a constant rotation from link orientation - com_quat_b = cube_object.data.body_com_quat_b.torch - com_quat_w = quat_mul(object_link_pose_w[..., 3:], com_quat_b) - torch.testing.assert_close(com_quat_w, object_com_pose_w[..., 3:]) - - # lin_vel will not match - # center of mass vel will be constant (i.e. spinning around com) - torch.testing.assert_close( - torch.zeros_like(object_com_vel_w[..., :3]), - object_com_vel_w[..., :3], - ) - - # link frame will be moving, and should be equal to input angular velocity cross offset - lin_vel_rel_object_gt = quat_apply_inverse(object_link_pose_w[..., 3:], object_link_vel_w[..., :3]) - lin_vel_rel_gt = torch.linalg.cross(spin_twist.repeat(num_envs, num_cubes, 1)[..., 3:], -offset) - torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_object_gt, atol=1e-4, rtol=1e-3) - - # ang_vel will always match - torch.testing.assert_close(object_com_vel_w[..., 3:], object_link_vel_w[..., 3:]) + 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 -@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. + torch.testing.assert_close(object_collection.data.body_com_acc_w.torch, gravity) - 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() +## +# 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) + ) + self.cubes.permanent_wrench_composer.reset() + self.cubes.instantaneous_wrench_composer.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) - ) +@pytest.fixture(scope="module", params=test_devices()) +def collection_scene(request) -> Iterator[_CollectionScene]: + """Initialize the composite collection scene once per device.""" + device = request.param + with build_simulation_context(device=device, gravity_enabled=True) as sim: + sim._app_control_on_stop_handle = None + cubes, cube_origins = generate_cubes_scene( + num_envs=_NUM_ENVS, num_cubes=_NUM_CUBES, disable_gravity=True, root="/World/Cubes" + ) + kinematic, kinematic_origins = generate_cubes_scene( + num_envs=1, num_cubes=1, kinematic_enabled=True, root="/World/Kinematic", y_offset=12.0 + ) + 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, + ) - 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 - ) +def _rest_poses(scene: _CollectionScene, yaw: float = 0.0) -> torch.Tensor: + """Return the default body poses of the dynamic cubes at their environment origins with a given yaw.""" + poses = scene.cubes.data.default_body_pose.torch.clone() + poses[..., :2] += scene.origins["cubes"].unsqueeze(1)[..., :2] + poses[..., 3:] = torch.tensor(_yaw_quat(yaw), device=scene.device) + return poses + + +def test_collection_initialization(collection_scene: _CollectionScene) -> None: + """Initialize local collections, including a single-cube one; 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) -> None: + """Body state writes reach the selected body-major view entries in the frame they are given in.""" + scene = collection_scene + device = scene.device + collection = scene.cubes + scene.place_cubes_at_rest(_rest_poses(scene)) + + # A non-sorted partial write reaches only the selected view entries. + env_ids = torch.tensor([1, 0], dtype=torch.int32, device=device) + body_ids = torch.tensor([2], dtype=torch.int32, device=device) + initial_pose = collection.data.body_link_pose_w.torch.clone() + initial_velocity = collection.data.body_com_vel_w.torch.clone() + target_pose = initial_pose[env_ids][:, body_ids].clone() + target_pose[0, 0, :3] += torch.tensor([0.2, 0.3, 0.4], device=device) + target_pose[1, 0, :3] += torch.tensor([-0.1, -0.2, 0.1], device=device) + target_pose[..., 3:] = torch.tensor(_yaw_quat(0.3), device=device) + target_velocity = torch.tensor([[[0.1, 0.2, 0.3, 0.4, 0.5, 0.6]], [[-0.6, -0.5, -0.4, -0.3, -0.2, -0.1]]]) + target_velocity = target_velocity.to(device) + collection.write_body_link_pose_to_sim_index(body_poses=target_pose, env_ids=env_ids, body_ids=body_ids) + collection.write_body_com_velocity_to_sim_index(body_velocities=target_velocity, env_ids=env_ids, body_ids=body_ids) + expected_pose = initial_pose.clone() + expected_pose[env_ids[:, None], body_ids] = target_pose + expected_velocity = initial_velocity.clone() + expected_velocity[env_ids[:, None], body_ids] = target_velocity + torch.testing.assert_close(collection.data.body_link_pose_w.torch, expected_pose) + torch.testing.assert_close(collection.data.body_com_vel_w.torch, expected_velocity) + for raw, expected in ( + (collection.root_view.get_transforms(), expected_pose), + (collection.root_view.get_velocities(), expected_velocity), + ): + raw_by_env = wp.to_torch(raw).to(device).reshape(_NUM_CUBES, _NUM_ENVS, -1).transpose(0, 1) + torch.testing.assert_close(raw_by_env, expected) + + # With a center-of-mass offset, every writer round-trips through its own frame. + scene.place_cubes_at_rest(_rest_poses(scene)) + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(_NUM_ENVS, _NUM_CUBES, 1) + coms = collection.data.body_com_pose_b.torch.clone() + coms[..., :3] = offset + coms[..., 3:] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) + collection.set_coms_index(coms=coms) # random pose and velocity so frame conversions see a non-trivial rotation - rand_state = torch.rand(num_envs, num_cubes, 13, device=device) - rand_state[..., :3] += cube_object.data.body_link_pos_w.torch - # make quaternion a unit vector + 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) - - env_ids = env_ids.to(device) - object_ids = object_ids.to(device) - for i in range(10): - sim.step() - cube_object.update(sim.cfg.dt) - - ids = {} if i % 2 == 0 else {"env_ids": env_ids, "body_ids": object_ids} - if state_location == "com": - cube_object.write_body_com_pose_to_sim_index(body_poses=rand_state[..., :7], **ids) - cube_object.write_body_com_velocity_to_sim_index(body_velocities=rand_state[..., 7:], **ids) - elif state_location == "link": - cube_object.write_body_link_pose_to_sim_index(body_poses=rand_state[..., :7], **ids) - cube_object.write_body_link_velocity_to_sim_index(body_velocities=rand_state[..., 7:], **ids) - elif state_location == "root": - cube_object.write_body_link_pose_to_sim_index(body_poses=rand_state[..., :7], **ids) - cube_object.write_body_com_velocity_to_sim_index(body_velocities=rand_state[..., 7:], **ids) - - if state_location == "com": - torch.testing.assert_close(rand_state[..., :7], cube_object.data.body_com_pose_w.torch) - torch.testing.assert_close(rand_state[..., 7:], cube_object.data.body_com_vel_w.torch) - elif state_location == "link": - torch.testing.assert_close(rand_state[..., :7], cube_object.data.body_link_pose_w.torch) - torch.testing.assert_close(rand_state[..., 7:], cube_object.data.body_link_vel_w.torch) - elif state_location == "root": - torch.testing.assert_close(rand_state[..., :7], cube_object.data.body_link_pose_w.torch) - torch.testing.assert_close(rand_state[..., 7:], cube_object.data.body_com_vel_w.torch) - - # the frame that was not written must follow through the center-of-mass offset - link_pose_w = cube_object.data.body_link_pose_w.torch - body_com_pose_b = cube_object.data.body_com_pose_b.torch - expected_com_pos, expected_com_quat = combine_frame_transforms( - link_pose_w[..., :3].reshape(-1, 3), - link_pose_w[..., 3:].reshape(-1, 4), - body_com_pose_b[..., :3].reshape(-1, 3), - body_com_pose_b[..., 3:].reshape(-1, 4), - ) - expected_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1).view(num_envs, num_cubes, 7) - torch.testing.assert_close(expected_com_pose, cube_object.data.body_com_pose_w.torch) - torch.testing.assert_close( - cube_object.data.body_com_vel_w.torch[..., 3:], cube_object.data.body_link_vel_w.torch[..., 3:] + all_ids = {"env_ids": torch.arange(_NUM_ENVS, device=device), "body_ids": torch.arange(_NUM_CUBES, device=device)} + for state_location in ("com", "link", "root"): + for ids in ({}, all_ids): + if state_location == "com": + collection.write_body_com_pose_to_sim_index(body_poses=rand_state[..., :7], **ids) + collection.write_body_com_velocity_to_sim_index(body_velocities=rand_state[..., 7:], **ids) + torch.testing.assert_close(rand_state[..., :7], collection.data.body_com_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], collection.data.body_com_vel_w.torch) + elif state_location == "link": + collection.write_body_link_pose_to_sim_index(body_poses=rand_state[..., :7], **ids) + collection.write_body_link_velocity_to_sim_index(body_velocities=rand_state[..., 7:], **ids) + torch.testing.assert_close(rand_state[..., :7], collection.data.body_link_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], collection.data.body_link_vel_w.torch) + # PhysX holds the center-of-mass velocity: the link velocity plus the angular velocity crossed + # with the world-frame center-of-mass offset. + com_offset_w = quat_apply(rand_state[..., 3:7], offset) + expected_com_velocity = rand_state[..., 7:].clone() + expected_com_velocity[..., :3] += torch.linalg.cross(rand_state[..., 10:], com_offset_w, dim=-1) + raw_velocity = wp.to_torch(collection.root_view.get_velocities()).to(device) + raw_velocity = raw_velocity.reshape(_NUM_CUBES, _NUM_ENVS, 6).transpose(0, 1) + torch.testing.assert_close(raw_velocity, expected_com_velocity) + else: + collection.write_body_link_pose_to_sim_index(body_poses=rand_state[..., :7], **ids) + collection.write_body_com_velocity_to_sim_index(body_velocities=rand_state[..., 7:], **ids) + torch.testing.assert_close(rand_state[..., :7], collection.data.body_link_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], collection.data.body_com_vel_w.torch) + + # the frame that was not written must follow through the center-of-mass offset + link_pose_w = collection.data.body_link_pose_w.torch + body_com_pose_b = collection.data.body_com_pose_b.torch + expected_com_pos, expected_com_quat = combine_frame_transforms( + link_pose_w[..., :3].reshape(-1, 3), + link_pose_w[..., 3:].reshape(-1, 4), + body_com_pose_b[..., :3].reshape(-1, 3), + body_com_pose_b[..., 3:].reshape(-1, 4), + ) + expected_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1).view(_NUM_ENVS, _NUM_CUBES, 7) + torch.testing.assert_close(expected_com_pose, collection.data.body_com_pose_w.torch) + torch.testing.assert_close( + collection.data.body_com_vel_w.torch[..., 3:], collection.data.body_link_vel_w.torch[..., 3:] + ) + # Move away from the written state so that the next pass writes it again. + scene.step() + + # Spinning about the center of mass keeps it in place while the link frame orbits it. + scene.place_cubes_at_rest(_rest_poses(scene)) + spin_twist = torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0, 1.3], device=device).repeat(_NUM_ENVS, _NUM_CUBES, 1) + init_com = collection.data.body_com_pose_w.torch[..., :3].clone() + for _ in range(10): + collection.write_body_com_velocity_to_sim_index(body_velocities=spin_twist) + scene.step() + object_link_pose_w = collection.data.body_link_pose_w.torch + object_link_vel_w = collection.data.body_link_vel_w.torch + object_com_pose_w = collection.data.body_com_pose_w.torch + object_com_vel_w = collection.data.body_com_vel_w.torch + # center of mass position will be constant (i.e. spinning around com) + torch.testing.assert_close(init_com, object_com_pose_w[..., :3]) + # link position will be moving but should stay constant away from center of mass + object_link_state_pos_rel_com = quat_apply_inverse( + object_link_pose_w[..., 3:], object_link_pose_w[..., :3] - object_com_pose_w[..., :3] ) - - -@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) -> None: + """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) -> None: + """External wrenches act on the selected bodies in the frame they are given in; reset clears them.""" + scene = collection_scene + device = scene.device + collection = scene.cubes + rest_poses = _rest_poses(scene, yaw=0.5 * math.pi) + # Give every body a 1 kg mass and the center of mass and inertia of the first body so that their responses are + # comparable. + collection.set_masses_index(masses=torch.ones((_NUM_ENVS, _NUM_CUBES), device=device)) + coms = collection.data.body_com_pose_b.torch[:1, :1] + collection.set_coms_index(coms=coms.expand(_NUM_ENVS, _NUM_CUBES, 7).contiguous()) + inertias = collection.data.body_inertia.torch[:1, :1] + collection.set_inertias_index(inertias=inertias.expand(_NUM_ENVS, _NUM_CUBES, 9).contiguous()) + + # A body-frame force on bodies 0 and 2 accelerates only those bodies, along the rotated force direction; the + # same force given in the world frame produces the same response. + object_ids, _ = collection.find_bodies(".*") + for local_wrench, global_wrench, response, is_linear in ( + ( + {"forces": [[6.0, 0.0, 0.0]]}, + {"forces": [[0.0, 6.0, 0.0]]}, + lambda: collection.data.body_com_lin_vel_w, + True, + ), + ( + {"forces": [[0.0, 0.0, 6.0]], "positions": [[0.0, 0.1, 0.0]]}, + {"forces": [[0.0, 0.0, 6.0]], "positions": [[-0.1, 0.0, 0.0]]}, + lambda: collection.data.body_com_ang_vel_b, + False, + ), + ): + scene.place_cubes_at_rest(rest_poses) + for env_id, wrench, is_global in ((1, local_wrench, False), (0, global_wrench, True)): + wrench = {key: torch.tensor(value, device=device).expand(1, 2, 3).clone() for key, value in wrench.items()} + if is_global and "positions" in wrench: + # World-frame positions add the rotated body-frame lever arm to the center of mass. + wrench["positions"] += collection.data.body_com_pos_w.torch[env_id : env_id + 1, 0::2] + collection.permanent_wrench_composer.set_forces_and_torques_index( + body_ids=object_ids[0::2], env_ids=[env_id], is_global=is_global, **wrench + ) + scene.step() + response_value = response().torch + torch.testing.assert_close(response_value[0], response_value[1], atol=1e-4, rtol=1e-3) + torch.testing.assert_close(response_value[:, 1], torch.zeros_like(response_value[:, 1]), atol=1e-5, rtol=0) + if is_linear: + assert torch.all(response_value[:, 0::2, 1] > 1e-2), response_value + # The permanent force persists across steps, so a second step doubles the velocity. + first_step_value = response_value.clone() + scene.step() + torch.testing.assert_close(response().torch, 2.0 * first_step_value, atol=1e-4, rtol=1e-3) + # An upward force 0.1 m along the body y-axis rolls the cube about its x-axis. + assert torch.all(collection.data.body_com_ang_vel_b.torch[:, 0::2, 0] > 0.1) + heading = quat_apply( + collection.data.body_link_quat_w.torch, torch.tensor([1.0, 0.0, 0.0], device=device).expand(2, 3, 3) + ) + assert torch.all(heading[..., 1] > 0.99) + + # A partial reset clears the external wrenches of the selected environments only. + composers = (collection.instantaneous_wrench_composer, collection.permanent_wrench_composer) + ones = torch.ones((_NUM_ENVS, _NUM_CUBES, 3), device=device) + collection.permanent_wrench_composer.set_forces_and_torques_index(forces=ones, torques=ones) + collection.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones) + collection.reset(env_ids=torch.tensor([0], device=device)) for composer in composers: assert composer.active for buffer in (composer.out_force_b.torch, composer.out_torque_b.torch): @@ -611,40 +523,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)