[OvPhysX] Snapshot env_0 stage pre-clone to speed up OvPhysX warm-up

source/isaaclab/changelog.d/ovphysx-frameview-dispatch.rst

@@ -0,0 +1,8 @@
+Fixed
+^^^^^
+
+* Fixed :class:`~isaaclab.sim.views.FrameView` dispatch under the OVPhysX
+  backend. ``FrameView(...)`` now routes to
+  :class:`~isaaclab_ovphysx.sim.views.OvPhysxFrameView` instead of silently
+  falling through to ``FabricFrameView``, which returned stale USD spawn
+  poses for sensor frames riding on physics bodies.

source/isaaclab/isaaclab/scene/interactive_scene.py

@@ -166,10 +166,14 @@ def __init__(self, cfg: InteractiveSceneCfg):
             clone_in_fabric=self.cfg.clone_in_fabric,
             device=self.device,
             physics_clone_fn=physics_clone_fn,
-            # For ovphysx: env_1..N are created by physx.clone() in the physics
-            # runtime after add_usd().  USD replication of the asset hierarchy
-            # to env_1..N is skipped — only env_0 needs physics prims in the USD.
-            clone_usd=not self.physics_backend.startswith("ovphysx"),
+            # 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,
         )
 
         # create source prim

source/isaaclab/isaaclab/sim/views/frame_view.py

@@ -25,11 +25,17 @@ class FrameView(FactoryBase, BaseFrameView):
 
     - **PhysX / no backend**: :class:`~isaaclab_physx.sim.views.FabricFrameView`
       (Fabric GPU acceleration with USD fallback).
+    - **OVPhysX**: :class:`~isaaclab_ovphysx.sim.views.OvPhysxFrameView`
+      (Warp-native, reads ``body_q`` via the OVPhysX scene data provider).
     - **Newton**: :class:`~isaaclab_newton.sim.views.NewtonSiteFrameView`
-      (GPU-resident site-based transforms).
+      (Warp-native, reads ``body_q`` from the Newton state).
     """
 
-    _backend_class_names = {"physx": "FabricFrameView", "newton": "NewtonSiteFrameView"}
+    _backend_class_names = {
+        "physx": "FabricFrameView",
+        "ovphysx": "OvPhysxFrameView",
+        "newton": "NewtonSiteFrameView",
+    }
 
     @classmethod
     def _get_backend(cls, *args, **kwargs) -> str:
@@ -41,6 +47,8 @@ def _get_backend(cls, *args, **kwargs) -> str:
         manager_name = ctx.physics_manager.__name__.lower()
         if "newton" in manager_name:
             return "newton"
+        if "ovphysx" in manager_name:
+            return "ovphysx"
         return "physx"
 
     def __new__(cls, *args, **kwargs) -> BaseFrameView:

source/isaaclab_ovphysx/changelog.d/ovphysx-frameview.minor.rst

@@ -0,0 +1,11 @@
+Added
+^^^^^
+
+* Added :class:`~isaaclab_ovphysx.sim.views.OvPhysxFrameView`, a
+  Warp-native batched-prim view that reads world poses from the OVPhysX
+  scene data provider's ``body_q`` array. Mirrors
+  :class:`~isaaclab_newton.sim.views.NewtonSiteFrameView` in semantics
+  and API: ``set_world_poses`` / ``set_local_poses`` update the view's
+  internal ``site_local`` buffer and never mutate the physics state.
+  Scales and visibility delegate to a lazy internal
+  :class:`~isaaclab.sim.views.UsdFrameView`.

source/isaaclab_ovphysx/changelog.d/ovphysx-pre-clone-snapshot.rst

@@ -0,0 +1,13 @@
+Fixed
+^^^^^
+
+* Fixed slow OVPhysX warmup at large env counts. The previous
+  :meth:`~isaaclab_ovphysx.physics.OvPhysxManager._warmup_and_load`
+  flattened the entire post-clone USD stage to disk before stripping
+  ``env_1..N`` from the resulting file — a ~31 s flatten at 4096 envs
+  on Anymal-D Rough. The manager now snapshots the live stage from
+  inside :func:`~isaaclab_ovphysx.cloner.ovphysx_replicate` (which runs
+  before :func:`cloner.usd_replicate` inflates the stage), and
+  :meth:`_warmup_and_load` consumes the snapshot directly. The old
+  export-and-strip path is preserved as a fallback for callers that do
+  not go through :meth:`isaaclab.scene.InteractiveScene.clone_environments`.

source/isaaclab_ovphysx/isaaclab_ovphysx/cloner/ovphysx_replicate.py

@@ -69,6 +69,13 @@ def ovphysx_replicate(
     # Deferred import to avoid circular dependency at module load time.
     from isaaclab_ovphysx.physics.ovphysx_manager import OvPhysxManager
 
+    # Snapshot the live stage now -- USD cloning hasn't fired yet, so the
+    # stage holds only env_0's authored content and ``stage.Export`` flattens
+    # a small stage rather than the post-clone 4096-env version.
+    # :meth:`OvPhysxManager._warmup_and_load` consumes this snapshot directly,
+    # bypassing the slower export+strip fallback.
+    OvPhysxManager.register_pre_clone_stage_snapshot(stage)
+
     for i, src in enumerate(sources):
         active_env_ids = env_ids[mapping[i]].tolist()
 

source/isaaclab_ovphysx/isaaclab_ovphysx/physics/ovphysx_manager.py

@@ -215,6 +215,13 @@ class OvPhysxManager(PhysicsManager):
     # 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]]]]] = []
+    # Pre-clone stage snapshot path.  Populated by
+    # :meth:`register_pre_clone_stage_snapshot` from inside
+    # :func:`~isaaclab_ovphysx.cloner.ovphysx_replicate`, at the moment the
+    # live USD stage carries only env_0 (USD-clone hasn't fired yet).
+    # :meth:`_warmup_and_load` consumes the snapshot directly, avoiding the
+    # post-clone full-stage flatten cost (~31s at 4096 envs on Anymal-D).
+    _pre_clone_stage_path: ClassVar[str | None] = None
     _atexit_registered: ClassVar[bool] = False
     _scene_data_backend: ClassVar[OvPhysxSceneDataBackend | None] = None
 
@@ -245,6 +252,66 @@ def register_clone(
         """
         cls._pending_clones.append((source, targets, parent_positions or []))
 
