Randomization Snippets

Examples of randomization using USD and Isaac Sim APIs. These examples demonstrate how to randomize scenes for synthetic data generation (SDG) in scenarios where default replicator randomizers are not sufficient or applicable.

The snippets are designed to align with the structure and function names used in the replicator example snippets. In comparison they also have the option to write the data to disk by stetting write_data=True.

Prerequisites:

  • Familiarity with USD.

  • Ability to execute code from the Script Editor.

  • Understanding basic replicator concepts, such as subframes.

Randomizing Light Sources

This snippet sets up a new environment containing a cube and a sphere; it then spawns a given number of lights and randomizes selected attributes for these lights over a specified number of frames.

../_images/isaac_tutorial_replicator_randomization_lights.gif
Randomizing Light Sources
 1import asyncio
 2import os
 3
 4import numpy as np
 5import omni.kit.commands
 6import omni.replicator.core as rep
 7import omni.usd
 8from omni.isaac.core.utils.semantics import add_update_semantics
 9from pxr import Gf, Sdf, UsdGeom
10
11omni.usd.get_context().new_stage()
12stage = omni.usd.get_context().get_stage()
13
14sphere = stage.DefinePrim("/World/Sphere", "Sphere")
15UsdGeom.Xformable(sphere).AddTranslateOp().Set((0.0, 1.0, 1.0))
16add_update_semantics(sphere, "sphere", "class")
17
18cube = stage.DefinePrim("/World/Cube", "Cube")
19UsdGeom.Xformable(cube).AddTranslateOp().Set((0.0, -2.0, 2.0))
20add_update_semantics(cube, "cube", "class")
21
22plane_path = "/World/Plane"
23omni.kit.commands.execute("CreateMeshPrimWithDefaultXform", prim_path=plane_path, prim_type="Plane")
24plane_prim = stage.GetPrimAtPath(plane_path)
25plane_prim.CreateAttribute("xformOp:scale", Sdf.ValueTypeNames.Double3, False).Set(Gf.Vec3d(10, 10, 1))
26
27
28def sphere_lights(num):
29    lights = []
30    for i in range(num):
31        # "CylinderLight", "DiskLight", "DistantLight", "DomeLight", "RectLight", "SphereLight"
32        prim_type = "SphereLight"
33        next_free_path = omni.usd.get_stage_next_free_path(stage, f"/World/{prim_type}", False)
34        light_prim = stage.DefinePrim(next_free_path, prim_type)
35        UsdGeom.Xformable(light_prim).AddTranslateOp().Set((0.0, 0.0, 0.0))
36        UsdGeom.Xformable(light_prim).AddRotateXYZOp().Set((0.0, 0.0, 0.0))
37        UsdGeom.Xformable(light_prim).AddScaleOp().Set((1.0, 1.0, 1.0))
38        light_prim.CreateAttribute("inputs:enableColorTemperature", Sdf.ValueTypeNames.Bool).Set(True)
39        light_prim.CreateAttribute("inputs:colorTemperature", Sdf.ValueTypeNames.Float).Set(6500.0)
40        light_prim.CreateAttribute("inputs:radius", Sdf.ValueTypeNames.Float).Set(0.5)
41        light_prim.CreateAttribute("inputs:intensity", Sdf.ValueTypeNames.Float).Set(30000.0)
42        light_prim.CreateAttribute("inputs:color", Sdf.ValueTypeNames.Color3f).Set((1.0, 1.0, 1.0))
43        light_prim.CreateAttribute("inputs:exposure", Sdf.ValueTypeNames.Float).Set(0.0)
44        light_prim.CreateAttribute("inputs:diffuse", Sdf.ValueTypeNames.Float).Set(1.0)
45        light_prim.CreateAttribute("inputs:specular", Sdf.ValueTypeNames.Float).Set(1.0)
46        lights.append(light_prim)
47    return lights
48
49
50async def run_randomizations_async(num_frames, lights, write_data=True, delay=0):
51    if write_data:
52        writer = rep.WriterRegistry.get("BasicWriter")
53        out_dir = os.getcwd() + "/_out_rand_lights"
54        print(f"Writing data to {out_dir}..")
55        writer.initialize(output_dir=out_dir, rgb=True)
56        rp = rep.create.render_product("/OmniverseKit_Persp", (512, 512))
57        writer.attach(rp)
58
59    for _ in range(num_frames):
60        for light in lights:
61            light.GetAttribute("xformOp:translate").Set(
62                (np.random.uniform(-5, 5), np.random.uniform(-5, 5), np.random.uniform(4, 6))
63            )
64            scale_rand = np.random.uniform(0.5, 1.5)
65            light.GetAttribute("xformOp:scale").Set((scale_rand, scale_rand, scale_rand))
66            light.GetAttribute("inputs:colorTemperature").Set(np.random.normal(4500, 1500))
67            light.GetAttribute("inputs:intensity").Set(np.random.normal(25000, 5000))
68            light.GetAttribute("inputs:color").Set(
69                (np.random.uniform(0.1, 0.9), np.random.uniform(0.1, 0.9), np.random.uniform(0.1, 0.9))
70            )
71
72        if write_data:
73            await rep.orchestrator.step_async(rt_subframes=16)
74        else:
75            await omni.kit.app.get_app().next_update_async()
76        if delay > 0:
77            await asyncio.sleep(delay)
78
79
80num_frames = 10
81lights = sphere_lights(10)
82asyncio.ensure_future(run_randomizations_async(num_frames=num_frames, lights=lights, delay=0.2))

Randomizing Textures

