|
19 | 19 |
|
20 | 20 | import carb |
21 | 21 | import omni.physics.tensors.impl.api as physx |
| 22 | +import omni.usd |
| 23 | +from pxr import Gf, Sdf, UsdGeom, Vt |
22 | 24 |
|
23 | 25 | import isaaclab.sim as sim_utils |
24 | 26 | import isaaclab.utils.math as math_utils |
|
31 | 33 | from isaaclab.envs import ManagerBasedEnv |
32 | 34 |
|
33 | 35 |
|
| 36 | +def randomize_rigid_body_scale( |
| 37 | + env: ManagerBasedEnv, |
| 38 | + env_ids: torch.Tensor | None, |
| 39 | + scale_range: tuple[float, float] | dict[str, tuple[float, float]], |
| 40 | + asset_cfg: SceneEntityCfg, |
| 41 | + relative_child_path: str | None = None, |
| 42 | +): |
| 43 | + """Randomize the scale of a rigid body asset in the USD stage. |
| 44 | +
|
| 45 | + This function modifies the "xformOp:scale" property of all the prims corresponding to the asset. |
| 46 | +
|
| 47 | + It takes a tuple or dictionary for the scale ranges. If it is a tuple, then the scaling along |
| 48 | + individual axis is performed equally. If it is a dictionary, the scaling is independent across each dimension. |
| 49 | + The keys of the dictionary are ``x``, ``y``, and ``z``. The values are tuples of the form ``(min, max)``. |
| 50 | +
|
| 51 | + If the dictionary does not contain a key, the range is set to one for that axis. |
| 52 | +
|
| 53 | + Relative child path can be used to randomize the scale of a specific child prim of the asset. |
| 54 | + For example, if the asset at prim path expression "/World/envs/env_.*/Object" has a child |
| 55 | + with the path "/World/envs/env_.*/Object/mesh", then the relative child path should be "mesh" or |
| 56 | + "/mesh". |
| 57 | +
|
| 58 | + .. attention:: |
| 59 | + Since this function modifies USD properties that are parsed by the physics engine once the simulation |
| 60 | + starts, the term should only be used before the simulation starts playing. This corresponds to the |
| 61 | + event mode named "usd". Using it at simulation time, may lead to unpredictable behaviors. |
| 62 | +
|
| 63 | + .. note:: |
| 64 | + When randomizing the scale of individual assets, please make sure to set |
| 65 | + :attr:`isaaclab.scene.InteractiveSceneCfg.replicate_physics` to False. This ensures that physics |
| 66 | + parser will parse the individual asset properties separately. |
| 67 | + """ |
| 68 | + # check if sim is running |
| 69 | + if env.sim.is_playing(): |
| 70 | + raise RuntimeError( |
| 71 | + "Randomizing scale while simulation is running leads to unpredictable behaviors." |
| 72 | + " Please ensure that the event term is called before the simulation starts by using the 'usd' mode." |
| 73 | + ) |
| 74 | + |
| 75 | + # extract the used quantities (to enable type-hinting) |
| 76 | + asset: RigidObject = env.scene[asset_cfg.name] |
| 77 | + |
| 78 | + if isinstance(asset, Articulation): |
| 79 | + raise ValueError( |
| 80 | + "Scaling an articulation randomly is not supported, as it affects joint attributes and can cause" |
| 81 | + " unexpected behavior. To achieve different scales, we recommend generating separate USD files for" |
| 82 | + " each version of the articulation and using multi-asset spawning. For more details, refer to:" |
| 83 | + " https://isaac-sim.github.io/IsaacLab/main/source/how-to/multi_asset_spawning.html" |
| 84 | + ) |
| 85 | + |
| 86 | + # resolve environment ids |
| 87 | + if env_ids is None: |
| 88 | + env_ids = torch.arange(env.scene.num_envs, device="cpu") |
| 89 | + else: |
| 90 | + env_ids = env_ids.cpu() |
| 91 | + |
| 92 | + # acquire stage |
| 93 | + stage = omni.usd.get_context().get_stage() |
| 94 | + # resolve prim paths for spawning and cloning |
| 95 | + prim_paths = sim_utils.find_matching_prim_paths(asset.cfg.prim_path) |
| 96 | + |
| 97 | + # sample scale values |
| 98 | + if isinstance(scale_range, dict): |
| 99 | + range_list = [scale_range.get(key, (1.0, 1.0)) for key in ["x", "y", "z"]] |
| 100 | + ranges = torch.tensor(range_list, device="cpu") |
| 101 | + rand_samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (len(env_ids), 3), device="cpu") |
| 102 | + else: |
| 103 | + rand_samples = math_utils.sample_uniform(*scale_range, (len(env_ids), 1), device="cpu") |
| 104 | + rand_samples = rand_samples.repeat(1, 3) |
| 105 | + # convert to list for the for loop |
| 106 | + rand_samples = rand_samples.tolist() |
| 107 | + |
| 108 | + # apply the randomization to the parent if no relative child path is provided |
| 109 | + # this might be useful if user wants to randomize a particular mesh in the prim hierarchy |
| 110 | + if relative_child_path is None: |
| 111 | + relative_child_path = "" |
| 112 | + elif not relative_child_path.startswith("/"): |
| 113 | + relative_child_path = "/" + relative_child_path |
| 114 | + |
| 115 | + # use sdf changeblock for faster processing of USD properties |
| 116 | + with Sdf.ChangeBlock(): |
| 117 | + for i, env_id in enumerate(env_ids): |
| 118 | + # path to prim to randomize |
| 119 | + prim_path = prim_paths[env_id] + relative_child_path |
| 120 | + # spawn single instance |
| 121 | + prim_spec = Sdf.CreatePrimInLayer(stage.GetRootLayer(), prim_path) |
| 122 | + |
| 123 | + # get the attribute to randomize |
| 124 | + scale_spec = prim_spec.GetAttributeAtPath(prim_path + ".xformOp:scale") |
| 125 | + # if the scale attribute does not exist, create it |
| 126 | + has_scale_attr = scale_spec is not None |
| 127 | + if not has_scale_attr: |
| 128 | + scale_spec = Sdf.AttributeSpec(prim_spec, prim_path + ".xformOp:scale", Sdf.ValueTypeNames.Double3) |
| 129 | + |
| 130 | + # set the new scale |
| 131 | + scale_spec.default = Gf.Vec3f(*rand_samples[i]) |
| 132 | + |
| 133 | + # ensure the operation is done in the right ordering if we created the scale attribute. |
| 134 | + # otherwise, we assume the scale attribute is already in the right order. |
| 135 | + # note: by default isaac sim follows this ordering for the transform stack so any asset |
| 136 | + # created through it will have the correct ordering |
| 137 | + if not has_scale_attr: |
| 138 | + op_order_spec = prim_spec.GetAttributeAtPath(prim_path + ".xformOpOrder") |
| 139 | + if op_order_spec is None: |
| 140 | + op_order_spec = Sdf.AttributeSpec( |
| 141 | + prim_spec, UsdGeom.Tokens.xformOpOrder, Sdf.ValueTypeNames.TokenArray |
| 142 | + ) |
| 143 | + op_order_spec.default = Vt.TokenArray(["xformOp:translate", "xformOp:orient", "xformOp:scale"]) |
| 144 | + |
| 145 | + |
34 | 146 | class randomize_rigid_body_material(ManagerTermBase): |
35 | 147 | """Randomize the physics materials on all geometries of the asset. |
36 | 148 |
|
|
0 commit comments