Adds USD-level randomization mode to event manager

scripts/tutorials/03_envs/create_cube_base_env.py

@@ -9,7 +9,15 @@
 
 While going through this tutorial, we recommend you to pay attention to how a custom action term
 is defined. The action term is responsible for processing the raw actions and applying them to the
-scene entities. The rest of the environment is similar to the previous tutorials.
+scene entities.
+
+We also define an event term called 'randomize_scale' that randomizes the scale of
+the cube. This event term has the mode 'usd', which means that it is applied on the USD stage
+before the simulation starts. Additionally, the flag 'replicate_physics' is set to False,
+which means that the cube is not replicated across multiple environments but rather each
+environment gets its own cube instance.
+
+The rest of the environment is similar to the previous tutorials.
 
 .. code-block:: bash
 
@@ -223,6 +231,9 @@ def __post_init__(self):
 class EventCfg:
     """Configuration for events."""
 
+    # This event term resets the base position of the cube.
+    # The mode is set to 'reset', which means that the base position is reset whenever
+    # the environment instance is reset (because of terminations defined in 'TerminationCfg').
     reset_base = EventTerm(
         func=mdp.reset_root_state_uniform,
         mode="reset",
@@ -237,6 +248,19 @@ class EventCfg:
         },
     )
 
+    # This event term randomizes the scale of the cube.
+    # The mode is set to 'usd', which means that the scale is randomize on the USD stage before the
+    # simulation starts.
+    # Note: USD-level randomizations require the flag 'replicate_physics' to be set to False.
+    randomize_scale = EventTerm(
+        func=mdp.randomize_rigid_body_scale,
+        mode="usd",
+        params={
+            "scale_range": {"x": (0.5, 1.5), "y": (0.5, 1.5), "z": (0.5, 1.5)},
+            "asset_cfg": SceneEntityCfg("cube"),
+        },
+    )
+
 
 ##
 # Environment configuration
@@ -248,7 +272,11 @@ class CubeEnvCfg(ManagerBasedEnvCfg):
     """Configuration for the locomotion velocity-tracking environment."""
 
     # Scene settings
-    scene: MySceneCfg = MySceneCfg(num_envs=args_cli.num_envs, env_spacing=2.5)
+    # The flag 'replicate_physics' is set to False, which means that the cube is not replicated
+    # across multiple environments but rather each environment gets its own cube instance.
+    # This allows modifying the cube's properties independently for each environment.
+    scene: MySceneCfg = MySceneCfg(num_envs=args_cli.num_envs, env_spacing=2.5, replicate_physics=False)
+
     # Basic settings
     observations: ObservationsCfg = ObservationsCfg()
     actions: ActionsCfg = ActionsCfg()
@@ -261,6 +289,7 @@ def __post_init__(self):
         # simulation settings
         self.sim.dt = 0.01
         self.sim.physics_material = self.scene.terrain.physics_material
+        self.sim.render_interval = 2  # render steps should be a multiple of decimation
 
 
 def main():

source/isaaclab/isaaclab/envs/direct_marl_env.py