The snippet sets up an environment, spawns a given number of cubes and spheres, and randomizes their textures for the given number of frames. After the randomizations their original materials are reassigned. The snippet also showcases how to create a new material and assign it to a prim.

../_images/isaac_tutorial_replicator_randomization_textures.gif
Randomizing Textures
  1import asyncio
  2import os
  3
  4import numpy as np
  5import omni.replicator.core as rep
  6import omni.usd
  7from omni.isaac.nucleus import get_assets_root_path
  8from omni.isaac.core.utils.semantics import add_update_semantics, get_semantics
  9from pxr import Gf, Sdf, UsdGeom, UsdShade
 10
 11omni.usd.get_context().new_stage()
 12stage = omni.usd.get_context().get_stage()
 13dome_light = stage.DefinePrim("/World/DomeLight", "DomeLight")
 14dome_light.CreateAttribute("inputs:intensity", Sdf.ValueTypeNames.Float).Set(1000.0)
 15
 16sphere = stage.DefinePrim("/World/Sphere", "Sphere")
 17UsdGeom.Xformable(sphere).AddTranslateOp().Set((0.0, 0.0, 1.0))
 18add_update_semantics(sphere, "sphere", "class")
 19
 20num_cubes = 10
 21for _ in range(num_cubes):
 22    prim_type = "Cube"
 23    next_free_path = omni.usd.get_stage_next_free_path(stage, f"/World/{prim_type}", False)
 24    cube = stage.DefinePrim(next_free_path, prim_type)
 25    UsdGeom.Xformable(cube).AddTranslateOp().Set((np.random.uniform(-3.5, 3.5), np.random.uniform(-3.5, 3.5), 1))
 26    scale_rand = np.random.uniform(0.25, 0.5)
 27    UsdGeom.Xformable(cube).AddScaleOp().Set((scale_rand, scale_rand, scale_rand))
 28    add_update_semantics(cube, "cube", "class")
 29
 30plane_path = "/World/Plane"
 31omni.kit.commands.execute("CreateMeshPrimWithDefaultXform", prim_path=plane_path, prim_type="Plane")
 32plane_prim = stage.GetPrimAtPath(plane_path)
 33plane_prim.CreateAttribute("xformOp:scale", Sdf.ValueTypeNames.Double3, False).Set(Gf.Vec3d(10, 10, 1))
 34
 35
 36def get_shapes():
 37    stage = omni.usd.get_context().get_stage()
 38    shapes = []
 39    for prim in stage.Traverse():
 40        sem_dict = get_semantics(prim)
 41        sem_values = sem_dict.values()
 42        if ("class", "cube") in sem_values or ("class", "sphere") in sem_values:
 43            shapes.append(prim)
 44    return shapes
 45
 46
 47shapes = get_shapes()
 48
 49
 50def create_omnipbr_material(mtl_url, mtl_name, mtl_path):
 51    stage = omni.usd.get_context().get_stage()
 52    omni.kit.commands.execute("CreateMdlMaterialPrim", mtl_url=mtl_url, mtl_name=mtl_name, mtl_path=mtl_path)
 53    material_prim = stage.GetPrimAtPath(mtl_path)
 54    shader = UsdShade.Shader(omni.usd.get_shader_from_material(material_prim, get_prim=True))
 55
 56    # Add value inputs
 57    shader.CreateInput("diffuse_color_constant", Sdf.ValueTypeNames.Color3f)
 58    shader.CreateInput("reflection_roughness_constant", Sdf.ValueTypeNames.Float)
 59    shader.CreateInput("metallic_constant", Sdf.ValueTypeNames.Float)
 60
 61    # Add texture inputs
 62    shader.CreateInput("diffuse_texture", Sdf.ValueTypeNames.Asset)
 63    shader.CreateInput("reflectionroughness_texture", Sdf.ValueTypeNames.Asset)
 64    shader.CreateInput("metallic_texture", Sdf.ValueTypeNames.Asset)
 65
 66    # Add other attributes
 67    shader.CreateInput("project_uvw", Sdf.ValueTypeNames.Bool)
 68
 69    # Add texture scale and rotate
 70    shader.CreateInput("texture_scale", Sdf.ValueTypeNames.Float2)
 71    shader.CreateInput("texture_rotate", Sdf.ValueTypeNames.Float)
 72
 73    material = UsdShade.Material(material_prim)
 74    return material
 75
 76
 77def create_materials(num):
 78    MDL = "OmniPBR.mdl"
 79    mtl_name, _ = os.path.splitext(MDL)
 80    MAT_PATH = "/World/Looks"
 81    materials = []
 82    for _ in range(num):
 83        prim_path = omni.usd.get_stage_next_free_path(stage, f"{MAT_PATH}/{mtl_name}", False)
 84        mat = create_omnipbr_material(mtl_url=MDL, mtl_name=mtl_name, mtl_path=prim_path)
 85        materials.append(mat)
 86    return materials
 87
 88
 89materials = create_materials(len(shapes))
 90
 91
 92async def run_randomizations_async(num_frames, materials, textures, write_data=True, delay=0):
 93    if write_data:
 94        writer = rep.WriterRegistry.get("BasicWriter")
 95        out_dir = os.getcwd() + "/_out_rand_textures"
 96        print(f"Writing data to {out_dir}..")
 97        writer.initialize(output_dir=out_dir, rgb=True)
 98        rp = rep.create.render_product("/OmniverseKit_Persp", (512, 512))
 99        writer.attach(rp)
