[Newton] Delegate homogeneous world prefix generation to ModelBuilder.replicate()

pyproject.toml

@@ -394,7 +394,7 @@ override-dependencies = [
     "numpy>=2",
     "mujoco~=3.11.0",
     "mujoco-warp~=3.11.0",
-    "newton[sim]==1.5.1",
+    "newton[sim] @ git+https://github.com/newton-physics/newton.git@24bd863528d6b91137408930d0fbe8fa216ad962",
     # Force the Newton-matched schemas over isaacsim's ==0.2.0 pin.
     "newton-usd-schemas>=0.4.1",
     "torch==2.11.0",

source/isaaclab/changelog.d/brewster-replication-names-copies.rst

@@ -0,0 +1,9 @@
+Added
+^^^^^
+
+* Added :attr:`~isaaclab.cloner.ClonePlan.env_template`, the destination template for one
+  environment. Every row's destination is that template followed by the asset's path below the
+  environment, so it names the part a clone varies while the remainder is shared. It was
+  previously a constructor argument that the plan discarded, leaving a consumer holding a row
+  unable to recover it -- a destination carries no mark of where the environment ends. Backend
+  replication contexts receive it alongside ``global_paths``.

source/isaaclab/isaaclab/cloner/clone_plan.py

@@ -64,6 +64,13 @@ class ClonePlan:
     global_paths: tuple[str, ...] = ()
     """Unique prim paths for scene assets shared by every environment."""
 
+    env_template: str = DEFAULT_ENV_TEMPLATE
+    """Destination template for one environment, with ``"{}"`` for the env id.
+
+    Every row's destination is this template followed by the asset's path below the
+    environment, so this names the part a clone varies from the part it shares.
+    """
+
 
 def grid_transforms(N: int, spacing: float = 1.0, up_axis: str = "z", device="cpu"):
     """Create a centered grid of transforms for ``N`` instances.
@@ -293,6 +300,7 @@ def set_spawn_paths(spawn_cfg: Any, paths: list[str | None]) -> None:
             positions=positions,
             cfg_rows={},
             global_paths=global_paths,
+            env_template=env_template,
         )
 
     # 3) Homogeneous (every cfg is single-variant): emit the simpler env-root plan.
@@ -308,6 +316,7 @@ def set_spawn_paths(spawn_cfg: Any, paths: list[str | None]) -> None:
             positions=positions,
             cfg_rows=cfg_rows,
             global_paths=global_paths,
+            env_template=env_template,
         )
 
     # 4) Heterogeneous: enumerate prototype combos, build per-row mask, mutate spawn paths.
@@ -377,6 +386,7 @@ def validate_combo_tensor(combos: torch.Tensor, name: str, expected_rows: int |
         positions=positions,
         cfg_rows=cfg_rows,
         global_paths=global_paths,
+        env_template=env_template,
     )
 
 
@@ -419,4 +429,5 @@ def clone_plan_from_env_0(
         positions=positions,
         cfg_rows=cfg_rows,
         global_paths=global_paths,
+        env_template=destination,
     )

source/isaaclab/isaaclab/cloner/replicate_session.py

@@ -99,7 +99,7 @@ def replicate(plan: ClonePlan, *, stage: Usd.Stage, replicate_physics: bool = Tr
 
     backend_ctxs: dict[type, Any] = {}
     for BackendCtxCls, row_set in backend_rows.items():
-        ctx = BackendCtxCls(stage, global_paths=plan.global_paths)
+        ctx = BackendCtxCls(stage, global_paths=plan.global_paths, env_template=plan.env_template)
         backend_ctxs[BackendCtxCls] = ctx
         row_list = sorted(row_set)
         ctx.queue_mapping(

source/isaaclab/isaaclab/cloner/usd.py

@@ -12,6 +12,7 @@
 from pxr import Gf, Sdf, Usd, UsdGeom, Vt
 
 from ._fabric_notices import disabled_fabric_change_notifies
+from .cloner_cfg import DEFAULT_ENV_TEMPLATE
 from .path import split
 
 
@@ -30,7 +31,7 @@ class UsdReplicateContext:
 
     replicate_priority = 100
 
-    def __init__(self, stage: Usd.Stage, global_paths: tuple[str, ...] = ()):
+    def __init__(self, stage: Usd.Stage, global_paths: tuple[str, ...] = (), env_template: str = DEFAULT_ENV_TEMPLATE):
         """Initialize the context.
 
         Args:

source/isaaclab/test/cloner/test_replicate_session.py

@@ -35,8 +35,9 @@ class FakeUsdContext:
         replicate_priority = 100
         instances: list["FakeUsdContext"] = []
 
-        def __init__(self, stage, *, global_paths):
+        def __init__(self, stage, *, global_paths, env_template):
             self.global_paths = global_paths
+            self.env_template = env_template
             FakeUsdContext.instances.append(self)
 
         def queue_mapping(self, sources, destinations, env_ids, mask, *, positions=None):

source/isaaclab/test/sim/test_newton_manager_visualization_state.py

@@ -836,7 +836,7 @@ def test_clone_visualization_builder_ignores_non_env_deformables_on_world_import
     monkeypatch.setattr(vb, "_restore_visible_colliders_without_visual_shapes", lambda *args, **kwargs: None)
     monkeypatch.setattr(vb, "import_builder_visual_material_paths", lambda *args, **kwargs: None)
     monkeypatch.setattr(vb, "build_source_builders", lambda *args, **kwargs: {})
-    monkeypatch.setattr(vb, "replicate_builder_mapping", lambda *args, **kwargs: None)
+    monkeypatch.setattr(vb, "replicate_builder_mapping", lambda *args, **kwargs: ({}, [], False))
     monkeypatch.setattr(vb, "rename_builder_labels", lambda *args, **kwargs: None)
 
     _builder, (shadow_entities, registry_groups) = vb.build_visualization_builder_from_stage_envs(

source/isaaclab_newton/changelog.d/brewster-env-root-site-labels.rst

@@ -0,0 +1,7 @@
+Fixed
+^^^^^
+
+* Fixed a bodyless per-environment site carrying the label ``ft_0`` in every environment. Such a
+  site is now registered with the destination template of the clone-plan row that requested it and
+  labelled from the environment it lands in, so it reads e.g. ``/World/envs/env_3/ft_0``. Sites are
+  still resolved by index, so consumers that look them up by label and index are unaffected.

source/isaaclab_newton/changelog.d/brewster-replication-names-copies.rst

@@ -0,0 +1,10 @@
+Changed
+^^^^^^^
+
+* Changed the homogeneous Newton cloning path to let replication name each cloned entity for
+  the environment it lands in, instead of rewriting every replicated label afterwards. The
+  prototype's labels are rebased once -- a few hundred entries -- and
+  :meth:`~newton.ModelBuilder.replicate` is given the per-env roots, replacing a pass over
+  every label in every world that cost 215 ms on ``Isaac-Velocity-Flat-G1`` at 4096
+  environments. The labels are identical either way; a prototype whose labels a per-world
+  prefix cannot spell keeps the previous path.

source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py

@@ -16,6 +16,7 @@
 from pxr import Usd, UsdGeom, UsdPhysics
 
 from isaaclab.cloner import path as clone_path
+from isaaclab.cloner.cloner_cfg import DEFAULT_ENV_TEMPLATE
 from isaaclab.sim.utils.newton_model_utils import replace_newton_builder_shape_colors
 
 from isaaclab_newton.renderers.visual_material import import_builder_visual_material_paths
@@ -186,6 +187,37 @@ def _invert_xform(xform: Sequence[float] | np.ndarray) -> np.ndarray:
     return np.concatenate([-_quat_rotate(quat_inv, xform[:3]), quat_inv])
 
 
+def _site_label(env_root: str | None, label: str) -> str:
+    """Site label, beneath *env_root* when it has one and bare otherwise."""
+    return f"{env_root}/{label}" if env_root else label
+
+
+def _rebase_to_env(builder: ModelBuilder, env_root: str) -> bool:
+    """Rewrite every entity label relative to its environment root, in place.
+
+    Replication makes N copies that differ only in the environment they sit in, so a label is
+    ``<env><within-env>`` and only the first part varies. Returns whether every label could be
+    written that way: one outside the environment, or one naming the environment root itself,
+    has no within-environment part a per-world prefix could carry.
+    """
+    rebased = []
+    for labels in (
+        builder.body_label,
+        builder.joint_label,
+        builder.shape_label,
+        builder.articulation_label,
+        builder.constraint_mimic_label,
+    ):
+        for index, label in enumerate(labels):
+            suffix = clone_path.relative_to(label, env_root) if isinstance(label, str) else None
+            if not suffix:
+                return False
+            rebased.append((labels, index, suffix.lstrip("/")))
+    for labels, index, suffix in rebased:
+        labels[index] = suffix
+    return True
+
+
 def replicate_builder_mapping(
     builder: ModelBuilder,
     sources: Sequence[str],
@@ -195,12 +227,26 @@ def replicate_builder_mapping(
     source_builders: dict[str, ModelBuilder],
     *,
     source_site_indices: dict[int, dict[str, list[int]]] | None = None,
-    env_root_sites: dict[str, wp.transform] | None = None,
+    env_root_sites: dict[str, tuple[wp.transform, str | None]] | None = None,
+    env_ids: torch.Tensor | None = None,
+    env_template: str = DEFAULT_ENV_TEMPLATE,
     per_world_builder_hooks: Sequence[Callable[[ModelBuilder, int, list[float], list[float]], None]] = (),
-) -> tuple[dict[str, list[list[int]]], list[wp.transform]]:
-    """Replicate source builders into per-env Newton worlds."""
+) -> tuple[dict[str, list[list[int]]], list[wp.transform], bool]:
+    """Replicate source builders into per-env Newton worlds.
+
+    Returns the per-env site indices, the per-env world transforms, and whether replication
+    already named each entity for the env it landed in.
+
+    Args:
+        env_root_sites: Site transform and the destination template naming its env, per label.
+        env_ids: Environment ids for the destination worlds. Given with an environment that
+            owns the prototype, replication names each copy and the caller does not have to
+            rewrite the entity labels afterwards.
+        env_template: Destination template for one environment, from the clone plan.
+    """
     source_site_indices = source_site_indices or {}
     env_root_sites = env_root_sites or {}
+    env_ids_list = env_ids.tolist() if env_ids is not None else None
     num_worlds = mapping.size(1)
     local_site_map: dict[str, list[list[int]]] = {}
     positions_np = positions.detach().cpu().numpy().astype(np.float32, copy=False)
@@ -223,8 +269,12 @@ def replicate_builder_mapping(
         # by world_xforms[0] so R_w = world_xform_w * inv(world_xform_0) lands each
         # copy at world_xform_w * xform.
         site_local_indices: dict[str, list[int]] = {}
-        for label, xform in env_root_sites.items():
-            idx = source_builder.add_site(body=-1, xform=wp.transform_multiply(world_xforms[0], xform), label=label)
+        for label, (xform, destination_template) in env_root_sites.items():
+            # Every copy shares one label; the env name is applied after, by ``label_prefixes``
+            # below or by ``rename_builder_labels``. ``can_batch`` makes either one exact.
+            root = sources[0] if destination_template else None
+            site_xform = wp.transform_multiply(world_xforms[0], xform)
+            idx = source_builder.add_site(body=-1, xform=site_xform, label=_site_label(root, label))
             site_local_indices.setdefault(label, []).append(idx)
         for label, indices in source_site_indices.get(id(source_builder), {}).items():
             site_local_indices.setdefault(label, []).extend(indices)
@@ -234,21 +284,30 @@ def replicate_builder_mapping(
         stride = source_builder.shape_count
         source_xform_inv = _invert_xform(xforms_np[0])
         xforms = _compose_world_xforms(positions_np, quaternions_np, source_xform_inv)
-        builder.replicate(source_builder, num_worlds, xforms=xforms)
+
+        # One source populating every world is the shape replication can name itself: rebase the
+        # prototype's labels once -- a few hundred entries -- and let each copy carry its own
+        # env root, instead of rewriting every label in every world afterwards.
+        label_prefixes = None
+        prototype_env = clone_path.match(sources[0], env_template) if env_ids is not None else None
+        if prototype_env is not None and _rebase_to_env(source_builder, env_template.format(prototype_env.instance)):
+            label_prefixes = [env_template.format(env_id) for env_id in env_ids.tolist()]
+        builder.replicate(source_builder, num_worlds, xforms=xforms, label_prefixes=label_prefixes)
 
         for label, local_indices in site_local_indices.items():
             local_site_map[label] = [
                 [base_shape + world * stride + local for local in local_indices] for world in range(num_worlds)
             ]
 
-        return local_site_map, world_xforms
+        return local_site_map, world_xforms, label_prefixes is not None
 
     source_world_indices = mapping.to(dtype=torch.int64).argmax(dim=1).tolist()
 
     # Per-world placements for every env-root site, composed up front so the per-world loop
     # below only indexes rows.
     root_site_xforms = {
-        label: _compose_world_xforms(positions_np, quaternions_np, xform) for label, xform in env_root_sites.items()
+        label: (_compose_world_xforms(positions_np, quaternions_np, xform), destination_template)
+        for label, (xform, destination_template) in env_root_sites.items()
     }
     # Same for the source placements, but only for the occupied ``(row, col)`` pairs of the
     # mapping: composing a dense ``num_rows x num_worlds`` table would blow up on heterogeneous
@@ -274,8 +333,11 @@ def replicate_builder_mapping(
     for col in range(num_worlds):
         builder.begin_world()
 
-        for label, world_site_xforms in root_site_xforms.items():
-            site_idx = builder.add_site(body=-1, xform=world_site_xforms[col], label=label)
+        for label, (world_site_xforms, destination_template) in root_site_xforms.items():
+            # Named here, not by ``rename_builder_labels``: that only rewrites labels in the
+            # worlds the requesting row covers, and an env-root site sits in every world.
+            env_root = destination_template.format(env_ids_list[col]) if destination_template and env_ids_list else None
+            site_idx = builder.add_site(body=-1, xform=world_site_xforms[col], label=_site_label(env_root, label))
             local_site_map.setdefault(label, [[] for _ in range(num_worlds)])[col].append(site_idx)
 
         for row in rows_per_world[col]:
@@ -291,7 +353,7 @@ def replicate_builder_mapping(
             hook(builder, col, xform_rows[col][:3], xform_rows[col][3:])
         builder.end_world()
 
-    return local_site_map, world_xforms
+    return local_site_map, world_xforms, False
 
 
 _BUILTIN_LABEL_TYPES: tuple[str, ...] = (
@@ -310,8 +372,17 @@ def rename_builder_labels(
     destinations: Sequence[str],
     env_ids: torch.Tensor,
     mapping: torch.Tensor,
+    *,
+    skip_entity_labels: bool = False,
 ) -> list[tuple[str, int]]:
-    """Rewrite source-root labels to per-env destination roots and return Fabric body bindings."""
+    """Rewrite source-root labels to per-env destination roots and return Fabric body bindings.
+
+    Args:
+        skip_entity_labels: Whether the entity labels already name the env they are in, as they
+            do when replication was given the per-env prefixes. Only the string custom
+            attributes are rewritten then, since Newton cannot tell which of those name
+            entities.
+    """
     fabric_body_bindings: list[tuple[str, int]] = []
     bound_body_indices: set[int] = set()
     env_ids_list = env_ids.tolist()
@@ -346,14 +417,15 @@ def _rename_pair(values, worlds, src_root=source_root, roots=world_roots, *, col
                         fabric_body_bindings.append((renamed_value, index))
                         bound_body_indices.add(index)
 
-        for labels, worlds, collect_body_bindings in (
-            (builder.body_label, builder.body_world, True),
-            (builder.joint_label, builder.joint_world, False),
-            (builder.shape_label, builder.shape_world, False),
-            (builder.articulation_label, builder.articulation_world, False),
-            (builder.constraint_mimic_label, builder.constraint_mimic_world, False),
-        ):
-            _rename_pair(labels, worlds, collect_body_bindings=collect_body_bindings)
+        if not skip_entity_labels:
+            for labels, worlds, collect_body_bindings in (
+                (builder.body_label, builder.body_world, True),
+                (builder.joint_label, builder.joint_world, False),
+                (builder.shape_label, builder.shape_world, False),
+                (builder.articulation_label, builder.articulation_world, False),
+                (builder.constraint_mimic_label, builder.constraint_mimic_world, False),
+            ):
+                _rename_pair(labels, worlds, collect_body_bindings=collect_body_bindings)
 
         custom_attrs = builder.custom_attributes.values()
         worlds_by_freq = {attr.frequency: attr.values for attr in custom_attrs if attr.references == "world"}