Fix Newton visualization ClonePlan handling

source/isaaclab/changelog.d/nblauch-newton-visualization-cloneplan.skip

@@ -0,0 +1 @@
+Internal cloner utility extracted for Newton visualization clone-plan fix.

source/isaaclab/isaaclab/cloner/cloner_utils.py

@@ -76,6 +76,18 @@ def get_suffix(path_expr: str, destination_template: str) -> str | None:
     return None if suffix and not suffix.startswith("/") else suffix
 
 
+def replace_path_prefix(path: str, source_prefix: str, destination_prefix: str) -> str:
+    """Replace ``source_prefix`` in ``path`` with ``destination_prefix`` on a path boundary."""
+    source_prefix = source_prefix.rstrip("/") or "/"
+    destination_prefix = destination_prefix.rstrip("/") or "/"
+    if not path.startswith(source_prefix):
+        return path
+    suffix = path[len(source_prefix) :]
+    if suffix and not suffix.startswith("/"):
+        return path
+    return destination_prefix + suffix
+
+
 def resolve_clone_plan_source(path_expr: str, plan: ClonePlan) -> tuple[str, str, str] | None:
     """Resolve a destination path expression to its row's source path, destination glob, and asset suffix.
 

source/isaaclab/test/sim/test_newton_manager_visualization_state.py

@@ -7,8 +7,7 @@
 
 When the active sim backend is PhysX and a Newton-native visualizer/renderer is in
 use, :meth:`NewtonManager._ensure_visualization_model` must build the manager's
-``_model`` / ``_state_0`` directly from the USD stage (via
-:meth:`NewtonManager._build_visualization_model_from_stage`), and
+``_model`` / ``_state_0`` directly from the USD stage, and
 :meth:`NewtonManager.update_visualization_state` must copy fresh transforms into
 ``_state_0.body_q`` via the new
 :class:`~isaaclab.scene_data.SceneDataProvider`.
@@ -30,6 +29,17 @@ def _reset_newton_manager_state():
     NewtonManager._scene_data_mapping = None
 
 
+def _make_env_stage(num_envs: int = 1):
+    from pxr import Usd, UsdGeom
+
+    stage = Usd.Stage.CreateInMemory()
+    UsdGeom.Xform.Define(stage, "/World")
+    UsdGeom.Xform.Define(stage, "/World/envs")
+    for env_id in range(num_envs):
+        UsdGeom.Xform.Define(stage, f"/World/envs/env_{env_id}")
+    return stage
+
+
 def test_ensure_visualization_model_noop_when_backend_is_newton(monkeypatch):
     """When sim backend is Newton, the manager keeps its own model/state untouched."""
     from isaaclab_newton.physics import NewtonManager
@@ -48,7 +58,9 @@ def test_ensure_visualization_model_builds_from_stage_when_backend_is_physx(monk
 
     _reset_newton_manager_state()
     monkeypatch.setattr(NewtonManager, "_backend_is_newton", classmethod(lambda cls, scene_data_provider=None: False))
-    monkeypatch.setattr(nm, "get_current_stage", lambda *args, **kwargs: SimpleNamespace())
+    monkeypatch.setattr(nm, "get_current_stage", lambda *args, **kwargs: _make_env_stage())
+    monkeypatch.setattr(nm.PhysicsManager, "_sim", SimpleNamespace(get_clone_plan=lambda: SimpleNamespace()))
+    monkeypatch.setattr(nm.PhysicsManager, "_device", "cpu", raising=False)
     monkeypatch.setattr(nm, "replace_newton_shape_colors", lambda model, *a, **kw: 0)
 
     finalize_calls: list[str] = []
@@ -60,12 +72,7 @@ def finalize(self, device):
             finalize_calls.append(device)
             return SimpleNamespace(state=lambda: SimpleNamespace(body_q=None))
 
-    monkeypatch.setattr(
-        NewtonManager,
-        "_build_visualization_model_from_stage",
-        classmethod(lambda cls, stage: _FakeBuilder()),
-    )
-    monkeypatch.setattr(nm.PhysicsManager, "_device", "cpu", raising=False)
+    monkeypatch.setattr(nm, "build_visualization_builder_from_stage_envs", lambda *args, **kwargs: _FakeBuilder())
 
     NewtonManager._ensure_visualization_model()
 
@@ -81,16 +88,13 @@ def test_ensure_visualization_model_empty_builder_logs_and_skips(monkeypatch, ca
 
     _reset_newton_manager_state()
     monkeypatch.setattr(NewtonManager, "_backend_is_newton", classmethod(lambda cls, scene_data_provider=None: False))
-    monkeypatch.setattr(nm, "get_current_stage", lambda *args, **kwargs: SimpleNamespace())
+    monkeypatch.setattr(nm, "get_current_stage", lambda *args, **kwargs: _make_env_stage())
+    monkeypatch.setattr(nm.PhysicsManager, "_sim", SimpleNamespace(get_clone_plan=lambda: SimpleNamespace()))
 
     class _EmptyBuilder:
         body_count = 0
 
-    monkeypatch.setattr(
-        NewtonManager,
-        "_build_visualization_model_from_stage",
-        classmethod(lambda cls, stage: _EmptyBuilder()),
-    )
+    monkeypatch.setattr(nm, "build_visualization_builder_from_stage_envs", lambda *args, **kwargs: _EmptyBuilder())
 
     with caplog.at_level("ERROR"):
         NewtonManager._ensure_visualization_model()
@@ -107,7 +111,9 @@ def test_ensure_visualization_model_populates_num_envs_when_backend_is_physx(mon
 
     _reset_newton_manager_state()
     monkeypatch.setattr(NewtonManager, "_backend_is_newton", classmethod(lambda cls, scene_data_provider=None: False))
-    monkeypatch.setattr(nm, "get_current_stage", lambda *args, **kwargs: SimpleNamespace())
+    monkeypatch.setattr(nm, "get_current_stage", lambda *args, **kwargs: _make_env_stage(num_envs=4))
+    monkeypatch.setattr(nm.PhysicsManager, "_sim", SimpleNamespace(get_clone_plan=lambda: SimpleNamespace()))
+    monkeypatch.setattr(nm.PhysicsManager, "_device", "cpu", raising=False)
     monkeypatch.setattr(nm, "replace_newton_shape_colors", lambda model, *a, **kw: 0)
 
     class _FakeBuilder:
@@ -116,13 +122,7 @@ class _FakeBuilder:
         def finalize(self, device):
             return SimpleNamespace(state=lambda: SimpleNamespace(body_q=None))
 
-    def _fake_build(cls, stage):
-        # Mirror the real shadow-build behaviour: writes the env count discovered on the stage.
-        NewtonManager._num_envs = 4
-        return _FakeBuilder()
-
-    monkeypatch.setattr(NewtonManager, "_build_visualization_model_from_stage", classmethod(_fake_build))
-    monkeypatch.setattr(nm.PhysicsManager, "_device", "cpu", raising=False)
+    monkeypatch.setattr(nm, "build_visualization_builder_from_stage_envs", lambda *args, **kwargs: _FakeBuilder())
 
     NewtonManager._ensure_visualization_model()
 

source/isaaclab_newton/changelog.d/nblauch-newton-visualization-cloneplan.rst

@@ -0,0 +1,4 @@
+Fixed
+^^^^^
+
+* Fixed Newton visualization model construction for heterogeneous ClonePlans so PhysX-backed Newton renderers use each destination asset's selected source and transform.

source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py

@@ -104,6 +104,8 @@ def __init__(self, cfg: MPMObjectCfg):
         queue_newton_physics_replication(cfg)
         self._registry_entry = MPMObjectRegistryEntry(self.cfg)
         SimulationManager._mpm_object_registry.append(self._registry_entry)
+        if add_registered_mpm_objects_to_builder not in SimulationManager._per_world_builder_hooks:
+            SimulationManager._per_world_builder_hooks.append(add_registered_mpm_objects_to_builder)
         self._physics_ready_handle = None
 
     @property

source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py

@@ -0,0 +1,156 @@
+# 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
+
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from typing import Any
+
+import torch
+import warp as wp
+from newton import ModelBuilder, solvers
+
+from pxr import Usd
+
+from isaaclab.cloner.cloner_utils import replace_path_prefix
+
+
+def build_source_builders(
+    stage: Usd.Stage,
+    sources: Sequence[str],
+    create_builder: Callable[[], ModelBuilder],
+    schema_resolvers: Sequence[Any],
+    *,
+    ignore_paths: Sequence[str] | None = None,
+    simplify_meshes: bool = True,
+) -> dict[str, ModelBuilder]:
+    """Build one Newton builder for each clone source prim path."""
+    builders: dict[str, ModelBuilder] = {}
+    for source in sources:
+        builder = create_builder()
+        solvers.SolverMuJoCo.register_custom_attributes(builder)
+        builder.add_usd(
+            stage,
+            root_path=source,
+            load_visual_shapes=True,
+            skip_mesh_approximation=True,
+            schema_resolvers=schema_resolvers,
+            ignore_paths=ignore_paths,
+        )
+        if simplify_meshes:
+            builder.approximate_meshes("convex_hull", keep_visual_shapes=True)
+        builders[source] = builder
+    return builders
+
+
+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]] = (),
+    post_replicate_hooks: Sequence[Callable[[ModelBuilder], None]] = (),
+) -> tuple[dict[str, list[list[int]]], list[wp.transform]]:
+    """Replicate source builders into per-env Newton worlds."""
+    source_site_indices = source_site_indices or {}
+    env_root_sites = env_root_sites or {}
+    num_worlds = mapping.size(1)
+    local_site_map: dict[str, list[list[int]]] = {}
+    world_xforms: list[wp.transform] = []
+    source_world_indices = mapping.to(dtype=torch.int64).argmax(dim=1)
+
+    for col in range(num_worlds):
+        builder.begin_world()
+        world_xform = wp.transform(positions[col], quaternions[col])
+        world_xforms.append(world_xform)
+
+        for label, xform in env_root_sites.items():
+            site_idx = builder.add_site(body=-1, xform=wp.transform_multiply(world_xform, xform), label=label)
+            local_site_map.setdefault(label, [[] for _ in range(num_worlds)])[col].append(site_idx)
+
+        for row in torch.nonzero(mapping[:, col], as_tuple=True)[0].tolist():
+            source_builder = source_builders[sources[int(row)]]
+            offset = builder.shape_count
+            source_col = int(source_world_indices[int(row)])
+            source_xform = wp.transform(positions[source_col], quaternions[source_col])
+            builder.add_builder(
+                source_builder, xform=wp.transform_multiply(world_xform, wp.transform_inverse(source_xform))
+            )
+
+            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, positions[col].tolist(), quaternions[col].tolist())
+        builder.end_world()
+
+    for hook in post_replicate_hooks:
+        hook(builder)
+    return local_site_map, world_xforms
+
+
+_BUILTIN_LABEL_TYPES: tuple[str, ...] = (
+    "body",
+    "joint",
+    "shape",
+    "articulation",
+    "constraint_mimic",
+    "equality_constraint",
+)
+
+
+def rename_builder_labels(
+    builder: ModelBuilder,
+    sources: Sequence[str],
+    destinations: Sequence[str],
+    env_ids: torch.Tensor,
+    mapping: torch.Tensor,
+) -> 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]] = []
+    bound_body_indices: set[int] = set()
+
+    for source_index, source in enumerate(sources):
+        source_root = source.rstrip("/")
+        world_cols = torch.nonzero(mapping[source_index], as_tuple=True)[0].tolist()
+        world_roots = {int(env_ids[col]): destinations[source_index].format(int(env_ids[col])) for col in world_cols}
+
+        def _rename_pair(values, worlds, *, collect_body_bindings: bool = False):
+            for index, (value, world) in enumerate(zip(values, worlds, strict=True)):
+                world_root = world_roots.get(int(world))
+                if isinstance(value, str) and world_root is not None:
+                    renamed_value = replace_path_prefix(value, source_root, world_root)
+                    if renamed_value != value:
+                        values[index] = renamed_value
+                        if collect_body_bindings:
+                            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),
+            (builder.equality_constraint_label, builder.equality_constraint_world, False),
+        ):
+            _rename_pair(labels, worlds, collect_body_bindings=collect_body_bindings)
+
+        custom_attrs = builder.custom_attributes.values()
+        worlds_by_freq = {attr.frequency: attr.values for attr in custom_attrs if attr.references == "world"}
+        for attr in custom_attrs:
+            if attr.dtype is str and attr.values and (worlds := worlds_by_freq.get(attr.frequency)):
+                _rename_pair(attr.values, worlds)
+
+    fabric_body_bindings.extend(
+        (label, index) for index, label in enumerate(builder.body_label) if index not in bound_body_indices
+    )
+    return fabric_body_bindings