Fixes the event for randomizing rigid body material (#1140)

# Description

The current friction randomization event only selects a single random
number in the given range and does not vary them. With the given PR,
this is getting fixed, and there is a sampling of the entire given
range.

## Type of change

- Bug fix (non-breaking change which fixes an issue)

## Screenshots

| Before | After |
| ------ | ----- |
| ![Screenshot from 2024-10-03
22-26-49](https://github.com/user-attachments/assets/d13f86ee-c776-4046-af2e-46be8f271a00)
| ![Screenshot from 2024-10-03
22-27-15](https://github.com/user-attachments/assets/cf0a536d-20d0-47f1-b580-25241049cdd4)
|


## 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
- [x] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [x] 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

---------

Co-authored-by: Mayank Mittal <[email protected]>

source/extensions/omni.isaac.lab/config/extension.toml

@@ -1,7 +1,7 @@
 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.24.19"
+version = "0.24.20"
 
 # Description
 title = "Isaac Lab framework for Robot Learning"

source/extensions/omni.isaac.lab/docs/CHANGELOG.rst

@@ -1,6 +1,16 @@
 Changelog
 ---------
 
+0.24.20 (2024-10-07)
+~~~~~~~~~~~~~~~~~~~~
+
+Fixes
+^^^^^
+
+* Fixed the :meth:`omni.isaac.lab.envs.mdp.events.randomize_rigid_body_material` function to
+  correctly sample friction and restitution from the given ranges.
+
+
 0.24.19 (2024-10-05)
 ~~~~~~~~~~~~~~~~~~~~
 

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

@@ -14,7 +14,6 @@
 
 from __future__ import annotations
 
-import numpy as np
 import torch
 from typing import TYPE_CHECKING, Literal
 
@@ -25,22 +24,14 @@
 import omni.isaac.lab.utils.math as math_utils
 from omni.isaac.lab.actuators import ImplicitActuator
 from omni.isaac.lab.assets import Articulation, DeformableObject, RigidObject
-from omni.isaac.lab.managers import SceneEntityCfg
+from omni.isaac.lab.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg
 from omni.isaac.lab.terrains import TerrainImporter
 
 if TYPE_CHECKING:
     from omni.isaac.lab.envs import ManagerBasedEnv
 
 
-def randomize_rigid_body_material(
-    env: ManagerBasedEnv,
-    env_ids: torch.Tensor | None,
-    static_friction_range: tuple[float, float],
-    dynamic_friction_range: tuple[float, float],
-    restitution_range: tuple[float, float],
-    num_buckets: int,
-    asset_cfg: SceneEntityCfg,
-):
+class randomize_rigid_body_material(ManagerTermBase):
     """Randomize the physics materials on all geometries of the asset.
 
     This function creates a set of physics materials with random static friction, dynamic friction, and restitution
@@ -53,76 +44,122 @@ def randomize_rigid_body_material(
     all bodies). The integer values are used as indices to select the material properties from the
     material buckets.
 
+    If the flag ``make_consistent`` is set to ``True``, the dynamic friction is set to be less than or equal to
+    the static friction. This obeys the physics constraint on friction values. However, it may not always be
+    essential for the application. Thus, the flag is set to ``False`` by default.
+
     .. attention::
         This function uses CPU tensors to assign the material properties. It is recommended to use this function
         only during the initialization of the environment. Otherwise, it may lead to a significant performance
         overhead.
 
     .. note::
         PhysX only allows 64000 unique physics materials in the scene. If the number of materials exceeds this
-        limit, the simulation will crash.
+        limit, the simulation will crash. Due to this reason, we sample the materials only once during initialization.
+        Afterwards, these materials are randomly assigned to the geometries of the asset.
     """
-    # extract the used quantities (to enable type-hinting)
-    asset: RigidObject | Articulation = env.scene[asset_cfg.name]
 
-    if not isinstance(asset, (RigidObject, Articulation)):
-        raise ValueError(
-            f"Randomization term 'randomize_rigid_body_material' not supported for asset: '{asset_cfg.name}'"
-            f" with type: '{type(asset)}'."
-        )
+    def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv):
+        """Initialize the term.
 
-    # resolve environment ids
-    if env_ids is None:
-        env_ids = torch.arange(env.scene.num_envs, device="cpu")
-    else:
-        env_ids = env_ids.cpu()
-
-    # retrieve material buffer
-    materials = asset.root_physx_view.get_material_properties()
+        Args:
+            cfg: The configuration of the event term.
+            env: The environment instance.
 
-    # sample material properties from the given ranges
-    material_samples = np.zeros(materials[env_ids].shape)
-    material_samples[..., 0] = np.random.uniform(*static_friction_range)
-    material_samples[..., 1] = np.random.uniform(*dynamic_friction_range)
-    material_samples[..., 2] = np.random.uniform(*restitution_range)
+        Raises:
+            ValueError: If the asset is not a RigidObject or an Articulation.
+        """
+        super().__init__(cfg, env)
 
-    # create uniform range tensor for bucketing
-    lo = np.array([static_friction_range[0], dynamic_friction_range[0], restitution_range[0]])
-    hi = np.array([static_friction_range[1], dynamic_friction_range[1], restitution_range[1]])
+        # extract the used quantities (to enable type-hinting)
+        self.asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"]
+        self.asset: RigidObject | Articulation = env.scene[self.asset_cfg.name]
 
-    # to avoid 64k material limit in physx, we bucket materials by binning randomized material properties
-    # into buckets based on the number of buckets specified
-    for d in range(3):
-        buckets = np.array([(hi[d] - lo[d]) * i / num_buckets + lo[d] for i in range(num_buckets)])
-        material_samples[..., d] = buckets[np.searchsorted(buckets, material_samples[..., d]) - 1]
+        if not isinstance(self.asset, (RigidObject, Articulation)):
+            raise ValueError(
+                f"Randomization term 'randomize_rigid_body_material' not supported for asset: '{self.asset_cfg.name}'"
+                f" with type: '{type(self.asset)}'."
+            )
 
-    # update material buffer with new samples
-    if isinstance(asset, Articulation) and asset_cfg.body_ids != slice(None):
         # obtain number of shapes per body (needed for indexing the material properties correctly)
         # note: this is a workaround since the Articulation does not provide a direct way to obtain the number of shapes
         #  per body. We use the physics simulation view to obtain the number of shapes per body.
-        num_shapes_per_body = []
-        for link_path in asset.root_physx_view.link_paths[0]:
-            link_physx_view = asset._physics_sim_view.create_rigid_body_view(link_path)  # type: ignore
-            num_shapes_per_body.append(link_physx_view.max_shapes)
+        if isinstance(self.asset, Articulation) and self.asset_cfg.body_ids != slice(None):
+            self.num_shapes_per_body = []
+            for link_path in self.asset.root_physx_view.link_paths[0]:
+                link_physx_view = self.asset._physics_sim_view.create_rigid_body_view(link_path)  # type: ignore
+                self.num_shapes_per_body.append(link_physx_view.max_shapes)
+            # ensure the parsing is correct
+            num_shapes = sum(self.num_shapes_per_body)
+            expected_shapes = self.asset.root_physx_view.max_shapes
+            if num_shapes != expected_shapes:
+                raise ValueError(
+                    "Randomization term 'randomize_rigid_body_material' failed to parse the number of shapes per body."
+                    f" Expected total shapes: {expected_shapes}, but got: {num_shapes}."
+                )
+        else:
+            # in this case, we don't need to do special indexing
+            self.num_shapes_per_body = None
+
+        # obtain parameters for sampling friction and restitution values
+        static_friction_range = cfg.params.get("static_friction_range", (1.0, 1.0))
+        dynamic_friction_range = cfg.params.get("dynamic_friction_range", (1.0, 1.0))
+        restitution_range = cfg.params.get("restitution_range", (0.0, 0.0))
+        num_buckets = int(cfg.params.get("num_buckets", 1))
 
         # sample material properties from the given ranges
-        for body_id in asset_cfg.body_ids:
-            # start index of shape
-            start_idx = sum(num_shapes_per_body[:body_id])
-            # end index of shape
-            end_idx = start_idx + num_shapes_per_body[body_id]
-            # assign the new materials
-            # material ids are of shape: num_env_ids x num_shapes
-            # material_buckets are of shape: num_buckets x 3
-            materials[env_ids, start_idx:end_idx] = torch.from_numpy(material_samples[:, start_idx:end_idx]).to(
-                dtype=torch.float
-            )
-    else:
-        materials[env_ids] = torch.from_numpy(material_samples).to(dtype=torch.float)
+        # note: we only sample the materials once during initialization
+        #   afterwards these are randomly assigned to the geometries of the asset
+        range_list = [static_friction_range, dynamic_friction_range, restitution_range]
+        ranges = torch.tensor(range_list, device="cpu")
+        self.material_buckets = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (num_buckets, 3), device="cpu")
+
+        # ensure dynamic friction is always less than static friction
+        make_consistent = cfg.params.get("make_consistent", False)
+        if make_consistent:
+            self.material_buckets[:, 1] = torch.min(self.material_buckets[:, 0], self.material_buckets[:, 1])
+
+    def __call__(
+        self,
+        env: ManagerBasedEnv,
+        env_ids: torch.Tensor | None,
+        static_friction_range: tuple[float, float],
+        dynamic_friction_range: tuple[float, float],
+        restitution_range: tuple[float, float],
+        num_buckets: int,
+        asset_cfg: SceneEntityCfg,
+        make_consistent: bool = False,
+    ):
+        # resolve environment ids
+        if env_ids is None:
+            env_ids = torch.arange(env.scene.num_envs, device="cpu")
+        else:
+            env_ids = env_ids.cpu()
+
+        # randomly assign material IDs to the geometries
+        total_num_shapes = self.asset.root_physx_view.max_shapes
+        bucket_ids = torch.randint(0, num_buckets, (len(env_ids), total_num_shapes), device="cpu")
+        material_samples = self.material_buckets[bucket_ids]
+
+        # retrieve material buffer from the physics simulation
+        materials = self.asset.root_physx_view.get_material_properties()
+
+        # update material buffer with new samples
+        if self.num_shapes_per_body is not None:
+            # sample material properties from the given ranges
+            for body_id in self.asset_cfg.body_ids:
+                # obtain indices of shapes for the body
+                start_idx = sum(self.num_shapes_per_body[:body_id])
+                end_idx = start_idx + self.num_shapes_per_body[body_id]
+                # assign the new materials
+                # material samples are of shape: num_env_ids x total_num_shapes x 3
+                materials[env_ids, start_idx:end_idx] = material_samples[:, start_idx:end_idx]
+        else:
+            # assign all the materials
+            materials[env_ids] = material_samples[:]
 
-    # apply to simulation
-    asset.root_physx_view.set_material_properties(materials, env_ids)
+        # apply to simulation
+        self.asset.root_physx_view.set_material_properties(materials, env_ids)
 
 
 def randomize_rigid_body_mass(