# 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
import contextlib
import copy
import re
from collections.abc import Callable, Iterator, Sequence
from typing import TYPE_CHECKING, TypeAlias
import torch
import warp as wp
from newton import ModelBuilder
from newton._src.usd.schemas import SchemaResolverNewton, SchemaResolverPhysx
from pxr import Usd
from isaaclab.physics import PhysicsManager
from isaaclab.sim.utils.newton_model_utils import replace_newton_builder_shape_colors
from isaaclab_newton.cloner.newton_clone_utils import (
_restore_visible_colliders_without_visual_shapes,
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 copy_newton_clone_source(source_path: str, xform: wp.transform | None = None) -> ModelBuilder:
"""Copy a retained clone-source builder without sharing mutable shape geometry.
Args:
source_path: Clone-plan source prim path retained during Newton replication.
xform: Optional transform applied while copying the source.
Returns:
An independent builder that is safe to finalize or extend.
Raises:
RuntimeError: If Newton replication did not retain the requested source.
"""
source = NewtonManager._cl_protos.get(source_path)
if source is None:
raise RuntimeError(f"No retained Newton clone source for {source_path!r}.")
builder = ModelBuilder(up_axis=source.up_axis)
if xform is None:
builder.add_builder(source)
else:
builder.add_builder(source, xform=xform)
builder.shape_source = [
value.copy() if callable(getattr(value, "copy", None)) else copy.copy(value) for value in builder.shape_source
]
return builder
[docs]
@contextlib.contextmanager
def newton_builder_world_hook(
hook: Callable[[ModelBuilder, int, list[float], list[float]], None],
) -> Iterator[None]:
"""Temporarily extend every world built by Newton replication.
The callback must not already be registered. On exit, the context removes
only its callback and preserves hooks owned by other callers.
Args:
hook: Callback receiving the builder, world index, world position [m],
and world orientation quaternion in xyzw order during replication.
Yields:
Control while the callback is registered.
Raises:
RuntimeError: If the callback is already registered.
"""
hooks = NewtonManager._per_world_builder_hooks
if hook in hooks:
raise RuntimeError("Newton world-builder hook is already registered.")
hooks.append(hook)
try:
yield
finally:
if hook in hooks:
hooks.remove(hook)
def _build_newton_builder_from_mapping(
stage: Usd.Stage,
sources: Sequence[str],
destinations: Sequence[str],
env_ids: torch.Tensor,
mapping: torch.Tensor,
positions: torch.Tensor | None = None,
quaternions: torch.Tensor | None = None,
up_axis: str = "Z",
load_visual_shapes: bool = True,
) -> tuple[ModelBuilder, object, dict, list, dict[str, ModelBuilder]]:
"""Build a Newton model builder from clone mapping inputs.
Also returns the per-source builders (``{source_path: ModelBuilder}``) so the
committing path can retain them for single-model consumers such as the
batched Newton IK action.
"""
if positions is None:
positions = torch.zeros((mapping.size(1), 3), device=mapping.device, dtype=torch.float32)
if quaternions is None:
quaternions = torch.zeros((mapping.size(1), 4), device=mapping.device, dtype=torch.float32)
quaternions[:, 3] = 1.0
schema_resolvers = [SchemaResolverNewton(), SchemaResolverPhysx()]
manager_cls = PhysicsManager._sim.physics_manager
builder = manager_cls.create_builder(up_axis=up_axis)
# Swap height-field-tagged terrain colliders for Newton heightfields before the
# mesh import, and skip those prims in add_usd so the terrain is not imported twice.
hf_ignore_paths = manager_cls._inject_terrain_heightfields(stage, builder)
stage_info = builder.add_usd(
stage,
ignore_paths=["/World/envs", *sources, *hf_ignore_paths],
schema_resolvers=schema_resolvers,
load_visual_shapes=load_visual_shapes,
)
_restore_visible_colliders_without_visual_shapes(builder, stage, stage_info["path_shape_map"], load_visual_shapes)
replace_newton_builder_shape_colors(builder, stage)
# 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.
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,
load_visual_shapes=load_visual_shapes,
)
# 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)
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,
)
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, source_builders
def _renderer_wants_visual_shapes() -> bool:
"""Whether anything in this run will draw the Newton model's visual-only shapes.
Visual shapes are consumed by the viewers, offscreen ``rgb_array`` capture, and camera
sensors on any renderer backend. A headless training run without cameras draws none of
them, so importing them only costs USD parse time and memory.
"""
from isaaclab.sim import SimulationContext
sim = SimulationContext.instance()
if sim is None:
return True
return bool(sim.is_rendering or sim.can_render_rgb_array() or sim.visual_shapes_required)
[docs]
class NewtonReplicateContext:
"""Queue and run Newton replication work for one stage."""
[docs]
def __init__(
self,
stage: Usd.Stage,
*,
device: str = "cpu",
up_axis: str = "Z",
load_visual_shapes: bool | None = None,
commit_to_manager: bool = True,
):
"""Initialize the context.
Args:
stage: USD stage containing source assets.
device: Device used by the finalized Newton model builder.
up_axis: Up axis for the Newton model builder.
load_visual_shapes: Whether to import visual-only geometry. If ``None``,
read from the active :class:`NewtonCfg`, which itself defaults to
importing them only when a renderer or visualizer is active.
commit_to_manager: Whether :meth:`replicate` should publish the builder to
:class:`NewtonManager`.
"""
self.stage = stage
self.device = device
self.up_axis = up_axis
if load_visual_shapes is None:
from isaaclab_newton.physics import NewtonCfg
cfg = PhysicsManager._cfg
load_visual_shapes = cfg.load_visual_shapes if isinstance(cfg, NewtonCfg) else None
self.load_visual_shapes = _renderer_wants_visual_shapes() if load_visual_shapes is None else load_visual_shapes
self.commit_to_manager = commit_to_manager
self._queue: list[_MappingBatch] = []
def queue_mapping(
self,
sources: Sequence[str],
destinations: Sequence[str],
env_ids: torch.Tensor,
mapping: torch.Tensor,
*,
positions: torch.Tensor | None = None,
quaternions: torch.Tensor | None = None,
) -> None:
"""Queue replication rows from the current flat clone mapping.
Args:
sources: Source prim paths used for cloning.
destinations: Destination prim path templates.
env_ids: Environment ids for destination worlds.
mapping: Boolean source-to-environment mapping matrix.
positions: Optional per-environment world positions [m].
quaternions: Optional per-environment orientations in xyzw order.
"""
self._queue.append((tuple(sources), tuple(destinations), env_ids, mapping, positions, quaternions))
@staticmethod
def _merge_optional_tensor(
name: str, current: torch.Tensor | None, incoming: torch.Tensor | None
) -> torch.Tensor | None:
"""Merge optional tensors, requiring equal values when both are present."""
if current is None:
return incoming
if incoming is None:
return current
if current.device != incoming.device or current.shape != incoming.shape or not torch.equal(current, incoming):
raise ValueError(f"Queued Newton mappings must use the same {name} tensor.")
return current
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.")
sources: list[str] = []
destinations: list[str] = []
mappings: list[torch.Tensor] = []
env_ids = self._queue[0][2]
positions = self._queue[0][4]
quaternions = self._queue[0][5]
for (
queued_sources,
queued_destinations,
queued_env_ids,
mapping,
queued_positions,
queued_quaternions,
) in self._queue:
if (
env_ids.device != queued_env_ids.device
or env_ids.shape != queued_env_ids.shape
or not torch.equal(env_ids, queued_env_ids)
):
raise ValueError("Queued Newton mappings must use the same env_ids tensor.")
sources.extend(queued_sources)
destinations.extend(queued_destinations)
mappings.append(mapping)
positions = self._merge_optional_tensor("positions", positions, queued_positions)
quaternions = self._merge_optional_tensor("quaternions", quaternions, queued_quaternions)
return tuple(sources), tuple(destinations), env_ids, torch.cat(mappings, dim=0), positions, quaternions
def replicate(self) -> tuple[ModelBuilder, object, dict]:
"""Build the Newton model builder from queued mappings and optionally publish it."""
sources, destinations, env_ids, mapping, positions, quaternions = self._merged_mapping()
builder, stage_info, site_index_map, world_xforms, source_builders = _build_newton_builder_from_mapping(
stage=self.stage,
sources=sources,
destinations=destinations,
env_ids=env_ids,
mapping=mapping,
positions=positions,
quaternions=quaternions,
up_axis=self.up_axis,
load_visual_shapes=self.load_visual_shapes,
)
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
NewtonManager._world_xforms = world_xforms
NewtonManager._cl_protos = source_builders
NewtonManager.set_builder(builder)
NewtonManager._num_envs = mapping.size(1)
self._queue.clear()
return builder, stage_info, site_index_map
PHYSICS_CONTEXT = NewtonReplicateContext
"""Physics-only replication context for Newton assets. USD replication is added automatically
by :func:`~isaaclab.cloner.replicate` when the asset has a spawner and Kit is available."""
def newton_physics_replicate(
stage: Usd.Stage,
sources: Sequence[str],
destinations: Sequence[str],
env_ids: torch.Tensor,
mapping: torch.Tensor,
positions: torch.Tensor | None = None,
quaternions: torch.Tensor | None = None,
device: str = "cpu",
up_axis: str = "Z",
):
"""Replicate prims into a Newton ``ModelBuilder`` using a per-source mapping.
Args:
stage: USD stage containing source assets.
sources: Source prim paths used for cloning.
destinations: Destination prim path templates.
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.
device: Device used by the finalized Newton model builder.
up_axis: Up axis for the Newton model builder.
Returns:
Tuple of the populated Newton model builder and stage metadata.
"""
ctx = NewtonReplicateContext(stage, device=device, up_axis=up_axis, commit_to_manager=True)
ctx.queue_mapping(sources, destinations, env_ids, mapping, positions=positions, quaternions=quaternions)
builder, stage_info, _site_index_map = ctx.replicate()
return builder, stage_info