100
101    # Apply the new materials and store the initial ones to reassign later
102    initial_materials = {}
103    for i, shape in enumerate(shapes):
104        cur_mat, _ = UsdShade.MaterialBindingAPI(shape).ComputeBoundMaterial()
105        initial_materials[shape] = cur_mat
106        UsdShade.MaterialBindingAPI(shape).Bind(materials[i], UsdShade.Tokens.strongerThanDescendants)
107
108    for _ in range(num_frames):
109        for mat in materials:
110            shader = UsdShade.Shader(omni.usd.get_shader_from_material(mat, get_prim=True))
111            diffuse_texture = np.random.choice(textures)
112            shader.GetInput("diffuse_texture").Set(diffuse_texture)
113            project_uvw = np.random.choice([True, False], p=[0.9, 0.1])
114            shader.GetInput("project_uvw").Set(bool(project_uvw))
115            texture_scale = np.random.uniform(0.1, 1)
116            shader.GetInput("texture_scale").Set((texture_scale, texture_scale))
117            texture_rotate = np.random.uniform(0, 45)
118            shader.GetInput("texture_rotate").Set(texture_rotate)
119
120        if write_data:
121            await rep.orchestrator.step_async(rt_subframes=4)
122        else:
123            await omni.kit.app.get_app().next_update_async()
124        if delay > 0:
125            await asyncio.sleep(delay)
126
127    # Reassign the initial materials
128    for shape, mat in initial_materials.items():
129        if mat:
130            UsdShade.MaterialBindingAPI(shape).Bind(mat, UsdShade.Tokens.strongerThanDescendants)
131        else:
132            UsdShade.MaterialBindingAPI(shape).UnbindAllBindings()
133
134
135assets_root_path = get_assets_root_path()
136textures = [
137    assets_root_path + "/NVIDIA/Materials/vMaterials_2/Ground/textures/aggregate_exposed_diff.jpg",
138    assets_root_path + "/NVIDIA/Materials/vMaterials_2/Ground/textures/gravel_track_ballast_diff.jpg",
139    assets_root_path + "/NVIDIA/Materials/vMaterials_2/Ground/textures/gravel_track_ballast_multi_R_rough_G_ao.jpg",
140    assets_root_path + "/NVIDIA/Materials/vMaterials_2/Ground/textures/rough_gravel_rough.jpg",
141]
142
143num_frames = 10
144asyncio.ensure_future(run_randomizations_async(num_frames, materials, textures, delay=0.2))

Sequential Randomizations

The snippet provides an example of more complex randomizations, where the results of the first randomization are used to determine the next randomization. It uses a custom sampler function to set the location of the camera by iterating over (almost) equidistant points on a sphere. The snippet starts by setting up the environment, a forklift, a pallet, a bin, and a dome light. For every randomization frame, it cycles through the dome light textures, moves the pallet to a random location, and then moves the bin so that it is fully on top of the pallet. Finally, it moves the camera to a new location on the sphere, ensuring it faces the bin.

