[Bug Report] OVRTX cloning and OVSTAGE cloning lose Franka visuals with link-level instance groups

OLD
panda_link0
├── link0_0 [instance]
│   └── Mesh
├── link0_1 [instance]
│   └── Mesh
└── ...

NEW
panda_link0
└── physx_visuals [instance]
    ├── link0_0 [ordinary Xform]
    │   └── Mesh
    ├── link0_1 [ordinary Xform]
    │   └── Mesh
    └── ...

tar -xzf franka-old.tar.gz
tar -xzf franka-new.tar.gz

CUDA_VISIBLE_DEVICES=0 PXR_WORK_THREAD_LIMIT=1 uv run --extra ov python franka_instancing_repro.py \
    --path legacy --asset franka-new/franka_panda.usda \
    --old-asset franka-old/franka_old.usda --output /tmp/franka-instancing-legacy

CUDA_VISIBLE_DEVICES=0 PXR_WORK_THREAD_LIMIT=1 uv run --extra ov python franka_instancing_repro.py \
    --path ovstage --asset franka-new/franka_panda.usda \
    --old-asset franka-old/franka_old.usda --output /tmp/franka-instancing-ovstage

# 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

"""Render a Franka source and native clone, with controlled instance/collider toggles.

Run from an Isaac Lab environment with the ov extra installed. No Kit or physics simulation.
The default downloads the public Franka USD dependencies; --asset accepts a local copy instead.
"""

import argparse
import faulthandler
import hashlib
import importlib.metadata
import json
import urllib.parse
import urllib.request
from pathlib import Path

import numpy as np
import ovrtx
import ovstage
import torch
from PIL import Image
from pxr import Gf, Sdf, Usd, UsdGeom, UsdPhysics

from isaaclab_ov.renderers.ovrtx_usd import build_render_scope_usd
from isaaclab_ov.stage import create_ovstage, xform_tensor_from_numpy

ASSET_URL = (
    "https://omniverse-content-production.s3-us-west-2.amazonaws.com/Assets/Isaac/6.1/Isaac/IsaacLab/Robots/"
    "FrankaEmika/franka_panda.usda"
)


def download_asset(url: str, directory: Path, downloaded: set[str]) -> Path:
    """Download USD composition dependencies, preserving their relative paths."""
    relative = url.removeprefix(ASSET_URL.rsplit("/", 1)[0] + "/")
    if relative == url or ".." in Path(relative).parts:
        raise ValueError(f"Unexpected dependency outside the Franka asset directory: {url}")
    target = directory / relative
    if url in downloaded:
        return target
    downloaded.add(url)
    target.parent.mkdir(parents=True, exist_ok=True)
    with urllib.request.urlopen(url, timeout=60) as response:
        target.write_bytes(response.read())
    layer = Sdf.Layer.FindOrOpen(str(target))
    for reference in layer.GetExternalReferences():
        if Path(reference).suffix in {".usd", ".usda", ".usdc"}:
            download_asset(urllib.parse.urljoin(url, reference), directory, downloaded)
    return target


parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--path", choices=("legacy", "ovstage"), required=True)
parser.add_argument("--asset", type=Path, help="Local franka_panda.usda with its composition dependencies.")
parser.add_argument("--old-asset", type=Path, help="Optional retained old-layout asset, not needed to reproduce.")
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
args.output.mkdir(parents=True, exist_ok=True)
asset = args.asset.resolve() if args.asset else download_asset(ASSET_URL, args.output.resolve() / "asset", set())
versions = {name: importlib.metadata.version(name) for name in ("ovrtx", "ovstage", "torch")}
versions.update(usd=Usd.GetVersion(), asset_sha256=hashlib.sha256(asset.read_bytes()).hexdigest(), path=args.path)
(args.output / "versions.json").write_text(json.dumps(versions, indent=2))
print("VERSIONS", json.dumps(versions), flush=True)
faulthandler.enable()
faulthandler.dump_traceback_later(240, exit=True)
torch.cuda.init()
renderer = ovrtx.Renderer(ovrtx.RendererConfig(
    log_file_path=str(args.output.resolve() / "renderer.log"), log_level="warn", active_cuda_gpus="0",
    keep_system_alive=True, suppress_deprecation_warnings=True, read_gpu_transforms=True,
))
results = []
images = {}
cases = [
    ("new", asset, False, False),
    ("new_arm_colliders_off", asset, True, False),
    ("new_groups_only", asset, False, True),
    ("new_arm_off_groups", asset, True, True),
]
if args.old_asset:
    cases.insert(0, ("old", args.old_asset.resolve(), False, False))

