diff --git a/python/bricksim/cli/assemble_brick_expert.py b/python/bricksim/cli/assemble_brick_expert.py index f6f1006..2709437 100644 --- a/python/bricksim/cli/assemble_brick_expert.py +++ b/python/bricksim/cli/assemble_brick_expert.py @@ -90,6 +90,26 @@ def build_argument_parser() -> argparse.ArgumentParser: action="store_true", help="Show assembly goal marker visualizer.", ) + parser.add_argument( + "--randomize-dome-light", + action="store_true", + help="Randomize dome-light HDR texture, intensity, and yaw on reset.", + ) + parser.add_argument( + "--randomize-table-material", + action="store_true", + help="Randomize table MDL material on reset.", + ) + parser.add_argument( + "--randomize-background-material", + action="store_true", + help="Randomize ground/background MDL material on reset.", + ) + parser.add_argument( + "--randomize-scene-background", + action="store_true", + help="Randomize visual-only factory/kitchen/outdoor background presets.", + ) return parser @@ -177,6 +197,7 @@ def main() -> int: simulation_app = app_launcher.app from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.managers import EventTermCfg from bricksim.envs.assemble_brick.env import ( AssembleBrickBaseEnvCfg, @@ -185,6 +206,12 @@ def main() -> int: AssembleBrickRGBEnvCfg, ) from bricksim.envs.assemble_brick.expert import AssembleBrickExpert + from bricksim.mdp.events import ( + randomize_dome_light_viral, + randomize_global_material_viral, + randomize_scene_background_viral, + randomize_table_material_viral, + ) env_cfg: AssembleBrickBaseEnvCfg if args_cli.obs_mode == "state": @@ -201,6 +228,86 @@ def main() -> int: env_cfg.commands.assembly_goal.goals = args_cli.goals if args_cli.show_goals: env_cfg.commands.assembly_goal.debug_vis = True + if args_cli.randomize_dome_light: + env_cfg.events.randomize_dome_light_viral = EventTermCfg( + func=randomize_dome_light_viral, + mode="reset", + params={ + "prim_path": "/World/light", + "texture_folder_paths": ( + "https://omniverse-content-production.s3-us-west-2.amazonaws.com/" + "Assets/Isaac/4.5/NVIDIA/Assets/Skies/Indoor/", + "https://omniverse-content-production.s3-us-west-2.amazonaws.com/" + "Assets/Isaac/4.5/NVIDIA/Assets/Skies/Clear/", + "https://omniverse-content-production.s3-us-west-2.amazonaws.com/" + "Assets/Isaac/4.5/NVIDIA/Assets/Skies/Cloudy/", + "https://omniverse-content-production.s3-us-west-2.amazonaws.com/" + "Assets/Isaac/4.5/NVIDIA/Assets/Skies/Night/", + "https://omniverse-content-production.s3-us-west-2.amazonaws.com/" + "Assets/Isaac/4.5/NVIDIA/Assets/Skies/Studio/", + "https://omniverse-content-production.s3-us-west-2.amazonaws.com/" + "Assets/Isaac/4.5/NVIDIA/Environments/2024_1/DomeLights/Clear/", + "https://omniverse-content-production.s3-us-west-2.amazonaws.com/" + "Assets/Isaac/4.5/NVIDIA/Environments/2024_1/DomeLights/Cloudy/", + "https://omniverse-content-production.s3-us-west-2.amazonaws.com/" + "Assets/Isaac/4.5/NVIDIA/Environments/2024_1/DomeLights/Indoor/", + ), + "intensity_range": (800.0, 2000.0), + "yaw_range": (-3.142, 3.142), + }, + ) + print("[INFO]: VIRAL-style dome-light randomization enabled") + if args_cli.randomize_table_material: + env_cfg.events.randomize_table_material_viral = EventTermCfg( + func=randomize_table_material_viral, + mode="reset", + params={ + "env_relative_prim_paths": ("Table",), + "material_root_path": "/World/Looks/TableVisualRandomization", + "material_folder_paths": ( + "https://omniverse-content-production.s3-us-west-2.amazonaws.com/" + "Assets/Isaac/4.5/NVIDIA/Materials/Base/Wood/", + ), + "material_name_prefix": "TableMaterial", + }, + ) + print("[INFO]: VIRAL-style table material randomization enabled") + if args_cli.randomize_background_material: + env_cfg.events.randomize_background_material_viral = EventTermCfg( + func=randomize_global_material_viral, + mode="reset", + params={ + "prim_paths": ("/World/GroundPlane",), + "material_root_path": "/World/Looks/BackgroundVisualRandomization", + "material_folder_paths": ( + "https://omniverse-content-production.s3-us-west-2.amazonaws.com/" + "Assets/Isaac/4.5/NVIDIA/Materials/Base/Carpet/", + "https://omniverse-content-production.s3-us-west-2.amazonaws.com/" + "Assets/Isaac/4.5/NVIDIA/Materials/Base/Masonry/", + "https://omniverse-content-production.s3-us-west-2.amazonaws.com/" + "Assets/Isaac/4.5/NVIDIA/Materials/Base/Natural/", + "https://omniverse-content-production.s3-us-west-2.amazonaws.com/" + "Assets/Isaac/4.5/NVIDIA/Materials/Base/Stone/", + "https://omniverse-content-production.s3-us-west-2.amazonaws.com/" + "Assets/Isaac/4.5/NVIDIA/Materials/Base/Wall_Board/", + "https://omniverse-content-production.s3-us-west-2.amazonaws.com/" + "Assets/Isaac/4.5/NVIDIA/Materials/Base/Wood/", + ), + "material_name_prefix": "BackgroundMaterial", + "label": "background", + }, + ) + print("[INFO]: VIRAL-style background material randomization enabled") + if args_cli.randomize_scene_background: + env_cfg.events.randomize_scene_background_viral = EventTermCfg( + func=randomize_scene_background_viral, + mode="reset", + params={ + "root_path": "/World/VisualSceneBackground", + "preset_names": ("factory", "kitchen", "outdoor"), + }, + ) + print("[INFO]: VIRAL-style scene background randomization enabled") env = ManagerBasedRLEnv(cfg=env_cfg) expert = AssembleBrickExpert() diff --git a/python/bricksim/mdp/events.py b/python/bricksim/mdp/events.py index 99d2e14..1c728a3 100644 --- a/python/bricksim/mdp/events.py +++ b/python/bricksim/mdp/events.py @@ -1,8 +1,16 @@ """Event terms for BrickSim-managed Isaac Lab environments.""" +import asyncio +import math +import os + +import omni.client +import omni.kit.commands import torch from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedEnv +from isaaclab.sim import get_current_stage +from pxr import Gf, Sdf, Usd, UsdGeom, UsdShade, UsdLux from bricksim.core import deallocate_all_managed @@ -128,3 +136,483 @@ def reset_bricksim_managed(env: ManagerBasedEnv, env_ids: torch.Tensor) -> None: env_ids = env_ids.detach().cpu() for env_id in env_ids: deallocate_all_managed(env_id.item()) + + +def _cached_dome_light_texture_files( + env: ManagerBasedEnv, + texture_folder_paths: tuple[str, ...], +) -> list[str]: + cache = getattr(env, "_bricksim_dome_light_texture_cache", {}) + cached_texture_files = cache.get(texture_folder_paths) + if cached_texture_files is not None: + return cached_texture_files + + texture_files: list[str] = [] + skipped_folders: list[str] = [] + for texture_folder_path in texture_folder_paths: + result, entries = omni.client.list(texture_folder_path) + if result != omni.client.Result.OK: + skipped_folders.append(f"{texture_folder_path} ({result})") + continue + texture_files.extend( + os.path.join(texture_folder_path, entry.relative_path) + for entry in entries + if entry.relative_path.endswith(".hdr") + ) + + if skipped_folders: + print( + "[WARN]: Skipped unavailable dome-light texture folders: " + + "; ".join(skipped_folders) + ) + + if not texture_files: + raise RuntimeError( + "No .hdr dome-light texture files found in configured folders: " + f"{texture_folder_paths}" + ) + + cache[texture_folder_paths] = texture_files + env._bricksim_dome_light_texture_cache = cache + print(f"[INFO]: Loaded {len(texture_files)} dome-light HDR textures") + return texture_files + + +def randomize_dome_light_viral( + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + prim_path: str, + texture_folder_paths: tuple[str, ...], + intensity_range: tuple[float, float], + yaw_range: tuple[float, float], +) -> None: + """Randomize a DomeLight using VIRAL's asset-folder style. + + HDR texture files are discovered from configured Omniverse folders and + cached on the environment. On each reset, the existing DomeLight prim is + updated in place with a random texture, intensity, and yaw rotation. + """ + del env_ids + texture_files = _cached_dome_light_texture_files(env, texture_folder_paths) + + if not hasattr(env, "_bricksim_cached_dome_light"): + stage = get_current_stage() + prim = stage.GetPrimAtPath(prim_path) + env._bricksim_cached_dome_light = UsdLux.DomeLight(prim) + dome_light = env._bricksim_cached_dome_light + if not dome_light.GetPrim().IsValid(): + return + + texture_id = int( + torch.randint(0, len(texture_files), (1,), device=env.device).item() + ) + texture_file = texture_files[texture_id] + if dome_light.GetTextureFileAttr().Get() != texture_file: + dome_light.GetTextureFileAttr().Set(texture_file) + + intensity = torch.empty((), device=env.device).uniform_( + intensity_range[0], intensity_range[1] + ) + dome_light.GetIntensityAttr().Set(float(intensity.item())) + + yaw = torch.empty((), device=env.device).uniform_(yaw_range[0], yaw_range[1]) + rotation = Gf.Rotation(Gf.Vec3d(0.0, 0.0, 1.0), math.degrees(float(yaw.item()))) + transform = Gf.Matrix4d(1.0) + transform.SetRotate(rotation) + omni.kit.commands.execute( + "TransformPrimCommand", + path=dome_light.GetPath(), + new_transform_matrix=transform, + ) + + + +def _env_ids_to_cpu_list(env: ManagerBasedEnv, env_ids: torch.Tensor | None): + if env_ids is None: + return range(env.num_envs) + if env_ids.device.type != "cpu": + env_ids = env_ids.detach().cpu() + return [int(env_id.item()) for env_id in env_ids] + + +def _list_available_mdl_material_files( + material_folder_paths: tuple[str, ...], +) -> list[str]: + material_files: list[str] = [] + skipped_folders: list[str] = [] + for material_folder_path in material_folder_paths: + result, entries = omni.client.list(material_folder_path) + if result != omni.client.Result.OK: + skipped_folders.append(f"{material_folder_path} ({result})") + continue + material_files.extend( + os.path.join(material_folder_path, entry.relative_path) + for entry in entries + if entry.relative_path.endswith(".mdl") + ) + + if skipped_folders: + print( + "[WARN]: Skipped unavailable MDL material folders: " + + "; ".join(skipped_folders) + ) + return material_files + + +def _cached_mdl_material_prim_paths( + env: ManagerBasedEnv, + material_root_path: str, + material_folder_paths: tuple[str, ...], + material_name_prefix: str, +) -> list[str]: + import omni.kit.material.library + + cache = getattr(env, "_bricksim_mdl_material_cache", {}) + cache_key = (material_root_path, material_folder_paths, material_name_prefix) + cached_material_paths = cache.get(cache_key) + if cached_material_paths is not None: + return cached_material_paths + + stage = get_current_stage() + material_root = Sdf.Path(material_root_path) + stage.DefinePrim(material_root.GetParentPath(), "Scope") + stage.DefinePrim(material_root, "Scope") + + material_files = _list_available_mdl_material_files(material_folder_paths) + if not material_files: + raise RuntimeError( + "No .mdl material files found in configured folders: " + f"{material_folder_paths}" + ) + + material_prim_paths: list[str] = [] + for material_file in material_files: + subidentifiers = asyncio.run( + omni.kit.material.library.get_subidentifier_from_mdl(material_file) + ) + for subidentifier in subidentifiers: + material_path = material_root.AppendChild( + f"{material_name_prefix}_{len(material_prim_paths)}" + ) + success, _ = omni.kit.commands.execute( + "CreateMdlMaterialPrimCommand", + mtl_url=str(material_file), + mtl_name=str(subidentifier), + mtl_path=str(material_path), + ) + if not success: + continue + shader_prim = stage.GetPrimAtPath(material_path.AppendChild("Shader")) + if shader_prim.IsValid(): + project_uvw = shader_prim.GetAttribute("inputs:project_uvw") + if not project_uvw: + project_uvw = shader_prim.CreateAttribute( + "inputs:project_uvw", Sdf.ValueTypeNames.Bool + ) + project_uvw.Set(True) + world_or_object = shader_prim.GetAttribute("inputs:world_or_object") + if not world_or_object: + world_or_object = shader_prim.CreateAttribute( + "inputs:world_or_object", Sdf.ValueTypeNames.Bool + ) + world_or_object.Set(True) + material_prim_paths.append(str(material_path)) + + if not material_prim_paths: + raise RuntimeError( + "Omniverse material folders were readable, but no MDL materials were " + "created." + ) + + cache[cache_key] = material_prim_paths + env._bricksim_mdl_material_cache = cache + print( + f"[INFO]: Loaded {len(material_prim_paths)} MDL materials under " + f"{material_root_path}" + ) + return material_prim_paths + + + +def _bind_material_to_prim_tree( + root_prim: Usd.Prim, + material: UsdShade.Material, + root_prim_path: str, + label: str, +) -> int: + if root_prim.IsInstanceable(): + root_prim.SetInstanceable(False) + + binding_api = UsdShade.MaterialBindingAPI.Apply(root_prim) + binding_api.Bind( + material, + bindingStrength=UsdShade.Tokens.strongerThanDescendants, + ) + + bound_gprim_count = 0 + for prim in Usd.PrimRange(root_prim): + if prim.IsInstanceable(): + prim.SetInstanceable(False) + if prim.IsInstanceProxy(): + continue + if prim.IsA(UsdGeom.Gprim): + UsdShade.MaterialBindingAPI.Apply(prim).Bind( + material, + bindingStrength=UsdShade.Tokens.strongerThanDescendants, + ) + bound_gprim_count += 1 + + if bound_gprim_count == 0: + print(f"[WARN]: No {label} geometry prims found under {root_prim_path}") + else: + print( + f"[INFO]: Bound {label} material {material.GetPath()} to " + f"{bound_gprim_count} geometry prims under {root_prim_path}" + ) + return bound_gprim_count + + +def randomize_global_material_viral( + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + prim_paths: tuple[str, ...], + material_root_path: str, + material_folder_paths: tuple[str, ...], + material_name_prefix: str = "Material", + label: str = "global", +) -> None: + """Bind random MDL materials to global scene prims using VIRAL's reset style.""" + del env_ids + stage = get_current_stage() + material_prim_paths = _cached_mdl_material_prim_paths( + env, + material_root_path, + material_folder_paths, + material_name_prefix, + ) + + for prim_path in prim_paths: + root_prim = stage.GetPrimAtPath(prim_path) + if not root_prim.IsValid(): + print(f"[WARN]: Skipping missing {label} prim: {prim_path}") + continue + + material_id = int( + torch.randint(0, len(material_prim_paths), (1,), device=env.device).item() + ) + material = UsdShade.Material.Get( + stage, Sdf.Path(material_prim_paths[material_id]) + ) + if not material: + raise RuntimeError( + f"Could not resolve material: {material_prim_paths[material_id]}" + ) + _bind_material_to_prim_tree(root_prim, material, prim_path, label) + +def randomize_table_material_viral( + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + env_relative_prim_paths: tuple[str, ...], + material_root_path: str, + material_folder_paths: tuple[str, ...], + material_name_prefix: str = "TableMaterial", +) -> None: + """Bind random MDL materials to table prims using VIRAL's reset style.""" + stage = get_current_stage() + material_prim_paths = _cached_mdl_material_prim_paths( + env, + material_root_path, + material_folder_paths, + material_name_prefix, + ) + + for env_id in _env_ids_to_cpu_list(env, env_ids): + for relative_prim_path in env_relative_prim_paths: + table_prim_path = f"/World/envs/env_{env_id}/{relative_prim_path}" + table_prim = stage.GetPrimAtPath(table_prim_path) + if not table_prim.IsValid(): + print(f"[WARN]: Skipping missing table prim: {table_prim_path}") + continue + + material_id = int( + torch.randint( + 0, len(material_prim_paths), (1,), device=env.device + ).item() + ) + material = UsdShade.Material.Get( + stage, Sdf.Path(material_prim_paths[material_id]) + ) + if not material: + raise RuntimeError( + f"Could not resolve material: {material_prim_paths[material_id]}" + ) + + _bind_material_to_prim_tree(table_prim, material, table_prim_path, "table") + + + +def _define_preview_material( + stage: Usd.Stage, + material_path: str, + color: tuple[float, float, float], + roughness: float = 0.7, +) -> UsdShade.Material: + material = UsdShade.Material.Define(stage, material_path) + shader = UsdShade.Shader.Define(stage, f"{material_path}/PreviewSurface") + shader.CreateIdAttr("UsdPreviewSurface") + shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set( + Gf.Vec3f(*color) + ) + shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(roughness) + shader_output = shader.CreateOutput("surface", Sdf.ValueTypeNames.Token) + material.CreateSurfaceOutput().ConnectToSource(shader_output) + return material + + +def _set_xform( + prim: Usd.Prim, + translate: tuple[float, float, float], + scale: tuple[float, float, float], +) -> None: + xformable = UsdGeom.Xformable(prim) + xformable.ClearXformOpOrder() + xformable.AddTranslateOp().Set(Gf.Vec3d(*translate)) + xformable.AddScaleOp().Set(Gf.Vec3f(*scale)) + + +def _define_colored_cube( + stage: Usd.Stage, + path: str, + translate: tuple[float, float, float], + scale: tuple[float, float, float], + material: UsdShade.Material, +) -> None: + cube = UsdGeom.Cube.Define(stage, path) + cube.CreateSizeAttr(1.0) + _set_xform(cube.GetPrim(), translate, scale) + UsdShade.MaterialBindingAPI.Apply(cube.GetPrim()).Bind(material) + + +def _define_colored_sphere( + stage: Usd.Stage, + path: str, + translate: tuple[float, float, float], + scale: tuple[float, float, float], + material: UsdShade.Material, +) -> None: + sphere = UsdGeom.Sphere.Define(stage, path) + sphere.CreateRadiusAttr(0.5) + _set_xform(sphere.GetPrim(), translate, scale) + UsdShade.MaterialBindingAPI.Apply(sphere.GetPrim()).Bind(material) + + +def _create_factory_background(stage: Usd.Stage, root_path: str) -> None: + root = UsdGeom.Xform.Define(stage, root_path).GetPrim() + metals = _define_preview_material(stage, f"{root_path}/Looks/BrushedMetal", (0.42, 0.45, 0.46), 0.35) + dark = _define_preview_material(stage, f"{root_path}/Looks/DarkPanels", (0.16, 0.18, 0.18), 0.55) + yellow = _define_preview_material(stage, f"{root_path}/Looks/SafetyYellow", (0.92, 0.68, 0.12), 0.45) + crate = _define_preview_material(stage, f"{root_path}/Looks/Cardboard", (0.55, 0.36, 0.18), 0.8) + + _define_colored_cube(stage, f"{root_path}/BackWall", (0.5, 1.65, 0.65), (2.4, 0.04, 1.4), metals) + _define_colored_cube(stage, f"{root_path}/SideWall", (-0.78, 0.35, 0.65), (0.04, 2.4, 1.4), dark) + _define_colored_cube(stage, f"{root_path}/RackLeft", (-0.45, 1.1, 0.45), (0.06, 0.05, 0.9), yellow) + _define_colored_cube(stage, f"{root_path}/RackRight", (1.45, 1.1, 0.45), (0.06, 0.05, 0.9), yellow) + for index, z in enumerate((0.25, 0.55, 0.85)): + _define_colored_cube(stage, f"{root_path}/RackShelf_{index}", (0.5, 1.1, z), (1.9, 0.05, 0.035), yellow) + for index, x in enumerate((-0.25, 0.05, 1.05, 1.35)): + _define_colored_cube(stage, f"{root_path}/Crate_{index}", (x, 1.0, 0.16), (0.22, 0.18, 0.16), crate) + for index, z in enumerate((0.9, 1.08)): + _define_colored_cube(stage, f"{root_path}/Pipe_{index}", (0.5, 1.52, z), (2.2, 0.035, 0.035), dark) + UsdGeom.Imageable(root).MakeInvisible() + + +def _create_kitchen_background(stage: Usd.Stage, root_path: str) -> None: + root = UsdGeom.Xform.Define(stage, root_path).GetPrim() + wall = _define_preview_material(stage, f"{root_path}/Looks/TileWall", (0.82, 0.84, 0.80), 0.4) + cabinet = _define_preview_material(stage, f"{root_path}/Looks/CabinetWood", (0.50, 0.30, 0.14), 0.65) + counter = _define_preview_material(stage, f"{root_path}/Looks/Countertop", (0.18, 0.20, 0.22), 0.25) + steel = _define_preview_material(stage, f"{root_path}/Looks/Steel", (0.62, 0.64, 0.63), 0.2) + + _define_colored_cube(stage, f"{root_path}/TileWall", (0.5, 1.55, 0.65), (2.4, 0.035, 1.4), wall) + for index, x in enumerate((-0.15, 0.35, 0.85, 1.35)): + _define_colored_cube(stage, f"{root_path}/BaseCabinet_{index}", (x, 1.18, 0.25), (0.44, 0.26, 0.5), cabinet) + _define_colored_cube(stage, f"{root_path}/UpperCabinet_{index}", (x, 1.25, 0.95), (0.42, 0.18, 0.32), cabinet) + _define_colored_cube(stage, f"{root_path}/Countertop", (0.6, 1.03, 0.53), (2.1, 0.32, 0.06), counter) + _define_colored_cube(stage, f"{root_path}/Sink", (1.05, 0.88, 0.58), (0.38, 0.20, 0.035), steel) + _define_colored_cube(stage, f"{root_path}/Fridge", (-0.65, 1.18, 0.65), (0.34, 0.28, 1.1), steel) + UsdGeom.Imageable(root).MakeInvisible() + + +def _create_outdoor_background(stage: Usd.Stage, root_path: str) -> None: + root = UsdGeom.Xform.Define(stage, root_path).GetPrim() + grass = _define_preview_material(stage, f"{root_path}/Looks/Grass", (0.20, 0.43, 0.16), 0.85) + trunk = _define_preview_material(stage, f"{root_path}/Looks/TreeTrunk", (0.34, 0.20, 0.10), 0.75) + leaves = _define_preview_material(stage, f"{root_path}/Looks/Leaves", (0.12, 0.35, 0.12), 0.9) + rock = _define_preview_material(stage, f"{root_path}/Looks/Rock", (0.38, 0.38, 0.35), 0.8) + sky = _define_preview_material(stage, f"{root_path}/Looks/SkyPanel", (0.52, 0.70, 0.88), 0.6) + + _define_colored_cube(stage, f"{root_path}/SkyPanel", (0.5, 1.75, 0.85), (2.6, 0.035, 1.6), sky) + _define_colored_cube(stage, f"{root_path}/GrassPatch", (0.5, 0.95, -0.035), (2.5, 1.8, 0.025), grass) + for index, x in enumerate((-0.55, 1.45)): + _define_colored_cube(stage, f"{root_path}/TreeTrunk_{index}", (x, 1.15, 0.35), (0.08, 0.08, 0.7), trunk) + _define_colored_sphere(stage, f"{root_path}/TreeCanopy_{index}", (x, 1.15, 0.92), (0.55, 0.45, 0.38), leaves) + for index, x in enumerate((-0.25, 0.25, 1.05, 1.35)): + _define_colored_sphere(stage, f"{root_path}/Rock_{index}", (x, 1.0, 0.08), (0.18, 0.13, 0.08), rock) + UsdGeom.Imageable(root).MakeInvisible() + + +def _ensure_scene_background_presets( + env: ManagerBasedEnv, + root_path: str, + preset_names: tuple[str, ...], +) -> list[str]: + cache_key = (root_path, preset_names) + cache = getattr(env, "_bricksim_scene_background_cache", {}) + cached_paths = cache.get(cache_key) + if cached_paths is not None: + return cached_paths + + stage = get_current_stage() + stage.DefinePrim(root_path, "Xform") + preset_paths: list[str] = [] + for preset_name in preset_names: + preset_path = f"{root_path}/{preset_name}" + if preset_name == "factory": + _create_factory_background(stage, preset_path) + elif preset_name == "kitchen": + _create_kitchen_background(stage, preset_path) + elif preset_name == "outdoor": + _create_outdoor_background(stage, preset_path) + else: + raise ValueError(f"Unknown scene background preset: {preset_name}") + preset_paths.append(preset_path) + + cache[cache_key] = preset_paths + env._bricksim_scene_background_cache = cache + print(f"[INFO]: Created scene background presets: {', '.join(preset_names)}") + return preset_paths + + +def randomize_scene_background_viral( + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + root_path: str = "/World/VisualSceneBackground", + preset_names: tuple[str, ...] = ("factory", "kitchen", "outdoor"), +) -> None: + """Randomly activate one visual-only scene background preset on reset.""" + del env_ids + stage = get_current_stage() + preset_paths = _ensure_scene_background_presets(env, root_path, preset_names) + preset_id = int(torch.randint(0, len(preset_paths), (1,), device=env.device).item()) + active_path = preset_paths[preset_id] + + for preset_path in preset_paths: + prim = stage.GetPrimAtPath(preset_path) + if not prim.IsValid(): + continue + imageable = UsdGeom.Imageable(prim) + if preset_path == active_path: + imageable.MakeVisible() + else: + imageable.MakeInvisible() + print(f"[INFO]: Activated scene background preset: {active_path}")