../_images/isaac_tutorial_replicator_randomization_chained_persp.gif ../_images/isaac_tutorial_replicator_randomization_chained_sphere.gif
Sequential Randomizations
  1import asyncio
  2import itertools
  3import os
  4
  5import numpy as np
  6import omni.replicator.core as rep
  7import omni.usd
  8from omni.isaac.nucleus import get_assets_root_path
  9from pxr import Gf, Usd, UsdGeom, UsdLux
 10
 11
 12# https://stackoverflow.com/questions/9600801/evenly-distributing-n-points-on-a-sphere
 13# https://arxiv.org/pdf/0912.4540.pdf
 14def next_point_on_sphere(idx, num_points, radius=1, origin=(0, 0, 0)):
 15    offset = 2.0 / num_points
 16    inc = np.pi * (3.0 - np.sqrt(5.0))
 17    z = ((idx * offset) - 1) + (offset / 2)
 18    phi = ((idx + 1) % num_points) * inc
 19    r = np.sqrt(1 - pow(z, 2))
 20    y = np.cos(phi) * r
 21    x = np.sin(phi) * r
 22    return [(x * radius) + origin[0], (y * radius) + origin[1], (z * radius) + origin[2]]
 23
 24
 25assets_root_path = get_assets_root_path()
 26FORKLIFT_PATH = assets_root_path + "/Isaac/Props/Forklift/forklift.usd"
 27PALLET_PATH = assets_root_path + "/Isaac/Props/Pallet/pallet.usd"
 28BIN_PATH = assets_root_path + "/Isaac/Props/KLT_Bin/small_KLT_visual.usd"
 29
 30omni.usd.get_context().new_stage()
 31stage = omni.usd.get_context().get_stage()
 32
 33dome_light = UsdLux.DomeLight.Define(stage, "/World/Lights/DomeLight")
 34dome_light.GetIntensityAttr().Set(1000)
 35
 36forklift_prim = stage.DefinePrim("/World/Forklift", "Xform")
 37forklift_prim.GetReferences().AddReference(FORKLIFT_PATH)
 38if not forklift_prim.GetAttribute("xformOp:translate"):
 39    UsdGeom.Xformable(forklift_prim).AddTranslateOp()
 40forklift_prim.GetAttribute("xformOp:translate").Set((-4.5, -4.5, 0))
 41
 42pallet_prim = stage.DefinePrim("/World/Pallet", "Xform")
 43pallet_prim.GetReferences().AddReference(PALLET_PATH)
 44if not pallet_prim.GetAttribute("xformOp:translate"):
 45    UsdGeom.Xformable(pallet_prim).AddTranslateOp()
 46if not pallet_prim.GetAttribute("xformOp:rotateXYZ"):
 47    UsdGeom.Xformable(pallet_prim).AddRotateXYZOp()
 48
 49bin_prim = stage.DefinePrim("/World/Bin", "Xform")
 50bin_prim.GetReferences().AddReference(BIN_PATH)
 51if not bin_prim.GetAttribute("xformOp:translate"):
 52    UsdGeom.Xformable(bin_prim).AddTranslateOp()
 53if not bin_prim.GetAttribute("xformOp:rotateXYZ"):
 54    UsdGeom.Xformable(bin_prim).AddRotateXYZOp()
 55
 56cam = stage.DefinePrim("/World/Camera", "Camera")
 57if not cam.GetAttribute("xformOp:translate"):
 58    UsdGeom.Xformable(cam).AddTranslateOp()
 59if not cam.GetAttribute("xformOp:orient"):
 60    UsdGeom.Xformable(cam).AddOrientOp()
 61
 62
 63async def run_randomizations_async(
 64    num_frames, dome_light, dome_textures, pallet_prim, bin_prim, write_data=True, delay=0
 65):
 66    if write_data:
 67        writer = rep.WriterRegistry.get("BasicWriter")
 68        out_dir = os.getcwd() + "/_out_rand_sphere_scan"
 69        print(f"Writing data to {out_dir}..")
 70        writer.initialize(output_dir=out_dir, rgb=True)
 71        rp_persp = rep.create.render_product("/OmniverseKit_Persp", (512, 512), name="PerspView")
 72        rp_cam = rep.create.render_product(str(cam.GetPath()), (512, 512), name="SphereView")
 73        writer.attach([rp_cam, rp_persp])
 74
 75    textures_cycle = itertools.cycle(dome_textures)
 76
 77    bb_cache = UsdGeom.BBoxCache(time=Usd.TimeCode.Default(), includedPurposes=[UsdGeom.Tokens.default_])
 78    pallet_size = bb_cache.ComputeWorldBound(pallet_prim).GetRange().GetSize()
 79    pallet_length = pallet_size.GetLength()
 80    bin_size = bb_cache.ComputeWorldBound(bin_prim).GetRange().GetSize()
 81
 82    for i in range(num_frames):
 83        # Set next background texture every nth frame and run an app update
 84        if i % 5 == 0:
 85            dome_light.GetTextureFileAttr().Set(next(textures_cycle))
 86            await omni.kit.app.get_app().next_update_async()
 87
 88        # Randomize pallet pose
 89        pallet_prim.GetAttribute("xformOp:translate").Set(
 90            Gf.Vec3d(np.random.uniform(-1.5, 1.5), np.random.uniform(-1.5, 1.5), 0)
 91        )
 92        rand_z_rot = np.random.uniform(-90, 90)
 93        pallet_prim.GetAttribute("xformOp:rotateXYZ").Set(Gf.Vec3d(0, 0, rand_z_rot))
 94        pallet_tf_mat = omni.usd.get_world_transform_matrix(pallet_prim)
 95        pallet_rot = pallet_tf_mat.ExtractRotation()
 96        pallet_pos = pallet_tf_mat.ExtractTranslation()
 97
 98        # Randomize bin position on top of the rotated pallet area making sure the bin is fully on the pallet
 99        rand_transl_x = np.random.uniform(-pallet_size[0] / 2 + bin_size[0] / 2, pallet_size[0] / 2 - bin_size[0] / 2)
100        rand_transl_y = np.random.uniform(-pallet_size[1] / 2 + bin_size[1] / 2, pallet_size[1] / 2 - bin_size[1] / 2)
101
102        # Adjust bin position to account for the random rotation of the pallet
103        rand_z_rot_rad = np.deg2rad(rand_z_rot)
104        rot_adjusted_transl_x = rand_transl_x * np.cos(rand_z_rot_rad) - rand_transl_y * np.sin(rand_z_rot_rad)
105        rot_adjusted_transl_y = rand_transl_x * np.sin(rand_z_rot_rad) + rand_transl_y * np.cos(rand_z_rot_rad)
106        bin_prim.GetAttribute("xformOp:translate").Set(
107            Gf.Vec3d(
108                pallet_pos[0] + rot_adjusted_transl_x,
109                pallet_pos[1] + rot_adjusted_transl_y,
110                pallet_pos[2] + pallet_size[2] + bin_size[2] / 2,
111            )
112        )
113        # Keep bin rotation aligned with pallet
114        bin_prim.GetAttribute("xformOp:rotateXYZ").Set(pallet_rot.GetAxis() * pallet_rot.GetAngle())
115
116        # Get next camera position on a sphere looking at the bin with a randomized distance
117        rand_radius = np.random.normal(3, 0.5) * pallet_length
118        bin_pos = omni.usd.get_world_transform_matrix(bin_prim).ExtractTranslation()
119        cam_pos = next_point_on_sphere(i, num_points=num_frames, radius=rand_radius, origin=bin_pos)
120        cam.GetAttribute("xformOp:translate").Set(Gf.Vec3d(*cam_pos))
121
122        eye = Gf.Vec3d(*cam_pos)
123        target = Gf.Vec3d(*bin_pos)
124        up_axis = Gf.Vec3d(0, 0, 1)
125        look_at_quatd = Gf.Matrix4d().SetLookAt(eye, target, up_axis).GetInverse().ExtractRotation().GetQuat()
126        cam.GetAttribute("xformOp:orient").Set(Gf.Quatf(look_at_quatd))
127
128        if write_data:
129            await rep.orchestrator.step_async(rt_subframes=4)
130        else:
131            await omni.kit.app.get_app().next_update_async()
132        if delay > 0:
133            await asyncio.sleep(delay)
134
135
136num_frames = 90
137dome_textures = [
138    assets_root_path + "/NVIDIA/Assets/Skies/Cloudy/champagne_castle_1_4k.hdr",
139    assets_root_path + "/NVIDIA/Assets/Skies/Clear/evening_road_01_4k.hdr",
140    assets_root_path + "/NVIDIA/Assets/Skies/Clear/mealie_road_4k.hdr",
141    assets_root_path + "/NVIDIA/Assets/Skies/Clear/qwantani_4k.hdr",
142]
143asyncio.ensure_future(run_randomizations_async(num_frames, dome_light, dome_textures, pallet_prim, bin_prim, delay=0.2))