+    @classmethod
+    def register_pre_clone_stage_snapshot(cls, stage: Any) -> None:
+        """Snapshot the live USD stage to disk before USD cloning inflates it.
+
+        Called by :func:`~isaaclab_ovphysx.cloner.ovphysx_replicate`, which
+        runs *before* :func:`cloner.usd_replicate` in
+        :meth:`isaaclab.scene.InteractiveScene.clone_environments`. At that
+        moment, the live stage carries only ``env_0``'s authored content (plus
+        globals); ``stage.Export`` therefore flattens a small stage rather than
+        the post-clone 4096-env version.
+
+        :meth:`_warmup_and_load` consumes the snapshot path directly, bypassing
+        the slow ``_export_env0_only_stage`` fallback (which exports the full
+        post-clone stage and then strips ``env_1..N`` from the resulting file —
+        a ~31s flatten at 4096 envs on Anymal-D).
+
+        Idempotent: subsequent calls within the same setup are no-ops.
+
+        Args:
+            stage: Live USD stage held by :class:`SimulationContext`.
+        """
+        if cls._pre_clone_stage_path is not None:
+            return
+        if cls._tmp_dir is None:
+            cls._tmp_dir = tempfile.TemporaryDirectory(prefix="isaaclab_ovphysx_")
+        snapshot_path = os.path.join(cls._tmp_dir.name, "scene_env0.usda")
+        stage.Export(snapshot_path)
+        cls._pre_clone_stage_path = snapshot_path
+        logger.info("OvPhysxManager: pre-clone stage snapshot exported to %s", snapshot_path)
+
+    _physx_schemas_registered: ClassVar[bool] = False
+
+    @classmethod
+    def _ensure_physx_schemas_registered(cls) -> None:
+        """Register the ``PhysxSchema`` USD plugin shipped with the ovphysx wheel.
+
+        In Kit-based runs ``omni.physx`` registers the schema; in kitless
+        runs it must be registered manually before the wheel can match
+        ``PhysxContactReportAPI`` and friends on the stage.  The wheel
+        bundles the plugin under ``ovphysx/plugins/usd/PhysxSchema``.  This
+        method is idempotent — :meth:`pxr.Plug.Registry.RegisterPlugins`
+        is a no-op once the plugin is registered.
+        """
+        if cls._physx_schemas_registered:
+            return
+        try:
+            import os  # noqa: PLC0415
+
+            import ovphysx  # noqa: PLC0415
+
+            from pxr import Plug  # noqa: PLC0415
+        except Exception:
+            return
+        plugin_root = os.path.join(os.path.dirname(ovphysx.__file__), "plugins", "usd")
+        for sub in ("PhysxSchema/resources", "PhysxSchemaAddition/resources"):
+            path = os.path.join(plugin_root, sub)
+            if os.path.isdir(path):
+                Plug.Registry().RegisterPlugins(path)
+        cls._physx_schemas_registered = True
+
     @classmethod
     def initialize(cls, sim_context: SimulationContext) -> None:
         """Initialize the physics manager with simulation context.
@@ -265,6 +332,7 @@ def initialize(cls, sim_context: SimulationContext) -> None:
         cls._warmup_done = False
         cls._usd_handle = None
         cls._stage_path = None
+        cls._pre_clone_stage_path = None
         cls._pending_clones = []
         # Construct the SceneDataBackend eagerly so :class:`SimulationContext`
         # captures a real instance (not ``None``) when it builds the central
@@ -374,6 +442,84 @@ def get_scene_data_backend(cls) -> SceneDataBackend:
     # Internal helpers
     # ------------------------------------------------------------------
 
+    @staticmethod
+    def _export_env0_only_stage(sim_stage: Any, target_file: str) -> None:
+        """Export the simulation stage to ``target_file`` with env_1..N stripped.
+
+        Writes a USD file containing every prim under the live stage **except**
+        ``/World/envs/env_<i>`` for ``i != 0``. Globals (``/physicsScene``,
+        ``/World/ground``, lights, materials, etc.) and ``/World/envs/env_0`` are
+        retained.  ``physx.clone()`` is then expected to repopulate env_1..N at
+        the physics layer with proper clone lineage so that subsequent
+        ``create_tensor_binding`` calls hit the wheel's fast path.
+
+        Implementation: export the full stage to disk, then re-open the result
+        as an :class:`Sdf.Layer` and delete env_1..N prim specs in place.  This
+        avoids mutating the live stage (which other consumers -- sensors,
+        visualizers -- still see in its full N-env form).
+
+        Limitations:
+            * **Homogeneous-env assumption.** Every env is treated as an
+              identical copy of env_0 from the physics runtime's point of view.
+              Anything authored *only* under ``/World/envs/env_<i>`` for
+              ``i != 0`` (per-env mass overrides, per-env friction, per-env
+              collision filters, etc.) is dropped from the file handed to
+              ``physx.add_usd`` and therefore not seen by PhysX. Sensors and
+              visualizers still see those overrides in USD (the live stage is
+              unmodified), so a divergence is possible.  Per-env physics state
+              must instead be written via the runtime APIs
+              (``RigidObject.write_root_state_to_sim_index``, etc.).
+            * **Global path convention.** Any physics-relevant prim that lives
+              under ``/World/envs/env_<i!=0>/`` (e.g. an asset-specific
+              ``PhysicsScene``, a per-env material) gets stripped. Globals must
+              live outside ``/World/envs`` (or under ``/World/envs/env_0``) to
+              survive the export.
+            * **Static topology.** Envs added or removed at runtime after
+              warmup are not supported by ``physx.clone()`` lineage and would
+              require a re-warmup with a re-exported stage.
+
+        Args:
+            sim_stage: Live USD stage held by ``SimulationContext``.
+            target_file: Output ``.usda`` file path.  Overwritten if it exists.
+        """
+        from pxr import Sdf  # noqa: PLC0415
+
+        # Step 1: full flatten-export of the live stage.  We pass the full file
+        # to ``Sdf.Layer.OpenAsAnonymous`` so the edits below don't write back
+        # to the source layer on disk.
+        sim_stage.Export(target_file)
+
+        # Step 2: open the exported file as an editable Sdf layer and delete
+        # ``/World/envs/env_<digits>`` children for digits != 0.  Walking the
+        # ``/World/envs`` ``PrimSpec``'s ``nameChildren`` keeps us scoped to
+        # the env-namespace and leaves the rest of the stage untouched.
+        layer = Sdf.Layer.FindOrOpen(target_file)
+        if layer is None:
+            raise RuntimeError(
+                f"OvPhysxManager: failed to re-open exported USD layer at {target_file!r} for env-scoping."
+            )
+        envs_spec = layer.GetPrimAtPath("/World/envs")
+        if envs_spec is None or not envs_spec:
+            # No /World/envs in the stage (single-env or non-IsaacLab scene); nothing to scope.
+            logger.debug("OvPhysxManager: no /World/envs prim — exported stage as-is.")
+            return
+
+        env_name_re = re.compile(r"^env_(\d+)$")
+        names_to_remove = [
+            child_name
+            for child_name in list(envs_spec.nameChildren.keys())
+            if (match := env_name_re.match(child_name)) and match.group(1) != "0"
+        ]
+        for child_name in names_to_remove:
+            del envs_spec.nameChildren[child_name]
+
+        if names_to_remove:
+            layer.Export(target_file)
+            logger.info(
+                "OvPhysxManager: stripped %d env_<i!=0> subtrees from exported USD (kept env_0 + globals)",
+                len(names_to_remove),
+            )
+
     @classmethod
     def _warmup_and_load(cls) -> None:
         """Export the USD stage and load it into the ovphysx runtime.
