From 9902c1826497243121cabd7137b99f20100c7e9a Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 6 May 2026 21:40:40 +0000 Subject: [PATCH 1/4] Rename per-env labels in Newton physics replication After Newton's add_builder copies proto rows into the merged builder for each environment, every row's label still references the source proto path. Newton tracks env identity in *_world int companion columns; IsaacLab keys data flow off USD prim paths, so each row's label needs to be rewritten to its per-env destination path. _rename_builder_labels now walks two kinds of label-bearing columns: * Built-in label arrays (body, joint, shape, articulation, constraint_mimic, equality_constraint), each paired with its matching *_world int column. * String-typed custom-attribute columns (e.g. mujoco:tendon_label) paired with a references="world" companion at the same frequency. Any future solver-registered string column at a frequency that has a references="world" companion is handled automatically. The prefix match uses a path-separator boundary (src_path.rstrip("/") + "/") so a source path that is a string prefix of another (/Sources/protoA vs /Sources/protoAB) does not cross-contaminate when both feed the same envs. The function raises ValueError if the parallel arrays of labels and worlds have mismatched lengths, instead of silently truncating. Adds test_rename_builder_labels.py with 10 cases covering both passes, multi-source prefix-overlap regression, sparse env ids, multi-frequency, multiple string columns at one frequency, and empty-values pass-through. Stacks on top of #5523 (Newton v1.2.0rc2 bump). The Newton 1.2 release also includes upstream newton-physics/newton#2659 which scopes parse_usd's custom-frequency walk natively, removing the cross-source MjcTendon contamination IsaacLab previously had to work around at the framework layer. --- .../jichuanh-newton-replicate-tendon-fix.rst | 14 + .../cloner/newton_replicate.py | 58 +++- .../test/cloner/test_rename_builder_labels.py | 279 ++++++++++++++++++ 3 files changed, 341 insertions(+), 10 deletions(-) create mode 100644 source/isaaclab_newton/changelog.d/jichuanh-newton-replicate-tendon-fix.rst create mode 100644 source/isaaclab_newton/test/cloner/test_rename_builder_labels.py diff --git a/source/isaaclab_newton/changelog.d/jichuanh-newton-replicate-tendon-fix.rst b/source/isaaclab_newton/changelog.d/jichuanh-newton-replicate-tendon-fix.rst new file mode 100644 index 000000000000..12a62ab4d414 --- /dev/null +++ b/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. diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py b/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py index 544756858d51..f2924faf97a9 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py @@ -127,6 +127,15 @@ 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 + 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. + Args: builder: Newton model builder to update in-place. sources: Source prim root paths. @@ -136,21 +145,50 @@ def _rename_builder_labels( """ # per-source, per-world renaming (strict prefix swap), compact style preserved for i, src_path in enumerate(sources): - src_prefix_len = len(src_path.rstrip("/")) + # 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 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} - for t in ("body", "joint", "shape", "articulation"): - labels = getattr(builder, f"{t}_label", None) - if labels is None: - labels = getattr(builder, f"{t}_key") - worlds_arr = getattr(builder, f"{t}_world") - for k, w in enumerate(worlds_arr): - world_id = int(w) - if world_id in world_roots and labels[k].startswith(src_path): - labels[k] = swap(labels[k], world_roots[world_id]) + 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)): + 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]) + + # Pass 1: built-in label arrays. Each has a paired ``*_world`` int column. + for t in ("body", "joint", "shape", "articulation", "constraint_mimic", "equality_constraint"): + labels = getattr(builder, f"{t}_label", None) or getattr(builder, f"{t}_key", None) + worlds_arr = getattr(builder, f"{t}_world", None) + if labels is None or worlds_arr is None: + continue + _rename_pair(labels, worlds_arr) + + # Pass 2: string-typed custom-attribute columns (e.g. ``mujoco:tendon_label``) + # paired with a world companion declared via ``references="world"``. Index + # world companions by frequency for O(1) lookup, then walk the str columns. + custom = builder.custom_attributes + world_by_freq: dict[str, ModelBuilder.CustomAttribute] = {} + for attr in custom.values(): + if getattr(attr, "references", None) == "world": + world_by_freq[attr.frequency] = attr + for attr in custom.values(): + if attr.dtype is not str: + continue + world_attr = world_by_freq.get(attr.frequency) + if world_attr is None: + continue + values = attr.values + worlds = world_attr.values + if not values or not worlds: + continue + _rename_pair(values, worlds) def newton_physics_replicate( diff --git a/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py b/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py new file mode 100644 index 000000000000..dbc710af6155 --- /dev/null +++ b/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py @@ -0,0 +1,279 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Unit tests for ``_rename_builder_labels``. + +Covers both passes of the rewrite: + + * Pass 1 — built-in label arrays (``body``, ``joint``, ``shape``, + ``articulation``, ``constraint_mimic``). + * Pass 2 — any string-typed custom-attribute column whose frequency declares a + sibling ``references="world"`` companion (e.g. ``mujoco:tendon_label``). + +The contract under test: every label whose row maps to a world in ``env_ids`` +and whose value starts with the source root is rewritten to the destination +template's per-env path; everything else is left alone. +""" + +import unittest + +import newton +import torch +from isaaclab_newton.cloner.newton_replicate import _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: + labels = getattr(builder, f"{t}_label") + worlds_arr = getattr(builder, f"{t}_world") + for w in worlds: + labels.append(f"{src_path}/{t}_{w}") + worlds_arr.append(w) + + +def _inject_tendon_strings(builder: newton.ModelBuilder, src_path: str, worlds: list[int]) -> None: + """Append synthetic ``mujoco:tendon_label`` + ``mujoco:tendon_world`` rows.""" + label_attr = builder.custom_attributes["mujoco:tendon_label"] + world_attr = builder.custom_attributes["mujoco:tendon_world"] + if label_attr.values is None: + label_attr.values = [] + if world_attr.values is None: + world_attr.values = [] + for w in worlds: + label_attr.values.append(f"{src_path}/Tendon_{w}") + world_attr.values.append(w) + builder._custom_frequency_counts[_TENDON_FREQ] = builder._custom_frequency_counts.get(_TENDON_FREQ, 0) + len(worlds) + + +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"), _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.""" + + def setUp(self): + self.worlds = [0, 1, 2] + self.env_ids = torch.tensor(self.worlds, dtype=torch.int32) + self.mapping = torch.ones(1, len(self.worlds), dtype=torch.bool) + + 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"): + 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)}", + 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) + labels = b.custom_attributes["mujoco:tendon_label"].values + worlds_arr = b.custom_attributes["mujoco:tendon_world"].values + 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_/``.""" + 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"): + 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_non_path_string_left_untouched(self): + """Strings that don't start with ``src_path`` must pass through unchanged.""" + b = _make_builder_with_entries(self.worlds) + # Inject one tendon row whose label is an opaque identifier, not a path. + b.custom_attributes["mujoco:tendon_label"].values.append("named_tendon") + b.custom_attributes["mujoco:tendon_world"].values.append(self.worlds[0]) + self._rename(b) + self.assertEqual(b.custom_attributes["mujoco:tendon_label"].values[-1], "named_tendon") + + def test_world_outside_env_ids_left_untouched(self): + """A row whose world is not in ``env_ids`` must keep its original label.""" + b = _make_builder_with_entries(self.worlds) + # Inject one extra row tagged with a world id not present in env_ids. + b.body_label.append(f"{_SRC}/body_99") + b.body_world.append(99) + self._rename(b) + self.assertEqual(b.body_label[-1], f"{_SRC}/body_99") + + def test_sparse_env_ids(self): + """Non-contiguous ``env_ids`` (e.g. [10, 20, 30]) must rewrite using the right per-env root.""" + worlds = [10, 20, 30] + 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.""" + + def setUp(self): + self.worlds = [0, 1] + self.env_ids = torch.tensor(self.worlds, dtype=torch.int32) + self.mapping = torch.ones(1, len(self.worlds), dtype=torch.bool) + + def _register_synthetic_freq(self, builder, freq_name, world_attr_name, str_attr_names): + """Register a ``syn:`` frequency with one world int column and N string columns.""" + freq = f"syn:{freq_name}" + builder.add_custom_frequency(newton.ModelBuilder.CustomFrequency(name=freq_name, namespace="syn")) + builder.add_custom_attribute( + newton.ModelBuilder.CustomAttribute( + name=world_attr_name, + frequency=freq, + dtype=int, + default=0, + namespace="syn", + references="world", + ) + ) + for n in str_attr_names: + builder.add_custom_attribute( + newton.ModelBuilder.CustomAttribute( + name=n, + frequency=freq, + dtype=str, + default="", + namespace="syn", + ) + ) + + def _populate(self, builder, freq, world_attr_name, str_attr_names, worlds): + wa = builder.custom_attributes[f"syn:{world_attr_name}"] + if wa.values is None: + wa.values = [] + for w in worlds: + wa.values.append(w) + for n in str_attr_names: + sa = builder.custom_attributes[f"syn:{n}"] + if sa.values is None: + sa.values = [] + for w in worlds: + sa.values.append(f"{_SRC}/{n}_{w}") + builder._custom_frequency_counts[freq] = builder._custom_frequency_counts.get(freq, 0) + len(worlds) + + def test_two_coexisting_custom_frequencies(self): + """Each registered ``references='world'`` companion must drive its own frequency's str columns.""" + b = newton.ModelBuilder() + self._register_synthetic_freq(b, "freqA", "freqA_world", ["freqA_label"]) + self._register_synthetic_freq(b, "freqB", "freqB_world", ["freqB_label"]) + self._populate(b, "syn:freqA", "freqA_world", ["freqA_label"], self.worlds) + self._populate(b, "syn:freqB", "freqB_world", ["freqB_label"], self.worlds) + _rename_builder_labels(b, [_SRC], [_DST], self.env_ids, self.mapping) + for n in ("freqA_label", "freqB_label"): + wa = b.custom_attributes[f"syn:{n.split('_')[0]}_world"].values + sa = b.custom_attributes[f"syn:{n}"].values + for k, w in enumerate(wa): + self.assertEqual(sa[k], f"/World/envs/env_{int(w)}/{n}_{int(w)}") + + def test_multiple_string_columns_at_one_frequency(self): + """Two str columns sharing one frequency must both be rewritten using the shared world companion.""" + b = newton.ModelBuilder() + self._register_synthetic_freq(b, "freqA", "freqA_world", ["freqA_label", "freqA_alt"]) + self._populate(b, "syn:freqA", "freqA_world", ["freqA_label", "freqA_alt"], self.worlds) + _rename_builder_labels(b, [_SRC], [_DST], self.env_ids, self.mapping) + wa = b.custom_attributes["syn:freqA_world"].values + for n in ("freqA_label", "freqA_alt"): + sa = b.custom_attributes[f"syn:{n}"].values + for k, w in enumerate(wa): + self.assertEqual(sa[k], f"/World/envs/env_{int(w)}/{n}_{int(w)}") + + def test_empty_values_pass_through(self): + """A registered-but-empty string column must not crash the rename pass.""" + b = newton.ModelBuilder() + self._register_synthetic_freq(b, "freqA", "freqA_world", ["freqA_label"]) + # values stay None (registered, never populated) + _rename_builder_labels(b, [_SRC], [_DST], self.env_ids, self.mapping) + # Fully populate after the no-op rename: ensures the early-return guard didn't corrupt state. + self._populate(b, "syn:freqA", "freqA_world", ["freqA_label"], self.worlds) + self.assertEqual(len(b.custom_attributes["syn:freqA_label"].values), len(self.worlds)) + + +class TestRenameMultiSource(unittest.TestCase): + """Multi-source handling must not cross-contaminate when source paths share a string prefix.""" + + def test_prefix_overlap_does_not_cross_contaminate(self): + """Sources whose paths share a string prefix and that both feed the same envs must not cross-rename. + + Common IL pattern: a robot proto and an object proto both feed every env. If the two source + paths share a string prefix (``/Sources/protoA`` and ``/Sources/protoAB``), iter 0 + (``src=protoA``) sees the protoAB rows for the same world ids it owns and would over-match + them under a non-boundary ``startswith``. The world-id guard alone does not catch this case + because both sources contribute to the same set of worlds. + """ + sources = ["/Sources/protoA", "/Sources/protoAB"] + # 2 envs, both fed by both sources. + env_ids = torch.tensor([0, 1], dtype=torch.int32) + mapping = torch.tensor([[1, 1], [1, 1]], dtype=torch.bool) + b = newton.ModelBuilder() + SolverMuJoCo.register_custom_attributes(b) + # One body row from each source per env: 4 rows total, world ids interleaved. + b.body_label.extend( + [ + f"{sources[0]}/body", # row 0: protoA, world 0 + f"{sources[1]}/body", # row 1: protoAB, world 0 + f"{sources[0]}/body", # row 2: protoA, world 1 + f"{sources[1]}/body", # row 3: protoAB, world 1 + ] + ) + b.body_world.extend([0, 0, 1, 1]) + _rename_builder_labels(b, sources, ["/World/envs/env_{}", "/World/envs/env_{}"], env_ids, mapping) + # Each row must end up under its own per-env root with the suffix preserved verbatim. + # Without the "/" boundary on ``startswith``, iter 0 (src=protoA) would match rows 1 and 3 + # because ``/Sources/protoAB/body``.startswith(``/Sources/protoA``) is True, rewriting them + # to ``/World/envs/env_/B/body`` (wrong suffix). + self.assertEqual(b.body_label[0], "/World/envs/env_0/body") + self.assertEqual(b.body_label[1], "/World/envs/env_0/body") + self.assertEqual(b.body_label[2], "/World/envs/env_1/body") + self.assertEqual(b.body_label[3], "/World/envs/env_1/body") + + +if __name__ == "__main__": + unittest.main() From 64e88740fb2fa6705c3c289dd256fe04cdbbfa9a Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 12 May 2026 00:11:24 +0000 Subject: [PATCH 2/4] Address review: empty-list guard + equality_constraint coverage * newton_replicate.py: use ``is None`` instead of ``or`` when picking between ``*_label`` and ``*_key`` so an empty-but-defined ``*_label`` list is recognized (the ``or`` fallback over-matched it as falsy). * test_rename_builder_labels.py: add ``equality_constraint`` to the Pass 1 fixture, both per-type loops, and the module docstring so the test surface matches the impl's tuple at line 167. --- .../isaaclab_newton/cloner/newton_replicate.py | 7 ++++++- .../test/cloner/test_rename_builder_labels.py | 10 ++++++---- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py b/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py index f2924faf97a9..46d4f967d51f 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py @@ -163,8 +163,13 @@ def _rename_pair(values, worlds): values[k] = swap(values[k], world_roots[world_id]) # 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"): - labels = getattr(builder, f"{t}_label", None) or getattr(builder, f"{t}_key", None) + labels = getattr(builder, f"{t}_label", None) + if labels is None: + labels = getattr(builder, f"{t}_key", None) worlds_arr = getattr(builder, f"{t}_world", None) if labels is None or worlds_arr is None: continue diff --git a/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py b/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py index dbc710af6155..2fe930f8b520 100644 --- a/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py +++ b/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py @@ -8,7 +8,7 @@ Covers both passes of the rewrite: * Pass 1 — built-in label arrays (``body``, ``joint``, ``shape``, - ``articulation``, ``constraint_mimic``). + ``articulation``, ``constraint_mimic``, ``equality_constraint``). * Pass 2 — any string-typed custom-attribute column whose frequency declares a sibling ``references="world"`` companion (e.g. ``mujoco:tendon_label``). @@ -60,7 +60,9 @@ 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"), _SRC, worlds) + _inject_builtins( + b, ("body", "joint", "shape", "articulation", "constraint_mimic", "equality_constraint"), _SRC, worlds + ) _inject_tendon_strings(b, _SRC, worlds) return b @@ -84,7 +86,7 @@ def _rename(self, builder): 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"): + for t in ("body", "joint", "shape", "articulation", "constraint_mimic", "equality_constraint"): labels = getattr(b, f"{t}_label") worlds_arr = getattr(b, f"{t}_world") for k, w in enumerate(worlds_arr): @@ -111,7 +113,7 @@ def test_all_renamed_labels_share_the_per_env_root(self): 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"): + for t in ("body", "joint", "shape", "articulation", "constraint_mimic", "equality_constraint"): 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 From 3fcefb20e52e3491cdd24ad939422af898f42f2b Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 14 May 2026 17:35:55 +0000 Subject: [PATCH 3/4] Address review: exact-root match, hoist constant, drop banner comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * newton_replicate.py: - Hoist the built-in label types as the module-level ``_BUILTIN_LABEL_TYPES`` constant so tests can import it instead of duplicating the tuple. - Replace the ``src_prefix`` / ``src_prefix_len`` / ``swap`` plumbing with a single ``removeprefix`` step inside ``_rename_pair``. The new form accepts both an exact source-root match (label value equals ``src_path``) and a child path (``src_path + "/..."``), while still rejecting boundary-bleed siblings like ``/Sources/protoAB/x`` when ``src_path`` is ``/Sources/protoA``. Fixes a regression where a label whose value equals the source root was silently dropped after the trailing-slash boundary fix. * test_rename_builder_labels.py: - Import ``_BUILTIN_LABEL_TYPES`` and collapse the three duplicated inline tuples. - Drop the box-drawing banner comments (``# ─── helpers ───``, ``# ─── tests ───``) and the in-class ``# Pass 1`` / ``# Pass 2`` / ``# Cross-pass consistency`` / ``# Guards`` separators — function and class names already partition the file. - Add ``test_label_equal_to_source_root_is_rewritten`` (the exact-match regression). - Add ``test_trailing_slash_on_source_path_is_canonicalized`` (a ``sources=["/X/"]`` argument must rewrite identically to ``"/X"``). - Add ``test_large_world_ids`` ([0, 1_000_000, 2_147_000_000]) documenting that dispatch is dict-keyed (constant cost) and the only hard cap is Newton's int32 ``*_world`` column type. --- .../cloner/newton_replicate.py | 52 ++++++++++----- .../test/cloner/test_rename_builder_labels.py | 63 +++++++++++++------ 2 files changed, 82 insertions(+), 33 deletions(-) diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py b/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py index 46d4f967d51f..1d41aa4acf0b 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py +++ b/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 ``_label`` (or ``_key``) string column +# and a ``_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,28 @@ 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 + # ``removeprefix`` returns the suffix after src_root, or the original + # string when src_root isn't a prefix at all. Two legitimate match + # shapes: suffix == "" (exact root match) or suffix starts with "/" + # (descendant). Anything else is a boundary-bleed sibling like + # "/Sources/protoAB/x" (suffix "B/x" when src_root is "/Sources/protoA") + # — skip it. + suffix = v.removeprefix(src_root) + 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) diff --git a/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py b/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py index 2fe930f8b520..5ecf162fbcab 100644 --- a/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py +++ b/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py @@ -21,7 +21,7 @@ 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" @@ -29,9 +29,6 @@ _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,14 +94,12 @@ 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_/``.""" 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 @@ -121,7 +107,29 @@ def test_all_renamed_labels_share_the_per_env_root(self): 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.""" From e5a2af2d8bcdc53c82a35e8144227e174e24e858 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 14 May 2026 21:18:46 +0000 Subject: [PATCH 4/4] Fix multi-source re-prefix regression in _rename_pair The previous ``removeprefix`` form had a hidden trap: when the source root is not a prefix of the value, ``str.removeprefix`` returns the *original string unchanged*. For an already-rewritten label like ``/World/envs/env_0/body``, that returned string still starts with ``/``, so the boundary check ``suffix and not suffix.startswith("/")`` treated it as a legitimate child and re-prepended the dst root, yielding ``/World/envs/env_0/World/envs/env_0/body`` on the second outer-loop iteration. Gate on ``startswith(src_root)`` before slicing so "no prefix at all" exits cleanly. Boundary-bleed siblings (e.g. ``/Sources/protoAB/x`` when ``src_root`` is ``/Sources/protoA``) are still rejected by the slash-boundary check on the resulting suffix. Caught by ``TestRenameMultiSource::test_prefix_overlap_does_not_cross_contaminate``. --- .../cloner/newton_replicate.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py b/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py index 1d41aa4acf0b..e2751e2274f4 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py @@ -175,13 +175,18 @@ def _rename_pair(values, worlds): world_id = int(worlds[k]) if world_id not in world_roots: continue - # ``removeprefix`` returns the suffix after src_root, or the original - # string when src_root isn't a prefix at all. Two legitimate match - # shapes: suffix == "" (exact root match) or suffix starts with "/" - # (descendant). Anything else is a boundary-bleed sibling like - # "/Sources/protoAB/x" (suffix "B/x" when src_root is "/Sources/protoA") - # — skip it. - suffix = v.removeprefix(src_root) + # 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