for name, filename, disable_arm_colliders, deinstance_groups in cases:
    stage = Usd.Stage.CreateInMemory()
    UsdGeom.SetStageUpAxis(stage, "Z")
    UsdGeom.SetStageMetersPerUnit(stage, 1.0)
    for env in range(2):
        UsdGeom.Xform.Define(stage, f"/World/envs/env_{env}")
    camera = UsdGeom.Camera.Define(stage, "/World/Camera")
    camera.CreateProjectionAttr("orthographic")
    camera.CreateHorizontalApertureAttr(60.0)
    camera.CreateVerticalApertureAttr(30.0)
    camera.CreateClippingRangeAttr((0.01, 100.0))
    camera.AddTransformOp().Set(Gf.Matrix4d().SetLookAt(
        Gf.Vec3d(1.5, -5, 0.6), Gf.Vec3d(1.5, 0, 0.6), Gf.Vec3d(0, 0, 1)
    ).GetInverse())
    robot = UsdGeom.Xform.Define(stage, "/World/envs/env_0/Robot").GetPrim()
    robot.GetReferences().AddReference(str(filename))
    if disable_arm_colliders:
        for prim in stage.Traverse():
            if prim.HasAPI(UsdPhysics.CollisionAPI) and "/panda_hand/" not in str(prim.GetPath()):
                UsdPhysics.CollisionAPI(prim).CreateCollisionEnabledAttr(False)
    prims = list(Usd.PrimRange(robot, Usd.TraverseInstanceProxies()))
    enabled = sum(
        bool(UsdPhysics.CollisionAPI(prim).GetCollisionEnabledAttr().Get())
        for prim in prims if prim.HasAPI(UsdPhysics.CollisionAPI)
    )
    visual_meshes = sum(
        prim.IsA(UsdGeom.Mesh) and UsdGeom.Imageable(prim).ComputePurpose() != "guide" for prim in prims
    )
    original_layers = [layer.ExportToString() for layer in stage.GetLayerStack()]
    session = Sdf.Layer.CreateAnonymous()
    session.subLayerPaths = [stage.GetSessionLayer().identifier]
    export_stage = Usd.Stage.Open(stage.GetRootLayer(), session)
    groups = [prim.GetPath() for prim in export_stage.Traverse() if prim.GetName() == "physx_visuals" and prim.IsInstance()]
    if deinstance_groups:
        assert groups, "Asset no longer has the affected instanced physx_visuals grouping."
        with Usd.EditContext(export_stage, session):
            for path in groups:
                export_stage.GetPrimAtPath(path).SetInstanceable(False)
    exported = export_stage.ExportToString()
    exported += build_render_scope_usd(
        ["/World/Camera"], "Product", "/Render/Vars/LdrColor", "LdrColor", "LdrColor", 512, 256,
        minimal_mode=1, background_color=(0.0, 0.0, 0.0), device_id=0,
    )
    assert original_layers == [layer.ExportToString() for layer in stage.GetLayerStack()]
    transforms = np.tile(np.eye(4, dtype=np.float64), (2, 1, 1))
    transforms[1, 3, 0] = 3.0
    if args.path == "ovstage":
        native_stage = create_ovstage("franka-instancing-" + name)
        ovstage.population.open_usd_from_string(native_stage, exported, ordinal=1, domains=ovstage.PopulationDomain.RENDERING)
        native_stage.clone(str(robot.GetPath()), ["/World/envs/env_1/Robot"], ordinal=1)
        with ovstage.PathDictionary(native_stage) as path_dict:
            path_list = path_dict.create_path_list_from_strings(["/World/envs/env_0", "/World/envs/env_1"])
            with native_stage.query_from_path_list(path_list) as query:
                native_stage.write_attribute(
                    query, "omni:xform", ordinal=1, tensors=xform_tensor_from_numpy(transforms),
                    is_array=False, semantic=ovstage.AttributeSemantic.MATRIX,
                ).wait()
            path_dict.destroy_path_list(path_list)
        native_stage.advance_write_floor(ordinal=1).wait()
        renderer.attach_ovstage(native_stage)
    else:
        renderer.reset_stage()
        renderer.open_usd_from_string(exported)
        renderer.clone_usd(str(robot.GetPath()), ["/World/envs/env_1/Robot"])
        renderer.write_attribute(
            ["/World/envs/env_0", "/World/envs/env_1"], "omni:xform", transforms, semantic=ovrtx.Semantic.XFORM_MAT4x4
        )
    for ordinal in range(2, 6):
        if args.path == "ovstage":
            native_stage.advance_write_floor(ordinal=ordinal).wait()
        products = renderer.step({"/Render/Product"}, 1.0 / 60.0, **({"ordinal": ordinal} if args.path == "ovstage" else {}))
        frame = products["/Render/Product"].frames[0]
        with frame.render_vars["LdrColor"].map(device=ovrtx.Device.CPU) as mapped:
            view = np.from_dlpack(mapped)
            pixels = view.copy()
            del view
        del mapped, frame, products
    Image.fromarray(pixels).save(args.output / f"{name}.png")
    images[name] = pixels
    record = dict(
        path=args.path, case=name, usd_enabled_colliders=enabled, usd_visual_meshes=visual_meshes,
        visual_instance_groups=len(groups),
        source_visible_pixels=int(np.any(pixels[:, :256, :3] > 5, axis=-1).sum()),
        clone_visible_pixels=int(np.any(pixels[:, 256:, :3] > 5, axis=-1).sum()),
    )
    results.append(record)
    print("RESULT", json.dumps(record), flush=True)
    if args.path == "ovstage":
        renderer.detach_ovstage()
        native_stage.destroy()

comparisons = {
    "collider_toggle_identical": bool(np.array_equal(images["new"], images["new_arm_colliders_off"])),
    "fixed_collider_toggle_identical": bool(np.array_equal(images["new_groups_only"], images["new_arm_off_groups"])),
}
if "old" in images:
    comparisons["old_and_fixed_identical"] = bool(np.array_equal(images["old"], images["new_groups_only"]))
(args.output / "results.json").write_text(json.dumps(results, indent=2))
(args.output / "comparisons.json").write_text(json.dumps(comparisons, indent=2))
print("COMPARISONS", json.dumps(comparisons), flush=True)
renderer.destroy()
faulthandler.cancel_dump_traceback_later()