Physics-based Randomized Volume Filling

The snippet randomizes the stacking of objects on multiple surfaces. It randomly spawns a given number of pallets in the selected areas and then spawns physically simulated boxes on top of them. A temporary collision box area is created around the pallets to prevent the boxes from falling off. Once all the boxes have been dropped, they are moved in various directions and finally pulled towards the center of the pallet for more stable stacking. Finally, the collision area is removed, after which the boxes can also fall to the ground. To allow easier sliding of the boxes into more stable positions, their friction is temporarily reduced during the simulation.

../_images/isaac_tutorial_replicator_randomization_volume_fill.gif ../_images/isaac_tutorial_replicator_randomization_volume_fill_warehouse.gif
Physics-based Randomized Volume Filling
  1import asyncio
  2import random
  3from itertools import chain
  4
  5import carb
  6import omni.kit.app
  7import omni.usd
  8from omni.isaac.core.utils.bounds import compute_aabb, compute_obb, create_bbox_cache
  9from omni.isaac.nucleus import get_assets_root_path
 10from omni.physx import get_physx_simulation_interface
 11from pxr import (
 12    Gf,
 13    PhysicsSchemaTools,
 14    PhysxSchema,
 15    Sdf,
 16    Usd,
 17    UsdGeom,
 18    UsdPhysics,
 19    UsdShade,
 20    UsdUtils,
 21)
 22
 23
 24# Add transformation properties to the prim (if not already present)
 25def set_transform_attributes(prim, location=None, orientation=None, rotation=None, scale=None):
 26    if location is not None:
 27        if not prim.HasAttribute("xformOp:translate"):
 28            UsdGeom.Xformable(prim).AddTranslateOp()
 29        prim.GetAttribute("xformOp:translate").Set(location)
 30    if orientation is not None:
 31        if not prim.HasAttribute("xformOp:orient"):
 32            UsdGeom.Xformable(prim).AddOrientOp()
 33        prim.GetAttribute("xformOp:orient").Set(orientation)
 34    if rotation is not None:
 35        if not prim.HasAttribute("xformOp:rotateXYZ"):
 36            UsdGeom.Xformable(prim).AddRotateXYZOp()
 37        prim.GetAttribute("xformOp:rotateXYZ").Set(rotation)
 38    if scale is not None:
 39        if not prim.HasAttribute("xformOp:scale"):
 40            UsdGeom.Xformable(prim).AddScaleOp()
 41        prim.GetAttribute("xformOp:scale").Set(scale)
 42
 43
 44# Enables collisions with the asset (without rigid body dynamics the asset will be static)
 45def add_colliders(prim):
 46    # Iterate descendant prims (including root) and add colliders to mesh or primitive types
 47    for desc_prim in Usd.PrimRange(prim):
 48        if desc_prim.IsA(UsdGeom.Mesh) or desc_prim.IsA(UsdGeom.Gprim):
 49            # Physics
 50            if not desc_prim.HasAPI(UsdPhysics.CollisionAPI):
 51                collision_api = UsdPhysics.CollisionAPI.Apply(desc_prim)
 52            else:
 53                collision_api = UsdPhysics.CollisionAPI(desc_prim)
 54            collision_api.CreateCollisionEnabledAttr(True)
 55
 56        # Add mesh specific collision properties only to mesh types
 57        if desc_prim.IsA(UsdGeom.Mesh):
 58            if not desc_prim.HasAPI(UsdPhysics.MeshCollisionAPI):
 59                mesh_collision_api = UsdPhysics.MeshCollisionAPI.Apply(desc_prim)
 60            else:
 61                mesh_collision_api = UsdPhysics.MeshCollisionAPI(desc_prim)
 62            mesh_collision_api.CreateApproximationAttr().Set("convexHull")
 63
 64
 65# Enables rigid body dynamics (physics simulation) on the prim (having valid colliders is recommended)
 66def add_rigid_body_dynamics(prim, disable_gravity=False, angular_damping=None):
 67    # Physics
 68    if not prim.HasAPI(UsdPhysics.RigidBodyAPI):
 69        rigid_body_api = UsdPhysics.RigidBodyAPI.Apply(prim)
 70    else:
 71        rigid_body_api = UsdPhysics.RigidBodyAPI(prim)
 72    rigid_body_api.CreateRigidBodyEnabledAttr(True)
 73    # PhysX
 74    if not prim.HasAPI(PhysxSchema.PhysxRigidBodyAPI):
 75        physx_rigid_body_api = PhysxSchema.PhysxRigidBodyAPI.Apply(prim)
 76    else:
 77        physx_rigid_body_api = PhysxSchema.PhysxRigidBodyAPI(prim)
 78    physx_rigid_body_api.GetDisableGravityAttr().Set(disable_gravity)
 79    if angular_damping is not None:
 80        physx_rigid_body_api.CreateAngularDampingAttr().Set(angular_damping)
 81
 82
 83# Create a new prim with the provided asset URL and transform properties
 84def create_asset(stage, asset_url, path, location=None, rotation=None, orientation=None, scale=None):
 85    prim_path = omni.usd.get_stage_next_free_path(stage, path, False)
 86    reference_url = asset_url if asset_url.startswith("omniverse://") else get_assets_root_path() + asset_url
 87    prim = stage.DefinePrim(prim_path, "Xform")
 88    prim.GetReferences().AddReference(reference_url)
 89    set_transform_attributes(prim, location=location, rotation=rotation, orientation=orientation, scale=scale)
 90    return prim
 91
 92
 93# Create a new prim with the provided asset URL and transform properties including colliders
 94def create_asset_with_colliders(stage, asset_url, path, location=None, rotation=None, orientation=None, scale=None):
 95    prim = create_asset(stage, asset_url, path, location, rotation, orientation, scale)
 96    add_colliders(prim)
 97    return prim
 98
 99
