Complete OVRTX cloning (#5781)

# Description

Complete OVRTX cloning

## Type of change

- Bug fix (non-breaking change which fixes an issue)

## Checklist

- [x] I have read and understood the [contribution
guidelines](https://isaac-sim.github.io/IsaacLab/main/source/refs/contributing.html)
- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

Author: huidongc

.github/workflows/build.yaml

@@ -681,7 +681,12 @@ jobs:
         isaacsim-base-image: ${{ needs.config.outputs.isaacsim_image_name }}
         isaacsim-version: ${{ needs.config.outputs.isaacsim_image_tag }}
         filter-pattern: "isaaclab_tasks"
-        include-files: "test_rendering_cartpole.py,test_rendering_dexsuite_kuka.py,test_rendering_registered_tasks.py,test_rendering_shadow_hand.py"
+        include-files: >-
+          test_rendering_cartpole.py,
+          test_rendering_dexsuite_kuka_homo.py,
+          test_rendering_dexsuite_kuka_hetero.py,
+          test_rendering_registered_tasks.py,
+          test_rendering_shadow_hand.py
         container-name: isaac-lab-rendering-correctness-test
 
   test-rendering-correctness-kitless:
@@ -703,7 +708,11 @@ jobs:
         isaacsim-version: ${{ needs.config.outputs.isaacsim_image_tag }}
         filter-pattern: "isaaclab_tasks"
         extra-pip-packages: "ovrtx ovphysx==0.4.13"
-        include-files: "test_rendering_cartpole_kitless.py,test_rendering_dexsuite_kuka_kitless.py,test_rendering_shadow_hand_kitless.py"
+        include-files: >-
+          test_rendering_cartpole_kitless.py,
+          test_rendering_dexsuite_kuka_homo_kitless.py,
+          test_rendering_dexsuite_kuka_hetero_kitless.py,
+          test_rendering_shadow_hand_kitless.py
         container-name: isaac-lab-rendering-correctness-kitless-test
   #endregion
 

scripts/benchmarks/nsys_trace.json

@@ -162,5 +162,25 @@
             {"function": "NewtonWarpRenderer.render", "color": "0x4CAF50"},
             {"function": "NewtonWarpRenderer.read_output", "color": "0xFF9800"}
         ]
+    },
+    {
+        "domain": "OVRTXRenderer",
+        "color": "0x76B800",
+        "module": "isaaclab_ov.renderers.ovrtx_renderer",
+        "functions": [
+            "OVRTXRenderer.prepare_stage",
+            "OVRTXRenderer.create_render_data",
+            "OVRTXRenderer._clone_sources_in_ovrtx",
+            "OVRTXRenderer._update_scene_partitions_after_clone",
+            "OVRTXRenderer._setup_object_bindings",
+            "OVRTXRenderer.update_transforms",
+            "OVRTXRenderer.render",
+            "OVRTXRenderer.read_output",
+            {"module": "isaaclab_ov.renderers.ovrtx_usd", "function": "create_scene_partition_attributes"},
+            {"module": "isaaclab_ov.renderers.ovrtx_usd", "function": "export_stage_to_string"},
+            {"module": "ovrtx", "function": "Renderer.open_usd_from_string", "color": "0x87CEEB"},
+            {"module": "ovrtx", "function": "Renderer.bind_attribute", "color": "0x87CEEB"},
+            {"module": "ovrtx", "function": "Renderer.write_attribute", "color": "0x87CEEB"}
+        ]
     }
 ]

source/isaaclab_ov/changelog.d/ovrtx-heterogeneous-cloning.rst

@@ -0,0 +1,8 @@
+Changed
+^^^^^^^
+
+* Extended the :attr:`~isaaclab_ov.renderers.OVRTXRendererCfg.use_ovrtx_cloning` path to support
+  heterogeneous scenes as well as homogeneous ones. :meth:`~isaaclab_ov.renderers.OVRTXRenderer.prepare_stage`
+  now exports only :class:`~isaaclab.cloner.ClonePlan` source prototypes plus global stage metadata, and
+  replication uses OVRTX cloning API (``clone_usd``) for all rows in the published clone plan instead of
+  cloning only ``/World/envs/env_0``.

source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py

@@ -22,6 +22,7 @@
 import logging
 import math
 import os
+from itertools import compress
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, NoReturn
 
@@ -51,6 +52,7 @@
         "(or, manually: pip install --extra-index-url https://pypi.nvidia.com -e 'source/isaaclab_ov[ovrtx]')."
     ) from exc
 
+from isaaclab.cloner.clone_plan import ClonePlan
 from isaaclab.renderers import BaseRenderer, RenderBufferKind, RenderBufferSpec
 from isaaclab.sim import SimulationContext
 from isaaclab.utils.warp.warp_math import convert_camera_frame_orientation_convention_wp
@@ -130,6 +132,67 @@ def _resolve_rtx_minimal_mode(data_types: list[str]) -> int | None:
     return _RTX_MINIMAL_MODES[filtered_data_types[0]]
 
 
