create random assignment for heterogenous cloning

ooctipus · ooctipus · commit ba5c1199e4f0 · 2025-11-04T23:58:40.000-08:00
diff --git a/scripts/demos/multi_asset.py b/scripts/demos/multi_asset.py
@@ -105,7 +105,7 @@ class MultiObjectSceneCfg(InteractiveSceneCfg):
 
     # rigid object
     object: RigidObjectCfg = RigidObjectCfg(
-        prim_path="/World/envs/env_.*/Object",
+        prim_path="{ENV_REGEX_NS}/Object",
         spawn=sim_utils.MultiAssetSpawnerCfg(
             assets_cfg=[
                 sim_utils.ConeCfg(
@@ -136,7 +136,7 @@ class MultiObjectSceneCfg(InteractiveSceneCfg):
     object_collection: RigidObjectCollectionCfg = RigidObjectCollectionCfg(
         rigid_objects={
             "object_A": RigidObjectCfg(
-                prim_path="/World/envs/env_.*/Object_A",
+                prim_path="{ENV_REGEX_NS}/Object_A",
                 spawn=sim_utils.SphereCfg(
                     radius=0.1,
                     visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), metallic=0.2),
@@ -149,7 +149,7 @@ class MultiObjectSceneCfg(InteractiveSceneCfg):
                 init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, -0.5, 2.0)),
             ),
             "object_B": RigidObjectCfg(
-                prim_path="/World/envs/env_.*/Object_B",
+                prim_path="{ENV_REGEX_NS}/Object_B",
                 spawn=sim_utils.CuboidCfg(
                     size=(0.1, 0.1, 0.1),
                     visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), metallic=0.2),
@@ -162,7 +162,7 @@ class MultiObjectSceneCfg(InteractiveSceneCfg):
                 init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.5, 2.0)),
             ),
             "object_C": RigidObjectCfg(
-                prim_path="/World/envs/env_.*/Object_C",
+                prim_path="{ENV_REGEX_NS}/Object_C",
                 spawn=sim_utils.ConeCfg(
                     radius=0.1,
                     height=0.3,
@@ -180,7 +180,7 @@ class MultiObjectSceneCfg(InteractiveSceneCfg):
 
     # articulation
     robot: ArticulationCfg = ArticulationCfg(
-        prim_path="/World/envs/env_.*/Robot",
+        prim_path="{ENV_REGEX_NS}/Robot",
         spawn=sim_utils.MultiUsdFileCfg(
             usd_path=[
                 f"{ISAACLAB_NUCLEUS_DIR}/Robots/ANYbotics/ANYmal-C/anymal_c.usd",
diff --git a/source/isaaclab/isaaclab/scene/cloner.py b/source/isaaclab/isaaclab/scene/cloner.py
@@ -9,38 +9,79 @@ def usd_replicate(
     sources: list[str],
     destinations: list[str],
     env_ids: torch.Tensor,
-    mapping: torch.Tensor,
+    mask: torch.Tensor | None = None,
     positions: torch.Tensor | None = None,
     quaternions: torch.Tensor | None = None,
 ) -> None:
+    """Replicate USD prims to per-environment destinations.
+
+    Copies each source prim spec to destination templates for selected environments
+    (``mask``). Optionally authors translate/orient from position/quaternion buffers.
+    Replication runs in path-depth order (parents before children) for robust composition.
+
+    Args:
+        stage: USD stage.
+        sources: Source prim paths.
+        destinations: Destination formattable templates with ``"{}"`` for env index.
+        env_ids: Environment indices.
+        mask: Optional per-source or shared mask. ``None`` selects all.
+        positions: Optional positions (``[E, 3]``) -> ``xformOp:translate``.
+        quaternions: Optional orientations (``[E, 4]``) in ``wxyz`` -> ``xformOp:orient``.
+
+    Returns:
+        None
+    """
     rl = stage.GetRootLayer()
-    with Sdf.ChangeBlock():
-        for i, src in enumerate(sources):
-            tmpl = destinations[i]
-            for wid in env_ids[mapping[i]].tolist():
-                dp = tmpl.format(wid)
-                Sdf.CreatePrimInLayer(rl, dp)
-                Sdf.CopySpec(rl, Sdf.Path(src), rl, Sdf.Path(dp))
-
-                if positions is not None or quaternions is not None:
-                    ps = rl.GetPrimAtPath(dp)
-                    # translate
-                    if positions is not None:
-                        p = positions[wid]
-                        t_attr = ps.GetAttributeAtPath(dp + ".xformOp:translate")
-                        if t_attr is None:
-                            t_attr = Sdf.AttributeSpec(ps, "xformOp:translate", Sdf.ValueTypeNames.Double3)
-                        t_attr.default = Gf.Vec3d(float(p[0]), float(p[1]), float(p[2]))
-                    # orient (always set as Quatd)
-                    if quaternions is not None:
-                        q = quaternions[wid]
-                        o_attr = ps.GetAttributeAtPath(dp + ".xformOp:orient")
-                        if o_attr is None:
-                            o_attr = Sdf.AttributeSpec(ps, "xformOp:orient", Sdf.ValueTypeNames.Quatd)
-                        o_attr.default = Gf.Quatd(float(q[0]), Gf.Vec3d(float(q[1]), float(q[2]), float(q[3])))
-                    # op order
-                    order = ps.GetAttributeAtPath(dp + ".xformOpOrder") or Sdf.AttributeSpec(ps, UsdGeom.Tokens.xformOpOrder, Sdf.ValueTypeNames.TokenArray)
-                    order.default = Vt.TokenArray(["xformOp:translate", "xformOp:orient", "xformOp:scale"])
+
+    # Group replication by destination path depth so ancestors land before deeper paths.
+    # This avoids composition issues for nested or interdependent specs.
+    def dp_depth(template: str) -> int:
+        dp = template.format(0)
+        return Sdf.Path(dp).pathElementCount
+
+    order = sorted(range(len(sources)), key=lambda i: dp_depth(destinations[i]))
+
+    # Process in layers of equal depth, committing at each depth to stabilize composition
+    depth_to_indices: dict[int, list[int]] = {}
+    for i in order:
+        d = dp_depth(destinations[i])
+        depth_to_indices.setdefault(d, []).append(i)
+
+    for depth in sorted(depth_to_indices.keys()):
+        with Sdf.ChangeBlock():
+            for i in depth_to_indices[depth]:
+                src = sources[i]
+                tmpl = destinations[i]
+                # Select target environments for this source (supports None, [E], or [S, E])
+                target_envs = env_ids if mask is None else env_ids[mask[i]]
+                for wid in target_envs.tolist():
+                    dp = tmpl.format(wid)
+                    Sdf.CreatePrimInLayer(rl, dp)
+                    Sdf.CopySpec(rl, Sdf.Path(src), rl, Sdf.Path(dp))
+
+                    if positions is not None or quaternions is not None:
+                        ps = rl.GetPrimAtPath(dp)
+                        op_names = []
+                        if positions is not None:
+                            p = positions[wid]
+                            t_attr = ps.GetAttributeAtPath(dp + ".xformOp:translate")
+                            if t_attr is None:
+                                t_attr = Sdf.AttributeSpec(ps, "xformOp:translate", Sdf.ValueTypeNames.Double3)
+                            t_attr.default = Gf.Vec3d(float(p[0]), float(p[1]), float(p[2]))
+                            op_names.append("xformOp:translate")
+                        if quaternions is not None:
+                            q = quaternions[wid]
+                            o_attr = ps.GetAttributeAtPath(dp + ".xformOp:orient")
+                            if o_attr is None:
+                                o_attr = Sdf.AttributeSpec(ps, "xformOp:orient", Sdf.ValueTypeNames.Quatd)
+                            o_attr.default = Gf.Quatd(float(q[0]), Gf.Vec3d(float(q[1]), float(q[2]), float(q[3])))
+                            op_names.append("xformOp:orient")
+                        # Only author xformOpOrder for the ops we actually authored
+                        if op_names:
+                            op_order = ps.GetAttributeAtPath(dp + ".xformOpOrder") or Sdf.AttributeSpec(
+                                ps, UsdGeom.Tokens.xformOpOrder, Sdf.ValueTypeNames.TokenArray
+                            )
+                            op_order.default = Vt.TokenArray(op_names)
 
 
 def physx_replicate(
@@ -49,16 +90,24 @@ def physx_replicate(
     destinations: list[str],       # e.g. ["/World/envs/env_{}/Robot", "/World/envs/env_{}/Object"]
     env_ids: torch.Tensor,         # env_ids
     mapping: torch.Tensor,         # (num_sources, num_envs) bool; True -> place sources[i] into world=j
-    use_env_ids: bool = True,      # hint: only used when a row covers all worlds
     use_fabric: bool = False,
 ) -> None:
-    """
-    PhysX replication that matches the (sources, destinations, mapping) convention.
-    For each source i, we build the per-world destination list:
-        dests_i = [ destinations[i].format(w) for w in worlds where mapping[i, w] is True ]
+    """Replicate prims via PhysX replicator with per-row mapping.
+
+    Builds per-source destination lists from ``mapping`` and calls PhysX ``replicate``.
+    Rows covering all environments use ``useEnvIds=True``; partial rows use ``False``.
+    The replicator is registered for the call and then unregistered.
 
-    If a row covers ALL worlds, we pass useEnvIds=True to PhysX so per-env filtering is automatic.
-    If a row is partial, we force useEnvIds=False so cross-row collisions within the same env remain possible.
+    Args:
+        stage: USD stage.
+        sources: Source prim paths (``S``).
+        destinations: Destination templates (``S``) with ``"{}"`` for env index.
+        env_ids: Environment indices (``[E]``).
+        mapping: Bool/int mask (``[S, E]``) selecting envs per source.
+        use_fabric: Use Fabric for replication.
+
+    Returns:
+        None
     """
 
     stage_id = UsdUtils.StageCache.Get().Insert(stage).ToLongInt()
@@ -86,16 +135,27 @@ def attach_end_fn(_stage_id: int):
 
 
 def filter_collisions(
-    stage: Usd.Stage, physicsscene_path: str, collision_root_path: str, prim_paths: list[str], global_paths: list[str] = []
-):
-    """Filters collisions between clones. Clones will not collide with each other, but can collide with objects specified in global_paths.
+    stage: Usd.Stage,
+    physicsscene_path: str,
+    collision_root_path: str,
+    prim_paths: list[str],
+    global_paths: list[str] = []
+) -> None:
+    """Create inverted collision groups for clones.
 
-    Args:
-        physicsscene_path (str): Path to PhysicsScene object in stage.
-        collision_root_path (str): Path to place collision groups under.
-        prim_paths (List[str]): Paths of objects to filter out collision.
-        global_paths (List[str]): Paths of objects to generate collision (e.g. ground plane).
+    Creates one PhysicsCollisionGroup per prim under ``collision_root_path``, enabling
+    inverted filtering so clones don't collide across groups. Optionally adds a global
+    group that collides with all.
 
+    Args:
+        stage: USD stage.
+        physicsscene_path: Path to PhysicsScene prim.
+        collision_root_path: Root scope for collision groups.
+        prim_paths: Per-clone prim paths.
+        global_paths: Optional global-collider paths.
+
+    Returns:
+        None
     """
 
     physx_scene = PhysxSchema.PhysxSceneAPI(stage.GetPrimAtPath(physicsscene_path))
@@ -181,6 +241,23 @@ def filter_collisions(
 
 
 def grid_transforms(N: int, spacing: float = 1.0, up_axis: str = "z", device="cpu"):
+    """Create a centered grid of transforms for ``N`` instances.
+
+    Computes ``(x, y)`` coordinates in a roughly square grid centered at the origin
+    with the provided spacing, places the third coordinate according to ``up_axis``,
+    and returns identity orientations (``wxyz``) for each instance.
+
+    Args:
+        N: Number of instances.
+        spacing: Distance between neighboring grid positions.
+        up_axis: Up axis for positions ("z", "y", or "x").
+        device: Torch device for returned tensors.
+
+    Returns:
+        A tuple ``(pos, ori)`` where:
+            - ``pos`` is a tensor of shape ``(N, 3)`` with positions.
+            - ``ori`` is a tensor of shape ``(N, 4)`` with identity quaternions in ``(w, x, y, z)``.
+    """
     # rows/cols
     rows = int(math.ceil(math.sqrt(N)))
     cols = int(math.ceil(N / rows))
diff --git a/source/isaaclab/isaaclab/scene/interactive_scene.py b/source/isaaclab/isaaclab/scene/interactive_scene.py
@@ -30,8 +30,8 @@
 from isaaclab.sim import SimulationContext
 from isaaclab.sim.utils import get_current_stage_id
 from isaaclab.terrains import TerrainImporter, TerrainImporterCfg
-from isaaclab.sim.spawners.cloner import physx_replicate, filter_collisions, usd_replicate, grid_transforms
 
+from .cloner import physx_replicate, filter_collisions, usd_replicate, grid_transforms
 from .interactive_scene_cfg import InteractiveSceneCfg
 
 
@@ -145,14 +145,7 @@ def __init__(self, cfg: InteractiveSceneCfg):
             prim.CreateAttribute("physxScene:envIdInBoundsBitCount", Sdf.ValueTypeNames.Int).Set(4)
         self._default_env_origins, _ = grid_transforms(self.cfg.num_envs, self.cfg.env_spacing, device=self.device)
         # copy empty prim of env_0 to env_1, env_2, ..., env_{num_envs-1} with correct location.
-        usd_replicate(
-            stage=self.stage,
-            sources=[self.env_fmt.format(0)],
-            destinations=[self.env_fmt],
-            env_ids=self._ALL_INDICES,
-            mapping=torch.ones((1, self.num_envs), dtype=torch.bool),
-            positions=self._default_env_origins
-        )
+        usd_replicate(self.stage, [self.env_fmt.format(0)], [self.env_fmt], self._ALL_INDICES, positions=self._default_env_origins)
 
         self._global_prim_paths = list()
         if self._is_scene_setup_from_cfg():
@@ -182,15 +175,20 @@ def clone_environments(self, copy_from_source: bool = False):
                 self.template_path,
                 predicate=lambda prim: str(prim.GetPath()).split("/")[-1].startswith(self.prototype_name)
             )
-            prototype_set = set(["/".join(str(prototype.GetPath()).split("/")[:-1]) for prototype in prototypes])
-            for prototype_path in prototype_set:
-                protos = sim_utils.find_matching_prim_paths(f"{prototype_path}/.*")
+            prototype_root_set = set(["/".join(str(prototype.GetPath()).split("/")[:-1]) for prototype in prototypes])
+            for prototype_root in prototype_root_set:
+                protos = sim_utils.find_matching_prim_paths(f"{prototype_root}/.*")
                 protos = [proto for proto in protos if proto.split("/")[-1].startswith(self.prototype_name)]
                 m = torch.zeros((len(protos), self.num_envs), dtype=torch.bool, device=self.device)
-                m[self._ALL_INDICES % len(protos), self._ALL_INDICES] = True
+                # Optionally select prototypes randomly per environment; else round-robin by modulo
+                if self.cfg.random_heterogenous_cloning:
+                    rand_idx = torch.randint(len(protos), (self.num_envs,), device=self.device)
+                    m[rand_idx, self._ALL_INDICES] = True
+                else:
+                    m[self._ALL_INDICES % len(protos), self._ALL_INDICES] = True
 
                 clone_plan['src'].extend(protos)
-                clone_plan['dest'].extend([prototype_path.replace(self.template_path, self.env_fmt)] * len(protos))
+                clone_plan['dest'].extend([prototype_root.replace(self.template_path, self.env_fmt)] * len(protos))
                 clone_plan['mapping'] = torch.cat((clone_plan['mapping'].reshape(-1, m.size(1)), m), dim=0)
 
             proto_idx = clone_plan['mapping'].to(torch.int32).argmax(dim=1)
@@ -205,14 +203,14 @@ def clone_environments(self, copy_from_source: bool = False):
             else:
                 src = [tpl.format(int(idx)) for tpl, idx in zip(clone_plan['dest'], proto_idx.tolist())]
                 replicate_args = src, clone_plan['dest'], self._ALL_INDICES, clone_plan['mapping']
-
-            physx_replicate(self.stage, *replicate_args)
-            usd_replicate(self.stage, *replicate_args, positions=self._default_env_origins)
         else:
             mapping = torch.ones((1, self.num_envs), device=self.device, dtype=torch.bool)
             replicate_args = [self.env_fmt.format(0)], [self.env_fmt], self._ALL_INDICES, mapping
+
+        if not copy_from_source:
+            # skip physx cloning, this means physx will walk and parse the stage one by one faithfully
             physx_replicate(self.stage, *replicate_args)
-            usd_replicate(self.stage, *replicate_args, positions=self._default_env_origins)
+        usd_replicate(self.stage, *replicate_args, positions=self._default_env_origins)
 
     def filter_collisions(self, global_prim_paths: list[str] | None = None):
         """Filter environments collisions.
diff --git a/source/isaaclab/isaaclab/scene/interactive_scene_cfg.py b/source/isaaclab/isaaclab/scene/interactive_scene_cfg.py
@@ -124,3 +124,11 @@ class MySceneCfg(InteractiveSceneCfg):
         default to ``False``.
 
     """
+
+    random_heterogenous_cloning: bool = True
+    """Randomly assign prototypes to environments. Default is True.
+
+    When enabled, each environment selects a prototype at random from the available
+    prototypes under a given template root. When disabled, environments use a
+    round-robin assignment based on ``env_index % num_prototypes``.
+    """
diff --git a/source/isaaclab/isaaclab/sim/spawners/wrappers/wrappers.py b/source/isaaclab/isaaclab/sim/spawners/wrappers/wrappers.py
@@ -14,7 +14,6 @@
 from isaacsim.core.utils.stage import get_current_stage
 from pxr import Usd
 
-import isaaclab.sim as sim_utils
 from isaaclab.sim.spawners.from_files import UsdFileCfg
 
 if TYPE_CHECKING:
@@ -124,8 +123,6 @@ def spawn_multi_usd_file(
     for usd_path in usd_paths:
         usd_cfg = usd_template_cfg.replace(usd_path=usd_path)
         multi_asset_cfg.assets_cfg.append(usd_cfg)
-    # set random choice
-    multi_asset_cfg.random_choice = cfg.random_choice
 
     # propagate the contact sensor settings
     # note: the default value for activate_contact_sensors in MultiAssetSpawnerCfg is False.
diff --git a/source/isaaclab/isaaclab/sim/spawners/wrappers/wrappers_cfg.py b/source/isaaclab/isaaclab/sim/spawners/wrappers/wrappers_cfg.py
@@ -64,4 +64,8 @@ class MultiUsdFileCfg(UsdFileCfg):
 
     If False, the asset configurations are spawned in the order they are provided in the list.
     If True, a random asset configuration is selected for each spawn.
+    .. warning::
+
+        This attribute is deprecated. Use :attr:`~isaaclab.scene.InteractiveSceneCfg.random_heterogenous_spawning`
+        instead.
     """
