Adds flag to recompute inertia when randomizing the mass of a rigid body (#989)

# Description

Previously, the method for randomizing masses of rigid bodies only set
the masses. However, the inertia tensors were not automatically updated,
and their original values were used inside the solver. This MR adds a
flag to recompute the inertia when mass is randomized by assuming a
uniform-density object. We make this an optional flag in case users want
to handle inertia tensors explicitly on their own.

## Type of change

- Bug fix (non-breaking change which fixes an issue)
- Breaking change (fix or feature that would cause existing
functionality to not work as expected)

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [x] I have made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [ ] 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

source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation.py

@@ -954,6 +954,7 @@ class Articulation(AssetBase):
 
         # -- bodies
         self._data.default_mass = self.root_physx_view.get_masses().clone()
+        self._data.default_inertia = self.root_physx_view.get_inertias().clone()
 
         # -- default joint state
         self._data.default_joint_pos = torch.zeros(self.num_instances, self.num_joints, device=self.device)

source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation_data.py

@@ -104,6 +104,13 @@ class ArticulationData:
     default_mass: torch.Tensor = None
     """Default mass read from the simulation. Shape is (num_instances, num_bodies)."""
 
+    default_inertia: torch.Tensor = None
+    """Default inertia read from the simulation. Shape is (num_instances, num_bodies, 9).
+
+    The inertia is the inertia tensor relative to the center of mass frame. The values are stored in
+    the order :math:`[I_{xx}, I_{xy}, I_{xz}, I_{yx}, I_{yy}, I_{yz}, I_{zx}, I_{zy}, I_{zz}]`.
+    """
+
     default_joint_pos: torch.Tensor = None
     """Default joint positions of all joints. Shape is (num_instances, num_joints)."""
 

source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object/rigid_object.py

@@ -327,6 +327,7 @@ class RigidObject(AssetBase):
         # set information about rigid body into data
         self._data.body_names = self.body_names
         self._data.default_mass = self.root_physx_view.get_masses().clone()
+        self._data.default_inertia = self.root_physx_view.get_inertias().clone()
 
     def _process_cfg(self):
         """Post processing of configuration parameters."""

source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object/rigid_object_data.py

@@ -96,6 +96,13 @@ class RigidObjectData:
     default_mass: torch.Tensor = None
     """Default mass read from the simulation. Shape is (num_instances, 1)."""
 
+    default_inertia: torch.Tensor = None
+    """Default inertia tensor read from the simulation. Shape is (num_instances, 9).
+
+    The inertia is the inertia tensor relative to the center of mass frame. The values are stored in
+    the order :math:`[I_{xx}, I_{xy}, I_{xz}, I_{yx}, I_{yy}, I_{yz}, I_{zx}, I_{zy}, I_{zz}]`.
+    """
+
     ##
     # Properties.
     ##

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/events.py

@@ -132,12 +132,18 @@ def randomize_rigid_body_mass(
     mass_distribution_params: tuple[float, float],
     operation: Literal["add", "scale", "abs"],
     distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform",
+    recompute_inertia: bool = True,
 ):
     """Randomize the mass of the bodies by adding, scaling, or setting random values.
 
     This function allows randomizing the mass of the bodies of the asset. The function samples random values from the
     given distribution parameters and adds, scales, or sets the values into the physics simulation based on the operation.
 
+    If the ``recompute_inertia`` flag is set to ``True``, the function recomputes the inertia tensor of the bodies
+    after setting the mass. This is useful when the mass is changed significantly, as the inertia tensor depends
+    on the mass. It assumes the body is a uniform density object. If the body is not a uniform density object,
+    the inertia tensor may not be accurate.
+
     .. tip::
         This function uses CPU tensors to assign the body masses. It is recommended to use this function
         only during the initialization of the environment.
@@ -159,7 +165,10 @@ def randomize_rigid_body_mass(
 
     # get the current masses of the bodies (num_assets, num_bodies)
     masses = asset.root_physx_view.get_masses()
+
     # apply randomization on default values
+    # this is to make sure when calling the function multiple times, the randomization is applied on the
+    # default values and not the previously randomized values
     masses[env_ids[:, None], body_ids] = asset.data.default_mass[env_ids[:, None], body_ids].clone()
 
     # sample from the given range
@@ -172,6 +181,24 @@ def randomize_rigid_body_mass(
     # set the mass into the physics simulation
     asset.root_physx_view.set_masses(masses, env_ids)
 
+    # recompute inertia tensors if needed
+    if recompute_inertia:
+        # compute the ratios of the new masses to the initial masses
+        ratios = masses[env_ids[:, None], body_ids] / asset.data.default_mass[env_ids[:, None], body_ids]
+        # scale the inertia tensors by the the ratios
+        # since mass randomization is done on default values, we can use the default inertia tensors
+        inertias = asset.root_physx_view.get_inertias()
+        if isinstance(asset, Articulation):
+            # inertia has shape: (num_envs, num_bodies, 9) for articulation
+            inertias[env_ids[:, None], body_ids] = (
+                asset.data.default_inertia[env_ids[:, None], body_ids] * ratios[..., None]
+            )
+        else:
+            # inertia has shape: (num_envs, 9) for rigid object
+            inertias[env_ids] = asset.data.default_inertia[env_ids] * ratios
+        # set the inertia tensors into the physics simulation
+        asset.root_physx_view.set_inertias(inertias, env_ids)
+
 
 def randomize_physics_scene_gravity(
     env: ManagerBasedEnv,

source/extensions/omni.isaac.lab/test/assets/test_articulation.py

@@ -199,6 +199,12 @@ class TestArticulation(unittest.TestCase):
                         self.assertEqual(articulation.data.root_pos_w.shape, (num_articulations, 3))
                         self.assertEqual(articulation.data.root_quat_w.shape, (num_articulations, 4))
                         self.assertEqual(articulation.data.joint_pos.shape, (num_articulations, 12))
+                        self.assertEqual(
+                            articulation.data.default_mass.shape, (num_articulations, articulation.num_bodies)
+                        )
+                        self.assertEqual(
+                            articulation.data.default_inertia.shape, (num_articulations, articulation.num_bodies, 9)
+                        )
 
                         # Check some internal physx data for debugging
                         # -- joint related
@@ -246,6 +252,12 @@ class TestArticulation(unittest.TestCase):
                         self.assertEqual(articulation.data.root_pos_w.shape, (num_articulations, 3))
                         self.assertEqual(articulation.data.root_quat_w.shape, (num_articulations, 4))
                         self.assertEqual(articulation.data.joint_pos.shape, (num_articulations, 9))
+                        self.assertEqual(
+                            articulation.data.default_mass.shape, (num_articulations, articulation.num_bodies)
+                        )
+                        self.assertEqual(
+                            articulation.data.default_inertia.shape, (num_articulations, articulation.num_bodies, 9)
+                        )
 
                         # Check some internal physx data for debugging
                         # -- joint related
@@ -299,6 +311,12 @@ class TestArticulation(unittest.TestCase):
                         self.assertEqual(articulation.data.root_pos_w.shape, (num_articulations, 3))
                         self.assertEqual(articulation.data.root_quat_w.shape, (num_articulations, 4))
                         self.assertEqual(articulation.data.joint_pos.shape, (num_articulations, 1))
+                        self.assertEqual(
+                            articulation.data.default_mass.shape, (num_articulations, articulation.num_bodies)
+                        )
+                        self.assertEqual(
+                            articulation.data.default_inertia.shape, (num_articulations, articulation.num_bodies, 9)
+                        )
 
                         # Check some internal physx data for debugging
                         # -- joint related
@@ -352,6 +370,12 @@ class TestArticulation(unittest.TestCase):
                         self.assertTrue(articulation.data.root_pos_w.shape == (num_articulations, 3))
                         self.assertTrue(articulation.data.root_quat_w.shape == (num_articulations, 4))
                         self.assertTrue(articulation.data.joint_pos.shape == (num_articulations, 24))
+                        self.assertEqual(
+                            articulation.data.default_mass.shape, (num_articulations, articulation.num_bodies)
+                        )
+                        self.assertEqual(
+                            articulation.data.default_inertia.shape, (num_articulations, articulation.num_bodies, 9)
+                        )
 
                         # Check some internal physx data for debugging
                         # -- joint related

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

@@ -108,6 +108,8 @@ class TestRigidObject(unittest.TestCase):
                         # Check buffers that exists and have correct shapes
                         self.assertEqual(cube_object.data.root_pos_w.shape, (num_cubes, 3))
                         self.assertEqual(cube_object.data.root_quat_w.shape, (num_cubes, 4))
+                        self.assertEqual(cube_object.data.default_mass.shape, (num_cubes, 1))
+                        self.assertEqual(cube_object.data.default_inertia.shape, (num_cubes, 9))
 
                         # Simulate physics
                         for _ in range(2):