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

[AntoineRichard]

Description

Cuts ~31 seconds off OVPhysX warmup at 4096 envs.

The OvPhysX warmup path in :meth:~isaaclab_ovphysx.physics.OvPhysxManager._warmup_and_load previously called stage.Export(target_file) on the live USD stage and then opened the resulting file as an :class:Sdf.Layer to delete /World/envs/env_<i!=0> from it. With clone_usd=True (introduced by PR #5678), the live stage carries every per-env subtree authored, so the export step had to flatten 4096 envs to disk before the strip could discard most of them. Measured cost on Isaac-Velocity-Rough-Anymal-D-v0 at 4096 envs: ~31 s of pure setup.

:func:~isaaclab_ovphysx.cloner.ovphysx_replicate runs from :meth:~isaaclab.scene.InteractiveScene.clone_environments before :func:cloner.usd_replicate inflates the stage. At that moment the live stage holds only env_0's authored content (plus globals), and stage.Export is cheap. This PR snapshots the stage there and has :meth:_warmup_and_load consume the snapshot directly. The old export-and-strip path is preserved as a fallback for callers that bypass InteractiveScene.clone_environments (single-env tests, etc.).

Depends on PR #5678 (FrameView) for the existing _export_env0_only_stage fallback path.

Fixes # (issue)

Type of change

• New feature (non-breaking change which adds functionality)

Screenshots

N/A — backend init perf, no UI surface.

Checklist

• I have read and understood the contribution guidelines
• I have run the pre-commit checks with ./isaaclab.sh --format
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• I have added a changelog fragment under source/<pkg>/changelog.d/ for every touched package (do not edit CHANGELOG.rst or bump extension.toml — CI handles that)
• I have added my name to the CONTRIBUTORS.md or my name already exists there

[isaaclab-review-bot]

🤖 Isaac Lab Review Bot — Code Review Summary

Overview

This PR introduces a significant performance optimization (~31s speedup at 4096 envs) by snapshotting the USD stage before clone inflation, along with a new OvPhysxFrameView for GPU-native pose queries.

---

✅ Strengths

1. Sound Architecture: The pre-clone snapshot approach correctly identifies that ovphysx_replicate runs before usd_replicate, making it the ideal interception point.

2. Defensive Fallback: The _export_env0_only_stage fallback preserves compatibility for callers that bypass InteractiveScene.clone_environments.

3. Well-Documented Limitations: The docstring for _export_env0_only_stage clearly outlines the homogeneous-env assumption, global path convention, and static topology constraints.

4. Comprehensive OvPhysxFrameView: The new FrameView implementation includes proper handling of both clone_usd=True and clone_usd=False modes with synthetic path expansion.

5. Solid Test Coverage: Contract-suite integration and specific deferred-init tests are included.

---

🔴 Critical Issues

None identified — the implementation is sound.

---

🟡 Important Suggestions

1. State Reset Completeness (ovphysx_manager.py)

   In initialize(), _pre_clone_stage_path is reset to None, but if a previous run's temp file still exists (e.g., the process didn't clean up), and _tmp_dir is reused, there could be stale file confusion. Consider explicitly cleaning the snapshot file in initialize():

   if cls._pre_clone_stage_path and os.path.exists(cls._pre_clone_stage_path):
       try:
           os.remove(cls._pre_clone_stage_path)
       except OSError:
           pass
   cls._pre_clone_stage_path = None

2. Idempotency Edge Case (register_pre_clone_stage_snapshot)

   The method is documented as idempotent, but the guard if cls._pre_clone_stage_path is not None: return doesn't verify the file still exists. If the file was deleted between calls, subsequent warmup would fail. Consider:

   if cls._pre_clone_stage_path is not None:
       if os.path.exists(cls._pre_clone_stage_path):
           return
       # File was deleted; allow re-snapshot
       cls._pre_clone_stage_path = None

3. Layer Caching Clarification (_export_env0_only_stage)

   Using Sdf.Layer.FindOrOpen may return a cached layer from a previous session. Consider adding a comment explaining this is intentional, or use Sdf.Layer.OpenAsAnonymous if cache isolation is preferred.

---

🟢 Minor Notes

1. Import Style: The import os inside _ensure_physx_schemas_registered() is redundant since os is already imported at module level. Consider removing the local import.

