Unverified Commit 6b794ac2 authored by Mayank Mittal's avatar Mayank Mittal Committed by GitHub

Adds USD-level randomization mode to event manager (#2040)

# Description

Certain scene-level randomizations (such as randomizing the scale) must
happen before the simulation starts playing. To this end, the MR adds a
new event mode called "prestartup," which gets called right after the
scene design is complete and before the simulation is played.

Since the scene entities cannot be resolved before the simulation starts
playing (as we currently rely on PhysX to provide us with the joint/body
ordering), the MR adds a callback to resolve the scene entity
configurations separately once the simulation plays.

This MR replaces the prior implementation in #1165

## Type of change

- New feature (non-breaking change which adds functionality)
- This change requires a documentation update

## Screenshots

If you execute:

```bash
./isaaclab.sh -p scripts/tutorials/03_envs/create_cube_base_env.py --num_envs 32
```

Output:


![image](https://github.com/user-attachments/assets/a4b3a5e7-ff5c-4593-97ce-3b5af19bfde2)

## 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

---------
Signed-off-by: 's avatarMayank Mittal <12863862+Mayankm96@users.noreply.github.com>
Co-authored-by: 's avatarJames Smith <142246516+jsmith-bdai@users.noreply.github.com>
Co-authored-by: 's avatarKelly Guo <kellyguo123@hotmail.com>
parent 80e3964d
...@@ -56,7 +56,8 @@ ...@@ -56,7 +56,8 @@
"discretization", "discretization",
"trimesh", "trimesh",
"uninstanceable", "uninstanceable",
"coeff" "coeff",
"prestartup"
], ],
// This enables python language server. Seems to work slightly better than jedi: // This enables python language server. Seems to work slightly better than jedi:
"python.languageServer": "Pylance", "python.languageServer": "Pylance",
......
...@@ -9,7 +9,15 @@ controller to track an arbitrary target position. ...@@ -9,7 +9,15 @@ controller to track an arbitrary target position.
While going through this tutorial, we recommend you to pay attention to how a custom action term 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 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 'prestartup', 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 .. code-block:: bash
...@@ -223,6 +231,9 @@ class ObservationsCfg: ...@@ -223,6 +231,9 @@ class ObservationsCfg:
class EventCfg: class EventCfg:
"""Configuration for events.""" """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( reset_base = EventTerm(
func=mdp.reset_root_state_uniform, func=mdp.reset_root_state_uniform,
mode="reset", mode="reset",
...@@ -237,6 +248,19 @@ class EventCfg: ...@@ -237,6 +248,19 @@ class EventCfg:
}, },
) )
# This event term randomizes the scale of the cube.
# The mode is set to 'prestartup', 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="prestartup",
params={
"scale_range": {"x": (0.5, 1.5), "y": (0.5, 1.5), "z": (0.5, 1.5)},
"asset_cfg": SceneEntityCfg("cube"),
},
)
## ##
# Environment configuration # Environment configuration
...@@ -248,7 +272,11 @@ class CubeEnvCfg(ManagerBasedEnvCfg): ...@@ -248,7 +272,11 @@ class CubeEnvCfg(ManagerBasedEnvCfg):
"""Configuration for the locomotion velocity-tracking environment.""" """Configuration for the locomotion velocity-tracking environment."""
# Scene settings # 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 # Basic settings
observations: ObservationsCfg = ObservationsCfg() observations: ObservationsCfg = ObservationsCfg()
actions: ActionsCfg = ActionsCfg() actions: ActionsCfg = ActionsCfg()
...@@ -261,6 +289,7 @@ class CubeEnvCfg(ManagerBasedEnvCfg): ...@@ -261,6 +289,7 @@ class CubeEnvCfg(ManagerBasedEnvCfg):
# simulation settings # simulation settings
self.sim.dt = 0.01 self.sim.dt = 0.01
self.sim.physics_material = self.scene.terrain.physics_material self.sim.physics_material = self.scene.terrain.physics_material
self.sim.render_interval = 2 # render interval should be a multiple of decimation
def main(): def main():
......
[package] [package]
# Note: Semantic Versioning is used: https://semver.org/ # Note: Semantic Versioning is used: https://semver.org/
version = "0.36.1" version = "0.36.2"
# Description # Description
title = "Isaac Lab framework for Robot Learning" title = "Isaac Lab framework for Robot Learning"
......
Changelog Changelog
--------- ---------
0.36.2 (2025-03-12)
~~~~~~~~~~~~~~~~~~~
Added
^^^^^
* Added a new event mode called "prestartup", which gets called right after the scene design is complete
and before the simulation is played.
* Added a callback to resolve the scene entity configurations separately once the simulation plays,
since the scene entities cannot be resolved before the simulation starts playing
(as we currently rely on PhysX to provide us with the joint/body ordering)
0.36.1 (2025-03-10) 0.36.1 (2025-03-10)
~~~~~~~~~~~~~~~~~~~ ~~~~~~~~~~~~~~~~~~~
......
...@@ -95,6 +95,10 @@ class DirectMARLEnv(gym.Env): ...@@ -95,6 +95,10 @@ class DirectMARLEnv(gym.Env):
else: else:
raise RuntimeError("Simulation context already exists. Cannot create a new one.") 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 useful information
print("[INFO]: Base environment:") print("[INFO]: Base environment:")
print(f"\tEnvironment device : {self.device}") print(f"\tEnvironment device : {self.device}")
...@@ -126,6 +130,17 @@ class DirectMARLEnv(gym.Env): ...@@ -126,6 +130,17 @@ class DirectMARLEnv(gym.Env):
else: else:
self.viewport_camera_controller = None self.viewport_camera_controller = None
# create event manager
# note: this is needed here (rather than after simulation play) to allow USD-related randomization events
# that must happen before the simulation starts. Example: randomizing mesh scale
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 "prestartup" in self.event_manager.available_modes:
self.event_manager.apply(mode="prestartup")
# play the simulator to activate physics handles # play the simulator to activate physics handles
# note: this activates the physics simulation view that exposes TensorAPIs # 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 # note: when started in extension mode, first call sim.reset_async() and then initialize the managers
...@@ -138,15 +153,6 @@ class DirectMARLEnv(gym.Env): ...@@ -138,15 +153,6 @@ class DirectMARLEnv(gym.Env):
# this shouldn't cause an issue since later on, users do a reset over all the environments so the lazy buffers would be reset. # 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) 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 # check if debug visualization is has been implemented by the environment
source_code = inspect.getsource(self._set_debug_vis_impl) source_code = inspect.getsource(self._set_debug_vis_impl)
self.has_debug_vis_implementation = "NotImplementedError" not in source_code self.has_debug_vis_implementation = "NotImplementedError" not in source_code
......
...@@ -101,6 +101,10 @@ class DirectRLEnv(gym.Env): ...@@ -101,6 +101,10 @@ class DirectRLEnv(gym.Env):
else: else:
raise RuntimeError("Simulation context already exists. Cannot create a new one.") 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 useful information
print("[INFO]: Base environment:") print("[INFO]: Base environment:")
print(f"\tEnvironment device : {self.device}") print(f"\tEnvironment device : {self.device}")
...@@ -132,6 +136,17 @@ class DirectRLEnv(gym.Env): ...@@ -132,6 +136,17 @@ class DirectRLEnv(gym.Env):
else: else:
self.viewport_camera_controller = None self.viewport_camera_controller = None
# create event manager
# note: this is needed here (rather than after simulation play) to allow USD-related randomization events
# that must happen before the simulation starts. Example: randomizing mesh scale
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 "prestartup" in self.event_manager.available_modes:
self.event_manager.apply(mode="prestartup")
# play the simulator to activate physics handles # play the simulator to activate physics handles
# note: this activates the physics simulation view that exposes TensorAPIs # 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 # note: when started in extension mode, first call sim.reset_async() and then initialize the managers
...@@ -144,15 +159,6 @@ class DirectRLEnv(gym.Env): ...@@ -144,15 +159,6 @@ class DirectRLEnv(gym.Env):
# this shouldn't cause an issue since later on, users do a reset over all the environments so the lazy buffers would be reset. # 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) 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 # check if debug visualization is has been implemented by the environment
source_code = inspect.getsource(self._set_debug_vis_impl) source_code = inspect.getsource(self._set_debug_vis_impl)
self.has_debug_vis_implementation = "NotImplementedError" not in source_code self.has_debug_vis_implementation = "NotImplementedError" not in source_code
...@@ -199,7 +205,8 @@ class DirectRLEnv(gym.Env): ...@@ -199,7 +205,8 @@ class DirectRLEnv(gym.Env):
if "startup" in self.event_manager.available_modes: if "startup" in self.event_manager.available_modes:
self.event_manager.apply(mode="startup") 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 self.metadata["render_fps"] = 1 / self.step_dt
# print the environment information # print the environment information
......
...@@ -99,6 +99,10 @@ class ManagerBasedEnv: ...@@ -99,6 +99,10 @@ class ManagerBasedEnv:
raise RuntimeError("Simulation context already exists. Cannot create a new one.") raise RuntimeError("Simulation context already exists. Cannot create a new one.")
self.sim: SimulationContext = SimulationContext.instance() 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 useful information
print("[INFO]: Base environment:") print("[INFO]: Base environment:")
print(f"\tEnvironment device : {self.device}") print(f"\tEnvironment device : {self.device}")
...@@ -132,6 +136,16 @@ class ManagerBasedEnv: ...@@ -132,6 +136,16 @@ class ManagerBasedEnv:
else: else:
self.viewport_camera_controller = None self.viewport_camera_controller = None
# create event manager
# note: this is needed here (rather than after simulation play) to allow USD-related randomization events
# that must happen before the simulation starts. Example: randomizing mesh scale
self.event_manager = EventManager(self.cfg.events, self)
print("[INFO] Event Manager: ", self.event_manager)
# apply USD-related randomization events
if "prestartup" in self.event_manager.available_modes:
self.event_manager.apply(mode="prestartup")
# play the simulator to activate physics handles # play the simulator to activate physics handles
# note: this activates the physics simulation view that exposes TensorAPIs # 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 # note: when started in extension mode, first call sim.reset_async() and then initialize the managers
...@@ -146,10 +160,6 @@ class ManagerBasedEnv: ...@@ -146,10 +160,6 @@ class ManagerBasedEnv:
# add timeline event to load managers # add timeline event to load managers
self.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 # extend UI elements
# we need to do this here after all the managers are initialized # we need to do this here after all the managers are initialized
# this is because they dictate the sensors and commands right now # this is because they dictate the sensors and commands right now
...@@ -231,9 +241,6 @@ class ManagerBasedEnv: ...@@ -231,9 +241,6 @@ class ManagerBasedEnv:
# -- observation manager # -- observation manager
self.observation_manager = ObservationManager(self.cfg.observations, self) self.observation_manager = ObservationManager(self.cfg.observations, self)
print("[INFO] Observation Manager:", self.observation_manager) 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 # perform events at the start of the simulation
# in-case a child implementation creates other managers, the randomization should happen # in-case a child implementation creates other managers, the randomization should happen
......
...@@ -19,6 +19,8 @@ from typing import TYPE_CHECKING, Literal ...@@ -19,6 +19,8 @@ from typing import TYPE_CHECKING, Literal
import carb import carb
import omni.physics.tensors.impl.api as physx 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.sim as sim_utils
import isaaclab.utils.math as math_utils import isaaclab.utils.math as math_utils
...@@ -31,6 +33,116 @@ if TYPE_CHECKING: ...@@ -31,6 +33,116 @@ if TYPE_CHECKING:
from isaaclab.envs import ManagerBasedEnv 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): class randomize_rigid_body_material(ManagerTermBase):
"""Randomize the physics materials on all geometries of the asset. """Randomize the physics materials on all geometries of the asset.
......
...@@ -38,7 +38,9 @@ class EventManager(ManagerBase): ...@@ -38,7 +38,9 @@ class EventManager(ManagerBase):
For a typical training process, you may want to apply events in the following modes: For a typical training process, you may want to apply events in the following modes:
- "startup": Event is applied once at the beginning of the training. - "prestartup": Event is applied once at the beginning of the training 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 once simulation is started.
- "reset": Event is applied at every reset. - "reset": Event is applied at every reset.
- "interval": Event is applied at pre-specified intervals of time. - "interval": Event is applied at pre-specified intervals of time.
...@@ -184,6 +186,16 @@ class EventManager(ManagerBase): ...@@ -184,6 +186,16 @@ class EventManager(ManagerBase):
if mode not in self._mode_term_names: if mode not in self._mode_term_names:
omni.log.warn(f"Event mode '{mode}' is not defined. Skipping event.") omni.log.warn(f"Event mode '{mode}' is not defined. Skipping event.")
return return
# check if mode is pre-startup and scene replication is enabled
if mode == "prestartup" 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 # check if mode is interval and dt is not provided
if mode == "interval" and dt is None: if mode == "interval" and dt is None:
raise ValueError(f"Event mode '{mode}' requires the time-step of the environment.") raise ValueError(f"Event mode '{mode}' requires the time-step of the environment.")
......
...@@ -7,11 +7,13 @@ from __future__ import annotations ...@@ -7,11 +7,13 @@ from __future__ import annotations
import copy import copy
import inspect import inspect
import weakref
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from collections.abc import Sequence from collections.abc import Sequence
from typing import TYPE_CHECKING, Any from typing import TYPE_CHECKING, Any
import omni.log import omni.log
import omni.timeline
import isaaclab.utils.string as string_utils import isaaclab.utils.string as string_utils
from isaaclab.utils import string_to_callable from isaaclab.utils import string_to_callable
...@@ -119,6 +121,13 @@ class ManagerBase(ABC): ...@@ -119,6 +121,13 @@ class ManagerBase(ABC):
def __init__(self, cfg: object, env: ManagerBasedEnv): def __init__(self, cfg: object, env: ManagerBasedEnv):
"""Initialize the manager. """Initialize the manager.
This function is responsible for parsing the configuration object and creating the terms.
If the simulation is not playing, the scene entities are not resolved immediately.
Instead, the resolution is deferred until the simulation starts. This is done to ensure
that the scene entities are resolved even if the manager is created after the simulation
has already started.
Args: Args:
cfg: The configuration object. If None, the manager is initialized without any terms. cfg: The configuration object. If None, the manager is initialized without any terms.
env: The environment instance. env: The environment instance.
...@@ -126,10 +135,34 @@ class ManagerBase(ABC): ...@@ -126,10 +135,34 @@ class ManagerBase(ABC):
# store the inputs # store the inputs
self.cfg = copy.deepcopy(cfg) self.cfg = copy.deepcopy(cfg)
self._env = env self._env = env
# parse config to create terms information # parse config to create terms information
if self.cfg: if self.cfg:
self._prepare_terms() self._prepare_terms()
# if the simulation is not playing, we use callbacks to trigger the resolution of the scene
# entities configuration. this is needed for cases where the manager is created after the
# simulation, but before the simulation is playing.
if not self._env.sim.is_playing():
# note: Use weakref on all callbacks to ensure that this object can be deleted when its destructor
# is called
# The order is set to 20 to allow asset/sensor initialization to complete before the scene entities
# are resolved. Those have the order 10.
timeline_event_stream = omni.timeline.get_timeline_interface().get_timeline_event_stream()
self._resolve_scene_entities_handle = timeline_event_stream.create_subscription_to_pop_by_type(
int(omni.timeline.TimelineEventType.PLAY),
lambda event, obj=weakref.proxy(self): obj._resolve_scene_entities_callback(event),
order=20,
)
else:
self._resolve_scene_entities_handle = None
def __del__(self):
"""Delete the manager."""
if self._resolve_scene_entities_handle:
self._resolve_scene_entities_handle.unsubscribe()
self._resolve_scene_entities_handle = None
""" """
Properties. Properties.
""" """
...@@ -212,6 +245,30 @@ class ManagerBase(ABC): ...@@ -212,6 +245,30 @@ class ManagerBase(ABC):
"""Prepare terms information from the configuration object.""" """Prepare terms information from the configuration object."""
raise NotImplementedError raise NotImplementedError
"""
Internal callbacks.
"""
def _resolve_scene_entities_callback(self, event):
"""Resolve the scene entities configuration.
This callback is called when the simulation starts. It is used to resolve the
scene entities configuration for the terms.
"""
# check if config is dict already
if isinstance(self.cfg, dict):
cfg_items = self.cfg.items()
else:
cfg_items = self.cfg.__dict__.items()
# iterate over all the terms
for term_name, term_cfg in cfg_items:
# check for non config
if term_cfg is None:
continue
# resolve the scene entity configuration
self._resolve_scene_entity_cfg(term_name, term_cfg)
""" """
Helper functions. Helper functions.
""" """
...@@ -221,7 +278,6 @@ class ManagerBase(ABC): ...@@ -221,7 +278,6 @@ class ManagerBase(ABC):
Usually, called by the :meth:`_prepare_terms` method to resolve common term configuration. Usually, called by the :meth:`_prepare_terms` method to resolve common term configuration.
Note:
By default, all term functions are expected to have at least one argument, which is the By default, all term functions are expected to have at least one argument, which is the
environment object. Some other managers may expect functions to take more arguments, for environment object. Some other managers may expect functions to take more arguments, for
instance, the environment indices as the second argument. In such cases, the instance, the environment indices as the second argument. In such cases, the
...@@ -246,25 +302,10 @@ class ManagerBase(ABC): ...@@ -246,25 +302,10 @@ class ManagerBase(ABC):
f"Configuration for the term '{term_name}' is not of type ManagerTermBaseCfg." f"Configuration for the term '{term_name}' is not of type ManagerTermBaseCfg."
f" Received: '{type(term_cfg)}'." f" Received: '{type(term_cfg)}'."
) )
# iterate over all the entities and parse the joint and body names # iterate over all the entities and parse the joint and body names
for key, value in term_cfg.params.items(): if self._env.sim.is_playing():
# deal with string self._resolve_scene_entity_cfg(term_name, term_cfg)
if isinstance(value, SceneEntityCfg):
# load the entity
try:
value.resolve(self._env.scene)
except ValueError as e:
raise ValueError(f"Error while parsing '{term_name}:{key}'. {e}")
# log the entity for checking later
msg = f"[{term_cfg.__class__.__name__}:{term_name}] Found entity '{value.name}'."
if value.joint_ids is not None:
msg += f"\n\tJoint names: {value.joint_names} [{value.joint_ids}]"
if value.body_ids is not None:
msg += f"\n\tBody names: {value.body_names} [{value.body_ids}]"
# print the information
omni.log.info(msg)
# store the entity
term_cfg.params[key] = value
# get the corresponding function or functional class # get the corresponding function or functional class
if isinstance(term_cfg.func, str): if isinstance(term_cfg.func, str):
...@@ -296,3 +337,28 @@ class ManagerBase(ABC): ...@@ -296,3 +337,28 @@ class ManagerBase(ABC):
f"The term '{term_name}' expects mandatory parameters: {args_without_defaults[min_argc:]}" f"The term '{term_name}' expects mandatory parameters: {args_without_defaults[min_argc:]}"
f" and optional parameters: {args_with_defaults}, but received: {term_params}." f" and optional parameters: {args_with_defaults}, but received: {term_params}."
) )
def _resolve_scene_entity_cfg(self, term_name: str, term_cfg: ManagerTermBaseCfg):
"""Resolve the scene entity configuration for the term.
Args:
term_name: The name of the term.
term_cfg: The term configuration.
"""
for key, value in term_cfg.params.items():
if isinstance(value, SceneEntityCfg):
# load the entity
try:
value.resolve(self._env.scene)
except ValueError as e:
raise ValueError(f"Error while parsing '{term_name}:{key}'. {e}")
# log the entity for checking later
msg = f"[{term_cfg.__class__.__name__}:{term_name}] Found entity '{value.name}'."
if value.joint_ids is not None:
msg += f"\n\tJoint names: {value.joint_names} [{value.joint_ids}]"
if value.body_ids is not None:
msg += f"\n\tBody names: {value.body_names} [{value.body_ids}]"
# print the information
omni.log.info(msg)
# store the entity
term_cfg.params[key] = value
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment