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

pyproject.toml

@@ -85,9 +85,8 @@ dependencies = [
     "rsl-rl-lib==5.4.1",  # default RL framework
     "onnxscript>=0.5",
     # ----- newton (default physics engine) -----
-    # Loose bound so the wheel co-resolves with isaacsim's newton[sim]==1.2.0 pin; the
-    # exact PyPI release is forced via [tool.uv].override-dependencies (uv sync only).
-    "newton[sim]>=1.2.0",
+    # Pin the Newton build used by both workspace and wheel installs until the 1.6 release.
+    "newton[sim] @ git+https://github.com/newton-physics/newton.git@24bd863528d6b91137408930d0fbe8fa216ad962",
     # Import and mesh-processing packages used by Newton, including the ones that honoring
     # USD-authored ``physics:approximation`` requires. Keep these explicit instead of selecting
     # newton[importers], whose standalone USD dependency would overlap with usd-exchange.
@@ -394,7 +393,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/replication-names-copies.skip

@@ -0,0 +1 @@
+Dependency pin coverage is documented by the Newton package fragment.

source/isaaclab/test/install_ci/uv_pip/test_uv_pip_install_isaaclab_all_trains_cartpole.py

@@ -15,8 +15,8 @@
          Reinstall AFTER the wheel install: unsafe-best-match re-resolves torch from PyPI to CPU.)
     - (aarch64 only) export LD_PRELOAD=/lib/aarch64-linux-gnu/libgomp.so.1
 Tests:
-    - python -c "import importlib.metadata as m; assert m.version('newton') == '1.5.1'"
-        -> verify the wheel resolves Newton 1.5
+    - python -c "import importlib.metadata as m; assert m.version('newton') == '1.6.0.dev0'"
+        -> verify the wheel resolves the pinned Newton 1.6 development build
     - uv run isaaclab train --rl_library rsl_rl --task Isaac-Cartpole-Direct --num_envs 16
         presets=newton_mjwarp --max_iterations 5; uv run isaaclab train --rl_library rsl_rl
         --task Isaac-Cartpole-Camera-Direct --num_envs 16 presets=newton_mjwarp,newton_renderer --max_iterations 2
