Honor replicate_physics via cfg-registered replication lists

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,16 @@ 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 perform physics
+    replication. If None, :func:`~isaaclab.cloner.replicate` uses the backend's default
+    physics context. :class:`~isaaclab.cloner.UsdReplicateContext` is always added
+    automatically when ``spawn`` is set and Kit is available. An empty tuple suppresses
+    all cloning.
+    """
+
     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/cloner_utils.py

@@ -106,17 +106,20 @@ def resolve_clone_plan_source(path_expr: str, plan: ClonePlan) -> tuple[str, str
         Three-tuple of ``(source_asset_path, dest_glob_prefix, asset_suffix)``. The
         ``asset_suffix`` is the part of ``path_expr`` beyond the matching row's
         destination template (empty when ``path_expr`` equals the row's template).
-        Returns ``None`` when ``path_expr`` matches no row in the plan, letting
-        callers fall back to direct stage resolution (e.g. for sensor frames
-        mounted at the env root rather than under a planned asset).
+        Returns ``None`` when ``path_expr`` matches no row in the plan, or when the
+        matching rows have no active env, letting callers fall back to direct stage
+        resolution (e.g. for sensor frames mounted at the env root rather than under
+        a planned asset).
+
+        Partial-env coverage is supported: when the matching rows cover only a subset
+        of envs (an asset present in some envs but not others, as in heterogeneous
+        scenes), the returned destination glob resolves to just those envs.
 
     Raises:
         ValueError: When ``path_expr`` is owned by multiple distinct, equally
             specific destination templates (a genuine ambiguity). Nested
             templates do not conflict: the most specific (longest-matching) one
             wins, mirroring :func:`iter_clone_plan_matches`.
-        NotImplementedError: When the union of matching rows' clone masks does not
-            cover every env (partial-env heterogeneous coverage is unsupported).
     """
     # Collect every template that owns ``path_expr`` together with the suffix below it.
     # A shorter suffix means a longer matched prefix, i.e. a more specific (nearer) owner.
@@ -140,12 +143,15 @@ def resolve_clone_plan_source(path_expr: str, plan: ClonePlan) -> tuple[str, str
     matching_template = next(iter(owning_templates))
     matching_rows = [index for template, _, index in candidates if template == matching_template]
     matching_suffix = next(suffix for template, suffix, _ in candidates if template == matching_template)
-    if not plan.clone_mask[matching_rows].any(dim=0).all():
-        raise NotImplementedError(
-            f"path_expr {path_expr!r}: partial-env heterogeneous coverage is unsupported;"
-            " matching rows must collectively cover all envs."
-        )
-    return plan.sources[matching_rows[0]], matching_template.replace("{}", "*"), matching_suffix or ""
+    # Partial-env coverage (the union of matching rows misses some envs) is expected for
+    # heterogeneous scenes: an asset present in only a subset of envs. The destination glob
+    # below resolves only to the envs that actually received the asset. Resolution must
+    # still walk a source that exists on stage, so prefer the first matching row with at
+    # least one active env over an inactive fallback source.
+    active_rows = [index for index in matching_rows if plan.clone_mask[index].any()]
+    if not active_rows:
+        return None
+    return plan.sources[active_rows[0]], matching_template.replace("{}", "*"), matching_suffix or ""
 
 
 def iter_clone_plan_matches(plan: ClonePlan, path_expr: str) -> Iterator[tuple[str, str, str, tuple[int, ...]]]:

source/isaaclab/isaaclab/cloner/replicate_session.py

@@ -7,10 +7,16 @@
 
 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 isaaclab.utils.version import has_kit
+
 from .cloner_strategies import sequential
+from .usd import UsdReplicateContext
 
 if TYPE_CHECKING:
     import torch
@@ -20,32 +26,71 @@
     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.
 
+    Physics contexts come from :attr:`~isaaclab.assets.AssetBaseCfg.cloning_contexts` when
+    set, otherwise from the backend's ``PHYSICS_CONTEXT`` class.
+    :class:`~isaaclab.cloner.UsdReplicateContext` is added automatically when the cfg has a
+    spawner and Kit is available. With ``replicate_physics=False`` physics contexts are
+    dropped; USD replication still fires when the spawner+Kit condition is met.
+
     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)
+        if cfg.cloning_contexts is None:
+            physics_ctx = getattr(
+                importlib.import_module(f"isaaclab_{FactoryBase._get_backend()}.cloner"), "PHYSICS_CONTEXT", None
+            )
+            contexts = [physics_ctx] if physics_ctx else []
+        else:
+            contexts = [string_to_callable(c) if isinstance(c, str) else c for c in cfg.cloning_contexts]
+        if not replicate_physics:
+            contexts = [c for c in contexts if c is UsdReplicateContext]
+        ctx_set = dict.fromkeys(contexts)
+        if getattr(cfg, "spawn", None) is not None and has_kit():
+            ctx_set.setdefault(UsdReplicateContext, None)
+        for BackendCtxCls in ctx_set:
+            backend_rows.setdefault(BackendCtxCls, set()).update(rows)
 
     backend_ctxs: dict[type, Any] = {}
     for BackendCtxCls, row_set in backend_rows.items():
@@ -70,7 +115,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 +137,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 +150,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 +174,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()