Skip PhysX replicator for fully-heterogeneous 1:1 layouts (fixes mGPU double-free)

docs/source/how-to/cloning.rst

@@ -283,27 +283,35 @@ Every backend supplies its own :class:`~isaaclab.cloner.UsdReplicateContext` /
 ``PhysxReplicateContext`` / ``NewtonReplicateContext``, a class that hides the
 timing and granularity differences above behind a single uniform interface. A
 shared :data:`~isaaclab.cloner.REPLICATION_QUEUE` then remembers which asset
-belongs to which backend's context until it is time to run. The three
+cfgs participate until it is time to run. The three
 subsections below explain the queue, the contexts, and the function that joins
 them against a plan.
 
 The registration queue
 ~~~~~~~~~~~~~~~~~~~~~~
 
-Asset constructors do not replicate inline. They register their intent with
-:data:`~isaaclab.cloner.REPLICATION_QUEUE` and the framework defers the actual
-work to the drain. The queue ends up holding one entry per ``(asset, backend)``
-pair:
+Assets do not replicate inline. Construction only registers the asset's cfg
+through :func:`~isaaclab.cloner.queue_replication`; the queue records *which*
+cfgs participate:
 
 .. code-block:: text
 
     REPLICATION_QUEUE
-        (cartpole_cfg, PhysxReplicateContext)
-        (cartpole_cfg, UsdReplicateContext)
-        (cube_cfg,     UsdReplicateContext)
-        (light_cfg,    UsdReplicateContext)
+        cartpole_cfg
+        cube_cfg
+        camera_cfg
         ...
 
+*How* each cfg is cloned is resolved by :func:`~isaaclab.cloner.replicate` at
+dispatch: the cfg's :attr:`~isaaclab.assets.AssetBaseCfg.cloning_contexts`
+when set, otherwise the active backend's default stack. Each backend cloner
+exports that stack under the conventional name ``REPLICATION`` — PhysX pairs
+its native replication with USD clones, Newton includes USD only under Kit,
+and OvPhysX replicates alone because its clone replay authors USD itself.
+With :attr:`~isaaclab.cloner.CloneCfg.replicate_physics` disabled, cloning is
+USD-only: every context except ``UsdReplicateContext`` is dropped and the
+physics engine parses the per-env USD prims directly.
+
 Deferring the work like this buys three things at once:
 
 * Replication can wait until the plan is fully built, so the final layout is
@@ -313,6 +321,11 @@ Deferring the work like this buys three things at once:
 * Asset code stays free of any branching on which backend is active — it just
   registers and lets the framework take it from there.
 
+:attr:`~isaaclab.scene.InteractiveSceneCfg.replicate_physics` is piped into
+:attr:`~isaaclab.cloner.CloneCfg.replicate_physics` and applied at dispatch;
+an asset whose only cloning mechanism is physics replication is then simply
+not cloned.
+
 Backend contexts
 ~~~~~~~~~~~~~~~~
 
@@ -326,12 +339,12 @@ runtime:
     PhysxReplicateContext    # replicates PhysX rigid bodies and articulations
     NewtonReplicateContext   # replicates Newton bodies in its parallel pipeline
 
-A single asset can register more than one context — a PhysX articulation
-registers a PhysX context and a USD context so physics and visuals both follow,
-a Newton articulation registers a Newton context plus a USD context only if it
-owns visual prims. This is where backend differences are absorbed: swapping a
-scene from PhysX to Newton swaps which context an asset registers with, while
-the cfgs and the rest of the user code stay unchanged.
+A single asset usually resolves to more than one context — PhysX pairs its
+context with USD so physics and visuals both follow; Newton's default stack
+includes USD only under Kit, so kitless runs skip the authoring cost. This is
+where backend differences are absorbed: swapping a scene from PhysX to Newton
+swaps which default stack resolves at dispatch, while the cfgs and the rest of
+the user code stay unchanged.
 
 Running replication
 ~~~~~~~~~~~~~~~~~~~