@@ -95,6 +95,10 @@ def __init__(self, cfg: DirectMARLEnvCfg, render_mode: str | None = None, **kwar
         else:
             raise RuntimeError("Simulation context already exists. Cannot create a new one.")
 
+        # make sure torch is running on the correct device
+        if "cuda" in self.device:
+            torch.cuda.set_device(self.device)
+
         # print useful information
         print("[INFO]: Base environment:")
         print(f"\tEnvironment device    : {self.device}")
@@ -126,6 +130,17 @@ def __init__(self, cfg: DirectMARLEnvCfg, render_mode: str | None = None, **kwar
         else:
             self.viewport_camera_controller = None
 
+        # create event manager
+        # note: this is moved here to allow USD-related randomization events that need to be
+        #   applied before the simulation starts.
+        if self.cfg.events:
+            self.event_manager = EventManager(self.cfg.events, self)
+            print("[INFO] Event Manager: ", self.event_manager)
+
+            # apply USD-related randomization events
+            if "usd" in self.event_manager.available_modes:
+                self.event_manager.apply(mode="usd")
+
         # play the simulator to activate physics handles
         # note: this activates the physics simulation view that exposes TensorAPIs
         # note: when started in extension mode, first call sim.reset_async() and then initialize the managers
@@ -138,15 +153,6 @@ def __init__(self, cfg: DirectMARLEnvCfg, render_mode: str | None = None, **kwar
                 # this shouldn't cause an issue since later on, users do a reset over all the environments so the lazy buffers would be reset.
                 self.scene.update(dt=self.physics_dt)
 
-        # -- event manager used for randomization
-        if self.cfg.events:
-            self.event_manager = EventManager(self.cfg.events, self)
-            print("[INFO] Event Manager: ", self.event_manager)
-
-        # make sure torch is running on the correct device
-        if "cuda" in self.device:
-            torch.cuda.set_device(self.device)
-
         # check if debug visualization is has been implemented by the environment
         source_code = inspect.getsource(self._set_debug_vis_impl)
         self.has_debug_vis_implementation = "NotImplementedError" not in source_code

source/isaaclab/isaaclab/envs/direct_rl_env.py

@@ -101,6 +101,10 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
         else:
             raise RuntimeError("Simulation context already exists. Cannot create a new one.")
 
+        # make sure torch is running on the correct device
+        if "cuda" in self.device:
+            torch.cuda.set_device(self.device)
+
         # print useful information
         print("[INFO]: Base environment:")
         print(f"\tEnvironment device    : {self.device}")
@@ -132,6 +136,17 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
         else:
             self.viewport_camera_controller = None
 
+        # create event manager
+        # note: this is moved here to allow USD-related randomization events that need to be
+        #   applied before the simulation starts.
+        if self.cfg.events:
+            self.event_manager = EventManager(self.cfg.events, self)
+            print("[INFO] Event Manager: ", self.event_manager)
+
+            # apply USD-related randomization events
+            if "usd" in self.event_manager.available_modes:
+                self.event_manager.apply(mode="usd")
+
         # play the simulator to activate physics handles
         # note: this activates the physics simulation view that exposes TensorAPIs
         # note: when started in extension mode, first call sim.reset_async() and then initialize the managers
@@ -144,15 +159,6 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
                 # this shouldn't cause an issue since later on, users do a reset over all the environments so the lazy buffers would be reset.
                 self.scene.update(dt=self.physics_dt)
 
-        # -- event manager used for randomization
-        if self.cfg.events:
-            self.event_manager = EventManager(self.cfg.events, self)
-            print("[INFO] Event Manager: ", self.event_manager)
-
-        # make sure torch is running on the correct device
-        if "cuda" in self.device:
-            torch.cuda.set_device(self.device)
-
         # check if debug visualization is has been implemented by the environment
         source_code = inspect.getsource(self._set_debug_vis_impl)
         self.has_debug_vis_implementation = "NotImplementedError" not in source_code
@@ -199,7 +205,8 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
             if "startup" in self.event_manager.available_modes:
                 self.event_manager.apply(mode="startup")
 
-        # -- set the framerate of the gym video recorder wrapper so that the playback speed of the produced video matches the simulation
+        # set the framerate of the gym video recorder wrapper so that the playback speed of the produced
+        # video matches the simulation
         self.metadata["render_fps"] = 1 / self.step_dt
 
         # print the environment information

source/isaaclab/isaaclab/envs/manager_based_env.py

@@ -99,6 +99,10 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
                 raise RuntimeError("Simulation context already exists. Cannot create a new one.")
             self.sim: SimulationContext = SimulationContext.instance()
 
+        # make sure torch is running on the correct device
+        if "cuda" in self.device:
+            torch.cuda.set_device(self.device)
+
         # print useful information
         print("[INFO]: Base environment:")
         print(f"\tEnvironment device    : {self.device}")
@@ -132,6 +136,16 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
         else:
             self.viewport_camera_controller = None
 
+        # create event manager
+        # note: this is moved here to allow USD-related randomization events that need to be
+        #   applied before the simulation starts.
+        self.event_manager = EventManager(self.cfg.events, self)
+        print("[INFO] Event Manager: ", self.event_manager)
+
+        # apply USD-related randomization events
+        if "usd" in self.event_manager.available_modes:
+            self.event_manager.apply(mode="usd")
+
         # play the simulator to activate physics handles
         # note: this activates the physics simulation view that exposes TensorAPIs
         # note: when started in extension mode, first call sim.reset_async() and then initialize the managers
@@ -146,10 +160,6 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
             # add timeline event to load managers
             self.load_managers()
 
-        # make sure torch is running on the correct device
-        if "cuda" in self.device:
-            torch.cuda.set_device(self.device)
-
         # extend UI elements
         # we need to do this here after all the managers are initialized
         # this is because they dictate the sensors and commands right now
@@ -231,9 +241,6 @@ def load_managers(self):
         # -- observation manager
         self.observation_manager = ObservationManager(self.cfg.observations, self)
         print("[INFO] Observation Manager:", self.observation_manager)
-        # -- event manager
-        self.event_manager = EventManager(self.cfg.events, self)
-        print("[INFO] Event Manager: ", self.event_manager)
 
         # perform events at the start of the simulation
         # in-case a child implementation creates other managers, the randomization should happen

source/isaaclab/isaaclab/envs/mdp/events.py

@@ -19,6 +19,8 @@
 
 import carb
 import omni.physics.tensors.impl.api as physx
+import omni.usd
+from pxr import Gf, Sdf, UsdGeom, Vt
 
 import isaaclab.sim as sim_utils
 import isaaclab.utils.math as math_utils
@@ -31,6 +33,116 @@
     from isaaclab.envs import ManagerBasedEnv
 
 
+def randomize_rigid_body_scale(
+    env: ManagerBasedEnv,
+    env_ids: torch.Tensor | None,
+    scale_range: tuple[float, float] | dict[str, tuple[float, float]],
+    asset_cfg: SceneEntityCfg,
+    relative_child_path: str | None = None,
+):
+    """Randomize the scale of a rigid body asset in the USD stage.
+
+    This function modifies the "xformOp:scale" property of all the prims corresponding to the asset.
+
+    It takes a tuple or dictionary for the scale ranges. If it is a tuple, then the scaling along
+    individual axis is performed equally. If it is a dictionary, the scaling is independent across each dimension.
+    The keys of the dictionary are ``x``, ``y``, and ``z``. The values are tuples of the form ``(min, max)``.
+
+    If the dictionary does not contain a key, the range is set to one for that axis.
+
+    Relative child path can be used to randomize the scale of a specific child prim of the asset.
+    For example, if the asset at prim path expression "/World/envs/env_.*/Object" has a child
+    with the path "/World/envs/env_.*/Object/mesh", then the relative child path should be "mesh" or
+    "/mesh".
+
+    .. attention::
+        Since this function modifies USD properties that are parsed by the physics engine once the simulation
+        starts, the term should only be used before the simulation starts playing. This corresponds to the
+        event mode named "usd". Using it at simulation time, may lead to unpredictable behaviors.
+
+    .. note::
+        When randomizing the scale of individual assets, please make sure to set
+        :attr:`isaaclab.scene.InteractiveSceneCfg.replicate_physics` to False. This ensures that physics
+        parser will parse the individual asset properties separately.
+    """
+    # check if sim is running
+    if env.sim.is_playing():
+        raise RuntimeError(
+            "Randomizing scale while simulation is running leads to unpredictable behaviors."
+            " Please ensure that the event term is called before the simulation starts by using the 'usd' mode."
+        )
+
+    # extract the used quantities (to enable type-hinting)
+    asset: RigidObject = env.scene[asset_cfg.name]
+
+    if isinstance(asset, Articulation):
+        raise ValueError(
+            "Scaling an articulation randomly is not supported, as it affects joint attributes and can cause"
+            " unexpected behavior. To achieve different scales, we recommend generating separate USD files for"
+            " each version of the articulation and using multi-asset spawning. For more details, refer to:"
+            " https://isaac-sim.github.io/IsaacLab/main/source/how-to/multi_asset_spawning.html"
+        )
+
+    # resolve environment ids
+    if env_ids is None:
+        env_ids = torch.arange(env.scene.num_envs, device="cpu")
+    else:
+        env_ids = env_ids.cpu()
+
+    # acquire stage
+    stage = omni.usd.get_context().get_stage()
+    # resolve prim paths for spawning and cloning
+    prim_paths = sim_utils.find_matching_prim_paths(asset.cfg.prim_path)
+
+    # sample scale values
+    if isinstance(scale_range, dict):
+        range_list = [scale_range.get(key, (1.0, 1.0)) for key in ["x", "y", "z"]]
+        ranges = torch.tensor(range_list, device="cpu")
+        rand_samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (len(env_ids), 3), device="cpu")
+    else:
+        rand_samples = math_utils.sample_uniform(*scale_range, (len(env_ids), 1), device="cpu")
+        rand_samples = rand_samples.repeat(1, 3)
+    # convert to list for the for loop
+    rand_samples = rand_samples.tolist()
+
+    # apply the randomization to the parent if no relative child path is provided
+    # this might be useful if user wants to randomize a particular mesh in the prim hierarchy
+    if relative_child_path is None:
+        relative_child_path = ""
+    elif not relative_child_path.startswith("/"):
+        relative_child_path = "/" + relative_child_path
+
+    # use sdf changeblock for faster processing of USD properties
+    with Sdf.ChangeBlock():
+        for i, env_id in enumerate(env_ids):
+            # path to prim to randomize
+            prim_path = prim_paths[env_id] + relative_child_path
+            # spawn single instance
+            prim_spec = Sdf.CreatePrimInLayer(stage.GetRootLayer(), prim_path)
+
+            # get the attribute to randomize
+            scale_spec = prim_spec.GetAttributeAtPath(prim_path + ".xformOp:scale")
+            # if the scale attribute does not exist, create it
+            has_scale_attr = scale_spec is not None
+            if not has_scale_attr:
+                scale_spec = Sdf.AttributeSpec(prim_spec, prim_path + ".xformOp:scale", Sdf.ValueTypeNames.Double3)
+
+            # set the new scale
+            scale_spec.default = Gf.Vec3f(*rand_samples[i])
+
+            # ensure the operation is done in the right ordering if we created the scale attribute.
+            # otherwise, we assume the scale attribute is already in the right order.
+            # note: by default isaac sim follows this ordering for the transform stack so any asset
+            #   created through it will have the correct ordering
+            if not has_scale_attr:
+                op_order_spec = prim_spec.GetAttributeAtPath(prim_path + ".xformOpOrder")
+                if op_order_spec is None:
+                    op_order_spec = Sdf.AttributeSpec(
+                        prim_spec, UsdGeom.Tokens.xformOpOrder, Sdf.ValueTypeNames.TokenArray
+                    )
+                op_order_spec.default = Vt.TokenArray(["xformOp:translate", "xformOp:orient", "xformOp:scale"])
+
+
 class randomize_rigid_body_material(ManagerTermBase):
     """Randomize the physics materials on all geometries of the asset.
 

source/isaaclab/isaaclab/managers/event_manager.py

@@ -38,6 +38,8 @@ class EventManager(ManagerBase):
 
     For a typical training process, you may want to apply events in the following modes:
 
+    - "usd": Event is applied before the simulation starts. This is used to randomize USD-level properties
+      of the simulation stage.
     - "startup": Event is applied once at the beginning of the training.
     - "reset": Event is applied at every reset.
     - "interval": Event is applied at pre-specified intervals of time.
@@ -184,6 +186,16 @@ def apply(
         if mode not in self._mode_term_names:
             omni.log.warn(f"Event mode '{mode}' is not defined. Skipping event.")
             return
+        # check if mode is usd and scene replication is enabled
+        if mode == "usd" and self._env.scene.cfg.replicate_physics:
+            omni.log.warn(
+                "Scene replication is enabled, which may affect USD-level randomization."
+                " When assets are replicated, their properties are shared across instances,"
+                " potentially leading to unintended behavior."
+                " For stable USD-level randomization, consider disabling scene replication"
+                " by setting 'replicate_physics' to False in 'InteractiveSceneCfg'."
+            )
+
         # check if mode is interval and dt is not provided
         if mode == "interval" and dt is None:
             raise ValueError(f"Event mode '{mode}' requires the time-step of the environment.")