Shimmed Ovstage integration into OVRTX Renderer

[rilei-nvidia]

Description

A shimmed ovstage implementation that we can optionally merge into develop to make it easier to collectively test and iterate, at the cost of supporting 2 separate code paths in the meantime (with agents it is not too much work). Once ovstage is considered ready for IsaacLab, it should be very straightforward to remove the legacy codepaths. This work was mostly agentic - this skills and context in https://github.com/NVIDIA-Omniverse/ovrtx made it straightforward.

ovstage code path is disabled by default and can be enabled via ISAAC_LAB_OVRTX_USE_OVSTAGE=1.

Example of running the cartpole rendering correctness tests with ovrtx + ovstage:

ISAAC_LAB_OVRTX_USE_OVSTAGE=1 ./isaaclab.sh -p -m pytest -rA source/isaaclab_tasks/test/core/test_rendering_cartpole_kitless.py

The ovstage implementation still suffers from the tiled rendering bug introduced in ovrtx-0.4 that @nvsekkin was looking at:

In order to further test ovstage without being bitten by the tiled rendering bug, I introduced a separate ISAAC_LAB_OVRTX_ENABLE_CLONE=0 env var to disable OVRTX cloning (ovrtx.clone_usd / ovstage.clone) which allows us to at least have working renders we can benchmark. (@r-schmitt @pbarejko)

Summarized differences

Feature	Legacy	Ovstage
Scene ownership	Renderer.open_usd_from_string(usd) — renderer owns scene data internally	ovstage.Stage + PathDictionary created first, USD loaded via population.open_usd_from_string, renderer attached via attach_ovstage. Scene data lives in ovstage; renderer is a consumer
Stage export	Non-source env subtrees temporarily deactivated on root layer; env_1…N-1 root prims remain in USD as inactive stubs	Flattened to anonymous Sdf.Layer copy; non-source env roots removed entirely via Sdf.BatchNamespaceEdit. Required because stage.clone rejects pre-existing targets
Cloning	Renderer.clone_usd(source, targets) — single call per row, renderer handles internally, stubs may already exist	Env-root xforms pre-captured from the live USD stage before export via UsdGeom.XformCache (_capture_env_root_xforms_ovstage) — export removes non-source env roots, so they can't be read back from ovstage. Then stage.clone(source, targets, ordinal) per row (shared ordinal for atomicity), and the captured xforms are written back with write_attribute("omni:xform", semantic=MATRIX) so hierarchy-positioned static prims stay correctly placed per env
Rigid body / camera transforms	bind_attribute returns persistent binding; each frame maps into GPU memory via binding.map(Device.CUDA) and Warp kernel writes directly — zero-copy, GPU-resident	Persistent query via query_from_path_list; each frame: kernel → wp.synchronize_device → .numpy() CPU copy → stage.write_attribute with _xform_tensor_from_numpy
Deformable writes	bind_array_attribute + binding.write — GPU-to-GPU, cross-stream dependency, no host stall	query_from_path_list + wp.synchronize_device (host stall) + particle_q.numpy() + stage.write_attribute — host copy required to rebuild the warp vec3f slices as lanes=3 DLTensors (semantic POINT) matching the point3f[] column; make_dltensor won't override dtype on GPU DLPack producers, so no zero-copy path
Sequencing	None — renderer serializes internally	Monotonically increasing _current_ordinal; all init writes share one ordinal committed via advance_write_floor before attach_ovstage; all perordinal, committed in render before step(ordinal=N), then incremented for the next frame
Cleanup	binding.unbind() per binding, Renderer.reset_stage()	stage.release_query per query, PathDictionary.destroy_path_list per path list, renderer.detach_ovstage() before `ExitStack.c
Zero-copy GPU tensor ingestion	Transform writes arpersistent binding map	GPU→CPU copy required every frame — ovstage make_dltensor doesn't support lane-packed dtype :xformmatrices, lanes=3 forpoints`) on GPU DLPackproducers

Notable detailed differences

Env-root transform snapshot: legacy vs ovstage

After cloning, we need to restore each env root's original transform so that static prims (tables, fixtures, etc.) that aren't driven by physics remain correctly placed per environment.

The two paths diverge in when that env root xform snapshot can be taken:

