Expands on unit test for rigid object asset class (#379)

# Description
We want to expand our test coverage, so this PR aims to increase
coverage of rigid object.

This PR ensures the following test cases are covered by unit tests for
rigid_object.py:
* external forces are applied correctly
* friction/mass/collider offsets randomization works correctly
* setting states work correctly
* reset works correctly

All test cases pass 👍 

Note: I wasn't able to find `collider offsets` in the code base but
happy to add a test for this too if someone can point me in the right
direction.

Fixes #363 

## Type of change

- New feature (non-breaking change which adds functionality)

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./orbit.sh --format`
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] I have run all the tests with `./orbit.sh --test` and they pass
- [ ] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

<!--
As you go through the checklist above, you can mark something as done by
putting an x character in it

For example,
- [x] I have done this task
- [ ] I have not done this task
-->

---------
Signed-off-by: jsmith-bdai <142246516+jsmith-bdai@users.noreply.github.com>
Signed-off-by: James Smith <142246516+jsmith-bdai@users.noreply.github.com>
Co-authored-by: AutonomousHansen <50837800+AutonomousHansen@users.noreply.github.com>
Co-authored-by: Mayank Mittal <12863862+Mayankm96@users.noreply.github.com>

source/extensions/omni.isaac.orbit/test/assets/test_rigid_object.py

@@ -12,8 +12,12 @@ from __future__ import annotations
 
 from omni.isaac.orbit.app import AppLauncher
 
+# Can set this to False to see the GUI for debugging
+# This will also add lights to the scene
+HEADLESS = True
+
 # launch omniverse app
-app_launcher = AppLauncher(headless=True)
+app_launcher = AppLauncher(headless=HEADLESS)
 simulation_app = app_launcher.app
 
 """Rest everything follows."""
@@ -23,120 +27,538 @@ import torch
 import unittest
 
 import omni.isaac.core.utils.prims as prim_utils
-import omni.isaac.core.utils.stage as stage_utils
 
 import omni.isaac.orbit.sim as sim_utils
 from omni.isaac.orbit.assets import RigidObject, RigidObjectCfg
+from omni.isaac.orbit.sim import build_simulation_context
+from omni.isaac.orbit.sim.spawners import materials
 from omni.isaac.orbit.utils.assets import ISAAC_NUCLEUS_DIR
+from omni.isaac.orbit.utils.math import default_orientation, random_orientation
 
 
-class TestRigidObject(unittest.TestCase):
-    """Test for rigid object class."""
+def generate_cubes_scene(num_cubes: int = 1, height=1.0) -> tuple[RigidObject, torch.Tensor]:
+    """Generate a scene with the provided number of cubes.
 
-    def setUp(self):
-        """Create a blank new stage for each test."""
-        # Create a new stage
-        stage_utils.create_new_stage()
-        # Simulation time-step
-        self.dt = 0.01
-        # Load kit helper
-        sim_cfg = sim_utils.SimulationCfg(dt=self.dt, device="cuda:0")
-        self.sim = sim_utils.SimulationContext(sim_cfg)
-
-    def tearDown(self):
-        """Stops simulator after each test."""
-        # stop simulation
-        self.sim.stop()
-        # clear the stage
-        self.sim.clear_instance()
+    Args:
+        num_cubes: Number of cubes to generate.
+        height: Height of the cubes.
+
+    Returns:
+        RigidObject: The rigid object representing the cubes.
+        origins: The origins of the cubes.
 
     """
-    Tests
-    """
+    origins = torch.tensor([(i * 1.0, 0, height) for i in range(num_cubes)])
+    # Create Top-level Xforms, one for each cube
+    for i, origin in enumerate(origins):
+        prim_utils.create_prim(f"/World/Table_{i}", "Xform", translation=origin)
 
-    def test_initialization(self):
-        """Test initialization for with rigid body API at the provided prim path."""
     # Create rigid object
     cube_object_cfg = RigidObjectCfg(
-            prim_path="/World/Object",
+        prim_path="/World/Table_.*/Object",
         spawn=sim_utils.UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd"),
-            init_state=RigidObjectCfg.InitialStateCfg(),
+        init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, height)),
     )
     cube_object = RigidObject(cfg=cube_object_cfg)
 
-        # Check that boundedness of articulation is correct
+    return cube_object, origins
+
+
+class TestRigidObject(unittest.TestCase):
+    """Test for rigid object class."""
+
+    """
+    Tests
+    """
+
+    def test_initialization(self):
+        """Test initialization for with rigid body API at the provided prim path."""
+        for num_cubes in (1, 2):
+            for device in ("cuda:0", "cpu"):
+                with self.subTest(num_cubes=num_cubes, device=device):
+                    with build_simulation_context(device=device, auto_add_lighting=True) as sim:
+                        cube_object, _ = generate_cubes_scene(num_cubes=num_cubes)
+
+                        # Check that boundedness of rigid object is correct
                         self.assertEqual(ctypes.c_long.from_address(id(cube_object)).value, 1)
 
                         # Play sim
-        self.sim.reset()
+                        sim.reset()
+
                         # Check if object is initialized
                         self.assertTrue(cube_object._is_initialized)
                         self.assertEqual(len(cube_object.body_names), 1)
+
                         # Check buffers that exists and have correct shapes
-        self.assertTrue(cube_object.data.root_pos_w.shape == (1, 3))
-        self.assertTrue(cube_object.data.root_quat_w.shape == (1, 4))
+                        self.assertEqual(cube_object.data.root_pos_w.shape, (num_cubes, 3))
+                        self.assertEqual(cube_object.data.root_quat_w.shape, (num_cubes, 4))
 
                         # Simulate physics
-        for _ in range(20):
+                        for _ in range(2):
                             # perform rendering
-            self.sim.step()
+                            sim.step()
                             # update object
-            cube_object.update(self.dt)
+                            cube_object.update(sim.cfg.dt)
 
     def test_external_force_on_single_body(self):
         """Test application of external force on the base of the object.
 
-        In this test, we apply a force equal to the weight of the object on the base of
+        In this test, we apply a force equal to the weight of an object on the base of
         one of the objects. We check that the object does not move. For the other object,
         we do not apply any force and check that it falls down.
         """
-        # Create parent prims
-        prim_utils.create_prim("/World/Table_1", "Xform", translation=(0.0, -1.0, 0.0))
-        prim_utils.create_prim("/World/Table_2", "Xform", translation=(0.0, 1.0, 0.0))
-
-        # Create rigid object
-        cube_object_cfg = RigidObjectCfg(
-            prim_path="/World/Table_.*/Object",
-            spawn=sim_utils.UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd"),
-            init_state=RigidObjectCfg.InitialStateCfg(),
-        )
-        # create handles for the objects
-        cube_object = RigidObject(cfg=cube_object_cfg)
+        for num_cubes in (2, 4):
+            for device in ("cuda:0", "cpu"):
+                with self.subTest(num_cubes=num_cubes, device=device):
+                    with build_simulation_context(device=device, add_ground_plane=True, auto_add_lighting=True) as sim:
+                        cube_object, origins = generate_cubes_scene(num_cubes=num_cubes)
 
                         # Play the simulator
-        self.sim.reset()
+                        sim.reset()
 
                         # Find bodies to apply the force
-        body_ids, _ = cube_object.find_bodies(".*")
+                        body_ids, body_names = cube_object.find_bodies(".*")
 
-        # Sample a large force
-        external_wrench_b = torch.zeros(cube_object.num_instances, len(body_ids), 6, device=self.sim.device)
-        external_wrench_b[0, 0, 2] = 9.81 * cube_object.root_physx_view.get_masses()[0]
+                        # Sample a force equal to the weight of the object
+                        external_wrench_b = torch.zeros(cube_object.num_instances, len(body_ids), 6, device=sim.device)
+                        # Every 2nd cube should have a force applied to it
+                        external_wrench_b[0::2, :, 2] = 9.81 * cube_object.root_physx_view.get_masses()[0]
 
                         # Now we are ready!
                         for _ in range(5):
                             # reset root state
                             root_state = cube_object.data.default_root_state.clone()
-            root_state[0, :2] = torch.tensor([0.0, -1.0], device=self.sim.device)
-            root_state[1, :2] = torch.tensor([0.0, 1.0], device=self.sim.device)
+
+                            # need to shift the position of the cubes otherwise they will be on top of each other
+                            root_state[:, :3] = origins
                             cube_object.write_root_state_to_sim(root_state)
+
                             # reset object
                             cube_object.reset()
+
                             # apply force
                             cube_object.set_external_force_and_torque(
                                 external_wrench_b[..., :3], external_wrench_b[..., 3:], body_ids=body_ids
                             )
                             # perform simulation
-            for _ in range(50):
+                            for _ in range(5):
                                 # apply action to the object
                                 cube_object.write_data_to_sim()
+
                                 # perform step
-                self.sim.step()
+                                sim.step()
+
                                 # update buffers
-                cube_object.update(self.dt)
-            # check condition that the objects have fallen down
-            self.assertLess(abs(cube_object.data.root_pos_w[0, 2].item()), 1e-3)
-            self.assertLess(cube_object.data.root_pos_w[1, 2].item(), -1.0)
+                                cube_object.update(sim.cfg.dt)
+
+                            # First object should still be at the same Z position (1.0)
+                            torch.testing.assert_close(
+                                cube_object.data.root_pos_w[0::2, 2], torch.ones(num_cubes // 2, device=sim.device)
+                            )
+                            # Second object should have fallen, so it's Z height should be less than initial height of 1.0
+                            self.assertTrue(torch.all(cube_object.data.root_pos_w[1::2, 2] < 1.0))
+
+    def test_set_rigid_object_state(self):
+        """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.
+        """
+        for num_cubes in (1, 2):
+            for device in ("cuda:0", "cpu"):
+                with self.subTest(num_cubes=num_cubes, device=device):
+                    # 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:
+                        cube_object, _ = generate_cubes_scene(num_cubes=num_cubes)
+
+                        # 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, 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, device=sim.device),
+                                "root_ang_vel_w": torch.zeros_like(cube_object.data.root_ang_vel_w, 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_state = torch.cat(
+                                        [
+                                            state_dict["root_pos_w"],
+                                            state_dict["root_quat_w"],
+                                            state_dict["root_lin_vel_w"],
+                                            state_dict["root_ang_vel_w"],
+                                        ],
+                                        dim=-1,
+                                    )
+                                    # reset root state
+                                    cube_object.write_root_state_to_sim(root_state=root_state)
+
+                                    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.testing.assert_allclose(value, expected_value, rtol=1e-5, atol=1e-5)
+
+                                    cube_object.update(sim.cfg.dt)
+
+    def test_reset_rigid_object(self):
+        """Test resetting the state of the rigid object."""
+        for num_cubes in (1, 2):
+            for device in ("cuda:0", "cpu"):
+                with self.subTest(num_cubes=num_cubes, device=device):
+                    with build_simulation_context(device=device, gravity_enabled=True, auto_add_lighting=True) as sim:
+                        cube_object, _ = generate_cubes_scene(num_cubes=num_cubes)
+
+                        # Play the simulator
+                        sim.reset()
+
+                        for i in range(5):
+                            # perform rendering
+                            sim.step()
+
+                            # update object
+                            cube_object.update(sim.cfg.dt)
+
+                            # Move the object to a random position
+                            root_state = cube_object.data.default_root_state.clone()
+                            root_state[:, :3] = torch.randn(num_cubes, 3, device=sim.device)
+
+                            # Random orientation
+                            root_state[:, 3:7] = random_orientation(num=num_cubes, device=sim.device)
+                            cube_object.write_root_state_to_sim(root_state)
+
+                            if i % 2 == 0:
+                                # reset object
+                                cube_object.reset()
+
+                                # Reset should zero external forces and torques and set last body velocity to zero
+                                self.assertFalse(cube_object.has_external_wrench)
+                                self.assertEqual(torch.count_nonzero(cube_object._external_force_b), 0)
+                                self.assertEqual(torch.count_nonzero(cube_object._external_torque_b), 0)
+                                self.assertEqual(torch.count_nonzero(cube_object._last_body_vel_w), 0)
+
+    def test_rigid_body_set_material_properties(self):
+        """Test getting and setting material properties of rigid object."""
+        for num_cubes in (1, 2):
+            for device in ("cuda:0", "cpu"):
+                with self.subTest(num_cubes=num_cubes, device=device):
+                    with build_simulation_context(
+                        device=device, gravity_enabled=True, add_ground_plane=True, auto_add_lighting=True
+                    ) as sim:
+                        # Create rigid object(s)
+                        cube_object, _ = generate_cubes_scene(num_cubes=num_cubes)
+
+                        # Play sim
+                        sim.reset()
+
+                        # Set material properties
+                        static_friction = torch.FloatTensor(num_cubes, 1).uniform_(0.4, 0.8)
+                        dynamic_friction = torch.FloatTensor(num_cubes, 1).uniform_(0.4, 0.8)
+                        restitution = torch.FloatTensor(num_cubes, 1).uniform_(0.0, 0.2)
+
+                        materials = torch.cat([static_friction, dynamic_friction, restitution], dim=-1)
+
+                        indices = torch.tensor(range(num_cubes), dtype=torch.int)
+                        # Add friction to cube
+                        cube_object.body_physx_view.set_material_properties(materials, indices)
+
+                        # Simulate physics
+                        # perform rendering
+                        sim.step()
+                        # update object
+                        cube_object.update(sim.cfg.dt)
+
+                        # Get material properties
+                        materials_to_check = cube_object.body_physx_view.get_material_properties()
+
+                        # Check if material properties are set correctly
+                        torch.testing.assert_close(materials_to_check.reshape(num_cubes, 3), materials)
+
+    def test_rigid_body_no_friction(self):
+        """Test that a rigid object with no friction will maintain it's velocity when sliding across a plane."""
+        for num_cubes in (1, 2):
+            for device in ("cuda:0", "cpu"):
+                with self.subTest(num_cubes=num_cubes, device=device):
+                    with build_simulation_context(device=device, auto_add_lighting=True) as sim:
+                        cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=0.0)
+
+                        # 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.int)
+
+                        cube_object.body_physx_view.set_material_properties(cube_object_materials, indices)
+
+                        # 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(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 == "cuda:0":
+                                tolerance = 1e-2
+                            else:
+                                tolerance = 1e-5
+
+                            torch.testing.assert_close(
+                                cube_object.data.root_lin_vel_w, initial_velocity[:, :3], rtol=1e-5, atol=tolerance
+                            )
+
+    def test_rigid_body_with_static_friction(self):
+        """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.
+        """
+        for num_cubes in (1, 2):
+            for device in ("cuda:0", "cpu"):
+                with self.subTest(num_cubes=num_cubes, device=device):
+                    with build_simulation_context(device=device, add_ground_plane=True, auto_add_lighting=True) as sim:
+                        cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=0.03125)
+
+                        # Create ground plane with no friction
+                        cfg = sim_utils.GroundPlaneCfg(
+                            physics_material=materials.RigidBodyMaterialCfg(
+                                static_friction=0.0,
+                                dynamic_friction=0.0,
+                            )
+                        )
+                        cfg.func("/World/GroundPlane", cfg)
+
+                        # Play sim
+                        sim.reset()
+
+                        # Set static friction to be non-zero
+                        static_friction_coefficient = 0.5
+                        static_friction = torch.Tensor([[static_friction_coefficient]] * num_cubes)
+                        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.int)
+
+                        # Add friction to cube
+                        cube_object.body_physx_view.set_material_properties(cube_object_materials, indices)
+
+                        # 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":
+                            with self.subTest(force=force):
+                                external_wrench_b = torch.zeros((num_cubes, 1, 6), device=sim.device)
+
+                                if force == "below_mu":
+                                    external_wrench_b[:, 0, 0] = static_friction_coefficient * 0.999
+                                else:
+                                    external_wrench_b[:, 0, 0] = static_friction_coefficient * 1.001
+
+                                cube_object.set_external_force_and_torque(
+                                    external_wrench_b[..., :3],
+                                    external_wrench_b[..., 3:],
+                                )
+
+                                # Get root state
+                                initial_root_state = cube_object.data.root_state_w
+
+                                # Simulate physics
+                                for _ in range(10):
+                                    # perform rendering
+                                    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_state_w, initial_root_state, rtol=1e-5, atol=1e-5
+                                        )
+                                    else:
+                                        torch.testing.assert_close(
+                                            cube_object.data.root_state_w, initial_root_state, rtol=1e-5, atol=1e-5
+                                        )
+
+    def test_rigid_body_with_restitution(self):
+        """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 1, the block should bounce
+        with the same energy. When the restitution is between 0 and 1, the block should bounce with less energy.
+
+        """
+        for num_cubes in (1, 2):
+            for device in ("cuda:0", "cpu"):
+                with self.subTest(num_cubes=num_cubes, device=device):
+                    with build_simulation_context(device=device, add_ground_plane=True, auto_add_lighting=True) as sim:
+                        cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=1.0)
+
+                        # Create ground plane such that has a restitution of 1.0 (perfectly elastic collision)
+                        cfg = sim_utils.GroundPlaneCfg(
+                            physics_material=materials.RigidBodyMaterialCfg(
+                                restitution=1.0,
+                            )
+                        )
+                        cfg.func("/World/GroundPlane", cfg)
+
+                        indices = torch.tensor(range(num_cubes), dtype=torch.int)
+
+                        # Play sim
+                        sim.reset()
+
+                        # 3 cases: inelastic, partially elastic, elastic
+                        # inelastic: resitution = 0, block should not bounce
+                        # partially elastic: 0 <= restitution <= 1, block should bounce with less energy
+                        # elastic: restitution = 1, block should bounce with same energy
+                        for expected_collision_type in "inelastic", "partially_elastic", "elastic":
+                            root_state = torch.zeros(1, 13, device=sim.device)
+                            root_state[0, 3] = 1.0  # To make orientation a quaternion
+                            root_state[0, 2] = 0.1  # Set an initial drop height
+                            root_state[0, 9] = -1.0  # Set an initial downward velocity
+
+                            cube_object.write_root_state_to_sim(root_state=root_state)
+
+                            prev_z_velocity = 0.0
+                            curr_z_velocity = 0.0
+
+                            with self.subTest(expected_collision_type=expected_collision_type):
+                                # cube_object.reset()
+                                # 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
+                                else:
+                                    restitution_coefficient = 1.0
+
+                                restitution = 0.5
+                                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 friction to cube
+                                cube_object.body_physx_view.set_material_properties(cube_object_materials, indices)
+
+                                curr_z_velocity = cube_object.data.root_lin_vel_w[:, 2]
+
+                                while torch.all(curr_z_velocity <= 0.0):
+                                    # Simulate physics
+                                    curr_z_velocity = cube_object.data.root_lin_vel_w[:, 2]
+
+                                    # perform rendering
+                                    sim.step()
+
+                                    # update object
+                                    cube_object.update(sim.cfg.dt)
+                                    if torch.all(curr_z_velocity <= 0.0):
+                                        # Still in the air
+                                        prev_z_velocity = curr_z_velocity
+
+                                # We have made contact with the ground and can verify expected collision type
+                                # based on how velocity has changed after the collision
+                                if expected_collision_type == "inelastic":
+                                    # Assert that the block has lost most energy by checking that the z velocity is < 1/2 previous
+                                    # velocity. This is because the floor's resitution means it will bounce back an object that itself
+                                    # has restitution set to 0.0
+                                    self.assertTrue(torch.all(torch.le(curr_z_velocity / 2, abs(prev_z_velocity))))
+                                elif expected_collision_type == "partially_elastic":
+                                    # Assert that the block has lost some energy by checking that the z velocity is less
+                                    self.assertTrue(torch.all(torch.le(abs(curr_z_velocity), abs(prev_z_velocity))))
+                                elif expected_collision_type == "elastic":
+                                    # Assert that the block has not lost any energy by checking that the z velocity is the same
+                                    torch.testing.assert_close(abs(curr_z_velocity), abs(prev_z_velocity))
+
+    def test_rigid_body_set_mass(self):
+        """Test getting and setting mass of rigid object."""
+        for num_cubes in (1, 2):
+            for device in ("cuda:0", "cpu"):
+                with self.subTest(num_cubes=num_cubes, device=device):
+                    with build_simulation_context(
+                        device=device, gravity_enabled=False, add_ground_plane=True, auto_add_lighting=True
+                    ) as sim:
+                        cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=1.0)
+
+                        # Play sim
+                        sim.reset()
+
+                        # Get masses before increasing
+                        original_masses = cube_object.body_physx_view.get_masses()
+
+                        self.assertEqual(original_masses.shape, (num_cubes, 1))
+
+                        # Randomize mass of the object
+                        masses = original_masses + torch.FloatTensor(num_cubes, 1).uniform_(4, 8)
+
+                        indices = torch.tensor(range(num_cubes), dtype=torch.int)
+
+                        # Add friction to cube
+                        cube_object.body_physx_view.set_masses(masses, indices)
+
+                        torch.testing.assert_close(cube_object.body_physx_view.get_masses(), masses)
+
+                        # Simulate physics
+                        # perform rendering
+                        sim.step()
+                        # update object
+                        cube_object.update(sim.cfg.dt)
+
+                        masses_to_check = cube_object.body_physx_view.get_masses()
+
+                        # Check if mass is set correctly
+                        torch.testing.assert_close(masses, masses_to_check)
 
 
 if __name__ == "__main__":