Use clone-plan positions for OVRTX roots

Make the published clone plan authoritative for environment root translation in both OVRTX scene paths. Remove fallback plan resolution and stage transform capture so clone ownership stays explicit.

Author: ooctipus

source/isaaclab_ov/changelog.d/ovrtx-clone-plan-position.rst

@@ -0,0 +1,5 @@
+Fixed
+^^^^^
+
+* Fixed OVRTX environment placement by authoring root translations from the clone plan after
+  cloning instead of capturing transforms from the USD stage.

source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py

@@ -24,7 +24,6 @@
 import math
 import os
 import re
-from itertools import compress
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, NoReturn, cast
 
@@ -251,49 +250,6 @@ def _write_file(output_dir: Path, file_name: str, content: str) -> None:
         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`."""
 
@@ -435,27 +391,20 @@ def prepare_stage(self, stage: Any, num_envs: int) -> None:
         if stage is None:
             return
 
+        self._clone_plan = SimulationContext.instance().get_clone_plan()
+        if self._clone_plan is None or self._clone_plan.positions is None:
+            raise RuntimeError("Clone plan with environment positions is required when preparing OVRTX stage")
+
         # 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())
 
         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)
-        if self._clone_plan is None:
-            raise RuntimeError("Clone plan is required when preparing OVRTX stage")
-
-        # The ovstage path cannot read env-root transforms back after load (see
-        # _clone_sources_ovstage), so snapshot them here while the full USD stage is still live.
-        if self._use_ovstage:
-            self._capture_env_root_xforms_ovstage(stage, num_envs)
-
         # keep_env_roots is False on the ovstage path: ovstage's ``Stage.clone`` requires each target
         # path to not already exist, so the non-source env roots must be trimmed from the exported
-        # stage for it to recreate them. Their xforms were captured just above, since trimming them
-        # is what makes them unreadable afterwards.
+        # stage for it to recreate them.
         self._exported_usd_string = export_stage_to_string(
             stage,
             num_envs,
@@ -563,35 +512,11 @@ def _initialize_from_spec_legacy(self, spec: CameraRenderSpec):
     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")
+        if clone_plan is None or clone_plan.positions is None:
+            raise RuntimeError("Clone plan with environment positions is required when using OVRTX cloning")
 
         num_envs = clone_plan.clone_mask.shape[1]
         env_prim_paths = [f"/World/envs/env_{i}" for i in range(num_envs)]
-        xform_attr_name = "omni:xform"
-
-        # Snapshot per-env root transforms before clone_usd overwrites them with the source env root.
-        #
-        # We only create xform bindings for prims in the body_label list, so their transforms are
-        # driven each frame from simulation data. Prims not in that list (e.g. static rigid objects
-        # such as tables) get no binding, and we never reset their xform stack. Their world placement
-        # therefore relies on every ancestor holding a correct xform in the ovrtx data. In
-        # particular, if an env root has no valid xform, these prims collapse toward env_0 frame
-        # (e.g. tables ending up at env_0 and missing from the other tiles).
-        #
-        # Snapshotting the env-root xforms here and restoring them after clone (see below) keeps those
-        # hierarchy-positioned prims correctly placed per env. This is a cheap one-off operation,
-        # preferable to forcing every static object into the body_label list just to bind its xform.
-        #
-        env_root_xforms = np.empty((num_envs, 4, 4), dtype=np.float64)
-        self._renderer.read_attribute(
-            xform_attr_name,
-            env_prim_paths,
-            prim_mode=PrimMode.MUST_EXIST,
-            dest=env_root_xforms,
-        )
-        logger.info("Captured per-env root transforms before cloning")
-
         logger.info("Cloning sources in OVRTX...")
         env_ids = torch.arange(num_envs, dtype=torch.int32, device=clone_plan.clone_mask.device)
 
@@ -618,15 +543,15 @@ def _clone_sources_in_ovrtx(self):
 
         logger.info("Cloned %d sources successfully in OVRTX", num_cloned_sources)
 
-        # Restore the pre-clone xforms.
+        env_root_xforms = np.tile(np.eye(4, dtype=np.float64), (num_envs, 1, 1))
+        env_root_xforms[:, 3, :3] = clone_plan.positions.cpu().numpy()
         self._renderer.write_attribute(
             prim_paths=env_prim_paths,
-            attribute_name=xform_attr_name,
+            attribute_name="omni:xform",
             tensor=env_root_xforms,
             semantic=Semantic.XFORM_MAT4x4,
             prim_mode=PrimMode.MUST_EXIST,
         )
-        logger.info("Restored per-env root transforms after cloning")
 
     def _update_scene_partitions_after_clone(self, num_envs: int):
         """Update scene partition attributes on cloned environments and cameras in OvRTX."""
@@ -1567,7 +1492,6 @@ def _init_fields_ovstage(self) -> None:
         self._deformable_paths_list = None
         self._particle_points_query = None
         self._particle_paths_list = None
-        self._env_root_xforms: np.ndarray | None = None
 
     def _initialize_from_spec_ovstage(self, spec: CameraRenderSpec) -> None:
         """Initialize the OVRTX renderer with internal environment cloning (ovstage path).
@@ -1678,55 +1602,15 @@ def _initialize_from_spec_ovstage(self, spec: CameraRenderSpec) -> None:
         logger.info("OVRTX loaded USD from string successfully via ovstage")
         self._current_ordinal += 1
 
-    def _capture_env_root_xforms_ovstage(self, stage: Any, num_envs: int) -> None:
-        """Capture per-env root transforms from the live USD stage before export.
-
-        Must be called before :func:`export_stage_to_string`, which trims non-source
-        env geometry so those transforms cannot be read back reliably from ovstage after load.
-        The captured array is consumed by :meth:`_clone_sources_ovstage` and cleared after use.
-        """
-        from pxr import UsdGeom
-
-        xform_cache = UsdGeom.XformCache()
-        self._env_root_xforms = np.empty((num_envs, 4, 4), dtype=np.float64)
-        for i in range(num_envs):
-            prim = stage.GetPrimAtPath(f"/World/envs/env_{i}")
-            self._env_root_xforms[i] = np.array(xform_cache.GetLocalToWorldTransform(prim), dtype=np.float64).reshape(
-                4, 4
-            )
-
     def _clone_sources_ovstage(self):
         """Clone sources in OVRTX using the scene :class:`~isaaclab.cloner.ClonePlan` (ovstage path)."""
         clone_plan = self._clone_plan
-        if clone_plan is None:
-            raise RuntimeError("Clone plan is required when using OVRTX cloning")
+        if clone_plan is None or clone_plan.positions is None:
+            raise RuntimeError("Clone plan with environment positions is required when using OVRTX cloning")
 
         num_envs = clone_plan.clone_mask.shape[1]
         env_prim_paths = [f"/World/envs/env_{i}" for i in range(num_envs)]
 
-        env_paths_list = self._stage_paths.create_path_list_from_strings(env_prim_paths)
-        env_query = self._stage.query_from_path_list(env_paths_list)
-
-        # Snapshot per-env root transforms before clone overwrites them with the source env root.
-        #
-        # We only create xform queries for prims in the body_label list, so their transforms are
-        # driven each frame from simulation data. Prims not in that list (e.g. static rigid objects
-        # such as tables) get no query, and we never reset their xform stack. Their world placement
-        # therefore relies on every ancestor holding a correct xform in the ovstage data. In
-        # particular, if an env root has no valid xform, these prims collapse toward env_0 frame
-        # (e.g. tables ending up at env_0 and missing from the other tiles).
-        #
-        # Snapshotting the env-root xforms here and restoring them after clone (see below) keeps those
-        # hierarchy-positioned prims correctly placed per env. This is a cheap one-off operation,
-        # preferable to forcing every static object into the body_label list just to bind its xform.
-        #
-        # Transforms were captured from the live USD stage in prepare_stage before export
-        # stripped non-source envs — they are not readable from ovstage at this point.
-        env_root_xforms = self._env_root_xforms
-        if env_root_xforms is None:
-            raise RuntimeError("env_root_xforms not captured; ensure prepare_stage was called first")
-        logger.info("Using pre-captured per-env root transforms for post-clone restore")
-
         logger.info("Cloning sources in OVRTX...")
         env_ids = torch.arange(num_envs, dtype=torch.int32, device=clone_plan.clone_mask.device)
 
@@ -1752,7 +1636,10 @@ def _clone_sources_ovstage(self):
 
         logger.info("Cloned %d sources successfully in OVRTX", num_cloned_sources)
 
-        # Restore the pre-clone xforms.
+        env_root_xforms = np.tile(np.eye(4, dtype=np.float64), (num_envs, 1, 1))
+        env_root_xforms[:, 3, :3] = clone_plan.positions.cpu().numpy()
+        env_paths_list = self._stage_paths.create_path_list_from_strings(env_prim_paths)
+        env_query = self._stage.query_from_path_list(env_paths_list)
         self._stage.write_attribute(
             env_query,
             "omni:xform",
@@ -1761,8 +1648,6 @@ def _clone_sources_ovstage(self):
             is_array=False,
             semantic=ovstage.AttributeSemantic.MATRIX,
         ).wait()
-        self._env_root_xforms = None
-        logger.info("Restored per-env root transforms after cloning")
 
         self._stage.release_query(env_query).wait()
         self._stage_paths.destroy_path_list(env_paths_list)
@@ -2220,7 +2105,6 @@ def _safe_destroy_path_list(path_list, name: str) -> None:
         self._deformable_particle_counts = []
         self._particle_visual_offsets = []
         self._particle_visual_counts = []
-        self._env_root_xforms = None
 
         # Detach before closing ExitStack: the renderer holds a live reference into the stage,
         # so detaching first avoids a use-after-free when ExitStack destroys Stage and PathDictionary.