source/isaaclab/changelog.d/zhengyuz-fix-replicate-physics-cfg-registration.major.rst

@@ -0,0 +1,35 @@
+Fixed
+^^^^^
+
+* Fixed :attr:`~isaaclab.scene.InteractiveSceneCfg.replicate_physics` being ignored since the
+  replication-session refactor: scenes configured with ``replicate_physics=False`` invoked
+  native physics replication anyway, silently discarding per-environment USD differences
+  (e.g. prestartup scale randomization). Cloning is now USD-only in that case, so the physics
+  engine parses each environment's USD prims directly; assets whose only cloning mechanism is
+  physics replication are not cloned.
+
+Added
+^^^^^
+
+* Added :attr:`~isaaclab.cloner.CloneCfg.replicate_physics` as the cloner-side home of the
+  policy; :class:`~isaaclab.scene.InteractiveScene` pipes the scene flag into it and
+  :func:`~isaaclab.cloner.replicate` applies it.
+* Added :attr:`~isaaclab.assets.AssetBaseCfg.cloning_contexts` to override, per asset, the
+  cloning contexts resolved at replication time; ``None`` uses the active backend's default
+  stack (``isaaclab_<backend>.cloner.REPLICATION``).
+
+Changed
+^^^^^^^
+
+* **Breaking:** Changed :data:`~isaaclab.cloner.REPLICATION_QUEUE` to hold asset cfgs instead
+  of ``(cfg, BackendCtxCls)`` pairs: construction registers *which* cfgs participate via
+  :func:`~isaaclab.cloner.queue_replication`, and *how* each is cloned resolves inside
+  :func:`~isaaclab.cloner.replicate`. Code appending pairs should register the cfg and direct
+  contexts through :attr:`~isaaclab.assets.AssetBaseCfg.cloning_contexts`.
+
+Removed
+^^^^^^^
+
+* Removed ``queue_usd_replication``: register the cfg with
+  :func:`~isaaclab.cloner.queue_replication` and direct contexts through the cfg's
+  ``cloning_contexts`` field instead.

source/isaaclab/isaaclab/assets/asset_base.py

@@ -15,6 +15,7 @@
 import warp as wp
 
 import isaaclab.sim as sim_utils
+from isaaclab.cloner import queue_replication
 from isaaclab.physics import PhysicsEvent, PhysicsManager
 from isaaclab.sim.simulation_context import SimulationContext
 from isaaclab.sim.utils.stage import get_current_stage
@@ -68,6 +69,9 @@ def __init__(self, cfg: AssetBaseCfg):
         """
         # check that the config is valid
         cfg.validate()
+        # register the original cfg object for cloning: the clone plan keys rows by the
+        # cfg identity the scene collected; contexts and policy resolve at replication time
+        queue_replication(cfg)
         # store inputs
         self.cfg = cfg.copy()
         # Resolve shape-check flag once: True means checks are active.

source/isaaclab/isaaclab/assets/asset_base_cfg.py

@@ -46,6 +46,15 @@ class InitialStateCfg:
     The class should inherit from :class:`isaaclab.assets.asset_base.AssetBase`.
     """
 