+def _write_file(output_dir: Path, file_name: str, content: str) -> None:
+    """Write ``content`` to ``output_dir / file_name``.
+
+    Creates ``output_dir`` and any missing parents when needed.
+
+    Args:
+        output_dir: Directory that receives the file.
+        file_name: Base name of the file to write.
+        content: Text content written with UTF-8 encoding.
+    """
+    output_dir.mkdir(parents=True, exist_ok=True)
+    output_path = output_dir / file_name
+
+    with open(output_path, "w", encoding="utf-8") as file:
+        file.write(content)
+        logger.info("Wrote USD file: %s", output_path)
+
+
+def _create_homogeneous_clone_plan(num_envs: int) -> ClonePlan:
+    """Create a homogeneous fallback plan that replicates ``env_0`` to every environment."""
+    return ClonePlan(
+        sources=("/World/envs/env_0",),
+        destinations=("/World/envs/env_{}",),
+        clone_mask=torch.ones((1, num_envs), dtype=torch.bool, device="cpu"),
+    )
+
+
+def _resolve_clone_plan(num_envs: int) -> ClonePlan:
+    """Resolve clone plan for local use.
+
+    If no clone plan is published by the scene or it has no active rows, returns a homogeneous clone plan.
+    If the clone plan has some inactive rows, returns a copy of the published clone plan with only the active rows.
+    If the clone plan has all active rows, returns the published clone plan (shallow copy).
+    """
+    published_clone_plan = SimulationContext.instance().get_clone_plan()
+
+    if published_clone_plan is None:
+        logger.warning("No clone plan is published by the scene; returning homogeneous clone plan")
+        return _create_homogeneous_clone_plan(num_envs)
+
+    active_rows = published_clone_plan.clone_mask.any(dim=1)
+
+    # If no rows are active, return a homogeneous clone plan.
+    if not active_rows.any():
+        logger.warning("Clone plan has no active rows, returning homogeneous clone plan")
+        return _create_homogeneous_clone_plan(num_envs)
+
+    # If some rows are inactive, return a copy of the published clone plan with only the active rows.
+    if not active_rows.all():
+        logger.warning("Clone plan has some inactive rows; returning a copy with only active rows")
+        active = active_rows.tolist()
+        return ClonePlan(
+            sources=tuple(compress(published_clone_plan.sources, active)),
+            destinations=tuple(compress(published_clone_plan.destinations, active)),
+            clone_mask=published_clone_plan.clone_mask[active_rows],
+        )
+
+    # If all rows are active, return the published clone plan (shallow copy).
+    return published_clone_plan
+
+
 class OVRTXRenderData:
     """OVRTX-specific RenderData. Holds warp output buffers sized from :class:`CameraRenderSpec`."""
 
@@ -197,14 +260,7 @@ def __init__(self, cfg: OVRTXRendererCfg):
         self._exported_usd_string: str | None = None
         self._camera_rel_path: str | None = None
         self._output_semantic_color_buffer: wp.array | None = None
