Adds performance optimizations for domain randomization (#494)

# Description

- Adds `min_step_count_between_reset` for Event terms to specify minimum
randomization frequency for reset randomizations
- Switch to numpy for material bucketing logic
- Updates SH Config with DR

## Type of change

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

## 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 run all the tests with `./isaaclab.sh --test` and they pass
- [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

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

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.20.6"
+version = "0.20.7"
 
 # Description
 title = "Isaac Lab framework for Robot Learning"

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

 Changelog
 ---------
 
+0.20.7 (2024-08-02)
+~~~~~~~~~~~~~~~~~~~
+
+Changed
+^^^^^^^
+
+* Performance improvements for material randomization in events.
+
+Added
+^^^^^
+
+* Added minimum randomization frequency for reset mode randomizations.
+
+
 0.20.6 (2024-08-02)
 ~~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/direct_rl_env.py

@@ -287,7 +287,8 @@ class DirectRLEnv(gym.Env):
         """
         # add action noise
         if self.cfg.action_noise_model:
-            action = self._action_noise_model.apply(action.clone())
+            action = self._action_noise_model.apply(action)
+
         # process actions
         self._pre_physics_step(action)
 
@@ -518,7 +519,8 @@ class DirectRLEnv(gym.Env):
         # apply events such as randomization for environments that need a reset
         if self.cfg.events:
             if "reset" in self.event_manager.available_modes:
-                self.event_manager.apply(env_ids=env_ids, mode="reset")
+                env_step_count = self._sim_step_counter // self.cfg.decimation
+                self.event_manager.apply(env_ids=env_ids, mode="reset", global_env_step_count=env_step_count)
 
         # reset noise models
         if self.cfg.action_noise_model:

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_env.py

@@ -341,7 +341,8 @@ class ManagerBasedEnv:
         self.scene.reset(env_ids)
         # apply events such as randomizations for environments that need a reset
         if "reset" in self.event_manager.available_modes:
-            self.event_manager.apply(env_ids=env_ids, mode="reset")
+            env_step_count = self._sim_step_counter // self.cfg.decimation
+            self.event_manager.apply(env_ids=env_ids, mode="reset", global_env_step_count=env_step_count)
 
         # iterate over all managers and reset them
         # this returns a dictionary of information which is stored in the extras

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_rl_env.py

@@ -319,7 +319,8 @@ class ManagerBasedRLEnv(ManagerBasedEnv, gym.Env):
         self.scene.reset(env_ids)
         # apply events such as randomizations for environments that need a reset
         if "reset" in self.event_manager.available_modes:
-            self.event_manager.apply(env_ids=env_ids, mode="reset")
+            env_step_count = self._sim_step_counter // self.cfg.decimation
+            self.event_manager.apply(env_ids=env_ids, mode="reset", global_env_step_count=env_step_count)
 
         # iterate over all managers and reset them
         # this returns a dictionary of information which is stored in the extras

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

@@ -14,6 +14,7 @@ the event introduced by the function.
 
 from __future__ import annotations
 
+import numpy as np
 import torch
 from typing import TYPE_CHECKING, Literal
 
@@ -80,20 +81,20 @@ def randomize_rigid_body_material(
     materials = asset.root_physx_view.get_material_properties()
 
     # sample material properties from the given ranges
-    material_samples = torch.zeros(materials[env_ids].shape)
-    material_samples[..., 0].uniform_(*static_friction_range)
-    material_samples[..., 1].uniform_(*dynamic_friction_range)
-    material_samples[..., 2].uniform_(*restitution_range)
+    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)
 
     # create uniform range tensor for bucketing
-    lo = torch.tensor([static_friction_range[0], dynamic_friction_range[0], restitution_range[0]], device="cpu")
-    hi = torch.tensor([static_friction_range[1], dynamic_friction_range[1], restitution_range[1]], device="cpu")
+    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]])
 
     # 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 = torch.tensor([(hi[d] - lo[d]) * i / num_buckets + lo[d] for i in range(num_buckets)], device="cpu")
-        material_samples[..., d] = buckets[torch.searchsorted(buckets, material_samples[..., d].contiguous()) - 1]
+        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]
 
     # update material buffer with new samples
     if isinstance(asset, Articulation) and asset_cfg.body_ids != slice(None):
@@ -114,9 +115,11 @@ def randomize_rigid_body_material(
             # 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] = material_samples[:, start_idx:end_idx]
+            materials[env_ids, start_idx:end_idx] = torch.from_numpy(material_samples[:, start_idx:end_idx]).to(
+                dtype=torch.float
+            )
     else:
-        materials[env_ids] = material_samples
+        materials[env_ids] = torch.from_numpy(material_samples).to(dtype=torch.float)
 
     # apply to simulation
     asset.root_physx_view.set_material_properties(materials, env_ids)

source/extensions/omni.isaac.lab/omni/isaac/lab/managers/event_manager.py

@@ -119,7 +119,13 @@ class EventManager(ManagerBase):
         # nothing to log here
         return {}
 
-    def apply(self, mode: str, env_ids: Sequence[int] | None = None, dt: float | None = None):
+    def apply(
+        self,
+        mode: str,
+        env_ids: Sequence[int] | None = None,
+        dt: float | None = None,
+        global_env_step_count: int | None = None,
+    ):
         """Calls each event term in the specified mode.
 
         Note:
@@ -132,6 +138,8 @@ class EventManager(ManagerBase):
                 Defaults to None, in which case the event is applied to all environments.
             dt: The time step of the environment. This is only used for the "interval" mode.
                 Defaults to None to simplify the call for other modes.
+            global_env_step_count: The environment step count of the task. This is only used for the "reset" mode.
+                Defaults to None to simplify the call for other modes.
 
         Raises:
             ValueError: If the mode is ``"interval"`` and the time step is not provided.
@@ -171,6 +179,23 @@ class EventManager(ManagerBase):
                     if len(env_ids) > 0:
                         lower, upper = term_cfg.interval_range_s
                         time_left[env_ids] = torch.rand(len(env_ids), device=self.device) * (upper - lower) + lower
+            # check for minimum frequency for reset
+            elif mode == "reset":
+                if global_env_step_count is None:
+                    raise ValueError(
+                        f"Event mode '{mode}' requires the step count of the environment"
+                        " to be passed to the event manager."
+                    )
+
+                if env_ids is not None and len(env_ids) > 0:
+                    last_reset_step = self._reset_mode_last_reset_step_count[index]
+                    steps_since_last_reset = global_env_step_count - last_reset_step
+                    env_ids = env_ids[steps_since_last_reset[env_ids] >= term_cfg.min_step_count_between_reset]
+                    if len(env_ids) > 0:
+                        last_reset_step[env_ids] = global_env_step_count
+                    else:
+                        # no need to call func to sample
+                        continue
             # call the event term
             term_cfg.func(self._env, env_ids, **term_cfg.params)
 
@@ -234,6 +259,8 @@ class EventManager(ManagerBase):
         self._interval_mode_time_left: list[torch.Tensor] = list()
         # global timer for "interval" mode for global properties
         self._interval_mode_time_global: list[torch.Tensor] = list()
+        # buffer to store the step count when reset was last performed for each environment for "reset" mode
+        self._reset_mode_last_reset_step_count: list[torch.Tensor] = list()
 
         # check if config is dict already
         if isinstance(self.cfg, dict):
@@ -251,6 +278,13 @@ class EventManager(ManagerBase):
                     f"Configuration for the term '{term_name}' is not of type EventTermCfg."
                     f" Received: '{type(term_cfg)}'."
                 )
+
+            if term_cfg.mode != "reset" and term_cfg.min_step_count_between_reset != 0:
+                carb.log_warn(
+                    f"Event term '{term_name}' has 'min_step_count_between_reset' set to a non-zero value"
+                    " but the mode is not 'reset'. Ignoring the 'min_step_count_between_reset' value."
+                )
+
             # resolve common parameters
             self._resolve_common_term_cfg(term_name, term_cfg, min_argc=2)
             # check if mode is a new mode
@@ -284,3 +318,7 @@ class EventManager(ManagerBase):
                     lower, upper = term_cfg.interval_range_s
                     time_left = torch.rand(self.num_envs, device=self.device) * (upper - lower) + lower
                     self._interval_mode_time_left.append(time_left)
+
+            elif term_cfg.mode == "reset":
+                step_count = torch.zeros(self.num_envs, device=self.device, dtype=torch.int32)
+                self._reset_mode_last_reset_step_count.append(step_count)

source/extensions/omni.isaac.lab/omni/isaac/lab/managers/manager_term_cfg.py

@@ -209,6 +209,17 @@ class EventTermCfg(ManagerTermBaseCfg):
         This is only used if the mode is ``"interval"``.
     """
 
+    min_step_count_between_reset: int = 0
+    """The minimum number of environment steps between when term is applied.
+
+    When mode is "reset", the term will not be applied on the next reset unless
+    the number of steps since the last application of the term has exceeded this.
+    This is useful to avoid calling this term too often and improve performance.
+
+    Note:
+        This is only used if the mode is ``"reset"``.
+    """
+
 
 ##
 # Reward manager.

source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand/shadow_hand_env_cfg.py

@@ -29,8 +29,9 @@ class EventCfg:
     robot_physics_material = EventTerm(
         func=mdp.randomize_rigid_body_material,
         mode="reset",
+        min_step_count_between_reset=720,
         params={
-            "asset_cfg": SceneEntityCfg("robot", body_names=".*"),
+            "asset_cfg": SceneEntityCfg("robot"),
             "static_friction_range": (0.7, 1.3),
             "dynamic_friction_range": (1.0, 1.0),
             "restitution_range": (1.0, 1.0),
@@ -39,6 +40,7 @@ class EventCfg:
     )
     robot_joint_stiffness_and_damping = EventTerm(
         func=mdp.randomize_actuator_gains,
+        min_step_count_between_reset=720,
         mode="reset",
         params={
             "asset_cfg": SceneEntityCfg("robot", joint_names=".*"),
@@ -50,6 +52,7 @@ class EventCfg:
     )
     robot_joint_limits = EventTerm(
         func=mdp.randomize_joint_parameters,
+        min_step_count_between_reset=720,
         mode="reset",
         params={
             "asset_cfg": SceneEntityCfg("robot", joint_names=".*"),
@@ -61,6 +64,7 @@ class EventCfg:
     )
     robot_tendon_properties = EventTerm(
         func=mdp.randomize_fixed_tendon_parameters,
+        min_step_count_between_reset=720,
         mode="reset",
         params={
             "asset_cfg": SceneEntityCfg("robot", fixed_tendon_names=".*"),
@@ -74,9 +78,10 @@ class EventCfg:
     # -- object
     object_physics_material = EventTerm(
         func=mdp.randomize_rigid_body_material,
+        min_step_count_between_reset=720,
         mode="reset",
         params={
-            "asset_cfg": SceneEntityCfg("object", body_names=".*"),
+            "asset_cfg": SceneEntityCfg("object"),
             "static_friction_range": (0.7, 1.3),
             "dynamic_friction_range": (1.0, 1.0),
             "restitution_range": (1.0, 1.0),
@@ -85,6 +90,7 @@ class EventCfg:
     )
     object_scale_mass = EventTerm(
         func=mdp.randomize_rigid_body_mass,
+        min_step_count_between_reset=720,
         mode="reset",
         params={
             "asset_cfg": SceneEntityCfg("object"),
@@ -223,9 +229,18 @@ class ShadowHandEnvCfg(DirectRLEnvCfg):
 
 @configclass
 class ShadowHandOpenAIEnvCfg(ShadowHandEnvCfg):
+    # env
+    decimation = 3
+    episode_length_s = 8.0
+    num_actions = 20
+    num_observations = 42
+    num_states = 187
+    asymmetric_obs = True
+    obs_type = "openai"
     # simulation
     sim: SimulationCfg = SimulationCfg(
         dt=1 / 60,
+        render_interval=decimation,
         physics_material=RigidBodyMaterialCfg(
             static_friction=1.0,
             dynamic_friction=1.0,
@@ -236,14 +251,6 @@ class ShadowHandOpenAIEnvCfg(ShadowHandEnvCfg):
             gpu_max_rigid_patch_count=2**23,
         ),
     )
-    # env
-    decimation = 3
-    episode_length_s = 8.0
-    num_actions = 20
-    num_observations = 42
-    num_states = 187
-    asymmetric_obs = True
-    obs_type = "openai"
     # reset
     reset_position_noise = 0.01  # range of position at reset
     reset_dof_pos_noise = 0.2  # range of dof pos at reset