+    cloning_contexts: tuple[str | type, ...] | None = None
+    """Cloning contexts for this asset. Defaults to None.
+
+    Entries are ``"module:ContextClass"`` references (or classes) that clone this asset
+    across environments. If None, :func:`~isaaclab.cloner.replicate` uses the active
+    physics backend's default stack (``isaaclab_<backend>.cloner.REPLICATION``); an empty
+    tuple means the asset is not cloned.
+    """
+
     prim_path: str = MISSING
     """Prim path (or expression) to the asset.
 

source/isaaclab/isaaclab/cloner/__init__.pyi

@@ -16,9 +16,9 @@ __all__ = [
     "ReplicateSession",
     "REPLICATION_QUEUE",
     "replicate",
+    "queue_replication",
     "resolve_clone_plan_source",
     "split_clone_template",
-    "queue_usd_replication",
     "sequential",
     "UsdReplicateContext",
     "usd_replicate",
@@ -40,10 +40,10 @@ from .cloner_utils import (
 from .replicate_session import (
     REPLICATION_QUEUE,
     ReplicateSession,
+    queue_replication,
     replicate,
 )
 from .usd import (
     UsdReplicateContext,
-    queue_usd_replication,
     usd_replicate,
 )

source/isaaclab/isaaclab/cloner/clone_plan.py

@@ -50,8 +50,8 @@ def from_env_0(
         Auto-populates :attr:`cfg_rows` from
         :data:`~isaaclab.cloner.REPLICATION_QUEUE`, including only cfgs whose
         ``prim_path`` falls under the env-root prefix of ``destination``. Must be
-        called *after* all asset constructors have run, so their replication entries
-        are already in the queue; otherwise those assets will be skipped by the
+        called *after* all asset constructors have run, so their cfgs are already
+        registered in the queue; otherwise those assets will be skipped by the
         subsequent :func:`~isaaclab.cloner.replicate` call.
 
         Args:
@@ -66,7 +66,7 @@ def from_env_0(
 
         prefix, _ = split_clone_template(destination)
         cfg_rows: dict[int, tuple[int, ...]] = {
-            id(cfg): (0,) for cfg, _ in REPLICATION_QUEUE if cfg.prim_path.startswith(prefix)
+            id(cfg): (0,) for cfg in REPLICATION_QUEUE if cfg.prim_path.startswith(prefix)
         }
         return cls(
             sources=(source,),

source/isaaclab/isaaclab/cloner/cloner_cfg.py

@@ -28,3 +28,10 @@ class CloneCfg:
 
     clone_regex: str = "/World/envs/env_.*"
     """Regex matching every replicated env prim. Used to expand ``{ENV_REGEX_NS}`` cfg macros."""
+
+    replicate_physics: bool = True
+    """Whether physics replication clones each environment. Default is True.
+
+    If False, cloning is USD-only: the physics engine parses the per-env USD prims directly
+    instead of replicating env_0's parsed structure. Applied by :func:`~isaaclab.cloner.replicate`.
+    """

source/isaaclab/isaaclab/cloner/replicate_session.py

@@ -7,10 +7,15 @@
 
 from __future__ import annotations
 
+import importlib
 from collections.abc import Callable, Iterable
 from typing import TYPE_CHECKING, Any
 
+from isaaclab.utils.backend_utils import FactoryBase
+from isaaclab.utils.string import string_to_callable
+
 from .cloner_strategies import sequential
+from .usd import UsdReplicateContext
 
 if TYPE_CHECKING:
     import torch
@@ -20,32 +25,67 @@
     from .clone_plan import ClonePlan
 
 
-REPLICATION_QUEUE: list[tuple[Any, type]] = []
-"""``(cfg, BackendCtxCls)`` pairs appended by ``queue_<backend>_replication`` and drained by :func:`replicate`."""
+REPLICATION_QUEUE: list[Any] = []
+"""Asset cfgs registered by :func:`queue_replication` and drained by :func:`replicate`.
+
+The queue only records *which* cfgs participate in cloning; how each cfg is cloned is
+resolved at dispatch from :attr:`~isaaclab.assets.AssetBaseCfg.cloning_contexts` or the
+active backend's default stack.
+"""
+
+
+def queue_replication(cfg: Any) -> None:
+    """Register ``cfg`` for cloning when :func:`replicate` next runs.
+
+    Args:
+        cfg: Asset cfg with resolved ``prim_path``.
+    """
+    REPLICATION_QUEUE.append(cfg)
 
 
-def replicate(plan: ClonePlan, *, stage: Usd.Stage) -> None:
+def replicate(plan: ClonePlan, *, stage: Usd.Stage, replicate_physics: bool = True) -> None:
     """Drain :data:`REPLICATION_QUEUE` against ``plan``, dispatch each backend, publish the plan.
 
