Add texture randomizations to objects

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:

isaaclab_arena/examples/policy_runner.py

@@ -35,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

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()
+

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(
@@ -100,3 +102,123 @@ def get_asset_usd_path_from_prim_path(prim_path: str, stage: Usd.Stage) -> str |
                 return reference_spec.assetPath
 
     return None
+
+
+def apply_material_variants_to_objects(
+    prim_paths: list[str],
+    stage: Usd.Stage,
+    randomize: bool = True,
+):
+    """
+    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.
+    """
+
+    for path in prim_paths:
+        prim = stage.GetPrimAtPath(path)
+        if not prim.IsValid():
+            print(f"Warning: Prim at path '{path}' does not exist. Skipping.")
+            continue
+
+        # 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)