100# Create collision walls around the top surface of the prim with the given height and thickness
101def create_collision_walls(stage, prim, bbox_cache=None, height=2, thickness=0.3, material=None, visible=False):
102    # Use the untransformed axis-aligned bounding box to calculate the prim surface size and center
103    if bbox_cache is None:
104        bbox_cache = UsdGeom.BBoxCache(Usd.TimeCode.Default(), includedPurposes=[UsdGeom.Tokens.default_])
105    local_range = bbox_cache.ComputeWorldBound(prim).GetRange()
106    width, depth, local_height = local_range.GetSize()
107    # Raise the midpoint height to the prim's surface
108    mid = local_range.GetMidpoint() + Gf.Vec3d(0, 0, local_height / 2)
109
110    # Define the walls (name, location, size) with the specified thickness added externally to the surface and height
111    walls = [
112        ("floor", (mid[0], mid[1], mid[2] - thickness / 2), (width, depth, thickness)),
113        ("ceiling", (mid[0], mid[1], mid[2] + height + thickness / 2), (width, depth, thickness)),
114        ("left_wall", (mid[0] - (width + thickness) / 2, mid[1], mid[2] + height / 2), (thickness, depth, height)),
115        ("right_wall", (mid[0] + (width + thickness) / 2, mid[1], mid[2] + height / 2), (thickness, depth, height)),
116        ("front_wall", (mid[0], mid[1] + (depth + thickness) / 2, mid[2] + height / 2), (width, thickness, height)),
117        ("back_wall", (mid[0], mid[1] - (depth + thickness) / 2, mid[2] + height / 2), (width, thickness, height)),
118    ]
119
120    # Use the parent prim path to create the walls as children (use local coordinates)
121    prim_path = prim.GetPath()
122    collision_walls = []
123    for name, location, size in walls:
124        prim = stage.DefinePrim(f"{prim_path}/{name}", "Cube")
125        scale = (size[0] / 2.0, size[1] / 2.0, size[2] / 2.0)
126        set_transform_attributes(prim, location=location, scale=scale)
127        add_colliders(prim)
128        if not visible:
129            UsdGeom.Imageable(prim).MakeInvisible()
130        if material is not None:
131            mat_binding_api = UsdShade.MaterialBindingAPI.Apply(prim)
132            mat_binding_api.Bind(material, UsdShade.Tokens.weakerThanDescendants, "physics")
133        collision_walls.append(prim)
134    return collision_walls
135
136
137# Slide the assets independently in perpendicular directions and then pull them all together towards the given center
138async def apply_forces_async(stage, boxes, pallet, strength=550, strength_center_multiplier=2):
139    timeline = omni.timeline.get_timeline_interface()
140    timeline.play()
141    # Get the pallet center and forward vector to apply forces in the perpendicular directions and towards the center
142    pallet_tf: Gf.Matrix4d = UsdGeom.Xformable(pallet).ComputeLocalToWorldTransform(Usd.TimeCode.Default())
143    pallet_center = pallet_tf.ExtractTranslation()
144    pallet_rot: Gf.Rotation = pallet_tf.ExtractRotation()
145    force_forward = Gf.Vec3d(pallet_rot.TransformDir(Gf.Vec3d(1, 0, 0))) * strength
146    force_right = Gf.Vec3d(pallet_rot.TransformDir(Gf.Vec3d(0, 1, 0))) * strength
147
148    physx_api = get_physx_simulation_interface()
149    stage_id = UsdUtils.StageCache.Get().GetId(stage).ToLongInt()
150    for box_prim in boxes:
151        body_path = PhysicsSchemaTools.sdfPathToInt(box_prim.GetPath())
152        forces = [force_forward, force_right, -force_forward, -force_right]
153        for force in chain(forces, forces):
154            box_tf: Gf.Matrix4d = UsdGeom.Xformable(box_prim).ComputeLocalToWorldTransform(Usd.TimeCode.Default())
155            box_position = carb.Float3(*box_tf.ExtractTranslation())
156            physx_api.apply_force_at_pos(stage_id, body_path, carb.Float3(force), box_position, "Force")
157            for _ in range(10):
158                await omni.kit.app.get_app().next_update_async()
159
160    # Pull all box at once to the pallet center
161    for box_prim in boxes:
162        body_path = PhysicsSchemaTools.sdfPathToInt(box_prim.GetPath())
163        box_tf: Gf.Matrix4d = UsdGeom.Xformable(box_prim).ComputeLocalToWorldTransform(Usd.TimeCode.Default())
164        box_location = box_tf.ExtractTranslation()
165        force_to_center = (pallet_center - box_location) * strength * strength_center_multiplier
166        physx_api.apply_force_at_pos(stage_id, body_path, carb.Float3(*force_to_center), carb.Float3(*box_location))
167    for _ in range(20):
168        await omni.kit.app.get_app().next_update_async()
169    timeline.pause()
170
171
172# Create a new stage and and run the example scenario
173async def stack_boxes_on_pallet_async(pallet_prim, boxes_urls_and_weights, num_boxes, drop_height=1.5, drop_margin=0.2):
174    pallet_path = pallet_prim.GetPath()
175    print(f"[BoxStacking] Running scenario for pallet {pallet_path} with {num_boxes} boxes..")
176    stage = omni.usd.get_context().get_stage()
177    bbox_cache = UsdGeom.BBoxCache(Usd.TimeCode.Default(), includedPurposes=[UsdGeom.Tokens.default_])
178
179    # Create a custom physics material to allow the boxes to easily slide into stacking positions
180    material_path = f"{pallet_path}/Looks/PhysicsMaterial"
181    default_material = UsdShade.Material.Define(stage, material_path)
182    physics_material = UsdPhysics.MaterialAPI.Apply(default_material.GetPrim())
183    physics_material.CreateRestitutionAttr().Set(0.0)  # Inelastic collision (no bouncing)
184    physics_material.CreateStaticFrictionAttr().Set(0.01)  # Small friction to allow sliding of stationary boxes
185    physics_material.CreateDynamicFrictionAttr().Set(0.01)  # Small friction to allow sliding of moving boxes
186
187    # Apply the physics material to the pallet
188    mat_binding_api = UsdShade.MaterialBindingAPI.Apply(pallet_prim)
189    mat_binding_api.Bind(default_material, UsdShade.Tokens.weakerThanDescendants, "physics")
190
191    # Create collision walls around the top of the pallet and apply the physics material to them
192    collision_walls = create_collision_walls(
193        stage, pallet_prim, bbox_cache, height=drop_height + drop_margin, material=default_material
194    )
195
196    # Create the random boxes (without physics) with the specified weights and sort them by size (volume)
197    box_urls, box_weights = zip(*boxes_urls_and_weights)
198    rand_boxes_urls = random.choices(box_urls, weights=box_weights, k=num_boxes)
199    boxes = [create_asset(stage, box_url, f"{pallet_path}_Boxes/Box_{i}") for i, box_url in enumerate(rand_boxes_urls)]
200    boxes.sort(key=lambda box: bbox_cache.ComputeLocalBound(box).GetVolume(), reverse=True)
201
202    # Calculate the drop area above the pallet taking into account the pallet surface, drop height and the margin
203    # Note: The boxes can be spawned colliding with the surrounding collision walls as they will be pushed inwards
204    pallet_range = bbox_cache.ComputeWorldBound(pallet_prim).GetRange()
205    pallet_width, pallet_depth, pallet_heigth = pallet_range.GetSize()
206    # Move the spawn center at the given height above the pallet surface
207    spawn_center = pallet_range.GetMidpoint() + Gf.Vec3d(0, 0, pallet_heigth / 2 + drop_height)
208    spawn_width, spawn_depth = pallet_width / 2 - drop_margin, pallet_depth / 2 - drop_margin
209
210    # Use the pallet local-to-world transform to apply the local random offsets relative to the pallet
211    pallet_tf: Gf.Matrix4d = UsdGeom.Xformable(pallet_prim).ComputeLocalToWorldTransform(Usd.TimeCode.Default())
212    pallet_rot: Gf.Rotation = pallet_tf.ExtractRotation()
213
214    # Simulate dropping the boxes from random poses on the pallet
215    timeline = omni.timeline.get_timeline_interface()
216    for box_prim in boxes:
217        # Create a random location and orientation for the box within the drop area in local frame
218        local_loc = spawn_center + Gf.Vec3d(
219            random.uniform(-spawn_width, spawn_width), random.uniform(-spawn_depth, spawn_depth), 0
220        )
221        axes = [Gf.Vec3d(1, 0, 0), Gf.Vec3d(0, 1, 0), Gf.Vec3d(0, 0, 1)]
222        angles = [random.choice([180, 90, 0, -90, -180]) + random.uniform(-3, 3) for _ in axes]
223        local_rot = Gf.Rotation()
224        for axis, angle in zip(axes, angles):
225            local_rot *= Gf.Rotation(axis, angle)
226
227        # Transform the local pose to the pallet's world coordinate system
228        world_loc = pallet_tf.Transform(local_loc)
229        world_quat = Gf.Quatf((pallet_rot * local_rot).GetQuat())
230
231        # Set the spawn pose and enable collisions and rigid body dynamics with dampened angular movements
232        set_transform_attributes(box_prim, location=world_loc, orientation=world_quat)
233        add_colliders(box_prim)
234        add_rigid_body_dynamics(box_prim, angular_damping=0.9)
235
236        # Bind the physics material to the box (allow frictionless sliding)
237        mat_binding_api = UsdShade.MaterialBindingAPI.Apply(box_prim)
238        mat_binding_api.Bind(default_material, UsdShade.Tokens.weakerThanDescendants, "physics")
239        # Wait for an app update to load the new attributes
240        await omni.kit.app.get_app().next_update_async()
241
242        # Play simulation for a few frames for each box
243        timeline.play()
244        for _ in range(20):
245            await omni.kit.app.get_app().next_update_async()
246        timeline.pause()
247
248    # Iteratively apply forces to the boxes to move them around then pull them all together towards the pallet center
249    await apply_forces_async(stage, boxes, pallet_prim)
250
251    # Remove rigid body dynamics of the boxes until all other scenarios are completed
252    for box in boxes:
253        UsdPhysics.RigidBodyAPI(box).GetRigidBodyEnabledAttr().Set(False)
254
255    # Increase the friction to prevent sliding of the boxes on the pallet before removing the collision walls
256    physics_material.CreateStaticFrictionAttr().Set(0.9)
257    physics_material.CreateDynamicFrictionAttr().Set(0.9)
258
259    # Remove collision walls
260    for wall in collision_walls:
261        stage.RemovePrim(wall.GetPath())
262    return boxes
263
264
265# Run the example scenario
266async def run_box_stacking_scenarios_async(num_pallets=1, env_url=None):
267    # List of pallets and boxes to randomly choose from with their respective weights
268    pallets_urls_and_weights = [
269        ("/Isaac/Environments/Simple_Warehouse/Props/SM_PaletteA_01.usd", 0.25),
270        ("/Isaac/Environments/Simple_Warehouse/Props/SM_PaletteA_02.usd", 0.75),
271    ]
272    boxes_urls_and_weights = [
273        ("/Isaac/Environments/Simple_Warehouse/Props/SM_CardBoxA_01.usd", 0.02),
274        ("/Isaac/Environments/Simple_Warehouse/Props/SM_CardBoxB_01.usd", 0.06),
275        ("/Isaac/Environments/Simple_Warehouse/Props/SM_CardBoxC_01.usd", 0.12),
276        ("/Isaac/Environments/Simple_Warehouse/Props/SM_CardBoxD_01.usd", 0.80),
277    ]
278
279    # Load a predefined or create a new stage
280    if env_url is not None:
281        env_path = env_url if env_url.startswith("omniverse://") else get_assets_root_path() + env_url
282        omni.usd.get_context().open_stage(env_path)
283        stage = omni.usd.get_context().get_stage()
284    else:
285        omni.usd.get_context().new_stage()
286        stage = omni.usd.get_context().get_stage()
287        distant_light = stage.DefinePrim("/World/Lights/DistantLight", "DistantLight")
288        distant_light.CreateAttribute("inputs:intensity", Sdf.ValueTypeNames.Float).Set(400.0)
289        if not distant_light.HasAttribute("xformOp:rotateXYZ"):
290            UsdGeom.Xformable(distant_light).AddRotateXYZOp()
291        distant_light.GetAttribute("xformOp:rotateXYZ").Set((0, 60, 0))
292        dome_light = stage.DefinePrim("/World/Lights/DomeLight", "DomeLight")
293        dome_light.CreateAttribute("inputs:intensity", Sdf.ValueTypeNames.Float).Set(500.0)
294
295    # Spawn the pallets
296    pallets = []
297    pallets_urls, pallets_weights = zip(*pallets_urls_and_weights)
298    rand_pallet_urls = random.choices(pallets_urls, weights=pallets_weights, k=num_pallets)
299    # Custom pallet poses for the evnironment
300    custom_pallet_locations = [
301        (-9.3, 5.3, 1.3),
302        (-9.3, 7.3, 1.3),
303        (-9.3, -0.6, 1.3),
304    ]
305    random.shuffle(custom_pallet_locations)
306    for i, pallet_url in enumerate(rand_pallet_urls):
307        # Use a custom location for every other pallet
308        if env_url is not None:
309            if i % 2 == 0 and custom_pallet_locations:
310                rand_loc = Gf.Vec3d(*custom_pallet_locations.pop())
311            else:
312                rand_loc = Gf.Vec3d(-6.5, i * 1.75, 0) + Gf.Vec3d(random.uniform(-0.2, 0.2), random.uniform(0, 0.2), 0)
313        else:
314            rand_loc = Gf.Vec3d(i * 1.5, 0, 0) + Gf.Vec3d(random.uniform(0, 0.2), random.uniform(-0.2, 0.2), 0)
315        rand_rot = (0, 0, random.choice([180, 90, 0, -90, -180]) + random.uniform(-15, 15))
316        pallet_prim = create_asset_with_colliders(
317            stage, pallet_url, f"/World/Pallet_{i}", location=rand_loc, rotation=rand_rot
318        )
319        pallets.append(pallet_prim)
320
321    # Stack the boxes on the pallets
322    total_boxes = []
323    for pallet in pallets:
324        if env_url is not None:
325            rand_num_boxes = random.randint(8, 15)
326            stacked_boxes = await stack_boxes_on_pallet_async(
327                pallet, boxes_urls_and_weights, num_boxes=rand_num_boxes, drop_height=1.0
328            )
329        else:
330            rand_num_boxes = random.randint(12, 20)
331            stacked_boxes = await stack_boxes_on_pallet_async(pallet, boxes_urls_and_weights, num_boxes=rand_num_boxes)
332        total_boxes.extend(stacked_boxes)
333
334    # Re-enable rigid body dynamics of the boxes and run the simulation for a while
335    for box in total_boxes:
336        UsdPhysics.RigidBodyAPI(box).GetRigidBodyEnabledAttr().Set(True)
337    timeline = omni.timeline.get_timeline_interface()
338    timeline.play()
339    for _ in range(200):
340        await omni.kit.app.get_app().next_update_async()
341    timeline.pause()
342
343
344async def run_scenarios_async():
345    await run_box_stacking_scenarios_async(num_pallets=6)
346    await run_box_stacking_scenarios_async(num_pallets=6, env_url="/Isaac/Environments/Simple_Warehouse/warehouse.usd")
347
348
349asyncio.ensure_future(run_scenarios_async())