diff --git a/source/isaaclab/changelog.d/fix-physx-newton-camera-pose-scaling.rst b/source/isaaclab/changelog.d/fix-physx-newton-camera-pose-scaling.rst new file mode 100644 index 000000000000..73108d32b7aa --- /dev/null +++ b/source/isaaclab/changelog.d/fix-physx-newton-camera-pose-scaling.rst @@ -0,0 +1,13 @@ +Fixed +^^^^^ + +* Fixed :class:`~isaaclab.scene_data.SceneDataProvider` transform mapping stalling + at high rigid-body counts, which delayed setup by minutes in scenes with + thousands of environments. + +Added +^^^^^ + +* Added :meth:`~isaaclab.sim.views.BaseFrameView.close` to release backend state + authored by a frame view. Backends also release best-effort on garbage + collection, but only an explicit close is deterministic. diff --git a/source/isaaclab/isaaclab/envs/utils/camera_view.py b/source/isaaclab/isaaclab/envs/utils/camera_view.py index a686e0e65b16..71aeae869890 100644 --- a/source/isaaclab/isaaclab/envs/utils/camera_view.py +++ b/source/isaaclab/isaaclab/envs/utils/camera_view.py @@ -250,10 +250,13 @@ def prim_world_positions( for env_id in env_indices: prim_path = env_path_from_template(prim_path_template, env_id) view = FrameView(prim_path, device="cpu", stage=stage) - if view.count != 1: - raise RuntimeError(f"expected one prim, got {view.count}") - pos_w, _ = view.get_world_poses() - pos = pos_w.torch[0].detach().cpu() + try: + if view.count != 1: + raise RuntimeError(f"expected one prim, got {view.count}") + pos_w, _ = view.get_world_poses() + pos = pos_w.torch[0].detach().cpu() + finally: + view.close() positions.append((float(pos[0]), float(pos[1]), float(pos[2]))) return torch.tensor(positions, dtype=torch.float32) except Exception: diff --git a/source/isaaclab/isaaclab/scene_data/scene_data_provider.py b/source/isaaclab/isaaclab/scene_data/scene_data_provider.py index 72263d6e6d1d..2442d37e4fd1 100644 --- a/source/isaaclab/isaaclab/scene_data/scene_data_provider.py +++ b/source/isaaclab/isaaclab/scene_data/scene_data_provider.py @@ -5,7 +5,6 @@ from __future__ import annotations -import contextlib import logging import re from collections import deque @@ -213,10 +212,13 @@ def create_mapping(self, paths: list[str | None]) -> wp.array(dtype=wp.int32) | paths or if no mapping is needed. """ if input_paths := self.backend.transform_paths: - mapping = [-1] * len(input_paths) - for i, path in enumerate(input_paths): - with contextlib.suppress(ValueError): - mapping[i] = paths.index(path) + # The map keeps resolution linear in the number of paths. For duplicate + # paths the first occurrence wins, matching ``list.index``. + path_to_out: dict[str | None, int] = {} + for out_idx, out_path in enumerate(paths): + if out_path not in path_to_out: + path_to_out[out_path] = out_idx + mapping = [path_to_out.get(path, -1) for path in input_paths] if not np.array_equal(mapping, np.arange(len(input_paths))): return wp.array(mapping, dtype=wp.int32) return None diff --git a/source/isaaclab/isaaclab/sensors/camera/camera.py b/source/isaaclab/isaaclab/sensors/camera/camera.py index 2484cffe00d2..a7892f58d35e 100644 --- a/source/isaaclab/isaaclab/sensors/camera/camera.py +++ b/source/isaaclab/isaaclab/sensors/camera/camera.py @@ -221,6 +221,12 @@ def __del__(self): """Unsubscribes from callbacks and cleans up renderer resources.""" # unsubscribe callbacks super().__del__() + # release the frame view's backend state (getattr: _view is assigned in + # _initialize_impl, so it is absent if construction failed earlier) + view = getattr(self, "_view", None) + if view is not None: + view.close() + self._view = None # cleanup render resources (renderer may be None if never initialized) if self._renderer is not None: self._renderer.cleanup(self._render_data) @@ -901,5 +907,7 @@ def _invalidate_initialize_callback(self, event): self._renderer = None # call parent super()._invalidate_initialize_callback(event) - # set all existing views to None to invalidate them - self._view = None + # release backend state deterministically, then invalidate the view + if self._view is not None: + self._view.close() + self._view = None diff --git a/source/isaaclab/isaaclab/sim/views/base_frame_view.py b/source/isaaclab/isaaclab/sim/views/base_frame_view.py index 79c672b7cfd0..8b14729177c8 100644 --- a/source/isaaclab/isaaclab/sim/views/base_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/base_frame_view.py @@ -62,6 +62,17 @@ def device(self) -> str: """Device where arrays are allocated (``"cpu"`` or ``"cuda:0"``).""" ... + def close(self) -> None: + """Release backend state authored by this view. The view must not be used afterwards. + + The base implementation is a no-op; backends that author persistent + state (e.g. the Fabric backend's per-view index attributes) override it. + Backends also release best-effort when the view is garbage collected, + but only an explicit :meth:`close` is deterministic -- collection + timing is up to the interpreter. Calling :meth:`close` more than once + is safe. + """ + # ------------------------------------------------------------------ # Write scope -- recommended API for all transform writes. # ------------------------------------------------------------------ diff --git a/source/isaaclab/isaaclab/utils/warp/fabric.py b/source/isaaclab/isaaclab/utils/warp/fabric.py index e0519d98c338..c2681552e1a5 100644 --- a/source/isaaclab/isaaclab/utils/warp/fabric.py +++ b/source/isaaclab/isaaclab/utils/warp/fabric.py @@ -21,6 +21,7 @@ IndexedFabricArrayMat44d = Any ArrayUInt32 = Any ArrayUInt32_1d = Any + ArrayInt32_1d = Any ArrayFloat32_2d = Any else: FabricArrayUInt32 = wp.fabricarray(dtype=wp.uint32) @@ -28,6 +29,7 @@ IndexedFabricArrayMat44d = wp.indexedfabricarray(dtype=wp.mat44d) ArrayUInt32 = wp.array(ndim=1, dtype=wp.uint32) ArrayUInt32_1d = wp.array(dtype=wp.uint32) + ArrayInt32_1d = wp.array(dtype=wp.int32) ArrayFloat32_2d = wp.array(ndim=2, dtype=wp.float32) @@ -46,6 +48,29 @@ def arange_k(a: ArrayUInt32_1d): a[tid] = wp.uint32(tid) +@wp.kernel(enable_backward=False) +def map_view_indices_to_fabric_slots(view_indices: FabricArrayUInt32, fabric_slots: ArrayInt32_1d): + """Invert a selection's per-prim view-index attribute into a slot lookup table. + + Inverts a permutation: ``view_indices`` holds each selected prim's view-side + index, and after the launch ``fabric_slots[view_index]`` is that prim's + fabric-side slot, ready to use as :class:`wp.indexedfabricarray` indices. + + The launch dimension must equal the selection's prim count, and the stored + view indices must cover ``0..dim-1`` exactly for the table to be complete. + """ + fabric_slot = int(wp.tid()) + view_index = int(view_indices[fabric_slot]) + fabric_slots[view_index] = fabric_slot + + +@wp.kernel(enable_backward=False) +def gather_fabric_slots(slots: ArrayInt32_1d, gather_map: ArrayUInt32_1d, out_slots: ArrayInt32_1d): + """Gather ``slots`` entries through ``gather_map``: ``out_slots[i] = slots[gather_map[i]]``.""" + i = int(wp.tid()) + out_slots[i] = slots[int(gather_map[i])] + + @wp.kernel(enable_backward=False) def decompose_fabric_transformation_matrix_to_warp_arrays( fabric_matrices: FabricArrayMat44d, diff --git a/source/isaaclab_mimic/changelog.d/fix-fabric-frameview-stall.rst b/source/isaaclab_mimic/changelog.d/fix-fabric-frameview-stall.rst new file mode 100644 index 000000000000..fb4ca7ed85b2 --- /dev/null +++ b/source/isaaclab_mimic/changelog.d/fix-fabric-frameview-stall.rst @@ -0,0 +1,5 @@ +Fixed +^^^^^ + +* Fixed :class:`SceneAsset` leaking its cached frame view when the view is rebuilt, + which left the view's backend state to be released on garbage collection. diff --git a/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py b/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py index 75b891de47d9..0b2506f48b3b 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py +++ b/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py @@ -112,6 +112,8 @@ def _get_xform_view(self) -> FrameView: cloned prims exist. """ if self._xform_view is None or self._xform_view.count == 0: + if self._xform_view is not None: + self._xform_view.close() entity = self.scene[self.entity_name] prim_path = ( entity.prim_path diff --git a/source/isaaclab_physx/changelog.d/fix-physx-newton-camera-pose-scaling.rst b/source/isaaclab_physx/changelog.d/fix-physx-newton-camera-pose-scaling.rst new file mode 100644 index 000000000000..72f5e00d6fe5 --- /dev/null +++ b/source/isaaclab_physx/changelog.d/fix-physx-newton-camera-pose-scaling.rst @@ -0,0 +1,14 @@ +Fixed +^^^^^ + +* Fixed camera world-pose resolution stalling at high environment counts under the + PhysX backend, which caused multi-second pauses between rendered frames and + benchmark timeouts. + +Added +^^^^^ + +* Added :meth:`close` to the PhysX Fabric frame view, removing its per-view Fabric + index attributes so that views recreated over the same prims no longer accumulate + attributes. Views dropped without closing are cleaned up on garbage collection, + with a warning. diff --git a/source/isaaclab_physx/isaaclab_physx/sim/views/fabric_frame_view.py b/source/isaaclab_physx/isaaclab_physx/sim/views/fabric_frame_view.py index 0faba8f08781..1788d2ea294a 100644 --- a/source/isaaclab_physx/isaaclab_physx/sim/views/fabric_frame_view.py +++ b/source/isaaclab_physx/isaaclab_physx/sim/views/fabric_frame_view.py @@ -7,7 +7,10 @@ from __future__ import annotations +import contextlib +import itertools import logging +import sys import torch import warp as wp @@ -24,6 +27,26 @@ logger = logging.getLogger(__name__) +def _parent_path(prim_path: str) -> str: + """Parent prim path of ``prim_path``. + + Args: + prim_path: Absolute prim path, so it always contains a separator. + + Raises: + RuntimeError: If the prim is directly under the stage root and thus has + no non-pseudoroot parent to read Fabric matrices from. + """ + parent = prim_path[: prim_path.rfind("/")] + if not parent: + raise RuntimeError( + f"Child prim '{prim_path}' is at the stage root and has no parent prim. " + "FabricFrameView requires every prim to have a non-pseudoroot parent " + "with Fabric world+local matrices." + ) + return parent + + def _to_float32_2d(a: wp.array | torch.Tensor) -> wp.array | torch.Tensor: """Ensure array is compatible with Fabric kernels (2-D float32). @@ -132,31 +155,23 @@ class FabricFrameView(BaseFrameView): :mod:`isaaclab.sim.views.xform_space_writer` for the full contract). The "torn data" concern is what motivates that no-step rule; it is separate from why the tracking pause exists. - * **Two persistent selections, flipped by the writer scope.** Two - selections are built once during ``_initialize_fabric`` and kept for - the view's lifetime: - - .. code-block:: text - - _sel_ro : worldMatrix=RO, localMatrix=RO (steady state) - _sel_rw : worldMatrix=RW, localMatrix=RW (inside writer scope) - - Each selection has its own bundle of indexed-fabric arrays - (``_world_ifa_*``, ``_local_ifa_*``, ``_parent_world_ifa_*``) cached - against the selection's path ordering. Writer ``__enter__`` flips an - ``_is_rw`` flag so subsequent get/set helpers resolve to the RW - bundle; ``__exit__`` flips back to RO. Nothing is rebuilt on the - flip -- both bundles are always kept consistent via independent - ``PrepareForReuse()`` polls in the accessors. - - The RO steady state tells Kit's next-tick - ``update_world_xforms()`` that no attribute is user-authored, so it - leaves both alone. Combined with the tracking pause and the - opposite-space derive at scope exit, this is what keeps the next - render tick from overwriting our writes. - * **Topology-adaptive.** Fabric topology changes are detected on each - access via per-selection ``PrepareForReuse()`` polls; the affected - indexed arrays rebuild automatically and no manual refresh is required. + * **Selections are scoped to the view, not the stage.** The view tags its + own prims (and their parents) with private per-view index attributes and + requires those attributes in every prim selection, so a selection resolves + to exactly the prims the view manages however large the stage grows. + Tag names are unique per view instance, so views never interfere with one + another. The tags are authored on first use and removed again by + :meth:`close` -- or, best-effort and with a warning, when the view is + garbage collected. Call :meth:`close` when done with a view; + collection timing is up to the interpreter, so relying on it can remove + the tags at an arbitrary point in the frame (or, on a leaked reference, + not at all). + * **Topology changes are absorbed, with no cache to invalidate.** The + view-to-Fabric mapping is re-derived from live Fabric data on every + access, so prims moving between Fabric buckets can never leave a stale + mapping behind. If a managed prim disappears (prim or attribute removed) + the next access raises :class:`RuntimeError` and the view must be + recreated. See ``_refresh_child_selection`` for how this is done. Pose getters return :class:`~isaaclab.utils.warp.ProxyArray`; the convenience :meth:`set_world_poses` / :meth:`set_local_poses` helpers accept @@ -168,6 +183,12 @@ class FabricFrameView(BaseFrameView): _WORLD_MATRIX_NAME = "omni:fabric:worldMatrix" _LOCAL_MATRIX_NAME = "omni:fabric:localMatrix" + # Process-wide uid source for per-view Fabric attribute names. A monotonic + # counter (NOT ``id(self)``/``hash(self)``) guarantees a name is never + # reused after a view is garbage-collected, so a dead view's leftover + # attributes can never satisfy a live view's selection. + _view_uid_counter = itertools.count() + def __init__( self, prim_path: str, @@ -203,33 +224,88 @@ def __init__( self._stage = None self._fabric_hierarchy = None - # Two persistent Fabric selections. ``_is_rw`` is True only inside - # an active writer scope; the accessors below resolve to the matching - # bundle of indexed arrays. + # Per-view Fabric index attributes (authored once in ``_initialize_fabric``). + self._child_index_attr: str | None = None + self._parent_index_attr: str | None = None + self._unique_parent_paths: list[str] = [] + + # Three persistent selections keyed on the index attributes: child RO + # (steady state), child RW (active inside a writer scope; ``_is_rw`` + # flips between them), and parent world (always read-only). self._sel_ro = None self._sel_rw = None + self._sel_parent = None self._is_rw: bool = False - # View-side indices array (shared across both bundles). + # View-side indices array. self._view_indices: wp.array | None = None - # Per-selection view->fabric mappings. - self._ro_fabric_indices: wp.array | None = None - self._rw_fabric_indices: wp.array | None = None - self._ro_parent_fabric_indices: wp.array | None = None - self._rw_parent_fabric_indices: wp.array | None = None - - # Indexed fabric arrays per (selection, attribute) pair. - self._world_ifa_ro = None - self._local_ifa_ro = None - self._parent_world_ifa_ro = None - self._world_ifa_rw = None - self._local_ifa_rw = None - self._parent_world_ifa_rw = None + # Kernel-built view->fabric slot mappings, refreshed on every selection + # access (see ``_refresh_child_selection``). ``_child_parent_map`` holds + # view-side indices (uint32, like the Fabric ``UInt`` index attributes); + # the ``*_slots_buf`` buffers hold Fabric slots and must be int32, the + # only dtype ``wp.indexedfabricarray`` accepts for indices. + self._child_parent_map: wp.array | None = None + self._child_slots_buf: wp.array | None = None + self._parent_slots_buf: wp.array | None = None + self._parent_slot_of_child_buf: wp.array | None = None # Sentinel passed to compose/decompose kernels for unused slots. self._fabric_empty_2d_array_sentinel: wp.array | None = None + # Index-attribute cleanup state (see ``close``): the ``(attribute, + # prims)`` groups authored by ``_initialize_fabric``, and the flag that + # makes ``close()`` idempotent and lets ``__del__`` warn when cleanup + # had to happen via garbage collection. + self._tagged_prims: list[tuple[str, list]] = [] + self._is_closed: bool = False + + def close(self) -> None: + """Remove this view's Fabric index attributes. The view must not be used afterwards. + + Calling :meth:`close` again is a no-op. If :meth:`close` is never + called, the same cleanup runs best-effort from ``__del__`` (with a + warning, since collection timing is up to the interpreter) -- except at + interpreter exit, where Fabric is being torn down anyway and the + attributes die with it. + """ + if self._is_closed: + return + self._is_closed = True + failed = total = 0 + for attr, prims in self._tagged_prims: + total += len(prims) + for prim in prims: + try: + prim.RemoveProperty(attr) + except Exception: # noqa: BLE001 -- one bad handle must not strand the remaining tags + failed += 1 + self._tagged_prims = [] + if failed: + logger.debug("FabricFrameView(%s): %d of %d tag removals failed", self._usd_view._prim_path, failed, total) + + def __del__(self, _sys=sys): + """Best-effort cleanup when the view is collected without :meth:`close`. + + Follows the repo's shutdown-safe ``__del__`` idiom (see + :meth:`~isaaclab.envs.ManagerBasedEnv.__del__`): ``sys`` is bound as a + default argument so it survives module teardown, and nothing runs during + interpreter finalization, when calling into Kit can crash and the + attributes die with Fabric anyway. + """ + # getattr: __init__ may have raised before the flag existed + if getattr(self, "_is_closed", True) or _sys.is_finalizing() or _sys.meta_path is None: + return + if self._tagged_prims: + logger.warning( + "FabricFrameView(%s) was garbage-collected without close(); its Fabric index " + "attributes were removed best-effort at an arbitrary point in the frame. Call " + "close() for deterministic cleanup.", + self._usd_view._prim_path, + ) + with contextlib.suppress(Exception): # never propagate from __del__ + self.close() + # ------------------------------------------------------------------ # Delegated properties # ------------------------------------------------------------------ @@ -278,7 +354,6 @@ def _make_local_space_writer(self) -> FrameViewLocalSpaceWriter: # ------------------------------------------------------------------ # Getter hooks -- read directly from Fabric (no lazy sync) # ------------------------------------------------------------------ - def _get_world_poses_impl(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: if not self._use_fabric: return self._usd_view._get_world_poses_impl(indices) @@ -433,13 +508,14 @@ def _recompute_local_from_world_all(self) -> None: Storage convention: see :func:`isaaclab.utils.warp.fabric.update_indexed_local_matrix_from_world`. """ + world_ifa, local_ifa = self._get_child_ifas() wp.launch( kernel=fabric_utils.update_indexed_local_matrix_from_world, dim=self.count, inputs=[ - self._get_world_ifa(), + world_ifa, self._get_parent_world_ifa(), - self._get_local_ifa(), + local_ifa, self._view_indices, ], device=self._device, @@ -453,91 +529,109 @@ def _recompute_world_from_local_all(self) -> None: Storage convention: see :func:`isaaclab.utils.warp.fabric.update_indexed_world_matrix_from_local`. """ + world_ifa, local_ifa = self._get_child_ifas() wp.launch( kernel=fabric_utils.update_indexed_world_matrix_from_local, dim=self.count, inputs=[ - self._get_local_ifa(), + local_ifa, self._get_parent_world_ifa(), - self._get_world_ifa(), + world_ifa, self._view_indices, ], device=self._device, ) # ------------------------------------------------------------------ - # Internal -- selection accessors with on-demand index rebuild + # Internal -- selection accessors (kernel-built slot mappings) # ------------------------------------------------------------------ - def _get_world_ifa(self): - self._refresh_active_bundle_if_needed() - return self._world_ifa_rw if self._is_rw else self._world_ifa_ro + def _get_world_ifa(self) -> wp.indexedfabricarray: + sel = self._refresh_child_selection() + return wp.indexedfabricarray(fa=wp.fabricarray(sel, self._WORLD_MATRIX_NAME), indices=self._child_slots_buf) - def _get_local_ifa(self): - self._refresh_active_bundle_if_needed() - return self._local_ifa_rw if self._is_rw else self._local_ifa_ro + def _get_local_ifa(self) -> wp.indexedfabricarray: + sel = self._refresh_child_selection() + return wp.indexedfabricarray(fa=wp.fabricarray(sel, self._LOCAL_MATRIX_NAME), indices=self._child_slots_buf) - def _get_parent_world_ifa(self): - self._refresh_active_bundle_if_needed() - return self._parent_world_ifa_rw if self._is_rw else self._parent_world_ifa_ro + def _get_child_ifas(self) -> tuple[wp.indexedfabricarray, wp.indexedfabricarray]: + """Return ``(world, local)`` child arrays from a single selection refresh. - def _refresh_active_bundle_if_needed(self) -> None: - """Rebuild the active bundle's indexed arrays if its selection's buckets changed.""" - if self._is_rw: - if self._world_ifa_rw is None or self._sel_rw.PrepareForReuse(): - self._rebuild_rw_arrays() - else: - if self._world_ifa_ro is None or self._sel_ro.PrepareForReuse(): - self._rebuild_ro_arrays() - - def _rebuild_ro_arrays(self) -> None: - """Rebuild the four ``_sel_ro``-keyed indexed arrays (children + parents).""" - self._ro_fabric_indices = self._compute_fabric_indices(self._sel_ro) - self._world_ifa_ro = self._build_indexed_array(self._sel_ro, self._WORLD_MATRIX_NAME, self._ro_fabric_indices) - self._local_ifa_ro = self._build_indexed_array(self._sel_ro, self._LOCAL_MATRIX_NAME, self._ro_fabric_indices) - self._ro_parent_fabric_indices = self._compute_parent_fabric_indices(self._sel_ro) - self._parent_world_ifa_ro = wp.indexedfabricarray( - fa=wp.fabricarray(self._sel_ro, self._WORLD_MATRIX_NAME), - indices=self._ro_parent_fabric_indices, + Callers that need both spaces must use this instead of calling + :meth:`_get_world_ifa` and :meth:`_get_local_ifa`, which would refresh + the same selection -- and re-run its mapping kernel -- twice. + """ + sel = self._refresh_child_selection() + return ( + wp.indexedfabricarray(fa=wp.fabricarray(sel, self._WORLD_MATRIX_NAME), indices=self._child_slots_buf), + wp.indexedfabricarray(fa=wp.fabricarray(sel, self._LOCAL_MATRIX_NAME), indices=self._child_slots_buf), ) - def _rebuild_rw_arrays(self) -> None: - """Rebuild the four ``_sel_rw``-keyed indexed arrays (children + parents).""" - self._rw_fabric_indices = self._compute_fabric_indices(self._sel_rw) - self._world_ifa_rw = self._build_indexed_array(self._sel_rw, self._WORLD_MATRIX_NAME, self._rw_fabric_indices) - self._local_ifa_rw = self._build_indexed_array(self._sel_rw, self._LOCAL_MATRIX_NAME, self._rw_fabric_indices) - self._rw_parent_fabric_indices = self._compute_parent_fabric_indices(self._sel_rw) - self._parent_world_ifa_rw = wp.indexedfabricarray( - fa=wp.fabricarray(self._sel_rw, self._WORLD_MATRIX_NAME), - indices=self._rw_parent_fabric_indices, + def _get_parent_world_ifa(self) -> wp.indexedfabricarray: + self._refresh_parent_selection() + return wp.indexedfabricarray( + fa=wp.fabricarray(self._sel_parent, self._WORLD_MATRIX_NAME), + indices=self._parent_slot_of_child_buf, ) - # ------------------------------------------------------------------ - # Internal -- index computation - # ------------------------------------------------------------------ + def _refresh_child_selection(self): + """Refresh the active child selection and rebuild its slot mapping on device. - def _compute_fabric_indices(self, selection) -> wp.array: - """View-side indices that map each managed prim into ``selection``.""" - return self._compute_fabric_indices_for(selection, list(self.prim_paths)) + Runs on every accessor call. ``PrepareForReuse`` lets the persistent + selection absorb Fabric bucket changes (and notifies the renderer for + the RW selection); a single Warp kernel launch over the selection's + index attribute then rebuilds ``_child_slots_buf`` so that entry ``i`` + is the fabric-side slot of view prim ``i``. Re-deriving the mapping + from live Fabric data on each access means bucket reorders can never + leave a stale mapping behind, with no host-side path resolution and no + cache to invalidate. - def _compute_parent_fabric_indices(self, selection) -> wp.array: - """View-side indices that map each managed prim's parent into ``selection``.""" + Returns: + The active (RO or RW) child prim selection. + """ + sel = self._sel_rw if self._is_rw else self._sel_ro + sel.PrepareForReuse() + self._check_selection_count(sel.GetCount(), self.count, self._child_index_attr) + wp.launch( + kernel=fabric_utils.map_view_indices_to_fabric_slots, + dim=self.count, + inputs=[wp.fabricarray(sel, self._child_index_attr), self._child_slots_buf], + device=self._device, + ) + return sel - def parent_path(prim_path: str) -> str: - p = prim_path.rsplit("/", 1)[0] - if not p: - raise RuntimeError( - f"Child prim '{prim_path}' is at stage root and has no parent prim. " - "FabricFrameView requires every prim to have a non-pseudoroot parent " - "with Fabric world+local matrices." - ) - return p + def _refresh_parent_selection(self) -> None: + """Refresh the parent selection and rebuild the per-child parent-slot mapping. - return self._compute_fabric_indices_for(selection, [parent_path(p) for p in self.prim_paths]) + Two kernel launches: the first inverts the parent index attribute into + per-ordinal fabric slots, the second gathers those slots per child + through ``_child_parent_map`` (children sharing a parent read the same + slot). + """ + num_parents = self._parent_slots_buf.shape[0] + self._sel_parent.PrepareForReuse() + self._check_selection_count(self._sel_parent.GetCount(), num_parents, self._parent_index_attr) + wp.launch( + kernel=fabric_utils.map_view_indices_to_fabric_slots, + dim=num_parents, + inputs=[wp.fabricarray(self._sel_parent, self._parent_index_attr), self._parent_slots_buf], + device=self._device, + ) + wp.launch( + kernel=fabric_utils.gather_fabric_slots, + dim=self.count, + inputs=[self._parent_slots_buf, self._child_parent_map, self._parent_slot_of_child_buf], + device=self._device, + ) - def _build_indexed_array(self, selection, attribute_name: str, fabric_indices: wp.array) -> wp.indexedfabricarray: - fa = wp.fabricarray(selection, attribute_name) - return wp.indexedfabricarray(fa=fa, indices=fabric_indices) + def _check_selection_count(self, found: int, expected: int, index_attr: str) -> None: + """Raise if a selection stopped matching exactly the view's tagged prims.""" + if found != expected: + raise RuntimeError( + f"FabricFrameView: selection on '{index_attr}' matched {found} prims, expected {expected}. " + "A prim managed by this view (or one of its Fabric matrix/index attributes) was removed " + "from the Fabric stage; recreate the view." + ) def _resolve_indices_wp(self, indices: wp.array | None) -> wp.array: """Resolve view indices as a Warp uint32 array.""" @@ -545,16 +639,19 @@ def _resolve_indices_wp(self, indices: wp.array | None) -> wp.array: if self._view_indices is None: raise RuntimeError("Fabric view indices are not initialized.") return self._view_indices - if indices.dtype != wp.uint32: - return wp.array(indices.numpy().astype("uint32"), dtype=wp.uint32, device=self._device) - return indices + if indices.dtype == wp.uint32: + return indices + if indices.dtype == wp.int32: + # Zero-copy reinterpret: callers (e.g. Camera) pass non-negative int32 indices. + # Device placement is not checked here; ``wp.launch`` validates it for every input. + return indices.view(wp.uint32) + return wp.array(indices.numpy().astype("uint32"), dtype=wp.uint32, device=self._device) # ------------------------------------------------------------------ # Internal -- Fabric initialization # ------------------------------------------------------------------ - def _initialize_fabric(self) -> None: - """One-time Fabric setup: hierarchy handle, attribute population, selections, indexed arrays.""" + """One-time Fabric setup: hierarchy handle, per-view index tagging, selections, buffers.""" import usdrt # noqa: PLC0415 # The hierarchy bindings are a separate submodule and are not loaded by ``import usdrt``. @@ -577,39 +674,72 @@ def _initialize_fabric(self) -> None: fabric_id, self._stage.GetStageIdAsStageId() ) - # Ensure each child prim AND its parent have BOTH Fabric world and local matrix - # attributes. ``Create*Attr`` calls are idempotent. - seen_paths: set[str] = set() - for child_path in self.prim_paths: - for path in (child_path, child_path.rsplit("/", 1)[0]): - if path in seen_paths: - continue - seen_paths.add(path) + # Per-view Fabric index attribute names (see ``_view_uid_counter``). + uid = next(FabricFrameView._view_uid_counter) + self._child_index_attr = f"isaaclab:fabricFrameView:{uid}:index" + self._parent_index_attr = f"isaaclab:fabricFrameView:{uid}:parentIndex" + + # Per-child parent paths, computed once and reused for the ordinal map + # below. Unique parents keep first-occurrence order; ``parent_ordinal`` + # maps a parent path to its position in that order. + child_parent_paths = [_parent_path(p) for p in self.prim_paths] + self._unique_parent_paths = list(dict.fromkeys(child_parent_paths)) + parent_ordinal = {path: i for i, path in enumerate(self._unique_parent_paths)} + + # Tag children and parents with their per-view index and ensure both + # carry the Fabric world+local matrix attributes (``Create*Attr`` calls + # are idempotent). The index attribute doubles as the selection filter: + # the selections below match ONLY tagged prims, so their size is + # O(view), not O(stage). A prim that is both a child and a parent of + # this view receives both index attributes. + tagged_prims: list[tuple[str, list]] = [] + for paths, index_attr in ( + (list(self.prim_paths), self._child_index_attr), + (self._unique_parent_paths, self._parent_index_attr), + ): + group_prims: list = [] + for i, path in enumerate(paths): rt_prim = self._stage.GetPrimAtPath(path) if not rt_prim.IsValid(): - continue + raise RuntimeError(f"FabricFrameView: prim '{path}' does not exist in the Fabric stage.") rt_xformable = Rt.Xformable(rt_prim) rt_xformable.CreateFabricHierarchyWorldMatrixAttr() rt_xformable.CreateFabricHierarchyLocalMatrixAttr() rt_xformable.SetLocalXformFromUsd() rt_xformable.SetWorldXformFromUsd() + rt_prim.CreateAttribute(index_attr, usdrt.Sdf.ValueTypeNames.UInt, custom=True) + rt_prim.GetAttribute(index_attr).Set(i) + group_prims.append(rt_prim) + tagged_prims.append((index_attr, group_prims)) + + # Remembered so ``close()`` / ``__del__`` can remove the tags again. + self._tagged_prims = tagged_prims - # Two persistent selections: all-RO (steady state) and all-RW (active - # only inside a writer scope). Each will own its own bundle of - # indexed-fabric arrays built lazily by ``_rebuild_{ro,rw}_arrays``. + # Three persistent selections keyed on the per-view index attributes: + # child RO (steady state), child RW (active only inside a writer + # scope), and parent world (always read-only). matrix = usdrt.Sdf.ValueTypeNames.Matrix4d + uint_type = usdrt.Sdf.ValueTypeNames.UInt ro = usdrt.Usd.Access.Read rw = usdrt.Usd.Access.ReadWrite + child_tag = (uint_type, self._child_index_attr, ro) + parent_tag = (uint_type, self._parent_index_attr, ro) wm_ro = (matrix, self._WORLD_MATRIX_NAME, ro) lm_ro = (matrix, self._LOCAL_MATRIX_NAME, ro) wm_rw = (matrix, self._WORLD_MATRIX_NAME, rw) lm_rw = (matrix, self._LOCAL_MATRIX_NAME, rw) - self._sel_ro = self._stage.SelectPrims(require_attrs=[wm_ro, lm_ro], device=self._device, want_paths=True) - self._sel_rw = self._stage.SelectPrims(require_attrs=[wm_rw, lm_rw], device=self._device, want_paths=True) + self._sel_ro = self._stage.SelectPrims(require_attrs=[child_tag, wm_ro, lm_ro], device=self._device) + self._sel_rw = self._stage.SelectPrims(require_attrs=[child_tag, wm_rw, lm_rw], device=self._device) + self._sel_parent = self._stage.SelectPrims(require_attrs=[parent_tag, wm_ro], device=self._device) + # View-side indices + kernel-built slot-mapping buffers. self._view_indices = wp.array(list(range(self.count)), dtype=wp.uint32, device=self._device) - self._rebuild_ro_arrays() - self._rebuild_rw_arrays() + self._child_parent_map = wp.array( + [parent_ordinal[p] for p in child_parent_paths], dtype=wp.uint32, device=self._device + ) + self._child_slots_buf = wp.empty((self.count,), dtype=wp.int32, device=self._device) + self._parent_slots_buf = wp.empty((len(self._unique_parent_paths),), dtype=wp.int32, device=self._device) + self._parent_slot_of_child_buf = wp.empty((self.count,), dtype=wp.int32, device=self._device) # Pre-allocated reusable output buffers (world + local + scales). self._fabric_positions_buf = wp.zeros((self.count, 3), dtype=wp.float32, device=self._device) @@ -628,8 +758,8 @@ def _initialize_fabric(self) -> None: self._fabric_initialized = True # Seed Fabric matrices from USD authoritatively. The seed writes, so - # flip into the RW bundle for its duration; flip back to RO afterwards - # so steady-state getters use the RO bundle. + # flip onto the RW selection for its duration; flip back afterwards so + # steady-state getters use the RO selection. self._is_rw = True try: self._sync_fabric_from_usd_initial() @@ -650,7 +780,7 @@ def _sync_fabric_from_usd_initial(self) -> None: kernel=fabric_utils.compose_indexed_fabric_transforms, dim=self.count, inputs=[ - self._local_ifa_rw, # explicit RW: init-time write, no scope yet + self._get_local_ifa(), # caller holds ``_is_rw=True``: init-time write, no scope yet _to_float32_2d(local_pos_ta.warp), _to_float32_2d(local_ori_ta.warp), _to_float32_2d(scales_wp), @@ -663,8 +793,10 @@ def _sync_fabric_from_usd_initial(self) -> None: ) # --- Parents (one entry per unique parent path) --- - unique_parent_paths = list(dict.fromkeys(p.rsplit("/", 1)[0] for p in self.prim_paths)) + unique_parent_paths = self._unique_parent_paths if unique_parent_paths: + import usdrt # noqa: PLC0415 + from isaaclab.sim.utils import get_current_stage # noqa: PLC0415 usd_stage = get_current_stage() @@ -707,9 +839,25 @@ def _sync_fabric_from_usd_initial(self) -> None: parent_pos_wp = wp.array(world_pos_rows, dtype=wp.float32, device=self._device) parent_ori_wp = wp.array(world_ori_rows, dtype=wp.float32, device=self._device) parent_scale_wp = wp.array(world_scale_rows, dtype=wp.float32, device=self._device) + # One-off RW selection on the parent tag for the initial seed; the + # persistent ``_sel_parent`` stays read-only for steady-state reads. + sel_parent_rw = self._stage.SelectPrims( + require_attrs=[ + (usdrt.Sdf.ValueTypeNames.UInt, self._parent_index_attr, usdrt.Usd.Access.Read), + (usdrt.Sdf.ValueTypeNames.Matrix4d, self._WORLD_MATRIX_NAME, usdrt.Usd.Access.ReadWrite), + ], + device=self._device, + ) + self._check_selection_count(sel_parent_rw.GetCount(), len(unique_parent_paths), self._parent_index_attr) + wp.launch( + kernel=fabric_utils.map_view_indices_to_fabric_slots, + dim=len(unique_parent_paths), + inputs=[wp.fabricarray(sel_parent_rw, self._parent_index_attr), self._parent_slots_buf], + device=self._device, + ) parent_world_rw = wp.indexedfabricarray( - fa=wp.fabricarray(self._sel_rw, self._WORLD_MATRIX_NAME), - indices=self._compute_fabric_indices_for(self._sel_rw, unique_parent_paths), + fa=wp.fabricarray(sel_parent_rw, self._WORLD_MATRIX_NAME), + indices=self._parent_slots_buf, ) wp.launch( kernel=fabric_utils.compose_indexed_fabric_transforms, @@ -733,24 +881,6 @@ def _sync_fabric_from_usd_initial(self) -> None: self._recompute_world_from_local_all() wp.synchronize() - def _compute_fabric_indices_for(self, selection, paths: list[str]) -> wp.array: - """Look up each path in ``selection`` and return the matching fabric-side indices. - - Shared primitive used by :meth:`_compute_fabric_indices` (children), - :meth:`_compute_parent_fabric_indices` (parents), and one-off - index arrays such as the parent-world seed in - :meth:`_sync_fabric_from_usd_initial`. - """ - path_to_idx = {str(p): i for i, p in enumerate(selection.GetPaths())} - - def lookup(path: str) -> int: - idx = path_to_idx.get(path) - if idx is None: - raise RuntimeError(f"Path '{path}' not found in Fabric selection.") - return idx - - return wp.array([lookup(p) for p in paths], dtype=wp.int32, device=self._device) - # ---------------------------------------------------------------------- # Concrete writer classes for FabricFrameView @@ -763,11 +893,11 @@ class _FabricWriterMixin: On enter: pauses ``track_local_xform_changes`` / ``track_world_xform_changes`` on the Fabric hierarchy (saving prior state) and flips the view's ``_is_rw`` so all get/set helpers resolve to the persistent RW selection - bundle (no rebuild -- both bundles are kept alive for the view's lifetime). + (both selections are kept alive for the view's lifetime). On exit (normal or via exception): runs a best-effort opposite-space derive + ``wp.synchronize()`` whenever any write happened inside the - scope, then flips ``_is_rw`` back to ``False`` (RO bundle for + scope, then flips ``_is_rw`` back to ``False`` (RO selection for steady-state reads) and restores hierarchy-tracking state. **Exception safety.** If the scope unwinds because of an exception @@ -838,7 +968,7 @@ def _derive_opposite(self) -> None: class _FabricWorldSpaceWriter(_FabricWriterMixin, FrameViewWorldSpaceWriter): """World-space writer for :class:`FabricFrameView`. - Writes flow through ``_world_ifa_rw`` (the RW-bundle worldMatrix array); + Writes flow through the RW selection's ``worldMatrix`` indexed array; on exit ``localMatrix`` is derived from the just-written ``worldMatrix`` via :func:`update_indexed_local_matrix_from_world`. """ @@ -899,7 +1029,7 @@ def get_scales(self, indices=None) -> ProxyArray: class _FabricLocalSpaceWriter(_FabricWriterMixin, FrameViewLocalSpaceWriter): """Local-space writer for :class:`FabricFrameView`. - Writes flow through ``_local_ifa_rw`` (the RW-bundle localMatrix array); + Writes flow through the RW selection's ``localMatrix`` indexed array; on exit ``worldMatrix`` is derived from the just-written ``localMatrix`` via :func:`update_indexed_world_matrix_from_local`. """ diff --git a/source/isaaclab_physx/test/sim/test_views_xform_prim_fabric.py b/source/isaaclab_physx/test/sim/test_views_xform_prim_fabric.py index 849e99bf0778..3cc92147827c 100644 --- a/source/isaaclab_physx/test/sim/test_views_xform_prim_fabric.py +++ b/source/isaaclab_physx/test/sim/test_views_xform_prim_fabric.py @@ -10,6 +10,7 @@ Camera prim type for Fabric SelectPrims compatibility). """ +import logging import sys from pathlib import Path @@ -94,7 +95,7 @@ def _set_parent_positions(positions, num_envs): @pytest.fixture -def view_factory(): +def view_factory(request): """Fabric factory: Camera child at CHILD_OFFSET under parent Xforms, with Fabric enabled.""" def factory(num_envs: int, device: str) -> ViewBundle: @@ -107,11 +108,15 @@ def factory(num_envs: int, device: str) -> ViewBundle: sim_utils.SimulationContext(sim_utils.SimulationCfg(dt=0.01, device=device, use_fabric=True)) view = FrameView("/World/Parent_.*/Child", device=device) + # close() is idempotent, so this is safe even for tests that close (or + # tear down) themselves; it keeps views from being reaped by garbage + # collection, which would log the missing-close() warning per test. + request.addfinalizer(view.close) return ViewBundle( view=view, get_parent_pos=_get_parent_positions, set_parent_pos=_set_parent_positions, - teardown=lambda: None, + teardown=view.close, ) return factory @@ -185,14 +190,11 @@ def test_fabric_set_world_does_not_write_back_to_usd(device, view_factory): @pytest.mark.parametrize("device", test_devices()) def test_fabric_rebuild_after_topology_change(device, view_factory): - """A simulated topology change rebuilds the indexed fabric arrays and leaves - the view in a state where subsequent writes/reads still produce correct data. - - Real ``PrimSelection.PrepareForReuse`` reports topology change only when Fabric - reallocates internally, which is hard to provoke from a unit test. Instead we - invoke :meth:`FabricFrameView._compute_fabric_indices` and rebuild the indexed - arrays manually, mimicking what ``_get_*_array`` would do on a real topology - event, then verify a roundtrip still works. + """Refreshing every selection mid-use leaves writes and reads correct. + + ``PrepareForReuse`` only reports a topology change when Fabric reallocates + internally, which this test does not provoke, so this is a smoke test of the + refresh paths rather than true topology-recovery coverage. """ bundle = view_factory(2, device) view = bundle.view @@ -203,10 +205,15 @@ def test_fabric_rebuild_after_topology_change(device, view_factory): with view.xform_world_space_writer() as w: w.set_poses(positions=initial) - # Simulate topology change: rebuild both selection bundles, mirroring the - # lazy paths in the ``_refresh_active_bundle_if_needed`` accessor. - view._rebuild_ro_arrays() - view._rebuild_rw_arrays() + # Simulate topology change: refresh both child selections and the parent + # selection, mirroring the accessor paths. + view._refresh_child_selection() # RO (steady state) + view._is_rw = True + try: + view._refresh_child_selection() # RW (writer scope) + finally: + view._is_rw = False + view._refresh_parent_selection() # Trigger another write through the rebuilt arrays. new = wp.zeros((2, 3), dtype=wp.float32, device=device) @@ -297,6 +304,90 @@ def test_prepare_for_reuse_detects_topology_change(device, view_factory): assert not result, "PrepareForReuse should return False when no topology change" +@pytest.mark.parametrize("device", test_devices()) +def test_selections_match_only_the_view_prims(device, view_factory): + """Selections contain only the managed child prims and their unique parents. + + Without the per-view index attribute in the selection predicate the child + selections also pick up the parents (and, on a real stage, every other + xformable), so this fails with "matched 8 prims, expected 4". + """ + num_envs = 4 + bundle = view_factory(num_envs, device) + view = bundle.view + view.get_world_poses() # trigger Fabric init + + for name in ("_sel_ro", "_sel_rw"): + count = getattr(view, name).GetCount() + assert count == view.count, ( + f"{name} matched {count} prims but the view manages {view.count}. " + "The selection is not scoped by the per-view index attribute, so it is " + "picking up unrelated prims from the stage." + ) + parent_count = view._sel_parent.GetCount() + assert parent_count == num_envs, f"parent selection matched {parent_count} prims, expected {num_envs}" + + +def _count_prims_with_tag(view, attr: str) -> int: + """Number of prims on the view's Fabric stage carrying ``attr``.""" + import usdrt # noqa: PLC0415 + + sel = view._stage.SelectPrims( + require_attrs=[(usdrt.Sdf.ValueTypeNames.UInt, attr, usdrt.Usd.Access.Read)], device="cpu" + ) + return sel.GetCount() + + +@pytest.mark.parametrize("device", ["cuda:0"]) +def test_close_removes_index_attributes(device, view_factory): + """close() removes the view's Fabric index tags; a second close is a no-op.""" + bundle = view_factory(2, device) + view = bundle.view + view.get_world_poses() # trigger Fabric init (authors the tags) + + child_attr = view._child_index_attr + assert _count_prims_with_tag(view, child_attr) == view.count + view.close() + assert _count_prims_with_tag(view, child_attr) == 0, "close() left index attributes behind" + view.close() # idempotent + + +@pytest.mark.parametrize("device", ["cuda:0"]) +def test_garbage_collection_removes_index_attributes_and_warns(device, caplog): + """Dropping a view without close() still removes its tags, with a warning. + + Builds the view directly instead of via ``view_factory``: the fixture + registers ``view.close`` as a finalizer, and that bound method would keep + the view alive past the ``del`` below. + """ + import gc # noqa: PLC0415 + + _skip_if_unavailable(device) + stage_usd = sim_utils.get_current_stage() + for i in range(2): + sim_utils.create_prim(f"/World/Parent_{i}", "Xform", translation=PARENT_POS, stage=stage_usd) + sim_utils.create_prim(f"/World/Parent_{i}/Child", "Camera", translation=CHILD_OFFSET, stage=stage_usd) + sim_utils.SimulationContext(sim_utils.SimulationCfg(dt=0.01, device=device, use_fabric=True)) + view = FrameView("/World/Parent_.*/Child", device=device) + view.get_world_poses() + + child_attr = view._child_index_attr + stage = view._stage # keep a stage handle to count tags after the view dies + assert _count_prims_with_tag(view, child_attr) == view.count + + with caplog.at_level(logging.WARNING, logger="isaaclab_physx.sim.views.fabric_frame_view"): + del view + gc.collect() + + import usdrt # noqa: PLC0415 + + sel = stage.SelectPrims( + require_attrs=[(usdrt.Sdf.ValueTypeNames.UInt, child_attr, usdrt.Usd.Access.Read)], device="cpu" + ) + assert sel.GetCount() == 0, "garbage collection left index attributes behind" + assert any("without close()" in r.message for r in caplog.records), "expected a close() warning" + + def _read_fabric_world_matrix_translation(view, prim_index=0): """Read cached Fabric worldMatrix directly, without FrameView getter sync.""" rt_prim = view._stage.GetPrimAtPath(view.prim_paths[prim_index]) @@ -484,6 +575,7 @@ def test_set_local_then_get_world_with_rotated_parent(device): world_pos, _ = view.get_world_poses() expected = torch.tensor([[0.0, 1.0, 1.0]], dtype=torch.float32, device=device) torch.testing.assert_close(torch.as_tensor(world_pos, device=device), expected, atol=1e-5, rtol=0) + view.close() @pytest.mark.parametrize("device", ["cpu", "cuda:0"]) @@ -506,6 +598,7 @@ def test_set_world_then_get_local_with_rotated_parent(device): local_pos, _ = view.get_local_poses() expected = torch.tensor([[0.0, -5.0, 1.0]], dtype=torch.float32, device=device) torch.testing.assert_close(torch.as_tensor(local_pos, device=device), expected, atol=1e-5, rtol=0) + view.close() @pytest.mark.parametrize("device", ["cpu", "cuda:0"]) @@ -554,6 +647,7 @@ def test_initial_seed_with_scaled_parent(device): atol=1e-5, rtol=0, ) + view.close() # ------------------------------------------------------------------ @@ -632,6 +726,8 @@ def test_multi_view_writer_isolation(device): assert view_b._active_writer is not None assert view_a._active_writer is None assert view_b._active_writer is None + view_a.close() + view_b.close() # ------------------------------------------------------------------ @@ -789,6 +885,7 @@ def test_sequential_world_then_local_scopes_partial_indices(device): atol=1e-5, rtol=0, ) + view.close() @pytest.mark.parametrize("device", ["cpu", "cuda:0"]) @@ -837,6 +934,7 @@ def test_sequential_local_then_world_scopes_partial_indices(device): atol=1e-5, rtol=0, ) + view.close() # ------------------------------------------------------------------