Avoid disk I/O when preparing USD stage for OVRTX renderer

source/isaaclab_ov/changelog.d/huidongc-avoid-ovrtx-disk-rw.rst

@@ -0,0 +1,17 @@
+Changed
+^^^^^^^
+
+* Changed :attr:`~isaaclab_ov.renderers.OVRTXRendererCfg.temp_usd_dir` defaults to ``None``. Set it to a writable
+  directory when you want the combined stage written to disk for debugging.
+
+Removed
+^^^^^^^
+
+* Removed :attr:`~isaaclab_ov.renderers.OVRTXRendererCfg.temp_usd_suffix`. When a temp file is written, the renderer
+  uses ``ovrtx_renderer_stage.usda`` filename under the configured temp directory.
+
+Fixed
+^^^^^
+
+* Avoided OVRTX staging disk I/O by exporting the prepared USD to memory and loading it with ``open_usd_from_string``
+  instead of always writing intermediate scene and combined USD files.

source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py

@@ -22,6 +22,8 @@
 import logging
 import math
 import os
+import tempfile
+from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
 logger = logging.getLogger(__name__)
@@ -58,9 +60,9 @@
     sync_newton_transforms_kernel,
 )
 from .ovrtx_usd import (
+    build_render_product_as_string,
     create_scene_partition_attributes,
-    export_stage_for_ovrtx,
-    inject_cameras_into_usd,
+    export_stage_to_string,
 )
 
 if TYPE_CHECKING:
@@ -164,7 +166,7 @@ def __init__(self, cfg: OVRTXRendererCfg):
         self._object_binding = None
         self._object_newton_indices: wp.array | None = None
         self._initialized_scene = False
-        self._exported_usd_path: str | None = None
+        self._exported_usd_string: str | None = None
         self._camera_rel_path: str | None = None
         self._output_semantic_color_buffer: wp.array | None = None
 
@@ -192,21 +194,18 @@ def __init__(self, cfg: OVRTXRendererCfg):
         logger.info("OVRTX renderer created successfully")
 
     def prepare_stage(self, stage: Any, num_envs: int) -> None:
-        """Export the USD stage for OVRTX before create_render_data.
+        """Prepare the USD stage for OVRTX before :meth:`create_render_data`.
 
-        Adds cloning attributes and exports the stage to a temporary file.
-        The exported path is used by create_render_data when loading into OVRTX.
+        Adds cloning 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, not _IS_OVRTX_0_3_0_OR_NEWER)
 