• Legacy - Renderer.clone_usd tolerates pre-existing target prims, so the stage exported to the renderer still contains env_1…N-1 as de-activated. read_attribute("omni:xform", env_prim_paths) can be called on the loaded stage immediately before cloning.
• Ovstage — stage.clone requires targets to be absent, so export_stage_to_string_ovstage removes env_1…N-1 entirely via Sdf.BatchNamespaceEdit before the string is loaded into ovstage. By the time _clone_sources_ovstage runs, those prims don't exist in the stage and can't be read back. The snapshot must be taken earlier, from the live USD stage in prepare_stage before the export strips them, and stored on the renderer until _clone_sources_ovstage consumes it.
  • The fact that for clone targets must be deleted before we execute clone means that the workaround for preserving camera prims (https://github.com/isaac-sim/IsaacLab/pull/6592/files) will not work for ovstage! FYI @nvsekkin

We'd need to do some profiling around this comparing the legacy way of traversing the env_1...N-1 descendent hierarchies and deactivating all the prims vs flatten + the batch namespace edit prune approach.

Newton GPU transform/points write path: zero-copy vs host round-trip

When the physics engine is Newton - each frame, Isaac Lab computes rigid body and camera world transforms on the GPU using Warp kernels that write into wp.array(dtype=wp.mat44d) buffers — one 4×4 float64 matrix per prim. These buffers stay GPU-resident and need to be handed to the renderer so it can update prim placements before the next step().

• Legacy - Renderer.bind_attribute maps a persistent GPU buffers for the camera & object xforms; the Warp kernel writes directly into it every frame and is picked up by the renderer.

• Ovstage path writes transforms via stage.write_attribute("omni:xform", ...). The omni:xform Fabric column is typed as float64, lanes=16 — a single 16-element vector per prim representing the flattened 4×4 matrix. The problem is that wp.mat44d exports through DLPack as shape (N, 4, 4), dtype=float64, lanes=1, which mismatches the existing column type and causes ovstage to reject the write with INVALID_ARGUMENT: existing attribute 'omni:xform' has a different type.

make_dltensor can reinterpret a buffer as lanes=16 via its dtype and shape override kwargs, but those overrides only work when the input is a numpy array. For any DLPack producer (warp, torch, cupy) it calls DLTensor.from_dlpack(array) directly and the overrides are rejected:

import warp as wp
from ovstage import make_dltensor, DLDataType, DLDataTypeCode

xform_dtype = DLDataType(code=DLDataTypeCode.kDLFloat, bits=64, lanes=16)
gpu_transforms = wp.zeros(4, dtype=wp.mat44d, device="cuda:0")

# Fails: overrides not accepted for DLPack producers
tensor = make_dltensor(gpu_transforms, dtype=xform_dtype, shape=[4])
# ValueError: shape/dtype/ndim/strides overrides are not supported for DLPack producers;
# pass a numpy array to use them

# Works, but forces GPU→CPU copy every frame
cpu_np = gpu_transforms.numpy().reshape(-1, 4, 4)  # synchronize + copy
tensor = make_dltensor(cpu_np.reshape(-1), dtype=xform_dtype, shape=[4])

The current implementation takes the numpy path: wp.synchronize_device + .numpy().reshape(-1, 4, 4) + _xform_tensor_from_numpy. This is a per-frame GPU→CPU copy for every camera and rigid body transform write — the main performance regression relative to the legacy binding path. Once ovstage supports dtype/shape overrides for GPU DLPack producers, the copy can be eliminated and the warp array passed directly.

Similarly for deformable points, Newton particle_q (warp vec3f) dlpack-exports as lanes=1, but the USD points column is lanes=3 so we must copy to host and rebuild it as a lanes=3 DLTensor (semantic POINT), since make_dltensor only honors the lanes override on numpy arrays, not GPU DLPack arrays (@huidongc):

import warp as wp
from ovstage import make_dltensor, DLDataType, DLDataTypeCode

point_dtype = DLDataType(code=DLDataTypeCode.kDLFloat, bits=32, lanes=3)
gpu_points = wp.zeros(4, dtype=wp.vec3f, device="cuda:0")

# Fails: overrides not accepted for DLPack producers
tensor = make_dltensor(gpu_points, dtype=point_dtype, shape=[4])
# ValueError: shape/dtype/ndim/strides overrides are not supported for DLPack producers;
# pass a numpy array to use them

# Works, but forces GPU→CPU copy every frame
cpu_np = gpu_points.numpy().reshape(-1, 3)  # synchronize + copy
tensor = make_dltensor(cpu_np.reshape(-1), dtype=point_dtype, shape=[4])

Type of change

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

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 updated the changelog and the corresponding version in the extension's config/extension.toml file
• I have added my name to the CONTRIBUTORS.md or my name already exists there

[greptile-apps]

Greptile Summary

This PR introduces a shimmed ovstage integration into the OVRTX renderer, adding a parallel code path where ovstage owns scene data and OVRTX consumes it as a rendering-only consumer. The new path is gated by ISAAC_LAB_OVRTX_USE_OVSTAGE=1 and disabled by default, preserving the existing "legacy" renderer-owned scene APIs.

• Dual code paths: All renderer operations (initialize, clone, setup_xform_bindings, update_transforms, update_geometries, update_camera, render, cleanup) are split into _ovstage and _legacy variants dispatched from shared public methods via self._use_ovstage.
• Stage export change: The ovstage variant flattens to an anonymous Sdf.Layer copy and removes non-source env subtrees via Sdf.BatchNamespaceEdit (rather than deactivating them), since stage.clone requires targets to be entirely absent.
• Per-frame GPU→CPU copy: Object and camera transforms must be synchronized and copied to NumPy before each ovstage write because make_dltensor does not support the lanes=16 dtype override for GPU DLPack producers.

Confidence Score: 3/5

The ovstage code path itself is well-structured and the cleanup ordering is correct, but the mandatory ovstage dep in isaaclab_ov/pyproject.toml forces it on all users even though the PR presents it as opt-in.

The mandatory ovstage dep in isaaclab_ov/pyproject.toml directly contradicts the "disabled by default, optional" framing. Every user who installs isaaclab_ov gets ovstage pulled in, even if they never set the env var. The dead null-guards and missing warning in _use_ovstage_enabled are lower-impact but add noise.

source/isaaclab_ov/pyproject.toml (mandatory vs. optional dep mismatch) and source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py (dead null-guards after write_attribute calls, silent env-var fallback).

Important Files Changed

Filename	Overview
source/isaaclab_ov/pyproject.toml	Adds ovstage>=0.1.0,<0.2.0 as a hard (non-optional) dependency — contradicts the optional import in renderer code, the root pyproject.toml optional extra, and the PR's "disabled by default" intent.
source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py	Core change: adds a full parallel ovstage code path (initialize, clone, xform bindings, render, cleanup) behind a _use_ovstage flag. Null guards on queries are placed after first use, making them dead code. Silent fallback when env var is set but ovstage is unavailable should emit a warning.
source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_usd.py	Renames existing export function to export_stage_to_string_legacy, adds export_stage_to_string_ovstage that flattens to anonymous Sdf.Layer and removes non-source env roots via BatchNamespaceEdit, and keeps a backward-compat alias. Logic is clean.
source/isaaclab_ov/test/test_ovrtx_clone_plan.py	Rewrites tests to cover both ovstage and legacy paths via skipif markers; introduces comprehensive fake stubs. Path-agnostic tests unintentionally hardcode the ovstage renderer variant, missing legacy export path coverage.
source/isaaclab_ov/test/test_ovrtx_deformable_bindings.py	Replaces binding-centric stubs with ovstage/legacy-specific fakes and guards all tests with skipif markers; coverage of both paths looks thorough.
pyproject.toml	Bumps ovrtx from >=0.3 to >=0.4, adds ovstage as an optional extra, and updates comments/conflict table accordingly. Changes are consistent with each other.
source/isaaclab/isaaclab/cli/commands/install.py	Adds ovstage to the valid selector set and wires it to the ovstage root extra; straightforward and consistent with existing pattern.
source/isaaclab/test/cli/test_uv_run_pyproject.py	Adds assertion that ovstage version spec in optional extra matches the versions table; consistent with existing test style.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Sim as Simulation
    participant Renderer as OVRTXRenderer
    participant OvStage as ovstage.Stage
    participant OVRTX as ovrtx.Renderer

    Note over Renderer: _use_ovstage = True path

    Sim->>Renderer: prepare_stage(stage, num_envs)
    Renderer->>Renderer: export_stage_to_string_ovstage()

    Sim->>Renderer: create_render_data(spec)
    Renderer->>OvStage: Stage(isaaclab.ovrtx)
    Renderer->>OvStage: "population.open_usd_from_string(ordinal=1)"
    Renderer->>OvStage: "advance_write_floor(ordinal=1)"
    Renderer->>OVRTX: attach_ovstage(stage)
    Renderer->>OvStage: stage.clone(source, targets, ordinal)
    Renderer->>OvStage: advance_write_floor(clone_ordinal)
    Renderer->>OvStage: write_attribute(env_query, omni:xform)

    loop Each Frame
        Sim->>Renderer: update_transforms()
        Renderer->>Renderer: wp.synchronize_device() GPU to CPU
        Renderer->>OvStage: "write_attribute(object_query, omni:xform, ordinal=N)"
        Renderer->>OvStage: advance_write_floor(N)

        Sim->>Renderer: update_camera()
        Renderer->>Renderer: wp.synchronize_device()
        Renderer->>OvStage: "write_attribute(camera_query, omni:xform, ordinal=N+1)"
        Renderer->>OvStage: advance_write_floor(N+1)

        Sim->>Renderer: render(render_data)
        Renderer->>OVRTX: "step(ordinal=N+1)"
        OVRTX-->>Renderer: frames[]
    end

    Sim->>Renderer: cleanup()
    Renderer->>OvStage: release_query for camera, object, deformable
    Renderer->>OVRTX: detach_ovstage()
    Renderer->>OvStage: ExitStack.close()

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Sim as Simulation
    participant Renderer as OVRTXRenderer
    participant OvStage as ovstage.Stage
    participant OVRTX as ovrtx.Renderer

    Note over Renderer: _use_ovstage = True path

    Sim->>Renderer: prepare_stage(stage, num_envs)
    Renderer->>Renderer: export_stage_to_string_ovstage()

    Sim->>Renderer: create_render_data(spec)
    Renderer->>OvStage: Stage(isaaclab.ovrtx)
    Renderer->>OvStage: "population.open_usd_from_string(ordinal=1)"
    Renderer->>OvStage: "advance_write_floor(ordinal=1)"
    Renderer->>OVRTX: attach_ovstage(stage)
    Renderer->>OvStage: stage.clone(source, targets, ordinal)
    Renderer->>OvStage: advance_write_floor(clone_ordinal)
    Renderer->>OvStage: write_attribute(env_query, omni:xform)

    loop Each Frame
        Sim->>Renderer: update_transforms()
        Renderer->>Renderer: wp.synchronize_device() GPU to CPU
        Renderer->>OvStage: "write_attribute(object_query, omni:xform, ordinal=N)"
        Renderer->>OvStage: advance_write_floor(N)

        Sim->>Renderer: update_camera()
        Renderer->>Renderer: wp.synchronize_device()
        Renderer->>OvStage: "write_attribute(camera_query, omni:xform, ordinal=N+1)"
        Renderer->>OvStage: advance_write_floor(N+1)

        Sim->>Renderer: render(render_data)
        Renderer->>OVRTX: "step(ordinal=N+1)"
        OVRTX-->>Renderer: frames[]
    end

    Sim->>Renderer: cleanup()
    Renderer->>OvStage: release_query for camera, object, deformable
    Renderer->>OVRTX: detach_ovstage()
    Renderer->>OvStage: ExitStack.close()

Loading

Comments Outside Diff (1)

1. source/isaaclab_ov/test/test_ovrtx_clone_plan.py, line 1711-1726 (link)

   "Path-agnostic" tests always use the ovstage renderer variant

   test_prepare_stage_writes_pre_ovrtx_stage_dump and test_prepare_stage_skips_temp_usd_write_when_temp_usd_dir_unset are grouped under "path-agnostic tests" and have no _skip_if_not_* guard. However, both call _make_ovrtx_renderer_without_backend() which hardcodes _use_ovstage = True, so prepare_stage always exercises export_stage_to_string_ovstage regardless of the env var. In legacy mode these tests silently cover the wrong export path. Consider adding a legacy variant or parametrizing them to cover both code paths.

_{Reviews (1): Last reviewed commit: "Added shimmed ovstage code-path to ovrtx..." | Re-trigger Greptile}

[greptile-apps]

ovstage is a hard dependency but is intended to be optional

ovstage is listed as a required dependency here, but the PR description explicitly states it is "disabled by default" and the ovrtx_renderer.py treats it as purely optional (try: import ovstage … except ModuleNotFoundError: pass). The root pyproject.toml correctly places it under an optional extra (ovstage = ["ovstage>=0.1.0,<0.2.0"]). Listing it here as a hard dep forces it onto every user who installs isaaclab_ov, including those who only want the legacy renderer — contradicting the intent and potentially blocking environments where ovstage cannot be installed.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[rilei-nvidia]

Yeah good call. Before we actually merge this PR we may want to back out that this change and also remove it from the "all" selector, so that this can go in and folks can install ovstage separately

[greptile-apps]

Null check after write_attribute use — dead guard

write_attribute is called on line 692 with self._object_xform_query before the if self._object_xform_query is None guard on line 701. If query_from_path_list ever returned None, the write call on line 692 would already raise, making the guard unreachable. The same pattern repeats in _setup_deformable_bindings_ovstage (for _deformable_points_query). Either move the guard before the first use, or remove it if the ovstage API guarantees a non-None return.

[greptile-apps]

Silent no-op when ovstage unavailable but explicitly requested

When a user sets ISAAC_LAB_OVRTX_USE_OVSTAGE=1 but _OVSTAGE_AVAILABLE is False (e.g., installed without the extra), this function silently returns False and falls back to the legacy path with no log message. The user gets no indication that their environment variable was ignored. Since ovstage is hard-required in isaaclab_ov/pyproject.toml, this is currently an edge case, but if that dep is made optional (as the PR description implies it should be), a warning here would prevent silent misconfiguration.