-
-        self._use_ovrtx_cloning = self.cfg.use_ovrtx_cloning
-
-        if self._use_ovrtx_cloning:
-            clone_plan = SimulationContext.instance().get_clone_plan()
-            if clone_plan and not clone_plan.clone_mask.all().item():
-                logger.warning("OVRTX cloning disabled because the simulation uses a heterogeneous env setup")
-                self._use_ovrtx_cloning = False
+        self._clone_plan: ClonePlan | None = None
 
         logger.info("Creating OVRTX renderer...")
         OVRTX_CONFIG = RendererConfig(
@@ -246,16 +302,27 @@ def prepare_cameras(self, stage: Any, spec: CameraRenderSpec) -> None:
     def prepare_stage(self, stage: Any, num_envs: int) -> None:
         """Prepare the USD stage for OVRTX before :meth:`create_render_data`.
 
-        Adds cloning attributes and exports the stage to a string held on the renderer until
+        Adds scene partition attributes and exports the stage to a string held on the renderer until
         :meth:`create_render_data` is called.
         """
         if stage is None:
             return
 
-        logger.info("Preparing stage for export (%d envs, cloning=%s)...", num_envs, self._use_ovrtx_cloning)
-        create_scene_partition_attributes(stage, num_envs, self._use_ovrtx_cloning)
+        # If temp_usd_dir is set, write the pre-ovrtx stage to a temporary file.
+        if self.cfg.temp_usd_dir is not None:
+            _write_file(Path(self.cfg.temp_usd_dir), "pre_ovrtx_renderer_stage.usda", stage.ExportToString())
 
-        self._exported_usd_string = export_stage_to_string(stage, num_envs, self._use_ovrtx_cloning)
+        logger.info("Preparing stage (%d envs)...", num_envs)
+        create_scene_partition_attributes(stage, num_envs)
+
+        # Resolve the clone plan for local use.
+        self._clone_plan = _resolve_clone_plan(num_envs)
+
+        self._exported_usd_string = export_stage_to_string(
+            stage,
+            num_envs,
+            source_paths=self._clone_plan.sources,
+        )
 
     def _initialize_from_spec(self, spec: CameraRenderSpec):
         """Initialize the OVRTX renderer with internal environment cloning.
@@ -294,12 +361,7 @@ def _initialize_from_spec(self, spec: CameraRenderSpec):
 
             # If temp_usd_dir is set, write the combined USD stage to a temporary file.
             if self.cfg.temp_usd_dir is not None:
-                temp_usd_dir = Path(self.cfg.temp_usd_dir)
-                temp_usd_dir.mkdir(parents=True, exist_ok=True)
-                temp_usd_path = temp_usd_dir / "ovrtx_renderer_stage.usda"
-                with open(temp_usd_path, "w", encoding="utf-8") as f:
-                    f.write(combined_usd_string)
-                    logger.info("Wrote combined USD stage to %s", temp_usd_path)
+                _write_file(Path(self.cfg.temp_usd_dir), "ovrtx_renderer_stage.usda", combined_usd_string)
 
             logger.info("Loading USD into OvRTX...")
             try:
@@ -309,9 +371,8 @@ def _initialize_from_spec(self, spec: CameraRenderSpec):
                 logger.exception("Error loading USD: %s", e)
                 raise
 
-            if self._use_ovrtx_cloning and num_envs > 1:
-                logger.info("Using OVRTX internal cloning")
-                self._clone_environments_in_ovrtx(num_envs)
+            if num_envs > 1:
+                self._clone_sources_in_ovrtx()
                 self._update_scene_partitions_after_clone(num_envs)
 
             self._initialized_scene = True
@@ -341,17 +402,39 @@ def _initialize_from_spec(self, spec: CameraRenderSpec):
 
             self._setup_object_bindings()
 
-    def _clone_environments_in_ovrtx(self, num_envs: int):
-        """Clone base environment (env_0) to all other environments using OvRTX."""
-        logger.info("Cloning base environment to %d targets...", num_envs - 1)
-        source_path = "/World/envs/env_0"
-        target_paths = [f"/World/envs/env_{i}" for i in range(1, num_envs)]
-        try:
-            self._renderer.clone_usd(source_path, target_paths)
-            logger.info("Cloned %d environments successfully", len(target_paths))
-        except Exception as e:
-            logger.error("Failed to clone environments: %s", e)
-            raise RuntimeError(f"OvRTX environment cloning failed: {e}")
+    def _clone_sources_in_ovrtx(self):
+        """Clone sources in OVRTX using the scene :class:`~isaaclab.cloner.ClonePlan`."""
+        clone_plan = self._clone_plan
+        if clone_plan is None:
+            raise RuntimeError("Clone plan is required when using OVRTX cloning")
+
+        logger.info("Cloning sources in OVRTX...")
+
+        num_envs = clone_plan.clone_mask.shape[1]
+        env_ids = torch.arange(num_envs, dtype=torch.int32, device=clone_plan.clone_mask.device)
+
+        num_cloned_sources = 0
+
+        for row_idx, (source, destination) in enumerate(zip(clone_plan.sources, clone_plan.destinations, strict=True)):
+            target_env_ids = env_ids[clone_plan.clone_mask[row_idx]].tolist()
+
+            target_paths = []
+            for env_id in target_env_ids:
+                resolved_destination = destination.format(env_id)
+                if resolved_destination != source:
+                    target_paths.append(resolved_destination)
+
+            if target_paths:
+                logger.debug("Cloning row %d: %s -> %d target(s)", row_idx, source, len(target_paths))
+                try:
+                    self._renderer.clone_usd(source, target_paths)
+                    num_cloned_sources += 1
+                except Exception as e:
+                    error_msg = f"Failed to clone row {row_idx} from {source}: {e}"
+                    logger.error(error_msg)
+                    raise RuntimeError(error_msg)
+
+        logger.info("Cloned %d sources successfully in OVRTX", num_cloned_sources)
 
     def _update_scene_partitions_after_clone(self, num_envs: int):
         """Update scene partition attributes on cloned environments and cameras in OvRTX."""

source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer_cfg.py

@@ -27,17 +27,11 @@ class OVRTXRendererCfg(RendererCfg):
     """Type identifier for OVRTX renderer."""
 
     temp_usd_dir: str | None = None
-    """Directory for temporary combined USD files (scene + injected cameras).
-    Used by the OVRTX renderer when building the render scope; must be writable.
-    """
-
-    use_ovrtx_cloning: bool = True
-    """When True, export only env_0 and use OVRTX ``clone_usd``. When False, export full multi-environment stage.
-
-    OVRTX cloning is only supported in OVRTX 0.3.0 or newer.
+    """Directory for temporary USD debug dumps written during OVRTX stage preparation.
 
-    If the simulation uses a heterogeneous env setup, the renderer disables this path and exports the full
-    multi-environment stage instead (same effect as setting this to ``False`` for that run).
+    When set, the renderer writes ``pre_ovrtx_renderer_stage.usda`` (raw stage before
+    partition attributes and export trimming) and ``ovrtx_renderer_stage.usda`` (exported
+    stage plus injected render products) under this directory. Must be writable.
     """
 
     log_level: str = "verbose"