2. Logging Consistency: logger.info is used for snapshot creation but logger.debug for the "no /World/envs" case. Consider consistent log levels for similar operational messages.

3. Magic Constant: WORLD_BODY_INDEX = -1 is well-named but could benefit from a brief docstring explaining it's used in Warp kernels to indicate world-attached sites.

4. Test Fixture Clarity: The test file comment about detaching the binding for synthetic pose injection is helpful — consider adding a similar note in the OvPhysxFrameView class docstring about the binding lifecycle.

---

📋 Test Coverage Assessment

Covered:

• Factory dispatch to OvPhysxFrameView ✓
• Deferred initialization guard ✓
• Shared FrameView contract suite ✓

Gaps to Consider (non-blocking):

• No explicit test for the pre-clone snapshot path vs. fallback path
• No test for _export_env0_only_stage with edge cases (single env, missing /World/envs)
• No benchmark test validating the 31s speedup claim

---

Verdict: Approve ✅

This is a well-designed performance optimization with proper fallback handling and good test coverage. The suggestions above are minor improvements that don't block merging. The ~31s speedup at scale is a meaningful improvement for training workflows.

[greptile-apps]

Greptile Summary

This PR cuts ~31 s from OVPhysX warmup at large env counts by snapshotting the USD stage inside ovphysx_replicate — before usd_replicate inflates it to 4096 envs — and having _warmup_and_load consume the small env_0 snapshot directly instead of flattening the full post-clone stage. It also adds OvPhysxFrameView, wires it into the FrameView factory, and changes clone_usd to True unconditionally for all backends.

• Pre-clone snapshot fast path: register_pre_clone_stage_snapshot writes scene_env0.usda before USD cloning fires; _warmup_and_load uses this path when available and falls back to the slower _export_env0_only_stage strip-and-re-export for callers that bypass InteractiveScene.
• OvPhysxFrameView: Warp-native, GPU-resident site-pose view using a RIGID_BODY_POSE tensor binding; supports both clone_usd=True and clone_usd=False expansion modes.
• clone_usd=True for OVPhysX: Previously skipped USD replication for OVPhysX; now runs it for all backends. The PR's own comment flags this as an in-progress assumption that may need reverting.

Confidence Score: 3/5

The warmup snapshot logic is sound for the happy path, but two defects in the fallback and kitless code paths need fixing before merge.

The new _ensure_physx_schemas_registered method is added but never called anywhere in the codebase. Its docstring says kitless runs require it to resolve PhysxContactReportAPI and related schemas — without a call site, those runs silently lose schema resolution. The _export_env0_only_stage fallback uses Sdf.Layer.FindOrOpen without a subsequent Reload(), so a re-warmup in the same process with a reused tmpdir path could feed the wrong prim hierarchy to physx.add_usd. The clone_usd=True change for OVPhysX is acknowledged as exploratory without CI test coverage for the failure path.

ovphysx_manager.py (call site for _ensure_physx_schemas_registered missing; Sdf.Layer.FindOrOpen stale-cache risk in _export_env0_only_stage); interactive_scene.py (unconditional clone_usd=True for OVPhysX needs test coverage)

Important Files Changed

Filename	Overview
source/isaaclab_ovphysx/isaaclab_ovphysx/physics/ovphysx_manager.py	Core change: adds register_pre_clone_stage_snapshot fast path and _export_env0_only_stage fallback. Issues: _ensure_physx_schemas_registered is added but never called; Sdf.Layer.FindOrOpen in the fallback may return a stale cached layer on re-warmup.
source/isaaclab_ovphysx/isaaclab_ovphysx/cloner/ovphysx_replicate.py	Adds a single call to OvPhysxManager.register_pre_clone_stage_snapshot(stage) before the USD-clone loop. Idempotency handled in the callee. Change is minimal and correct.
source/isaaclab/isaaclab/scene/interactive_scene.py	Unconditionally sets clone_usd=True for all backends including OVPhysX. The comment acknowledges this is exploratory and may need reverting.
source/isaaclab_ovphysx/isaaclab_ovphysx/sim/views/ovphysx_frame_view.py	New 890-line OvPhysxFrameView class. Warp kernels are correct; deferred init pattern is sound; clone_usd=False expansion logic correctly handles env_1..N via env_token substitution.
source/isaaclab/isaaclab/sim/views/frame_view.py	Adds "ovphysx": "OvPhysxFrameView" to the factory dispatch map and a _get_backend check. Straightforward and correct.
source/isaaclab_ovphysx/test/sim/test_views_xform_prim_ovphysx.py	Tests factory dispatch and deferred-init error path. The fixture detaches the live binding after one read to allow synchronous get/set testing — a reasonable test scaffold.