@@ -416,12 +562,36 @@ def _warmup_and_load(cls) -> None:
         if scene_prim.IsValid():
             cls._configure_physx_scene_prim(scene_prim, PhysicsManager._cfg, ovphysx_device)
 
-        # Export the current USD stage to a temporary file so ovphysx can load it.
-        cls._tmp_dir = tempfile.TemporaryDirectory(prefix="isaaclab_ovphysx_")
-        stage_file = os.path.join(cls._tmp_dir.name, "scene.usda")
-        sim.stage.Export(stage_file)
+        # Resolve the USD stage path handed to ``physx.add_usd``.
+        #
+        # When ``InteractiveScene`` runs with ``clone_usd=True``, the live USD
+        # stage carries env_0..N's full asset subtrees as authored copies.
+        # Handing that stage to ``physx.add_usd`` would make the wheel ingest
+        # all 4096 envs as independent USD-defined bodies, defeating the
+        # ``physx.clone()`` fast path and turning every subsequent
+        # ``create_tensor_binding`` call into an O(N) USD enumeration -- the
+        # hang you'd see at large env counts.
+        #
+        # Fast path: ``ovphysx_replicate`` runs *before* ``cloner.usd_replicate``
+        # in ``InteractiveScene.clone_environments``, so when it called
+        # :meth:`register_pre_clone_stage_snapshot` the live stage carried only
+        # env_0. That snapshot is the env_0-scoped USD the wheel needs; use it
+        # directly when available.
+        #
+        # Fallback: if no pre-clone snapshot was registered (e.g. tests that
+        # don't go through ``InteractiveScene.clone_environments``), export the
+        # current (potentially post-clone) stage and strip ``env_1..N`` from
+        # the resulting file -- correct, but with a flatten cost proportional
+        # to the post-clone stage size.
+        if cls._pre_clone_stage_path is not None and os.path.exists(cls._pre_clone_stage_path):
+            stage_file = cls._pre_clone_stage_path
+            logger.info("OvPhysxManager: using pre-clone stage snapshot at %s", stage_file)
+        else:
+            cls._tmp_dir = tempfile.TemporaryDirectory(prefix="isaaclab_ovphysx_")
+            stage_file = os.path.join(cls._tmp_dir.name, "scene.usda")
+            cls._export_env0_only_stage(sim.stage, stage_file)
+            logger.info("OvPhysxManager: exported env_0-scoped USD stage to %s", stage_file)
         cls._stage_path = stage_file
-        logger.info("OvPhysxManager: exported USD stage to %s", stage_file)
 
         if cls._physx is None:
             cls._construct_physx(ovphysx_device, gpu_index)

source/isaaclab_ovphysx/isaaclab_ovphysx/sim/__init__.py

@@ -0,0 +1,10 @@
+# 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 simulation views."""
+
+from isaaclab.utils.module import lazy_export
+
+lazy_export()

source/isaaclab_ovphysx/isaaclab_ovphysx/sim/__init__.pyi

@@ -0,0 +1,6 @@
+# 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
+
+__all__: list[str] = []

source/isaaclab_ovphysx/isaaclab_ovphysx/sim/views/__init__.py

@@ -0,0 +1,10 @@
+# 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 simulation views."""
+
+from isaaclab.utils.module import lazy_export
+
+lazy_export()

source/isaaclab_ovphysx/isaaclab_ovphysx/sim/views/__init__.pyi

@@ -0,0 +1,10 @@
+# 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
+
+__all__ = [
+    "OvPhysxFrameView",
+]
+
+from .ovphysx_frame_view import OvPhysxFrameView