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 46d4f967d51f..e2751e2274f4 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,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) 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."""