Sequence Diagram

sequenceDiagram
    participant IS as InteractiveScene
    participant OR as ovphysx_replicate
    participant UR as usd_replicate
    participant OM as OvPhysxManager
    participant Stage as USD Stage

    IS->>OR: "physics_clone_fn() [stage = env_0 only]"
    OR->>OM: register_pre_clone_stage_snapshot(stage)
    OM->>Stage: stage.Export(scene_env0.usda)
    Note over OM: _pre_clone_stage_path set
    OR->>OM: register_clone(source, targets, positions)
    IS->>UR: usd_replicate() [inflates stage to env_0..N]
    Note over Stage: Stage now holds env_0..N USD prims
    IS->>OM: "reset() -> _warmup_and_load()"
    alt fast path
        OM-->>OM: "stage_file = scene_env0.usda"
    else fallback
        OM->>Stage: Export full stage
        OM->>OM: _export_env0_only_stage() strip env_1..N
    end
    OM->>OM: physx.add_usd(stage_file)
    loop pending_clones
        OM->>OM: physx.clone(source, targets)
    end
    OM->>OM: physx.warmup_gpu()
    OM-->>IS: PhysicsEvent.PHYSICS_READY

Loading

Comments Outside Diff (1)

1. source/isaaclab/isaaclab/scene/interactive_scene.py, line 26-34 (link)

   Uncertainty baked into production via "revert if it breaks" comment

   The comment "Probing whether this assumption holds in practice; revert to not startswith("ovphysx") if physx.clone() errors on already-populated targets" signals that the behavioral change from clone_usd=False to clone_usd=True for the OVPhysX backend has not been validated in all important cases. Embedding an exploratory probe in production code without test coverage for the failure path makes regressions hard to detect. At a minimum this should be gated by a CI test for the OVPhysX + clone_usd=True path, or the comment should be converted to a tracked issue.

_{Reviews (1): Last reviewed commit: "Snapshot env_0 stage pre-clone to skip f..." | Re-trigger Greptile}

[greptile-apps]

_ensure_physx_schemas_registered has no call site

This method was added in this PR but is never called anywhere in the codebase (confirmed via repo-wide search). Its docstring explicitly states that in kitless runs "it must be registered manually before the wheel can match PhysxContactReportAPI and friends on the stage." Without a call in initialize() or _warmup_and_load(), kitless runs would silently skip schema registration and any USD stage with PhysxContactReportAPI, PhysxJointAPI, etc. would not resolve those schemas — causing silent failures when loading physics prims.

[greptile-apps]

After sim_stage.Export(target_file), using Sdf.Layer.FindOrOpen(target_file) looks up the path in the Sdf layer registry first. On a re-warmup in the same process (where _tmp_dir is reused and the path is identical), the registry would return the cached layer with old content rather than the freshly exported file. Sdf.Layer.Reload() must be called on the found layer if it was already open, or the file should be opened anonymously to bypass the registry entirely.

Suggested change

	layer = Sdf.Layer.FindOrOpen(target_file)
	if layer is None:
	# Use FindOrOpen and unconditionally reload so that a re-warmup in the
	# same process (reusing the same tmp path) sees the freshly-exported file
	# rather than a stale cached layer from the Sdf registry.
	layer = Sdf.Layer.FindOrOpen(target_file)
	if layer is not None:
	layer.Reload()
	if layer is None:

[greptile-apps]

import os is a standard-library module that will never fail to import. Including it inside the broad except Exception: return block means any ImportError from ovphysx or pxr.Plug silently suppresses schema registration, but the root cause (missing wheel or missing pxr) is completely hidden. Move import os outside the try block.

Suggested change

	try:
	import os # noqa: PLC0415
	import ovphysx # noqa: PLC0415
	from pxr import Plug # noqa: PLC0415
	except Exception:
	return
	try:
	import ovphysx # noqa: PLC0415
	from pxr import Plug # noqa: PLC0415
	except Exception:
	return