From 3cf13890283721aebb7345119f205ca2778b4046 Mon Sep 17 00:00:00 2001 From: zrene Date: Tue, 9 Jul 2024 19:30:16 +0200 Subject: [PATCH 1/2] Adds multi asset randomization --- .../lab/sim/spawners/multi_asset/__init__.py | 9 + .../spawners/multi_asset/asset_randomizer.py | 39 ++++ .../multi_asset/asset_randomizer_cfg.py | 36 ++++ .../sim/spawners/multi_asset/multi_asset.py | 185 +++++++++++++++++ .../spawners/multi_asset/multi_asset_cfg.py | 42 ++++ .../multi_asset/randomizations/__init__.py | 12 ++ .../randomizations/randomizations.py | 79 +++++++ .../randomizations/randomizations_cfg.py | 64 ++++++ .../omni/isaac/lab/sim/utils.py | 7 +- source/standalone/demos/multi_object.py | 194 ++++++++++++++++++ 10 files changed, 666 insertions(+), 1 deletion(-) create mode 100644 source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/__init__.py create mode 100644 source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/asset_randomizer.py create mode 100644 source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/asset_randomizer_cfg.py create mode 100644 source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/multi_asset.py create mode 100644 source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/multi_asset_cfg.py create mode 100644 source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/randomizations/__init__.py create mode 100644 source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/randomizations/randomizations.py create mode 100644 source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/randomizations/randomizations_cfg.py create mode 100644 source/standalone/demos/multi_object.py diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/__init__.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/__init__.py new file mode 100644 index 0000000000..b0669a6e61 --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-module containing spawners for spawning multiple assets randomly and randomizing them.""" + +from .asset_randomizer_cfg import AssetRandomizerCfg +from .multi_asset_cfg import MultiAssetCfg diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/asset_randomizer.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/asset_randomizer.py new file mode 100644 index 0000000000..3a6f2eb12f --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/asset_randomizer.py @@ -0,0 +1,39 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from collections.abc import Iterable +from typing import TYPE_CHECKING + +from pxr import Usd + +if TYPE_CHECKING: + from .asset_randomizer_cfg import AssetRandomizerCfg + + +def randomize_and_spawn_asset( + prim_path: str, + cfg: AssetRandomizerCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, +) -> Usd.Prim: + """Randomizes and spawns an asset. + Args: + prim_path: The path to the asset to spawn + cfg: The configuration for the asset randomization + translation: The translation to apply to the spawned asset, gets passed to the child spawner defined in the configuration + orientation: The orientation to apply to the spawned asset, gets passed to the child spawner defined in the configuration + """ + prim = cfg.child_spawner_cfg.func(prim_path, cfg.child_spawner_cfg, translation, orientation) + randomizations = cfg.randomization_cfg + + if not isinstance(randomizations, Iterable): + randomizations = [randomizations] + + for randomization_cfg in randomizations: + prim = randomization_cfg.func(prim, randomization_cfg) + + return prim diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/asset_randomizer_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/asset_randomizer_cfg.py new file mode 100644 index 0000000000..227c0184ba --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/asset_randomizer_cfg.py @@ -0,0 +1,36 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configuration for loading multiple assets randomly.""" + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import MISSING + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.utils import configclass + +from .asset_randomizer import randomize_and_spawn_asset +from .randomizations.randomizations_cfg import RandomizationCfg + + +@configclass +class AssetRandomizerCfg(sim_utils.SpawnerCfg): + """Configuration for loading multiple assets randomly.""" + + func: sim_utils.SpawnerCfg.func = randomize_and_spawn_asset + + num_random_assets: int = 1 + """The number of random assets to spawn""" + + base_seed: int = 0 + """The seed to use for randomization for replicability""" + + child_spawner_cfg: sim_utils.SpawnerCfg = MISSING + """The configuration to use to spawn each asset""" + + randomization_cfg: RandomizationCfg | Iterable[RandomizationCfg] = MISSING + """Optional randomization configuration that is applied at the time of spawning prims.""" diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/multi_asset.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/multi_asset.py new file mode 100644 index 0000000000..4bee494fd4 --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/multi_asset.py @@ -0,0 +1,185 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Spawns multiple objects randomly in the scene.""" + +from __future__ import annotations + +import random +import re +import torch +from typing import TYPE_CHECKING + +import omni.isaac.core.utils.prims as prim_utils +import omni.usd +from pxr import Gf, Sdf, Semantics, Usd, UsdGeom, Vt + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.sim.spawners.multi_asset.asset_randomizer_cfg import AssetRandomizerCfg +from omni.isaac.lab.sim.spawners.spawner_cfg import SpawnerCfg + +if TYPE_CHECKING: + from .multi_asset_cfg import MultiAssetCfg + + +def spawn_multi_object_randomly_sdf( + prim_path: str, + cfg: MultiAssetCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, +) -> Usd.Prim: + """Spawns multiple objects randomly in the scene. + Args: + prim_path: The path to the asset to spawn (e.g. "/World/env_*/Object") + cfg: The configuration for the asset spawning + translation: The translation to apply to the spawned asset + orientation: The orientation to apply to the spawned asset + """ + + # resolve: {SPAWN_NS}/AssetName + # note: this assumes that the spawn namespace already exists in the stage + root_path, asset_path = prim_path.rsplit("/", 1) + + # check if input is a regex expression + # note: a valid prim path can only contain alphanumeric characters, underscores, and forward slashes + is_regex_expression = re.match(r"^[a-zA-Z0-9/_]+$", root_path) is None + + # resolve matching prims for source prim path expression + if is_regex_expression and root_path != "": + source_prim_paths = sim_utils.find_matching_prim_paths(root_path) + # if no matching prims are found, raise an error + if len(source_prim_paths) == 0: + raise RuntimeError( + f"Unable to find source prim path: '{root_path}'. Please create the prim before spawning." + ) + else: + source_prim_paths = [root_path] + + # spawn everything first in a "Dataset" prim + prim_utils.create_prim("/World/Dataset", "Scope") + proto_prim_paths = list() + + object_id = 0 + for asset_cfg in cfg.assets_cfg: + if isinstance(asset_cfg, AssetRandomizerCfg): + for i in range(asset_cfg.num_random_assets): + # set the seeds + random.seed(asset_cfg.base_seed + i) + proto_prim_paths.append(_spawn_dataset_objects(object_id, asset_cfg)) + object_id += 1 + else: + proto_prim_paths.append(_spawn_dataset_objects(object_id, asset_cfg)) + object_id += 1 + + # resolve prim paths for spawning and cloning + prim_paths = [f"{source_prim_path}/{asset_path}" for source_prim_path in source_prim_paths] + # acquire stage + stage = omni.usd.get_context().get_stage() + + if orientation is not None: + # convert orientation ordering (wxyz to xyzw) + orientation = (orientation[1], orientation[2], orientation[3], orientation[0]) + if isinstance(orientation[0], torch.Tensor): + orientation = [x.cpu().item() for x in orientation] + + # manually clone prims if the source prim path is a regex expression + with Sdf.ChangeBlock(): + for i, prim_path in enumerate(prim_paths): + # spawn single instance + env_spec = Sdf.CreatePrimInLayer(stage.GetRootLayer(), prim_path) + if cfg.randomize: + # randomly select an asset configuration + proto_path = random.choice(proto_prim_paths) + else: + proto_path = proto_prim_paths[i % len(proto_prim_paths)] + + Sdf.CopySpec( + env_spec.layer, + Sdf.Path(proto_path), + env_spec.layer, + Sdf.Path(prim_path), + ) + + # set the translation and orientation + _ = UsdGeom.Xform(stage.GetPrimAtPath(proto_path)).GetPrim().GetPrimStack() + + translate_spec = env_spec.GetAttributeAtPath(prim_path + ".xformOp:translate") + if translate_spec is None: + translate_spec = Sdf.AttributeSpec(env_spec, "xformOp:translate", Sdf.ValueTypeNames.Double3) + if translation is not None: + translate_spec.default = Gf.Vec3d(*translation) + + orient_spec = env_spec.GetAttributeAtPath(prim_path + ".xformOp:orient") + if orient_spec is None: + orient_spec = Sdf.AttributeSpec(env_spec, "xformOp:orient", Sdf.ValueTypeNames.Quatd) + + if orientation is not None: + orient_spec.default = Gf.Quatd(*orientation) + + op_order_spec = env_spec.GetAttributeAtPath(prim_path + ".xformOpOrder") + if op_order_spec is None: + op_order_spec = Sdf.AttributeSpec(env_spec, UsdGeom.Tokens.xformOpOrder, Sdf.ValueTypeNames.TokenArray) + op_order_spec.default = Vt.TokenArray(["xformOp:translate", "xformOp:orient", "xformOp:scale"]) + + if cfg.postprocess_func is not None: + cfg.postprocess_func(proto_path, prim_path, stage) + + # delete the dataset prim after spawning + prim_utils.delete_prim("/World/Dataset") + + # return the prim + return prim_utils.get_prim_at_path(prim_paths[0]) + + +### +# Internal Helper functions +### + + +def _spawn_dataset_objects( + object_id: int, + asset_cfg: SpawnerCfg, +): + """Helper function to spawn a single instance. + This functions expects a prototype prim to be located at "/World/Dataset/Object_{object_id:02d} + Args: + object_id: The unique identifier for the object + asset_cfg: The configuration for the asset to spawn + """ + + # spawn single instance + proto_prim_path = f"/World/Dataset/Object_{object_id:02d}" + + prim = asset_cfg.func(proto_prim_path, asset_cfg) + + # set the prim visibility + if hasattr(asset_cfg, "visible"): + imageable = UsdGeom.Imageable(prim) + if asset_cfg.visible: + imageable.MakeVisible() + else: + imageable.MakeInvisible() + + # set the semantic annotations + if hasattr(asset_cfg, "semantic_tags") and asset_cfg.semantic_tags is not None: + # note: taken from replicator scripts.utils.utils.py + for semantic_type, semantic_value in asset_cfg.semantic_tags: + # deal with spaces by replacing them with underscores + semantic_type_sanitized = semantic_type.replace(" ", "_") + semantic_value_sanitized = semantic_value.replace(" ", "_") + # set the semantic API for the instance + instance_name = f"{semantic_type_sanitized}_{semantic_value_sanitized}" + sem = Semantics.SemanticsAPI.Apply(prim, instance_name) + # create semantic type and data attributes + sem.CreateSemanticTypeAttr() + sem.CreateSemanticDataAttr() + sem.GetSemanticTypeAttr().Set(semantic_type) + sem.GetSemanticDataAttr().Set(semantic_value) + + # activate rigid body contact sensors + if hasattr(asset_cfg, "activate_contact_sensors") and asset_cfg.activate_contact_sensors: + sim_utils.activate_contact_sensors(proto_prim_path, asset_cfg.activate_contact_sensors) + + return proto_prim_path diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/multi_asset_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/multi_asset_cfg.py new file mode 100644 index 0000000000..53806235ef --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/multi_asset_cfg.py @@ -0,0 +1,42 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""This script demonstrates how to spawn multiple objects in multiple environments. +.. code-block:: bash + # Usage + ././isaaclab.sh -p source/standalone/demos/multi_object.py --num_envs 2048 +""" + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import MISSING + +from pxr import Usd + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.utils import configclass + +from .multi_asset import spawn_multi_object_randomly_sdf + + +@configclass +class MultiAssetCfg(sim_utils.SpawnerCfg): + """Configuration parameters for loading multiple assets randomly.""" + + func: sim_utils.SpawnerCfg.func = spawn_multi_object_randomly_sdf + + assets_cfg: Iterable[sim_utils.SpawnerCfg] = MISSING + """Iterable of asset configurations to spawn.""" + + postprocess_func: callable[[str, str, Usd.Stage], None] | None = None + """Optional postprocess function to call after spawning each assets. + The function should take two arguments: The path to the prototype prim, the path to the spawned prim in the environment, and the stage. + """ + + randomize: bool = True + """Whether to randomly spawn assets, or spawn them in the order they are defined in the configuration. + If True, the assets are spawned in random order by using random.choice (i.e. it is not guaranteed that all assets will be spawned the same number of times). + """ diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/randomizations/__init__.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/randomizations/__init__.py new file mode 100644 index 0000000000..f00d7f4a8d --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/randomizations/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-module containing randomizations that are applied at the time of spawning prims. +Note that this is different from the randomizations that are applied at the time of simulation reset, and allows +to modify properties such as scale, joint offsets, etc. at the time of spawning the prims. +""" + +from .randomizations import * # noqa: F401, F403 +from .randomizations_cfg import * # noqa: F401, F403 diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/randomizations/randomizations.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/randomizations/randomizations.py new file mode 100644 index 0000000000..188055ae62 --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/randomizations/randomizations.py @@ -0,0 +1,79 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Implementations for randomizations that are applied at the time of spawning prims.""" + +from __future__ import annotations + +import random +from typing import TYPE_CHECKING + +import carb +import omni.isaac.core.utils.prims as prim_utils +from pxr import Usd + +import omni.isaac.lab.sim as sim_utils + +if TYPE_CHECKING: + from .randomizations_cfg import RandomizeJointOffsetsCfg, RandomizeScaleCfg + + +def randomize_joint_offset(prim: Usd.Prim, cfg: RandomizeJointOffsetsCfg) -> Usd.Prim: + """Randomize the joint offsets of a joint in an articulation by modifying the physics:localPos0 property. + Args: + prim: The root prim of the articulation to randomize + cfg: The configuration for the randomization + """ + + # Find matching prims + matching_prims = sim_utils.find_matching_prim_paths(str(prim.GetPath()) + cfg.joint_name) + + if len(matching_prims) == 0: + return prim + + for prim_path in matching_prims: + prim = prim_utils.get_prim_at_path(prim_path) + prop = prim.GetProperty("physics:localPos0") + value = prop.Get() + value[0] += random.random() * (cfg.x_range[1] - cfg.x_range[0]) + cfg.x_range[0] + value[1] += random.random() * (cfg.y_range[1] - cfg.y_range[0]) + cfg.y_range[0] + value[2] += random.random() * (cfg.z_range[1] - cfg.z_range[0]) + cfg.z_range[0] + + # Set property + if not prop.Set(value): + carb.log_error("Failed to set property", prop, "for", prim_path) + return prim + + +def randomize_scale(prim: Usd.Prim, cfg: RandomizeScaleCfg) -> Usd.Prim: + """Randomize the scale of a prim by modifying the xformOp:scale property. + Args: + prim: The prim to randomize + cfg: The configuration for the randomization + """ + + matching_prims = sim_utils.find_matching_prim_paths(str(prim.GetPath())) + + if len(matching_prims) == 0: + return prim + + for prim_path in matching_prims: + prim = prim_utils.get_prim_at_path(prim_path) + prop = prim.GetProperty("xformOp:scale") + value = prop.Get() + if cfg.equal_scale: + scale = cfg.x_range[0] + random.random() * (cfg.x_range[1] - cfg.x_range[0]) + value[0] = scale * value[0] + value[1] = scale * value[1] + value[2] = scale * value[2] + else: + value[0] *= random.random() * (cfg.x_range[1] - cfg.x_range[0]) + cfg.x_range[0] + value[1] *= random.random() * (cfg.y_range[1] - cfg.y_range[0]) + cfg.y_range[0] + value[2] *= random.random() * (cfg.z_range[1] - cfg.z_range[0]) + cfg.z_range[0] + + # Set property + if not prop.Set(value): + carb.log_error("Failed to set property", prop, "for", prim_path) + return prim diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/randomizations/randomizations_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/randomizations/randomizations_cfg.py new file mode 100644 index 0000000000..256bfb5d57 --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/randomizations/randomizations_cfg.py @@ -0,0 +1,64 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configurations for randomizations that are applied at the time of spawning prims.""" + +from __future__ import annotations + +from dataclasses import MISSING + +from pxr import Usd + +from omni.isaac.lab.utils import configclass + +from .randomizations import randomize_joint_offset, randomize_scale + + +@configclass +class RandomizationCfg: + """Base class for randomization configurations.""" + + func: callable[..., Usd.Prim] = MISSING + + +@configclass +class RandomizeJointOffsetsCfg(RandomizationCfg): + """Randomize the joint offsets of a joint in an articulation + The joint offsets are randomized in the x, y, and z axis by adding random value within the specified range. + """ + + x_range: tuple[float, float] = (0, 0) + """The range of the randomization in the x-axis""" + + y_range: tuple[float, float] = (0, 0) + """The range of the randomization in the y-axis""" + + z_range: tuple[float, float] = (0, 0) + """The range of the randomization in the z-axis""" + + joint_name: str = MISSING + """Joint name to randomize. Note that the joint prim is assumed to be located + at /""" + + func: callable[..., Usd.Prim] = randomize_joint_offset + + +@configclass +class RandomizeScaleCfg(RandomizationCfg): + """Randomize the scale of a prim. + The scale is randomized in the x, y, and z axis by multiplying the current scale by a random value within the specified range. + """ + + x_range: tuple[float, float] = (1, 1) + """The range of the randomization in the x-axis.""" + y_range: tuple[float, float] = (1, 1) + """The range of the randomization in the y-axis.""" + z_range: tuple[float, float] = (1, 1) + """The range of the randomization in the z-axis.""" + + equal_scale: bool = False + """If True, the sampled x_range is used to scale all axes.""" + + func: callable[..., Usd.Prim] = randomize_scale diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/utils.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/utils.py index d7a06de866..7f949fe08e 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/utils.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/utils.py @@ -272,7 +272,12 @@ def wrapper(prim_path: str | Sdf.Path, cfg: SpawnerCfg, *args, **kwargs): if len(prim_paths) > 1: cloner = Cloner() # clone the prim - cloner.clone(prim_paths[0], prim_paths[1:], replicate_physics=False, copy_from_source=cfg.copy_from_source) + cloner.clone( + prim_paths[0], + prim_paths[1:], + replicate_physics=False, + copy_from_source=cfg.copy_from_source if hasattr(cfg, "copy_from_source") else False, + ) # return the source prim return prim diff --git a/source/standalone/demos/multi_object.py b/source/standalone/demos/multi_object.py new file mode 100644 index 0000000000..ba5fcad036 --- /dev/null +++ b/source/standalone/demos/multi_object.py @@ -0,0 +1,194 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""This script demonstrates how to spawn multiple objects in multiple environments. +.. code-block:: bash + # Usage + ./isaaclab.sh -p source/standalone/demos/multi_object.py --num_envs 512 +""" + +from __future__ import annotations + +"""Launch Isaac Sim Simulator first.""" + + +import argparse + +from omni.isaac.lab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Demo on spawning different objects in multiple environments.") +parser.add_argument("--num_envs", type=int, default=64, help="Number of environments to spawn.") +parser.add_argument("--randomize", action="store_true", help="Randomize the objects scale.") + +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli = parser.parse_args() + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import traceback + +import carb + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.assets import AssetBaseCfg, RigidObjectCfg +from omni.isaac.lab.scene import InteractiveScene, InteractiveSceneCfg +from omni.isaac.lab.sim import SimulationContext +from omni.isaac.lab.sim.spawners.multi_asset.asset_randomizer_cfg import AssetRandomizerCfg +from omni.isaac.lab.sim.spawners.multi_asset.multi_asset_cfg import MultiAssetCfg +from omni.isaac.lab.sim.spawners.multi_asset.randomizations import RandomizeScaleCfg +from omni.isaac.lab.utils import configclass + +check_size = True + + +def get_assets(): + kwargs = { + "rigid_props": sim_utils.RigidBodyPropertiesCfg(), + "mass_props": sim_utils.MassPropertiesCfg(mass=0.05), + "collision_props": sim_utils.CollisionPropertiesCfg(), + } + + return [ + sim_utils.SphereCfg( + radius=0.25, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 0.0, 1.0)), + **kwargs, + ), + sim_utils.CuboidCfg( + size=(0.25, 0.25, 0.25), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 1.0, 0.0)), + **kwargs, + ), + sim_utils.CylinderCfg( + radius=0.2, + height=0.3, + axis="Y", + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0)), + **kwargs, + ), + ] + + +def get_randomized_assets(): + assets = get_assets() + + return [ + AssetRandomizerCfg( + child_spawner_cfg=assets[0], + randomization_cfg=RandomizeScaleCfg( + x_range=(0.5, 1.25), + equal_scale=True, + ), + num_random_assets=args_cli.num_envs // 3, + ), + AssetRandomizerCfg( + child_spawner_cfg=assets[1], + randomization_cfg=RandomizeScaleCfg( + x_range=(0.5, 1.25), + equal_scale=True, + ), + num_random_assets=args_cli.num_envs // 3, + ), + AssetRandomizerCfg( + child_spawner_cfg=assets[2], + randomization_cfg=RandomizeScaleCfg( + x_range=(0.5, 1.25), + equal_scale=True, + ), + num_random_assets=args_cli.num_envs // 3, + ), + ] + + +@configclass +class MultiObjectSceneCfg(InteractiveSceneCfg): + """Configuration for a multi-object scene.""" + + # ground plane + ground = AssetBaseCfg(prim_path="/World/defaultGroundPlane", spawn=sim_utils.GroundPlaneCfg()) + + # lights + dome_light = AssetBaseCfg( + prim_path="/World/Light", + spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75)), + ) + + object: RigidObjectCfg = RigidObjectCfg( + prim_path="/World/envs/env_.*/Objects", + spawn=MultiAssetCfg(assets_cfg=get_randomized_assets() if args_cli.randomize else get_assets()), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 0.3)), + ) + + +def run_simulator(sim: SimulationContext, scene: InteractiveScene): + """Runs the simulation loop.""" + # Extract scene entities + rigid_object = scene["object"] + # Define simulation stepping + sim_dt = sim.get_physics_dt() + count = 0 + # Simulation loop + while simulation_app.is_running(): + # Reset + if count % 500 == 0: + # reset counter + count = 0 + # reset the scene entities + # root state + root_state = rigid_object.data.default_root_state.clone() + root_state[:, :3] += scene.env_origins + rigid_object.write_root_state_to_sim(root_state) + # clear internal buffers + scene.reset() + print("[INFO]: Resetting robot state...") + # Write data to sim + scene.write_data_to_sim() + # Perform step + sim.step() + # Increment counter + count += 1 + # Update buffers + scene.update(sim_dt) + + +def main(): + """Main function.""" + # Load kit helper + sim_cfg = sim_utils.SimulationCfg(device="cpu", use_gpu_pipeline=False) + sim_cfg.use_fabric = False + sim_cfg.physx.use_gpu = False + sim = SimulationContext(sim_cfg) + # Set main camera + sim.set_camera_view([3.0, 0.0, 3.0], [0.0, 0.0, 0.0]) + # Design scene + scene_cfg = MultiObjectSceneCfg(num_envs=args_cli.num_envs, env_spacing=1.5, replicate_physics=False) + scene = InteractiveScene(scene_cfg) + + # Play the simulator + sim.reset() + # Now we are ready! + print("[INFO]: Setup complete...") + # Run the simulator + run_simulator(sim, scene) + + +if __name__ == "__main__": + try: + # run the main execution + main() + except Exception as err: + carb.log_error(err) + carb.log_error(traceback.format_exc()) + raise + finally: + # close sim app + simulation_app.close() From fd051443b33a35c26d0886c0901891b7583d9a74 Mon Sep 17 00:00:00 2001 From: Pascal Roth <57946385+pascal-roth@users.noreply.github.com> Date: Sun, 1 Sep 2024 15:27:35 +0200 Subject: [PATCH 2/2] Apply suggestions from code review Co-authored-by: jtigue-bdai <166445701+jtigue-bdai@users.noreply.github.com> Signed-off-by: Pascal Roth <57946385+pascal-roth@users.noreply.github.com> --- .../isaac/lab/sim/spawners/multi_asset/asset_randomizer.py | 4 ++-- .../omni/isaac/lab/sim/spawners/multi_asset/multi_asset.py | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/asset_randomizer.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/asset_randomizer.py index 3a6f2eb12f..bebe42b7fb 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/asset_randomizer.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/asset_randomizer.py @@ -24,8 +24,8 @@ def randomize_and_spawn_asset( Args: prim_path: The path to the asset to spawn cfg: The configuration for the asset randomization - translation: The translation to apply to the spawned asset, gets passed to the child spawner defined in the configuration - orientation: The orientation to apply to the spawned asset, gets passed to the child spawner defined in the configuration + translation: The translation to apply to the prim w.r.t. its parent prim in meters. Defaults to None, in which case this is set to the origin. Gets passed to the child spawner defined in the configuration + orientation: The orientation in (w, x, y, z) to apply to the prim w.r.t. its parent prim. Defaults to None, in which case this is set to identity (1, 0, 0, 0). Gets passed to the child spawner defined in the configuration """ prim = cfg.child_spawner_cfg.func(prim_path, cfg.child_spawner_cfg, translation, orientation) randomizations = cfg.randomization_cfg diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/multi_asset.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/multi_asset.py index 4bee494fd4..353510020d 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/multi_asset.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/multi_asset/multi_asset.py @@ -34,8 +34,9 @@ def spawn_multi_object_randomly_sdf( Args: prim_path: The path to the asset to spawn (e.g. "/World/env_*/Object") cfg: The configuration for the asset spawning - translation: The translation to apply to the spawned asset - orientation: The orientation to apply to the spawned asset + translation: The translation to apply to the prim w.r.t. its parent prim in meters. Defaults to None, in which case this is set to the origin. + orientation: The orientation in (w, x, y, z) to apply to the prim w.r.t. its parent prim. Defaults to None, in which case this is set to identity (1, 0, 0, 0). + """ # resolve: {SPAWN_NS}/AssetName