-        export_path = "/tmp/stage_before_ovrtx.usda"
-        export_stage_for_ovrtx(stage, export_path, num_envs, self._use_ovrtx_cloning)
-        self._exported_usd_path = export_path
-        logger.info("Exported to %s", export_path)
+        self._exported_usd_string = export_stage_to_string(stage, num_envs, self._use_ovrtx_cloning)
 
     def _initialize_from_spec(self, spec: CameraRenderSpec):
         """Initialize the OVRTX renderer with internal environment cloning.
@@ -225,14 +224,10 @@ def _initialize_from_spec(self, spec: CameraRenderSpec):
             raise RuntimeError(f"Expected camera prim under '{env_0_prefix}', got '{first_cam_path}'")
         self._camera_rel_path = spec.camera_path_relative_to_env_0
 
-        usd_scene_path = self._exported_usd_path
-
-        if usd_scene_path is not None:
+        if self._exported_usd_string is not None:
             logger.info("Injecting camera definitions...")
 
-            combined_usd_path, render_product_path = inject_cameras_into_usd(
-                usd_scene_path,
-                self.cfg,
+            render_product_string, render_product_path = build_render_product_as_string(
                 width=width,
                 height=height,
                 num_envs=num_envs,
@@ -242,23 +237,44 @@ def _initialize_from_spec(self, spec: CameraRenderSpec):
             )
             self._render_product_paths.append(render_product_path)
 
+            combined_usd_string = self._exported_usd_string + "\n\n" + render_product_string
+            self._exported_usd_string = None  # Free memory
+
+            if self.cfg.temp_usd_dir is not None:
+                temp_usd_dir = Path(self.cfg.temp_usd_dir)
+            elif not _IS_OVRTX_0_3_0_OR_NEWER:
+                # OVRTX 0.2.0 is not able to load USD from a string, so we need to write to a temporary file.
+                temp_usd_dir = Path(tempfile.gettempdir()) / "ovrtx"
+            else:
+                temp_usd_dir = None
+
+            if temp_usd_dir is not None:
+                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)
+            else:
+                temp_usd_path = None
+
             logger.info("Loading USD into OvRTX...")
             try:
                 if _IS_OVRTX_0_3_0_OR_NEWER:
-                    self._renderer.open_usd(combined_usd_path)
-                    logger.info("USD loaded as root layer (path: %s)", combined_usd_path)
+                    self._renderer.open_usd_from_string(combined_usd_string)
+                    logger.info("OVRTX loaded USD from string successfully")
                 else:
-                    handle = self._renderer.add_usd(combined_usd_path, path_prefix=None)
+                    assert temp_usd_path is not None  # OVRTX < 0.3.0 always materializes combined USD on disk.
+                    handle = self._renderer.add_usd(str(temp_usd_path), path_prefix=None)
                     self._usd_handles.append(handle)
-                    logger.info("USD loaded (path: %s, handle: %s)", combined_usd_path, handle)
+                    logger.info("OVRTX loaded USD from file successfully (path: %s, handle: %s)", temp_usd_path, handle)
             except Exception as e:
                 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)
-                self._update_scene_partitions_after_clone(combined_usd_path, num_envs)
+                self._update_scene_partitions_after_clone(num_envs)
 
             self._initialized_scene = True
 
@@ -299,7 +315,7 @@ def _clone_environments_in_ovrtx(self, num_envs: int):
             logger.error("Failed to clone environments: %s", e)
             raise RuntimeError(f"OvRTX environment cloning failed: {e}")
 
-    def _update_scene_partitions_after_clone(self, usd_file_path: str, num_envs: int):
+    def _update_scene_partitions_after_clone(self, num_envs: int):
         """Update scene partition attributes on cloned environments and cameras in OvRTX."""
         logger.info("Writing scene partitions for %d environments...", num_envs)
         partition_tokens = [f"env_{i}" for i in range(num_envs)]

source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer_cfg.py

@@ -5,9 +5,6 @@
 
 """Configuration for OVRTX Renderer."""
 
-import tempfile
-from pathlib import Path
-
 from isaaclab.renderers.renderer_cfg import RendererCfg
 from isaaclab.utils.configclass import configclass
 
@@ -26,14 +23,11 @@ class OVRTXRendererCfg(RendererCfg):
     renderer_type: str = "ovrtx"
     """Type identifier for OVRTX renderer."""
 
-    temp_usd_dir: str = str(Path(tempfile.gettempdir()) / "ovrtx")
+    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.
     """
 
-    temp_usd_suffix: str = ".usda"
-    """File suffix for temporary combined USD files (e.g. '.usda' or '.usdc')."""
-
     use_ovrtx_cloning: bool = True
     """When True, export only env_0 and use OVRTX ``clone_usd``. When False, export full multi-environment stage.
 

source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_usd.py

@@ -7,15 +7,9 @@
 
 import logging
 import math
-import tempfile
-from pathlib import Path
-from typing import TYPE_CHECKING
 
 from pxr import Sdf, Usd, UsdGeom
 
-if TYPE_CHECKING:
-    from .ovrtx_renderer_cfg import OVRTXRendererCfg
-
 logger = logging.getLogger(__name__)
 
 
@@ -105,34 +99,29 @@ def _tiled_resolution(num_envs: int, width: int, height: int) -> tuple[int, int]
     return num_cols * width, num_rows * height
 
 
-def inject_cameras_into_usd(
-    usd_scene_path: str,
-    cfg: "OVRTXRendererCfg",
+def build_render_product_as_string(
     width: int,
     height: int,
     num_envs: int,
     data_types: list[str],
     minimal_mode: int | None = None,
     camera_rel_path: str = "Camera",
 ) -> tuple[str, str]:
-    """Inject camera and render product definitions into an existing USD file.
+    """Build the render product USD snippet as a string.
 
-    Reads the USD file, appends a Render scope (cameras + RenderProduct + Vars),
-    writes to a temp file in cfg.temp_usd_dir, and returns (path_to_combined_usd, render_product_path).
+    This string is meant to be appended to an exported stage (ASCII) before loading into OVRTX.
 
     Args:
-        usd_scene_path: Path to the base USD scene.
-        cfg: OVRTX renderer config (simple_shading_mode, temp_usd_dir, temp_usd_suffix).
-        width: Tile width from sensor config.
-        height: Tile height from sensor config.
+        width: Tile width from sensor config [px].
+        height: Tile height from sensor config [px].
         num_envs: Number of environments from scene.
         data_types: Data types from sensor config.
         minimal_mode: RTX minimal mode. None if not requested. Valid values are 1, 2, 3.
         camera_rel_path: Camera prim path relative to the env root (e.g. ``"Camera"`` or ``"Robot/head_cam"``).
-    """
-    with open(usd_scene_path) as f:
-        original_usd = f.read()
 