@@ -57,11 +57,11 @@ def test_uv_pip_install_isaaclab_all_trains_cartpole(self, isaaclab_root, wheel,
             assert result.returncode == 0, f"uv pip install {wheel}[all] failed:\n{result.stdout}\n{result.stderr}"
 
             result = self.run_in_uv_env(
-                ["python", "-c", "import importlib.metadata as m; assert m.version('newton') == '1.5.1'"],
+                ["python", "-c", "import importlib.metadata as m; assert m.version('newton') == '1.6.0.dev0'"],
                 cwd=isaaclab_root,
             )
             assert result.returncode == 0, (
-                f"isaaclab[all] did not resolve Newton 1.5:\n{result.stdout}\n{result.stderr}"
+                f"isaaclab[all] did not resolve the pinned Newton build:\n{result.stdout}\n{result.stderr}"
             )
 
             # Restore the CUDA build selected for this architecture.

source/isaaclab/test/install_ci/uv_pip/uv-overrides.txt

@@ -1,7 +1,7 @@
 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
 newton-usd-schemas>=0.4.1
 torch==2.11.0
 torchvision==0.26.0

source/isaaclab/test/sim/test_newton_manager_visualization_state.py

@@ -836,8 +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, "rename_builder_labels", lambda *args, **kwargs: None)
+    monkeypatch.setattr(vb, "replicate_builder_mapping", lambda *args, **kwargs: ({}, [], []))
 
     _builder, (shadow_entities, registry_groups) = vb.build_visualization_builder_from_stage_envs(
         stage,

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

@@ -0,0 +1,4 @@
+Changed
+^^^^^^^
+
+* Changed homogeneous Newton cloning to assign labels during replication, avoiding a second full-model pass.

source/isaaclab_newton/changelog.d/root-joint-names.rst

@@ -0,0 +1,5 @@
+Changed
+^^^^^^^
+
+* Changed importer-generated floating-base root joints to use ``{body}_free_joint`` instead of
+  generated ``joint_<n>`` labels, giving replicated environments stable body-derived joint paths.

source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py

@@ -11,7 +11,7 @@
 import numpy as np
 import torch
 import warp as wp
-from newton import GeoType, ModelBuilder, ShapeFlags
+from newton import GeoType, JointType, ModelBuilder, ShapeFlags
 
 from pxr import Usd, UsdGeom, UsdPhysics
 
@@ -148,9 +148,22 @@ def _build_source_builder(
     replace_newton_builder_shape_colors(builder, stage)
     if load_visual_shapes:
         import_builder_visual_material_paths(builder, stage)
+    _name_root_joints_after_their_body(builder)
     return builder
 
 
+def _name_root_joints_after_their_body(builder: ModelBuilder) -> None:
+    """Name importer-generated free root joints after their child bodies, in place."""
+    for index, label in enumerate(builder.joint_label):
+        if not isinstance(label, str) or not label.startswith("joint_") or not label[6:].isdigit():
+            continue
+        if builder.joint_type[index] != JointType.FREE or builder.joint_parent[index] != -1:
+            continue
+        body_label = builder.body_label[builder.joint_child[index]]
+        if isinstance(body_label, str) and body_label.startswith("/"):
+            builder.joint_label[index] = f"{body_label}_free_joint"
+
+
 def _quat_multiply(a: np.ndarray, b: np.ndarray) -> np.ndarray:
     """Hamilton product of xyzw quaternion arrays, broadcast over the leading axes."""
     ax, ay, az, aw = a[..., 0], a[..., 1], a[..., 2], a[..., 3]
@@ -186,19 +199,52 @@ def _invert_xform(xform: Sequence[float] | np.ndarray) -> np.ndarray:
     return np.concatenate([-_quat_rotate(quat_inv, xform[:3]), quat_inv])
 
 
+def _label_groups(builder: ModelBuilder) -> dict[str, list]:
+    """Return every entity-label container owned by a Newton builder."""
+    groups = {
+        name: value for name, value in vars(builder).items() if name.endswith("_label") and isinstance(value, list)
+    }
+    groups["mujoco:equality_constraint_label"] = builder.custom_attributes["mujoco:equality_constraint_label"].values
+    return groups
+
+
+def _rebase_labels(builder: ModelBuilder, source: str, destination: str) -> str:
+    """Make entity labels relative to the nearest templated destination ancestor."""
+    source = source.rstrip("/") or "/"
+    destination = destination.rstrip("/") or "/"
+    prefix, _, destination_name = destination.rpartition("/")
+    if "{}" in destination_name:
+        prefix, destination_name = destination, ""
+    for labels in _label_groups(builder).values():
+        for index, label in enumerate(labels):
+            if not isinstance(label, str) or not label or not label.startswith("/"):
+                continue
+            suffix = clone_path.relative_to(label, source)
+            if suffix is None:
+                suffix = label[len(source) :] if label.startswith(source + "_") else None
+            if suffix is None:
+                raise ValueError(f"Newton label {label!r} is outside clone source {source!r}.")
+            labels[index] = (destination_name + suffix).lstrip("/")
+            if not labels[index]:
+                raise ValueError(f"Newton label {label!r} cannot be prefixed by destination {destination!r}.")
+    return prefix
+
+
 def replicate_builder_mapping(
     builder: ModelBuilder,
     sources: Sequence[str],
     mapping: torch.Tensor,
     positions: torch.Tensor,
     quaternions: torch.Tensor,
     source_builders: dict[str, ModelBuilder],
+    destinations: Sequence[str] | None = None,
+    env_ids: torch.Tensor | None = None,
     *,
     source_site_indices: dict[int, dict[str, list[int]]] | None = None,
     env_root_sites: dict[str, wp.transform] | None = None,
     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], list[tuple[str, int]]]:
+    """Replicate source builders, naming homogeneous copies at their destinations."""
     source_site_indices = source_site_indices or {}
     env_root_sites = env_root_sites or {}
     num_worlds = mapping.size(1)
@@ -215,6 +261,8 @@ def replicate_builder_mapping(
         and num_worlds > 0
         and bool(mapping.all().item())
         and not per_world_builder_hooks
+        and bool(destinations)
+        and env_ids is not None
     )
     if can_batch:
         source_builder = source_builders[sources[0]]
@@ -234,14 +282,24 @@ 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)
+
+        label_groups = _label_groups(source_builder)
+        original_labels = {name: list(labels) for name, labels in label_groups.items()}
+        try:
+            prefix = _rebase_labels(source_builder, sources[0], destinations[0])
+            prefixes = [prefix.format(env_id) for env_id in env_ids.tolist()]
+            builder.replicate(source_builder, num_worlds, xforms=xforms, label_prefixes=prefixes)
+        finally:
+            for name, labels in original_labels.items():
+                label_groups[name][:] = labels
 
         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
