diff --git a/source/isaaclab/changelog.d/nblauch-newton-visualization-cloneplan.skip b/source/isaaclab/changelog.d/nblauch-newton-visualization-cloneplan.skip new file mode 100644 index 000000000000..7e151a33e97e --- /dev/null +++ b/source/isaaclab/changelog.d/nblauch-newton-visualization-cloneplan.skip @@ -0,0 +1 @@ +Internal cloner utility extracted for Newton visualization clone-plan fix. diff --git a/source/isaaclab/isaaclab/cloner/cloner_utils.py b/source/isaaclab/isaaclab/cloner/cloner_utils.py index e905253d5029..2b25a60de3aa 100644 --- a/source/isaaclab/isaaclab/cloner/cloner_utils.py +++ b/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. diff --git a/source/isaaclab/test/sim/test_newton_manager_visualization_state.py b/source/isaaclab/test/sim/test_newton_manager_visualization_state.py index 919d4f986a4c..0b1d2359d363 100644 --- a/source/isaaclab/test/sim/test_newton_manager_visualization_state.py +++ b/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() diff --git a/source/isaaclab_newton/changelog.d/nblauch-newton-visualization-cloneplan.rst b/source/isaaclab_newton/changelog.d/nblauch-newton-visualization-cloneplan.rst new file mode 100644 index 000000000000..282fd857ce1e --- /dev/null +++ b/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. diff --git a/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py index c5007f69cfc1..4866706743d2 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py +++ b/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 diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py b/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py new file mode 100644 index 000000000000..f85a8f6bddd9 --- /dev/null +++ b/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 diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py b/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py index df4daca53148..b19050da70d2 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py @@ -5,22 +5,33 @@ from __future__ import annotations +import re from collections.abc import Sequence -from typing import Any +from typing import TYPE_CHECKING, Any, TypeAlias import torch -import warp as wp -from newton import ModelBuilder, solvers +from newton import ModelBuilder from newton._src.usd.schemas import SchemaResolverNewton, SchemaResolverPhysx from pxr import Usd from isaaclab.cloner.replicate_session import REPLICATION_QUEUE from isaaclab.physics import PhysicsManager -from isaaclab.sim import SimulationContext +from isaaclab_newton.cloner.newton_clone_utils import ( + build_source_builders, + rename_builder_labels, + replicate_builder_mapping, +) from isaaclab_newton.physics import NewtonManager +if TYPE_CHECKING: + _MappingBatch: TypeAlias = tuple[ + tuple[str, ...], tuple[str, ...], torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None + ] +else: + _MappingBatch = tuple + def _build_newton_builder_from_mapping( stage: Usd.Stage, @@ -33,25 +44,7 @@ def _build_newton_builder_from_mapping( up_axis: str = "Z", simplify_meshes: bool = True, ) -> tuple[ModelBuilder, object, dict, list]: - """Build a Newton model builder from clone mapping inputs. - - Args: - stage: USD stage containing source assets. - sources: Source prim paths used for cloning. - destinations: Destination path templates with one ``"{}"`` slot per source row. - env_ids: Environment ids for destination worlds. - mapping: Boolean source-to-environment mapping matrix. - positions: Optional per-environment world positions. - quaternions: Optional per-environment orientations in xyzw order. - up_axis: Up axis for the Newton model builder. - simplify_meshes: Whether to run convex-hull mesh approximation. - - Returns: - Tuple of the populated Newton model builder, stage metadata returned - by ``add_usd``, a site index map for - :attr:`NewtonManager._cl_site_index_map`, and the absolute per-world - transforms for :attr:`NewtonManager._world_xforms`. - """ + """Build a Newton model builder from clone mapping inputs.""" if positions is None: positions = torch.zeros((mapping.size(1), 3), device=mapping.device, dtype=torch.float32) if quaternions is None: @@ -59,8 +52,7 @@ def _build_newton_builder_from_mapping( quaternions[:, 3] = 1.0 schema_resolvers = [SchemaResolverNewton(), SchemaResolverPhysx()] - sim_ctx = SimulationContext.instance() - manager_cls = sim_ctx.physics_manager if sim_ctx is not None else NewtonManager + manager_cls = PhysicsManager._sim.physics_manager builder = manager_cls.create_builder(up_axis=up_axis) stage_info = builder.add_usd( @@ -72,241 +64,41 @@ def _build_newton_builder_from_mapping( # Deformable prim paths are handled by per_world_builder_hooks, not add_usd. # Resolve the regex prim_path patterns to concrete env_0 paths so add_usd # can skip them via ignore_paths. - import re - - _deformable_ignore_paths: list[str] = [] - if hasattr(NewtonManager, "_deformable_registry"): - for entry in NewtonManager._deformable_registry: - pat = re.compile(entry.prim_path.replace(".*", "[^/]*") + "$") - for src_path in sources: - # Check if any prim under this source matches the deformable pattern - prim = stage.GetPrimAtPath(src_path) - if prim.IsValid(): - for child in Usd.PrimRange(prim): - child_path = str(child.GetPath()) - if pat.match(child_path): - _deformable_ignore_paths.append(child_path) - - protos: dict[str, ModelBuilder] = {} - for src_path in sources: - p = manager_cls.create_builder(up_axis=up_axis) - solvers.SolverMuJoCo.register_custom_attributes(p) - p.add_usd( - stage, - root_path=src_path, - load_visual_shapes=True, - skip_mesh_approximation=True, - schema_resolvers=schema_resolvers, - ignore_paths=_deformable_ignore_paths if _deformable_ignore_paths else None, - ) - if simplify_meshes: - p.approximate_meshes("convex_hull", keep_visual_shapes=True) - protos[src_path] = p - - # Inject registered sites into prototypes (and global sites into main builder) - global_sites, proto_sites, world_sites = NewtonManager._cl_inject_sites(builder, protos) - - # Global sites: (int, None) - global_site_map: dict[str, tuple[int, None]] = {label: (idx, None) for label, idx in global_sites.items()} - - # Local sites: per-world sublists, populated in the loop below - num_worlds = mapping.size(1) - local_site_map: dict[str, list[list[int]]] = {} - # Absolute per-world transforms (env-root local-to-world). Consumed by - # FrameView to place non-physics frames (e.g. cameras) relative to each - # cloned env, mirroring the legacy ``_replicate_from_stage`` path. - world_xforms: list[wp.transform] = [] - - # Heterogeneous clone-plan rows spawn their prototype in the first active environment - # for that row, then reuse that prototype for every other active environment. - # NOTE: None is used to indicate that the source does not map to any environment. - source_world_indices = [] - for mapping_row in mapping: - nz = torch.nonzero(mapping_row, as_tuple=True)[0] - source_world_indices.append(int(nz[0]) if nz.numel() > 0 else None) - - # create a separate world for each environment (heterogeneous spawning) - # Newton assigns sequential world IDs (0, 1, 2, ...), so we need to track the mapping - for col, _ in enumerate(env_ids.tolist()): - # begin a new world context (Newton assigns world ID = col) - builder.begin_world() - - world_xform = wp.transform(positions[col], quaternions[col]) - world_xforms.append(world_xform) - - # Per-world bodyless sites are placed in each world's (global) frame. - for label, xform in world_sites.items(): - if label not in local_site_map: - local_site_map[label] = [[] for _ in range(num_worlds)] - site_idx = builder.add_site(body=-1, xform=wp.transform_multiply(world_xform, xform), label=label) - local_site_map[label][col].append(site_idx) - - for row in torch.nonzero(mapping[:, col], as_tuple=True)[0].tolist(): - source = sources[int(row)] - proto = protos[source] - offset = builder.shape_count - - source_world_index = source_world_indices[int(row)] - source_world_xform = wp.transform(positions[source_world_index], quaternions[source_world_index]) - builder.add_builder( - proto, xform=wp.transform_multiply(world_xform, wp.transform_inverse(source_world_xform)) - ) - - # Compute final shape indices for sites in this proto - for label, proto_shape_indices in proto_sites.get(id(proto), {}).items(): - if label not in local_site_map: - local_site_map[label] = [[] for _ in range(num_worlds)] - for proto_shape_idx in proto_shape_indices: - local_site_map[label][col].append(offset + proto_shape_idx) - - # Emit registered MPM particle objects into this Newton world. - if NewtonManager._mpm_object_registry: - from isaaclab_newton.assets.mpm_object.mpm_object import add_registered_mpm_objects_to_builder - - add_registered_mpm_objects_to_builder(builder, col, positions[col].tolist(), quaternions[col].tolist()) - - # Run per-world builder hooks (e.g. deformable body registration). - if hasattr(NewtonManager, "_per_world_builder_hooks"): - for hook in NewtonManager._per_world_builder_hooks: - hook(builder, col, positions[col].tolist(), quaternions[col].tolist()) - - # end the world context - builder.end_world() - - # Run post-replicate hooks (e.g. builder.color() for deformable coloring). - if hasattr(NewtonManager, "_post_replicate_hooks"): - for hook in NewtonManager._post_replicate_hooks: - hook(builder) - - site_index_map = { - **global_site_map, - **{label: (None, per_world) for label, per_world in local_site_map.items()}, - } - - return builder, stage_info, site_index_map, world_xforms - - -# 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], - destinations: Sequence[str], - env_ids: torch.Tensor, - mapping: torch.Tensor, -) -> list[tuple[str, str, int]]: - """Rename builder labels/keys from source roots to destination roots. + deformable_patterns = tuple( + re.compile(entry.prim_path.replace(".*", "[^/]*")) for entry in NewtonManager._deformable_registry + ) + deformable_ignore_paths = [] + if deformable_patterns: + for source in sources: + for child in Usd.PrimRange(stage.GetPrimAtPath(source)): + child_path = str(child.GetPath()) + if any(pattern.fullmatch(child_path) for pattern in deformable_patterns): + deformable_ignore_paths.append(child_path) + + source_builders = build_source_builders( + stage, + sources, + lambda: manager_cls.create_builder(up_axis=up_axis), + schema_resolvers, + ignore_paths=deformable_ignore_paths or None, + simplify_meshes=simplify_meshes, + ) - 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 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. + # Inject registered sites into source builders (and global sites into main builder). + global_sites, source_sites, root_sites = NewtonManager._cl_inject_sites(builder, source_builders) - Args: - builder: Newton model builder to update in-place. - sources: Source prim root paths. - destinations: Destination prim path templates. - env_ids: Environment ids corresponding to mapping columns. - mapping: Boolean source-to-environment mapping matrix. + replicate_args = (builder, sources, mapping, positions, quaternions, source_builders) + local_site_map, world_xforms = replicate_builder_mapping( + *replicate_args, + source_site_indices=source_sites, + env_root_sites=root_sites, + per_world_builder_hooks=NewtonManager._per_world_builder_hooks, + post_replicate_hooks=NewtonManager._post_replicate_hooks, + ) - Returns: - Fabric body binding records as ``(fabric_body_path, body_index)``. - """ - fabric_body_bindings: list[tuple[str, int]] = [] - bound_body_indices: set[int] = set() - - # per-source, per-world renaming (strict prefix swap), compact style preserved - for i, src_path in enumerate(sources): - # 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} - - def _rename_pair(values, worlds, *, collect_body_bindings: bool = False): - 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 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 - renamed_value = world_roots[world_id] + suffix - if collect_body_bindings: - fabric_body_bindings.append((renamed_value, k)) - bound_body_indices.add(k) - values[k] = renamed_value - - # 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 _BUILTIN_LABEL_TYPES: - 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 - _rename_pair(labels, worlds_arr, collect_body_bindings=t == "body") - - # 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) - - for index, label in enumerate(builder.body_label): - if index not in bound_body_indices: - fabric_body_bindings.append((label, index)) - - return fabric_body_bindings + site_index_map = {label: (idx, None) for label, idx in global_sites.items()} + site_index_map.update((label, (None, per_world)) for label, per_world in local_site_map.items()) + return builder, stage_info, site_index_map, world_xforms class NewtonReplicateContext: @@ -342,16 +134,7 @@ def __init__( simplify_meshes = cfg.simplify_meshes if isinstance(cfg, NewtonCfg) else True self.simplify_meshes = simplify_meshes self.commit_to_manager = commit_to_manager - self._queue: list[ - tuple[ - tuple[str, ...], - tuple[str, ...], - torch.Tensor, - torch.Tensor, - torch.Tensor | None, - torch.Tensor | None, - ] - ] = [] + self._queue: list[_MappingBatch] = [] def queue_mapping( self, @@ -388,9 +171,7 @@ def _merge_optional_tensor( raise ValueError(f"Queued Newton mappings must use the same {name} tensor.") return current - def _merged_mapping( - self, - ) -> tuple[tuple[str, ...], tuple[str, ...], torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + def _merged_mapping(self) -> _MappingBatch: """Merge queued mapping batches into the legacy flat mapping shape.""" if not self._queue: raise RuntimeError("Cannot replicate without queued Newton mappings.") @@ -438,7 +219,7 @@ def replicate(self) -> tuple[ModelBuilder, object, dict]: up_axis=self.up_axis, simplify_meshes=self.simplify_meshes, ) - fabric_body_bindings = _rename_builder_labels(builder, sources, destinations, env_ids, mapping) + fabric_body_bindings = rename_builder_labels(builder, sources, destinations, env_ids, mapping) if self.commit_to_manager: NewtonManager._cl_site_index_map = site_index_map NewtonManager._cl_fabric_body_bindings = fabric_body_bindings @@ -492,13 +273,6 @@ def newton_physics_replicate( ctx = NewtonReplicateContext( stage, device=device, up_axis=up_axis, simplify_meshes=simplify_meshes, commit_to_manager=True ) - ctx.queue_mapping( - sources, - destinations, - env_ids, - mapping, - positions=positions, - quaternions=quaternions, - ) + ctx.queue_mapping(sources, destinations, env_ids, mapping, positions=positions, quaternions=quaternions) builder, stage_info, _site_index_map = ctx.replicate() return builder, stage_info diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index af2a2b37e53e..bb5279e7d8be 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -10,11 +10,13 @@ import contextlib import ctypes import logging +import re from abc import abstractmethod from collections.abc import Callable, Iterable, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING +import torch import warp as wp # Load CUDA runtime for relaxed-mode graph capture (RTX-compatible). @@ -34,6 +36,8 @@ from newton.sensors import SensorIMU as NewtonSensorIMU from newton.solvers import SolverBase, SolverKamino, SolverNotifyFlags +from pxr import UsdGeom + from isaaclab.physics import CallbackHandle, PhysicsEvent, PhysicsManager from isaaclab.scene_data import SceneDataBackend, SceneDataFormat, SceneDataProvider from isaaclab.sim.utils.newton_model_utils import replace_newton_shape_colors @@ -42,6 +46,9 @@ from isaaclab.utils.string import resolve_matching_names from isaaclab.utils.timer import Timer +from isaaclab_newton.cloner.newton_clone_utils import replicate_builder_mapping +from isaaclab_newton.physics.visualization_builder import build_visualization_builder_from_stage_envs + from .newton_manager_cfg import NewtonCfg, NewtonShapeCfg if TYPE_CHECKING: @@ -304,6 +311,9 @@ class NewtonManager(PhysicsManager): _cl_site_index_map: dict[str, _SiteEntry] = {} _cl_fabric_body_bindings: list[tuple[str, int]] | None = None _world_xforms: list[wp.transform] | None = None + _deformable_registry: list = [] + _per_world_builder_hooks: list[Callable[[ModelBuilder, int, list[float], list[float]], None]] = [] + _post_replicate_hooks: list[Callable[[ModelBuilder], None]] = [] @classmethod def initialize(cls, sim_context: SimulationContext) -> None: @@ -766,6 +776,9 @@ def clear(cls): NewtonManager._particles_dirty = False NewtonManager._particle_visual_prims = {} NewtonManager._mpm_object_registry = [] + NewtonManager._deformable_registry = [] + NewtonManager._per_world_builder_hooks = [] + NewtonManager._post_replicate_hooks = [] NewtonManager._up_axis = "Z" NewtonManager._scene_data = None NewtonManager._scene_data_mapping = None @@ -902,72 +915,72 @@ def request_extended_contact_attribute(cls, attr: str) -> None: def _cl_inject_sites( cls, main_builder: ModelBuilder, - proto_builders: dict[str, ModelBuilder], + source_builders: dict[str, ModelBuilder], ) -> tuple[dict[str, int], dict[int, dict[str, list[int]]], dict[str, wp.transform]]: - """Inject registered sites into prototype builders before replication. + """Inject registered sites into source builders before replication. - Non-global sites are matched against prototype body labels using + Non-global sites are matched against source builder body labels using :func:`resolve_matching_names` (regex). Global sites (``body_pattern is None``) are added to *main_builder* with ``body=-1``. - Returns proto-local shape indices so that ``newton_replicate`` can + Returns source-builder-local shape indices so that ``newton_replicate`` can compute final indices during replication without a second pattern match. Pending requests are cleared after processing. Args: main_builder: Top-level builder that receives global sites. - proto_builders: ``{src_path: ModelBuilder}`` prototype builders. + source_builders: ``{source_path: ModelBuilder}`` source builders. Returns: - Tuple of ``(global_sites, proto_sites, world_sites)`` where - *global_sites* maps ``{label: main_builder_shape_idx}``, - *proto_sites* maps ``{id(proto): {label: [proto_local_shape_idx, ...]}}``, - and *world_sites* maps ``{label: env_root_relative_transform}``. + Tuple of ``(global_site_indices, source_site_indices, env_root_sites)`` where + *global_site_indices* maps ``{label: main_builder_shape_idx}``, + *source_site_indices* maps ``{id(source_builder): {label: [source_local_shape_idx, ...]}}``, + and *env_root_sites* maps ``{label: env_root_relative_transform}``. """ - global_sites: dict[str, int] = {} - proto_sites: dict[int, dict[str, list[int]]] = {} + global_site_indices: dict[str, int] = {} + source_site_indices: dict[int, dict[str, list[int]]] = {} - world_sites: dict[str, wp.transform] = {} + env_root_sites: dict[str, wp.transform] = {} for (body_pattern, per_world, _xform_key), (label, xform) in cls._cl_pending_sites.items(): if per_world: - world_sites[label] = xform + env_root_sites[label] = xform continue if body_pattern is None: site_idx = main_builder.add_site(body=-1, xform=xform, label=label) - global_sites[label] = site_idx + global_site_indices[label] = site_idx continue any_matched = False - for src_prefix, proto in proto_builders.items(): - body_labels = list(proto.body_label) + for _source_path, source_builder in source_builders.items(): + body_labels = list(source_builder.body_label) matched_indices, matched_names = resolve_matching_names( body_pattern, body_labels, raise_when_no_match=False ) - if not matched_indices: # Pattern has no matches in this prototype + if not matched_indices: # Pattern has no matches in this source builder continue any_matched = True - proto_id = id(proto) + source_builder_id = id(source_builder) site_indices: list[int] = [] for body_idx, body_name in zip(matched_indices, matched_names): site_label = f"{body_name}/{label}" - proto_site_idx = proto.add_site(body=body_idx, xform=xform, label=site_label) - site_indices.append(proto_site_idx) - logger.debug(f"Injected site '{site_label}' into prototype") - proto_sites.setdefault(proto_id, {})[label] = site_indices + source_site_idx = source_builder.add_site(body=body_idx, xform=xform, label=site_label) + site_indices.append(source_site_idx) + logger.debug(f"Injected site '{site_label}' into source builder") + source_site_indices.setdefault(source_builder_id, {})[label] = site_indices if not any_matched: raise ValueError( - f"Site '{label}' with body_pattern '{body_pattern}' matched no prototype bodies " - f"across {len(proto_builders)} prototype(s). " - f"Check that the pattern matches a body label in the prototype builder." + f"Site '{label}' with body_pattern '{body_pattern}' matched no source-builder bodies " + f"across {len(source_builders)} source builder(s). " + f"Check that the pattern matches a body label in a source builder." ) cls._cl_pending_sites.clear() - return global_sites, proto_sites, world_sites + return global_site_indices, source_site_indices, env_root_sites @classmethod def _cl_inject_sites_fallback(cls) -> None: @@ -1219,70 +1232,47 @@ def instantiate_builder_from_stage(cls): # No env Xforms — flat loading builder.add_usd(stage, schema_resolvers=schema_resolvers) NewtonManager._world_xforms = [wp.transform()] - if cls._mpm_object_registry: - from isaaclab_newton.assets.mpm_object.mpm_object import add_registered_mpm_objects_to_builder - - add_registered_mpm_objects_to_builder(builder, 0, [0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0]) + for hook in cls._per_world_builder_hooks: + hook(builder, 0, [0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0]) else: # Load everything except the env subtrees (ground plane, lights, etc.) ignore_paths = [path for _, path in env_paths] builder.add_usd(stage, ignore_paths=ignore_paths, schema_resolvers=schema_resolvers) - # Build a prototype from the first env (all envs assumed identical) _, proto_path = env_paths[0] - proto = cls.create_builder(up_axis=up_axis) - proto.add_usd( - stage, - root_path=proto_path, - schema_resolvers=schema_resolvers, - ) + source_builders = {proto_path: cls.create_builder(up_axis=up_axis)} + source_builders[proto_path].add_usd(stage, root_path=proto_path, schema_resolvers=schema_resolvers) - # Inject registered sites into the proto before replication - global_sites, proto_sites, world_sites = cls._cl_inject_sites(builder, {proto_path: proto}) - global_site_map: dict[str, tuple[int, None]] = {label: (idx, None) for label, idx in global_sites.items()} - num_worlds = len(env_paths) - local_site_map: dict[str, list[list[int]]] = {} - site_entries = proto_sites.get(id(proto), {}) - world_xforms: list[wp.transform] = [] - - # Add each env as a separate Newton world + global_site_indices, source_site_indices, env_root_sites = cls._cl_inject_sites(builder, source_builders) xform_cache = UsdGeom.XformCache() - for col, (_, env_path) in enumerate(env_paths): - builder.begin_world() - offset = builder.shape_count + poses = [] + for _, env_path in env_paths: world_xform = xform_cache.GetLocalToWorldTransform(stage.GetPrimAtPath(env_path)) translation = world_xform.ExtractTranslation() rotation = world_xform.ExtractRotationQuat() - pos = (translation[0], translation[1], translation[2]) - quat = ( - rotation.GetImaginary()[0], - rotation.GetImaginary()[1], - rotation.GetImaginary()[2], - rotation.GetReal(), + imag = rotation.GetImaginary() + poses.append( + ( + (translation[0], translation[1], translation[2]), + (imag[0], imag[1], imag[2], rotation.GetReal()), + ) ) - env_xform = wp.transform(pos, quat) - world_xforms.append(env_xform) - builder.add_builder(proto, xform=env_xform) - for label, xform in world_sites.items(): - if label not in local_site_map: - local_site_map[label] = [[] for _ in range(num_worlds)] - site_idx = builder.add_site(body=-1, xform=wp.transform_multiply(env_xform, xform), label=label) - local_site_map[label][col].append(site_idx) - for label, proto_shape_indices in site_entries.items(): - if label not in local_site_map: - local_site_map[label] = [[] for _ in range(num_worlds)] - for proto_shape_idx in proto_shape_indices: - local_site_map[label][col].append(offset + proto_shape_idx) - if cls._mpm_object_registry: - from isaaclab_newton.assets.mpm_object.mpm_object import add_registered_mpm_objects_to_builder - - add_registered_mpm_objects_to_builder(builder, col, list(pos), list(quat)) - builder.end_world() - - NewtonManager._cl_site_index_map = { - **global_site_map, - **{label: (None, per_world) for label, per_world in local_site_map.items()}, - } + + positions = torch.tensor([pos for pos, _ in poses], dtype=torch.float32) + 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 = replicate_builder_mapping( + *replicate_args, + source_site_indices=source_site_indices, + env_root_sites=env_root_sites, + per_world_builder_hooks=cls._per_world_builder_hooks, + ) + + NewtonManager._cl_site_index_map = {label: (idx, None) for label, idx in global_site_indices.items()} + NewtonManager._cl_site_index_map.update( + (label, (None, per_world)) for label, per_world in local_site_map.items() + ) NewtonManager._world_xforms = world_xforms NewtonManager._num_envs = len(env_paths) @@ -1773,8 +1763,7 @@ def _ensure_visualization_model(cls) -> None: built. This is the entry point that makes :meth:`get_model` / :meth:`get_state` work uniformly across both sim backends. - The shadow model is built by walking the USD stage via - :meth:`_build_visualization_model_from_stage` and finalizing the resulting + The shadow model is built by walking the USD stage and finalizing the resulting :class:`~newton.ModelBuilder`. Per-frame body transforms are pushed into ``_state_0.body_q`` by :meth:`update_visualization_state` using the new :class:`~isaaclab.scene_data.SceneDataProvider`. @@ -1794,16 +1783,28 @@ def _ensure_visualization_model(cls) -> None: ) return - try: - builder = cls._build_visualization_model_from_stage(stage) - except Exception: - logger.exception( - "[NewtonManager] Failed to build a Newton ModelBuilder from the USD stage " - "for visualization (sim backend is PhysX)." + up_axis_token = UsdGeom.GetStageUpAxis(stage) + up_axis = Axis.from_string(str(up_axis_token)) + + env_pattern = re.compile(r"^env_(\d+)$") + env_paths = sorted( + (int(match.group(1)), child.GetPath().pathString) + for child in stage.GetPrimAtPath("/World/envs").GetChildren() + if (match := env_pattern.match(child.GetName())) + ) + if not env_paths: + logger.error( + "[NewtonManager] No /World/envs/env_ prims found; cannot build a " + "Newton visualization model from the cloned Isaac Lab scene." ) return - if builder is None or builder.body_count == 0: + NewtonManager._num_envs = len(env_paths) + builder = build_visualization_builder_from_stage_envs( + stage, env_paths, PhysicsManager._sim.get_clone_plan(), up_axis=up_axis + ) + + if builder.body_count == 0: logger.error( "[NewtonManager] USD stage walk produced no Newton bodies; the shadow " "Newton model for visualization will be empty. Common causes: the cloned " @@ -1828,141 +1829,6 @@ def _ensure_visualization_model(cls) -> None: NewtonManager._model = None NewtonManager._state_0 = None - @classmethod - def _build_visualization_model_from_stage(cls, stage) -> ModelBuilder | None: - """Build a fresh Newton ``ModelBuilder`` from the USD stage for visualization. - - Walks IsaacLab's ``/World/envs/env_`` convention and adds each env as - its own Newton world. When the env subtree is identical across envs (the - common cloned-scene case) a single env_0 prototype is built once and - replicated via :meth:`ModelBuilder.add_builder`; otherwise each env is - ingested independently with :meth:`ModelBuilder.add_usd`. - - This routine is intentionally independent of - :meth:`instantiate_builder_from_stage` (which targets the live-sim path - and uses a different naming convention and writes into ``cls._builder`` - and ``cls._cl_site_index_map``). The visualization shadow path must not - pollute those live-sim slots. ``cls._num_envs`` is populated here too so - :meth:`get_num_envs` returns the env count when the sim backend is PhysX - (the live-sim path never runs in that configuration, so there is no slot - to collide with). - - Args: - stage: USD stage to inspect. - - Returns: - A populated :class:`~newton.ModelBuilder`, or ``None`` when no - ``/World/envs/env_`` prims exist on the stage. - """ - import re - - from pxr import UsdGeom - - up_axis_token = UsdGeom.GetStageUpAxis(stage) - up_axis = Axis.from_string(str(up_axis_token)) - schema_resolvers = [SchemaResolverNewton(), SchemaResolverPhysx()] - - env_pattern = re.compile(r"^env_(\d+)$") - env_paths: list[tuple[int, str]] = [] - envs_root = stage.GetPrimAtPath("/World/envs") - if envs_root and envs_root.IsValid(): - for child in envs_root.GetChildren(): - if match := env_pattern.match(child.GetName()): - env_paths.append((int(match.group(1)), child.GetPath().pathString)) - env_paths.sort(key=lambda x: x[0]) - - builder = ModelBuilder(up_axis=up_axis) - - if not env_paths: - # Fallback: ingest the whole stage as a single world. - builder.add_usd(stage, schema_resolvers=schema_resolvers) - NewtonManager._num_envs = 1 - return builder - - NewtonManager._num_envs = len(env_paths) - - # Ingest stage-level (non-env) geometry into the global world (``current_world == -1``) - # so visualization sees the ground plane, ceilings, fixed props, etc. The legacy - # cloner-based prebuild did this via ``add_usd(stage, ignore_paths=["/World/envs"], ...)`` - # before adding the per-env worlds; without this, renderers/visualizers driven off the - # shadow Newton model are missing every shape authored outside the env hierarchy. - builder.add_usd( - stage, - ignore_paths=[r"/World/envs($|/.*)"], - schema_resolvers=schema_resolvers, - ) - - # Build env_0 as a prototype, then replicate across envs. - proto_env_path = env_paths[0][1] - proto = ModelBuilder(up_axis=up_axis) - proto.add_usd( - stage, - root_path=proto_env_path, - schema_resolvers=schema_resolvers, - ) - - xform_cache = UsdGeom.XformCache() - - # ``add_builder`` copies the prototype's ``body_label`` (and sibling label arrays) - # verbatim into each replicated world, so all worlds end up with prim paths under - # the prototype env (e.g. ``/World/envs/env_0/...``). The visualization sync uses - # these labels to map PhysX transforms (which carry distinct per-env paths) into - # ``state.body_q``; without rewriting, ``paths.index()`` resolves every match to - # world 0 and worlds 1..N never receive fresh poses. Rewrite the newly-added - # labels after each ``add_builder`` so each world references its own env prim path. - label_attrs = ("body_label", "articulation_label", "joint_label", "shape_label") - label_starts = {attr: len(getattr(builder, attr)) for attr in label_attrs} - - # ``proto.add_usd`` ingests env_0's bodies at their absolute world positions - # (``UsdPhysics.LoadUsdPhysicsFromRange`` reports world-space transforms), so - # ``proto.body_q`` already encodes env_0's world transform. ``add_builder`` - # composes its ``xform`` onto every imported body, so passing each env's - # absolute world transform here would double the offset; the correct xform is - # the env's pose relative to the prototype (identity for env_0, env_X * env_0^-1 - # for the rest). Dynamic bodies are overwritten in ``update_visualization_state`` - # via the PhysX sync, but static bodies (e.g. the table) keep this initial pose - # and render at the wrong position when env_0 is not at the world origin. - proto_world_gf = xform_cache.GetLocalToWorldTransform(stage.GetPrimAtPath(proto_env_path)) - proto_translation = proto_world_gf.ExtractTranslation() - proto_rotation = proto_world_gf.ExtractRotationQuat() - proto_world_tf = wp.transform( - (proto_translation[0], proto_translation[1], proto_translation[2]), - ( - proto_rotation.GetImaginary()[0], - proto_rotation.GetImaginary()[1], - proto_rotation.GetImaginary()[2], - proto_rotation.GetReal(), - ), - ) - proto_world_tf_inv = wp.transform_inverse(proto_world_tf) - - for _, env_path in env_paths: - world_xform = xform_cache.GetLocalToWorldTransform(stage.GetPrimAtPath(env_path)) - translation = world_xform.ExtractTranslation() - rotation = world_xform.ExtractRotationQuat() - env_world_tf = wp.transform( - (translation[0], translation[1], translation[2]), - ( - rotation.GetImaginary()[0], - rotation.GetImaginary()[1], - rotation.GetImaginary()[2], - rotation.GetReal(), - ), - ) - relative_tf = wp.transform_multiply(env_world_tf, proto_world_tf_inv) - builder.begin_world() - builder.add_builder(proto, xform=relative_tf) - if env_path != proto_env_path: - for attr in label_attrs: - labels = getattr(builder, attr) - for i in range(label_starts[attr], len(labels)): - labels[i] = labels[i].replace(proto_env_path, env_path, 1) - for attr in label_attrs: - label_starts[attr] = len(getattr(builder, attr)) - builder.end_world() - - return builder - @classmethod def get_scene_data_provider(cls) -> SceneDataProvider: """Return the active scene data provider, or None if unavailable. diff --git a/source/isaaclab_newton/isaaclab_newton/physics/visualization_builder.py b/source/isaaclab_newton/isaaclab_newton/physics/visualization_builder.py new file mode 100644 index 000000000000..1dbad55df711 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/physics/visualization_builder.py @@ -0,0 +1,56 @@ +# 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 Sequence +from typing import Any + +import torch +from newton import ModelBuilder +from newton._src.usd.schemas import SchemaResolverNewton, SchemaResolverPhysx + +from pxr import Usd + +from isaaclab.sim.utils.transforms import resolve_prim_pose + +from isaaclab_newton.cloner.newton_clone_utils import ( + build_source_builders, + rename_builder_labels, + replicate_builder_mapping, +) + + +def build_visualization_builder_from_stage_envs( + stage: Usd.Stage, + env_paths: Sequence[tuple[int, str]], + clone_plan: Any, + *, + up_axis: str = "Z", +) -> ModelBuilder: + """Build the Newton shadow visualization builder from cloned USD environments.""" + env_path_by_id = dict(env_paths) + + sources = tuple(clone_plan.sources) + destinations = tuple(clone_plan.destinations) + env_ids = clone_plan.env_ids.detach().cpu() + mapping = clone_plan.clone_mask.detach().cpu() + + poses = [resolve_prim_pose(stage.GetPrimAtPath(env_path_by_id[int(env_id)])) for env_id in env_ids.tolist()] + positions = torch.tensor([pos for pos, _ in poses], dtype=torch.float32) + quaternions = torch.tensor([quat for _, quat in poses], dtype=torch.float32) + schema_resolvers = [SchemaResolverNewton(), SchemaResolverPhysx()] + builder = ModelBuilder(up_axis=up_axis) + builder.add_usd(stage, ignore_paths=["/World/envs", *sources], schema_resolvers=schema_resolvers) + source_builders = build_source_builders( + stage, + sources, + lambda: ModelBuilder(up_axis=up_axis), + schema_resolvers, + simplify_meshes=False, + ) + replicate_builder_mapping(builder, sources, mapping, positions, quaternions, source_builders) + rename_builder_labels(builder, sources, destinations, env_ids, mapping) + return builder 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 f9913d58ec95..4a45b9213272 100644 --- a/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py +++ b/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py @@ -3,303 +3,327 @@ # # 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``, ``equality_constraint``). - * 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. -""" +"""Unit tests for Newton clone label rewriting and visualization clone-plan sources.""" import unittest +from unittest import mock import newton import torch -from isaaclab_newton.cloner.replicate import _BUILTIN_LABEL_TYPES, _rename_builder_labels +from isaaclab_newton.cloner import newton_clone_utils as newton_clone_utils_module +from isaaclab_newton.cloner.newton_clone_utils import ( + _BUILTIN_LABEL_TYPES, + rename_builder_labels, + replicate_builder_mapping, +) +from isaaclab_newton.physics import visualization_builder as visualization_builder_module from newton.solvers import SolverMuJoCo +from pxr import Usd, UsdGeom + +from isaaclab.cloner import ClonePlan + +_VIS_LABEL_SUFFIXES = { + "body_label": "Body", + "joint_label": "Joint", + "shape_label": "Shape", + "articulation_label": "Articulation", + "constraint_mimic_label": "ConstraintMimic", + "equality_constraint_label": "EqualityConstraint", +} +_VIS_LABEL_ATTRS = tuple(_VIS_LABEL_SUFFIXES) + _TENDON_FREQ = "mujoco:tendon" _SRC = "/Sources/protoA" _DST = "/World/envs/env_{}" +class _FakeVisualizationModelBuilder: + def __init__(self, up_axis=None): + self.up_axis = up_axis + for attr in _VIS_LABEL_ATTRS: + setattr(self, attr, []) + setattr(self, attr.replace("_label", "_world"), []) + self.custom_attributes = {} + self.geometry_sources = [] + self.world_slices = [] + self._current_world = None + + @property + def shape_count(self): + return len(self.shape_label) + + def begin_world(self): + self._current_world = len(self.world_slices) + self.world_slices.append([]) + + def end_world(self): + self._current_world = None + + def add_usd(self, stage, root_path=None, ignore_paths=None, schema_resolvers=None, **kwargs): + del stage, ignore_paths, schema_resolvers, kwargs + if root_path is None: + return + label_start = len(self.body_label) + geometry_start = len(self.geometry_sources) + for attr, suffix in _VIS_LABEL_SUFFIXES.items(): + getattr(self, attr).append(f"{root_path}/{suffix}") + getattr(self, attr.replace("_label", "_world")).append(self._current_world or 0) + self.geometry_sources.append(root_path) + self._record_world_slice(label_start, len(self.body_label), geometry_start, len(self.geometry_sources)) + + def add_builder(self, builder, xform=None): + del xform + label_start = len(self.body_label) + geometry_start = len(self.geometry_sources) + for attr in _VIS_LABEL_ATTRS: + labels = getattr(builder, attr) + getattr(self, attr).extend(labels) + getattr(self, attr.replace("_label", "_world")).extend([self._current_world] * len(labels)) + self.geometry_sources.extend(builder.geometry_sources) + self._record_world_slice(label_start, len(self.body_label), geometry_start, len(self.geometry_sources)) + + def labels_for_world(self, world_id, attr): + labels = getattr(self, attr) + return [label for start, end, _, _ in self.world_slices[world_id] for label in labels[start:end]] + + def geometry_sources_for_world(self, world_id): + return [ + source for _, _, start, end in self.world_slices[world_id] for source in self.geometry_sources[start:end] + ] + + def _record_world_slice(self, label_start, label_end, geometry_start, geometry_end): + if self._current_world is not None: + self.world_slices[self._current_world].append((label_start, label_end, geometry_start, geometry_end)) + + 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, _BUILTIN_LABEL_TYPES, _SRC, worlds) - _inject_tendon_strings(b, _SRC, worlds) - return b + for kind in types: + for world in worlds: + getattr(builder, f"{kind}_label").append(f"{src_path}/{kind}_{world}") + getattr(builder, f"{kind}_world").append(world) + + +def _inject_tendons(builder: newton.ModelBuilder, src_path: str, worlds: list[int]) -> None: + labels = builder.custom_attributes["mujoco:tendon_label"].values = [] + world_ids = builder.custom_attributes["mujoco:tendon_world"].values = [] + for world in worlds: + labels.append(f"{src_path}/Tendon_{world}") + world_ids.append(world) + builder._custom_frequency_counts[_TENDON_FREQ] = len(worlds) + + +def _make_builder(worlds: list[int]) -> newton.ModelBuilder: + builder = newton.ModelBuilder() + SolverMuJoCo.register_custom_attributes(builder) + _inject_builtins(builder, _BUILTIN_LABEL_TYPES, _SRC, worlds) + _inject_tendons(builder, _SRC, worlds) + return builder + + +def _add_custom_frequency(builder, freq_name, 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=f"{freq_name}_world", frequency=freq, dtype=int, default=0, namespace="syn", references="world" + ) + ) + for column in string_columns: + builder.add_custom_attribute( + newton.ModelBuilder.CustomAttribute(name=column, frequency=freq, dtype=str, default="", namespace="syn") + ) -class TestRenameBuilderLabels(unittest.TestCase): - """Both passes rewrite to the same per-env destination pattern.""" +def _populate_custom_frequency(builder, freq_name, string_columns, worlds): + builder.custom_attributes[f"syn:{freq_name}_world"].values = list(worlds) + for column in string_columns: + builder.custom_attributes[f"syn:{column}"].values = [f"{_SRC}/{column}_{world}" for world in worlds] + builder._custom_frequency_counts[f"syn:{freq_name}"] = len(worlds) + +class TestRenameBuilderLabels(unittest.TestCase): 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) - - def test_builtin_labels_rewritten_per_world(self): - b = _make_builder_with_entries(self.worlds) - self._rename(b) - 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)}", - msg=f"{t}_label[{k}] not rewritten correctly", - ) + rename_builder_labels(builder, [_SRC], [_DST], self.env_ids, self.mapping) + + def _assert_builtins(self, builder, types=_BUILTIN_LABEL_TYPES): + for kind in types: + worlds = getattr(builder, f"{kind}_world") + self.assertEqual( + getattr(builder, f"{kind}_label"), [f"{_DST.format(int(w))}/{kind}_{int(w)}" for w in worlds] + ) + + def test_builtin_and_tendon_labels_rewritten_per_world(self): + builder = _make_builder(self.worlds) + self._rename(builder) + self._assert_builtins(builder) + tendon_worlds = builder.custom_attributes["mujoco:tendon_world"].values + self.assertEqual( + builder.custom_attributes["mujoco:tendon_label"].values, + [f"{_DST.format(int(w))}/Tendon_{int(w)}" for w in tendon_worlds], + ) - 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)}") - - 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 _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}") - - 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.""" - 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_source_root_boundary_cases(self): + builder = _make_builder(self.worlds) + builder.body_label.append(_SRC) + builder.body_world.append(self.worlds[0]) + self._rename(builder) + self.assertEqual(builder.body_label[-1], _DST.format(self.worlds[0])) + + builder = _make_builder(self.worlds) + rename_builder_labels(builder, [f"{_SRC}/"], [_DST], self.env_ids, self.mapping) + self._assert_builtins(builder) + + def test_unmatched_rows_left_untouched(self): + builder = _make_builder(self.worlds) + builder.body_label.append(f"{_SRC}/body_99") + builder.body_world.append(99) + builder.custom_attributes["mujoco:tendon_label"].values.append("named_tendon") + builder.custom_attributes["mujoco:tendon_world"].values.append(self.worlds[0]) + self._rename(builder) + self.assertEqual(builder.body_label[-1], f"{_SRC}/body_99") + self.assertEqual(builder.custom_attributes["mujoco:tendon_label"].values[-1], "named_tendon") 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)}") - - 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.""" + for worlds in ([10, 20, 30], [0, 1_000_000, 2_147_000_000]): + builder = newton.ModelBuilder() + SolverMuJoCo.register_custom_attributes(builder) + _inject_builtins(builder, ("body",), _SRC, worlds) + env_ids = torch.tensor(worlds, dtype=torch.int32) + rename_builder_labels(builder, [_SRC], [_DST], env_ids, torch.ones(1, len(worlds), dtype=torch.bool)) + self._assert_builtins(builder, ("body",)) + +class TestRenameCustomAttributes(unittest.TestCase): 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 test_custom_string_columns_follow_frequency_worlds(self): + builder = newton.ModelBuilder() + _add_custom_frequency(builder, "freqA", ["freqA_label", "freqA_alt"]) + _add_custom_frequency(builder, "freqB", ["freqB_label"]) + _populate_custom_frequency(builder, "freqA", ["freqA_label", "freqA_alt"], self.worlds) + _populate_custom_frequency(builder, "freqB", ["freqB_label"], self.worlds) + rename_builder_labels(builder, [_SRC], [_DST], self.env_ids, self.mapping) + + for freq, columns in {"freqA": ("freqA_label", "freqA_alt"), "freqB": ("freqB_label",)}.items(): + worlds = builder.custom_attributes[f"syn:{freq}_world"].values + for column in columns: + self.assertEqual( + builder.custom_attributes[f"syn:{column}"].values, + [f"{_DST.format(int(w))}/{column}_{int(w)}" for w in worlds], ) - ) - 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)) + def test_empty_custom_string_column_passes_through(self): + builder = newton.ModelBuilder() + _add_custom_frequency(builder, "freqA", ["freqA_label"]) + rename_builder_labels(builder, [_SRC], [_DST], self.env_ids, self.mapping) + _populate_custom_frequency(builder, "freqA", ["freqA_label"], self.worlds) + self.assertEqual(len(builder.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 - ] + builder = newton.ModelBuilder() + SolverMuJoCo.register_custom_attributes(builder) + builder.body_label.extend([f"{sources[0]}/body", f"{sources[1]}/body"] * 2) + builder.body_world.extend([0, 0, 1, 1]) + rename_builder_labels( + builder, + sources, + ["/World/envs/env_{}", "/World/envs/env_{}"], + torch.tensor([0, 1], dtype=torch.int32), + torch.tensor([[1, 1], [1, 1]], dtype=torch.bool), ) - 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") + self.assertEqual( + builder.body_label, + ["/World/envs/env_0/body", "/World/envs/env_0/body", "/World/envs/env_1/body", "/World/envs/env_1/body"], + ) + + +class TestReplicateBuilderMapping(unittest.TestCase): + @staticmethod + def _source_builder(root_path: str): + builder = _FakeVisualizationModelBuilder() + builder.add_usd(None, root_path=root_path) + return builder + + def test_inactive_source_rows_are_ignored(self): + sources = ("/Sources/inactive", "/Sources/active") + source_builders = {source: self._source_builder(source) for source in sources} + builder = _FakeVisualizationModelBuilder() + + replicate_builder_mapping( + builder, + sources, + torch.tensor([[False, False], [True, False]], dtype=torch.bool), + torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]), + torch.tensor([[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0, 1.0]]), + source_builders, + ) + + self.assertEqual(builder.geometry_sources_for_world(0), ["/Sources/active"]) + self.assertEqual(builder.geometry_sources_for_world(1), []) + + +class TestVisualizationClonePlan(unittest.TestCase): + @staticmethod + def _define_xform(stage, path, translation=None): + xform = UsdGeom.Xform.Define(stage, path) + if translation is not None: + xform.AddTranslateOp().Set(translation) + + def test_visualization_builder_uses_clone_plan_sources_and_rewrites_labels(self): + stage = Usd.Stage.CreateInMemory() + UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.z) + self._define_xform(stage, "/World") + self._define_xform(stage, "/World/envs") + env_paths = [(env_id, f"/World/envs/env_{env_id}") for env_id in (0, 1, 2)] + for env_id, env_path in env_paths: + self._define_xform(stage, env_path, (float(env_id) * 3.0, 0.0, 0.0)) + self._define_xform(stage, f"{env_path}/Object") + self._define_xform(stage, "/World/envs/env_0/Object/source_0_visual") + self._define_xform(stage, "/World/envs/env_1/Object/source_1_visual") + + clone_plan = ClonePlan( + sources=("/World/envs/env_0/Object", "/World/envs/env_1/Object"), + destinations=("/World/envs/env_{}/Object", "/World/envs/env_{}/Object"), + clone_mask=torch.tensor([[True, False, True], [False, True, False]], dtype=torch.bool), + env_ids=torch.tensor([0, 1, 2], dtype=torch.long), + ) + + with ( + mock.patch.object(visualization_builder_module, "ModelBuilder", _FakeVisualizationModelBuilder), + mock.patch.object(newton_clone_utils_module, "ModelBuilder", _FakeVisualizationModelBuilder), + mock.patch.object(visualization_builder_module, "SchemaResolverNewton", lambda: object()), + mock.patch.object(visualization_builder_module, "SchemaResolverPhysx", lambda: object()), + mock.patch.object(newton_clone_utils_module.solvers.SolverMuJoCo, "register_custom_attributes"), + ): + builder = visualization_builder_module.build_visualization_builder_from_stage_envs( + stage, env_paths, clone_plan + ) + + self.assertEqual( + [builder.geometry_sources_for_world(i) for i in range(3)], + [["/World/envs/env_0/Object"], ["/World/envs/env_1/Object"], ["/World/envs/env_0/Object"]], + ) + for attr, suffix in _VIS_LABEL_SUFFIXES.items(): + self.assertEqual( + [builder.labels_for_world(i, attr) for i in range(3)], + [ + [f"/World/envs/env_0/Object/{suffix}"], + [f"/World/envs/env_1/Object/{suffix}"], + [f"/World/envs/env_2/Object/{suffix}"], + ], + ) if __name__ == "__main__":