From 3c9fce5ebca5f3788ea0face1eee83068c59b209 Mon Sep 17 00:00:00 2001 From: Xinjie Yao Date: Wed, 17 Dec 2025 13:49:55 -0800 Subject: [PATCH 1/2] wip --- isaaclab_arena/examples/policy_runner.py | 8 +++++ isaaclab_arena/utils/usd_helpers.py | 43 ++++++++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/isaaclab_arena/examples/policy_runner.py b/isaaclab_arena/examples/policy_runner.py index 4d42e06d..d6ca1be9 100644 --- a/isaaclab_arena/examples/policy_runner.py +++ b/isaaclab_arena/examples/policy_runner.py @@ -28,6 +28,14 @@ def main(): # Build scene arena_builder = get_arena_builder_from_cli(args_cli) env = arena_builder.make_registered() + from isaaclab_arena.utils.usd_helpers import apply_render_variants_to_scene + from isaaclab.sim.utils import get_current_stage + # convert to list of prim paths by speicfy the env_index + expanded_paths = [ + f"{env.env_ns.format(i)}/cracker_box" + for i in range(args_cli.num_envs) + ] + apply_render_variants_to_scene(env, prim_paths=expanded_paths, stage=get_current_stage()) if args_cli.seed is not None: env.seed(args_cli.seed) diff --git a/isaaclab_arena/utils/usd_helpers.py b/isaaclab_arena/utils/usd_helpers.py index bbf788dd..41fdc2ee 100644 --- a/isaaclab_arena/utils/usd_helpers.py +++ b/isaaclab_arena/utils/usd_helpers.py @@ -100,3 +100,46 @@ def get_asset_usd_path_from_prim_path(prim_path: str, stage: Usd.Stage) -> str | return reference_spec.assetPath return None + +from pxr import Usd, Sdf, UsdShade + +def apply_render_variants_to_scene(env, prim_paths: list[str], stage: Usd.Stage): + """ + Injected after the scene is initialized in an Isaac Lab environment. + """ + + # Identify the asset you want to 'wrap' in variants + + # Get the prim paths for all environments via Isaac Lab's scene entity + # This automatically handles the {ENV_REGEX_NS} for you + + for path in prim_paths: + print(f"Applying render variants to {path}") + prim = stage.GetPrimAtPath(path) + if not prim.IsValid(): + continue + + # --- INJECTION START --- + # 1. Create the VariantSet on the live Stage + vset = prim.GetVariantSets().AddVariantSet("bringup_style") + + # 2. Author 'Normal' vs 'Debug' rendering styles + # Standard Style + vset.AddVariant("standard") + vset.SetVariantSelection("standard") + # (Standard keeps the USD's original look) + + # 3. Create a 'Synthetic' style (e.g., for depth/segmentation bring-up) + vset.AddVariant("synthetic_check") + vset.SetVariantSelection("synthetic_check") + with vset.GetVariantEditContext(): + # Force a specific display color to check segmentation + attr = prim.CreateAttribute("primvars:displayColor", Sdf.ValueTypeNames.Color3fArray) + attr.Set([(0.0, 1.0, 1.0)]) # Cyan + + # You can also toggle visibility/purpose here + prim.CreateAttribute("purpose", Sdf.ValueTypeNames.Token).Set("guide") + + # Finalize: Default everything to standard + for path in prim_paths: + stage.GetPrimAtPath(path).GetVariantSets().GetVariantSet("bringup_style").SetVariantSelection("standard") \ No newline at end of file From 8bf85fdd113d26e27a3351e6ff63b55dd02e6213 Mon Sep 17 00:00:00 2001 From: Xinjie Yao Date: Thu, 18 Dec 2025 00:14:07 -0800 Subject: [PATCH 2/2] add texture post-spawn --- isaaclab_arena/cli/isaaclab_arena_cli.py | 7 ++ isaaclab_arena/examples/policy_runner.py | 21 ++-- isaaclab_arena/tests/test_usd_helpers.py | 74 +++++++++++ isaaclab_arena/utils/usd_helpers.py | 149 +++++++++++++++++------ 4 files changed, 208 insertions(+), 43 deletions(-) create mode 100644 isaaclab_arena/tests/test_usd_helpers.py diff --git a/isaaclab_arena/cli/isaaclab_arena_cli.py b/isaaclab_arena/cli/isaaclab_arena_cli.py index 3cc11fbf..e9acd511 100644 --- a/isaaclab_arena/cli/isaaclab_arena_cli.py +++ b/isaaclab_arena/cli/isaaclab_arena_cli.py @@ -39,6 +39,13 @@ def add_isaac_lab_cli_args(parser: argparse.ArgumentParser) -> None: help="Disable Pinocchio.", ) isaac_lab_group.add_argument("--mimic", action="store_true", default=False, help="Enable mimic environment.") + isaac_lab_group.add_argument( + "--randomize_object_texture_names", + type=str, + nargs="+", + default=[], + help="List of object names to randomize texture of.", + ) def add_external_environments_cli_args(parser: argparse.ArgumentParser) -> None: diff --git a/isaaclab_arena/examples/policy_runner.py b/isaaclab_arena/examples/policy_runner.py index d6ca1be9..5482080d 100644 --- a/isaaclab_arena/examples/policy_runner.py +++ b/isaaclab_arena/examples/policy_runner.py @@ -28,14 +28,6 @@ def main(): # Build scene arena_builder = get_arena_builder_from_cli(args_cli) env = arena_builder.make_registered() - from isaaclab_arena.utils.usd_helpers import apply_render_variants_to_scene - from isaaclab.sim.utils import get_current_stage - # convert to list of prim paths by speicfy the env_index - expanded_paths = [ - f"{env.env_ns.format(i)}/cracker_box" - for i in range(args_cli.num_envs) - ] - apply_render_variants_to_scene(env, prim_paths=expanded_paths, stage=get_current_stage()) if args_cli.seed is not None: env.seed(args_cli.seed) @@ -43,6 +35,19 @@ def main(): np.random.seed(args_cli.seed) random.seed(args_cli.seed) + # Post-spawn injection + if args_cli.randomize_object_texture_names is not None and len(args_cli.randomize_object_texture_names) > 0: + from isaaclab.sim.utils import get_current_stage + + from isaaclab_arena.utils.usd_helpers import randomize_objects_texture + + randomize_objects_texture( + object_names=args_cli.randomize_object_texture_names, + num_envs=args_cli.num_envs, + env_ns=env.scene.env_ns, + stage=get_current_stage(), + ) + obs, _ = env.reset() # NOTE(xinjieyao, 2025-09-29): General rule of thumb is to have as many non-standard python diff --git a/isaaclab_arena/tests/test_usd_helpers.py b/isaaclab_arena/tests/test_usd_helpers.py new file mode 100644 index 00000000..3ad756e7 --- /dev/null +++ b/isaaclab_arena/tests/test_usd_helpers.py @@ -0,0 +1,74 @@ +# Copyright (c) 2025, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function + +HEADLESS = True + + +def _test_apply_material_variants_to_objects(simulation_app) -> bool: + """Test applying UsdPreviewSurface materials to objects with randomization.""" + from pxr import Usd, UsdShade + + from isaaclab_arena.assets.asset_registry import AssetRegistry + from isaaclab_arena.utils.usd_helpers import apply_material_variants_to_objects + + stage = Usd.Stage.CreateInMemory() + + root = stage.DefinePrim("/World", "Xform") + stage.SetDefaultPrim(root) + + # Get asset registry and reference two cracker boxes + asset_registry = AssetRegistry() + cracker_box = asset_registry.get_asset_by_name("cracker_box")() + + # Create two cracker box prims by referencing the USD + box1_prim = stage.DefinePrim("/World/cracker_box_1", "Xform") + box1_prim.GetReferences().AddReference(cracker_box.usd_path) + + box2_prim = stage.DefinePrim("/World/cracker_box_2", "Xform") + box2_prim.GetReferences().AddReference(cracker_box.usd_path) + + # Apply randomized materials + prim_paths = ["/World/cracker_box_1", "/World/cracker_box_2"] + apply_material_variants_to_objects( + prim_paths=prim_paths, + stage=stage, + randomize=True, + ) + + # Verify materials were created under each object's prim path + material_paths = [ + "/World/cracker_box_1/MaterialVariants", + "/World/cracker_box_2/MaterialVariants", + ] + for material_path in material_paths: + material_prim = stage.GetPrimAtPath(material_path) + assert material_prim.IsValid(), f"Material prim not created at {material_path}" + + # Verify shader has UsdPreviewSurface ID + shader = UsdShade.Shader.Get(stage, f"{material_path}/Shader") + shader_id = shader.GetIdAttr().Get() + assert shader_id == "UsdPreviewSurface", f"Shader ID is {shader_id}, expected 'UsdPreviewSurface'" + + # Verify shader inputs exist + assert shader.GetInput("diffuseColor"), "diffuseColor not found" + assert shader.GetInput("roughness"), "roughness not found" + assert shader.GetInput("metallic"), "metallic not found" + + return True + + +def test_apply_material_variants_to_objects(): + result = run_simulation_app_function( + _test_apply_material_variants_to_objects, + headless=HEADLESS, + ) + assert result, "Test failed" + + +if __name__ == "__main__": + test_apply_material_variants_to_objects() + diff --git a/isaaclab_arena/utils/usd_helpers.py b/isaaclab_arena/utils/usd_helpers.py index 41fdc2ee..e60aff2b 100644 --- a/isaaclab_arena/utils/usd_helpers.py +++ b/isaaclab_arena/utils/usd_helpers.py @@ -3,9 +3,11 @@ # # SPDX-License-Identifier: Apache-2.0 +import colorsys +import random from contextlib import contextmanager -from pxr import Usd, UsdLux, UsdPhysics +from pxr import Gf, Sdf, Usd, UsdGeom, UsdLux, UsdPhysics, UsdShade def get_all_prims( @@ -101,45 +103,122 @@ def get_asset_usd_path_from_prim_path(prim_path: str, stage: Usd.Stage) -> str | return None -from pxr import Usd, Sdf, UsdShade -def apply_render_variants_to_scene(env, prim_paths: list[str], stage: Usd.Stage): +def apply_material_variants_to_objects( + prim_paths: list[str], + stage: Usd.Stage, + randomize: bool = True, +): """ - Injected after the scene is initialized in an Isaac Lab environment. + Apply UsdPreviewSurface materials to objects with optional randomization. + Uses standard USD shaders for maximum compatibility. + + Args: + prim_paths: List of USD prim paths to apply material to. + stage: The USD stage + randomize: If True, randomizes color, roughness, and metallic for each prim. Otherwise, uses default values. """ - - # Identify the asset you want to 'wrap' in variants - # Get the prim paths for all environments via Isaac Lab's scene entity - # This automatically handles the {ENV_REGEX_NS} for you - for path in prim_paths: - print(f"Applying render variants to {path}") prim = stage.GetPrimAtPath(path) if not prim.IsValid(): + print(f"Warning: Prim at path '{path}' does not exist. Skipping.") continue - - # --- INJECTION START --- - # 1. Create the VariantSet on the live Stage - vset = prim.GetVariantSets().AddVariantSet("bringup_style") - - # 2. Author 'Normal' vs 'Debug' rendering styles - # Standard Style - vset.AddVariant("standard") - vset.SetVariantSelection("standard") - # (Standard keeps the USD's original look) - - # 3. Create a 'Synthetic' style (e.g., for depth/segmentation bring-up) - vset.AddVariant("synthetic_check") - vset.SetVariantSelection("synthetic_check") - with vset.GetVariantEditContext(): - # Force a specific display color to check segmentation - attr = prim.CreateAttribute("primvars:displayColor", Sdf.ValueTypeNames.Color3fArray) - attr.Set([(0.0, 1.0, 1.0)]) # Cyan - - # You can also toggle visibility/purpose here - prim.CreateAttribute("purpose", Sdf.ValueTypeNames.Token).Set("guide") - - # Finalize: Default everything to standard - for path in prim_paths: - stage.GetPrimAtPath(path).GetVariantSets().GetVariantSet("bringup_style").SetVariantSelection("standard") \ No newline at end of file + + # Generate material properties + if randomize: + hue = random.random() + saturation = random.random() + value = random.random() + rgb = colorsys.hsv_to_rgb(hue, saturation, value) + mat_color = Gf.Vec3f(rgb[0], rgb[1], rgb[2]) + # roughness is a float between 0 and 1, 0 is smooth, 1 is rough + mat_roughness = random.choice([random.uniform(0.1, 0.3), random.uniform(0.7, 1.0)]) + # metallic is a float between 0 and 1, 0 is dielectric, 1 is metal + mat_metallic = random.choice([0.0, random.uniform(0.8, 1.0)]) + else: + mat_color = Gf.Vec3f(0.0, 1.0, 1.0) + mat_roughness = 0.5 + mat_metallic = 0.0 + + # Create and bind material for this prim + material_path = create_usdpreviewsurface_material(stage, prim.GetPath(), mat_color, mat_roughness, mat_metallic) + bind_material_to_object(prim, material_path, stage) + + +def create_usdpreviewsurface_material( + stage: Usd.Stage, prim_path: Sdf.Path, color: Gf.Vec3f, roughness: float, metallic: float +) -> str: + """ + Create a UsdPreviewSurface material with specified properties under the object's prim path. + + Args: + stage: The USD stage + prim_path: Path of the prim this material will be bound to + color: Diffuse color (RGB, 0-1 range) + roughness: Reflection roughness (0-1) + metallic: Metallic value (0-1) + + Returns: + The material path as string + """ + # Create material under the object's prim path + material_path = f"{str(prim_path)}/MaterialVariants" + + # Always create a new material (or update if exists) + material = UsdShade.Material.Define(stage, material_path) + shader_path = f"{material_path}/Shader" + shader = UsdShade.Shader.Define(stage, shader_path) + + shader.CreateIdAttr("UsdPreviewSurface") + + shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set(color) + shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(roughness) + shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(metallic) + + # Set opacity to fully opaque + shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(1.0) + + # Connect shader output to material surface + shader_output = shader.CreateOutput("surface", Sdf.ValueTypeNames.Token) + material.CreateSurfaceOutput().ConnectToSource(shader_output) + + print( + f"Created UsdPreviewSurface material at {material_path} (color: {color}, roughness: {roughness:.2f}, metallic:" + f" {metallic:.2f})" + ) + + return material_path + + +def bind_material_to_object(prim: Usd.Prim, material_path: str, stage: Usd.Stage): + """ + Recursively bind a material to an object and all its children. + + Args: + prim: The object to bind the material to + material_path: USD path to the material to bind + stage: The USD stage + """ + if prim.IsA(UsdGeom.Mesh): + # Bind the material to this object with strong binding + binding_api = UsdShade.MaterialBindingAPI.Apply(prim) + material = UsdShade.Material(stage.GetPrimAtPath(material_path)) + + # Unbind any existing material first + binding_api.UnbindAllBindings() + + # Note (xinjieyao, 2025.12.17): Bind with "strongerThanDescendants" strength to override child materials + binding_api.Bind(material, bindingStrength=UsdShade.Tokens.strongerThanDescendants) + print(f"Bound material (strong) to mesh: {prim.GetPath()}") + + # Recursively apply to children + for child in prim.GetChildren(): + bind_material_to_object(child, material_path, stage) + + +def randomize_objects_texture(object_names: list[str], num_envs: int, env_ns: str, stage: Usd.Stage): + assert object_names is not None and len(object_names) > 0 + for object_name in object_names: + expanded_paths = [f"{env_ns}/env_{i}/{object_name}" for i in range(num_envs)] + apply_material_variants_to_objects(prim_paths=expanded_paths, stage=stage, randomize=True)