Restore the two-tuple contract of replicate_builder_mapping

Retaining import provenance made replicate_builder_mapping return the
world-0 landing offsets as a third element against its declared two-tuple,
which broke four develop cloner tests and every caller that unpacks two
values. The offsets are recorded only on request by a private
implementation; replicate_builder_mapping keeps its contract and the two
provenance callers use replicate_builder_mapping_with_provenance.

Author: hujc7

source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py

@@ -253,7 +253,7 @@ def _invert_xform(xform: Sequence[float] | np.ndarray) -> np.ndarray:
     return np.concatenate([-_quat_rotate(quat_inv, xform[:3]), quat_inv])
 
 
-def replicate_builder_mapping(
+def _replicate_builder_mapping(
     builder: ModelBuilder,
     sources: Sequence[str],
     mapping: torch.Tensor,
@@ -264,8 +264,13 @@ def replicate_builder_mapping(
     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."""
+    record_offsets: bool = False,
+) -> tuple[dict[str, list[list[int]]], list[wp.transform], dict[int, tuple[int, int, int]]]:
+    """Replicate source builders into per-env Newton worlds.
+
+    With ``record_offsets`` the (body, shape, joint) counts of ``builder`` at the moment each source
+    lands in world 0 are recorded per source row; they locate a source's entities in the model.
+    """
     source_site_indices = source_site_indices or {}
     env_root_sites = env_root_sites or {}
     num_worlds = mapping.size(1)
@@ -300,7 +305,7 @@ def replicate_builder_mapping(
         base_shape = builder.shape_count
         stride = source_builder.shape_count
         # World 0's single source lands right after the global content.
-        world0_offsets = {0: (builder.body_count, builder.shape_count, builder.joint_count)}
+        world0_offsets = {0: (builder.body_count, builder.shape_count, builder.joint_count)} if record_offsets else {}
         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)
@@ -351,7 +356,7 @@ def replicate_builder_mapping(
         for row in rows_per_world[col]:
             source_builder = source_builders[sources[row]]
             offset = builder.shape_count
-            if col == 0:
+            if col == 0 and record_offsets:
                 world0_offsets[row] = (builder.body_count, builder.shape_count, builder.joint_count)
             builder.add_builder(source_builder, xform=source_xforms[row, col])
 
@@ -376,6 +381,65 @@ def replicate_builder_mapping(
 )
 
 
+def replicate_builder_mapping(
+    builder: ModelBuilder,
+    sources: Sequence[str],
+    mapping: torch.Tensor,
+    positions: torch.Tensor,
+    quaternions: torch.Tensor,
+    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,
+    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."""
+    local_site_map, world_xforms, _ = _replicate_builder_mapping(
+        builder,
+        sources,
+        mapping,
+        positions,
+        quaternions,
+        source_builders,
+        source_site_indices=source_site_indices,
+        env_root_sites=env_root_sites,
+        per_world_builder_hooks=per_world_builder_hooks,
+    )
+    return local_site_map, world_xforms
+
+
+def replicate_builder_mapping_with_provenance(
+    builder: ModelBuilder,
+    sources: Sequence[str],
+    mapping: torch.Tensor,
+    positions: torch.Tensor,
+    quaternions: torch.Tensor,
+    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,
+    per_world_builder_hooks: Sequence[Callable[[ModelBuilder, int, list[float], list[float]], None]] = (),
+) -> tuple[dict[str, list[list[int]]], list[wp.transform], dict[int, tuple[int, int, int]]]:
+    """Same as :func:`replicate_builder_mapping`, additionally returning each source's landing offsets.
+
+    The third element maps a source row to the (body, shape, joint) index at which that source's
+    entities start in world 0, which :func:`merge_import_results` needs to lift the importer's
+    source-local provenance onto the replicated model.
+    """
+    return _replicate_builder_mapping(
+        builder,
+        sources,
+        mapping,
+        positions,
+        quaternions,
+        source_builders,
+        source_site_indices=source_site_indices,
+        env_root_sites=env_root_sites,
+        per_world_builder_hooks=per_world_builder_hooks,
+        record_offsets=True,
+    )
+
+
 def rename_builder_labels(
     builder: ModelBuilder,
     sources: Sequence[str],

source/isaaclab_newton/isaaclab_newton/cloner/replicate.py

@@ -26,7 +26,7 @@
     build_source_builders_with_provenance,
     merge_import_results,
     rename_builder_labels,
-    replicate_builder_mapping,
+    replicate_builder_mapping_with_provenance,
 )
 from isaaclab_newton.physics import NewtonManager
 from isaaclab_newton.renderers.visual_material import import_builder_visual_material_paths
@@ -170,7 +170,7 @@ 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, world0_offsets = replicate_builder_mapping(
+    local_site_map, world_xforms, world0_offsets = replicate_builder_mapping_with_provenance(
         *replicate_args,
         source_site_indices=source_sites,
         env_root_sites=root_sites,

source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py

@@ -100,7 +100,7 @@ def _paused_gc():
 from isaaclab_newton.cloner.newton_clone_utils import (
     _restore_visible_colliders_without_visual_shapes,
     merge_import_results,
-    replicate_builder_mapping,
+    replicate_builder_mapping_with_provenance,
 )
 from isaaclab_newton.physics.featherstone_manager_cfg import FeatherstoneSolverCfg
 from isaaclab_newton.physics.mjwarp_manager_cfg import MJWarpSolverCfg
@@ -1952,7 +1952,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, world0_offsets = replicate_builder_mapping(
+            local_site_map, world_xforms, world0_offsets = replicate_builder_mapping_with_provenance(
                 *replicate_args,
                 source_site_indices=source_site_indices,
                 env_root_sites=env_root_sites,