# 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 TYPE_CHECKING
import numpy as np
from omni.physx import get_physx_replicator_interface
from pxr import Sdf, Usd, UsdUtils
from isaaclab import cloner
if TYPE_CHECKING:
from isaaclab.cloner import ClonePlan
[docs]
class PhysxReplicateContext:
"""Apply one clone plan through the PhysX replicator."""
replicate_priority = 0
[docs]
def __init__(self, stage: Usd.Stage):
"""Initialize the context.
Args:
stage: USD stage to register with the PhysX replicator.
"""
self.stage = stage
cache = UsdUtils.StageCache.Get()
cached_id = cache.GetId(stage)
self._stage_id = cached_id.ToLongInt() if cached_id.IsValid() else cache.Insert(stage).ToLongInt()
def replicate(self, plan: ClonePlan) -> None:
"""Register the PhysX replicator for this context's plan rows.
Args:
plan: Replication layout shared by every clone backend.
"""
if plan.env_ids is None:
raise ValueError("ClonePlan.env_ids is required for replication.")
rows = plan.context_rows[type(self)]
native_rows = set(rows)
other_rows = {
row for context, routed in plan.context_rows.items() if context is not type(self) for row in routed
}
self._replicate_mapping(
sources=tuple(plan.sources[row] for row in rows),
destinations=tuple(plan.destinations[row] for row in rows),
env_ids=plan.env_ids,
mapping=plan.clone_mask[list(rows)],
has_usd_only_rows=bool(other_rows - native_rows),
exclude_self_replication=True,
)
def _replicate_mapping(
self,
sources: Sequence[str],
destinations: Sequence[str],
env_ids: np.ndarray,
mapping: np.ndarray,
has_usd_only_rows: bool,
exclude_self_replication: bool,
) -> None:
"""Register one raw source-to-environment mapping with PhysX."""
physx_queue: list[tuple[str, str, tuple[int, ...]]] = []
expected_shape = (len(sources), len(env_ids))
if mapping.shape != expected_shape:
raise ValueError(f"mapping must have shape {expected_shape}, got {mapping.shape}.")
if mapping.shape[1] <= 1:
return
native_paths: list[str] = []
for i, src in enumerate(sources):
worlds = tuple(map(int, env_ids[np.flatnonzero(mapping[i])]))
if has_usd_only_rows:
native_paths.append(src)
native_paths.extend(destinations[i].format(world) for world in worlds)
if exclude_self_replication:
matched = cloner.path.match(src, destinations[i])
if matched is not None and matched.instance.isdigit():
filtered = tuple(world for world in worlds if world != int(matched.instance))
worlds = filtered if filtered else worlds
physx_queue.append((src, destinations[i], worlds))
# Fully-heterogeneous 1:1 layouts have every source mapped only to its own
# environment (no cross-env replication needed). Calling rep.replicate() once
# per source with a single self-target is known to trigger intermittent native
# heap corruption (double-free / SIGABRT) under mGPU, likely due to per-call
# PhysX-internal allocations summing to a problematic total across processes.
# For these layouts the source prims are already in their correct env positions
# and PhysX can parse them from the stage without any replicator registration.
if all(len(envs) == 1 and src == destination.format(envs[0]) for src, destination, envs in physx_queue):
return
physics_scene_prim = self.stage.GetPrimAtPath("/physicsScene")
if physics_scene_prim.IsValid():
physics_scene_prim.CreateAttribute("physxScene:envIdInBoundsBitCount", Sdf.ValueTypeNames.Int).Set(4)
current_worlds: list[int] = []
current_template: str = ""
prefixes = [cloner.path.split(destination)[0] for destination in destinations]
env_namespaces = [
prefix.rstrip("/") if prefix.endswith("/") else prefix.rsplit("/", 1)[0] for prefix in prefixes
]
excluded_paths = (
list(dict.fromkeys(native_paths))
if has_usd_only_rows
else list(dict.fromkeys(("/World/template", *env_namespaces)))
)
def attach_fn(_stage_id: int):
return excluded_paths
def rename_fn(_replicate_path: str, i: int):
return current_template.format(current_worlds[i])
def attach_end_fn(_stage_id: int):
nonlocal current_template
replicator = get_physx_replicator_interface()
for src, destination, target_envs in physx_queue:
current_template = destination
current_worlds[:] = target_envs
if not current_worlds:
continue
replicator.replicate(
_stage_id,
src,
len(current_worlds),
useEnvIds=False,
useFabricForReplication=False,
)
replicator.unregister_replicator(_stage_id)
get_physx_replicator_interface().register_replicator(self._stage_id, attach_fn, attach_end_fn, rename_fn)
def physx_replicate(
stage: Usd.Stage,
sources: Sequence[str],
destinations: Sequence[str],
env_ids: np.ndarray,
mapping: np.ndarray,
positions: np.ndarray | None = None,
quaternions: np.ndarray | None = None,
exclude_self_replication: bool = True,
) -> None:
"""Replicate a raw source-to-environment mapping through PhysX.
Args:
stage: USD stage containing the source prims.
sources: Source prim paths, one per mapping row.
destinations: Destination templates containing ``"{}"``, one per mapping row.
env_ids: Integer environment identifiers, shape ``[num_envs]``.
mapping: Boolean source-to-environment selection, shape ``[len(sources), num_envs]``.
positions: Optional environment positions [m], shape ``[num_envs, 3]``. Unused by PhysX.
quaternions: Optional environment orientations in xyzw order, shape ``[num_envs, 4]``. Unused by PhysX.
exclude_self_replication: Whether to omit a source environment from its own targets.
"""
del positions, quaternions
context = PhysxReplicateContext(stage)
context._replicate_mapping(
sources=sources,
destinations=destinations,
env_ids=env_ids,
mapping=mapping,
has_usd_only_rows=False,
exclude_self_replication=exclude_self_replication,
)