+    Returns:
+        Tuple of (render product USD snippet as a string, absolute render product prim path).
+    """
     data_types = data_types if data_types else ["rgb"]
     tiled_width, tiled_height = _tiled_resolution(num_envs, width, height)
 
@@ -152,14 +141,7 @@ def inject_cameras_into_usd(
         tiled_height,
         minimal_mode,
     )
-    combined_usd = original_usd.rstrip() + "\n\n" + camera_content
-
-    Path(cfg.temp_usd_dir).mkdir(parents=True, exist_ok=True)
-    with tempfile.NamedTemporaryFile(mode="w", suffix=cfg.temp_usd_suffix, delete=False, dir=cfg.temp_usd_dir) as f:
-        f.write(combined_usd)
-        temp_path = f.name
-    logger.info("Created combined USD: %s", temp_path)
-    return temp_path, render_product_path
+    return camera_content, render_product_path
 
 
 def create_scene_partition_attributes(
@@ -207,40 +189,38 @@ def create_scene_partition_attributes(
                 logger.debug("Set scene partition '%s' on prim '%s'", scene_partition, prim.GetPath())
 
 
-def export_stage_for_ovrtx(stage, export_path: str, num_envs: int, use_ovrtx_cloning: bool = True) -> str:
-    """Export the stage to a USD file; when num_envs > 1, only env_0 is exported for OVRTX cloning.
+def export_stage_to_string(stage, num_envs: int, use_ovrtx_cloning: bool = True) -> str:
+    """Export the stage to a string; when num_envs > 1, only env_0 is exported for OVRTX cloning.
 
     When num_envs > 1, deactivates env_1..env_{num_envs-1} before export and reactivates
-    them after, so the file contains only env_0. The stage is modified in place.
+    them after, so the exported content contains only env_0. The stage is modified in place.
 
     Args:
         stage: USD stage to export.
-        export_path: Path for the exported file.
         num_envs: Number of environments.
+        use_ovrtx_cloning: Whether OVRTX cloning is enabled.
 
     Returns:
-        export_path (same as input).
+        The exported stage as a string.
     """
-    deactivated = []
+    deactivated_prims = []
     if use_ovrtx_cloning and num_envs > 1:
-        logger.info("Deactivating %d cloned environments...", num_envs - 1)
+        logger.info("Deactivating %d environment roots...", num_envs - 1)
         for env_idx in range(1, num_envs):
             env_path = f"/World/envs/env_{env_idx}"
             prim = stage.GetPrimAtPath(env_path)
             if prim.IsValid() and prim.IsActive():
                 prim.SetActive(False)
-                deactivated.append(prim)
-                if env_idx <= 3 or env_idx == num_envs - 1:
-                    logger.info("Deactivated: %s", env_path)
-        if num_envs > 5:
-            logger.info("... (deactivated %d environments total)", len(deactivated))
+                deactivated_prims.append(prim)
+                logger.debug("Deactivated environment root: %s", env_path)
+
+        logger.info("Deactivated %d environment roots in total", len(deactivated_prims))
 
     try:
-        stage.Export(export_path)
-        return export_path
+        return stage.ExportToString()
     finally:
-        if deactivated:
-            logger.info("Reactivating %d environments...", len(deactivated))
-            for prim in deactivated:
+        if deactivated_prims:
+            logger.info("Reactivating %d environment roots...", len(deactivated_prims))
+            for prim in deactivated_prims:
                 if prim.IsValid():
                     prim.SetActive(True)