Publish clone plans before scene construction

docs/source/how-to/cloning.rst

@@ -216,7 +216,7 @@ them is purely about ergonomics:
   when you want the lifecycle hidden and you are authoring assets through a
   scene config.
 * The second spells the same flow out as plain function calls, leaving a moment
-  between the build and the drain where you can inspect or mutate the plan.
+  between the build and the drain where you can inspect the plan.
   Reach for it when you are assembling a scene outside
   :class:`~isaaclab.scene.InteractiveScene` or want fine control over timing.
 * The third is a one-shot shortcut for the case where every env is just a copy
@@ -227,9 +227,8 @@ them is purely about ergonomics:
 ~~~~~~~~~~~~~~~~~~~~
 
 :class:`~isaaclab.cloner.ReplicateSession` is a context manager that brackets the
-whole cloning lifecycle. Entering the block builds the plan, the body is where
-you construct your assets (each one registers itself as part of its constructor),
-and exiting the block clears those constructor registrations and dispatches the plan:
+whole cloning lifecycle. Entering the block builds and publishes the plan, the body
+constructs assets at their planned source paths, and exiting dispatches that same plan:
 
 .. code-block:: python
 
@@ -260,15 +259,14 @@ When envs need to differ across the population, use
 ``make_clone_plan`` + ``replicate``
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-The same two phases as the session, written as separate function calls. The plan
-is built first, asset construction happens in between, and the drain runs
-explicitly at the end. The gap between the build and the drain is the point —
-that is where you can read the plan back, mutate it, log it, or otherwise
-intervene before replication actually happens:
+The same lifecycle as the session, written as separate function calls. Publish the
+plan before construction so every participant observes the same layout, then
+dispatch it after the prototypes exist:
 
 .. code-block:: python
 
     plan = cloner.make_clone_plan(cfgs, num_clones=N, env_spacing=2.0)
+    sim.set_clone_plan(plan)
     for cfg in cfgs:
         cfg.class_type(cfg)
     cloner.replicate(plan)
@@ -323,18 +321,22 @@ execution contract:
 
 :func:`~isaaclab.cloner.replicate` resolves these types through the
 :class:`~isaaclab.sim.SimulationContext` backend registry, orders them by
-``replicate_priority``, and passes the same plan to each one:
+``replicate_priority``, and passes the published plan to each one:
 
 .. code-block:: python
 
-    def replicate(plan):
-        for context_type in plan.context_rows:
-            simulation_backends[context_type].replicate(plan)
-        publish(plan)
+    simulation.set_clone_plan(plan)
+    construct_prototypes()
+    for context_type in plan.context_rows:
+        simulation_backends[context_type].replicate(plan)
+
+The cfg-first lifecycle publishes before ``construct_prototypes()``. The direct
+single-source workflow remains post-construction and is published by
+:func:`~isaaclab.cloner.replicate` immediately before dispatch. In either form,
+the simulation accepts one plan and each backend receives that exact object once.
 
 USD runs before native physics contexts so the destination topology exists when
-they consume it. No fallback context is constructed during dispatch. The plan is
-then published to the simulation context for downstream consumers.
+they consume it. No fallback context is constructed during dispatch.
 
 Collision Filtering
 -------------------

source/isaaclab/changelog.d/ooctipus-single-clone-lifecycle.major.rst

@@ -0,0 +1,7 @@
+Changed
+^^^^^^^
+
+* **Breaking:** Published cfg-owned clone plans before scene construction and limited each
+  :class:`~isaaclab.sim.SimulationContext` to one plan and one dispatch. Custom scene composition
+  roots should build and publish one plan before constructing its participants, then pass that same
+  plan once to :func:`~isaaclab.cloner.replicate`.

source/isaaclab/isaaclab/cloner/clone_plan.py

