[Newton] Rename per-env labels in physics replication (depends on #5523)

source/isaaclab_newton/changelog.d/jichuanh-newton-replicate-tendon-fix.rst

@@ -0,0 +1,14 @@
+Fixed
+^^^^^
+
+* Fixed per-environment string identifiers (e.g. ``mujoco:tendon_label``)
+  keeping the source proto path after replication.
+  :func:`~isaaclab_newton.cloner.newton_replicate._rename_builder_labels`
+  now also walks string-typed custom-attribute columns whose frequency
+  declares a ``references="world"`` companion, rewriting their per-row
+  source-path prefix to the destination world root in the same pass that
+  handles built-in label arrays. Adds ``constraint_mimic`` and
+  ``equality_constraint`` to that built-in pass for completeness. The
+  prefix match uses a path-separator boundary so a source path that is a
+  string prefix of another (e.g. ``/Sources/protoA`` vs
+  ``/Sources/protoAB``) does not cross-contaminate during the rename.

source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py

@@ -118,6 +118,20 @@ def _build_newton_builder_from_mapping(
     return builder, stage_info, site_index_map
 
 
+# Built-in label arrays that ``_rename_builder_labels`` rewrites in Pass 1.
+# Each type ``t`` has a paired ``<t>_label`` (or ``<t>_key``) string column
+# and a ``<t>_world`` int column on Newton's ``ModelBuilder``. Exposed as a
+# module-level constant so tests can import it instead of duplicating.
+_BUILTIN_LABEL_TYPES: tuple[str, ...] = (
+    "body",
+    "joint",
+    "shape",
+    "articulation",
+    "constraint_mimic",
+    "equality_constraint",
+)
+
+
 def _rename_builder_labels(
     builder: ModelBuilder,
     sources: Sequence[str],
@@ -127,14 +141,13 @@ def _rename_builder_labels(
 ) -> None:
     """Rename builder labels/keys from source roots to destination roots.
 
-    Walks both built-in label arrays (``body``, ``joint``, ``shape``,
-    ``articulation``, ``constraint_mimic``, ``equality_constraint``) and any
+    Walks both built-in label arrays (see :data:`_BUILTIN_LABEL_TYPES`) and any
     string-typed custom-attribute column whose frequency declares a sibling
     world column (``references="world"``).
-    The ``startswith(src_prefix)`` guard makes the rewrite a no-op for strings that
-    are not paths under the source, so non-path custom string columns are passed
-    through untouched and any future solver-registered string column is handled
-    automatically without changes here.
+    The boundary-safe match (exact source root, or source root followed by ``/``)
+    makes the rewrite a no-op for strings that are not paths under the source.
+    Non-path custom string columns are passed through untouched and any future
+    solver-registered string column is handled automatically without changes here.
 
     Args:
         builder: Newton model builder to update in-place.
@@ -145,11 +158,9 @@ def _rename_builder_labels(
     """
     # per-source, per-world renaming (strict prefix swap), compact style preserved
     for i, src_path in enumerate(sources):
-        # Boundary-terminated prefix prevents over-matching when one source path is a
-        # prefix of another (e.g. ``/Sources/protoA`` vs ``/Sources/protoAB``).
-        src_prefix = src_path.rstrip("/") + "/"
-        src_prefix_len = len(src_prefix) - 1  # slice index keeps the leading "/" in the suffix
-        swap = lambda name, new_root: new_root + name[src_prefix_len:]  # noqa: E731
+        # Canonicalize the source root (drop any trailing ``/``) so the
+        # boundary-safe match logic in ``_rename_pair`` is unambiguous.
+        src_root = src_path.rstrip("/")
         world_cols = torch.nonzero(mapping[i], as_tuple=True)[0].tolist()
         # Map Newton world IDs (sequential) to destination paths using env_ids
         world_roots = {int(env_ids[c]): destinations[i].format(int(env_ids[c])) for c in world_cols}
@@ -158,15 +169,33 @@ def _rename_pair(values, worlds):
             if len(values) != len(worlds):
                 raise ValueError(f"label/world column length mismatch: {len(values)} vs {len(worlds)}")
             for k in range(len(values)):
+                v = values[k]
+                if not isinstance(v, str):
+                    continue
                 world_id = int(worlds[k])
-                if world_id in world_roots and isinstance(values[k], str) and values[k].startswith(src_prefix):
-                    values[k] = swap(values[k], world_roots[world_id])
+                if world_id not in world_roots:
+                    continue
+                # Gate on an explicit prefix test before slicing. ``str.removeprefix``
+                # is tempting but conflates "match with empty suffix" and "no match"
+                # (both return a string starting with "/"), so a label already
+                # rewritten in an earlier source-iteration would be re-prepended to
+                # the next iteration's dst root.
+                if not v.startswith(src_root):
+                    continue
+                suffix = v[len(src_root) :]
+                # ``suffix == ""``     -> exact source-root match (rewrite to dst root).
+                # ``suffix[0] == "/"`` -> child path under source.
+                # otherwise           -> boundary-bleed sibling like "/Sources/protoAB/x"
+                #                        when src_root is "/Sources/protoA" -> skip.
+                if suffix and not suffix.startswith("/"):
+                    continue
+                values[k] = world_roots[world_id] + suffix
 
         # Pass 1: built-in label arrays. Each has a paired ``*_world`` int column.
         # Use ``is None`` (not ``or``) so an empty-but-defined ``*_label`` column
         # is recognized — falling through to ``*_key`` would over-match a
         # builder that legitimately exposes both attributes.
-        for t in ("body", "joint", "shape", "articulation", "constraint_mimic", "equality_constraint"):
+        for t in _BUILTIN_LABEL_TYPES:
             labels = getattr(builder, f"{t}_label", None)
             if labels is None:
                 labels = getattr(builder, f"{t}_key", None)

source/isaaclab_newton/test/cloner/test_rename_builder_labels.py

@@ -21,17 +21,14 @@
 
 import newton
 import torch
-from isaaclab_newton.cloner.newton_replicate import _rename_builder_labels
+from isaaclab_newton.cloner.newton_replicate import _BUILTIN_LABEL_TYPES, _rename_builder_labels
 from newton.solvers import SolverMuJoCo
 
 _TENDON_FREQ = "mujoco:tendon"
 _SRC = "/Sources/protoA"
 _DST = "/World/envs/env_{}"
 
 
-# ─── helpers ─────────────────────────────────────────────────────────────────
-
-
 def _inject_builtins(builder: newton.ModelBuilder, types: tuple[str, ...], src_path: str, worlds: list[int]) -> None:
     """Append ``len(worlds)`` synthetic entries to each built-in ``*_label``/``*_world`` pair."""
     for t in types:
@@ -60,16 +57,11 @@ def _make_builder_with_entries(worlds: list[int]) -> newton.ModelBuilder:
     """Builder pre-populated with one row per world for every label class under test."""
     b = newton.ModelBuilder()
     SolverMuJoCo.register_custom_attributes(b)
-    _inject_builtins(
-        b, ("body", "joint", "shape", "articulation", "constraint_mimic", "equality_constraint"), _SRC, worlds
-    )
+    _inject_builtins(b, _BUILTIN_LABEL_TYPES, _SRC, worlds)
     _inject_tendon_strings(b, _SRC, worlds)
     return b
 
 
-# ─── tests ───────────────────────────────────────────────────────────────────
-
-
 class TestRenameBuilderLabels(unittest.TestCase):
     """Both passes rewrite to the same per-env destination pattern."""
 
@@ -81,12 +73,10 @@ def setUp(self):
     def _rename(self, builder):
         _rename_builder_labels(builder, [_SRC], [_DST], self.env_ids, self.mapping)
 
-    # Pass 1 ---------------------------------------------------------------
-
     def test_builtin_labels_rewritten_per_world(self):
         b = _make_builder_with_entries(self.worlds)
         self._rename(b)
-        for t in ("body", "joint", "shape", "articulation", "constraint_mimic", "equality_constraint"):
+        for t in _BUILTIN_LABEL_TYPES:
             labels = getattr(b, f"{t}_label")
             worlds_arr = getattr(b, f"{t}_world")
             for k, w in enumerate(worlds_arr):
@@ -96,8 +86,6 @@ def test_builtin_labels_rewritten_per_world(self):
                     msg=f"{t}_label[{k}] not rewritten correctly",
                 )
 
-    # Pass 2 ---------------------------------------------------------------
-
     def test_tendon_label_string_custom_attr_rewritten(self):
         b = _make_builder_with_entries(self.worlds)
         self._rename(b)
@@ -106,22 +94,42 @@ def test_tendon_label_string_custom_attr_rewritten(self):
         for k, w in enumerate(worlds_arr):
             self.assertEqual(labels[k], f"{_DST.format(int(w))}/Tendon_{int(w)}")
 
-    # Cross-pass consistency ----------------------------------------------
-
     def test_all_renamed_labels_share_the_per_env_root(self):
         """Every label written by either pass must live under ``/World/envs/env_<world>/``."""
         b = _make_builder_with_entries(self.worlds)
         self._rename(b)
         per_world = {int(w): _DST.format(int(w)) + "/" for w in self.env_ids.tolist()}
-        for t in ("body", "joint", "shape", "articulation", "constraint_mimic", "equality_constraint"):
+        for t in _BUILTIN_LABEL_TYPES:
             for label, w in zip(getattr(b, f"{t}_label"), getattr(b, f"{t}_world")):
                 self.assertTrue(label.startswith(per_world[int(w)]), msg=f"{t}: {label!r}")
         tendon_labels = b.custom_attributes["mujoco:tendon_label"].values
         tendon_worlds = b.custom_attributes["mujoco:tendon_world"].values
         for label, w in zip(tendon_labels, tendon_worlds):
             self.assertTrue(label.startswith(per_world[int(w)]), msg=f"tendon: {label!r}")
 
-    # Guards ---------------------------------------------------------------
+    def test_label_equal_to_source_root_is_rewritten(self):
+        """A label whose value is exactly ``src_path`` (no suffix) maps to the env root.
+
+        Newton may tag a proto's own root prim with a label/key whose value equals the
+        proto's source path. Regression: an earlier ``startswith(src_prefix)`` form
+        (where ``src_prefix = src_path + "/"``) silently dropped this case.
+        """
+        b = _make_builder_with_entries(self.worlds)
+        # Append an exact-root row to body_label (any builtin type would do).
+        b.body_label.append(_SRC)
+        b.body_world.append(self.worlds[0])
+        self._rename(b)
+        self.assertEqual(b.body_label[-1], _DST.format(self.worlds[0]))
+
+    def test_trailing_slash_on_source_path_is_canonicalized(self):
+        """``sources=["/Sources/protoA/"]`` (trailing /) must rewrite identically to no slash."""
+        b = _make_builder_with_entries(self.worlds)
+        _rename_builder_labels(b, [f"{_SRC}/"], [_DST], self.env_ids, self.mapping)
+        for t in _BUILTIN_LABEL_TYPES:
+            labels = getattr(b, f"{t}_label")
+            worlds_arr = getattr(b, f"{t}_world")
+            for k, w in enumerate(worlds_arr):
+                self.assertEqual(labels[k], f"{_DST.format(int(w))}/{t}_{int(w)}")
 
     def test_non_path_string_left_untouched(self):
         """Strings that don't start with ``src_path`` must pass through unchanged."""
@@ -153,6 +161,23 @@ def test_sparse_env_ids(self):
         for k, w in enumerate(b.body_world):
             self.assertEqual(b.body_label[k], f"/World/envs/env_{int(w)}/body_{int(w)}")
 
+    def test_large_world_ids(self):
+        """Large/sparse ``env_ids`` round-trip — dispatch is by dict, not array index.
+
+        ``world_roots`` is a dict keyed on the actual world id, so id magnitude
+        does not affect correctness or storage. Cap kept inside ``int32`` since
+        Newton's ``*_world`` columns are typed int32.
+        """
+        worlds = [0, 1_000_000, 2_147_000_000]  # last entry within int32 range
+        b = newton.ModelBuilder()
+        SolverMuJoCo.register_custom_attributes(b)
+        _inject_builtins(b, ("body",), _SRC, worlds)
+        env_ids = torch.tensor(worlds, dtype=torch.int32)
+        mapping = torch.ones(1, len(worlds), dtype=torch.bool)
+        _rename_builder_labels(b, [_SRC], [_DST], env_ids, mapping)
+        for k, w in enumerate(b.body_world):
+            self.assertEqual(b.body_label[k], f"/World/envs/env_{int(w)}/body_{int(w)}")
+
 
 class TestRenamePass2Generality(unittest.TestCase):
     """Pass 2 must generalize across coexisting frequencies and multiple string columns."""