+    Each cfg's cloning contexts come from
+    :attr:`~isaaclab.assets.AssetBaseCfg.cloning_contexts` when set, otherwise from the
+    active backend's default stack (``isaaclab_<backend>.cloner.REPLICATION``); entries may
+    be ``"module:ContextClass"`` references or classes. With ``replicate_physics=False``
+    cloning is USD-only: contexts other than
+    :class:`~isaaclab.cloner.UsdReplicateContext` are dropped, and the physics engine
+    parses the per-env USD prims directly.
+
     Cfgs absent from ``plan.cfg_rows`` are silently skipped. Backend contexts run in
     ascending ``replicate_priority`` order. The queue is cleared up front, so a backend
     failure cannot leak stale entries into the next call.
+
+    Args:
+        plan: Replication layout to dispatch.
+        stage: USD stage to author replicated prim specs into.
+        replicate_physics: Whether physics replication clones each environment. If False,
+            cloning is USD-only; an asset whose contexts are all physics-based is not cloned.
     """
     from isaaclab.sim import SimulationContext  # noqa: PLC0415
 
     queued = list(REPLICATION_QUEUE)
     REPLICATION_QUEUE.clear()
 
     # Group queued cfgs by backend, taking the union of row indices each backend owns.
-    # In the homogeneous plan every cfg maps to row 0, so multiple queue_<backend>_replication
+    # In the homogeneous plan every cfg maps to row 0, so multiple queue_replication
     # calls (e.g. one per body type in RigidObjectCollection) all contribute {0} and the set
     # union keeps it as a single row — no redundant copy specs are authored.
     backend_rows: dict[type, set[int]] = {}
-    for cfg, BackendCtxCls in queued:
+    for cfg in queued:
         rows = plan.cfg_rows.get(id(cfg))
         if rows is None:
             continue
-        backend_rows.setdefault(BackendCtxCls, set()).update(rows)
+        contexts = cfg.cloning_contexts
+        if contexts is None:
+            contexts = getattr(importlib.import_module(f"isaaclab_{FactoryBase._get_backend()}.cloner"), "REPLICATION")
+        contexts = [string_to_callable(c) if isinstance(c, str) else c for c in contexts]
+        if not replicate_physics:
+            contexts = [c for c in contexts if c is UsdReplicateContext]
+        for BackendCtxCls in contexts:
+            backend_rows.setdefault(BackendCtxCls, set()).update(rows)
 
     backend_ctxs: dict[type, Any] = {}
     for BackendCtxCls, row_set in backend_rows.items():
@@ -70,7 +110,7 @@ class ReplicateSession:
     """Folds :func:`make_clone_plan` and :func:`replicate` into a ``with`` block.
 
     ``__enter__`` builds the plan (and mutates each cfg's ``spawn_path``); asset
-    constructors inside the block register backend replication into
+    constructors inside the block register their cfgs into
     :data:`REPLICATION_QUEUE`; ``__exit__`` drains and dispatches.
 
     Example:
@@ -92,6 +132,7 @@ def __init__(
         stage: Usd.Stage,
         clone_strategy: Callable = sequential,
         valid_set: torch.Tensor | None = None,
+        replicate_physics: bool = True,
     ):
         """Capture arguments for :func:`make_clone_plan` and :func:`replicate`.
 
@@ -104,9 +145,12 @@ def __init__(
             clone_strategy: Prototype-to-env assignment function.
             valid_set: Optional ``[num_combos, num_groups]`` long tensor of valid
                 prototype combinations; ``None`` uses the full cartesian product.
+            replicate_physics: Whether physics replication clones each environment;
+                forwarded to :func:`replicate`.
         """
         self._cfgs = cfgs
         self._stage = stage
+        self._replicate_physics = replicate_physics
         self._kwargs = dict(
             num_clones=num_clones,
             env_spacing=env_spacing,
@@ -125,7 +169,7 @@ def __enter__(self) -> ReplicateSession:
     def __exit__(self, exc_type, exc_value, traceback) -> None:
         if exc_type is None:
             assert self._plan is not None
-            replicate(self._plan, stage=self._stage)
+            replicate(self._plan, stage=self._stage, replicate_physics=self._replicate_physics)
         else:
             # Drop cfgs registered before the failure so the next session is clean.
             REPLICATION_QUEUE.clear()

source/isaaclab/isaaclab/cloner/usd.py

@@ -6,15 +6,13 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import Any
 
 import torch
 
 from pxr import Gf, Sdf, Usd, UsdGeom, Vt
 
 from ._fabric_notices import disabled_fabric_change_notifies
 from .cloner_utils import split_clone_template
-from .replicate_session import REPLICATION_QUEUE
 
 
 def _select_env_ids(env_ids: torch.Tensor, mask: torch.Tensor | None, row: int) -> torch.Tensor:
@@ -155,16 +153,6 @@ def dp_depth(template: str) -> int:
                                 op_order.default = Vt.TokenArray(op_names)
 
 
-def queue_usd_replication(cfg: Any) -> None:
-    """Register ``cfg`` for USD replication when :func:`~isaaclab.cloner.replicate` next runs.
-
-    Appends ``(cfg, UsdReplicateContext)`` to :data:`~isaaclab.cloner.REPLICATION_QUEUE`.
-    The actual row resolution and dispatch happen inside :func:`~isaaclab.cloner.replicate`,
-    so this helper is safe to call from any asset constructor — no active session is required.
-    """
-    REPLICATION_QUEUE.append((cfg, UsdReplicateContext))
-
-
 def usd_replicate(
     stage: Usd.Stage,
     sources: Sequence[str],

source/isaaclab/isaaclab/scene/interactive_scene.py

@@ -156,7 +156,7 @@ def __init__(self, cfg: InteractiveSceneCfg):
         # physics scene path
         self._physics_scene_path = None
         # prepare cloner for environment replication
-        self.cloner_cfg = cloner.CloneCfg(device=self.device)
+        self.cloner_cfg = cloner.CloneCfg(device=self.device, replicate_physics=self.cfg.replicate_physics)
         env_root = self.cloner_cfg.clone_regex.rsplit("/", 1)[0]
         self.env_prim_paths = [f"{env_root}/env_{i}" for i in range(self.cfg.num_envs)]
 
@@ -184,6 +184,7 @@ def __init__(self, cfg: InteractiveSceneCfg):
             device=self.device,
             stage=self.stage,
             clone_strategy=self.cloner_cfg.clone_strategy,
+            replicate_physics=self.cloner_cfg.replicate_physics,
         ):
             if self._is_scene_setup_from_cfg():
                 self._add_entities_from_cfg()

source/isaaclab/isaaclab/scene/interactive_scene_cfg.py

@@ -108,6 +108,17 @@ class MySceneCfg(InteractiveSceneCfg):
     .. note::
         Optimized parsing of certain prim types (such as deformable objects) is not currently supported
         by the physics engine. In these cases, this flag needs to be set to False.
+
+    .. attention::
+        Setting this flag to False is currently not supported on the Newton physics backend:
+        Newton discovers the scene through its replication path, which stage parsing cannot
+        replace for cloned environments.
+
+    .. note::
+        The scene pipes this flag into :attr:`~isaaclab.cloner.CloneCfg.replicate_physics`;
+        the policy is applied by :func:`~isaaclab.cloner.replicate`. Direct workflows that
+        call :func:`~isaaclab.cloner.replicate` themselves pass ``replicate_physics``
+        explicitly.
     """
 
     filter_collisions: bool = True