+        bindings = rename_builder_labels(builder, sources, destinations, env_ids, mapping, skip_entity_labels=True)
+        return local_site_map, world_xforms, bindings
 
     source_world_indices = mapping.to(dtype=torch.int64).argmax(dim=1).tolist()
 
@@ -273,35 +331,22 @@ 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)
             local_site_map.setdefault(label, [[] for _ in range(num_worlds)])[col].append(site_idx)
-
         for row in rows_per_world[col]:
             source_builder = source_builders[sources[row]]
             offset = builder.shape_count
             builder.add_builder(source_builder, xform=source_xforms[row, col])
-
             for label, source_shape_indices in source_site_indices.get(id(source_builder), {}).items():
                 local_indices = local_site_map.setdefault(label, [[] for _ in range(num_worlds)])[col]
                 local_indices.extend(offset + shape_idx for shape_idx in source_shape_indices)
-
         for hook in per_world_builder_hooks:
             hook(builder, col, xform_rows[col][:3], xform_rows[col][3:])
         builder.end_world()
 
-    return local_site_map, world_xforms
-
-
-_BUILTIN_LABEL_TYPES: tuple[str, ...] = (
-    "body",
-    "joint",
-    "shape",
-    "articulation",
-    "constraint_mimic",
-    "equality_constraint",
-)
+    bindings = rename_builder_labels(builder, sources, destinations, env_ids, mapping) if destinations else []
+    return local_site_map, world_xforms, bindings
 
 
 def rename_builder_labels(
@@ -310,6 +355,8 @@ 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."""
     fabric_body_bindings: list[tuple[str, int]] = []
@@ -321,10 +368,7 @@ def rename_builder_labels(
         world_cols = torch.nonzero(mapping[source_index], as_tuple=True)[0].tolist()
         # Pre-normalize the destination roots
         destination = destinations[source_index]
-        world_roots = {
-            env_id: (destination.format(env_id).rstrip("/") or "/")
-            for env_id in (env_ids_list[col] for col in world_cols)
-        }
+        world_roots = {col: (destination.format(env_ids_list[col]).rstrip("/") or "/") for col in world_cols}
 
         def _rename_pair(values, worlds, src_root=source_root, roots=world_roots, *, collect_body_bindings=False):
             rows = (
@@ -346,14 +390,11 @@ 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 name, labels in vars(builder).items():
+                worlds = getattr(builder, f"{name[:-6]}_world", None) if name.endswith("_label") else None
+                if isinstance(labels, list) and worlds is not None:
+                    _rename_pair(labels, worlds, collect_body_bindings=name == "body_label")
 
         custom_attrs = builder.custom_attributes.values()
         worlds_by_freq = {attr.frequency: attr.values for attr in custom_attrs if attr.references == "world"}

source/isaaclab_newton/isaaclab_newton/cloner/replicate.py

@@ -24,7 +24,6 @@
 from isaaclab_newton.cloner.newton_clone_utils import (
     _restore_visible_colliders_without_visual_shapes,
     build_source_builders,
-    rename_builder_labels,
     replicate_builder_mapping,
 )
 from isaaclab_newton.physics import NewtonManager
@@ -106,13 +105,8 @@ def _build_newton_builder_from_mapping(
     up_axis: str = "Z",
     load_visual_shapes: bool = True,
     global_paths: tuple[str, ...] = (),
-) -> tuple[ModelBuilder, object, dict, list, dict[str, ModelBuilder]]:
-    """Build a Newton model builder from clone mapping inputs.
-
-    Also returns the per-source builders (``{source_path: ModelBuilder}``) so the
-    committing path can retain them for single-model consumers such as the
-    batched Newton IK action.
-    """
+) -> tuple[ModelBuilder, object, dict, list, dict[str, ModelBuilder], list[tuple[str, int]]]:
+    """Build a Newton model builder from clone mapping inputs and retain its source builders."""
     if positions is None:
         positions = torch.zeros((mapping.size(1), 3), device=mapping.device, dtype=torch.float32)
     if quaternions is None:
@@ -170,16 +164,18 @@ def _build_newton_builder_from_mapping(
     global_sites, source_sites, root_sites = NewtonManager._cl_inject_sites(builder, source_builders)
 
     replicate_args = (builder, sources, mapping, positions, quaternions, source_builders)
-    local_site_map, world_xforms = replicate_builder_mapping(
+    local_site_map, world_xforms, fabric_body_bindings = replicate_builder_mapping(
         *replicate_args,
+        destinations,
+        env_ids,
         source_site_indices=source_sites,
         env_root_sites=root_sites,
         per_world_builder_hooks=NewtonManager._per_world_builder_hooks,
     )
 
     site_index_map = {label: (idx, None) for label, idx in global_sites.items()}
     site_index_map.update((label, (None, per_world)) for label, per_world in local_site_map.items())
-    return builder, stage_info, site_index_map, world_xforms, source_builders
+    return builder, stage_info, site_index_map, world_xforms, source_builders, fabric_body_bindings
 
 
 def _renderer_wants_visual_shapes() -> bool:
@@ -309,19 +305,20 @@ def _merged_mapping(self) -> _MappingBatch:
     def replicate(self) -> tuple[ModelBuilder, object, dict]:
         """Build the Newton model builder from queued mappings and optionally publish it."""
         sources, destinations, env_ids, mapping, positions, quaternions = self._merged_mapping()
-        builder, stage_info, site_index_map, world_xforms, source_builders = _build_newton_builder_from_mapping(
-            stage=self.stage,
-            sources=sources,
-            destinations=destinations,
-            env_ids=env_ids,
-            mapping=mapping,
-            positions=positions,
-            quaternions=quaternions,
-            up_axis=self.up_axis,
-            load_visual_shapes=self.load_visual_shapes,
-            global_paths=self._global_paths,
+        builder, stage_info, site_index_map, world_xforms, source_builders, fabric_body_bindings = (
+            _build_newton_builder_from_mapping(
+                stage=self.stage,
+                sources=sources,
+                destinations=destinations,
+                env_ids=env_ids,
+                mapping=mapping,
+                positions=positions,
+                quaternions=quaternions,
+                up_axis=self.up_axis,
+                load_visual_shapes=self.load_visual_shapes,
+                global_paths=self._global_paths,
+            )
         )
-        fabric_body_bindings = rename_builder_labels(builder, sources, destinations, env_ids, mapping)
         if self.commit_to_manager:
             NewtonManager._cl_site_index_map = site_index_map
             NewtonManager._cl_fabric_body_bindings = fabric_body_bindings

source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py

@@ -1944,7 +1944,7 @@ def instantiate_builder_from_stage(cls):
             quaternions = torch.tensor([quat for _, quat in poses], dtype=torch.float32)
             mapping = torch.ones((1, len(env_paths)), dtype=torch.bool)
             replicate_args = (builder, (proto_path,), mapping, positions, quaternions, source_builders)
-            local_site_map, world_xforms = replicate_builder_mapping(
+            local_site_map, world_xforms, _ = replicate_builder_mapping(
                 *replicate_args,
                 source_site_indices=source_site_indices,
                 env_root_sites=env_root_sites,

source/isaaclab_newton/isaaclab_newton/physics/visualization_builder.py

@@ -20,7 +20,6 @@
 from isaaclab_newton.cloner.newton_clone_utils import (
     _restore_visible_colliders_without_visual_shapes,
     build_source_builders,
-    rename_builder_labels,
     replicate_builder_mapping,
 )
 from isaaclab_newton.physics.visualization_deformables import add_shadow_deformables_to_builder
@@ -148,8 +147,7 @@ def build_visualization_builder_from_stage_envs(
         schema_resolvers,
         ignore_paths=source_deformable_ignore_paths or None,
     )
-    replicate_builder_mapping(builder, sources, mapping, positions, quaternions, source_builders)
-    rename_builder_labels(builder, sources, destinations, env_ids, mapping)
+    replicate_builder_mapping(builder, sources, mapping, positions, quaternions, source_builders, destinations, env_ids)
     shadow_entities, registry_groups = add_shadow_deformables_to_builder(
         builder, stage, env_paths, device=device, entries=deformable_entries, clone_plan=clone_plan
     )