From fa5b06cfa058554742eebcaf2e12184903006c13 Mon Sep 17 00:00:00 2001 From: Octi Zhang Date: Mon, 25 May 2026 04:05:03 -0700 Subject: [PATCH 1/7] commit --- docs/source/how-to/cloning.rst | 569 +++++++++--------- .../migration/migrating_from_isaacgymenvs.rst | 15 +- .../migrating_from_omniisaacgymenvs.rst | 6 +- .../setup/walkthrough/api_env_design.rst | 11 +- .../walkthrough/technical_env_design.rst | 9 +- .../setup/walkthrough/training_jetbot_gt.rst | 9 +- scripts/demos/pick_and_place.py | 7 +- scripts/tutorials/06_deploy/anymal_c_env.py | 6 +- .../replication-session-redesign.minor.rst | 80 +++ source/isaaclab/isaaclab/cloner/__init__.pyi | 16 +- .../isaaclab/cloner/_fabric_notices.py | 81 +++ source/isaaclab/isaaclab/cloner/clone_plan.py | 63 +- source/isaaclab/isaaclab/cloner/cloner_cfg.py | 29 +- .../isaaclab/isaaclab/cloner/cloner_utils.py | 341 +++++------ .../isaaclab/cloner/replicate_session.py | 132 ++++ source/isaaclab/isaaclab/cloner/usd.py | 182 ++++++ .../isaaclab/scene/interactive_scene.py | 286 +++------ .../isaaclab/scene/interactive_scene_cfg.py | 13 +- .../isaaclab/sensors/camera/camera.py | 2 + .../sensors/ray_caster/base_ray_caster.py | 2 + .../isaaclab/isaaclab/sensors/sensor_base.py | 1 + .../isaaclab/sim/simulation_context.py | 8 +- .../test/scene/test_interactive_scene.py | 303 ++-------- .../test_multi_mesh_ray_caster_camera.py | 3 + .../sensors/test_ray_caster_integration.py | 3 + source/isaaclab/test/sim/test_cloner.py | 361 +++++++++-- .../replication-session-redesign.skip | 5 + .../assets/articulation/articulation.py | 2 + .../assets/rigid_object/rigid_object.py | 2 + .../rigid_object_collection.py | 5 +- .../isaaclab_newton/cloner/__init__.pyi | 9 +- .../{newton_replicate.py => replicate.py} | 211 +++++-- .../test/cloner/test_rename_builder_labels.py | 2 +- .../assets/articulation/articulation.py | 4 + .../assets/rigid_object/rigid_object.py | 4 + .../rigid_object_collection.py | 6 +- .../isaaclab_ovphysx/cloner/__init__.py | 4 +- .../cloner/ovphysx_replicate.py | 105 ---- .../isaaclab_ovphysx/cloner/replicate.py | 185 ++++++ .../physics/ovphysx_manager.py | 11 +- .../replication-session-redesign.skip | 5 + .../assets/articulation/articulation.py | 4 + .../deformable_object/deformable_object.py | 4 + .../assets/rigid_object/rigid_object.py | 4 + .../rigid_object_collection.py | 7 +- .../assets/surface_gripper/surface_gripper.py | 6 + .../isaaclab_physx/cloner/__init__.pyi | 4 +- .../isaaclab_physx/cloner/physx_replicate.py | 113 ---- .../isaaclab_physx/cloner/replicate.py | 196 ++++++ source/isaaclab_physx/test/sim/test_cloner.py | 89 ++- .../replication-session-redesign.skip | 4 + .../direct/anymal_c/anymal_c_env.py | 7 +- .../direct/automate/assembly_env.py | 6 +- .../direct/automate/disassembly_env.py | 6 +- .../cart_double_pendulum_env.py | 7 +- .../direct/cartpole/cartpole_camera_env.py | 7 +- .../direct/cartpole/cartpole_env.py | 7 +- .../direct/factory/factory_env.py | 6 +- .../franka_cabinet/franka_cabinet_env.py | 7 +- .../direct/humanoid_amp/humanoid_amp_env.py | 7 +- .../inhand_manipulation_env.py | 7 +- .../direct/locomotion/locomotion_env.py | 7 +- .../direct/quadcopter/quadcopter_env.py | 7 +- .../shadow_hand/shadow_hand_vision_env.py | 7 +- .../shadow_hand_over/shadow_hand_over_env.py | 7 +- .../direct/cartpole/cartpole_warp_env.py | 7 +- .../inhand_manipulation_warp_env.py | 7 +- .../direct/locomotion/locomotion_env_warp.py | 7 +- .../templates/tasks/direct_multi-agent/env | 8 +- .../templates/tasks/direct_single-agent/env | 8 +- 70 files changed, 2212 insertions(+), 1449 deletions(-) create mode 100644 source/isaaclab/changelog.d/replication-session-redesign.minor.rst create mode 100644 source/isaaclab/isaaclab/cloner/replicate_session.py create mode 100644 source/isaaclab/isaaclab/cloner/usd.py create mode 100644 source/isaaclab_newton/changelog.d/replication-session-redesign.skip rename source/isaaclab_newton/isaaclab_newton/cloner/{newton_replicate.py => replicate.py} (67%) delete mode 100644 source/isaaclab_ovphysx/isaaclab_ovphysx/cloner/ovphysx_replicate.py create mode 100644 source/isaaclab_ovphysx/isaaclab_ovphysx/cloner/replicate.py create mode 100644 source/isaaclab_physx/changelog.d/replication-session-redesign.skip delete mode 100644 source/isaaclab_physx/isaaclab_physx/cloner/physx_replicate.py create mode 100644 source/isaaclab_physx/isaaclab_physx/cloner/replicate.py create mode 100644 source/isaaclab_tasks/changelog.d/replication-session-redesign.skip diff --git a/docs/source/how-to/cloning.rst b/docs/source/how-to/cloning.rst index 4377923bfa45..8f3710acc11e 100644 --- a/docs/source/how-to/cloning.rst +++ b/docs/source/how-to/cloning.rst @@ -5,347 +5,361 @@ Cloning Environments .. currentmodule:: isaaclab -Isaac Lab creates many parallel environments by spawning representative source prims and -then cloning them to the remaining environment paths. This guide starts with direct cloning -so the primitive contract is clear, then shows how :class:`~isaaclab.cloner.ClonePlan` and -:class:`~isaaclab.scene.InteractiveScene` build on top of that contract. +Parallel simulation at scale needs many environments stepping side by side — +hundreds, sometimes tens of thousands per GPU — and authoring each of those envs +by hand would be hopelessly slow. Cloning is Isaac Lab's answer: you author a +small representative scene under ``/World/envs/env_n`` and the cloner expands it +across the rest of the env population for you, optionally with per-env variation. + +The expansion itself is performed by each physics backend's native replicator — +USD, PhysX, or Newton — wrapped by Isaac Lab's core :mod:`isaaclab.cloner` module +behind a single uniform surface so the same user code works regardless of which +backend is active. .. contents:: On this page :local: :depth: 2 -Direct Cloning --------------- +The Backend Layer +----------------- + +At the bottom of the stack, each backend exposes a single function that takes a +flat description of the world layout and materializes it on its runtime. The +signatures are deliberately parallel so the layers above can target every backend +through one interface: + +.. code-block:: text -Use direct cloning for custom scene pipelines, tooling, or tests that need explicit -control over the replication contract. + backend_replicate(stage, sources, destinations, env_ids, mask, positions=None, quaternions=None, ...) -The cloner operates on three pieces of data: +The arguments are parallel arrays describing the layout: -1. **Source prims** that already exist on the stage. -2. **Destination templates** containing ``{}``, which is formatted with each environment id. -3. **A boolean mask** with shape ``[len(sources), num_envs]`` that selects which source - populates each environment. +* ``sources`` — source prim paths already authored on the stage. +* ``destinations`` — destination templates containing ``"{}"``, formatted with each env id. +* ``env_ids`` — long tensor of target env indices. +* ``mask`` — bool tensor of shape ``[len(sources), num_envs]``; ``mask[i, j]`` is + ``True`` when env ``j`` should be populated from source ``i``. +* ``positions`` / ``quaternions`` — optional per-env world transforms. -The direct flow is: -1. Create the environment namespace prims. -2. Spawn representative source prims. -3. Call the physics replicate function for your backend. -4. Call :func:`~isaaclab.cloner.usd_replicate` with the same source-to-environment mapping. +Standalone Examples +~~~~~~~~~~~~~~~~~~~ + +Direct calls into the backend functions, for tooling or tests that need full +control. Production code reaches for one of the ways in +`Cloning in a Backend-Agnostic Way`_ instead. + +**USD** — clone a visual cube across envs: .. code-block:: python import torch - import isaaclab.sim as sim_utils from isaaclab.cloner import usd_replicate - from isaaclab_physx.cloner import physx_replicate num_envs = 128 stage = sim_utils.get_current_stage() - env_ids = torch.arange(num_envs, device="cuda:0") - - sim_utils.create_prim("/World/envs", "Xform") - for env_id in range(num_envs): - sim_utils.create_prim(f"/World/envs/env_{env_id}", "Xform") - - source = "/World/envs/env_0/Cube" - destination = "/World/envs/env_{}/Object" + cube_cfg = sim_utils.CuboidCfg(size=(0.1, 0.1, 0.1)) + cube_cfg.func("/World/envs/env_0/Cube", cube_cfg) + + usd_replicate( + stage, + sources=["/World/envs/env_0/Cube"], + destinations=["/World/envs/env_{}/Cube"], + env_ids=torch.arange(num_envs, device="cuda:0"), + mask=torch.ones((1, num_envs), dtype=torch.bool, device="cuda:0"), + ) - cube_cfg = sim_utils.CuboidCfg(size=(0.5, 0.5, 0.5)) - cube_cfg.func(source, cube_cfg) +**PhysX** — call PhysX and USD on the same sources and destinations (either order): - mask = torch.ones((1, num_envs), dtype=torch.bool, device="cuda:0") +.. code-block:: python - physx_replicate(stage, [source], [destination], env_ids, mask, device="cuda:0") - usd_replicate(stage, [source], [destination], env_ids, mask) + from isaaclab_physx.cloner import physx_replicate -This creates one source cube at ``/World/envs/env_0/Cube`` and clones it to -``/World/envs/env_1/Object`` through ``/World/envs/env_127/Object``. When a source path is -the same as the destination for an environment, ``usd_replicate`` skips the self-copy. + physx_replicate(stage, sources, destinations, env_ids, mask) + usd_replicate(stage, sources, destinations, env_ids, mask) -Direct heterogeneous cloning uses the same API with more source rows. Each row in ``mask`` -selects the environments that receive the matching source. For example, this explicit mask -clones a cone into environments 0 and 2, and a sphere into environments 1 and 3: +**Newton**: .. code-block:: python - env_ids = torch.arange(4, device="cuda:0") - sources = ["/World/envs/env_0/Cone", "/World/envs/env_1/Sphere"] - destinations = ["/World/envs/env_{}/Object", "/World/envs/env_{}/Object"] + from isaaclab_newton.cloner import newton_physics_replicate - cone_cfg = sim_utils.ConeCfg(radius=0.25, height=0.5) - sphere_cfg = sim_utils.SphereCfg(radius=0.25) - cone_cfg.func(sources[0], cone_cfg) - sphere_cfg.func(sources[1], sphere_cfg) + newton_physics_replicate(stage, sources, destinations, env_ids, mapping=mask) - mask = torch.tensor([[True, False, True, False], [False, True, False, True]], dtype=torch.bool) - physx_replicate(stage, sources, destinations, env_ids, mask, device="cuda:0") - usd_replicate(stage, sources, destinations, env_ids, mask) +Cloning in a Backend-Agnostic Way +--------------------------------- -The mask above reads as: +Authoring every prim in every env by hand would be prohibitively slow and would +also tie scene code to whichever physics engine happens to be active. Isaac Lab +sidesteps both problems with a single central abstraction: +:class:`~isaaclab.cloner.ClonePlan` — a compact description of how a small set of +prim-level prototypes maps onto the full population of envs, with each prototype +free to land in some envs and not others. A plan is built once, fed to each backend, and +lets every engine take its own fastest replication path: USD instancing for +visuals, PhysX's native replicator for rigid bodies and articulations, Newton's +world system for its parallel pipeline. The same plan drives all of them, so user +code never branches on the backend. + +ClonePlan +~~~~~~~~~ + +A plan holds the parallel arrays a backend replicate consumes — sources, +destinations, mask, env ids — in one place. Conceptually it is a small table +where each row describes one distinct prototype-to-destination mapping; the +fields listed below are that table's columns. Every entry point in +:mod:`isaaclab.cloner` either produces a plan, consumes a plan, or both, so a +quick look at the fields is the fastest way to build intuition for the rest of +this page: .. list-table:: :header-rows: 1 - :widths: 15 40 20 25 - - * - Source row - - Source path - - Env ids - - Destination path - * - ``0`` - - ``/World/envs/env_0/Cone`` - - ``0, 2`` - - ``/World/envs/env_{}/Object`` - * - ``1`` - - ``/World/envs/env_1/Sphere`` - - ``1, 3`` - - ``/World/envs/env_{}/Object`` - -``usd_replicate`` copies parent paths before children and supports optional ``positions`` -and ``quaternions`` buffers. If ``positions`` is provided, it authors -``xformOp:translate`` on each destination using the environment id. The helper -:func:`~isaaclab.cloner.grid_transforms` creates the same grid layout used by -:class:`~isaaclab.scene.InteractiveScene`. - -.. code-block:: python + :widths: 22 78 - from isaaclab.cloner import grid_transforms - - positions, orientations = grid_transforms( - N=num_envs, - spacing=2.0, - up_axis="z", - device="cuda:0", - ) - usd_replicate(stage, [source], [destination], env_ids, mask, positions=positions) - - -Clone Plans ------------ - -For one source row, passing ``sources``, ``destinations``, and ``mask`` by hand is simple. -For heterogeneous scenes, the mapping is easier to build with -:func:`~isaaclab.cloner.make_clone_plan`. - -:class:`~isaaclab.cloner.ClonePlan` stores the same flat contract used by direct cloning: + * - Field + - Meaning + * - ``sources`` + - Source prim paths, one per replication row. + * - ``destinations`` + - Destination templates with ``"{}"`` for the env id, one per row. + * - ``clone_mask`` + - Bool tensor ``[len(sources), num_envs]``; ``True`` when env ``j`` comes from row ``i``. + * - ``env_ids`` + - Long tensor of target env ids. + * - ``positions`` + - Optional per-env world positions [m], shape ``[num_envs, 3]``. + +The plan is stage-agnostic by design — the same instance can be replayed against a +different stage, inspected by tooling, or serialized. + +When every env is a copy of env_0: .. code-block:: text - sources = [source_0, source_1, ...] - destinations = [destination_0, destination_1, ...] - clone_mask = bool tensor, shape [len(sources), num_envs] + sources = ("/World/envs/env_0",) + destinations = ("/World/envs/env_{}",) + clone_mask = [[True, True, ..., True]] -``clone_mask[i, j]`` is ``True`` when environment ``j`` should receive source row ``i``. -The same plan can be passed to USD replication, physics replication, and scene-data -providers. +When envs differ — say a cartpole in every env plus a 2-variant obstacle (box into +envs 0/1, sphere into envs 2/3): -Homogeneous Plans -~~~~~~~~~~~~~~~~~ +.. code-block:: text -In a homogeneous scene, every environment receives the same asset layout. The default plan -is: + sources = ("/World/envs/env_0/Cartpole", + "/World/envs/env_0/Obstacle_0", # box prototype + "/World/envs/env_0/Obstacle_1") # sphere prototype + destinations = ("/World/envs/env_{}/Cartpole", + "/World/envs/env_{}/Obstacle", + "/World/envs/env_{}/Obstacle") + clone_mask = [[1, 1, 1, 1], + [1, 1, 0, 0], + [0, 0, 1, 1]] + +A plan is the *what*. Putting one together and handing it to the backends is +the *how*, and Isaac Lab exposes three idiomatic ways to do that. All three end +in the same ``cloner.replicate(plan, stage=...)`` call, so the choice between +them is purely about ergonomics: + +* The first wraps both phases in a context manager and is what + :class:`~isaaclab.scene.InteractiveScene` runs under the hood. Reach for it + 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. + 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 + of env_0. Reach for it in :class:`~isaaclab.envs.DirectRLEnv` and standalone + scripts that hand-build the env-0 prototype prim by prim. + +``ReplicateSession`` +~~~~~~~~~~~~~~~~~~~~ + +: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 drains every registration against the plan: -.. code-block:: text +.. code-block:: python - sources = ["/World/envs/env_0"] - destinations = ["/World/envs/env_{}"] - clone_mask = all True, shape [1, num_envs] + with cloner.ReplicateSession(cfgs, num_clones=N, env_spacing=2.0, + device=device, stage=stage): + for cfg in cfgs: + cfg.class_type(cfg) -This means the scene spawns everything for ``env_0`` and replicates that environment to -``env_1`` through ``env_N``. +This is what :class:`~isaaclab.scene.InteractiveScene` runs when you declare assets +in an :class:`~isaaclab.scene.InteractiveSceneCfg`: -Heterogeneous Plans -~~~~~~~~~~~~~~~~~~~ +.. code-block:: python -Heterogeneous cloning is used when different environments receive different prototypes. -For example, an object with three variants may have representative source prims at: + @configclass + class MySceneCfg(InteractiveSceneCfg): + robot = CARTPOLE_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + light = AssetBaseCfg( + prim_path="/World/Light", + spawn=sim_utils.DistantLightCfg(intensity=3000.0), + ) -.. code-block:: text + scene = InteractiveScene(MySceneCfg(num_envs=128, env_spacing=2.0)) - /World/envs/env_0/Object - /World/envs/env_1/Object - /World/envs/env_2/Object +When envs need to differ across the population, use +:class:`~isaaclab.sim.spawners.wrappers.MultiAssetSpawnerCfg` or +:class:`~isaaclab.sim.spawners.wrappers.MultiUsdFileCfg`; see +:doc:`multi_asset_spawning`. -These paths have the same leaf name because each variant will be cloned to -``/World/envs/env_{}/Object``, but their authored contents are different. For example, -``env_0/Object`` could be a cone, ``env_1/Object`` a cuboid, and ``env_2/Object`` a sphere. +``make_clone_plan`` + ``replicate`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -The plan maps those source rows to all environments: +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: .. code-block:: python - from isaaclab.cloner import make_clone_plan, sequential - - plan = make_clone_plan( - sources=[ - [ - "/World/envs/env_0/Object", - "/World/envs/env_1/Object", - "/World/envs/env_2/Object", - ] - ], - destinations=["/World/envs/env_{}/Object"], - num_clones=8, - clone_strategy=sequential, - device="cuda:0", - ) + plan = cloner.make_clone_plan(cfgs, num_clones=N, env_spacing=2.0, device=device) + for cfg in cfgs: + cfg.class_type(cfg) + cloner.replicate(plan, stage=stage) - # source row used by env: 0, 1, 2, 0, 1, 2, 0, 1 +``ClonePlan.from_env_0`` + ``replicate`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -Direct code can use the plan exactly like the hand-written direct example: +Shortcut for the case where every env is just a copy of env_0. +:meth:`~isaaclab.cloner.ClonePlan.from_env_0` builds the single-source plan in +one line by pointing at the prototype, and :func:`~isaaclab.cloner.replicate` +finishes the setup. This is the pattern most :class:`~isaaclab.envs.DirectRLEnv` +subclasses use — they author the env-0 prototype prim by prim in +``_setup_scene`` and end the method with these four lines: .. code-block:: python - physx_replicate(stage, plan.sources, plan.destinations, env_ids, plan.clone_mask, device="cuda:0") - usd_replicate(stage, plan.sources, plan.destinations, env_ids, plan.clone_mask) - -When variants span multiple groups, such as robot variants and object variants, -``make_clone_plan`` enumerates the Cartesian product of the groups and assigns one -combination per environment. Unused prototype rows may still appear in the plan with an -all-false mask row. + def _setup_scene(self): + self.cartpole = Articulation(self.cfg.robot_cfg) + spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) + # ... any other assets ... -.. _cloning-strategies: + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) -Clone Strategies -~~~~~~~~~~~~~~~~ +Every env receives the same prototype. When envs need to differ, use one of the +other two. -A clone strategy chooses prototype combinations for the environments: -* :func:`~isaaclab.cloner.random` samples combinations randomly and is the default. -* :func:`~isaaclab.cloner.sequential` assigns combinations in round-robin order, which is - useful for reproducible tests and balanced coverage. +Under the Hood +-------------- -Custom strategies are callables with this signature: +To see how the backend-agnostic surface works, follow one asset through the +system. Suppose you write ``Articulation(cfg)`` for a PhysX articulation +somewhere inside a :class:`~isaaclab.cloner.ReplicateSession`. The constructor +does not actually clone anything yet — at that moment the plan describing how +the full env population should be laid out may not even exist. Instead the +constructor *registers* the asset with the cloner, the cloner files the +registration into a queue, and later — when the session exits and the cloner +runs replication — that registration is handed to the backend code that knows +how to replicate a PhysX articulation, with the plan telling it where each +clone goes. + +The story has to look like this because the engines underneath disagree about +*when* and *how* replication actually happens: + +* **PhysX** defers the real work to physics runtime. At construction time the + only thing user code can do is register intent; PhysX replays those + registrations entity by entity when the simulation comes up. +* **USD** is declarative and immediate — calling :func:`~isaaclab.cloner.usd_replicate` + materializes the clones in place, right then and there. +* **Newton** is also declarative and immediate, but it insists on replicating + the whole world in one shot rather than asset by asset, so the framework + cannot just hand it one cfg at a time — everything Newton-related has to be + assembled first. + +Isaac Lab reconciles these into one surface with two small pieces of plumbing. +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 +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: -.. code-block:: python +.. code-block:: text - def my_strategy(combinations: torch.Tensor, num_clones: int, device: str) -> torch.Tensor: + REPLICATION_QUEUE + (cartpole_cfg, PhysxReplicateContext) + (cartpole_cfg, UsdReplicateContext) + (cube_cfg, UsdReplicateContext) + (light_cfg, UsdReplicateContext) ... -``combinations`` has shape ``[num_combinations, num_groups]`` and the return value must have -shape ``[num_clones, num_groups]``. - - -Common Workflow: ``InteractiveScene`` -------------------------------------- +Deferring the work like this buys three things at once: -:class:`~isaaclab.scene.InteractiveScene` automates the direct cloning flow for task scenes. -It inspects scene configuration, builds a :class:`~isaaclab.cloner.ClonePlan`, rewrites -spawner paths to the representative sources, spawns those sources, runs physics and USD -replication, and filters inter-environment collisions for PhysX when configured. +* Replication can wait until the plan is fully built, so the final layout is + known before any prims are spawned. +* Every asset's request is batched into a single backend call instead of one + call per asset. +* Asset code stays free of any branching on which backend is active — it just + registers and lets the framework take it from there. -Put per-environment assets under ``{ENV_REGEX_NS}`` and global assets under normal USD -paths: - -.. code-block:: python +Backend contexts +~~~~~~~~~~~~~~~~ - import isaaclab.sim as sim_utils - from isaaclab.assets import AssetBaseCfg - from isaaclab.scene import InteractiveScene, InteractiveSceneCfg - from isaaclab.utils.configclass import configclass - from isaaclab_assets.robots.cartpole import CARTPOLE_CFG +Each backend ships a small adapter class — its *replicate context* — that +knows how to take a registered cfg and replicate it on the backend's specific +runtime: +.. code-block:: text - @configclass - class MySceneCfg(InteractiveSceneCfg): - # Cloned once per environment. - robot = CARTPOLE_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + UsdReplicateContext # replicates USD prim subtrees + PhysxReplicateContext # replicates PhysX rigid bodies and articulations + NewtonReplicateContext # replicates Newton bodies in its parallel pipeline - # Authored once globally, not cloned per environment. - light = AssetBaseCfg( - prim_path="/World/Light", - spawn=sim_utils.DistantLightCfg(intensity=3000.0), - ) +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. +Running replication +~~~~~~~~~~~~~~~~~~~ - scene_cfg = MySceneCfg(num_envs=128, env_spacing=2.0, replicate_physics=True) - scene = InteractiveScene(cfg=scene_cfg) - -For heterogeneous scenes, use :class:`~isaaclab.sim.spawners.wrappers.MultiAssetSpawnerCfg` -or :class:`~isaaclab.sim.spawners.wrappers.MultiUsdFileCfg`. ``InteractiveScene`` assigns -representative source paths to the spawner and lets the clone strategy choose which -prototype each environment receives. See :doc:`multi_asset_spawning` for the asset -configuration details. +:func:`~isaaclab.cloner.replicate` is what actually runs the registered work. +The dispatch shape is roughly: -The most important scene options are on :class:`~isaaclab.scene.InteractiveSceneCfg`: +.. code-block:: python -.. list-table:: - :header-rows: 1 - :widths: 25 15 60 + def replicate(plan, stage): + for context_cls, rows in group_queue_by_context(plan): + context_cls().replicate(rows=rows, stage=stage) + publish(plan) - * - Field - - Default - - When to change it - * - ``replicate_physics`` - - ``True`` - - Keep enabled for homogeneous environments and fast startup. Disable it when each - environment needs independently authored physics or USD randomization. - * - ``filter_collisions`` - - ``True`` - - Keep enabled for parallel RL so cloned environments do not collide with each other. - This is automatic for PhysX-backed scene cloning. - * - ``clone_in_fabric`` - - ``False`` - - Enables the PhysX Fabric cloning path for faster scene creation. Use USDRT for stage - inspection when Fabric cloning is enabled. - - -Choosing an API ---------------- +Contexts run in a priority order that puts physics ahead of visuals, and the +plan is published to :class:`~isaaclab.sim.SimulationContext` so the rest of the +framework can read the per-env layout back. -.. list-table:: - :header-rows: 1 - :widths: 25 45 30 - - * - Goal - - Recommended API - - Notes - * - Build a custom cloning pipeline - - :func:`~isaaclab.cloner.usd_replicate` and a backend physics replicate function - - Useful for tests, tooling, or advanced scene construction. - * - Build complex direct mappings - - :func:`~isaaclab.cloner.make_clone_plan` - - Produces the same ``sources``, ``destinations``, and ``clone_mask`` used by direct cloning. - * - Build normal task scenes - - :class:`~isaaclab.scene.InteractiveScene` - - Preferred path. Configure assets with ``{ENV_REGEX_NS}`` and let the scene clone them. - * - Randomize which asset each environment receives - - ``InteractiveScene`` with :class:`~isaaclab.sim.spawners.wrappers.MultiAssetSpawnerCfg` or - :class:`~isaaclab.sim.spawners.wrappers.MultiUsdFileCfg` - - See :doc:`multi_asset_spawning` for the asset configuration details. - * - Use Isaac Sim's ``GridCloner`` - - Isaac Sim API - - Isaac Lab's tested path is the ``isaaclab.cloner`` API described here. - - -Migrating From Template Cloning -------------------------------- - -The template-root discovery API has been removed. Replace -``clone_from_template(...)`` calls with explicit source prims plus -:func:`~isaaclab.cloner.make_clone_plan`, a backend physics replicate function, and -:func:`~isaaclab.cloner.usd_replicate`. Replace ``TemplateCloneCfg`` with -:class:`~isaaclab.cloner.CloneCfg` for execution settings such as clone strategy, -Fabric cloning, and backend replication. - - -Collision Filtering and Isolation ---------------------------------- +Collision Filtering +------------------- -Some prims, such as terrain, are intentionally shared across environments and should collide -with every environment. These are modeled as global collision paths. The workaround is only -the per-environment filtering: when cloning is fully isolated per world, cloned environments -should not collide with each other and no manual per-environment filter should be needed. -Some PhysX cloning paths still rely on USD collision groups for that isolation fallback. In -the scene workflow this is handled by ``InteractiveScene`` when ``filter_collisions=True`` -and the backend is PhysX. +PhysX models per-env isolation through collision groups, so PhysX scenes need a +filtering pass after cloning to keep envs from colliding with each other while +still letting them collide with global prims (terrain, ground planes, lights). -For direct PhysX usage, call :func:`~isaaclab.cloner.filter_collisions` after cloning if -per-environment isolation is not already provided by the cloning backend: +:class:`~isaaclab.scene.InteractiveScene` runs that pass automatically when +``filter_collisions=True`` and the backend is PhysX. For direct PhysX pipelines, +call :func:`~isaaclab.cloner.filter_collisions` after the replicate: .. code-block:: python @@ -359,39 +373,4 @@ per-environment isolation is not already provided by the cloning backend: global_paths=["/World/ground"], ) -.. note:: - - Collision filtering uses PhysX collision groups. Newton handles per-environment isolation - through its own world system. - - -Backend and Option Notes ------------------------- - -**Physics replication.** :class:`~isaaclab.scene.InteractiveScene` selects the backend -replication function automatically. Direct PhysX users call -:func:`~isaaclab_physx.cloner.physx_replicate`; Newton users call -:func:`~isaaclab_newton.cloner.newton_physics_replicate`. - -**``replicate_physics=False``.** Disable physics replication when environments need -independent authored USD or physics state, such as some scale, texture, or color -randomization workflows. Startup and physics parsing are slower because the backend cannot -assume every environment is a clone of the same source. - -**``copy_from_source``.** ``InteractiveScene`` calls -``clone_environments(copy_from_source=True)`` when ``replicate_physics=False``. This skips -backend physics replication and leaves physics parsing to the backend. Spawner-level -``copy_from_source`` is a separate setting used by spawn functions that clone from a source -path matched by a regex. - -**Fabric cloning.** ``clone_in_fabric=True`` applies to PhysX replication. It can reduce -scene-creation time for large PhysX scenes, especially when many replicated rigid bodies are -authored. Fabric-backed stage data must be inspected through USDRT rather than normal USD -APIs. - - -See Also --------- - -* :doc:`multi_asset_spawning` -- configuring multi-asset and multi-USD spawners. -* :doc:`optimize_stage_creation` -- Fabric cloning and stage-in-memory optimizations. +Newton isolates envs through its world system and does not need this pass. diff --git a/docs/source/migration/migrating_from_isaacgymenvs.rst b/docs/source/migration/migrating_from_isaacgymenvs.rst index db6371c40c9d..a80b8a5b79bf 100644 --- a/docs/source/migration/migrating_from_isaacgymenvs.rst +++ b/docs/source/migration/migrating_from_isaacgymenvs.rst @@ -194,7 +194,7 @@ adding any other optional objects into the scene, such as lights. | self.up_axis = self.cfg["sim"]["up_axis"] | # add ground plane | | | spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg() | | self.sim = super().create_sim(self.device_id, self.graphics_device_id, | # clone, filter, and replicate | -| self.physics_engine, self.sim_params) | self.scene.clone_environments(copy_from_source=False) | +| self.physics_engine, self.sim_params) | # assets are built inside ReplicateSession | | self._create_ground_plane() | self.scene.filter_collisions(global_prim_paths=[]) | | self._create_envs(self.num_envs, self.cfg["env"]['envSpacing'], | # add articulation to scene | | int(np.sqrt(self.num_envs))) | self.scene.articulations["cartpole"] = self.cartpole | @@ -369,16 +369,15 @@ Isaac Lab eliminates the need for looping through the environments by using the The scene creation process is as follow: #. Construct a single environment (what the scene would look like if number of environments = 1) -#. Call ``clone_environments()`` to replicate the single environment +#. Use ``cloner.ReplicateSession`` to replicate the single environment #. Call ``filter_collisions()`` to filter out collision between environments (if required) .. code-block:: python - # construct a single environment with the Cartpole robot - self.cartpole = Articulation(self.cfg.robot_cfg) - # clone the environment - self.scene.clone_environments(copy_from_source=False) + # construct and replicate a single environment with the Cartpole robot + with cloner.ReplicateSession(): + self.cartpole = Articulation(self.cfg.robot_cfg) # filter collisions self.scene.filter_collisions(global_prim_paths=[self.cfg.terrain.prim_path]) @@ -660,8 +659,8 @@ the need to set simulation parameters for actors in the task implementation. | | spawn_ground_plane(prim_path="/World/ground", | | self.sim = super().create_sim(self.device_id, | cfg=GroundPlaneCfg()) | | self.graphics_device_id, self.physics_engine, | # clone, filter, and replicate | -| self.sim_params) | self.scene.clone_environments( | -| self._create_ground_plane() | copy_from_source=False) | +| self.sim_params) | # assets are built inside ReplicateSession | +| self._create_ground_plane() | | | self._create_envs(self.num_envs, | self.scene.filter_collisions( | | self.cfg["env"]['envSpacing'], | global_prim_paths=[]) | | int(np.sqrt(self.num_envs))) | # add articulation to scene | diff --git a/docs/source/migration/migrating_from_omniisaacgymenvs.rst b/docs/source/migration/migrating_from_omniisaacgymenvs.rst index b3a46f0a518f..47f4a6ffe18b 100644 --- a/docs/source/migration/migrating_from_omniisaacgymenvs.rst +++ b/docs/source/migration/migrating_from_omniisaacgymenvs.rst @@ -216,7 +216,7 @@ will automatically be created for the actor. This avoids the need to separately | super().set_up_scene(scene) | # add ground plane | | | spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg() | | self._cartpoles = ArticulationView( | # clone, filter, and replicate | -| prim_paths_expr="/World/envs/.*/Cartpole", | self.scene.clone_environments(copy_from_source=False) | +| prim_paths_expr="/World/envs/.*/Cartpole", | # assets are built inside ReplicateSession | | name="cartpole_view", reset_xform_properties=False | self.scene.filter_collisions(global_prim_paths=[]) | | ) | # add articulation to scene | | scene.add(self._cartpoles) | self.scene.articulations["cartpole"] = self.cartpole | @@ -633,8 +633,8 @@ Adding actors to the scene has been replaced by ``self.scene.articulations["cart | super().set_up_scene(scene) | spawn_ground_plane(prim_path="/World/ground", | | self._cartpoles = ArticulationView( | cfg=GroundPlaneCfg()) | | prim_paths_expr="/World/envs/.*/Cartpole", | # clone, filter, and replicate | -| name="cartpole_view", | self.scene.clone_environments( | -| reset_xform_properties=False | copy_from_source=False) | +| name="cartpole_view", | # assets are built inside ReplicateSession | +| reset_xform_properties=False | | | ) | self.scene.filter_collisions( | | scene.add(self._cartpoles) | global_prim_paths=[]) | | return | # add articulation to scene | diff --git a/docs/source/setup/walkthrough/api_env_design.rst b/docs/source/setup/walkthrough/api_env_design.rst index a669bc47fa1c..9e35e2d4bcd4 100644 --- a/docs/source/setup/walkthrough/api_env_design.rst +++ b/docs/source/setup/walkthrough/api_env_design.rst @@ -96,13 +96,12 @@ Next, let's take a look at the contents of the other python file in our task dir . . . def _setup_scene(self): - self.robot = Articulation(self.cfg.robot_cfg) - # add ground plane - spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) + with cloner.ReplicateSession(): + self.robot = Articulation(self.cfg.robot_cfg) + # add ground plane + spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) # add articulation to scene self.scene.articulations["robot"] = self.robot - # clone and replicate - self.scene.clone_environments(copy_from_source=False) # add lights light_cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75)) light_cfg.func("/World/Light", light_cfg) @@ -143,7 +142,7 @@ When the environment is initialized, it receives its own config as an argument, to initialize the ``DirectRLEnv``. This super call also calls ``_setup_scene``, which actually constructs the scene and clones it appropriately. Notably is how the robot is created and registered to the scene in ``_setup_scene``. First, the robot articulation is created by using the ``robot_config`` we defined in ``IsaacLabTutorialEnvCfg``: it doesn't exist before this point! When the -articulation is created, the robot exists on the stage at ``/World/envs/env_0/Robot``. The call to ``scene.clone_environments`` then +articulation is created, the robot exists on the stage at ``/World/envs/env_0/Robot``. The call to ``cloner.ReplicateSession`` then copies ``env_0`` appropriately. At this point the robot exists as many copies on the stage, so all that's left is to notify the ``scene`` object of the existence of this articulation to be tracked. The articulations of the scene are kept as a dictionary, so ``scene.articulations["robot"] = self.robot`` creates a new ``robot`` element of the ``articulations`` dictionary and sets the value to be ``self.robot``. diff --git a/docs/source/setup/walkthrough/technical_env_design.rst b/docs/source/setup/walkthrough/technical_env_design.rst index 41b99445c135..62a086343f9c 100644 --- a/docs/source/setup/walkthrough/technical_env_design.rst +++ b/docs/source/setup/walkthrough/technical_env_design.rst @@ -105,11 +105,10 @@ replace the contents of the ``__init__`` and ``_setup_scene`` methods with the f self.dof_idx, _ = self.robot.find_joints(self.cfg.dof_names) def _setup_scene(self): - self.robot = Articulation(self.cfg.robot_cfg) - # add ground plane - spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) - # clone and replicate - self.scene.clone_environments(copy_from_source=False) + with cloner.ReplicateSession(): + self.robot = Articulation(self.cfg.robot_cfg) + # add ground plane + spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) # add articulation to scene self.scene.articulations["robot"] = self.robot # add lights diff --git a/docs/source/setup/walkthrough/training_jetbot_gt.rst b/docs/source/setup/walkthrough/training_jetbot_gt.rst index 05e89ef45644..f49cc9e7d45c 100644 --- a/docs/source/setup/walkthrough/training_jetbot_gt.rst +++ b/docs/source/setup/walkthrough/training_jetbot_gt.rst @@ -67,11 +67,10 @@ Next, we need to expand the initialization and setup steps to construct the data .. code-block:: python def _setup_scene(self): - self.robot = Articulation(self.cfg.robot_cfg) - # add ground plane - spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) - # clone and replicate - self.scene.clone_environments(copy_from_source=False) + with cloner.ReplicateSession(): + self.robot = Articulation(self.cfg.robot_cfg) + # add ground plane + spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) # add articulation to scene self.scene.articulations["robot"] = self.robot # add lights diff --git a/scripts/demos/pick_and_place.py b/scripts/demos/pick_and_place.py index d5e77aa4a12a..aa99cb1480fd 100644 --- a/scripts/demos/pick_and_place.py +++ b/scripts/demos/pick_and_place.py @@ -35,6 +35,7 @@ import omni import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import ( Articulation, ArticulationCfg, @@ -221,8 +222,10 @@ def _setup_scene(self): self.gripper = SurfaceGripper(self.cfg.gripper) # add ground plane spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) - # clone and replicate - self.scene.clone_environments(copy_from_source=False) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) # add articulation to scene self.scene.articulations["pick_and_place"] = self.pick_and_place self.scene.rigid_objects["cube"] = self.cube diff --git a/scripts/tutorials/06_deploy/anymal_c_env.py b/scripts/tutorials/06_deploy/anymal_c_env.py index 94022e3a0956..23f35a93a243 100644 --- a/scripts/tutorials/06_deploy/anymal_c_env.py +++ b/scripts/tutorials/06_deploy/anymal_c_env.py @@ -12,6 +12,7 @@ import warp as wp import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import Articulation from isaaclab.envs import DirectRLEnv from isaaclab.sensors import ContactSensor, RayCaster @@ -63,7 +64,10 @@ def _setup_scene(self): self.cfg.terrain.num_envs = self.scene.cfg.num_envs self.cfg.terrain.env_spacing = self.scene.cfg.env_spacing self._terrain = self.cfg.terrain.class_type(self.cfg.terrain) - self.scene.clone_environments(copy_from_source=False) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) if self.device == "cpu": self.scene.filter_collisions(global_prim_paths=[self.cfg.terrain.prim_path]) light_cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75)) diff --git a/source/isaaclab/changelog.d/replication-session-redesign.minor.rst b/source/isaaclab/changelog.d/replication-session-redesign.minor.rst new file mode 100644 index 000000000000..c7578decf575 --- /dev/null +++ b/source/isaaclab/changelog.d/replication-session-redesign.minor.rst @@ -0,0 +1,80 @@ +Added +^^^^^ + +* Added :data:`~isaaclab.cloner.REPLICATION_QUEUE`, a module-level list that backend + ``queue__replication`` helpers append ``(cfg, BackendCtxCls)`` pairs to. + Drained by :func:`~isaaclab.cloner.replicate`. +* Added :func:`~isaaclab.cloner.replicate` as a free function that takes an explicit + :class:`~isaaclab.cloner.ClonePlan`, drains + :data:`~isaaclab.cloner.REPLICATION_QUEUE`, dispatches each backend context in + ascending ``replicate_priority`` order, publishes the plan via + :meth:`~isaaclab.sim.SimulationContext.set_clone_plan`, and clears the queue. +* Added :meth:`~isaaclab.cloner.ClonePlan.from_env_0` classmethod that + constructs a homogeneous single-source clone plan for direct envs and auto-populates + ``cfg_rows`` from :data:`~isaaclab.cloner.REPLICATION_QUEUE` filtered by env-root + prefix. +* Added :attr:`~isaaclab.cloner.CloneCfg.clone_regex` (default ``"/World/envs/env_.*"``) + as the single source of truth for the env-namespace convention. Consumers inside + :class:`~isaaclab.scene.InteractiveScene` read it directly when expanding + ``{ENV_REGEX_NS}`` cfg macros. + +Fixed +^^^^^ + +* Fixed :data:`~isaaclab.cloner.REPLICATION_QUEUE` leaking stale ``(cfg, BackendCtxCls)`` + entries when a backend raised mid-dispatch or when asset construction failed inside + :class:`~isaaclab.cloner.ReplicateSession`. :func:`~isaaclab.cloner.replicate` now + snapshots and clears the queue up front, and :meth:`ReplicateSession.__exit__` clears + it on the exception path. + +Changed +^^^^^^^ + +* **Breaking:** Rewrote :class:`~isaaclab.cloner.ReplicateSession` as a thin context + manager that calls :func:`~isaaclab.cloner.make_clone_plan` in ``__enter__`` and + :func:`~isaaclab.cloner.replicate` in ``__exit__``. The class no longer exposes + ``plan``, ``stage``, ``env_ids``, ``cfg_rows``, ``positions``, ``publish_clone_plan_fn``, + or ``replicate_on_exit`` fields, and the ``ReplicateSession()`` (no-arg) form is gone. + Direct envs migrate to ``cloner.replicate(cloner.ClonePlan.from_env_0(...))``. +* **Breaking:** Changed :func:`~isaaclab.cloner.make_clone_plan` signature from + ``make_clone_plan(sources, destinations, num_clones, clone_strategy, device)`` (pure + flat-mapping helper) to + ``make_clone_plan(cfgs, num_clones, env_spacing, device, *, clone_strategy=sequential, + valid_set=None, stage=None)``. The new function absorbs the cfg-driven planning logic + formerly in ``InteractiveScene._build_clone_plan_from_cfg``, returns a + self-contained :class:`~isaaclab.cloner.ClonePlan` with ``cfg_rows`` populated, and + mutates each cfg's ``spawn_path`` / ``spawn_paths``. +* **Breaking:** :class:`~isaaclab.cloner.ClonePlan` now carries ``env_ids``, + ``positions``, and ``cfg_rows`` fields in addition to ``sources``, ``destinations``, + ``clone_mask``. All fields are required; ``cfg_rows`` is never ``None``. The USD + stage is intentionally *not* part of the plan — pass it explicitly via the new + ``stage=`` keyword on :func:`~isaaclab.cloner.replicate` and + :class:`~isaaclab.cloner.ReplicateSession`. +* **Breaking:** :func:`~isaaclab.cloner.replicate` now requires a ``stage`` keyword + argument: ``cloner.replicate(plan, stage=scene.stage)``. Direct envs typically pass + ``self.scene.stage``. +* **Breaking:** :class:`~isaaclab.cloner.ReplicateSession` now requires a ``stage`` + keyword argument that is forwarded to :func:`replicate` in ``__exit__``. +* Changed :attr:`~isaaclab.scene.InteractiveScene.env_origins` to read its per-env + positions from the :class:`~isaaclab.cloner.ClonePlan` published to + :class:`~isaaclab.sim.SimulationContext`, making the published plan the single + source of truth for env placement. :class:`~isaaclab.scene.InteractiveScene` now + always enters :class:`~isaaclab.cloner.ReplicateSession` (even for scene cfgs + with no entities, where the resulting plan is empty but still carries + ``positions``), so the lookup is unconditional and the prior + ``InteractiveScene._default_env_origins`` cached tensor has been removed. + +Removed +^^^^^^^ + +* **Breaking:** Removed the module-level ``isaaclab.cloner.replicate_session_defaults`` + dictionary and the ``isaaclab.cloner.replicate_session`` active-session reference. + Replication is now driven by the explicit :data:`~isaaclab.cloner.REPLICATION_QUEUE` + registry and :func:`~isaaclab.cloner.replicate` drain. +* **Breaking:** Removed :meth:`InteractiveScene.clone_environments` deprecated shim. + Direct envs should use ``cloner.replicate(cloner.ClonePlan.from_env_0(...))`` + to publish the plan. +* **Breaking:** Removed :attr:`InteractiveScene.env_ns` and + :attr:`InteractiveScene.env_regex_ns` properties (only used internally). Read + :attr:`~isaaclab.cloner.CloneCfg.clone_regex` directly when the env-namespace + convention is needed. diff --git a/source/isaaclab/isaaclab/cloner/__init__.pyi b/source/isaaclab/isaaclab/cloner/__init__.pyi index 1ee123e7cf56..312b356bb595 100644 --- a/source/isaaclab/isaaclab/cloner/__init__.pyi +++ b/source/isaaclab/isaaclab/cloner/__init__.pyi @@ -12,16 +12,30 @@ __all__ = [ "filter_collisions", "grid_transforms", "make_clone_plan", + "ReplicateSession", + "REPLICATION_QUEUE", + "replicate", + "queue_usd_replication", + "UsdReplicateContext", "usd_replicate", ] from .clone_plan import ClonePlan from .cloner_cfg import CloneCfg from .cloner_strategies import random, sequential +from ._fabric_notices import disabled_fabric_change_notifies from .cloner_utils import ( - disabled_fabric_change_notifies, filter_collisions, grid_transforms, make_clone_plan, +) +from .replicate_session import ( + REPLICATION_QUEUE, + ReplicateSession, + replicate, +) +from .usd import ( + UsdReplicateContext, + queue_usd_replication, usd_replicate, ) diff --git a/source/isaaclab/isaaclab/cloner/_fabric_notices.py b/source/isaaclab/isaaclab/cloner/_fabric_notices.py index 4326f4ae3195..0c7c850e9964 100644 --- a/source/isaaclab/isaaclab/cloner/_fabric_notices.py +++ b/source/isaaclab/isaaclab/cloner/_fabric_notices.py @@ -17,9 +17,13 @@ from __future__ import annotations +import contextlib import ctypes import logging import threading +from collections.abc import Iterator + +from pxr import Usd, UsdUtils logger = logging.getLogger(__name__) @@ -149,3 +153,80 @@ def get_bindings() -> FabricNoticeBindings | None: return None _BINDINGS = b return _BINDINGS + + +@contextlib.contextmanager +def disabled_fabric_change_notifies(stage: Usd.Stage, *, restore: bool = True) -> Iterator[None]: + """Suspend the ``IFabricUsd`` USD notice listener for the body of the ``with`` block. + + Targets the same handler that :meth:`isaacsim.core.cloner.Cloner.disable_change_listener` + toggles, but goes through ``omni::fabric::IFabricUsd`` directly so we don't take an + ``isaacsim.core.simulation_manager`` dependency. + + The listener is a global ``TfNotice`` registered when ``omni.fabric`` loads; it + short-circuits via a soft flag (``IFabricUsd.cpp:739``). Toggling that flag is what + skips the per-``Sdf.CopySpec`` Fabric sync that dominates cloning time on large scenes. + + When this provides a measurable speedup + ---------------------------------------- + Bisection on the regression test (see ``test_cloner.py``) shows the listener cost is + only on the critical path when **all** of these hold: + + 1. The clone happens through the ``InteractiveScene`` path with backend physics + replication queued. Calling :func:`usd_replicate` directly on a stage produces + no measurable gap; without physics replication the gap drops to ~1.19x. The + PhysX replication path is what amplifies per-spec listener work. + 2. The cloned prims carry PhysX rigid-body schemas (e.g. ``UsdPhysics.RigidBodyAPI``, + authored via ``rigid_props`` on a spawn cfg). Plain Xforms or geometry without + physics schemas produce ~1.0x — the listener has no Fabric-tracked state to sync. + ``mass_props`` and ``collision_props`` add nothing beyond ``rigid_props``. + 3. Total per-``Sdf.CopySpec`` firings reach ~32K — i.e. ``num_bodies × num_envs`` is + large enough to dominate scene-init cost. Below this the speedup sinks into noise. + + Conditions outside this envelope (no PhysX schemas, single-env scenes, raw + ``usd_replicate`` calls, no physics replication) won't see a perf win — the + suspension is correct but its effect is lost in the rest of the work. + + Re-entrant: if the flag is already off on entry, ``__exit__`` leaves it off. Falls + through to a no-op if the Carbonite interface can't be acquired (e.g. outside a live + Kit application) — the caller never breaks, it just doesn't get the perf win. + + Args: + stage: USD stage whose Fabric notice handler should be suspended. + restore: When ``True`` (default), re-enable the handler on exit. Set to ``False`` + inside a known clone-then-``sim.reset`` window where the downstream Fabric + resync happens anyway and re-enabling here would trigger a redundant + ``forceMinimalPopulate`` batch — see ``PluginInterface.cpp:337``. + + Yields: + None. + """ + bindings = get_bindings() + if bindings is None: + yield + return + + # usdrt only works with a live Kit app — defer import so module load stays cheap. + import usdrt + + # Avoid leaking a strong reference into the global ``StageCache`` for stages we did not + # author into the cache: ``Insert`` keeps the stage alive for the rest of the process. + cache = UsdUtils.StageCache.Get() + cached_id = cache.GetId(stage) + stage_id = cached_id.ToLongInt() if cached_id.IsValid() else cache.Insert(stage).ToLongInt() + # ``FabricId`` wraps a uint64; the C ABI needs the raw integer. + fabric_id = usdrt.Usd.Stage.Attach(stage_id).GetFabricId().id + # First-call ABI sanity check — if the toggle doesn't actually round-trip the flag + # (e.g. Kit's vtable shifted), fall through to a no-op rather than corrupting state. + if not bindings.validate_with(fabric_id): + logger.warning("Fabric notice toggle failed round-trip check — suspension disabled") + yield + return + was_enabled = bindings.is_enabled(fabric_id) + if was_enabled: + bindings.set_enable(fabric_id, False) + try: + yield + finally: + if restore and was_enabled: + bindings.set_enable(fabric_id, True) diff --git a/source/isaaclab/isaaclab/cloner/clone_plan.py b/source/isaaclab/isaaclab/cloner/clone_plan.py index 973122e7744b..8abb00d5d303 100644 --- a/source/isaaclab/isaaclab/cloner/clone_plan.py +++ b/source/isaaclab/isaaclab/cloner/clone_plan.py @@ -12,22 +12,59 @@ @dataclass(frozen=True, eq=False) class ClonePlan: - """Flat cloning source of truth. - - Produced by scene planning after representative source prims are assigned. The - three fields are the same flat replication contract consumed by USD, physics, - and downstream scene-data providers: each source path maps to the destination - template at the same index, and :attr:`clone_mask` selects the environments - populated from that source. - """ + """Description of a single replication layout, consumed by :func:`~isaaclab.cloner.replicate`.""" sources: tuple[str, ...] - """Source prim paths used for replication.""" + """Source prim paths, one per replication row.""" destinations: tuple[str, ...] - """Destination path templates, one per source path.""" + """Destination path templates with ``"{}"`` for the env id, one per row.""" clone_mask: torch.Tensor - """Boolean tensor of shape ``[len(sources), num_envs]``; - ``clone_mask[i, j]`` is ``True`` if env ``j`` was populated from - :attr:`sources` ``[i]``.""" + """Bool tensor ``[len(sources), num_clones]``; ``True`` if env ``j`` comes from row ``i``.""" + + env_ids: torch.Tensor + """Long tensor ``[num_clones]`` of target env ids.""" + + positions: torch.Tensor | None + """Per-env world positions [m], shape ``[num_clones, 3]``, or ``None``.""" + + cfg_rows: dict[int, tuple[int, ...]] + """``id(cfg)`` to the row indices the cfg owns.""" + + @classmethod + def from_env_0( + cls, + source: str, + destination: str, + num_clones: int, + device: str, + positions: torch.Tensor | None = None, + ) -> ClonePlan: + """Build a single-source clone plan that targets every env from one source row. + + 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``. + + Args: + source: Source prim path (typically ``/World/envs/env_0``). + destination: Destination template with ``"{}"`` for the env id. + num_clones: Number of target envs. + device: Torch device for the mask and env id buffers. + positions: Optional per-env world positions [m], shape ``[num_clones, 3]``. + """ + from .replicate_session import REPLICATION_QUEUE # noqa: PLC0415 + + prefix, _, _ = destination.partition("{}") + cfg_rows: dict[int, tuple[int, ...]] = { + id(cfg): (0,) for cfg, _ in REPLICATION_QUEUE if cfg.prim_path.startswith(prefix) + } + return cls( + sources=(source,), + destinations=(destination,), + clone_mask=torch.ones((1, num_clones), dtype=torch.bool, device=device), + env_ids=torch.arange(num_clones, dtype=torch.long, device=device), + positions=positions, + cfg_rows=cfg_rows, + ) diff --git a/source/isaaclab/isaaclab/cloner/cloner_cfg.py b/source/isaaclab/isaaclab/cloner/cloner_cfg.py index 477b4bbfc2c2..464794e366ac 100644 --- a/source/isaaclab/isaaclab/cloner/cloner_cfg.py +++ b/source/isaaclab/isaaclab/cloner/cloner_cfg.py @@ -5,6 +5,8 @@ from __future__ import annotations +from collections.abc import Callable + from isaaclab.utils.configclass import configclass from .cloner_strategies import random @@ -14,32 +16,15 @@ class CloneCfg: """Configuration for environment replication. - The scene builds a :class:`~isaaclab.cloner.ClonePlan` directly from asset - configuration, spawns the representative source prims, and then uses this - configuration to dispatch USD and physics replication for that plan. - """ - - clone_regex: str = "/World/envs/env_.*" - """Destination template for per-environment paths. - - The substring ``".*"`` is replaced with ``"{}"`` internally and formatted with the - environment index (e.g., ``/World/envs/env_0``, ``/World/envs/env_1``). + Holds the knobs :class:`~isaaclab.scene.InteractiveScene` forwards to + :func:`~isaaclab.cloner.make_clone_plan` when building per-env layouts. """ - clone_usd: bool = True - """Enable USD-spec replication to author cloned prims and optional transforms.""" - - clone_physics: bool = True - """Enable PhysX replication for the same mapping to speed up physics setup.""" - - physics_clone_fn: callable | None = None - """Function used to perform physics replication.""" - - clone_strategy: callable = random + clone_strategy: Callable[..., object] = random """Function used to build prototype-to-environment mapping. Default is :func:`random`.""" device: str = "cpu" """Torch device on which mapping buffers are allocated.""" - clone_in_fabric: bool = False - """Enable/disable cloning in fabric for PhysX replication. Default is False.""" + clone_regex: str = "/World/envs/env_.*" + """Regex matching every replicated env prim. Used to expand ``{ENV_REGEX_NS}`` cfg macros.""" diff --git a/source/isaaclab/isaaclab/cloner/cloner_utils.py b/source/isaaclab/isaaclab/cloner/cloner_utils.py index 337fad42f45f..23b493f1bc32 100644 --- a/source/isaaclab/isaaclab/cloner/cloner_utils.py +++ b/source/isaaclab/isaaclab/cloner/cloner_utils.py @@ -5,235 +5,168 @@ from __future__ import annotations -import contextlib import itertools -import logging import math -from collections.abc import Iterator, Sequence +from collections.abc import Callable, Iterable +from typing import TYPE_CHECKING, Any import torch -from pxr import Gf, Sdf, Usd, UsdGeom, UsdUtils, Vt +from pxr import Sdf, Usd, UsdGeom -from . import _fabric_notices from .clone_plan import ClonePlan +from .cloner_strategies import sequential -logger = logging.getLogger(__name__) - - -@contextlib.contextmanager -def disabled_fabric_change_notifies(stage: Usd.Stage, *, restore: bool = True) -> Iterator[None]: - """Suspend the ``IFabricUsd`` USD notice listener for the body of the ``with`` block. - - Targets the same handler that :meth:`isaacsim.core.cloner.Cloner.disable_change_listener` - toggles, but goes through ``omni::fabric::IFabricUsd`` directly so we don't take an - ``isaacsim.core.simulation_manager`` dependency. - - The listener is a global ``TfNotice`` registered when ``omni.fabric`` loads; it - short-circuits via a soft flag (``IFabricUsd.cpp:739``). Toggling that flag is what - skips the per-``Sdf.CopySpec`` Fabric sync that dominates cloning time on large scenes. - - When this provides a measurable speedup - ---------------------------------------- - Bisection on the regression test (see ``test_cloner.py``) shows the listener cost is - only on the critical path when **all** of these hold: - - 1. The clone happens through the ``InteractiveScene`` path with ``replicate_physics=True``. - Calling :func:`usd_replicate` directly on a stage produces no measurable gap; with - ``replicate_physics=False`` the gap drops to ~1.19x. The PhysX replication path is - what amplifies per-spec listener work. - 2. The cloned prims carry PhysX rigid-body schemas (e.g. ``UsdPhysics.RigidBodyAPI``, - authored via ``rigid_props`` on a spawn cfg). Plain Xforms or geometry without - physics schemas produce ~1.0x — the listener has no Fabric-tracked state to sync. - ``mass_props`` and ``collision_props`` add nothing beyond ``rigid_props``. - 3. Total per-``Sdf.CopySpec`` firings reach ~32K — i.e. ``num_bodies × num_envs`` is - large enough to dominate scene-init cost. Below this the speedup sinks into noise. - - Conditions outside this envelope (no PhysX schemas, single-env scenes, raw - ``usd_replicate`` calls, ``replicate_physics=False``) won't see a perf win — the - suspension is correct but its effect is lost in the rest of the work. - - Re-entrant: if the flag is already off on entry, ``__exit__`` leaves it off. Falls - through to a no-op if the Carbonite interface can't be acquired (e.g. outside a live - Kit application) — the caller never breaks, it just doesn't get the perf win. - - Args: - stage: USD stage whose Fabric notice handler should be suspended. - restore: When ``True`` (default), re-enable the handler on exit. Set to ``False`` - inside a known clone-then-``sim.reset`` window where the downstream Fabric - resync happens anyway and re-enabling here would trigger a redundant - ``forceMinimalPopulate`` batch — see ``PluginInterface.cpp:337``. - - Yields: - None. - """ - bindings = _fabric_notices.get_bindings() - if bindings is None: - yield - return - - # usdrt only works with a live Kit app — defer import so module load stays cheap. - import usdrt - - # Avoid leaking a strong reference into the global ``StageCache`` for stages we did not - # author into the cache: ``Insert`` keeps the stage alive for the rest of the process. - cache = UsdUtils.StageCache.Get() - cached_id = cache.GetId(stage) - stage_id = cached_id.ToLongInt() if cached_id.IsValid() else cache.Insert(stage).ToLongInt() - # ``FabricId`` wraps a uint64; the C ABI needs the raw integer. - fabric_id = usdrt.Usd.Stage.Attach(stage_id).GetFabricId().id - # First-call ABI sanity check — if the toggle doesn't actually round-trip the flag - # (e.g. Kit's vtable shifted), fall through to a no-op rather than corrupting state. - if not bindings.validate_with(fabric_id): - logger.warning("Fabric notice toggle failed round-trip check — suspension disabled") - yield - return - was_enabled = bindings.is_enabled(fabric_id) - if was_enabled: - bindings.set_enable(fabric_id, False) - try: - yield - finally: - if restore and was_enabled: - bindings.set_enable(fabric_id, True) +if TYPE_CHECKING: + pass def make_clone_plan( - sources: Sequence[Sequence[str]], - destinations: Sequence[str], + cfgs: Iterable[Any], num_clones: int, - clone_strategy: callable, - device: str = "cpu", + env_spacing: float, + device: str, + *, + clone_strategy: Callable = sequential, + valid_set: torch.Tensor | None = None, ) -> ClonePlan: - """Construct a cloning plan mapping prototype prims to per-environment destinations. + """Build a :class:`~isaaclab.cloner.ClonePlan` from asset cfgs. + + Iterates ``cfgs``, identifies env-scoped cfgs with a spawn, expands + :class:`~isaaclab.sim.MultiAssetSpawnerCfg` / :class:`~isaaclab.sim.MultiUsdFileCfg` + into per-variant prototype rows, runs ``clone_strategy`` to assign prototypes to + envs, and returns a self-contained :class:`ClonePlan` with ``cfg_rows`` populated. - The plan enumerates all combinations of prototypes, selects a combination per environment using ``clone_strategy``, - and builds a boolean masking matrix indicating which prototype populates each environment slot. + Each input cfg's ``spawn_path`` / ``spawn_paths`` is mutated so the subsequent + asset constructor spawns the prototype into its first active environment. Cfgs + whose ``prim_path`` is global (not under the env root ``/World/envs/``) or that + lack a spawn are skipped — they do not appear in the plan and are not replicated. Args: - sources: Prototype prim paths grouped by asset type (e.g., [[robot_a, robot_b], [obj_x]]). - destinations: Destination path templates (one per group) with ``"{}"`` placeholder for env id. - num_clones: Number of environments to populate. - clone_strategy: Function that picks a prototype combo per environment; signature - ``clone_strategy(combos: Tensor, num_clones: int, device: str) -> Tensor[num_clones, num_groups]``. - device: Torch device for tensors in the plan. Defaults to ``"cpu"``. + cfgs: Asset cfgs with resolved ``prim_path`` (no ``{ENV_REGEX_NS}`` macros). + num_clones: Number of target envs. + env_spacing: Distance between neighboring grid env origins [m]. + device: Torch device for plan tensors. + clone_strategy: Function that assigns prototype combinations to envs. Defaults + to :func:`~isaaclab.cloner.sequential`. + valid_set: Optional ``[num_combos, num_groups]`` long tensor of valid prototype + combinations. ``None`` (default) uses the full cartesian product of every + group's prototype indices. Returns: - A :class:`ClonePlan` whose ``sources`` and ``destinations`` are flattened per-source rows and - whose ``clone_mask`` is a ``[num_src, num_clones]`` boolean tensor. + A :class:`ClonePlan` whose ``sources``/``destinations``/``clone_mask`` describe + the flat prototype-to-env mapping and whose ``cfg_rows`` maps each cfg to the + rows it owns. """ - if len(sources) != len(destinations): - raise ValueError(f"Expected one destination per source group, got {len(destinations)} and {len(sources)}.") - if not sources: - raise ValueError("Expected at least one source group.") - group_sizes = [len(group) for group in sources] - if any(size == 0 for size in group_sizes): - raise ValueError("Source groups must not be empty.") - - # 1) Flatten into src and dest lists - src = tuple(p for group in sources for p in group) - dest = tuple(dst for dst, group in zip(destinations, sources) for _ in group) - - # 2) Enumerate all combinations of "one prototype per group" - # all_combos: list of tuples (g0_idx, g1_idx, ..., g_{G-1}_idx) - all_combos = list(itertools.product(*[range(s) for s in group_sizes])) - combos = torch.tensor(all_combos, dtype=torch.long, device=device) - - # 3) Assign a combination to each environment + import isaaclab.sim as sim_utils # noqa: PLC0415 + + def num_variants(spawn_cfg: Any) -> int: + if isinstance(spawn_cfg, sim_utils.MultiAssetSpawnerCfg): + return len(spawn_cfg.assets_cfg) + if isinstance(spawn_cfg, sim_utils.MultiUsdFileCfg): + return 1 if isinstance(spawn_cfg.usd_path, str) else len(spawn_cfg.usd_path) + return 1 + + def set_spawn_paths(spawn_cfg: Any, paths: list[str | None]) -> None: + if isinstance(spawn_cfg, (sim_utils.MultiAssetSpawnerCfg, sim_utils.MultiUsdFileCfg)): + spawn_cfg.spawn_paths = paths + else: + active = [p for p in paths if p is not None] + if len(active) != 1: + raise ValueError("Single spawner expects exactly one planned source path.") + spawn_cfg.spawn_path = active[0] + + env_root_marker = "/World/envs/" + env_template = "/World/envs/env_{}" + + # 1) Build per-group records: (cfg, spawn_cfg, destination_template, num_variants). + groups: list[tuple[Any, Any, str, int]] = [] + for cfg in cfgs: + if not hasattr(cfg, "prim_path") or not hasattr(cfg, "spawn") or cfg.spawn is None: + continue + prim_path = cfg.prim_path + if env_root_marker not in prim_path: + continue + count = num_variants(cfg.spawn) + if count <= 0: + raise ValueError(f"Spawner at '{prim_path}' must have at least one variant.") + destination = prim_path.replace(".*", "{}") + groups.append((cfg, cfg.spawn, destination, count)) + + env_ids = torch.arange(num_clones, dtype=torch.long, device=device) + positions, _ = grid_transforms(num_clones, env_spacing, device=device) + + # 2) No env-scoped cfgs: emit an empty plan so the scene can still proceed. + if not groups: + empty_mask = torch.zeros((0, num_clones), dtype=torch.bool, device=device) + return ClonePlan( + sources=(), + destinations=(), + clone_mask=empty_mask, + env_ids=env_ids, + positions=positions, + cfg_rows={}, + ) + + # 3) Homogeneous (every cfg is single-variant): emit the simpler env-root plan. + if all(count == 1 for _, _, _, count in groups): + for cfg, spawn_cfg, destination, _ in groups: + set_spawn_paths(spawn_cfg, [destination.format(0)]) + cfg_rows = {id(cfg): (0,) for cfg, _, _, _ in groups} + return ClonePlan( + sources=(env_template.format(0),), + destinations=(env_template,), + clone_mask=torch.ones((1, num_clones), dtype=torch.bool, device=device), + env_ids=env_ids, + positions=positions, + cfg_rows=cfg_rows, + ) + + # 4) Heterogeneous: enumerate prototype combos, build per-row mask, mutate spawn paths. + group_sizes = [count for _, _, _, count in groups] + if valid_set is None: + all_combos = list(itertools.product(*[range(s) for s in group_sizes])) + combos = torch.tensor(all_combos, dtype=torch.long, device=device) + else: + combos = valid_set.to(device=device, dtype=torch.long) chosen = clone_strategy(combos, num_clones, device) - # 4) Build masking: [num_src, num_clones] boolean - # For each env, for each group, mark exactly one prototype row as True. group_offsets = torch.tensor([0] + list(itertools.accumulate(group_sizes[:-1])), dtype=torch.long, device=device) rows = (chosen + group_offsets).view(-1) cols = torch.arange(num_clones, device=device).view(-1, 1).expand(-1, len(group_sizes)).reshape(-1) - masking = torch.zeros((sum(group_sizes), num_clones), dtype=torch.bool, device=device) - masking[rows, cols] = True - return ClonePlan(sources=src, destinations=dest, clone_mask=masking) - - -def usd_replicate( - stage: Usd.Stage, - sources: Sequence[str], - destinations: Sequence[str], - env_ids: torch.Tensor, - mask: torch.Tensor | None = None, - positions: torch.Tensor | None = None, - quaternions: torch.Tensor | None = None, -) -> None: - """Replicate USD prims to per-environment destinations. - - Copies each source prim spec to destination templates for selected environments - (``mask``). Optionally authors translate/orient from position/quaternion buffers. - Replication runs in path-depth order (parents before children) for robust composition. - - Args: - stage: USD stage. - sources: Source prim paths. - destinations: Destination formattable templates with ``"{}"`` for env index. - env_ids: Environment indices. - mask: Optional per-source or shared mask. ``None`` selects all. - positions: Optional positions (``[E, 3]``) -> ``xformOp:translate``. - quaternions: Optional orientations (``[E, 4]``) in ``xyzw`` -> ``xformOp:orient``. - - """ - rl = stage.GetRootLayer() - - # Group replication by destination path depth so ancestors land before deeper paths. - # This avoids composition issues for nested or interdependent specs. - def dp_depth(template: str) -> int: - """Return destination prim path depth for stable parent-first replication.""" - dp = template.format(0) - return Sdf.Path(dp).pathElementCount - - order = sorted(range(len(sources)), key=lambda i: dp_depth(destinations[i])) - - # Process in layers of equal depth, committing at each depth to stabilize composition - depth_to_indices: dict[int, list[int]] = {} - for i in order: - d = dp_depth(destinations[i]) - depth_to_indices.setdefault(d, []).append(i) - - for depth in sorted(depth_to_indices.keys()): - with Sdf.ChangeBlock(): - for i in depth_to_indices[depth]: - src = sources[i] - tmpl = destinations[i] - # Select target environments for this source (supports None, [E], or [S, E]) - target_envs = env_ids if mask is None else env_ids[mask[i]] - for wid in target_envs.tolist(): - dp = tmpl.format(wid) - Sdf.CreatePrimInLayer(rl, dp) - if src == dp: - pass # self-copy: CreatePrimInLayer already ensures it exists; CopySpec would be destructive - else: - Sdf.CopySpec(rl, Sdf.Path(src), rl, Sdf.Path(dp)) - - if positions is not None or quaternions is not None: - ps = rl.GetPrimAtPath(dp) - op_names = [] - if positions is not None: - p = positions[wid] - t_attr = ps.GetAttributeAtPath(dp + ".xformOp:translate") - if t_attr is None: - t_attr = Sdf.AttributeSpec(ps, "xformOp:translate", Sdf.ValueTypeNames.Double3) - t_attr.default = Gf.Vec3d(float(p[0]), float(p[1]), float(p[2])) - op_names.append("xformOp:translate") - if quaternions is not None: - q = quaternions[wid] - o_attr = ps.GetAttributeAtPath(dp + ".xformOp:orient") - if o_attr is None: - o_attr = Sdf.AttributeSpec(ps, "xformOp:orient", Sdf.ValueTypeNames.Quatd) - # xyzw convention: q[3] is w, q[0:3] is xyz - o_attr.default = Gf.Quatd(float(q[3]), Gf.Vec3d(float(q[0]), float(q[1]), float(q[2]))) - op_names.append("xformOp:orient") - # Only author xformOpOrder for the ops we actually authored - if op_names: - op_order = ps.GetAttributeAtPath(dp + ".xformOpOrder") or Sdf.AttributeSpec( - ps, UsdGeom.Tokens.xformOpOrder, Sdf.ValueTypeNames.TokenArray - ) - op_order.default = Vt.TokenArray(op_names) + num_rows = sum(group_sizes) + clone_mask = torch.zeros((num_rows, num_clones), dtype=torch.bool, device=device) + clone_mask[rows, cols] = True + + sources_list: list[str] = [] + destinations_list: list[str] = [] + cfg_rows: dict[int, tuple[int, ...]] = {} + row = 0 + for cfg, spawn_cfg, destination, count in groups: + cfg_rows[id(cfg)] = tuple(range(row, row + count)) + group_mask = clone_mask[row : row + count] + env_ids_assigned = group_mask.to(torch.int).argmax(dim=1).tolist() + active = group_mask.any(dim=1).tolist() + paths = [ + destination.format(env_id) if is_active else None for env_id, is_active in zip(env_ids_assigned, active) + ] + for i, path in enumerate(paths): + destinations_list.append(destination) + # Inactive prototypes fall back to env-i so the source path stays valid even + # when the variant has no active environment (matches the legacy behavior). + sources_list.append(path if path is not None else destination.format(i)) + set_spawn_paths(spawn_cfg, paths) + row += count + + return ClonePlan( + sources=tuple(sources_list), + destinations=tuple(destinations_list), + clone_mask=clone_mask, + env_ids=env_ids, + positions=positions, + cfg_rows=cfg_rows, + ) def filter_collisions( diff --git a/source/isaaclab/isaaclab/cloner/replicate_session.py b/source/isaaclab/isaaclab/cloner/replicate_session.py new file mode 100644 index 000000000000..5341867e1a24 --- /dev/null +++ b/source/isaaclab/isaaclab/cloner/replicate_session.py @@ -0,0 +1,132 @@ +# 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 + +"""Replication queue, :func:`replicate` drain, and :class:`ReplicateSession` sugar.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable +from typing import TYPE_CHECKING, Any + +from .cloner_strategies import sequential + +if TYPE_CHECKING: + import torch + from pxr import Usd + + from .clone_plan import ClonePlan + + +REPLICATION_QUEUE: list[tuple[Any, type]] = [] +"""``(cfg, BackendCtxCls)`` pairs appended by ``queue__replication`` and drained by :func:`replicate`.""" + + +def replicate(plan: "ClonePlan", *, stage: "Usd.Stage") -> None: + """Drain :data:`REPLICATION_QUEUE` against ``plan``, dispatch each backend, publish the plan. + + 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. + """ + from isaaclab.sim import SimulationContext # noqa: PLC0415 + + queued = list(REPLICATION_QUEUE) + REPLICATION_QUEUE.clear() + + backend_ctxs: dict[type, Any] = {} + for cfg, BackendCtxCls in queued: + rows = plan.cfg_rows.get(id(cfg)) + if rows is None: + continue + ctx = backend_ctxs.get(BackendCtxCls) + if ctx is None: + ctx = BackendCtxCls(stage) + backend_ctxs[BackendCtxCls] = ctx + row_list = list(rows) + ctx.queue_mapping( + [plan.sources[i] for i in row_list], + [plan.destinations[i] for i in row_list], + plan.env_ids, + plan.clone_mask[row_list], + positions=plan.positions, + ) + + for ctx in sorted(backend_ctxs.values(), key=lambda c: getattr(c, "replicate_priority", 0)): + ctx.replicate() + + SimulationContext.instance().set_clone_plan(plan) + + +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 + :data:`REPLICATION_QUEUE`; ``__exit__`` drains and dispatches. + + Example: + + .. code-block:: python + + with cloner.ReplicateSession(cfgs, num_clones=128, env_spacing=2.0, + device="cuda:0", stage=sim.stage): + for cfg in cfgs: + cfg.class_type(cfg) + """ + + def __init__( + self, + cfgs: Iterable[Any], + num_clones: int, + env_spacing: float, + device: str, + *, + stage: "Usd.Stage", + clone_strategy: Callable = sequential, + valid_set: "torch.Tensor | None" = None, + ): + """Capture arguments for :func:`make_clone_plan` and :func:`replicate`. + + Args: + cfgs: Asset cfgs with resolved ``prim_path``. + num_clones: Number of target envs. + env_spacing: Grid spacing between env origins [m]. + device: Torch device for plan tensors. + stage: USD stage to author replicated prim specs into. + 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. + """ + self._cfgs = cfgs + self._stage = stage + self._kwargs = dict( + num_clones=num_clones, + env_spacing=env_spacing, + device=device, + clone_strategy=clone_strategy, + valid_set=valid_set, + ) + self._plan: "ClonePlan | None" = None + + def __enter__(self) -> "ReplicateSession": + from .cloner_utils import make_clone_plan # noqa: PLC0415 + + self._plan = make_clone_plan(self._cfgs, **self._kwargs) + return self + + 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) + else: + # Drop cfgs registered before the failure so the next session is clean. + REPLICATION_QUEUE.clear() + + @property + def plan(self) -> "ClonePlan": + """The :class:`~isaaclab.cloner.ClonePlan` produced in :meth:`__enter__`.""" + if self._plan is None: + raise RuntimeError("ReplicateSession.plan is only available inside the with block.") + return self._plan diff --git a/source/isaaclab/isaaclab/cloner/usd.py b/source/isaaclab/isaaclab/cloner/usd.py new file mode 100644 index 000000000000..070a43a2d4c8 --- /dev/null +++ b/source/isaaclab/isaaclab/cloner/usd.py @@ -0,0 +1,182 @@ +# 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 Any + +import torch + +from pxr import Gf, Sdf, Usd, UsdGeom, Vt + +from .replicate_session import REPLICATION_QUEUE + + +def _select_env_ids(env_ids: torch.Tensor, mask: torch.Tensor | None, row: int) -> torch.Tensor: + """Return the environment ids selected by a replication row.""" + if mask is None: + return env_ids + row_mask = mask if mask.dim() == 1 else mask[row] + if row_mask.dtype != torch.bool: + row_mask = row_mask.to(dtype=torch.bool) + return env_ids[row_mask] + + +class UsdReplicateContext: + """Queue and apply USD replication work for one stage.""" + + replicate_priority = 100 + + def __init__(self, stage: Usd.Stage): + """Initialize the context. + + Args: + stage: USD stage to author replicated prim specs into. + """ + self.stage = stage + self._queue: list[tuple[str, str, torch.Tensor, torch.Tensor | None, torch.Tensor | None]] = [] + + def queue( + self, + source: str, + destination: str, + env_ids: torch.Tensor, + *, + positions: torch.Tensor | None = None, + quaternions: torch.Tensor | None = None, + ) -> None: + """Queue one USD source row for replication. + + Args: + source: Source prim path. + destination: Destination path template with ``"{}"`` for env id. + env_ids: Environment ids selected for this source row. + positions: Optional per-environment world positions [m]. + quaternions: Optional per-environment orientations in xyzw order. + """ + self._queue.append((source, destination, env_ids, positions, quaternions)) + + def queue_mapping( + self, + sources: Sequence[str], + destinations: Sequence[str], + env_ids: torch.Tensor, + mask: torch.Tensor | None = None, + *, + 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. + destinations: Destination path templates with ``"{}"`` for env id. + env_ids: Environment indices. + mask: Optional per-source or shared mask. + positions: Optional per-environment world positions [m]. + quaternions: Optional per-environment orientations in xyzw order. + """ + for i, source in enumerate(sources): + self.queue( + source, + destinations[i], + _select_env_ids(env_ids, mask, i), + positions=positions, + quaternions=quaternions, + ) + + def replicate(self) -> None: + """Apply all queued USD copy specs in parent-before-child order.""" + if not self._queue: + return + + try: + rl = self.stage.GetRootLayer() + + def dp_depth(template: str) -> int: + """Return destination prim path depth for stable parent-first replication.""" + dp = template.format(0) + return Sdf.Path(dp).pathElementCount + + depth_to_items: dict[int, list[tuple[str, str, torch.Tensor, torch.Tensor | None, torch.Tensor | None]]] = {} + for item in self._queue: + depth_to_items.setdefault(dp_depth(item[1]), []).append(item) + + for depth in sorted(depth_to_items.keys()): + with Sdf.ChangeBlock(): + for src, tmpl, target_envs, positions, quaternions in depth_to_items[depth]: + for wid in target_envs.tolist(): + wid = int(wid) + dp = tmpl.format(wid) + Sdf.CreatePrimInLayer(rl, dp) + if src != dp: + Sdf.CopySpec(rl, Sdf.Path(src), rl, Sdf.Path(dp)) + + if positions is not None or quaternions is not None: + ps = rl.GetPrimAtPath(dp) + op_names = [] + if positions is not None: + p = positions[wid] + t_attr = ps.GetAttributeAtPath(dp + ".xformOp:translate") + if t_attr is None: + t_attr = Sdf.AttributeSpec(ps, "xformOp:translate", Sdf.ValueTypeNames.Double3) + t_attr.default = Gf.Vec3d(float(p[0]), float(p[1]), float(p[2])) + op_names.append("xformOp:translate") + if quaternions is not None: + q = quaternions[wid] + o_attr = ps.GetAttributeAtPath(dp + ".xformOp:orient") + if o_attr is None: + o_attr = Sdf.AttributeSpec(ps, "xformOp:orient", Sdf.ValueTypeNames.Quatd) + o_attr.default = Gf.Quatd( + float(q[3]), Gf.Vec3d(float(q[0]), float(q[1]), float(q[2])) + ) + op_names.append("xformOp:orient") + if op_names: + op_order = ps.GetAttributeAtPath(dp + ".xformOpOrder") or Sdf.AttributeSpec( + ps, UsdGeom.Tokens.xformOpOrder, Sdf.ValueTypeNames.TokenArray + ) + op_order.default = Vt.TokenArray(op_names) + finally: + self._queue.clear() + + +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], + destinations: Sequence[str], + env_ids: torch.Tensor, + mask: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + quaternions: torch.Tensor | None = None, +) -> None: + """Replicate USD prims to per-environment destinations. + + Copies each source prim spec to destination templates for selected environments + (``mask``). Optionally authors translate/orient from position/quaternion buffers. + Replication runs in path-depth order (parents before children) for robust composition. + + Args: + stage: USD stage. + sources: Source prim paths. + destinations: Destination formattable templates with ``"{}"`` for env index. + env_ids: Environment indices. + mask: Optional per-source or shared mask. ``None`` selects all. + positions: Optional positions [m], shape ``[E, 3]``, authored as ``xformOp:translate``. + quaternions: Optional orientations in xyzw order, shape ``[E, 4]``, authored as ``xformOp:orient``. + """ + ctx = UsdReplicateContext(stage) + ctx.queue_mapping(sources, destinations, env_ids, mask, positions=positions, quaternions=quaternions) + ctx.replicate() diff --git a/source/isaaclab/isaaclab/scene/interactive_scene.py b/source/isaaclab/isaaclab/scene/interactive_scene.py index 1d1afa2cfad9..368ca6ddc75c 100644 --- a/source/isaaclab/isaaclab/scene/interactive_scene.py +++ b/source/isaaclab/isaaclab/scene/interactive_scene.py @@ -17,8 +17,6 @@ import torch import warp as wp -from pxr import Sdf - import isaaclab.sim as sim_utils from isaaclab import cloner from isaaclab.assets import ( @@ -103,12 +101,20 @@ class MySceneCfg(InteractiveSceneCfg): robot = scene.articulations["robot"] If the :class:`InteractiveSceneCfg` class does not include asset entities, the cloning process - can still be triggered if assets were added to the stage outside of the :class:`InteractiveScene` class: + can still be triggered by constructing assets directly on the stage and then calling + :func:`isaaclab.cloner.replicate` with a single-source :class:`~isaaclab.cloner.ClonePlan`: .. code-block:: python + from isaaclab import cloner + from isaaclab.assets import Articulation + scene = InteractiveScene(cfg=InteractiveSceneCfg(num_envs=128, replicate_physics=True)) - scene.clone_environments() + robot = Articulation(robot_cfg) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(scene.num_envs, scene.cfg.env_spacing, device=scene.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, scene.num_envs, scene.device, pos) + cloner.replicate(plan, stage=scene.stage) .. note:: It is important to note that the scene only performs common operations on the entities. For example, @@ -144,209 +150,67 @@ def __init__(self, cfg: InteractiveSceneCfg): self.stage_id = get_current_stage_id() self.physics_backend = self.sim.physics_manager.__name__.lower() requested_viz_types = set(self.sim.resolve_visualizer_types()) - if self.physics_backend.startswith("ovphysx"): - from isaaclab_ovphysx.cloner import ovphysx_replicate - - physics_clone_fn = ovphysx_replicate - elif self.physics_backend.startswith("physx"): - from isaaclab_physx.cloner import physx_replicate - - physics_clone_fn = physx_replicate - elif self.physics_backend.startswith("newton"): - from isaaclab_newton.cloner import newton_physics_replicate - - physics_clone_fn = newton_physics_replicate - else: - raise ValueError(f"Unsupported physics backend: {self.physics_backend}") # physics scene path self._physics_scene_path = None # prepare cloner for environment replication - self.env_prim_paths = [f"{self.env_ns}/env_{i}" for i in range(self.cfg.num_envs)] - - self.cloner_cfg = cloner.CloneCfg( - clone_regex=self.env_regex_ns, - clone_in_fabric=self.cfg.clone_in_fabric, - device=self.device, - physics_clone_fn=physics_clone_fn, - # USD replication runs for every backend. PhysX/Newton need per-env - # USD prims for sensor discovery. For OVPhysX, the per-env USD - # subtrees are layered on TOP of the physics-side ``physx.clone()`` - # replicas -- PhysX is indifferent to additional USD content and - # the two layers don't conflict. Probing whether this assumption - # holds in practice; revert to ``not startswith("ovphysx")`` if - # ``physx.clone()`` errors on already-populated targets. - clone_usd=True, - ) + self.cloner_cfg = cloner.CloneCfg(device=self.device) + 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)] # create source prim self.stage.DefinePrim(self.env_prim_paths[0], "Xform") - self.env_fmt = self.env_regex_ns.replace(".*", "{}") # allocate env indices self._ALL_INDICES = torch.arange(self.cfg.num_envs, dtype=torch.long, device=self.device) - self._default_env_origins, _ = cloner.grid_transforms(self.num_envs, self.cfg.env_spacing, device=self.device) - # copy empty prim of env_0 to env_1, env_2, ..., env_{num_envs-1} with correct location. - # Suspend Fabric's USD notice listener: scene-init is followed by ``SimulationContext.reset``, - # which does the Fabric resync naturally — re-enabling here would just trigger a redundant batch. - # Note: ``restore=False`` means the listener stays disabled past this ``with`` block — through - # ``_add_entities_from_cfg`` and ``clone_environments`` below — until ``SimulationContext.reset`` - # re-enables it. The nested suspension inside ``clone_environments`` becomes a no-op as a result. - with cloner.disabled_fabric_change_notifies(self.stage, restore=False): - cloner.usd_replicate( - self.stage, - [self.env_fmt.format(0)], - [self.env_fmt], - self._ALL_INDICES, - positions=self._default_env_origins, - ) + # clone env_0 xform to env_1..env_{N-1} at grid origins + env_origins, _ = cloner.grid_transforms(self.num_envs, self.cfg.env_spacing, device=self.device) + cloner.usd_replicate( + self.stage, + ["/World/envs/env_0"], + ["/World/envs/env_{}"], + self._ALL_INDICES, + positions=env_origins, + ) + # Always enter so a ClonePlan is published even when the scene cfg has no entities. self._global_prim_paths = list() - has_scene_cfg_entities = self._is_scene_setup_from_cfg() - if has_scene_cfg_entities: - self._clone_plan = self._build_clone_plan_from_cfg() - self._add_entities_from_cfg() - else: - self._clone_plan = cloner.ClonePlan( - sources=(self.env_fmt.format(0),), - destinations=(self.env_fmt,), - clone_mask=torch.ones((1, self.num_envs), device=self.device, dtype=torch.bool), - ) + with cloner.ReplicateSession( + self._collect_asset_cfgs(), + num_clones=self.num_envs, + env_spacing=self.cfg.env_spacing, + device=self.device, + stage=self.stage, + clone_strategy=self.cloner_cfg.clone_strategy, + ): + if self._is_scene_setup_from_cfg(): + self._add_entities_from_cfg() - # Aggregate scene-data requirements from declared visualizers and constructed sensors, - # then publish to ``SimulationContext`` so downstream providers (constructed later by - # :meth:`SimulationContext.initialize_visualizers`) see the full picture in one read. self._aggregate_scene_data_requirements(requested_viz_types) - if has_scene_cfg_entities: - self.clone_environments(copy_from_source=(not self.cfg.replicate_physics)) - # Collision filtering is PhysX-specific (PhysxSchema.PhysxSceneAPI) - # Intentionally matches both physx and ovphysx (both are PhysX-based) - if self.cfg.filter_collisions and "physx" in self.physics_backend: - self.filter_collisions(self._global_prim_paths) + # 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(): + self.filter_collisions(self._global_prim_paths) - def _build_clone_plan_from_cfg(self) -> cloner.ClonePlan | None: - """Build a clone plan from scene cfg spawn variants and write planned spawn paths. + def _collect_asset_cfgs(self) -> list[Any]: + """Flatten user-declared cfgs for :func:`~isaaclab.cloner.make_clone_plan`. - Returns ``None`` when the cfg has no env-scoped spawned assets. + Expands :class:`~isaaclab.assets.RigidObjectCollectionCfg` into its members, + resolves ``{ENV_REGEX_NS}`` macros, and orders sensors after non-sensors. """ - - def num_variants(spawn_cfg) -> int: - if isinstance(spawn_cfg, sim_utils.MultiAssetSpawnerCfg): - return len(spawn_cfg.assets_cfg) - if isinstance(spawn_cfg, sim_utils.MultiUsdFileCfg): - return 1 if isinstance(spawn_cfg.usd_path, str) else len(spawn_cfg.usd_path) - return 1 - - def set_spawn_paths(spawn_cfg, paths: list[str | None]) -> None: - if isinstance(spawn_cfg, (sim_utils.MultiAssetSpawnerCfg, sim_utils.MultiUsdFileCfg)): - spawn_cfg.spawn_paths = paths - else: - active = [path for path in paths if path is not None] - if len(active) != 1: - raise ValueError("Single spawner expects exactly one planned source path.") - spawn_cfg.spawn_path = active[0] - cfg_fields = InteractiveSceneCfg.__dataclass_fields__ items = [(k, v) for k, v in self.cfg.__dict__.items() if k not in cfg_fields and v is not None] - ordered_items = [item for item in items if not isinstance(item[1], SensorBaseCfg)] - ordered_items += [item for item in items if isinstance(item[1], SensorBaseCfg)] - - # One group is one prim path template plus its spawn variants. - groups = [] - for _, asset_cfg in ordered_items: - cfgs = asset_cfg.rigid_objects.values() if isinstance(asset_cfg, RigidObjectCollectionCfg) else [asset_cfg] - for cfg in (cfg for cfg in cfgs if hasattr(cfg, "prim_path")): - prim_path = cfg.prim_path.format(ENV_REGEX_NS=self.env_regex_ns) - if not hasattr(cfg, "spawn") or cfg.spawn is None or self.env_ns not in prim_path: - continue - if (count := num_variants(cfg.spawn)) <= 0: - raise ValueError(f"Spawner at '{prim_path}' must have at least one variant.") - groups.append((cfg.spawn, prim_path.replace(self.env_regex_ns, self.env_fmt), count)) - - if not groups: - return None - - # Homogeneous scenes still spawn sources at env_0, but publish the simpler env-root plan. - if all(count == 1 for _, _, count in groups): - for spawn_cfg, destination, _ in groups: - set_spawn_paths(spawn_cfg, [destination.format(0)]) - return cloner.ClonePlan( - sources=(self.env_fmt.format(0),), - destinations=(self.env_fmt,), - clone_mask=torch.ones((1, self.num_envs), device=self.device, dtype=torch.bool), - ) - - plan = cloner.make_clone_plan( - [[destination.format(i) for i in range(count)] for _, destination, count in groups], - [destination for _, destination, _ in groups], - self.num_envs, - self.cloner_cfg.clone_strategy, - self.device, - ) - - # Move each planned source row to the first environment that actually uses it. - row = 0 - sources = list(plan.sources) - for spawn_cfg, destination, count in groups: - mask = plan.clone_mask[row : row + count] - env_ids = mask.to(torch.int).argmax(dim=1).tolist() - active = mask.any(dim=1).tolist() - paths = [destination.format(env_id) if is_active else None for env_id, is_active in zip(env_ids, active)] - for i, path in zip(range(row, row + count), paths): - if path is not None: - sources[i] = path - set_spawn_paths(spawn_cfg, paths) - row += count - - plan = cloner.ClonePlan(sources=tuple(sources), destinations=plan.destinations, clone_mask=plan.clone_mask) - logger.debug("Built heterogeneous ClonePlan with %d source rows.", len(plan.sources)) - return plan - - def clone_environments(self, copy_from_source: bool = False): - """Creates clones of the environment ``/World/envs/env_0``. - - Args: - copy_from_source: (bool): If set to False, clones inherit from /World/envs/env_0 and mirror its changes. - If True, clones are independent copies of the source prim and won't reflect its changes (start-up time - may increase). Defaults to False. - """ - plan = self._clone_plan - assert self.sim is not None - if plan is None: - self.sim.set_clone_plan(None) - return + ordered_items = [v for _, v in items if not isinstance(v, SensorBaseCfg)] + ordered_items += [v for _, v in items if isinstance(v, SensorBaseCfg)] - # PhysX-only: set env id bit count for replicated physics. Newton handles env separation in its own API. - # Intentionally matches both physx and ovphysx (both are PhysX-based) - if self.cfg.replicate_physics and "physx" in self.physics_backend: - prim = self.stage.GetPrimAtPath("/physicsScene") - prim.CreateAttribute("physxScene:envIdInBoundsBitCount", Sdf.ValueTypeNames.Int).Set(4) - - # Suspend Fabric's USD notice listener around bulk authoring. ``restore=False`` because the downstream - # ``SimulationContext.reset`` does the Fabric resync — re-enabling here would batch-resync everything - # we just authored, which is slower than the unsuppressed baseline. - with cloner.disabled_fabric_change_notifies(self.stage, restore=False): - replicate_args = (plan.sources, plan.destinations, self._ALL_INDICES, plan.clone_mask) - - if not copy_from_source and self.cloner_cfg.physics_clone_fn is not None: - self.cloner_cfg.physics_clone_fn( - self.stage, - *replicate_args, - positions=self._default_env_origins, - device=self.cloner_cfg.device, - ) - if self.cloner_cfg.clone_usd: - is_env_root_plan = ( - len(plan.sources) == 1 - and plan.sources[0] == self.env_fmt.format(0) - and plan.destinations == (self.env_fmt,) - ) - usd_positions = self._default_env_origins if is_env_root_plan else None - cloner.usd_replicate(self.stage, *replicate_args, positions=usd_positions) - - # Publish to ``SimulationContext`` (the canonical owner). The :attr:`clone_plan` - # property below forwards reads back through ``sim.get_clone_plan()`` so consumers - # holding a scene reference still see the published plan without a duplicate cache. - self.sim.set_clone_plan(plan) + cfgs: list[Any] = [] + for asset_cfg in ordered_items: + children = ( + asset_cfg.rigid_objects.values() if isinstance(asset_cfg, RigidObjectCollectionCfg) else [asset_cfg] + ) + for child in children: + if hasattr(child, "prim_path"): + child.prim_path = child.prim_path.format(ENV_REGEX_NS=self.cloner_cfg.clone_regex) + cfgs.append(child) + return cfgs def _aggregate_scene_data_requirements(self, visualizer_types=()) -> None: """Aggregate scene-data requirements from visualizers and sensor renderers. @@ -489,20 +353,6 @@ def device(self) -> str: """The device on which the scene is created.""" return sim_utils.SimulationContext.instance().device # pyright: ignore [reportOptionalMemberAccess] - @property - def env_ns(self) -> str: - """The namespace ``/World/envs`` in which all environments created. - - The environments are present w.r.t. this namespace under "env_{N}" prim, - where N is a natural number. - """ - return "/World/envs" - - @property - def env_regex_ns(self) -> str: - """The namespace ``/World/envs/env_.*`` in which all environments created.""" - return f"{self.env_ns}/env_.*" - @property def num_envs(self) -> int: """The number of environments handled by the scene.""" @@ -510,11 +360,10 @@ def num_envs(self) -> int: @property def env_origins(self) -> torch.Tensor: - """The origins of the environments in the scene. Shape is (num_envs, 3).""" + """Per-env world origins, shape ``(num_envs, 3)``. From the terrain when registered, else from the published :class:`~isaaclab.cloner.ClonePlan`.""" if self._terrain is not None: return self._terrain.env_origins - else: - return self._default_env_origins + return self.sim.get_clone_plan().positions @property def terrain(self) -> TerrainImporter | None: @@ -558,11 +407,11 @@ def surface_grippers(self) -> dict[str, SurfaceGripper]: @property def clone_plan(self) -> cloner.ClonePlan | None: - """Clone plan produced by :meth:`clone_environments`. + """Clone plan produced by the most recent replication. 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 :meth:`clone_environments` runs. + assignment mask. ``None`` until :func:`isaaclab.cloner.replicate` has run. """ return self.sim.get_clone_plan() @@ -888,6 +737,9 @@ def _add_entities_from_cfg(self): # noqa: C901 # store paths that are in global collision filter self._global_prim_paths = list() + # Resolve the env-namespace convention from the cloner cfg once for this pass. + env_regex_ns = self.cloner_cfg.clone_regex + env_root = env_regex_ns.rsplit("/", 1)[0] # Process non-sensor entities before sensors so that asset prims exist in the template # when sensors (e.g. cameras attached to robot links) need to spawn under them. all_items = [ @@ -902,13 +754,13 @@ def _add_entities_from_cfg(self): # noqa: C901 for asset_name, asset_cfg in ordered_items: # resolve prim_path with env regex if hasattr(asset_cfg, "prim_path"): - asset_cfg.prim_path = asset_cfg.prim_path.format(ENV_REGEX_NS=self.env_regex_ns) + asset_cfg.prim_path = asset_cfg.prim_path.format(ENV_REGEX_NS=env_regex_ns) # set spawn_path on spawner if cloning is needed if hasattr(asset_cfg, "spawn") and asset_cfg.spawn is not None: is_multi_spawner = isinstance( asset_cfg.spawn, (sim_utils.MultiAssetSpawnerCfg, sim_utils.MultiUsdFileCfg) ) - if self.env_ns not in asset_cfg.prim_path: + if env_root not in asset_cfg.prim_path: asset_cfg.spawn.spawn_path = asset_cfg.prim_path elif is_multi_spawner and not asset_cfg.spawn.spawn_paths: raise RuntimeError(f"Clone planning did not assign spawn_paths for '{asset_cfg.prim_path}'.") @@ -928,13 +780,13 @@ def _add_entities_from_cfg(self): # noqa: C901 self._rigid_objects[asset_name] = asset_cfg.class_type(asset_cfg) elif isinstance(asset_cfg, RigidObjectCollectionCfg): for rigid_object_cfg in asset_cfg.rigid_objects.values(): - rigid_object_cfg.prim_path = rigid_object_cfg.prim_path.format(ENV_REGEX_NS=self.env_regex_ns) + rigid_object_cfg.prim_path = rigid_object_cfg.prim_path.format(ENV_REGEX_NS=env_regex_ns) # set spawn_path on spawner if cloning is needed if hasattr(rigid_object_cfg, "spawn") and rigid_object_cfg.spawn is not None: is_multi_spawner = isinstance( rigid_object_cfg.spawn, (sim_utils.MultiAssetSpawnerCfg, sim_utils.MultiUsdFileCfg) ) - if self.env_ns not in rigid_object_cfg.prim_path: + if env_root not in rigid_object_cfg.prim_path: rigid_object_cfg.spawn.spawn_path = rigid_object_cfg.prim_path elif is_multi_spawner and not rigid_object_cfg.spawn.spawn_paths: raise RuntimeError( @@ -957,32 +809,32 @@ def _add_entities_from_cfg(self): # noqa: C901 if isinstance(asset_cfg, FrameTransformerCfg): updated_target_frames = [] for target_frame in asset_cfg.target_frames: - target_frame.prim_path = target_frame.prim_path.format(ENV_REGEX_NS=self.env_regex_ns) + target_frame.prim_path = target_frame.prim_path.format(ENV_REGEX_NS=env_regex_ns) updated_target_frames.append(target_frame) asset_cfg.target_frames = updated_target_frames elif isinstance(asset_cfg, ContactSensorCfg): asset_cfg.filter_prim_paths_expr = [ - p.format(ENV_REGEX_NS=self.env_regex_ns) for p in asset_cfg.filter_prim_paths_expr + p.format(ENV_REGEX_NS=env_regex_ns) for p in asset_cfg.filter_prim_paths_expr ] if hasattr(asset_cfg, "sensor_shape_prim_expr") and asset_cfg.sensor_shape_prim_expr: asset_cfg.sensor_shape_prim_expr = [ - p.format(ENV_REGEX_NS=self.env_regex_ns) for p in asset_cfg.sensor_shape_prim_expr + p.format(ENV_REGEX_NS=env_regex_ns) for p in asset_cfg.sensor_shape_prim_expr ] if hasattr(asset_cfg, "filter_shape_prim_expr") and asset_cfg.filter_shape_prim_expr: asset_cfg.filter_shape_prim_expr = [ - p.format(ENV_REGEX_NS=self.env_regex_ns) for p in asset_cfg.filter_shape_prim_expr + p.format(ENV_REGEX_NS=env_regex_ns) for p in asset_cfg.filter_shape_prim_expr ] elif isinstance(asset_cfg, VisuoTactileSensorCfg): if hasattr(asset_cfg, "camera_cfg") and asset_cfg.camera_cfg is not None: asset_cfg.camera_cfg.prim_path = asset_cfg.camera_cfg.prim_path.format( - ENV_REGEX_NS=self.env_regex_ns + ENV_REGEX_NS=env_regex_ns ) if ( hasattr(asset_cfg, "contact_object_prim_path_expr") and asset_cfg.contact_object_prim_path_expr is not None ): asset_cfg.contact_object_prim_path_expr = asset_cfg.contact_object_prim_path_expr.format( - ENV_REGEX_NS=self.env_regex_ns + ENV_REGEX_NS=env_regex_ns ) self._sensors[asset_name] = asset_cfg.class_type(asset_cfg) diff --git a/source/isaaclab/isaaclab/scene/interactive_scene_cfg.py b/source/isaaclab/isaaclab/scene/interactive_scene_cfg.py index 034307017219..c628d314d993 100644 --- a/source/isaaclab/isaaclab/scene/interactive_scene_cfg.py +++ b/source/isaaclab/isaaclab/scene/interactive_scene_cfg.py @@ -112,15 +112,8 @@ class MySceneCfg(InteractiveSceneCfg): """ clone_in_fabric: bool = False - """Enable/disable cloning in fabric. Default is False. - - Omniverse Fabric is a more optimized method for performing cloning in scene creation. This reduces the time - taken to create the scene. However, it limits flexibility in accessing the stage through USD APIs and instead, - the stage must be accessed through USDRT. - - .. note:: - Cloning in fabric can only be enabled if :attr:`replicated_physics` is also enabled. - If :attr:`replicated_physics` is ``False``, cloning in Fabric will automatically - default to ``False``. + """Deprecated legacy Fabric cloning flag. Default is False. + Queued replication no longer forwards this flag to the PhysX replicator; + ``useFabricForReplication`` is always ``False``. """ diff --git a/source/isaaclab/isaaclab/sensors/camera/camera.py b/source/isaaclab/isaaclab/sensors/camera/camera.py index a26961c6fd75..175012b10b62 100644 --- a/source/isaaclab/isaaclab/sensors/camera/camera.py +++ b/source/isaaclab/isaaclab/sensors/camera/camera.py @@ -17,6 +17,7 @@ import isaaclab.sim as sim_utils import isaaclab.utils.sensors as sensor_utils +from isaaclab.cloner import queue_usd_replication from isaaclab.renderers import BaseRenderer, CameraRenderSpec from isaaclab.sim.views import FrameView from isaaclab.utils import to_camel_case @@ -156,6 +157,7 @@ def __init__(self, cfg: CameraCfg): if self.cfg.spawn is not None and self.cfg.spawn.vertical_aperture is None: self.cfg.spawn.vertical_aperture = self.cfg.spawn.horizontal_aperture * self.cfg.height / self.cfg.width self._resolve_and_spawn("camera", translation=self.cfg.offset.pos, orientation=rot_offset) + queue_usd_replication(self._source_cfg) # An ISP (any ``isp_cfg`` other than ``None``) requires the HDR AOV; # an explicit ``"rgb_hdr"`` in ``data_types`` also requires the diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py index 27af7fa8b0be..a30510b12d93 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py @@ -17,6 +17,7 @@ import isaaclab.sim as sim_utils import isaaclab.utils.math as math_utils +from isaaclab.cloner import queue_usd_replication from isaaclab.markers import VisualizationMarkers from isaaclab.terrains.trimesh.utils import make_plane from isaaclab.utils.warp import ProxyArray, convert_to_warp_mesh @@ -72,6 +73,7 @@ def __init__(self, cfg: RayCasterCfg): # Resolve physics-body paths and spawn the sensor Xform child if needed. self._requested_prim_path = self.cfg.prim_path self._resolve_and_spawn("raycaster") + queue_usd_replication(self._source_cfg) self._data = RayCasterData() def __str__(self) -> str: diff --git a/source/isaaclab/isaaclab/sensors/sensor_base.py b/source/isaaclab/isaaclab/sensors/sensor_base.py index 15fca5ee4ad4..b266ce7cb425 100644 --- a/source/isaaclab/isaaclab/sensors/sensor_base.py +++ b/source/isaaclab/isaaclab/sensors/sensor_base.py @@ -53,6 +53,7 @@ def __init__(self, cfg: SensorBaseCfg): # check that the config is valid cfg.validate() # store inputs + self._source_cfg = cfg self.cfg = cfg.copy() # flag for whether the sensor is initialized self._is_initialized = False diff --git a/source/isaaclab/isaaclab/sim/simulation_context.py b/source/isaaclab/isaaclab/sim/simulation_context.py index 6cf6bfa6e1b0..8253839d8083 100644 --- a/source/isaaclab/isaaclab/sim/simulation_context.py +++ b/source/isaaclab/isaaclab/sim/simulation_context.py @@ -181,7 +181,7 @@ def __init__(self, cfg: SimulationCfg | None = None): self._scene_data_requirements = SceneDataRequirement() # Clone plan published by InteractiveScene after cloning. Providers (e.g. the # Newton visualizer model rebuilder on a PhysX backend) consume this to derive - # their own backend args. None until :meth:`InteractiveScene.clone_environments` runs. + # their own backend args. None until a replication session publishes a plan. self._clone_plan: ClonePlan | None = None # Default visualization dt used before/without visualizer initialization. physics_dt = getattr(self.cfg.physics, "dt", None) @@ -674,9 +674,9 @@ def update_scene_data_requirements(self, requirements: SceneDataRequirement) -> def get_clone_plan(self) -> ClonePlan | None: """Return the clone plan published by the scene. - Set by :meth:`InteractiveScene.clone_environments` after replication. Consumed by - scene data providers that build backend models (e.g. Newton visualizer model on a - PhysX backend) from the same plan the cloner used. ``None`` until the scene clones. + Set after replication. Consumed by scene data providers that build backend models + (e.g. Newton visualizer model on a PhysX backend) from the same plan the cloner used. + ``None`` until the scene replicates. """ return self._clone_plan diff --git a/source/isaaclab/test/scene/test_interactive_scene.py b/source/isaaclab/test/scene/test_interactive_scene.py index f56803ef5cf6..a2ddc05f5eae 100644 --- a/source/isaaclab/test/scene/test_interactive_scene.py +++ b/source/isaaclab/test/scene/test_interactive_scene.py @@ -12,7 +12,6 @@ """Rest everything follows.""" -import contextlib from types import SimpleNamespace import pytest @@ -21,6 +20,7 @@ import isaaclab.sim as sim_utils from isaaclab.actuators import ImplicitActuatorCfg from isaaclab.assets import ArticulationCfg, RigidObjectCfg, RigidObjectCollectionCfg +from isaaclab.cloner import CloneCfg from isaaclab.physics.scene_data_requirements import SceneDataRequirement from isaaclab.scene import InteractiveScene, InteractiveSceneCfg from isaaclab.sim import build_simulation_context @@ -130,245 +130,36 @@ def test_reset_to_env_ids_input_types(device, setup_scene): assert_state_equal(prev_state, scene.get_state()) -def test_clone_environments_executes_env_root_plan_with_positions(monkeypatch: pytest.MonkeyPatch): - """Env-root plans replicate the whole environment and keep grid positions.""" - from isaaclab.cloner import ClonePlan +def test_scene_publishes_plan_via_replicate(monkeypatch: pytest.MonkeyPatch): + """A cfg-driven scene forwards the right plan and stage to cloner.replicate. - scene = object.__new__(InteractiveScene) - scene.cfg = SimpleNamespace(replicate_physics=False, num_envs=3) - scene.stage = object() - scene.physics_backend = "physx" - scene._sensors = {} - - set_plan_calls: list = [] - sim_state: dict = {"plan": None} - - def _set_clone_plan(plan): - sim_state["plan"] = plan - set_plan_calls.append(plan) - - scene.sim = SimpleNamespace( - get_scene_data_requirements=lambda: SceneDataRequirement(), - update_scene_data_requirements=lambda requirements: None, - set_clone_plan=_set_clone_plan, - get_clone_plan=lambda: sim_state["plan"], - ) - scene.env_fmt = "/World/envs/env_{}" - scene._ALL_INDICES = torch.arange(3, dtype=torch.long) - scene._default_env_origins = torch.zeros((3, 3), dtype=torch.float32) - scene._clone_plan = ClonePlan( - sources=(scene.env_fmt.format(0),), - destinations=(scene.env_fmt,), - clone_mask=torch.ones((1, scene.num_envs), dtype=torch.bool), - ) - # Avoid binding this unit test to global SimulationContext singleton state. - monkeypatch.setattr(InteractiveScene, "device", property(lambda self: "cpu")) - - # ``disabled_fabric_change_notifies`` resolves the stage via UsdUtils.StageCache and would - # crash on the bare ``object()`` mocked above. This unit test exercises clone-dispatch - # logic only; the fabric notice path has its own coverage in ``test_cloner.py``. - @contextlib.contextmanager - def _noop_fabric_notices(stage, *, restore=True): - yield - - monkeypatch.setattr("isaaclab.scene.interactive_scene.cloner.disabled_fabric_change_notifies", _noop_fabric_notices) - - physics_calls = [] - usd_calls = [] - - def _physics_clone_fn(stage, *args, **kwargs): - physics_calls.append((stage, args, kwargs)) - - def _usd_replicate(stage, *args, **kwargs): - usd_calls.append((stage, args, kwargs)) - - scene.cloner_cfg = SimpleNamespace( - device="cpu", - physics_clone_fn=_physics_clone_fn, - clone_usd=True, - ) - monkeypatch.setattr("isaaclab.scene.interactive_scene.cloner.usd_replicate", _usd_replicate) - - scene.clone_environments(copy_from_source=False) - assert len(physics_calls) == 1 - assert len(usd_calls) == 1 - mapping = physics_calls[0][1][3] - assert mapping.dtype == torch.bool - assert mapping.shape == (1, scene.num_envs) - assert physics_calls[0][2]["positions"] is scene._default_env_origins - assert usd_calls[0][2]["positions"] is scene._default_env_origins - assert len(set_plan_calls) == 1 - plan = set_plan_calls[-1] - assert isinstance(plan, ClonePlan) - assert plan.sources == (scene.env_fmt.format(0),) - assert plan.destinations == (scene.env_fmt,) - assert plan.clone_mask.shape == (1, scene.num_envs) - assert scene.clone_plan is plan - - physics_calls.clear() - usd_calls.clear() - set_plan_calls.clear() - scene.clone_environments(copy_from_source=True) - assert len(physics_calls) == 0 - assert len(usd_calls) == 1 - assert len(set_plan_calls) == 1 - - -def test_clone_environments_skips_replication_without_plan(): - """Direct-path cfg scenes publish no plan and do not dispatch cloners.""" - scene = object.__new__(InteractiveScene) - scene._clone_plan = None - set_plan_calls = [] - scene.sim = SimpleNamespace(set_clone_plan=set_plan_calls.append) - - scene.clone_environments(copy_from_source=False) - - assert set_plan_calls == [None] - - -def test_clone_environments_executes_asset_level_plan_without_usd_positions(monkeypatch: pytest.MonkeyPatch): - """Asset-level plans preserve env-root transforms by skipping USD positions.""" - from isaaclab.cloner import ClonePlan - - scene = object.__new__(InteractiveScene) - scene.cfg = SimpleNamespace(replicate_physics=False, num_envs=2) - scene.stage = object() - scene.physics_backend = "physx" - scene._sensors = {} - scene.env_fmt = "/World/envs/env_{}" - scene._ALL_INDICES = torch.arange(2, dtype=torch.long) - scene._default_env_origins = torch.ones((2, 3), dtype=torch.float32) - scene._clone_plan = ClonePlan( - sources=("/World/envs/env_0/Object", "/World/envs/env_1/Object"), - destinations=("/World/envs/env_{}/Object", "/World/envs/env_{}/Object"), - clone_mask=torch.tensor([[True, False], [False, True]], dtype=torch.bool), - ) - - set_plan_calls: list = [] - scene.sim = SimpleNamespace(set_clone_plan=set_plan_calls.append) - monkeypatch.setattr(InteractiveScene, "device", property(lambda self: "cpu")) - - @contextlib.contextmanager - def _noop_fabric_notices(stage, *, restore=True): - yield - - monkeypatch.setattr("isaaclab.scene.interactive_scene.cloner.disabled_fabric_change_notifies", _noop_fabric_notices) - monkeypatch.setattr( - "isaaclab.scene.interactive_scene.cloner.usd_replicate", - lambda *args, **kwargs: usd_calls.append((args, kwargs)), - ) - - physics_calls = [] - usd_calls = [] - scene.cloner_cfg = SimpleNamespace( - device="cpu", - physics_clone_fn=lambda *args, **kwargs: physics_calls.append((args, kwargs)), - clone_usd=True, - ) - - scene.clone_environments(copy_from_source=False) - - assert len(physics_calls) == 1 - assert physics_calls[0][1]["positions"] is scene._default_env_origins - assert len(usd_calls) == 1 - assert usd_calls[0][1]["positions"] is None - assert set_plan_calls == [scene._clone_plan] - - -def test_build_clone_plan_from_cfg_plans_multi_and_single_spawners(monkeypatch: pytest.MonkeyPatch): - """Heterogeneous planning writes source paths for multi and single spawners.""" - from isaaclab.cloner import sequential - - scene = object.__new__(InteractiveScene) - scene.cfg = SimpleNamespace( - num_envs=4, - object=SimpleNamespace( - prim_path="{ENV_REGEX_NS}/Object", - spawn=sim_utils.MultiAssetSpawnerCfg( - assets_cfg=[ - sim_utils.ConeCfg(radius=0.1, height=0.2), - sim_utils.SphereCfg(radius=0.1), - ] - ), - ), - robot=SimpleNamespace( - prim_path="{ENV_REGEX_NS}/Robot", - spawn=sim_utils.CuboidCfg(size=(0.1, 0.1, 0.1)), - ), - ) - scene.env_fmt = "/World/envs/env_{}" - scene.cloner_cfg = SimpleNamespace(clone_strategy=sequential) - monkeypatch.setattr(InteractiveScene, "device", property(lambda self: "cpu")) - - plan = scene._build_clone_plan_from_cfg() - - assert plan is not None - assert plan.sources == ( - "/World/envs/env_0/Object", - "/World/envs/env_1/Object", - "/World/envs/env_0/Robot", - ) - assert plan.destinations == ( - "/World/envs/env_{}/Object", - "/World/envs/env_{}/Object", - "/World/envs/env_{}/Robot", - ) - assert scene.cfg.object.spawn.spawn_paths == ["/World/envs/env_0/Object", "/World/envs/env_1/Object"] - assert scene.cfg.robot.spawn.spawn_path == "/World/envs/env_0/Robot" - assert scene.cfg.object.prim_path == "{ENV_REGEX_NS}/Object" - assert scene.cfg.robot.prim_path == "{ENV_REGEX_NS}/Robot" - assert torch.equal(plan.clone_mask.to(torch.int).argmax(dim=0).cpu(), torch.tensor([0, 1, 0, 1])) + Uses a test-seam fake to isolate this unit test from real backend dispatch; queue + lifecycle is owned by :func:`replicate` itself (snapshot-and-clear) and does not + need any cleanup hook here. + """ + import isaaclab.cloner.replicate_session as replicate_session_module + captured: list = [] -def test_build_clone_plan_from_cfg_defaults_to_env0_plan(monkeypatch: pytest.MonkeyPatch): - """Homogeneous cfg scenes use the default env_0-to-all ClonePlan.""" - from isaaclab.cloner import sequential + def fake_replicate(plan, *, stage): + captured.append((plan, stage)) - scene = object.__new__(InteractiveScene) - scene.cfg = SimpleNamespace( - num_envs=3, - robot=SimpleNamespace( - prim_path="{ENV_REGEX_NS}/Robot", - spawn=sim_utils.CuboidCfg(size=(0.1, 0.1, 0.1)), - ), - ) - scene.env_fmt = "/World/envs/env_{}" - scene.cloner_cfg = SimpleNamespace(clone_strategy=sequential) - monkeypatch.setattr(InteractiveScene, "device", property(lambda self: "cpu")) + monkeypatch.setattr(replicate_session_module, "replicate", fake_replicate) - plan = scene._build_clone_plan_from_cfg() + with build_simulation_context(device="cpu", auto_add_lighting=False, add_ground_plane=False) as sim: + sim._app_control_on_stop_handle = None + scene = InteractiveScene(MySceneCfg(num_envs=4, env_spacing=1.0)) - assert plan is not None + assert len(captured) == 1 + plan, stage = captured[0] assert plan.sources == ("/World/envs/env_0",) - assert plan.destinations == (scene.env_fmt,) - assert plan.clone_mask.shape == (1, scene.num_envs) - assert scene.cfg.robot.spawn.spawn_path == "/World/envs/env_0/Robot" - + assert plan.destinations == ("/World/envs/env_{}",) + assert plan.clone_mask.shape == (1, 4) + assert stage is scene.stage -def test_build_clone_plan_from_cfg_returns_none_without_env_scoped_groups(monkeypatch: pytest.MonkeyPatch): - """Direct-path cfg scenes should not force env-root replication.""" - from isaaclab.cloner import sequential - - scene = object.__new__(InteractiveScene) - scene.cfg = SimpleNamespace( - num_envs=1, - robot=SimpleNamespace( - prim_path="/World/Robot", - spawn=sim_utils.CuboidCfg(size=(0.1, 0.1, 0.1)), - ), - ) - scene.env_fmt = "/World/envs/env_{}" - scene.cloner_cfg = SimpleNamespace(clone_strategy=sequential) - monkeypatch.setattr(InteractiveScene, "device", property(lambda self: "cpu")) - - assert scene._build_clone_plan_from_cfg() is None - assert scene.cfg.robot.spawn.spawn_path is None - - -def test_build_clone_plan_from_cfg_sets_collection_member_paths(monkeypatch: pytest.MonkeyPatch): - """Rigid object collection members are planned independently.""" - from isaaclab.cloner import sequential +def test_collect_asset_cfgs_resolves_env_regex_macros(): + """_collect_asset_cfgs rewrites {ENV_REGEX_NS} macros and expands collections.""" scene = object.__new__(InteractiveScene) cube_cfg = RigidObjectCfg( prim_path="{ENV_REGEX_NS}/Cube", @@ -381,51 +172,31 @@ def test_build_clone_plan_from_cfg_sets_collection_member_paths(monkeypatch: pyt ), ) scene.cfg = SimpleNamespace( - num_envs=4, + num_envs=2, objects=RigidObjectCollectionCfg(rigid_objects={"cube": cube_cfg, "shape": shape_cfg}), ) - scene.env_fmt = "/World/envs/env_{}" - scene.cloner_cfg = SimpleNamespace(clone_strategy=sequential) - monkeypatch.setattr(InteractiveScene, "device", property(lambda self: "cpu")) + scene.cloner_cfg = CloneCfg() - plan = scene._build_clone_plan_from_cfg() + cfgs = scene._collect_asset_cfgs() - assert plan is not None - planned_cube = scene.cfg.objects.rigid_objects["cube"] - planned_shape = scene.cfg.objects.rigid_objects["shape"] - assert planned_cube.spawn.spawn_path == "/World/envs/env_0/Cube" - assert planned_shape.spawn.spawn_paths == ["/World/envs/env_0/Shape", "/World/envs/env_1/Shape"] - assert "/World/envs/env_{}/Cube" in plan.destinations - assert "/World/envs/env_{}/Shape" in plan.destinations + prim_paths = sorted(c.prim_path for c in cfgs) + assert prim_paths == ["/World/envs/env_.*/Cube", "/World/envs/env_.*/Shape"] -def test_build_clone_plan_from_cfg_marks_unused_variants(monkeypatch: pytest.MonkeyPatch): - """Unused variants keep a mask row but do not get spawned.""" - from isaaclab.cloner import sequential +def test_collect_asset_cfgs_orders_sensors_last(): + """Non-sensor cfgs precede sensor cfgs in _collect_asset_cfgs output.""" + from isaaclab.sensors import ContactSensorCfg scene = object.__new__(InteractiveScene) - scene.cfg = SimpleNamespace( - num_envs=2, - object=SimpleNamespace( - prim_path="{ENV_REGEX_NS}/Object", - spawn=sim_utils.MultiAssetSpawnerCfg( - assets_cfg=[ - sim_utils.ConeCfg(radius=0.1, height=0.2), - sim_utils.CuboidCfg(size=(0.1, 0.1, 0.1)), - sim_utils.SphereCfg(radius=0.1), - ] - ), - ), - ) - scene.env_fmt = "/World/envs/env_{}" - scene.cloner_cfg = SimpleNamespace(clone_strategy=sequential) - monkeypatch.setattr(InteractiveScene, "device", property(lambda self: "cpu")) + sensor = ContactSensorCfg(prim_path="{ENV_REGEX_NS}/Robot") + body = SimpleNamespace(prim_path="{ENV_REGEX_NS}/Robot") + scene.cfg = SimpleNamespace(num_envs=1, sensor=sensor, body=body) + scene.cloner_cfg = CloneCfg() - plan = scene._build_clone_plan_from_cfg() + cfgs = scene._collect_asset_cfgs() - assert plan is not None - assert scene.cfg.object.spawn.spawn_paths == ["/World/envs/env_0/Object", "/World/envs/env_1/Object", None] - assert plan.clone_mask[2].sum() == 0 + # Sensors come after non-sensor entities so they can bind to spawned bodies. + assert cfgs.index(body) < cfgs.index(sensor) def test_aggregate_scene_data_requirements_merges_visualizers_and_renderers(monkeypatch: pytest.MonkeyPatch): diff --git a/source/isaaclab/test/sensors/test_multi_mesh_ray_caster_camera.py b/source/isaaclab/test/sensors/test_multi_mesh_ray_caster_camera.py index 8f5014b3a860..f7f7a610cbdd 100644 --- a/source/isaaclab/test/sensors/test_multi_mesh_ray_caster_camera.py +++ b/source/isaaclab/test/sensors/test_multi_mesh_ray_caster_camera.py @@ -514,6 +514,9 @@ def _create_heterogeneous_clone_scene(sim: sim_utils.SimulationContext, num_envs env_fmt + "/Object", ), clone_mask=torch.cat([robot_mask, object_mask], dim=0), + env_ids=env_ids, + positions=None, + cfg_rows={}, ) ) sim_utils.update_stage() diff --git a/source/isaaclab/test/sensors/test_ray_caster_integration.py b/source/isaaclab/test/sensors/test_ray_caster_integration.py index 91fa9c12fedc..d7bbe7e76e78 100644 --- a/source/isaaclab/test/sensors/test_ray_caster_integration.py +++ b/source/isaaclab/test/sensors/test_ray_caster_integration.py @@ -358,6 +358,9 @@ def _create_object_body(path: str) -> None: sources=("/World/envs/env_0/Object", "/World/envs/env_1/Object"), destinations=("/World/envs/env_{}/Object", "/World/envs/env_{}/Object"), clone_mask=torch.tensor([[True, False, True], [False, True, False]], dtype=torch.bool, device=sim.device), + env_ids=torch.arange(3, dtype=torch.long, device=sim.device), + positions=None, + cfg_rows={}, ) ) sim_utils.update_stage() diff --git a/source/isaaclab/test/sim/test_cloner.py b/source/isaaclab/test/sim/test_cloner.py index 1f526ac74584..3b2757484a31 100644 --- a/source/isaaclab/test/sim/test_cloner.py +++ b/source/isaaclab/test/sim/test_cloner.py @@ -14,13 +14,26 @@ """Rest everything follows.""" +from types import SimpleNamespace +from unittest.mock import MagicMock + import pytest import torch from pxr import UsdGeom import isaaclab.sim as sim_utils -from isaaclab.cloner import make_clone_plan, sequential, usd_replicate +from isaaclab.cloner import ( + REPLICATION_QUEUE, + ClonePlan, + UsdReplicateContext, + grid_transforms, + make_clone_plan, + queue_usd_replication, + replicate, + sequential, + usd_replicate, +) from isaaclab.sim import build_simulation_context pytestmark = pytest.mark.isaacsim_ci @@ -33,6 +46,14 @@ def sim(request): yield sim +@pytest.fixture(autouse=True) +def _drain_replication_queue(): + """Ensure REPLICATION_QUEUE starts empty for every test and is cleared after.""" + REPLICATION_QUEUE.clear() + yield + REPLICATION_QUEUE.clear() + + def test_usd_replicate_with_positions_and_mask(sim): """Replicate sources to selected envs and author translate ops from positions.""" # Prepare sources under /World/template @@ -75,6 +96,28 @@ def test_usd_replicate_with_positions_and_mask(sim): assert any(op.GetOpType() == UsdGeom.XformOp.TypeTranslate for op in ops) +def test_usd_replicate_context_queue_and_replicate(sim): + """UsdReplicateContext queues copy specs and applies them on replicate.""" + sim_utils.create_prim("/World/template", "Xform") + sim_utils.create_prim("/World/template/A", "Xform") + sim_utils.create_prim("/World/envs", "Xform") + sim_utils.create_prim("/World/envs/env_0", "Xform") + sim_utils.create_prim("/World/envs/env_1", "Xform") + + stage = sim_utils.get_current_stage() + ctx = UsdReplicateContext(stage) + ctx.queue_mapping( + sources=["/World/template/A"], + destinations=["/World/envs/env_{}/A"], + env_ids=torch.tensor([0, 1], dtype=torch.long), + ) + assert not stage.GetPrimAtPath("/World/envs/env_1/A").IsValid() + ctx.replicate() + + assert stage.GetPrimAtPath("/World/envs/env_0/A").IsValid() + assert stage.GetPrimAtPath("/World/envs/env_1/A").IsValid() + + def test_usd_replicate_depth_order_parent_child(sim): """Replicate parent and child when provided out of order; parent should exist before child.""" # Prepare sources @@ -103,12 +146,7 @@ def test_usd_replicate_depth_order_parent_child(sim): def test_usd_replicate_self_copy_skips_copy_spec(sim): - """usd_replicate must not call Sdf.CopySpec when source and destination paths are identical. - - Sdf.CopySpec(src, src) is a no-op in the current USD version so it does not corrupt children, - but the call is still wasteful. The guard ensures it is skipped entirely. This test mocks - Sdf.CopySpec to verify it is called exactly once (for env_1) and never for the self case (env_0). - """ + """usd_replicate must not call Sdf.CopySpec when source and destination paths are identical.""" from unittest.mock import patch import isaaclab.cloner.cloner_utils as _cloner_mod @@ -136,7 +174,6 @@ def capturing_copy_spec(src_layer, src_path, dst_layer, dst_path): mask=torch.ones((1, 2), dtype=torch.bool), ) - # CopySpec must be called for env_1 but never for env_0 (self-copy) assert all(src != dst for src, dst in copy_calls), f"Self-copy detected in CopySpec calls: {copy_calls}" assert any(dst == "/World/envs/env_1" for _, dst in copy_calls), "CopySpec was not called for env_1" @@ -180,19 +217,7 @@ def capturing_copy_spec(src_layer, src_path, dst_layer, dst_path): def test_clone_decorator_wildcard_patterns( sim, parent_paths, spawn_pattern, expected_child_paths, bad_path, match_expr ): - """The @clone decorator handles two distinct wildcard patterns correctly. - - Case A – ``.*`` in root_path (parent is a regex): the child prim is spawned at - ``source_prim_paths[0]`` as a prototype and then copied to every other matching - parent via ``Sdf.CopySpec``, so **all** parents end up with the child. The old - ``prim_path.replace(".*", "0")`` approach created spurious intermediate prims - that inflated ``find_matching_prims`` counts and broke tiled-camera initialization. - - Case B – ``.*`` only in asset_path (leaf): no parent regex, so - ``source_prim_paths == [root_path]`` (one entry, no copy step). Replacing - ``".*"`` → ``"0"`` in the asset name gives the intended prototype name - (e.g. ``proto_asset_0``) under the single real parent. - """ + """The @clone decorator handles two distinct wildcard patterns correctly.""" for path in parent_paths: sim_utils.create_prim(path, "Xform") @@ -201,19 +226,16 @@ def test_clone_decorator_wildcard_patterns( stage = sim_utils.get_current_stage() - # Every expected child path must exist for child_path in expected_child_paths: assert stage.GetPrimAtPath(child_path).IsValid(), ( f"Prim was not spawned at '{child_path}'. The @clone decorator may have used the wrong spawn path." ) - # The spurious path from the old replace(".*", "0") must NOT exist assert not stage.GetPrimAtPath(bad_path).IsValid(), ( f"Spurious prim found at '{bad_path}'. " "The @clone decorator incorrectly derived the spawn path by replacing '.*' with '0'." ) - # find_matching_prims must see exactly the original parents — no spurious extras all_matching = sim_utils.find_matching_prims(match_expr) assert len(all_matching) == len(parent_paths), ( f"Expected {len(parent_paths)} matching prims, got {len(all_matching)}. " @@ -221,20 +243,289 @@ def test_clone_decorator_wildcard_patterns( ) -def test_make_clone_plan_returns_flat_source_rows(sim): - """make_clone_plan exposes the flat source-to-env mask used by scene cloning.""" +def test_queue_usd_replication_only_appends(sim): + """queue_usd_replication must only append to REPLICATION_QUEUE — no other side effects.""" + cfg_a = SimpleNamespace(prim_path="/World/envs/env_.*/Robot") + cfg_b = SimpleNamespace(prim_path="/World/envs/env_.*/Object") + + queue_usd_replication(cfg_a) + queue_usd_replication(cfg_b) + + assert [(cfg_a, UsdReplicateContext), (cfg_b, UsdReplicateContext)] == REPLICATION_QUEUE + + +def test_make_clone_plan_homogeneous_returns_env_root_plan(sim): + """Homogeneous (single-variant) cfgs produce one source row at the env root.""" + cube = SimpleNamespace( + prim_path="/World/envs/env_.*/Robot", + spawn=sim_utils.CuboidCfg(size=(0.1, 0.1, 0.1)), + ) + plan = make_clone_plan( - [["/World/envs/env_0/Object", "/World/envs/env_1/Object"]], - ["/World/envs/env_{}/Object"], + cfgs=[cube], num_clones=4, + env_spacing=1.0, + device=sim.cfg.device, + ) + + assert plan.sources == ("/World/envs/env_0",) + assert plan.destinations == ("/World/envs/env_{}",) + assert plan.clone_mask.shape == (1, 4) + assert plan.clone_mask.all() + assert plan.cfg_rows[id(cube)] == (0,) + assert plan.env_ids.shape == (4,) + assert plan.positions.shape == (4, 3) + assert cube.spawn.spawn_path == "/World/envs/env_0/Robot" + + +def test_make_clone_plan_heterogeneous_mutates_spawn_paths(sim): + """Multi-variant spawners get per-variant spawn_paths and contribute multiple plan rows.""" + multi_cfg = SimpleNamespace( + prim_path="/World/envs/env_.*/Object", + spawn=sim_utils.MultiAssetSpawnerCfg( + assets_cfg=[ + sim_utils.ConeCfg(radius=0.1, height=0.2), + sim_utils.SphereCfg(radius=0.1), + ] + ), + ) + plain_cfg = SimpleNamespace( + prim_path="/World/envs/env_.*/Robot", + spawn=sim_utils.CuboidCfg(size=(0.1, 0.1, 0.1)), + ) + + plan = make_clone_plan( + cfgs=[multi_cfg, plain_cfg], + num_clones=4, + env_spacing=1.0, + device=sim.cfg.device, clone_strategy=sequential, + ) + + assert plan.destinations == ( + "/World/envs/env_{}/Object", + "/World/envs/env_{}/Object", + "/World/envs/env_{}/Robot", + ) + assert plan.cfg_rows[id(multi_cfg)] == (0, 1) + assert plan.cfg_rows[id(plain_cfg)] == (2,) + assert multi_cfg.spawn.spawn_paths == ["/World/envs/env_0/Object", "/World/envs/env_1/Object"] + assert plain_cfg.spawn.spawn_path == "/World/envs/env_0/Robot" + + +def test_make_clone_plan_skips_global_cfgs(sim): + """Cfgs whose prim_path is not under /World/envs/ are excluded from the plan.""" + global_cfg = SimpleNamespace( + prim_path="/World/global/Robot", + spawn=sim_utils.CuboidCfg(size=(0.1, 0.1, 0.1)), + ) + + plan = make_clone_plan( + cfgs=[global_cfg], + num_clones=3, + env_spacing=1.0, + device=sim.cfg.device, + ) + + assert plan.sources == () + assert plan.destinations == () + assert plan.clone_mask.shape == (0, 3) + assert plan.cfg_rows == {} + + +def test_clone_plan_from_env_0_populates_cfg_rows(sim): + """from_env_0 auto-maps queued env-scoped cfgs to row 0 and excludes global ones.""" + env_cfg_a = SimpleNamespace(prim_path="/World/envs/env_.*/Robot") + env_cfg_b = SimpleNamespace(prim_path="/World/envs/env_.*/Object") + global_cfg = SimpleNamespace(prim_path="/World/global/Light") + + queue_usd_replication(env_cfg_a) + queue_usd_replication(env_cfg_b) + queue_usd_replication(global_cfg) + + plan = ClonePlan.from_env_0( + source="/World/envs/env_0", + destination="/World/envs/env_{}", + num_clones=4, device=sim.cfg.device, + positions=grid_transforms(4, 1.0, device=sim.cfg.device)[0], ) - assert plan.sources == ("/World/envs/env_0/Object", "/World/envs/env_1/Object") - assert plan.destinations == ("/World/envs/env_{}/Object", "/World/envs/env_{}/Object") - assert plan.clone_mask.shape == (2, 4) - assert plan.clone_mask.dtype == torch.bool - assert torch.all(plan.clone_mask.sum(dim=0) == 1) - actual_source_idx = plan.clone_mask.to(torch.int).argmax(dim=0).cpu() - assert torch.equal(actual_source_idx, torch.tensor([0, 1, 0, 1])) + assert plan.sources == ("/World/envs/env_0",) + assert plan.destinations == ("/World/envs/env_{}",) + assert plan.cfg_rows == {id(env_cfg_a): (0,), id(env_cfg_b): (0,)} + assert plan.clone_mask.all() and plan.clone_mask.shape == (1, 4) + assert torch.equal(plan.env_ids, torch.arange(4, dtype=torch.long, device=sim.cfg.device)) + + +def test_replicate_drains_queue_dispatches_and_publishes(sim): + """replicate(plan) drains REPLICATION_QUEUE, calls each backend once, publishes, clears.""" + + class FakeCtx: + replicate_priority = 0 + instances: list["FakeCtx"] = [] + + def __init__(self, stage): + self.stage = stage + self.queue_calls: list[tuple] = [] + self.replicate_calls = 0 + FakeCtx.instances.append(self) + + def queue_mapping(self, sources, destinations, env_ids, mask, *, positions=None): + self.queue_calls.append((tuple(sources), tuple(destinations), mask.clone())) + + def replicate(self): + self.replicate_calls += 1 + + cfg_a = SimpleNamespace(prim_path="/World/envs/env_.*/Robot") + cfg_b = SimpleNamespace(prim_path="/World/envs/env_.*/Object") + REPLICATION_QUEUE.append((cfg_a, FakeCtx)) + REPLICATION_QUEUE.append((cfg_b, FakeCtx)) + + plan = ClonePlan( + sources=("/World/envs/env_0/Robot", "/World/envs/env_0/Object"), + destinations=("/World/envs/env_{}/Robot", "/World/envs/env_{}/Object"), + clone_mask=torch.ones((2, 4), dtype=torch.bool, device=sim.cfg.device), + env_ids=torch.arange(4, dtype=torch.long, device=sim.cfg.device), + positions=grid_transforms(4, 1.0, device=sim.cfg.device)[0], + cfg_rows={id(cfg_a): (0,), id(cfg_b): (1,)}, + ) + sim.set_clone_plan(None) + + replicate(plan, stage=sim_utils.get_current_stage()) + + # Exactly one FakeCtx instance is shared across both cfgs. + assert len(FakeCtx.instances) == 1 + ctx = FakeCtx.instances[0] + assert len(ctx.queue_calls) == 2 + assert ctx.queue_calls[0][0] == ("/World/envs/env_0/Robot",) + assert ctx.queue_calls[1][0] == ("/World/envs/env_0/Object",) + assert ctx.replicate_calls == 1 + assert sim.get_clone_plan() is plan + assert REPLICATION_QUEUE == [] + + +def test_replicate_runs_lower_priority_backends_first(sim): + """Sort order: lower replicate_priority runs first (physics before USD).""" + + call_order: list[str] = [] + + class LowPriority: + replicate_priority = 0 + + def __init__(self, stage): + pass + + def queue_mapping(self, *args, **kwargs): + pass + + def replicate(self): + call_order.append("low") + + class HighPriority: + replicate_priority = 100 + + def __init__(self, stage): + pass + + def queue_mapping(self, *args, **kwargs): + pass + + def replicate(self): + call_order.append("high") + + cfg = SimpleNamespace(prim_path="/World/envs/env_.*/Robot") + REPLICATION_QUEUE.append((cfg, HighPriority)) + REPLICATION_QUEUE.append((cfg, LowPriority)) + + plan = ClonePlan( + sources=("/World/envs/env_0",), + destinations=("/World/envs/env_{}",), + clone_mask=torch.ones((1, 2), dtype=torch.bool, device=sim.cfg.device), + env_ids=torch.arange(2, dtype=torch.long, device=sim.cfg.device), + positions=None, + cfg_rows={id(cfg): (0,)}, + ) + replicate(plan, stage=sim_utils.get_current_stage()) + + assert call_order == ["low", "high"] + + +def test_replicate_skips_cfgs_not_in_plan(sim): + """Cfgs absent from plan.cfg_rows are silently skipped.""" + sentinel = MagicMock() + sentinel.replicate_priority = 0 + sentinel.replicate.side_effect = lambda: None + sentinel_cls = MagicMock(return_value=sentinel) + + excluded_cfg = SimpleNamespace(prim_path="/World/global/Skip") + REPLICATION_QUEUE.append((excluded_cfg, sentinel_cls)) + + plan = ClonePlan( + sources=("/World/envs/env_0",), + destinations=("/World/envs/env_{}",), + clone_mask=torch.ones((1, 2), dtype=torch.bool, device=sim.cfg.device), + env_ids=torch.arange(2, dtype=torch.long, device=sim.cfg.device), + positions=None, + cfg_rows={}, + ) + replicate(plan, stage=sim_utils.get_current_stage()) + + sentinel_cls.assert_not_called() + + +def test_replicate_clears_queue_on_backend_failure(sim): + """REPLICATION_QUEUE is drained even when a backend ctx raises mid-dispatch.""" + + class ExplodingCtx: + replicate_priority = 0 + + def __init__(self, stage): + pass + + def queue_mapping(self, *args, **kwargs): + pass + + def replicate(self): + raise RuntimeError("backend boom") + + cfg = SimpleNamespace(prim_path="/World/envs/env_.*/Robot") + REPLICATION_QUEUE.append((cfg, ExplodingCtx)) + + plan = ClonePlan( + sources=("/World/envs/env_0",), + destinations=("/World/envs/env_{}",), + clone_mask=torch.ones((1, 2), dtype=torch.bool, device=sim.cfg.device), + env_ids=torch.arange(2, dtype=torch.long, device=sim.cfg.device), + positions=None, + cfg_rows={id(cfg): (0,)}, + ) + + with pytest.raises(RuntimeError, match="backend boom"): + replicate(plan, stage=sim_utils.get_current_stage()) + + assert REPLICATION_QUEUE == [] + + +def test_replicate_session_clears_queue_when_asset_init_fails(sim): + """ReplicateSession.__exit__ drops queued cfgs if the asset constructor body raises.""" + from isaaclab.cloner import ReplicateSession + + leaked_cfg = SimpleNamespace(prim_path="/World/envs/env_.*/Robot") + + sentinel = MagicMock() + sentinel_cls = MagicMock(return_value=sentinel) + + with pytest.raises(RuntimeError, match="asset boom"): + with ReplicateSession( + cfgs=[], + num_clones=2, + env_spacing=1.0, + device=sim.cfg.device, + stage=sim_utils.get_current_stage(), + ): + REPLICATION_QUEUE.append((leaked_cfg, sentinel_cls)) + raise RuntimeError("asset boom") + + assert REPLICATION_QUEUE == [] + sentinel_cls.assert_not_called() diff --git a/source/isaaclab_newton/changelog.d/replication-session-redesign.skip b/source/isaaclab_newton/changelog.d/replication-session-redesign.skip new file mode 100644 index 000000000000..4ff1bbe80281 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/replication-session-redesign.skip @@ -0,0 +1,5 @@ +Internal refactor: ``queue_newton_physics_replication`` now appends ``(cfg, +NewtonReplicateContext)`` to :data:`isaaclab.cloner.REPLICATION_QUEUE` instead of +reaching into the now-removed module-level replication session. No user-visible +behavior change — :func:`~isaaclab.cloner.replicate` still drives Newton replication +via :class:`NewtonReplicateContext` exactly as before. diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index f921001cb8cf..0628c78d0a62 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -38,6 +38,7 @@ from isaaclab_newton.assets import kernels as shared_kernels from isaaclab_newton.assets.articulation import kernels as articulation_kernels +from isaaclab_newton.cloner import queue_newton_physics_replication from isaaclab_newton.physics import NewtonManager as SimulationManager from .articulation_data import ArticulationData @@ -121,6 +122,7 @@ def __init__(self, cfg: ArticulationCfg): from isaaclab.sim import SimulationContext # noqa: PLC0415 super().__init__(cfg) + queue_newton_physics_replication(cfg) sim_ctx = SimulationContext.instance() self._sim_cfg = sim_ctx.cfg if sim_ctx is not None else None diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py index b93c9075393d..a985e0256fe2 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py @@ -24,6 +24,7 @@ from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_newton.assets import kernels as shared_kernels +from isaaclab_newton.cloner import queue_newton_physics_replication from isaaclab_newton.physics import NewtonManager as SimulationManager from .rigid_object_data import RigidObjectData @@ -59,6 +60,7 @@ def __init__(self, cfg: RigidObjectCfg): cfg: A configuration instance. """ super().__init__(cfg) + queue_newton_physics_replication(cfg) """ Properties diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py index b11415d48231..7b11629d4cd6 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py @@ -25,6 +25,7 @@ from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_newton.assets import kernels as shared_kernels +from isaaclab_newton.cloner import queue_newton_physics_replication from isaaclab_newton.physics import NewtonManager as SimulationManager from .rigid_object_collection_data import RigidObjectCollectionData @@ -75,7 +76,8 @@ def __init__(self, cfg: RigidObjectCollectionCfg): # flag for whether the asset is initialized self._is_initialized = False # spawn the rigid objects - for rigid_body_cfg in self.cfg.rigid_objects.values(): + source_rigid_object_cfgs = cfg.rigid_objects + for rigid_body_name, rigid_body_cfg in self.cfg.rigid_objects.items(): # spawn the asset if rigid_body_cfg.spawn is not None: spawn_path = rigid_body_cfg.spawn.spawn_path or rigid_body_cfg.prim_path @@ -89,6 +91,7 @@ def __init__(self, cfg: RigidObjectCollectionCfg): matching_prims = sim_utils.find_matching_prims(rigid_body_cfg.prim_path) if len(matching_prims) == 0: raise RuntimeError(f"Could not find prim with path {rigid_body_cfg.prim_path}.") + queue_newton_physics_replication(source_rigid_object_cfgs[rigid_body_name]) # stores object names self._body_names_list = [] diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/cloner/__init__.pyi index f55377295a3e..8bb3102fa395 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/__init__.pyi +++ b/source/isaaclab_newton/isaaclab_newton/cloner/__init__.pyi @@ -4,8 +4,13 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ + "NewtonReplicateContext", "newton_physics_replicate", - "newton_visualizer_prebuild", + "queue_newton_physics_replication", ] -from .newton_replicate import newton_physics_replicate, newton_visualizer_prebuild +from .replicate import ( + NewtonReplicateContext, + newton_physics_replicate, + queue_newton_physics_replication, +) diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py b/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py similarity index 67% rename from source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py rename to source/isaaclab_newton/isaaclab_newton/cloner/replicate.py index e99b1eb7abdd..50bf1d049d26 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py @@ -6,6 +6,7 @@ from __future__ import annotations from collections.abc import Sequence +from typing import Any import torch import warp as wp @@ -14,6 +15,8 @@ from pxr import Usd +from isaaclab.cloner.replicate_session import REPLICATION_QUEUE + from isaaclab_newton.physics import NewtonManager @@ -256,6 +259,146 @@ def _rename_pair(values, worlds): _rename_pair(values, worlds) +class NewtonReplicateContext: + """Queue and run Newton replication work for one stage.""" + + def __init__( + self, + stage: Usd.Stage, + *, + device: str = "cpu", + up_axis: str = "Z", + simplify_meshes: bool = True, + 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. + simplify_meshes: Whether to run convex-hull mesh approximation. + commit_to_manager: Whether :meth:`replicate` should publish the builder to + :class:`NewtonManager`. + """ + self.stage = stage + self.device = device + self.up_axis = up_axis + self.simplify_meshes = simplify_meshes + self.commit_to_manager = commit_to_manager + self._queue: list[ + tuple[ + tuple[str, ...], + tuple[str, ...], + torch.Tensor, + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + ] + ] = [] + + 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, + ) -> tuple[tuple[str, ...], tuple[str, ...], torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """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 = _build_newton_builder_from_mapping( + stage=self.stage, + sources=sources, + env_ids=env_ids, + mapping=mapping, + positions=positions, + quaternions=quaternions, + up_axis=self.up_axis, + simplify_meshes=self.simplify_meshes, + ) + _rename_builder_labels(builder, sources, destinations, env_ids, mapping) + if self.commit_to_manager: + NewtonManager._cl_site_index_map = site_index_map + NewtonManager.set_builder(builder) + NewtonManager._num_envs = mapping.size(1) + self._queue.clear() + return builder, stage_info, site_index_map + + +def queue_newton_physics_replication(cfg: Any) -> None: + """Register ``cfg`` for Newton replication when :func:`~isaaclab.cloner.replicate` next runs. + + Appends ``(cfg, NewtonReplicateContext)`` 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, NewtonReplicateContext)) + + def newton_physics_replicate( stage: Usd.Stage, sources: Sequence[str], @@ -285,66 +428,16 @@ def newton_physics_replicate( Returns: Tuple of the populated Newton model builder and stage metadata. """ - builder, stage_info, site_index_map = _build_newton_builder_from_mapping( - stage=stage, - sources=sources, - env_ids=env_ids, - mapping=mapping, - positions=positions, - quaternions=quaternions, - up_axis=up_axis, - simplify_meshes=simplify_meshes, + ctx = NewtonReplicateContext( + stage, device=device, up_axis=up_axis, simplify_meshes=simplify_meshes, commit_to_manager=True ) - _rename_builder_labels(builder, sources, destinations, env_ids, mapping) - NewtonManager._cl_site_index_map = site_index_map - NewtonManager.set_builder(builder) - NewtonManager._num_envs = mapping.size(1) - return builder, stage_info - - -def newton_visualizer_prebuild( - 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", - simplify_meshes: bool = True, -): - """Replicate a clone plan into a finalized Newton model/state for visualization. - - Unlike :func:`newton_physics_replicate`, this path does not mutate ``NewtonManager`` and is intended - for prebuilding visualizer-only artifacts that can be consumed by scene data providers. - - 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. - up_axis: Up axis for the Newton model builder. - simplify_meshes: Whether to run convex-hull mesh approximation. - - Returns: - Tuple of finalized Newton model and state. - """ - builder, _, _site_index_map = _build_newton_builder_from_mapping( - stage=stage, - sources=sources, - env_ids=env_ids, - mapping=mapping, + ctx.queue_mapping( + sources, + destinations, + env_ids, + mapping, positions=positions, quaternions=quaternions, - up_axis=up_axis, - simplify_meshes=simplify_meshes, ) - _rename_builder_labels(builder, sources, destinations, env_ids, mapping) - model = builder.finalize(device=device) - state = model.state() - return model, state + builder, stage_info, _site_index_map = ctx.replicate() + return builder, stage_info diff --git a/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py b/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py index 5ecf162fbcab..f9913d58ec95 100644 --- a/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py +++ b/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py @@ -21,7 +21,7 @@ import newton import torch -from isaaclab_newton.cloner.newton_replicate import _BUILTIN_LABEL_TYPES, _rename_builder_labels +from isaaclab_newton.cloner.replicate import _BUILTIN_LABEL_TYPES, _rename_builder_labels from newton.solvers import SolverMuJoCo _TENDON_FREQ = "mujoco:tendon" diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py index fdb836fc08b9..9e5726e173fc 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py @@ -23,6 +23,7 @@ import isaaclab.sim as sim_utils from isaaclab.assets.articulation.articulation_cfg import ArticulationCfg from isaaclab.assets.articulation.base_articulation import BaseArticulation +from isaaclab.cloner import queue_usd_replication from isaaclab.physics import PhysicsManager from isaaclab.utils.string import resolve_matching_names from isaaclab.utils.wrench_composer import WrenchComposer @@ -30,6 +31,7 @@ from isaaclab_ovphysx import tensor_types as TT from isaaclab_ovphysx.assets import kernels as shared_kernels from isaaclab_ovphysx.assets.kernels import _body_wrench_to_world +from isaaclab_ovphysx.cloner import queue_ovphysx_replication from isaaclab_ovphysx.physics import OvPhysxManager from .articulation_data import ArticulationData @@ -84,6 +86,8 @@ def __init__(self, cfg: ArticulationCfg): cfg: A configuration instance. """ super().__init__(cfg) + queue_usd_replication(cfg) + queue_ovphysx_replication(cfg) # bindings are populated eagerly in ``_initialize_impl``; the dict # also caches any tensor type the user explicitly queries later self._bindings: dict[int, Any] = {} diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/rigid_object.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/rigid_object.py index 015c8f102f44..ca5f8e048364 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/rigid_object.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/rigid_object.py @@ -21,12 +21,14 @@ import isaaclab.sim as sim_utils from isaaclab.assets.rigid_object.base_rigid_object import BaseRigidObject from isaaclab.assets.rigid_object.rigid_object_cfg import RigidObjectCfg +from isaaclab.cloner import queue_usd_replication from isaaclab.utils.string import resolve_matching_names from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_ovphysx import tensor_types as TT from isaaclab_ovphysx.assets import kernels as shared_kernels from isaaclab_ovphysx.assets.kernels import _body_wrench_to_world +from isaaclab_ovphysx.cloner import queue_ovphysx_replication from isaaclab_ovphysx.physics import OvPhysxManager from .rigid_object_data import RigidObjectData @@ -62,6 +64,8 @@ def __init__(self, cfg: RigidObjectCfg): cfg: A configuration instance. """ super().__init__(cfg) + queue_usd_replication(cfg) + queue_ovphysx_replication(cfg) # Bindings are created lazily (on first access) to avoid allocating # handles for tensor types the user never queries. self._bindings: dict[int, Any] = {} diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/rigid_object_collection.py index 7a1d4e21e535..b2f7b8bcfca4 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/rigid_object_collection.py @@ -18,12 +18,14 @@ import isaaclab.sim as sim_utils from isaaclab.assets.rigid_object_collection.base_rigid_object_collection import BaseRigidObjectCollection +from isaaclab.cloner import queue_usd_replication from isaaclab.utils.string import resolve_matching_names from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_ovphysx import tensor_types as TT from isaaclab_ovphysx.assets import kernels as shared_kernels from isaaclab_ovphysx.assets.kernels import _body_wrench_to_world, resolve_view_ids +from isaaclab_ovphysx.cloner import queue_ovphysx_replication from isaaclab_ovphysx.physics import OvPhysxManager from .rigid_object_collection_data import RigidObjectCollectionData @@ -74,7 +76,7 @@ def __init__(self, cfg: RigidObjectCollectionCfg): # flag for whether the asset is initialized self._is_initialized = False # spawn the rigid objects - for rigid_body_cfg in self.cfg.rigid_objects.values(): + for rigid_body_name, rigid_body_cfg in self.cfg.rigid_objects.items(): # spawn the asset if rigid_body_cfg.spawn is not None: rigid_body_cfg.spawn.func( @@ -87,6 +89,8 @@ def __init__(self, cfg: RigidObjectCollectionCfg): matching_prims = sim_utils.find_matching_prims(rigid_body_cfg.prim_path) if len(matching_prims) == 0: raise RuntimeError(f"Could not find prim with path {rigid_body_cfg.prim_path}.") + queue_usd_replication(cfg.rigid_objects[rigid_body_name]) + queue_ovphysx_replication(cfg.rigid_objects[rigid_body_name]) # stores object names self._body_names_list: list[str] = [] # one fused TensorBinding per tensor type, populated in _initialize_impl diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/cloner/__init__.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/cloner/__init__.py index 3b9a12007792..34358e402794 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/cloner/__init__.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/cloner/__init__.py @@ -3,6 +3,6 @@ # # SPDX-License-Identifier: BSD-3-Clause -from .ovphysx_replicate import ovphysx_replicate +from .replicate import OvPhysxReplicateContext, ovphysx_replicate, queue_ovphysx_replication -__all__ = ["ovphysx_replicate"] +__all__ = ["OvPhysxReplicateContext", "ovphysx_replicate", "queue_ovphysx_replication"] diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/cloner/ovphysx_replicate.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/cloner/ovphysx_replicate.py deleted file mode 100644 index d89a45280a50..000000000000 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/cloner/ovphysx_replicate.py +++ /dev/null @@ -1,105 +0,0 @@ -# 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 - -"""OvPhysX replication hook for IsaacLab's cloning pipeline. - -Called from the scene cloning path in place of immediate PhysX or Newton -replication. Unlike those replicators, ovphysx.PhysX does not exist yet at -this point in the scene setup — it is created lazily on the first -:meth:`~isaaclab_ovphysx.physics.OvPhysxManager.reset` call. - -This function records a *pending clone* on :class:`OvPhysxManager`. When -:meth:`~isaaclab_ovphysx.physics.OvPhysxManager._warmup_and_load` eventually -creates the ``PhysX`` instance and loads the USD stage (which contains only -``env_0`` physics — env_1..N are empty Xform containers), it replays every -pending clone via ``physx.clone(source, targets)`` to create the remaining -environments entirely inside the physics runtime without touching USD. -""" - -from __future__ import annotations - -from collections.abc import Sequence - -import torch - -from pxr import Usd - - -def ovphysx_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", -) -> None: - """Record a physics clone for later execution by OvPhysxManager. - - Translates the generic IsaacLab source/destination/mapping representation - into ``(source_path, [target_paths])`` pairs and registers them on - :class:`~isaaclab_ovphysx.physics.OvPhysxManager`. The actual - ``physx.clone()`` calls happen in ``_warmup_and_load()`` after the USD - stage has been loaded. - - The ``positions`` parameter contains the 2-D grid world positions for all - environments. They are forwarded to the C++ clone plugin so that the - parent Xform prim for each clone (e.g. ``/World/envs/env_N``) is placed at - the correct grid location in Fabric. The exported USD stage only contains - ``env_0``; without explicit positions all clone parents would be created at - the origin, causing all articulations to pile up and the GPU solver to - diverge on the first warmup step. - - Args: - stage: USD stage (not modified by this function). - sources: Source prim paths (one per prototype). - destinations: Destination path templates with ``"{}"`` for env index. - env_ids: Environment indices tensor. - mapping: ``(num_sources, num_envs)`` bool tensor; True selects which - environments receive each source. - positions: World (x, y, z) positions for every environment, shape - ``[num_envs, 3]``. Used to place clone parent Xform prims in - Fabric at correct grid locations. - quaternions: Ignored — orientations are set at first reset. - device: Torch device (unused; kept for API compatibility). - """ - # Deferred import to avoid circular dependency at module load time. - from isaaclab_ovphysx.physics.ovphysx_manager import OvPhysxManager - - for i, src in enumerate(sources): - active_env_ids = env_ids[mapping[i]].tolist() - - # Exclude the source environment from its own target list. - # physx.clone() is only needed for *other* envs; the source env_0 is - # already loaded from USD. We detect self by matching the source path - # against the destination template. - pre, _, suf = destinations[i].partition("{}") - self_env_id: int | None = None - candidate = src.removeprefix(pre).removesuffix(suf) - if candidate.isdigit(): - self_env_id = int(candidate) - - # Build parallel (targets, parent_positions) lists for non-self envs. - # parent_positions[j] is the world (x,y,z) for the parent Xform of - # targets[j] (e.g. /World/envs/env_N). These positions are passed to - # the C++ clone plugin so that env_N Xform prims — absent from the - # exported USD stage — are created at the correct 2-D grid location - # rather than the origin. Without this, all clones pile up at env_0's - # position during the warmup physics step and the GPU solver diverges. - targets: list[str] = [] - parent_positions: list[tuple[float, float, float]] = [] - for e in active_env_ids: - if e == self_env_id: - continue - targets.append(destinations[i].format(e)) - if positions is not None and e < len(positions): - pos = positions[e] - parent_positions.append((float(pos[0]), float(pos[1]), float(pos[2]))) - else: - parent_positions.append((0.0, 0.0, 0.0)) - - if targets: - OvPhysxManager.register_clone(src, targets, parent_positions) diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/cloner/replicate.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/cloner/replicate.py new file mode 100644 index 000000000000..d1fbe451ea8f --- /dev/null +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/cloner/replicate.py @@ -0,0 +1,185 @@ +# 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 + +"""OvPhysX replication hook for IsaacLab's cloning pipeline. + +Called from the scene cloning path in place of immediate PhysX or Newton +replication. Unlike those replicators, ovphysx.PhysX does not exist yet at +this point in the scene setup — it is created lazily on the first +:meth:`~isaaclab_ovphysx.physics.OvPhysxManager.reset` call. + +This function records a *pending clone* on :class:`OvPhysxManager`. When +:meth:`~isaaclab_ovphysx.physics.OvPhysxManager._warmup_and_load` eventually +creates the ``PhysX`` instance and loads the USD stage (which contains only +``env_0`` physics — env_1..N are empty Xform containers), it replays every +pending clone via ``physx.clone(source, targets)`` to create the remaining +environments entirely inside the physics runtime without touching USD. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import torch + +from pxr import Sdf, Usd + +from isaaclab.cloner.replicate_session import REPLICATION_QUEUE + + +def _select_env_ids(env_ids: torch.Tensor, mapping: torch.Tensor, row: int) -> torch.Tensor: + """Return the environment ids selected by a replication row.""" + row_mask = mapping[row] + if row_mask.dtype != torch.bool: + row_mask = row_mask.to(dtype=torch.bool) + return env_ids[row_mask] + + +class OvPhysxReplicateContext: + """Queue and run OvPhysX clone operations for one stage.""" + + def __init__(self, stage: Usd.Stage): + """Initialize the context. + + Args: + stage: USD stage associated with the pending clone operations. + """ + self.stage = stage + physics_scene_prim = self.stage.GetPrimAtPath("/physicsScene") + if physics_scene_prim.IsValid(): + physics_scene_prim.CreateAttribute("physxScene:envIdInBoundsBitCount", Sdf.ValueTypeNames.Int).Set(4) + self._queue: list[tuple[str, list[str], list[tuple[float, float, float]]]] = [] + + def queue( + self, source: str, targets: Sequence[str], parent_positions: Sequence[tuple[float, float, float]] + ) -> None: + """Queue one pending OvPhysX clone operation. + + Args: + source: Source prim path. + targets: Destination prim paths. + parent_positions: Parent Xform positions [m] for each destination. + """ + self._queue.append((source, list(targets), list(parent_positions))) + + 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 clone operations from the current flat clone mapping. + + Args: + sources: Source prim paths. + destinations: Destination path templates with ``"{}"`` for env id. + env_ids: Environment indices. + mapping: Bool/int mask selecting envs per source. + positions: Optional per-environment world positions [m]. + quaternions: Optional per-environment orientations, unused by OvPhysX. + """ + del quaternions + + for i, src in enumerate(sources): + active_env_ids = _select_env_ids(env_ids, mapping, i).tolist() + + pre, _, suf = destinations[i].partition("{}") + self_env_id: int | None = None + candidate = src.removeprefix(pre).removesuffix(suf) + if candidate.isdigit(): + self_env_id = int(candidate) + + targets: list[str] = [] + parent_positions: list[tuple[float, float, float]] = [] + for env_id in active_env_ids: + env_id = int(env_id) + if env_id == self_env_id: + continue + targets.append(destinations[i].format(env_id)) + if positions is not None and env_id < len(positions): + pos = positions[env_id] + parent_positions.append((float(pos[0]), float(pos[1]), float(pos[2]))) + else: + parent_positions.append((0.0, 0.0, 0.0)) + + if targets: + self.queue(src, targets, parent_positions) + + def replicate(self) -> None: + """Publish all queued clones to :class:`OvPhysxManager`.""" + from isaaclab_ovphysx.physics.ovphysx_manager import OvPhysxManager + + for source, targets, parent_positions in self._queue: + OvPhysxManager.register_clone(source, targets, parent_positions) + self._queue.clear() + + +def queue_ovphysx_replication(cfg: Any) -> None: + """Register ``cfg`` for OvPhysX replication when :func:`~isaaclab.cloner.replicate` next runs. + + Appends ``(cfg, OvPhysxReplicateContext)`` 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, OvPhysxReplicateContext)) + + +def ovphysx_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", +) -> None: + """Record a physics clone for later execution by OvPhysxManager. + + Translates the generic IsaacLab source/destination/mapping representation + into ``(source_path, [target_paths])`` pairs and registers them on + :class:`~isaaclab_ovphysx.physics.OvPhysxManager`. The actual + ``physx.clone()`` calls happen in ``_warmup_and_load()`` after the USD + stage has been loaded. + + The ``positions`` parameter contains the 2-D grid world positions for all + environments. They are forwarded to the C++ clone plugin so that the + parent Xform prim for each clone (e.g. ``/World/envs/env_N``) is placed at + the correct grid location in Fabric. The exported USD stage only contains + ``env_0``; without explicit positions all clone parents would be created at + the origin, causing all articulations to pile up and the GPU solver to + diverge on the first warmup step. + + Args: + stage: USD stage (not modified by this function). + sources: Source prim paths (one per prototype). + destinations: Destination path templates with ``"{}"`` for env index. + env_ids: Environment indices tensor. + mapping: ``(num_sources, num_envs)`` bool tensor; True selects which + environments receive each source. + positions: World (x, y, z) positions [m] for every environment, shape + ``[num_envs, 3]``. Used to place clone parent Xform prims in + Fabric at correct grid locations. + quaternions: Ignored — orientations are set at first reset. + device: Torch device (unused; kept for API compatibility). + """ + del device + + ctx = OvPhysxReplicateContext(stage) + ctx.queue_mapping( + sources, + destinations, + env_ids, + mapping, + positions=positions, + quaternions=quaternions, + ) + ctx.replicate() diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/physics/ovphysx_manager.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/physics/ovphysx_manager.py index 8f7e1e576f17..8fc63190a795 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/physics/ovphysx_manager.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/physics/ovphysx_manager.py @@ -211,8 +211,8 @@ class OvPhysxManager(PhysicsManager): # :meth:`_release_physx`); we mirror it here so a clear Python error is raised # if a later :class:`~isaaclab.sim.SimulationContext` requests a different device. _locked_device: ClassVar[str | None] = None - # Pending (source, targets, parent_positions) triples registered by - # ovphysx_replicate() before the PhysX instance exists. Replayed via + # Pending (source, targets, parent_positions) triples queued by + # queue_ovphysx_replication() before the PhysX instance exists. Replayed via # physx.clone() in _warmup_and_load(). # parent_positions is a list of (x, y, z) tuples — one per target. _pending_clones: ClassVar[list[tuple[str, list[str], list[tuple[float, float, float]]]]] = [] @@ -572,10 +572,9 @@ def _warmup_and_load(cls) -> None: # Xform containers. physx.clone() creates the remaining environments # in the physics runtime without modifying the USD file. if cls._pending_clones: - # ovphysx_replicate() only registers pending clones when clone_usd=False, - # meaning the USD contains only env_0 physics and physx.clone() is required - # to populate env_1..N in the physics runtime. Execute unconditionally — - # no USD content heuristic is needed. + # The cfg-level OvPhysX replicator registers pending clones for physics + # regardless of whether USD copies were also queued for rendering. Execute + # unconditionally — no USD content heuristic is needed. for source, targets, parent_positions in cls._pending_clones: logger.info( "OvPhysxManager: cloning %s -> %d targets (%s ... %s)", diff --git a/source/isaaclab_physx/changelog.d/replication-session-redesign.skip b/source/isaaclab_physx/changelog.d/replication-session-redesign.skip new file mode 100644 index 000000000000..8b9f8ee958aa --- /dev/null +++ b/source/isaaclab_physx/changelog.d/replication-session-redesign.skip @@ -0,0 +1,5 @@ +Internal refactor: ``queue_physx_replication`` now appends ``(cfg, +PhysxReplicateContext)`` to :data:`isaaclab.cloner.REPLICATION_QUEUE` instead of +reaching into the now-removed module-level replication session. No user-visible +behavior change — :func:`~isaaclab.cloner.replicate` still drives PhysX replication +via :class:`PhysxReplicateContext` exactly as before. diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 78fdb387aa04..997ce9f12e33 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -23,6 +23,7 @@ from isaaclab.actuators import ActuatorBase, ActuatorBaseCfg, ImplicitActuator from isaaclab.assets.articulation.base_articulation import BaseArticulation +from isaaclab.cloner import queue_usd_replication _HAS_NEWTON_ACTUATORS = importlib.util.find_spec("isaaclab_newton.actuators") is not None @@ -35,6 +36,7 @@ from isaaclab_physx.assets import kernels as shared_kernels from isaaclab_physx.assets.articulation import kernels as articulation_kernels +from isaaclab_physx.cloner import queue_physx_replication from isaaclab_physx.physics import PhysxManager as SimulationManager from .articulation_data import ArticulationData @@ -121,6 +123,8 @@ def __init__(self, cfg: ArticulationCfg): from isaaclab.sim import SimulationContext # noqa: PLC0415 super().__init__(cfg) + queue_usd_replication(cfg) + queue_physx_replication(cfg) sim_ctx = SimulationContext.instance() self._sim_cfg = sim_ctx.cfg if sim_ctx is not None else None diff --git a/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py b/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py index 759e31d9fb87..59d0d6641567 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py @@ -20,9 +20,11 @@ import isaaclab.sim as sim_utils import isaaclab.utils.math as math_utils from isaaclab.assets.asset_base import AssetBase +from isaaclab.cloner import queue_usd_replication from isaaclab.markers import VisualizationMarkers from isaaclab.utils.warp import ProxyArray +from isaaclab_physx.cloner import queue_physx_replication from isaaclab_physx.physics import PhysxManager as SimulationManager from .deformable_object_data import DeformableObjectData @@ -79,6 +81,8 @@ def __init__(self, cfg: DeformableObjectCfg): cfg: A configuration instance. """ super().__init__(cfg) + queue_usd_replication(cfg) + queue_physx_replication(cfg) # Register custom vec6f type for nodal state validation. self._DTYPE_TO_TORCH_TRAILING_DIMS = {**self._DTYPE_TO_TORCH_TRAILING_DIMS, vec6f: (6,)} # initialize deformable type to None, should be set to either surface or volume on initialization diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py index 930b8836859a..7f7c90dcb909 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py @@ -18,9 +18,11 @@ import isaaclab.sim as sim_utils import isaaclab.utils.string as string_utils from isaaclab.assets.rigid_object.base_rigid_object import BaseRigidObject +from isaaclab.cloner import queue_usd_replication from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_physx.assets import kernels as shared_kernels +from isaaclab_physx.cloner import queue_physx_replication from isaaclab_physx.physics import PhysxManager as SimulationManager from .rigid_object_data import RigidObjectData @@ -65,6 +67,8 @@ def __init__(self, cfg: RigidObjectCfg): cfg: A configuration instance. """ super().__init__(cfg) + queue_usd_replication(cfg) + queue_physx_replication(cfg) """ Properties diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py index 2031ded53b2f..c7072244d168 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py @@ -21,9 +21,11 @@ import isaaclab.sim as sim_utils import isaaclab.utils.string as string_utils from isaaclab.assets.rigid_object_collection.base_rigid_object_collection import BaseRigidObjectCollection +from isaaclab.cloner import queue_usd_replication from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_physx.assets import kernels as shared_kernels +from isaaclab_physx.cloner import queue_physx_replication from isaaclab_physx.physics import PhysxManager as SimulationManager from .kernels import resolve_view_ids @@ -78,7 +80,8 @@ def __init__(self, cfg: RigidObjectCollectionCfg): # flag for whether the asset is initialized self._is_initialized = False # spawn the rigid objects - for rigid_body_cfg in self.cfg.rigid_objects.values(): + source_rigid_object_cfgs = cfg.rigid_objects + for rigid_body_name, rigid_body_cfg in self.cfg.rigid_objects.items(): # spawn the asset if rigid_body_cfg.spawn is not None: spawn_path = rigid_body_cfg.spawn.spawn_path or rigid_body_cfg.prim_path @@ -92,6 +95,8 @@ def __init__(self, cfg: RigidObjectCollectionCfg): matching_prims = sim_utils.find_matching_prims(rigid_body_cfg.prim_path) if len(matching_prims) == 0: raise RuntimeError(f"Could not find prim with path {rigid_body_cfg.prim_path}.") + queue_usd_replication(source_rigid_object_cfgs[rigid_body_name]) + queue_physx_replication(source_rigid_object_cfgs[rigid_body_name]) # stores object names self._body_names_list = [] diff --git a/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper.py b/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper.py index 6662582dac7c..25b0e31722c1 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper.py @@ -17,8 +17,11 @@ import isaaclab.sim as sim_utils from isaaclab.assets import AssetBase +from isaaclab.cloner import queue_usd_replication from isaaclab.utils.version import get_isaac_sim_version, has_kit +from isaaclab_physx.cloner import queue_physx_replication + if TYPE_CHECKING: from isaacsim.robot.surface_gripper import GripperView @@ -98,6 +101,9 @@ def __init__(self, cfg: SurfaceGripperCfg): self._is_initialized = False self._debug_vis_handle = None + queue_usd_replication(cfg) + queue_physx_replication(cfg) + # register various callback functions self._register_callbacks() diff --git a/source/isaaclab_physx/isaaclab_physx/cloner/__init__.pyi b/source/isaaclab_physx/isaaclab_physx/cloner/__init__.pyi index ecc5b6e3923e..d865b4f56e80 100644 --- a/source/isaaclab_physx/isaaclab_physx/cloner/__init__.pyi +++ b/source/isaaclab_physx/isaaclab_physx/cloner/__init__.pyi @@ -4,7 +4,9 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ + "PhysxReplicateContext", "physx_replicate", + "queue_physx_replication", ] -from .physx_replicate import physx_replicate +from .replicate import PhysxReplicateContext, physx_replicate, queue_physx_replication diff --git a/source/isaaclab_physx/isaaclab_physx/cloner/physx_replicate.py b/source/isaaclab_physx/isaaclab_physx/cloner/physx_replicate.py deleted file mode 100644 index dcc5cc6d9677..000000000000 --- a/source/isaaclab_physx/isaaclab_physx/cloner/physx_replicate.py +++ /dev/null @@ -1,113 +0,0 @@ -# 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 - -import torch - -from omni.physx import get_physx_replicator_interface -from pxr import Usd, UsdUtils - - -def physx_replicate( - stage: Usd.Stage, - sources: Sequence[str], # e.g. ["/World/Template/A", "/World/Template/B"] - destinations: Sequence[str], # e.g. ["/World/envs/env_{}/Robot", "/World/envs/env_{}/Object"] - env_ids: torch.Tensor, # env_ids - mapping: torch.Tensor, # (num_sources, num_envs) bool; True -> place sources[i] into world=j - positions: torch.Tensor | None = None, - quaternions: torch.Tensor | None = None, - use_fabric: bool = False, - device: str = "cpu", - exclude_self_replication: bool = True, -) -> None: - """Replicate prims via PhysX replicator with per-row mapping. - - Builds per-source destination lists from ``mapping`` and calls PhysX ``replicate``. - Rows covering all environments use ``useEnvIds=True``; partial rows use ``False``. - The replicator is registered for the call and then unregistered. - - ``attach_fn`` excludes ``/World/template`` and ``/World/envs`` so that PhysX does - not independently parse prims that the replicator will handle. The source prim - receives its physics body as a side-effect of ``rep.replicate()`` (which always - parses the source internally), so every source must appear in at least one - ``replicate`` call. - - When ``exclude_self_replication`` is True (default), each source environment is - removed from its own replication targets so the replicator only creates bodies at - non-self destinations. If removing self would leave the world list empty (i.e. the - source maps only to its own environment), self is kept so that ``rep.replicate()`` - is still called and the source prim gets its physics body. - - Args: - stage: USD stage. - sources: Source prim paths (``S``). - destinations: Destination templates (``S``) with ``"{}"`` for env index. - env_ids: Environment indices (``[E]``). - mapping: Bool/int mask (``[S, E]``) selecting envs per source. - positions: Optional positions (unused, for API compatibility). - quaternions: Optional orientations (unused, for API compatibility). - use_fabric: Use Fabric for replication. - device: Torch device for determining replication mode. - exclude_self_replication: If True, skip replicating a source prim onto itself - when the source also maps to other environments. Default is True. - Self-only sources always keep self so that ``rep.replicate()`` fires. - - Returns: - None - """ - # Note: positions and quaternions are unused by PhysX replicator - # They are included for API compatibility with other backends (e.g., Newton) - del positions, quaternions - - stage_id = UsdUtils.StageCache.Get().Insert(stage).ToLongInt() - current_worlds: list[int] = [] - current_template: str = "" - num_envs = mapping.size(1) - - if num_envs > 1: - # Pre-compute effective world lists after self-exclusion. - # Self is only removed when the source also maps to other environments; - # if it is the sole destination we must keep it so that rep.replicate() - # is still called (the source gets its physics body from that call). - effective_worlds: list[list[int]] = [] - for i, src in enumerate(sources): - worlds = env_ids[mapping[i]].tolist() - if exclude_self_replication: - pre, _, suf = destinations[i].partition("{}") - self_id = src.removeprefix(pre).removesuffix(suf) - if self_id.isdigit(): - filtered = [w for w in worlds if w != int(self_id)] - worlds = filtered if filtered else worlds - effective_worlds.append(worlds) - - def attach_fn(_stage_id: int): - return ["/World/template", "/World/envs"] - - 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 - rep = get_physx_replicator_interface() - for i, src in enumerate(sources): - current_template = destinations[i] - current_worlds[:] = effective_worlds[i] - if not current_worlds: - continue - rep.replicate( - _stage_id, - src, - len(current_worlds), - # TODO: envIds needs to support heterogeneous setup. for now, we rely on USD collision filtering - useEnvIds=False, # (len(current_worlds) == num_envs - 1) and device != "cpu", - useFabricForReplication=use_fabric, - ) - # unregister only AFTER all replicate() calls completed - rep.unregister_replicator(_stage_id) - - get_physx_replicator_interface().register_replicator(stage_id, attach_fn, attach_end_fn, rename_fn) diff --git a/source/isaaclab_physx/isaaclab_physx/cloner/replicate.py b/source/isaaclab_physx/isaaclab_physx/cloner/replicate.py new file mode 100644 index 000000000000..9fa00e8c3b75 --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/cloner/replicate.py @@ -0,0 +1,196 @@ +# 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 Any + +import torch + +from omni.physx import get_physx_replicator_interface +from pxr import Sdf, Usd, UsdUtils + +from isaaclab.cloner.replicate_session import REPLICATION_QUEUE + + +def _select_env_ids(env_ids: torch.Tensor, mapping: torch.Tensor, row: int) -> torch.Tensor: + """Return the environment ids selected by a replication row.""" + row_mask = mapping[row] + if row_mask.dtype != torch.bool: + row_mask = row_mask.to(dtype=torch.bool) + return env_ids[row_mask] + + +class PhysxReplicateContext: + """Queue and run PhysX replication work for one stage.""" + + def __init__(self, stage: Usd.Stage): + """Initialize the context. + + Args: + stage: USD stage to register with the PhysX replicator. + """ + self.stage = stage + self._stage_id = UsdUtils.StageCache.Get().Insert(stage).ToLongInt() + physics_scene_prim = self.stage.GetPrimAtPath("/physicsScene") + if physics_scene_prim.IsValid(): + physics_scene_prim.CreateAttribute("physxScene:envIdInBoundsBitCount", Sdf.ValueTypeNames.Int).Set(4) + self._queue: list[tuple[str, str, tuple[int, ...]]] = [] + + def queue(self, source: str, destination: str, target_envs: Sequence[int]) -> None: + """Queue one PhysX source row for replication. + + Args: + source: Source prim path. + destination: Destination path template with ``"{}"`` for env id. + target_envs: Environment ids selected for this source row. + """ + self._queue.append((source, destination, tuple(int(env_id) for env_id in target_envs))) + + 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, + exclude_self_replication: bool = True, + ) -> None: + """Queue replication rows from the current flat clone mapping. + + Args: + sources: Source prim paths. + destinations: Destination path templates with ``"{}"`` for env id. + env_ids: Environment indices. + mapping: Bool/int mask selecting envs per source. + positions: Optional per-environment world positions [m], unused by PhysX. + quaternions: Optional per-environment orientations, unused by PhysX. + exclude_self_replication: Whether to skip replicating a source prim onto itself + when it also maps to other environments. + """ + del positions, quaternions + + if mapping.size(1) <= 1: + return + + for i, src in enumerate(sources): + worlds = _select_env_ids(env_ids, mapping, i).tolist() + if exclude_self_replication: + pre, _, suf = destinations[i].partition("{}") + self_id = src.removeprefix(pre).removesuffix(suf) + if self_id.isdigit(): + filtered = [w for w in worlds if w != int(self_id)] + worlds = filtered if filtered else worlds + self.queue(src, destinations[i], worlds) + + def replicate(self) -> None: + """Register the PhysX replicator and run queued rows from ``attach_end_fn``.""" + if not self._queue: + return + + physx_queue = tuple(self._queue) + current_worlds: list[int] = [] + current_template: str = "" + + def attach_fn(_stage_id: int): + return ["/World/template", "/World/envs"] + + 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 + rep = get_physx_replicator_interface() + for src, destination, target_envs in physx_queue: + current_template = destination + current_worlds[:] = target_envs + if not current_worlds: + continue + rep.replicate( + _stage_id, + src, + len(current_worlds), + # TODO: envIds needs to support heterogeneous setup. for now, we rely on USD collision filtering + useEnvIds=False, + useFabricForReplication=False, + ) + rep.unregister_replicator(_stage_id) + + get_physx_replicator_interface().register_replicator(self._stage_id, attach_fn, attach_end_fn, rename_fn) + self._queue.clear() + + +def queue_physx_replication(cfg: Any) -> None: + """Register ``cfg`` for PhysX replication when :func:`~isaaclab.cloner.replicate` next runs. + + Appends ``(cfg, PhysxReplicateContext)`` 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, PhysxReplicateContext)) + + +def physx_replicate( + stage: Usd.Stage, + sources: Sequence[str], # e.g. ["/World/Template/A", "/World/Template/B"] + destinations: Sequence[str], # e.g. ["/World/envs/env_{}/Robot", "/World/envs/env_{}/Object"] + env_ids: torch.Tensor, # env_ids + mapping: torch.Tensor, # (num_sources, num_envs) bool; True -> place sources[i] into world=j + positions: torch.Tensor | None = None, + quaternions: torch.Tensor | None = None, + device: str = "cpu", + exclude_self_replication: bool = True, +) -> None: + """Replicate prims via PhysX replicator with per-row mapping. + + Builds per-source destination lists from ``mapping`` and calls PhysX ``replicate``. + The replicator is registered for the call and then unregistered. Heterogeneous + rows currently use ``useEnvIds=False`` and rely on USD collision filtering. + + ``attach_fn`` excludes ``/World/template`` and ``/World/envs`` so that PhysX does + not independently parse prims that the replicator will handle. The source prim + receives its physics body as a side-effect of ``rep.replicate()`` (which always + parses the source internally), so every source must appear in at least one + ``replicate`` call. + + When ``exclude_self_replication`` is True (default), each source environment is + removed from its own replication targets so the replicator only creates bodies at + non-self destinations. If removing self would leave the world list empty (i.e. the + source maps only to its own environment), self is kept so that ``rep.replicate()`` + is still called and the source prim gets its physics body. + + Args: + stage: USD stage. + sources: Source prim paths (``S``). + destinations: Destination templates (``S``) with ``"{}"`` for env index. + env_ids: Environment indices (``[E]``). + mapping: Bool/int mask (``[S, E]``) selecting envs per source. + positions: Optional positions (unused, for API compatibility). + quaternions: Optional orientations (unused, for API compatibility). + device: Unused legacy argument retained for API compatibility. + exclude_self_replication: If True, skip replicating a source prim onto itself + when the source also maps to other environments. Default is True. + Self-only sources always keep self so that ``rep.replicate()`` fires. + + Returns: + None + """ + del device + + ctx = PhysxReplicateContext(stage) + ctx.queue_mapping( + sources, + destinations, + env_ids, + mapping, + positions=positions, + quaternions=quaternions, + exclude_self_replication=exclude_self_replication, + ) + ctx.replicate() diff --git a/source/isaaclab_physx/test/sim/test_cloner.py b/source/isaaclab_physx/test/sim/test_cloner.py index a4f1c730fc2d..68b7a36e60be 100644 --- a/source/isaaclab_physx/test/sim/test_cloner.py +++ b/source/isaaclab_physx/test/sim/test_cloner.py @@ -17,18 +17,38 @@ import pytest import torch import warp as wp -from isaaclab_physx.cloner import physx_replicate +from isaaclab_physx.cloner import PhysxReplicateContext, physx_replicate import isaaclab.sim as sim_utils from isaaclab.cloner import ( _fabric_notices, disabled_fabric_change_notifies, - make_clone_plan, sequential, usd_replicate, ) from isaaclab.sim import build_simulation_context + +def _make_flat_clone_plan(num_variants: int, num_clones: int, destination: str, device: str): + """Build a flat (sources, destinations, clone_mask) tuple for tests using sequential mapping. + + The PhysX test_cloner tests intentionally bypass cfg-driven planning and exercise + physx_replicate / usd_replicate against a hand-built per-variant mask. This helper + captures the small amount of flat-plan logic the tests need without re-introducing + the legacy ``make_clone_plan(sources, destinations, num_clones, ...)`` signature. + """ + chosen = sequential( + torch.arange(num_variants, dtype=torch.long, device=device).unsqueeze(1), + num_clones, + device, + ).view(-1) + mask = torch.zeros((num_variants, num_clones), dtype=torch.bool, device=device) + mask[chosen, torch.arange(num_clones, device=device)] = True + sources = tuple(destination.format(i) for i in range(num_variants)) + destinations = tuple([destination] * num_variants) + return sources, destinations, mask + + wp.init() pytestmark = pytest.mark.isaacsim_ci @@ -107,6 +127,29 @@ def _fake_register(_stage_id, attach_fn, attach_end_fn, rename_fn): return mock_rep, replicate_calls, attach_excluded +def test_physx_replicate_context_queue_and_replicate(sim): + """PhysxReplicateContext queues mapping rows and replicates them through attach_end_fn.""" + from unittest.mock import patch + + stage = sim_utils.get_current_stage() + sim_utils.create_prim("/World/envs", "Xform") + for i in range(3): + sim_utils.create_prim(f"/World/envs/env_{i}", "Xform") + + mock_rep, replicate_calls = _make_mock_physx_rep() + with patch("isaaclab_physx.cloner.replicate.get_physx_replicator_interface", return_value=mock_rep): + ctx = PhysxReplicateContext(stage) + ctx.queue_mapping( + sources=["/World/envs/env_0/Object"], + destinations=["/World/envs/env_{}/Object"], + env_ids=torch.arange(3, dtype=torch.long), + mapping=torch.ones((1, 3), dtype=torch.bool), + ) + ctx.replicate() + + assert replicate_calls == [2] + + @pytest.mark.parametrize( "num_envs,src,expected_worlds", [ @@ -135,7 +178,7 @@ def test_physx_replicate_world_counts(sim, num_envs, src, expected_worlds): sim_utils.create_prim(f"/World/envs/env_{i}", "Xform") mock_rep, replicate_calls = _make_mock_physx_rep() - with patch("isaaclab_physx.cloner.physx_replicate.get_physx_replicator_interface", return_value=mock_rep): + with patch("isaaclab_physx.cloner.replicate.get_physx_replicator_interface", return_value=mock_rep): physx_replicate( stage, sources=[src], @@ -217,7 +260,7 @@ def test_physx_replicate_heterogeneous_isolated_sources(sim, device): mapping[2, [7, 11]] = True mock_rep, replicate_calls, attach_excluded = _make_mock_physx_rep_detailed() - with patch("isaaclab_physx.cloner.physx_replicate.get_physx_replicator_interface", return_value=mock_rep): + with patch("isaaclab_physx.cloner.replicate.get_physx_replicator_interface", return_value=mock_rep): physx_replicate( stage, sources=["/World/envs/env_0/Object", "/World/envs/env_5/Object", "/World/envs/env_7/Object"], @@ -271,22 +314,20 @@ def test_direct_clone_plan_multi_asset(sim): mass_props=sim_utils.MassPropertiesCfg(mass=1.0), collision_props=sim_utils.CollisionPropertiesCfg(), ) - plan = make_clone_plan( - [[f"/World/envs/env_{i}/Object" for i in range(len(cfg.assets_cfg))]], - ["/World/envs/env_{}/Object"], - num_clones, - sequential, - sim.cfg.device, + sources, destinations, clone_mask = _make_flat_clone_plan( + num_variants=len(cfg.assets_cfg), + num_clones=num_clones, + destination="/World/envs/env_{}/Object", + device=sim.cfg.device, ) - spawn_paths: list[str | None] = list(plan.sources) - cfg.spawn_paths = spawn_paths + cfg.spawn_paths = list(sources) prim = cfg.func("/World/unused", cfg) assert prim.IsValid() stage = sim_utils.get_current_stage() env_ids = torch.arange(num_clones, dtype=torch.long, device=sim.cfg.device) - physx_replicate(stage, plan.sources, plan.destinations, env_ids, plan.clone_mask, device=sim.cfg.device) - usd_replicate(stage, plan.sources, plan.destinations, env_ids, plan.clone_mask) + physx_replicate(stage, sources, destinations, env_ids, clone_mask, device=sim.cfg.device) + usd_replicate(stage, sources, destinations, env_ids, clone_mask) primitive_prims = sim_utils.get_all_matching_child_prims( "/World/envs", predicate=lambda prim: prim.GetTypeName() in ["Cone", "Cube", "Sphere"] @@ -315,25 +356,23 @@ def _run_colocation_collision_filter(sim, asset_cfg, expected_types, assert_coun sim_utils.create_prim(f"/World/envs/env_{i}", "Xform", translation=(0, 0, 0)) num_variants = len(asset_cfg.assets_cfg) if isinstance(asset_cfg, sim_utils.MultiAssetSpawnerCfg) else 1 - plan = make_clone_plan( - [[f"/World/envs/env_{i}/Object" for i in range(num_variants)]], - ["/World/envs/env_{}/Object"], - num_clones, - sequential, - sim.cfg.device, + sources, destinations, clone_mask = _make_flat_clone_plan( + num_variants=num_variants, + num_clones=num_clones, + destination="/World/envs/env_{}/Object", + device=sim.cfg.device, ) if isinstance(asset_cfg, sim_utils.MultiAssetSpawnerCfg): - spawn_paths: list[str | None] = list(plan.sources) - asset_cfg.spawn_paths = spawn_paths + asset_cfg.spawn_paths = list(sources) prim = asset_cfg.func("/World/unused", asset_cfg) else: - prim = asset_cfg.func(plan.sources[0], asset_cfg) + prim = asset_cfg.func(sources[0], asset_cfg) assert prim.IsValid() stage = sim_utils.get_current_stage() env_ids = torch.arange(num_clones, dtype=torch.long, device=sim.cfg.device) - physx_replicate(stage, plan.sources, plan.destinations, env_ids, plan.clone_mask, device=sim.cfg.device) - usd_replicate(stage, plan.sources, plan.destinations, env_ids, plan.clone_mask) + physx_replicate(stage, sources, destinations, env_ids, clone_mask, device=sim.cfg.device) + usd_replicate(stage, sources, destinations, env_ids, clone_mask) primitive_prims = sim_utils.get_all_matching_child_prims( "/World/envs", predicate=lambda prim: prim.GetTypeName() in expected_types diff --git a/source/isaaclab_tasks/changelog.d/replication-session-redesign.skip b/source/isaaclab_tasks/changelog.d/replication-session-redesign.skip new file mode 100644 index 000000000000..c34b8831ae8d --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/replication-session-redesign.skip @@ -0,0 +1,4 @@ +Internal refactor: every direct env's ``_setup_scene`` now calls +``cloner.replicate(cloner.ClonePlan.from_env_0(...))`` after asset +construction in place of the prior ``with cloner.ReplicateSession():`` wrapper. +No user-visible behavior change. diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/anymal_c/anymal_c_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/anymal_c/anymal_c_env.py index 3c1c0aa0ef2f..f932022481e0 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/anymal_c/anymal_c_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/anymal_c/anymal_c_env.py @@ -10,6 +10,7 @@ import warp as wp import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import Articulation from isaaclab.envs import DirectRLEnv from isaaclab.sensors import ContactSensor, RayCaster @@ -72,8 +73,10 @@ def _setup_scene(self): self.cfg.terrain.num_envs = self.scene.cfg.num_envs self.cfg.terrain.env_spacing = self.scene.cfg.env_spacing self._terrain = self.cfg.terrain.class_type(self.cfg.terrain) - # clone and replicate - self.scene.clone_environments(copy_from_source=False) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) # we need to explicitly filter collisions for CPU simulation if self.device == "cpu": self.scene.filter_collisions(global_prim_paths=[self.cfg.terrain.prim_path]) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_env.py index e3037d67712d..10d9774f0361 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_env.py @@ -10,6 +10,7 @@ import warp as wp import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import DirectRLEnv from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane @@ -260,8 +261,11 @@ def _setup_scene(self): self._robot = Articulation(self.cfg.robot) self._fixed_asset = Articulation(self.cfg_task.fixed_asset) self._held_asset = RigidObject(self.cfg_task.held_asset) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) - self.scene.clone_environments(copy_from_source=False) self.scene.filter_collisions() self.scene.articulations["robot"] = self._robot diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/disassembly_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/disassembly_env.py index 1982786f65a7..4d85fb4a2476 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/disassembly_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/disassembly_env.py @@ -11,6 +11,7 @@ import warp as wp import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import DirectRLEnv from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane @@ -185,8 +186,11 @@ def _setup_scene(self): # self._held_asset = Articulation(self.cfg_task.held_asset) # self._fixed_asset = RigidObject(self.cfg_task.fixed_asset) self._held_asset = RigidObject(self.cfg_task.held_asset) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) - self.scene.clone_environments(copy_from_source=False) self.scene.filter_collisions() self.scene.articulations["robot"] = self._robot diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py index 47ec6d2da69a..2d7de6df1a4d 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py @@ -11,6 +11,7 @@ import torch import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import Articulation, ArticulationCfg from isaaclab.envs import DirectMARLEnv, DirectMARLEnvCfg from isaaclab.scene import InteractiveSceneCfg @@ -81,8 +82,10 @@ def _setup_scene(self): self.robot = Articulation(self.cfg.robot_cfg) # add ground plane spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) - # clone and replicate - self.scene.clone_environments(copy_from_source=False) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) # we need to explicitly filter collisions for CPU simulation if self.device == "cpu": self.scene.filter_collisions(global_prim_paths=[]) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_env.py index 807c0cc09d91..c1f88f7e558e 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_env.py @@ -12,6 +12,7 @@ import torch import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import Articulation from isaaclab.envs import DirectRLEnv from isaaclab.sensors import Camera, save_images_to_file @@ -74,9 +75,11 @@ def _setup_scene(self): """Setup the scene with the cartpole and camera.""" self._cartpole = Articulation(self.cfg.robot_cfg) self._tiled_camera = Camera(self.cfg.tiled_camera) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) - # clone and replicate - self.scene.clone_environments(copy_from_source=False) if self.device == "cpu": # we need to explicitly filter collisions for CPU simulation self.scene.filter_collisions(global_prim_paths=[]) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env.py index 10442ef08cc8..9b322f740d51 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env.py @@ -12,6 +12,7 @@ import torch import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import Articulation from isaaclab.envs import DirectRLEnv from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane @@ -38,8 +39,10 @@ def _setup_scene(self): self.cartpole = Articulation(self.cfg.robot_cfg) # add ground plane spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) - # clone and replicate - self.scene.clone_environments(copy_from_source=False) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) # we need to explicitly filter collisions for CPU simulation if self.device == "cpu": self.scene.filter_collisions(global_prim_paths=[]) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env.py index c38fe071b161..4860be520982 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env.py @@ -9,6 +9,7 @@ import carb import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import Articulation from isaaclab.envs import DirectRLEnv from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane @@ -97,8 +98,11 @@ def _setup_scene(self): if self.cfg_task.name == "gear_mesh": self._small_gear_asset = Articulation(self.cfg_task.small_gear_cfg) self._large_gear_asset = Articulation(self.cfg_task.large_gear_cfg) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) - self.scene.clone_environments(copy_from_source=False) if self.device == "cpu": # we need to explicitly filter collisions for CPU simulation self.scene.filter_collisions() diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/franka_cabinet/franka_cabinet_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/franka_cabinet/franka_cabinet_env.py index 3f3d341c1aad..a6a1d028c0f9 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/franka_cabinet/franka_cabinet_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/franka_cabinet/franka_cabinet_env.py @@ -11,6 +11,7 @@ from pxr import UsdGeom import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import Articulation from isaaclab.envs import DirectRLEnv from isaaclab.sim.utils.stage import get_current_stage @@ -134,9 +135,11 @@ def _setup_scene(self): self.cfg.terrain.num_envs = self.scene.cfg.num_envs self.cfg.terrain.env_spacing = self.scene.cfg.env_spacing self._terrain = self.cfg.terrain.class_type(self.cfg.terrain) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) - # clone and replicate - self.scene.clone_environments(copy_from_source=False) # we need to explicitly filter collisions for CPU simulation if self.device == "cpu": self.scene.filter_collisions(global_prim_paths=[self.cfg.terrain.prim_path]) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/humanoid_amp/humanoid_amp_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/humanoid_amp/humanoid_amp_env.py index 7c0fe5036132..630ce90d032b 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/humanoid_amp/humanoid_amp_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/humanoid_amp/humanoid_amp_env.py @@ -11,6 +11,7 @@ import warp as wp import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import Articulation from isaaclab.envs import DirectRLEnv from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane @@ -64,8 +65,10 @@ def _setup_scene(self): ), ), ) - # clone and replicate - self.scene.clone_environments(copy_from_source=False) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) # we need to explicitly filter collisions for CPU simulation if self.device == "cpu": self.scene.filter_collisions(global_prim_paths=["/World/ground"]) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py index bd178f3745ea..8a9fa92f1e84 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py @@ -13,6 +13,7 @@ import torch import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import DirectRLEnv from isaaclab.markers import VisualizationMarkers @@ -101,8 +102,10 @@ def _setup_scene(self): self._joint_wrench_sensor = self._create_joint_wrench_sensor() # add ground plane spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) - # clone and replicate (no need to filter for this environment) - self.scene.clone_environments(copy_from_source=False) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) # add articulation to scene - we must register to scene to randomize with EventManager self.scene.articulations["robot"] = self.hand self.scene.rigid_objects["object"] = self.object diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env.py index f97871ca83bf..56a0fd754f9c 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env.py @@ -12,6 +12,7 @@ from isaaclab_physx.physics import PhysxCfg import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import Articulation from isaaclab.envs import DirectRLEnv, DirectRLEnvCfg from isaaclab.utils.math import ( @@ -113,8 +114,10 @@ def _setup_scene(self): self.cfg.terrain.num_envs = self.scene.cfg.num_envs self.cfg.terrain.env_spacing = self.scene.cfg.env_spacing self.terrain = self.cfg.terrain.class_type(self.cfg.terrain) - # clone and replicate - self.scene.clone_environments(copy_from_source=False) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) # we need to explicitly filter collisions for CPU simulation if self.device == "cpu": self.scene.filter_collisions(global_prim_paths=[self.cfg.terrain.prim_path]) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/quadcopter/quadcopter_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/quadcopter/quadcopter_env.py index c1e3a1b3c9ee..0b26c4c9c9de 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/quadcopter/quadcopter_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/quadcopter/quadcopter_env.py @@ -10,6 +10,7 @@ import warp as wp import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import Articulation from isaaclab.envs import DirectRLEnv from isaaclab.envs.ui import BaseEnvWindow @@ -82,8 +83,10 @@ def _setup_scene(self): self.cfg.terrain.num_envs = self.scene.cfg.num_envs self.cfg.terrain.env_spacing = self.scene.cfg.env_spacing self._terrain = self.cfg.terrain.class_type(self.cfg.terrain) - # clone and replicate - self.scene.clone_environments(copy_from_source=False) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) # we need to explicitly filter collisions for CPU simulation if self.device == "cpu": self.scene.filter_collisions(global_prim_paths=[self.cfg.terrain.prim_path]) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env.py index 7da4db48e853..b133e6ab47d1 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env.py @@ -11,6 +11,7 @@ import torch import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import Articulation, RigidObject from isaaclab.sensors import Camera from isaaclab.utils.math import quat_apply @@ -51,8 +52,10 @@ def _setup_scene(self): self.object: Articulation | RigidObject = self.cfg.object_cfg.class_type(self.cfg.object_cfg) self._joint_wrench_sensor = self._create_joint_wrench_sensor() self._tiled_camera = Camera(self.cfg.tiled_camera) - # clone and replicate (no need to filter for this environment) - self.scene.clone_environments(copy_from_source=False) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) # add articulation to scene - we must register to scene to randomize with EventManager self.scene.articulations["robot"] = self.hand self.scene.rigid_objects["object"] = self.object diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py index 66b9012d036b..9e1e56a3f5cc 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py @@ -12,6 +12,7 @@ import torch import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import DirectMARLEnv from isaaclab.markers import VisualizationMarkers @@ -93,8 +94,10 @@ def _setup_scene(self): self.object = RigidObject(self.cfg.object_cfg) # add ground plane spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) - # clone and replicate (no need to filter for this environment) - self.scene.clone_environments(copy_from_source=False) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) # add articulation to scene - we must register to scene to randomize with EventManager self.scene.articulations["right_robot"] = self.right_hand self.scene.articulations["left_robot"] = self.left_hand diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py index 0826f6d98225..357ee522a60a 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py @@ -12,6 +12,7 @@ from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import Articulation, ArticulationCfg from isaaclab.envs import DirectRLEnvCfg from isaaclab.scene import InteractiveSceneCfg @@ -258,8 +259,10 @@ def _setup_scene(self) -> None: self.cartpole = Articulation(self.cfg.robot_cfg) # add ground plane spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) - # clone and replicate - self.scene.clone_environments(copy_from_source=False) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) # we need to explicitly filter collisions for CPU simulation if self.device == "cpu": self.scene.filter_collisions(global_prim_paths=[]) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py index d5bc6a2ad529..2082d470ff9b 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py @@ -14,6 +14,7 @@ from isaaclab_experimental.envs import DirectRLEnvWarp import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import Articulation # , RigidObject from isaaclab.markers import VisualizationMarkers from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane @@ -673,8 +674,10 @@ def _setup_scene(self): self.object = Articulation(self.cfg.object_cfg) # add ground plane spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) - # clone and replicate (no need to filter for this environment) - self.scene.clone_environments(copy_from_source=False) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) # add articulation to scene - we must register to scene to randomize with EventManager self.scene.articulations["robot"] = self.hand self.scene.articulations["object"] = self.object diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py index d8d712a655e9..62400c4d48fe 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py @@ -9,6 +9,7 @@ from isaaclab_experimental.envs import DirectRLEnvWarp import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import Articulation from isaaclab.envs import DirectRLEnvCfg @@ -415,8 +416,10 @@ def _setup_scene(self) -> None: self.cfg.terrain.num_envs = self.scene.cfg.num_envs self.cfg.terrain.env_spacing = self.scene.cfg.env_spacing self.terrain = self.cfg.terrain.class_type(self.cfg.terrain) - # clone and replicate - self.scene.clone_environments(copy_from_source=False) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) # add articulation to scene self.scene.articulations["robot"] = self.robot # add lights diff --git a/tools/template/templates/tasks/direct_multi-agent/env b/tools/template/templates/tasks/direct_multi-agent/env index eec2331722e8..0ca0ce4585ee 100644 --- a/tools/template/templates/tasks/direct_multi-agent/env +++ b/tools/template/templates/tasks/direct_multi-agent/env @@ -10,6 +10,7 @@ import torch from collections.abc import Sequence import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import Articulation from isaaclab.envs import DirectMARLEnv from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane @@ -35,8 +36,11 @@ class {{ task.classname }}Env(DirectMARLEnv): self.robot = Articulation(self.cfg.robot_cfg) # add ground plane spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) - # clone and replicate - self.scene.clone_environments(copy_from_source=False) + # build a homogeneous clone plan and replicate the env_0 layout to every env + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) # we need to explicitly filter collisions for CPU simulation if self.device == "cpu": self.scene.filter_collisions(global_prim_paths=[]) diff --git a/tools/template/templates/tasks/direct_single-agent/env b/tools/template/templates/tasks/direct_single-agent/env index e6f47fd3366b..8f956c845b86 100644 --- a/tools/template/templates/tasks/direct_single-agent/env +++ b/tools/template/templates/tasks/direct_single-agent/env @@ -10,6 +10,7 @@ import torch from collections.abc import Sequence import isaaclab.sim as sim_utils +from isaaclab import cloner from isaaclab.assets import Articulation from isaaclab.envs import DirectRLEnv from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane @@ -34,8 +35,11 @@ class {{ task.classname }}Env(DirectRLEnv): self.robot = Articulation(self.cfg.robot_cfg) # add ground plane spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) - # clone and replicate - self.scene.clone_environments(copy_from_source=False) + # build a homogeneous clone plan and replicate the env_0 layout to every env + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) # we need to explicitly filter collisions for CPU simulation if self.device == "cpu": self.scene.filter_collisions(global_prim_paths=[]) From fdb8e735b33741d94bdd7ccd2e16a71f41a0dae2 Mon Sep 17 00:00:00 2001 From: Octi Zhang Date: Mon, 25 May 2026 04:15:29 -0700 Subject: [PATCH 2/7] lint --- .../replication-session-redesign.minor.rst | 97 +++++++------------ .../isaaclab/cloner/_fabric_notices.py | 45 +++------ .../isaaclab/cloner/replicate_session.py | 16 +-- source/isaaclab/isaaclab/cloner/usd.py | 4 +- .../isaaclab/scene/interactive_scene.py | 4 +- .../isaaclab_newton/cloner/replicate.py | 6 +- 6 files changed, 62 insertions(+), 110 deletions(-) diff --git a/source/isaaclab/changelog.d/replication-session-redesign.minor.rst b/source/isaaclab/changelog.d/replication-session-redesign.minor.rst index c7578decf575..598ceefabb2a 100644 --- a/source/isaaclab/changelog.d/replication-session-redesign.minor.rst +++ b/source/isaaclab/changelog.d/replication-session-redesign.minor.rst @@ -1,80 +1,49 @@ Added ^^^^^ -* Added :data:`~isaaclab.cloner.REPLICATION_QUEUE`, a module-level list that backend - ``queue__replication`` helpers append ``(cfg, BackendCtxCls)`` pairs to. - Drained by :func:`~isaaclab.cloner.replicate`. -* Added :func:`~isaaclab.cloner.replicate` as a free function that takes an explicit - :class:`~isaaclab.cloner.ClonePlan`, drains - :data:`~isaaclab.cloner.REPLICATION_QUEUE`, dispatches each backend context in - ascending ``replicate_priority`` order, publishes the plan via - :meth:`~isaaclab.sim.SimulationContext.set_clone_plan`, and clears the queue. -* Added :meth:`~isaaclab.cloner.ClonePlan.from_env_0` classmethod that - constructs a homogeneous single-source clone plan for direct envs and auto-populates - ``cfg_rows`` from :data:`~isaaclab.cloner.REPLICATION_QUEUE` filtered by env-root - prefix. -* Added :attr:`~isaaclab.cloner.CloneCfg.clone_regex` (default ``"/World/envs/env_.*"``) - as the single source of truth for the env-namespace convention. Consumers inside - :class:`~isaaclab.scene.InteractiveScene` read it directly when expanding - ``{ENV_REGEX_NS}`` cfg macros. +* Added :data:`~isaaclab.cloner.REPLICATION_QUEUE` and the free function + :func:`~isaaclab.cloner.replicate`, the explicit registry-and-drain pair + that backends now hook into for replication. +* Added :meth:`~isaaclab.cloner.ClonePlan.from_env_0` for direct envs that + clone a single env-0 prototype across every env. +* Added :attr:`~isaaclab.cloner.CloneCfg.clone_regex` as the single source + of truth for the env-namespace convention (default ``"/World/envs/env_.*"``). Fixed ^^^^^ -* Fixed :data:`~isaaclab.cloner.REPLICATION_QUEUE` leaking stale ``(cfg, BackendCtxCls)`` - entries when a backend raised mid-dispatch or when asset construction failed inside - :class:`~isaaclab.cloner.ReplicateSession`. :func:`~isaaclab.cloner.replicate` now - snapshots and clears the queue up front, and :meth:`ReplicateSession.__exit__` clears - it on the exception path. +* Fixed :data:`~isaaclab.cloner.REPLICATION_QUEUE` leaking stale entries + when a backend or asset construction raised mid-session. Changed ^^^^^^^ -* **Breaking:** Rewrote :class:`~isaaclab.cloner.ReplicateSession` as a thin context - manager that calls :func:`~isaaclab.cloner.make_clone_plan` in ``__enter__`` and - :func:`~isaaclab.cloner.replicate` in ``__exit__``. The class no longer exposes - ``plan``, ``stage``, ``env_ids``, ``cfg_rows``, ``positions``, ``publish_clone_plan_fn``, - or ``replicate_on_exit`` fields, and the ``ReplicateSession()`` (no-arg) form is gone. - Direct envs migrate to ``cloner.replicate(cloner.ClonePlan.from_env_0(...))``. -* **Breaking:** Changed :func:`~isaaclab.cloner.make_clone_plan` signature from - ``make_clone_plan(sources, destinations, num_clones, clone_strategy, device)`` (pure - flat-mapping helper) to - ``make_clone_plan(cfgs, num_clones, env_spacing, device, *, clone_strategy=sequential, - valid_set=None, stage=None)``. The new function absorbs the cfg-driven planning logic - formerly in ``InteractiveScene._build_clone_plan_from_cfg``, returns a - self-contained :class:`~isaaclab.cloner.ClonePlan` with ``cfg_rows`` populated, and - mutates each cfg's ``spawn_path`` / ``spawn_paths``. -* **Breaking:** :class:`~isaaclab.cloner.ClonePlan` now carries ``env_ids``, - ``positions``, and ``cfg_rows`` fields in addition to ``sources``, ``destinations``, - ``clone_mask``. All fields are required; ``cfg_rows`` is never ``None``. The USD - stage is intentionally *not* part of the plan — pass it explicitly via the new - ``stage=`` keyword on :func:`~isaaclab.cloner.replicate` and - :class:`~isaaclab.cloner.ReplicateSession`. -* **Breaking:** :func:`~isaaclab.cloner.replicate` now requires a ``stage`` keyword - argument: ``cloner.replicate(plan, stage=scene.stage)``. Direct envs typically pass - ``self.scene.stage``. -* **Breaking:** :class:`~isaaclab.cloner.ReplicateSession` now requires a ``stage`` - keyword argument that is forwarded to :func:`replicate` in ``__exit__``. -* Changed :attr:`~isaaclab.scene.InteractiveScene.env_origins` to read its per-env - positions from the :class:`~isaaclab.cloner.ClonePlan` published to - :class:`~isaaclab.sim.SimulationContext`, making the published plan the single - source of truth for env placement. :class:`~isaaclab.scene.InteractiveScene` now - always enters :class:`~isaaclab.cloner.ReplicateSession` (even for scene cfgs - with no entities, where the resulting plan is empty but still carries - ``positions``), so the lookup is unconditional and the prior - ``InteractiveScene._default_env_origins`` cached tensor has been removed. +* **Breaking:** Rewrote :class:`~isaaclab.cloner.ReplicateSession` as a thin + context manager around :func:`~isaaclab.cloner.make_clone_plan` and + :func:`~isaaclab.cloner.replicate`. The no-arg form and the cached + ``plan`` / ``cfg_rows`` / ``replicate_on_exit`` fields are gone. Direct + envs migrate to ``cloner.replicate(cloner.ClonePlan.from_env_0(...))``. +* **Breaking:** Changed :func:`~isaaclab.cloner.make_clone_plan` to take + ``cfgs`` and absorb the cfg-driven planning logic previously inside + :class:`~isaaclab.scene.InteractiveScene`, returning a self-contained + :class:`~isaaclab.cloner.ClonePlan`. +* **Breaking:** :func:`~isaaclab.cloner.replicate` and + :class:`~isaaclab.cloner.ReplicateSession` now require an explicit + ``stage=`` keyword; the :class:`~isaaclab.cloner.ClonePlan` is + stage-agnostic. +* Changed :attr:`~isaaclab.scene.InteractiveScene.env_origins` to read from + the published :class:`~isaaclab.cloner.ClonePlan`, making the plan the + single source of truth for env placement. Removed ^^^^^^^ -* **Breaking:** Removed the module-level ``isaaclab.cloner.replicate_session_defaults`` - dictionary and the ``isaaclab.cloner.replicate_session`` active-session reference. - Replication is now driven by the explicit :data:`~isaaclab.cloner.REPLICATION_QUEUE` - registry and :func:`~isaaclab.cloner.replicate` drain. -* **Breaking:** Removed :meth:`InteractiveScene.clone_environments` deprecated shim. - Direct envs should use ``cloner.replicate(cloner.ClonePlan.from_env_0(...))`` - to publish the plan. +* **Breaking:** Removed ``isaaclab.cloner.replicate_session_defaults`` and + ``isaaclab.cloner.replicate_session``. Use + :data:`~isaaclab.cloner.REPLICATION_QUEUE` and + :func:`~isaaclab.cloner.replicate` instead. +* **Breaking:** Removed :meth:`InteractiveScene.clone_environments`; direct + envs should use ``cloner.replicate(cloner.ClonePlan.from_env_0(...))``. * **Breaking:** Removed :attr:`InteractiveScene.env_ns` and - :attr:`InteractiveScene.env_regex_ns` properties (only used internally). Read - :attr:`~isaaclab.cloner.CloneCfg.clone_regex` directly when the env-namespace - convention is needed. + :attr:`InteractiveScene.env_regex_ns`; read + :attr:`~isaaclab.cloner.CloneCfg.clone_regex` instead. diff --git a/source/isaaclab/isaaclab/cloner/_fabric_notices.py b/source/isaaclab/isaaclab/cloner/_fabric_notices.py index 0c7c850e9964..b3557ef4bd41 100644 --- a/source/isaaclab/isaaclab/cloner/_fabric_notices.py +++ b/source/isaaclab/isaaclab/cloner/_fabric_notices.py @@ -157,46 +157,23 @@ def get_bindings() -> FabricNoticeBindings | None: @contextlib.contextmanager def disabled_fabric_change_notifies(stage: Usd.Stage, *, restore: bool = True) -> Iterator[None]: - """Suspend the ``IFabricUsd`` USD notice listener for the body of the ``with`` block. + """Suspend Fabric's USD notice listener for the body of the ``with`` block. - Targets the same handler that :meth:`isaacsim.core.cloner.Cloner.disable_change_listener` - toggles, but goes through ``omni::fabric::IFabricUsd`` directly so we don't take an + The listener is a global ``TfNotice`` that fires on every ``Sdf.CopySpec`` during + cloning; toggling it off via ``IFabricUsd::setEnableChangeNotifies`` skips the + per-spec Fabric sync that dominates cloning time for large PhysX rigid-body scenes. + Same toggle that :meth:`isaacsim.core.cloner.Cloner.disable_change_listener` flips, + called directly through Carbonite so we don't take an ``isaacsim.core.simulation_manager`` dependency. - The listener is a global ``TfNotice`` registered when ``omni.fabric`` loads; it - short-circuits via a soft flag (``IFabricUsd.cpp:739``). Toggling that flag is what - skips the per-``Sdf.CopySpec`` Fabric sync that dominates cloning time on large scenes. - - When this provides a measurable speedup - ---------------------------------------- - Bisection on the regression test (see ``test_cloner.py``) shows the listener cost is - only on the critical path when **all** of these hold: - - 1. The clone happens through the ``InteractiveScene`` path with backend physics - replication queued. Calling :func:`usd_replicate` directly on a stage produces - no measurable gap; without physics replication the gap drops to ~1.19x. The - PhysX replication path is what amplifies per-spec listener work. - 2. The cloned prims carry PhysX rigid-body schemas (e.g. ``UsdPhysics.RigidBodyAPI``, - authored via ``rigid_props`` on a spawn cfg). Plain Xforms or geometry without - physics schemas produce ~1.0x — the listener has no Fabric-tracked state to sync. - ``mass_props`` and ``collision_props`` add nothing beyond ``rigid_props``. - 3. Total per-``Sdf.CopySpec`` firings reach ~32K — i.e. ``num_bodies × num_envs`` is - large enough to dominate scene-init cost. Below this the speedup sinks into noise. - - Conditions outside this envelope (no PhysX schemas, single-env scenes, raw - ``usd_replicate`` calls, no physics replication) won't see a perf win — the - suspension is correct but its effect is lost in the rest of the work. - - Re-entrant: if the flag is already off on entry, ``__exit__`` leaves it off. Falls - through to a no-op if the Carbonite interface can't be acquired (e.g. outside a live - Kit application) — the caller never breaks, it just doesn't get the perf win. + Falls through to a no-op if the Carbonite interface can't be acquired (e.g. outside + a live Kit application). Args: stage: USD stage whose Fabric notice handler should be suspended. - restore: When ``True`` (default), re-enable the handler on exit. Set to ``False`` - inside a known clone-then-``sim.reset`` window where the downstream Fabric - resync happens anyway and re-enabling here would trigger a redundant - ``forceMinimalPopulate`` batch — see ``PluginInterface.cpp:337``. + restore: When ``True`` (default), re-enable the handler on exit. Set to + ``False`` if a downstream Fabric resync is about to happen anyway (e.g. + right before ``sim.reset()``), to avoid a redundant re-enable. Yields: None. diff --git a/source/isaaclab/isaaclab/cloner/replicate_session.py b/source/isaaclab/isaaclab/cloner/replicate_session.py index 5341867e1a24..69e192068c6d 100644 --- a/source/isaaclab/isaaclab/cloner/replicate_session.py +++ b/source/isaaclab/isaaclab/cloner/replicate_session.py @@ -14,6 +14,7 @@ if TYPE_CHECKING: import torch + from pxr import Usd from .clone_plan import ClonePlan @@ -23,7 +24,7 @@ """``(cfg, BackendCtxCls)`` pairs appended by ``queue__replication`` and drained by :func:`replicate`.""" -def replicate(plan: "ClonePlan", *, stage: "Usd.Stage") -> None: +def replicate(plan: ClonePlan, *, stage: Usd.Stage) -> None: """Drain :data:`REPLICATION_QUEUE` against ``plan``, dispatch each backend, publish the plan. Cfgs absent from ``plan.cfg_rows`` are silently skipped. Backend contexts run in @@ -70,8 +71,7 @@ class ReplicateSession: .. code-block:: python - with cloner.ReplicateSession(cfgs, num_clones=128, env_spacing=2.0, - device="cuda:0", stage=sim.stage): + with cloner.ReplicateSession(cfgs, num_clones=128, env_spacing=2.0, device="cuda:0", stage=sim.stage): for cfg in cfgs: cfg.class_type(cfg) """ @@ -83,9 +83,9 @@ def __init__( env_spacing: float, device: str, *, - stage: "Usd.Stage", + stage: Usd.Stage, clone_strategy: Callable = sequential, - valid_set: "torch.Tensor | None" = None, + valid_set: torch.Tensor | None = None, ): """Capture arguments for :func:`make_clone_plan` and :func:`replicate`. @@ -108,9 +108,9 @@ def __init__( clone_strategy=clone_strategy, valid_set=valid_set, ) - self._plan: "ClonePlan | None" = None + self._plan: ClonePlan | None = None - def __enter__(self) -> "ReplicateSession": + def __enter__(self) -> ReplicateSession: from .cloner_utils import make_clone_plan # noqa: PLC0415 self._plan = make_clone_plan(self._cfgs, **self._kwargs) @@ -125,7 +125,7 @@ def __exit__(self, exc_type, exc_value, traceback) -> None: REPLICATION_QUEUE.clear() @property - def plan(self) -> "ClonePlan": + def plan(self) -> ClonePlan: """The :class:`~isaaclab.cloner.ClonePlan` produced in :meth:`__enter__`.""" if self._plan is None: raise RuntimeError("ReplicateSession.plan is only available inside the with block.") diff --git a/source/isaaclab/isaaclab/cloner/usd.py b/source/isaaclab/isaaclab/cloner/usd.py index 070a43a2d4c8..eb3c0bbd14d8 100644 --- a/source/isaaclab/isaaclab/cloner/usd.py +++ b/source/isaaclab/isaaclab/cloner/usd.py @@ -101,7 +101,9 @@ def dp_depth(template: str) -> int: dp = template.format(0) return Sdf.Path(dp).pathElementCount - depth_to_items: dict[int, list[tuple[str, str, torch.Tensor, torch.Tensor | None, torch.Tensor | None]]] = {} + depth_to_items: dict[ + int, list[tuple[str, str, torch.Tensor, torch.Tensor | None, torch.Tensor | None]] + ] = {} for item in self._queue: depth_to_items.setdefault(dp_depth(item[1]), []).append(item) diff --git a/source/isaaclab/isaaclab/scene/interactive_scene.py b/source/isaaclab/isaaclab/scene/interactive_scene.py index 368ca6ddc75c..0ac98e210ec3 100644 --- a/source/isaaclab/isaaclab/scene/interactive_scene.py +++ b/source/isaaclab/isaaclab/scene/interactive_scene.py @@ -360,7 +360,9 @@ 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)``. From the terrain when registered, + else from the published :class:`~isaaclab.cloner.ClonePlan`. + """ if self._terrain is not None: return self._terrain.env_origins return self.sim.get_clone_plan().positions diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py b/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py index 50bf1d049d26..6a28e353700a 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py @@ -354,8 +354,10 @@ def _merged_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 + 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) From 3fb6337670f73aa31759ce61229cecffe2d4c0c7 Mon Sep 17 00:00:00 2001 From: Octi Zhang Date: Mon, 25 May 2026 04:41:04 -0700 Subject: [PATCH 3/7] address some comments --- source/isaaclab/isaaclab/cloner/clone_plan.py | 5 +- source/isaaclab/isaaclab/cloner/usd.py | 105 +++++++++--------- 2 files changed, 59 insertions(+), 51 deletions(-) diff --git a/source/isaaclab/isaaclab/cloner/clone_plan.py b/source/isaaclab/isaaclab/cloner/clone_plan.py index 8abb00d5d303..53df4e128c65 100644 --- a/source/isaaclab/isaaclab/cloner/clone_plan.py +++ b/source/isaaclab/isaaclab/cloner/clone_plan.py @@ -45,7 +45,10 @@ 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``. + ``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 + subsequent :func:`~isaaclab.cloner.replicate` call. Args: source: Source prim path (typically ``/World/envs/env_0``). diff --git a/source/isaaclab/isaaclab/cloner/usd.py b/source/isaaclab/isaaclab/cloner/usd.py index eb3c0bbd14d8..b1a7ef9b755b 100644 --- a/source/isaaclab/isaaclab/cloner/usd.py +++ b/source/isaaclab/isaaclab/cloner/usd.py @@ -12,6 +12,7 @@ from pxr import Gf, Sdf, Usd, UsdGeom, Vt +from ._fabric_notices import disabled_fabric_change_notifies from .replicate_session import REPLICATION_QUEUE @@ -93,56 +94,60 @@ def replicate(self) -> None: if not self._queue: return - try: - rl = self.stage.GetRootLayer() - - def dp_depth(template: str) -> int: - """Return destination prim path depth for stable parent-first replication.""" - dp = template.format(0) - return Sdf.Path(dp).pathElementCount - - depth_to_items: dict[ - int, list[tuple[str, str, torch.Tensor, torch.Tensor | None, torch.Tensor | None]] - ] = {} - for item in self._queue: - depth_to_items.setdefault(dp_depth(item[1]), []).append(item) - - for depth in sorted(depth_to_items.keys()): - with Sdf.ChangeBlock(): - for src, tmpl, target_envs, positions, quaternions in depth_to_items[depth]: - for wid in target_envs.tolist(): - wid = int(wid) - dp = tmpl.format(wid) - Sdf.CreatePrimInLayer(rl, dp) - if src != dp: - Sdf.CopySpec(rl, Sdf.Path(src), rl, Sdf.Path(dp)) - - if positions is not None or quaternions is not None: - ps = rl.GetPrimAtPath(dp) - op_names = [] - if positions is not None: - p = positions[wid] - t_attr = ps.GetAttributeAtPath(dp + ".xformOp:translate") - if t_attr is None: - t_attr = Sdf.AttributeSpec(ps, "xformOp:translate", Sdf.ValueTypeNames.Double3) - t_attr.default = Gf.Vec3d(float(p[0]), float(p[1]), float(p[2])) - op_names.append("xformOp:translate") - if quaternions is not None: - q = quaternions[wid] - o_attr = ps.GetAttributeAtPath(dp + ".xformOp:orient") - if o_attr is None: - o_attr = Sdf.AttributeSpec(ps, "xformOp:orient", Sdf.ValueTypeNames.Quatd) - o_attr.default = Gf.Quatd( - float(q[3]), Gf.Vec3d(float(q[0]), float(q[1]), float(q[2])) - ) - op_names.append("xformOp:orient") - if op_names: - op_order = ps.GetAttributeAtPath(dp + ".xformOpOrder") or Sdf.AttributeSpec( - ps, UsdGeom.Tokens.xformOpOrder, Sdf.ValueTypeNames.TokenArray - ) - op_order.default = Vt.TokenArray(op_names) - finally: - self._queue.clear() + # Suspend Fabric's per-Sdf.CopySpec notice listener for the duration of the copy work; + # no-op outside a live Kit application. + with disabled_fabric_change_notifies(self.stage): + self._apply_queue() + + def _apply_queue(self) -> None: + """Author the queued copy specs into the stage's root layer.""" + rl = self.stage.GetRootLayer() + + def dp_depth(template: str) -> int: + """Return destination prim path depth for stable parent-first replication.""" + dp = template.format(0) + return Sdf.Path(dp).pathElementCount + + depth_to_items: dict[ + int, list[tuple[str, str, torch.Tensor, torch.Tensor | None, torch.Tensor | None]] + ] = {} + for item in self._queue: + depth_to_items.setdefault(dp_depth(item[1]), []).append(item) + + for depth in sorted(depth_to_items.keys()): + with Sdf.ChangeBlock(): + for src, tmpl, target_envs, positions, quaternions in depth_to_items[depth]: + for wid in target_envs.tolist(): + wid = int(wid) + dp = tmpl.format(wid) + Sdf.CreatePrimInLayer(rl, dp) + if src != dp: + Sdf.CopySpec(rl, Sdf.Path(src), rl, Sdf.Path(dp)) + + if positions is not None or quaternions is not None: + ps = rl.GetPrimAtPath(dp) + op_names = [] + if positions is not None: + p = positions[wid] + t_attr = ps.GetAttributeAtPath(dp + ".xformOp:translate") + if t_attr is None: + t_attr = Sdf.AttributeSpec(ps, "xformOp:translate", Sdf.ValueTypeNames.Double3) + t_attr.default = Gf.Vec3d(float(p[0]), float(p[1]), float(p[2])) + op_names.append("xformOp:translate") + if quaternions is not None: + q = quaternions[wid] + o_attr = ps.GetAttributeAtPath(dp + ".xformOp:orient") + if o_attr is None: + o_attr = Sdf.AttributeSpec(ps, "xformOp:orient", Sdf.ValueTypeNames.Quatd) + o_attr.default = Gf.Quatd( + float(q[3]), Gf.Vec3d(float(q[0]), float(q[1]), float(q[2])) + ) + op_names.append("xformOp:orient") + if op_names: + op_order = ps.GetAttributeAtPath(dp + ".xformOpOrder") or Sdf.AttributeSpec( + ps, UsdGeom.Tokens.xformOpOrder, Sdf.ValueTypeNames.TokenArray + ) + op_order.default = Vt.TokenArray(op_names) def queue_usd_replication(cfg: Any) -> None: From 080b727820f2331e514ea1ef694383d09b6e9b53 Mon Sep 17 00:00:00 2001 From: Octi Zhang Date: Mon, 25 May 2026 21:47:53 -0700 Subject: [PATCH 4/7] Fix CI: ruff-format usd.py and add ovphysx changelog skip - ruff-format collapsed two multi-line expressions in cloner/usd.py back to single lines. - isaaclab_ovphysx gained changes in this PR (cloner module rewrite) but had no changelog fragment; add a .skip mirroring the sibling backends. --- source/isaaclab/isaaclab/cloner/usd.py | 8 ++------ .../changelog.d/replication-session-redesign.skip | 0 2 files changed, 2 insertions(+), 6 deletions(-) create mode 100644 source/isaaclab_ovphysx/changelog.d/replication-session-redesign.skip diff --git a/source/isaaclab/isaaclab/cloner/usd.py b/source/isaaclab/isaaclab/cloner/usd.py index b1a7ef9b755b..ace2adffbcf0 100644 --- a/source/isaaclab/isaaclab/cloner/usd.py +++ b/source/isaaclab/isaaclab/cloner/usd.py @@ -108,9 +108,7 @@ def dp_depth(template: str) -> int: dp = template.format(0) return Sdf.Path(dp).pathElementCount - depth_to_items: dict[ - int, list[tuple[str, str, torch.Tensor, torch.Tensor | None, torch.Tensor | None]] - ] = {} + depth_to_items: dict[int, list[tuple[str, str, torch.Tensor, torch.Tensor | None, torch.Tensor | None]]] = {} for item in self._queue: depth_to_items.setdefault(dp_depth(item[1]), []).append(item) @@ -139,9 +137,7 @@ def dp_depth(template: str) -> int: o_attr = ps.GetAttributeAtPath(dp + ".xformOp:orient") if o_attr is None: o_attr = Sdf.AttributeSpec(ps, "xformOp:orient", Sdf.ValueTypeNames.Quatd) - o_attr.default = Gf.Quatd( - float(q[3]), Gf.Vec3d(float(q[0]), float(q[1]), float(q[2])) - ) + o_attr.default = Gf.Quatd(float(q[3]), Gf.Vec3d(float(q[0]), float(q[1]), float(q[2]))) op_names.append("xformOp:orient") if op_names: op_order = ps.GetAttributeAtPath(dp + ".xformOpOrder") or Sdf.AttributeSpec( diff --git a/source/isaaclab_ovphysx/changelog.d/replication-session-redesign.skip b/source/isaaclab_ovphysx/changelog.d/replication-session-redesign.skip new file mode 100644 index 000000000000..e69de29bb2d1 From 7349eb7cb8a69cef6a8c95c6cd94c0e7c943a3da Mon Sep 17 00:00:00 2001 From: Octi Zhang Date: Tue, 26 May 2026 00:00:47 -0700 Subject: [PATCH 5/7] fix bugs --- .../isaaclab/cloner/replicate_session.py | 15 +++-- .../isaaclab/isaaclab/sensors/sensor_base.py | 7 ++- source/isaaclab/test/sim/test_cloner.py | 57 +++++++++++++++++-- 3 files changed, 67 insertions(+), 12 deletions(-) diff --git a/source/isaaclab/isaaclab/cloner/replicate_session.py b/source/isaaclab/isaaclab/cloner/replicate_session.py index 69e192068c6d..d55c501b3bd9 100644 --- a/source/isaaclab/isaaclab/cloner/replicate_session.py +++ b/source/isaaclab/isaaclab/cloner/replicate_session.py @@ -36,16 +36,19 @@ def replicate(plan: ClonePlan, *, stage: Usd.Stage) -> None: queued = list(REPLICATION_QUEUE) REPLICATION_QUEUE.clear() - backend_ctxs: dict[type, Any] = {} + # Group queued cfgs by backend, taking the union of row indices each backend owns. + backend_rows: dict[type, set[int]] = {} for cfg, BackendCtxCls in queued: rows = plan.cfg_rows.get(id(cfg)) if rows is None: continue - ctx = backend_ctxs.get(BackendCtxCls) - if ctx is None: - ctx = BackendCtxCls(stage) - backend_ctxs[BackendCtxCls] = ctx - row_list = list(rows) + backend_rows.setdefault(BackendCtxCls, set()).update(rows) + + backend_ctxs: dict[type, Any] = {} + for BackendCtxCls, row_set in backend_rows.items(): + ctx = BackendCtxCls(stage) + backend_ctxs[BackendCtxCls] = ctx + row_list = sorted(row_set) ctx.queue_mapping( [plan.sources[i] for i in row_list], [plan.destinations[i] for i in row_list], diff --git a/source/isaaclab/isaaclab/sensors/sensor_base.py b/source/isaaclab/isaaclab/sensors/sensor_base.py index b266ce7cb425..7510c9521e60 100644 --- a/source/isaaclab/isaaclab/sensors/sensor_base.py +++ b/source/isaaclab/isaaclab/sensors/sensor_base.py @@ -217,10 +217,13 @@ def _initialize_impl(self): self._device = sim.device self._backend = sim.backend self._sim_physics_dt = sim.get_physics_dt() - # Count number of environments + # Count number of environments. Prefer the active simulation's clone plan, since + # USD may only carry the env_0 prototype under some backends (e.g. Newton clones + # solver-side without authoring per-env USD specs). env_prim_path_expr = self.cfg.prim_path.rsplit("/", 1)[0] self._parent_prims = sim_utils.find_matching_prims(env_prim_path_expr) - self._num_envs = len(self._parent_prims) + plan = sim.get_clone_plan() + self._num_envs = int(plan.env_ids.numel()) if plan is not None else len(self._parent_prims) # Create warp env mask arrays for "all envs" cases and resets. # Note: We use wp.to_torch() to create zero-copy torch tensor views of warp arrays. # This allows warp arrays to be passed to warp kernels while the corresponding torch diff --git a/source/isaaclab/test/sim/test_cloner.py b/source/isaaclab/test/sim/test_cloner.py index 3b2757484a31..0247d8d0efb3 100644 --- a/source/isaaclab/test/sim/test_cloner.py +++ b/source/isaaclab/test/sim/test_cloner.py @@ -394,17 +394,66 @@ def replicate(self): replicate(plan, stage=sim_utils.get_current_stage()) - # Exactly one FakeCtx instance is shared across both cfgs. + # Exactly one FakeCtx instance is shared across both cfgs, dispatched once per backend + # with the union of rows the cfgs own. assert len(FakeCtx.instances) == 1 ctx = FakeCtx.instances[0] - assert len(ctx.queue_calls) == 2 - assert ctx.queue_calls[0][0] == ("/World/envs/env_0/Robot",) - assert ctx.queue_calls[1][0] == ("/World/envs/env_0/Object",) + assert len(ctx.queue_calls) == 1 + sources, _destinations, mask = ctx.queue_calls[0] + assert sources == ("/World/envs/env_0/Robot", "/World/envs/env_0/Object") + assert mask.shape == (2, 4) assert ctx.replicate_calls == 1 assert sim.get_clone_plan() is plan assert REPLICATION_QUEUE == [] +def test_replicate_dedupes_shared_rows_across_cfgs(sim): + """Regression: multiple cfgs sharing the same row dispatch one mapping row, not N. + + In a homogeneous plan every cfg under the env root maps to row 0; without dedup, each + cfg would tell the backend to re-instantiate row 0 once more, multiplying the count + of articulations/rigids per world by the number of cfgs. + """ + + class FakeCtx: + replicate_priority = 0 + instances: list["FakeCtx"] = [] + + def __init__(self, stage): + self.queue_calls: list[tuple] = [] + self.replicate_calls = 0 + FakeCtx.instances.append(self) + + def queue_mapping(self, sources, destinations, env_ids, mask, *, positions=None): + self.queue_calls.append((tuple(sources), tuple(destinations), mask.clone())) + + def replicate(self): + self.replicate_calls += 1 + + cfgs = [SimpleNamespace(prim_path=f"/World/envs/env_.*/asset_{i}") for i in range(5)] + for cfg in cfgs: + REPLICATION_QUEUE.append((cfg, FakeCtx)) + + plan = ClonePlan( + sources=("/World/envs/env_0",), + destinations=("/World/envs/env_{}",), + clone_mask=torch.ones((1, 3), dtype=torch.bool, device=sim.cfg.device), + env_ids=torch.arange(3, dtype=torch.long, device=sim.cfg.device), + positions=grid_transforms(3, 1.0, device=sim.cfg.device)[0], + cfg_rows={id(cfg): (0,) for cfg in cfgs}, + ) + + replicate(plan, stage=sim_utils.get_current_stage()) + + assert len(FakeCtx.instances) == 1 + ctx = FakeCtx.instances[0] + assert len(ctx.queue_calls) == 1, "shared rows should collapse to one queue_mapping per backend" + sources, _destinations, mask = ctx.queue_calls[0] + assert sources == ("/World/envs/env_0",) + assert mask.shape == (1, 3) + assert ctx.replicate_calls == 1 + + def test_replicate_runs_lower_priority_backends_first(sim): """Sort order: lower replicate_priority runs first (physics before USD).""" From 6aef557fbcaefe3c06a23224695bd921dbde21e1 Mon Sep 17 00:00:00 2001 From: Octi Zhang Date: Tue, 26 May 2026 01:38:20 -0700 Subject: [PATCH 6/7] Exclude threedworld.org from lychee link checker Host returns HTTP 403 to lychee's user agent on every run, blocking the "Check for Broken Links" job on unrelated docs (`ecosystem.rst:212`). Same pattern used for other bot-blocking hosts already in the exclude list (stackoverflow.com, helm.ngc.nvidia.com, etc.). --- .github/workflows/check-links.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/check-links.yml b/.github/workflows/check-links.yml index edec5899aa3e..bf54b0d9c74a 100644 --- a/.github/workflows/check-links.yml +++ b/.github/workflows/check-links.yml @@ -104,6 +104,7 @@ jobs: --exclude 'docs\.ray\.io' --exclude 'docs\.conda\.io' --exclude 'stackoverflow\.com' + --exclude 'threedworld\.org' --max-retries 5 --retry-wait-time 10 --timeout 20 From 090cdccb95688c405dee75bc80efcf852c55f7dd Mon Sep 17 00:00:00 2001 From: Octi Zhang Date: Tue, 26 May 2026 14:58:53 -0700 Subject: [PATCH 7/7] fix guassian test --- .../test/sensors/generate_synthetic_gaussian_asset.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/source/isaaclab/test/sensors/generate_synthetic_gaussian_asset.py b/source/isaaclab/test/sensors/generate_synthetic_gaussian_asset.py index 9b67797dfca1..62f658b5607f 100644 --- a/source/isaaclab/test/sensors/generate_synthetic_gaussian_asset.py +++ b/source/isaaclab/test/sensors/generate_synthetic_gaussian_asset.py @@ -28,6 +28,7 @@ from pxr import Gf, Sdf, Usd, UsdGeom +import isaaclab.cloner as cloner import isaaclab.sim as sim_utils from isaaclab.assets import AssetBaseCfg, RigidObjectCfg from isaaclab.scene import InteractiveScene, InteractiveSceneCfg @@ -556,6 +557,14 @@ def render_synthetic_gaussian_scene( renderer_cfg=renderer_cfg, ) camera = Camera(cfg) + # Camera is constructed after the scene's ReplicateSession has exited, so its + # queued USD replication needs an explicit drain (Path B). Reuse the scene's + # env positions so env_origins stays consistent. + published = sim.get_clone_plan() + positions = published.positions if published is not None else None + src, dst = "/World/envs/env_0", "/World/envs/env_{}" + camera_plan = cloner.ClonePlan.from_env_0(src, dst, num_envs, str(sim.device), positions) + cloner.replicate(camera_plan, stage=sim.stage) sim.reset() for _ in range(stabilisation_steps): sim.step()