@@ -203,7 +203,10 @@ def make_valid_clone_combinations(
 
 
 def _context_rows(
-    cfgs: tuple[Any, ...], cfg_rows: dict[int, tuple[int, ...]], populated_rows: set[int]
+    cfgs: tuple[Any, ...],
+    cfg_rows: dict[int, tuple[int, ...]],
+    populated_rows: set[int],
+    global_paths: tuple[str, ...] = (),
 ) -> dict[type[object], tuple[int, ...]]:
     """Route plan rows to the clone contexts registered for this simulation."""
     sim = sim_utils.SimulationContext.instance()
@@ -213,7 +216,9 @@ def _context_rows(
     physics_context = sim.physics_manager.clone_context_type
     if physics_context is not None and not isinstance(physics_context, type):
         raise TypeError("PhysicsManager.clone_context_type must be a context class.")
-    rows_by_context: dict[type[object], set[int]] = {}
+    rows_by_context: dict[type[object], set[int]] = (
+        {} if physics_context is None or not global_paths else {physics_context: set()}
+    )
 
     for cfg in cfgs:
         rows = cfg_rows.get(id(cfg))
@@ -236,7 +241,7 @@ def _context_rows(
     return {
         context_type: tuple(sorted(rows & populated_rows))
         for context_type, rows in rows_by_context.items()
-        if rows & populated_rows
+        if rows & populated_rows or context_type is physics_context and bool(global_paths)
     }
 
 
@@ -313,6 +318,7 @@ def set_spawn_paths(spawn_cfg: Any, paths: list[str | None]) -> None:
             env_ids=env_ids,
             positions=positions,
             cfg_rows={},
+            context_rows=_context_rows(cfgs, {}, set(), global_paths),
             global_paths=global_paths,
         )
 
@@ -329,7 +335,7 @@ def set_spawn_paths(spawn_cfg: Any, paths: list[str | None]) -> None:
             env_ids=env_ids,
             positions=positions,
             cfg_rows=cfg_rows,
-            context_rows=_context_rows(cfgs, cfg_rows, {0}),
+            context_rows=_context_rows(cfgs, cfg_rows, {0}, global_paths),
             global_paths=global_paths,
         )
 
@@ -402,7 +408,7 @@ def validate_combinations(combos: np.ndarray, name: str, expected_rows: int | No
         env_ids=env_ids,
         positions=positions,
         cfg_rows=cfg_rows,
-        context_rows=_context_rows(cfgs, cfg_rows, populated_rows),
+        context_rows=_context_rows(cfgs, cfg_rows, populated_rows, global_paths),
         global_paths=global_paths,
     )
 
@@ -445,6 +451,6 @@ def clone_plan_from_env_0(
         env_ids=np.arange(num_clones, dtype=np.int64),
         positions=positions,
         cfg_rows=cfg_rows,
-        context_rows=_context_rows(queued, cfg_rows, {0}),
+        context_rows=_context_rows(queued, cfg_rows, {0}, global_paths),
         global_paths=global_paths,
     )

source/isaaclab/isaaclab/cloner/replicate_session.py

@@ -3,7 +3,7 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
-"""Post-construction clone-plan dispatch and :class:`ReplicateSession` sugar."""
+"""Clone-plan publication and dispatch."""
 
 from __future__ import annotations
 
@@ -17,6 +17,8 @@
 from .clone_plan import make_clone_plan
 from .cloner_cfg import DEFAULT_ENV_TEMPLATE
 from .cloner_strategies import sequential
+from .path import under
+from .query import path_to_source
 from .usd import UsdReplicateContext
 
 if TYPE_CHECKING:
@@ -32,16 +34,27 @@
 
 
 def queue_replication(cfg: Any) -> None:
-    """Register a constructed cfg for post-construction clone planning.
+    """Register a constructed cfg or verify that the active plan owns it.
 
     Args:
         cfg: Asset cfg with resolved ``prim_path``.
     """
-    REPLICATION_QUEUE.append(cfg)
+    sim = SimulationContext.instance()
+    plan = None if sim is None else sim.get_clone_plan()
+    if plan is None:
+        REPLICATION_QUEUE.append(cfg)
+        return
+
+    global_owned = any(under(cfg.prim_path, root) for root in plan.global_paths)
+    if not sim._clone_plan_consumed and (id(cfg) in plan.cfg_rows or global_owned):
+        return
+    if cfg.spawn is None and (global_owned or path_to_source(plan, cfg.prim_path) is not None):
+        return
+    raise RuntimeError(f"{type(cfg).__name__} at {cfg.prim_path!r} is not owned by the active ClonePlan.")
 
 
 def replicate(plan: ClonePlan, *, replicate_physics: bool = True) -> None:
-    """Dispatch a fully routed clone plan and publish it after replication.
+    """Publish and dispatch a fully routed clone plan.
 
     Planning derives routing from the input cfgs; dispatch does not rediscover or reshape that mapping.
     Every context is owned by the active :class:`~isaaclab.sim.SimulationContext` and receives
@@ -66,16 +79,16 @@ def replicate(plan: ClonePlan, *, replicate_physics: bool = True) -> None:
         raise RuntimeError(f"Clone contexts must be registered before plan dispatch: {names}.")
 
     contexts = [sim._backend_registry[context_type] for context_type in context_types]
+    sim._consume_clone_plan(plan)
     for context in sorted(contexts, key=lambda item: item.replicate_priority):
         context.replicate(plan)
-    sim.set_clone_plan(plan)
 
 
 class ReplicateSession:
     """Folds :func:`make_clone_plan` and :func:`replicate` into a ``with`` block.
 
-    ``__enter__`` builds the complete plan and mutates each cfg's ``spawn_path``;
-    ``__exit__`` clears constructor registrations and dispatches that same plan.
+    ``__enter__`` builds and publishes the complete plan while assigning each cfg's
+    ``spawn_path``; ``__exit__`` dispatches that same plan.
 
     Example:
 
@@ -125,7 +138,13 @@ def __init__(
         self._plan: ClonePlan | None = None
 
     def __enter__(self) -> ReplicateSession:
+        sim = SimulationContext.instance()
+        if sim is None:
+            raise RuntimeError("Clone planning requires an active SimulationContext.")
+        if sim.get_clone_plan() is not None:
+            raise RuntimeError("A SimulationContext owns exactly one clone lifecycle.")
         self._plan = make_clone_plan(self._cfgs, **self._kwargs)
+        sim.set_clone_plan(self._plan)
         return self
 
     def __exit__(self, exc_type, exc_value, traceback) -> None:
@@ -135,6 +154,9 @@ def __exit__(self, exc_type, exc_value, traceback) -> None:
         else:
             # Drop cfgs registered before the failure so the next session is clean.
             REPLICATION_QUEUE.clear()
+            sim = SimulationContext.instance()
+            if sim is not None and sim.get_clone_plan() is self._plan:
+                sim.set_clone_plan(None)
 
     @property
     def plan(self) -> ClonePlan:

source/isaaclab/isaaclab/scene/interactive_scene.py

@@ -171,34 +171,26 @@ def __init__(self, cfg: InteractiveSceneCfg):
         self._scene_asset_names: list[str] = []
         self._clone_valid_set: np.ndarray | None = None
 
-        self._ALL_INDICES = torch.arange(self.cfg.num_envs, dtype=torch.long, device=self.device)
-
-        # Always enter so a ClonePlan is published even when the scene cfg has no entities.
         self._global_prim_paths = list()
         clone_cfgs, global_paths = self._collect_asset_cfgs()
-        with cloner.ReplicateSession(
-            clone_cfgs,
-            num_clones=self.num_envs,
-            env_spacing=self.cfg.env_spacing,
-            global_paths=global_paths,
-            env_template=self._env_fmt,
-            clone_strategy=self.cloner_cfg.clone_strategy,
-            valid_set=self._clone_valid_set,
-            replicate_physics=self.cloner_cfg.replicate_physics,
-        ) as session:
-            self.stage.DefinePrim(self.env_prim_paths[0], "Xform")
-            with cloner.disabled_fabric_change_notifies(self.stage, restore=False):
-                cloner.usd_replicate(
-                    self.stage,
-                    [self.env_prim_paths[0]],
-                    [self._env_fmt],
-                    session.plan.env_ids,
-                    positions=session.plan.positions,
-                )
-            if self._is_scene_setup_from_cfg():
+        scene_from_cfg = bool(self._scene_asset_names)
+        if scene_from_cfg:
+            with cloner.ReplicateSession(
+                clone_cfgs,
+                num_clones=self.num_envs,
+                env_spacing=self.cfg.env_spacing,
+                global_paths=global_paths,
+                env_template=self._env_fmt,
+                clone_strategy=self.cloner_cfg.clone_strategy,
+                valid_set=self._clone_valid_set,
+                replicate_physics=self.cloner_cfg.replicate_physics,
+            ) as session:
+                self._author_envs(session.plan.env_ids, session.plan.positions)
                 self._add_entities_from_cfg()
-        self._env_origins_plan = session.plan
-        self._env_origins = torch.as_tensor(session.plan.positions, device=self.device)
+        else:
+            env_ids = np.arange(self.num_envs, dtype=np.int64)
+            env_origins = cloner.grid_transforms(self.num_envs, self.cfg.env_spacing)[0]
+            self._author_envs(env_ids, env_origins)
 
         # Every sensor exists by now, so all visualizer and camera-renderer requirements are visible.
         cam_types = [s.cfg.renderer_cfg.renderer_type for s in self._sensors.values() if isinstance(s.cfg, CameraCfg)]
@@ -208,7 +200,7 @@ def __init__(self, cfg: InteractiveSceneCfg):
             self.sim.requires_newton_model |= requires_model
 
         # Collision filtering is PhysX-only (matches both physx and ovphysx).
-        if self.cfg.filter_collisions and "physx" in self.physics_backend and self._is_scene_setup_from_cfg():
+        if self.cfg.filter_collisions and "physx" in self.physics_backend and scene_from_cfg:
             self.filter_collisions(self._global_prim_paths)
 
     def _collect_asset_cfgs(self) -> tuple[list[Any], tuple[str, ...]]:
@@ -375,17 +367,14 @@ def num_envs(self) -> int:
 
     @property
     def env_origins(self) -> torch.Tensor:
-        """Per-env world origins, shape ``(num_envs, 3)``. From the terrain when registered,
-        else from the published :class:`~isaaclab.cloner.ClonePlan`.
-        """
+        """Per-env world origins, shape ``(num_envs, 3)``."""
         if self._terrain is not None:
             return self._terrain.env_origins
-        plan = self.clone_plan
-        if plan is None or plan.positions is None:
-            raise RuntimeError("Environment origins require a published clone plan with positions.")
+        plan = self.sim.get_clone_plan()
         if plan is not self._env_origins_plan:
-            self._env_origins = torch.as_tensor(plan.positions, device=self.device)
             self._env_origins_plan = plan
+            if plan is not None and plan.positions is not None:
+                self._env_origins = torch.as_tensor(plan.positions, device=self.device)
         return self._env_origins
 
     @property
@@ -440,11 +429,12 @@ def visual_materials(self) -> dict[str, VisualMaterial]:
 
     @property
     def clone_plan(self) -> cloner.ClonePlan | None:
-        """Clone plan produced by the most recent replication.
+        """Clone plan owned by the active simulation.
 
         Forwards to :meth:`SimulationContext.get_clone_plan`, which is the canonical owner.
         The plan records the source paths, destination templates, and the per-env source
-        assignment mask. ``None`` until :func:`isaaclab.cloner.replicate` has run.
+        assignment mask. Cfg-owned scenes publish it before constructing their entities;
+        direct scenes publish it when their explicit clone lifecycle begins.
         """
         return self.sim.get_clone_plan()
 
@@ -789,16 +779,14 @@ def __getitem__(self, key: str) -> Any:
     Internal methods.
     """
 
-    def _is_scene_setup_from_cfg(self) -> bool:
-        """Check if scene entities are setup from the config or not.
-
-        Returns:
-            True if scene entities are setup from the config, False otherwise.
-        """
-        return any(
-            not (asset_name in InteractiveSceneCfg.__dataclass_fields__ or asset_cfg is None)
-            for asset_name, asset_cfg in self.cfg.__dict__.items()
-        )
+    def _author_envs(self, env_ids: np.ndarray, positions: np.ndarray) -> None:
+        """Author environment roots from the active layout."""
+        self._ALL_INDICES = torch.as_tensor(env_ids, device=self.device)
+        self._env_origins = torch.as_tensor(positions, device=self.device)
+        self._env_origins_plan = self.sim.get_clone_plan()
+        self.stage.DefinePrim(self.env_prim_paths[0], "Xform")
+        with cloner.disabled_fabric_change_notifies(self.stage, restore=False):
+            cloner.usd_replicate(self.stage, [self.env_prim_paths[0]], [self._env_fmt], env_ids, positions=positions)
 
     def _add_entities_from_cfg(self):  # noqa: C901
         """Add scene entities from the config."""