From b0fd052e46e790c594f111c158cb6a3bf21a3858 Mon Sep 17 00:00:00 2001 From: alex Date: Mon, 24 Aug 2026 16:33:37 +0200 Subject: [PATCH 1/4] Shim needed for Arena compatibility. --- .../isaaclab/isaaclab/cloner/cloner_utils.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/source/isaaclab/isaaclab/cloner/cloner_utils.py b/source/isaaclab/isaaclab/cloner/cloner_utils.py index 86d7232176f5..22ba6d8a9762 100644 --- a/source/isaaclab/isaaclab/cloner/cloner_utils.py +++ b/source/isaaclab/isaaclab/cloner/cloner_utils.py @@ -22,6 +22,27 @@ logger = logging.getLogger(__name__) +def split_clone_template(destination_template: str) -> tuple[str, str]: + """Split a clone destination template around its clone slot. + + The ``"{}"`` slot represents one concrete environment/instance path segment. + + Args: + destination_template: Destination path template with ``"{}"`` for the instance id. + + Returns: + The ``(prefix, suffix)`` around the clone slot. + + Raises: + ValueError: If ``destination_template`` does not contain a clone slot. + """ + destination_template = destination_template.rstrip("/") or "/" + prefix, slot, suffix = destination_template.partition("{}") + if slot != "{}": + raise ValueError(f"Clone destination template must contain '{{}}': {destination_template!r}.") + return prefix, suffix + + def get_suffix(path_expr: str, destination_template: str) -> str | None: """Return the part of ``path_expr`` below a destination template's env-instance root. From 4f5feefc847821a81af8658e2132561b8d8bea40 Mon Sep 17 00:00:00 2001 From: pv-nvidia <197907000+pv-nvidia@users.noreply.github.com> Date: Tue, 7 Jul 2026 14:01:02 +0200 Subject: [PATCH 2/4] feat: GPU-accelerated FrameView local poses + a new xform_space_writer API (#5677) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two things: 1. **Local-space pose and scale ops are now GPU-fast on the Fabric backend.** They go through the same Warp-kernel path that world-space ops already used. The old USD fallback for local-space is ~100× slower (see the table below). 2. **All transform writes go through a small context manager.** You open a scope, do your writes inside it, and the scope handles the cleanup: ```python with view.xform_world_space_writer() as writer: writer.set_poses(positions=p, orientations=o) writer.set_scales(scales=s) ``` Use `view.xform_local_space_writer()` for local-space writes. Builds on Piotr's prototype at [`bareya/pbarejko/camera-update`](https://github.com/bareya/IsaacLab/tree/pbarejko/camera-update). A FrameView keeps two copies of each prim's transform in Fabric storage: the `omni:fabric:worldMatrix` and `omni:fabric:localMatrix` attributes. When you write one, the other has to be recomputed so they stay consistent. Without a scope, that recompute has to run after every single write. With a scope, it runs **once**, when the scope closes. You can call `set_poses` and `set_scales` as many times as you want inside the scope; on exit, one Warp kernel derives the other space and one `wp.synchronize()` runs. Empty scopes cost nothing. The scope also pauses **Fabric Hierarchy**'s transform-change tracking while you're writing, then restores it on exit. (Fabric itself is just the data store; Fabric Hierarchy, exposed through USDRT's `IFabricHierarchy`, is the plugin that watches writes to `omni:fabric:localMatrix` / `omni:fabric:worldMatrix` and keeps them mutually consistent across the hierarchy.) Its tracking is pull-based: a per-attribute listener records writes into a changelog, and the plugin drains and processes that changelog on the next call to `IFabricHierarchy::update_world_xforms()`. Pausing the listener does not block the writes themselves — they still land in Fabric storage — it just keeps them out of the changelog. With the changelog empty, the next tick has nothing to "catch up on" for these prims, and can't decide one of the two spaces is canonical and derive the other from it. See the next section for why that matters. Rules: - Only one writer can be open on a view at a time. Opening a second one raises `RuntimeError`. - While a writer is open, the view's own getters (`get_world_poses`, etc.) raise. Read through the writer (`writer.get_poses()`) inside the scope, or close the scope first. - **Do not advance the simulation or render from inside the scope.** No `sim.step()`, no `world.render()`, no `SimulationApp.update()`. The matrices are mid-write until exit, and rendering against that state would read torn data. Keep scopes short and close them before stepping. - If something raises inside the scope (a user-code bug, a notebook cell interrupt, etc.), the scope still runs the opposite-space recompute on exit so the world and local matrices stay consistent prim-by-prim. The partial write itself is not rolled back — if you need all-or-nothing, snapshot the matrices yourself before entering. correctness) A writer scope is synchronous Python code, so the renderer can't see anything mid-scope — no tick runs until we exit. The risk is on the **next** render/sim tick after the scope closes. On that next tick, Fabric Hierarchy runs its `IFabricHierarchy::update_world_xforms()` step, which can recompute `omni:fabric:worldMatrix` from `omni:fabric:localMatrix` (or vice versa). If it thinks one of the two was the user's most recent edit, it derives the other from it; the derived value — not what we wrote — is then what the **Fabric Scene Delegate (FSD)** hands to RTX on the render path. Three things in this PR keep that next-tick recompute from clobbering our work: 1. **Opposite-space derive at scope exit.** By the time the scope closes, `omni:fabric:worldMatrix` and `omni:fabric:localMatrix` are mutually consistent prim-by-prim, so any recompute Fabric Hierarchy does is a no-op. 2. **Fabric-Hierarchy change tracking is paused for the whole scope.** The writer calls `track_local_xform_changes(False)` and `track_world_xform_changes(False)` on enter (saving the prior state) and restores them on exit. While tracking is paused, the listeners don't record our writes in their changelog, so when the next tick runs there's nothing for `update_world_xforms()` to "catch up on". 3. **Two persistent selections with explicit read-only / read-write flags.** Outside the scope, the view's selection is fully read-only — both `omni:fabric:worldMatrix` and `omni:fabric:localMatrix` are flagged RO. Inside the scope, the writer flips to a fully read-write selection. ``` _sel_ro : worldMatrix=RO, localMatrix=RO (steady state, between writes) _sel_rw : worldMatrix=RW, localMatrix=RW (active during a writer scope) ``` Both selections are built once when the view is initialized and kept alive for its lifetime. The writer scope flips a single flag (`_is_rw`) on enter/exit; nothing is rebuilt. The RO steady state tells Fabric Hierarchy's next-tick `update_world_xforms()` that no attribute is user-authored, so it leaves both alone. - **Fabric (PhysX):** the fast GPU path described above. - **USD, Newton, OVPhysX:** the writer is a thin wrapper. Writes go straight to the backend's existing storage. No batching, no extra recompute (there is no second matrix to keep in sync). Same API across all four backends. Newton has two different "scale" ideas, and the API keeps them apart: - `NewtonSiteFrameView.set_scales(...)` (deprecated) writes the **Newton collision-shape geometry** size. - The writer's `set_scales(...)` writes per-site **transform scale**, matching what USD FrameView does. They operate on different state and are not merged. **New (recommended):** - `view.xform_world_space_writer()` / `view.xform_local_space_writer()` — a context manager with `set_poses`, `set_scales`, `get_poses`, `get_scales`. **Deprecated (still works, warns once per class on first use):** - `set_world_poses` / `set_local_poses` — use the writer scope. - `get_scales` / `set_scales` — use `get_local_scales` / `get_world_scales`, or the writer's `set_scales`. Per-iteration timings (lower is better): | Operation | USD (ms) | Fabric (ms) | Newton Site (ms) | |----------------------------|---------:|------------:|-----------------:| | Get World Poses | 8.6148 | 0.0375 | 0.0307 | | Set World Poses | 19.1357 | 0.1090 | 0.0368 | | Get Local Poses | 5.9556 | 0.0374 | 0.0497 | | Set Local Poses | 7.9794 | 0.0979 | 0.0790 | | Get World Scales | 13.1193 | 0.0371 | 0.0015 | | Set World Scales | 21.4571 | 0.1014 | 0.0612 | | Get Local Scales | 3.2544 | 0.0375 | 0.0022 | | Set Local Scales | 3.8350 | 0.0937 | 0.0587 | | Get Both (World+Local) | 14.7506 | 0.0738 | 0.1302 | | Interleaved World Set→Get | 28.1046 | 0.1432 | 0.1588 | | **Per-iter total** | **126.2** | **0.77** | **0.61** | Speedup vs USD (per-iter ops; one-time view construction excluded): | Operation | Fabric × | Newton Site × | |----------------------------|---------:|--------------:| | Get World Poses | 229.6× | 280.2× | | Set World Poses | 175.6× | 519.8× | | Get Local Poses | 159.0× | 119.9× | | Set Local Poses | 81.5× | 101.0× | | Get World Scales | 353.9× | 8911.5× | | Set World Scales | 211.6× | 350.4× | | Get Local Scales | 86.8× | 1489.9× | | Set Local Scales | 40.9× | 65.4× | | Get Both (World+Local) | 200.0× | 113.3× | | Interleaved World Set→Get | 196.2× | 177.0× | | **Overall (per-iter)** | **164×** | **207×** | One-time view construction (reported separately, not part of the per-iter total): USD 4.6 ms, Fabric 4.4 ms, Newton Site 1013 ms — the Newton number is dominated by stage population on the first call. Steady-state per-iteration cost is what the speedup row reflects. - Getters call `wp.synchronize()` before returning, so the returned `ProxyArray` is always immediately readable from both GPU and host code — no caller-side sync needed. (Both the cached and the per-indices paths sync; this used to be asymmetric.) --------- Co-authored-by: pv-nvidia Co-authored-by: pv-nvidia <{"message":"Not Found","documentation_url":"https://docs.github.com/rest/users/emails#list-email-addresses-for-the-authenticated-user","status":"404"}> --- docs/source/api/lab/isaaclab.utils.rst | 13 + .../benchmarks/benchmark_view_comparison.py | 3 +- .../benchmarks/benchmark_xform_prim_view.py | 160 ++- .../changelog.d/fabric-local-poses.rst | 31 + .../changelog.d/xform-space-writer.rst | 35 + .../isaaclab/sensors/camera/camera.py | 14 +- .../isaaclab/isaaclab/sim/views/__init__.pyi | 4 + .../isaaclab/sim/views/base_frame_view.py | 271 ++++- .../isaaclab/sim/views/usd_frame_view.py | 246 +++-- .../isaaclab/sim/views/xform_space_writer.py | 145 +++ source/isaaclab/isaaclab/utils/warp/fabric.py | 176 ++++ .../test/sim/frame_view_contract_utils.py | 188 +++- .../test/sim/test_views_xform_prim.py | 70 +- .../test/terrains/check_terrain_importer.py | 3 +- .../test/terrains/test_terrain_importer.py | 3 +- .../changelog.d/xform-space-writer.skip | 1 + .../locomanipulation_sdg/scene_utils.py | 3 +- .../changelog.d/fabric-local-poses.rst | 19 + .../changelog.d/xform-space-writer.rst | 13 + .../sim/views/newton_site_frame_view.py | 240 ++++- .../test/sim/test_views_xform_prim_newton.py | 3 +- .../changelog.d/fabric-local-poses.rst | 19 + .../changelog.d/xform-space-writer.rst | 13 + .../sim/views/ovphysx_frame_view.py | 167 +-- .../changelog.d/fabric-local-poses.rst | 36 + .../changelog.d/xform-space-writer.rst | 36 + .../sim/views/fabric_frame_view.py | 968 ++++++++++++++---- .../test/sim/test_views_xform_prim_fabric.py | 827 +++++++++++++-- 28 files changed, 3122 insertions(+), 585 deletions(-) create mode 100644 source/isaaclab/changelog.d/fabric-local-poses.rst create mode 100644 source/isaaclab/changelog.d/xform-space-writer.rst create mode 100644 source/isaaclab/isaaclab/sim/views/xform_space_writer.py create mode 100644 source/isaaclab_mimic/changelog.d/xform-space-writer.skip create mode 100644 source/isaaclab_newton/changelog.d/fabric-local-poses.rst create mode 100644 source/isaaclab_newton/changelog.d/xform-space-writer.rst create mode 100644 source/isaaclab_ovphysx/changelog.d/fabric-local-poses.rst create mode 100644 source/isaaclab_ovphysx/changelog.d/xform-space-writer.rst create mode 100644 source/isaaclab_physx/changelog.d/fabric-local-poses.rst create mode 100644 source/isaaclab_physx/changelog.d/xform-space-writer.rst diff --git a/docs/source/api/lab/isaaclab.utils.rst b/docs/source/api/lab/isaaclab.utils.rst index 5b352152e0b5..59c78b3aef48 100644 --- a/docs/source/api/lab/isaaclab.utils.rst +++ b/docs/source/api/lab/isaaclab.utils.rst @@ -188,3 +188,16 @@ Warp operations :members: :imported-members: :show-inheritance: + +Warp Fabric kernels +^^^^^^^^^^^^^^^^^^^ + +Warp kernels for reading and writing Fabric ``Matrix4d`` attributes +(``omni:fabric:worldMatrix`` / ``omni:fabric:localMatrix``) via +:class:`wp.fabricarray` and :class:`wp.indexedfabricarray`. Used by +:class:`~isaaclab_physx.sim.views.FabricFrameView` to keep child world and +local matrices consistent without round-tripping through USD. + +.. automodule:: isaaclab.utils.warp.fabric + :members: + :show-inheritance: diff --git a/scripts/benchmarks/benchmark_view_comparison.py b/scripts/benchmarks/benchmark_view_comparison.py index aa5927e10b6a..2f19554d5ffc 100644 --- a/scripts/benchmarks/benchmark_view_comparison.py +++ b/scripts/benchmarks/benchmark_view_comparison.py @@ -284,7 +284,8 @@ def _run_pose_benchmarks( start_time = time.perf_counter() for _ in range(num_iterations): - view.set_world_poses(new_positions, orientations) + with view.xform_world_space_writer() as w: + w.set_poses(new_positions, orientations) timing_results["set_world_poses"] = (time.perf_counter() - start_time) / num_iterations ret_pos, ret_quat = view.get_world_poses() diff --git a/scripts/benchmarks/benchmark_xform_prim_view.py b/scripts/benchmarks/benchmark_xform_prim_view.py index fee3b9642c79..f5c13a0242c1 100644 --- a/scripts/benchmarks/benchmark_xform_prim_view.py +++ b/scripts/benchmarks/benchmark_xform_prim_view.py @@ -138,21 +138,35 @@ def benchmark_frame_view( # noqa: C901 is_newton = api == "isaaclab-newton-site" + # Synchronize around timed regions using Warp directly (rather than torch), + # since all backend kernels here are Warp launches and ``wp.synchronize()`` + # covers CPU/CUDA Warp devices consistently. We only need it for GPU runs, + # where kernel launches are asynchronous; guard on device to avoid paying + # it needlessly on CPU. + _needs_sync = str(device).startswith("cuda") + + def _sync() -> None: + if _needs_sync: + wp.synchronize() + def to_torch(a): - return wp.to_torch(a) if isinstance(a, wp.array) else a + if isinstance(a, wp.array): + return wp.to_torch(a) + if hasattr(a, "torch"): + return a.torch + return a try: # -- Warmup -------------------------------------------------------- xform_view.get_world_poses() + xform_view.get_world_scales() # -- get_world_poses ----------------------------------------------- - if is_newton: - torch.cuda.synchronize() + _sync() start_time = time.perf_counter() for _ in range(num_iterations): positions, orientations = xform_view.get_world_poses() - if is_newton: - torch.cuda.synchronize() + _sync() timing_results["get_world_poses"] = (time.perf_counter() - start_time) / num_iterations positions_t = to_torch(positions) @@ -161,20 +175,24 @@ def to_torch(a): computed_results["initial_world_orientations"] = orientations_t.clone() # -- set_world_poses ----------------------------------------------- + # ``.warp`` unwraps the ProxyArray returned by ``get_*_poses`` / + # ``get_*_scales`` to the underlying ``wp.array`` that ``wp.clone`` + # requires. ProxyArray was introduced in PR #5304 ("ProxyArray and + # Asset/Sensor level property caching") which changed the FrameView + # getter return type. Applies to every ``wp.clone`` call below. if is_newton: - new_positions = wp.clone(positions) + new_positions = wp.clone(positions.warp) wp.to_torch(new_positions)[:, 2] += 0.1 else: new_positions = positions_t.clone() new_positions[:, 2] += 0.1 - if is_newton: - torch.cuda.synchronize() + _sync() start_time = time.perf_counter() for _ in range(num_iterations): - xform_view.set_world_poses(new_positions, orientations) - if is_newton: - torch.cuda.synchronize() + with xform_view.xform_world_space_writer() as w: + w.set_poses(new_positions, orientations) + _sync() timing_results["set_world_poses"] = (time.perf_counter() - start_time) / num_iterations pa, oa = xform_view.get_world_poses() @@ -182,13 +200,11 @@ def to_torch(a): computed_results["world_orientations_after_set"] = to_torch(oa).clone() # -- get_local_poses ----------------------------------------------- - if is_newton: - torch.cuda.synchronize() + _sync() start_time = time.perf_counter() for _ in range(num_iterations): translations, orientations_local = xform_view.get_local_poses() - if is_newton: - torch.cuda.synchronize() + _sync() timing_results["get_local_poses"] = (time.perf_counter() - start_time) / num_iterations translations_t = to_torch(translations) @@ -198,45 +214,99 @@ def to_torch(a): # -- set_local_poses ----------------------------------------------- if is_newton: - new_translations = wp.clone(translations) + new_translations = wp.clone(translations.warp) wp.to_torch(new_translations)[:, 2] += 0.1 else: new_translations = translations_t.clone() new_translations[:, 2] += 0.1 - if is_newton: - torch.cuda.synchronize() + _sync() start_time = time.perf_counter() for _ in range(num_iterations): - xform_view.set_local_poses(new_translations, orientations_local) - if is_newton: - torch.cuda.synchronize() + with xform_view.xform_local_space_writer() as w: + w.set_poses(new_translations, orientations_local) + _sync() timing_results["set_local_poses"] = (time.perf_counter() - start_time) / num_iterations ta, ola = xform_view.get_local_poses() computed_results["local_translations_after_set"] = to_torch(ta).clone() computed_results["local_orientations_after_set"] = to_torch(ola).clone() - # -- get_both (world + local) -------------------------------------- + # -- get_world_scales ---------------------------------------------- + _sync() + start_time = time.perf_counter() + for _ in range(num_iterations): + world_scales = xform_view.get_world_scales() + _sync() + timing_results["get_world_scales"] = (time.perf_counter() - start_time) / num_iterations + + world_scales_t = to_torch(world_scales) + computed_results["initial_world_scales"] = world_scales_t.clone() + + # -- set_world_scales ---------------------------------------------- + if is_newton: + new_world_scales = wp.clone(world_scales.warp) + wp.to_torch(new_world_scales)[:] = 1.1 + else: + new_world_scales = world_scales_t.clone() + new_world_scales[:] = 1.1 + + _sync() + start_time = time.perf_counter() + for _ in range(num_iterations): + with xform_view.xform_world_space_writer() as w: + w.set_scales(new_world_scales) + _sync() + timing_results["set_world_scales"] = (time.perf_counter() - start_time) / num_iterations + + computed_results["world_scales_after_set"] = to_torch(xform_view.get_world_scales()).clone() + + # -- get_local_scales ---------------------------------------------- + _sync() + start_time = time.perf_counter() + for _ in range(num_iterations): + local_scales = xform_view.get_local_scales() + _sync() + timing_results["get_local_scales"] = (time.perf_counter() - start_time) / num_iterations + + local_scales_t = to_torch(local_scales) + computed_results["initial_local_scales"] = local_scales_t.clone() + + # -- set_local_scales ---------------------------------------------- if is_newton: - torch.cuda.synchronize() + new_local_scales = wp.clone(local_scales.warp) + wp.to_torch(new_local_scales)[:] = 0.9 + else: + new_local_scales = local_scales_t.clone() + new_local_scales[:] = 0.9 + + _sync() + start_time = time.perf_counter() + for _ in range(num_iterations): + with xform_view.xform_local_space_writer() as w: + w.set_scales(new_local_scales) + _sync() + timing_results["set_local_scales"] = (time.perf_counter() - start_time) / num_iterations + + computed_results["local_scales_after_set"] = to_torch(xform_view.get_local_scales()).clone() + + # -- get_both (world + local) -------------------------------------- + _sync() start_time = time.perf_counter() for _ in range(num_iterations): xform_view.get_world_poses() xform_view.get_local_poses() - if is_newton: - torch.cuda.synchronize() + _sync() timing_results["get_both"] = (time.perf_counter() - start_time) / num_iterations # -- interleaved set -> get ---------------------------------------- - if is_newton: - torch.cuda.synchronize() + _sync() start_time = time.perf_counter() for _ in range(num_iterations): - xform_view.set_world_poses(new_positions, orientations) + with xform_view.xform_world_space_writer() as w: + w.set_poses(new_positions, orientations) xform_view.get_world_poses() - if is_newton: - torch.cuda.synchronize() + _sync() timing_results["interleaved_world_set_get"] = (time.perf_counter() - start_time) / num_iterations finally: @@ -267,15 +337,26 @@ def print_results(results_dict: dict[str, dict[str, float]], num_prims: int, num print(header) print("-" * 120) - operations = [ - ("Initialization", "init"), + # ``init`` is the one-time view-construction cost. We display it in the + # per-operation table but EXCLUDE it from the steady-state totals and the + # overall speedup -- otherwise a backend whose construction is dominated + # by stage population (e.g. Newton, where the first call materializes the + # site cache) shows a misleading "0.00x" overall and crushes the rest of + # the table. The overall row is intended to compare per-iteration cost. + init_op = ("Initialization (one-time)", "init") + per_iter_operations = [ ("Get World Poses", "get_world_poses"), ("Set World Poses", "set_world_poses"), ("Get Local Poses", "get_local_poses"), ("Set Local Poses", "set_local_poses"), + ("Get World Scales", "get_world_scales"), + ("Set World Scales", "set_world_scales"), + ("Get Local Scales", "get_local_scales"), + ("Set Local Scales", "set_local_scales"), ("Get Both (World+Local)", "get_both"), ("Interleaved World Set->Get", "interleaved_world_set_get"), ] + operations = [init_op, *per_iter_operations] for op_name, op_key in operations: row = f"{op_name:<28}" @@ -286,15 +367,16 @@ def print_results(results_dict: dict[str, dict[str, float]], num_prims: int, num print("=" * 120) - total_row = f"{'Total':<28}" + total_row = f"{'Total (per-iter ops)':<28}" for name in api_names: - total_row += f" {sum(results_dict[name].values()) * 1000:>{col_width}.4f}" + per_iter_total = sum(results_dict[name].get(k, 0) for _, k in per_iter_operations) + total_row += f" {per_iter_total * 1000:>{col_width}.4f}" print(f"\n{total_row}") baseline = "isaaclab-usd" if baseline in results_dict and len(api_names) > 1: print("\n" + "=" * 120) - print(f"SPEEDUP vs {baseline.replace('-', ' ').title()}") + print(f"SPEEDUP vs {baseline.replace('-', ' ').title()} (per-iter ops; one-time init excluded)") print("=" * 120) header = f"{'Operation':<28}" for name in api_names: @@ -304,7 +386,7 @@ def print_results(results_dict: dict[str, dict[str, float]], num_prims: int, num print("-" * 120) base = results_dict[baseline] - for op_name, op_key in operations: + for op_name, op_key in per_iter_operations: row = f"{op_name:<28}" base_t = base.get(op_key, 0) for name in api_names: @@ -317,11 +399,11 @@ def print_results(results_dict: dict[str, dict[str, float]], num_prims: int, num print(row) print("=" * 120) - print(f"{'Overall':>28}", end="") - total_base = sum(base.values()) + print(f"{'Overall (per-iter ops)':>28}", end="") + total_base = sum(base.get(k, 0) for _, k in per_iter_operations) for name in api_names: if name != baseline: - total_impl = sum(results_dict[name].values()) + total_impl = sum(results_dict[name].get(k, 0) for _, k in per_iter_operations) if total_base > 0 and total_impl > 0: print(f" {total_base / total_impl:>{col_width}.2f}x", end="") else: diff --git a/source/isaaclab/changelog.d/fabric-local-poses.rst b/source/isaaclab/changelog.d/fabric-local-poses.rst new file mode 100644 index 000000000000..37a4f4ea6398 --- /dev/null +++ b/source/isaaclab/changelog.d/fabric-local-poses.rst @@ -0,0 +1,31 @@ +Added +^^^^^ + +* Added explicit local/world scale getters + :meth:`~isaaclab.sim.views.BaseFrameView.get_local_scales` and + :meth:`~isaaclab.sim.views.BaseFrameView.get_world_scales` to the FrameView + API, implemented for :class:`~isaaclab.sim.views.UsdFrameView`. Scale + writes go through the writer scope (see the ``xform-space-writer`` + fragment). + +* Added :func:`~isaaclab.utils.warp.fabric.decompose_indexed_fabric_transforms`, + :func:`~isaaclab.utils.warp.fabric.compose_indexed_fabric_transforms`, + :func:`~isaaclab.utils.warp.fabric.update_indexed_local_matrix_from_world`, and + :func:`~isaaclab.utils.warp.fabric.update_indexed_world_matrix_from_local` + Warp kernels operating on :class:`wp.indexedfabricarray` for reading and + writing Fabric ``Matrix4d`` attributes (``omni:fabric:worldMatrix`` / + ``omni:fabric:localMatrix``). + +Notes +^^^^^ + +* :meth:`~isaaclab.sim.views.BaseFrameView.get_scales` and + :meth:`~isaaclab.sim.views.BaseFrameView.set_scales` remain supported as + convenience helpers (not deprecated). For reads where the space matters, + prefer the explicit ``get_local_scales`` (operates on ``xformOp:scale``) or + ``get_world_scales`` (composed world-space scale). For writes that also + update poses, prefer batching inside one scope: + ``with view.xform_world_space_writer() as w: w.set_poses(...); w.set_scales(...)`` + (or ``xform_local_space_writer``). + :class:`~isaaclab.sim.views.UsdFrameView` preserves prior behavior by + defaulting :meth:`get_scales` / :meth:`set_scales` to local scales. diff --git a/source/isaaclab/changelog.d/xform-space-writer.rst b/source/isaaclab/changelog.d/xform-space-writer.rst new file mode 100644 index 000000000000..f3b2986beb6b --- /dev/null +++ b/source/isaaclab/changelog.d/xform-space-writer.rst @@ -0,0 +1,35 @@ +Added +^^^^^ + +* Added :class:`~isaaclab.sim.views.FrameViewSpaceWriterBase`, the new context-managed + write API for ``FrameView``-managed prim transforms. Open with + ``view.xform_world_space_writer()`` or ``view.xform_local_space_writer()`` and call + :meth:`~isaaclab.sim.views.FrameViewSpaceWriterBase.set_poses` / + :meth:`~isaaclab.sim.views.FrameViewSpaceWriterBase.set_scales` inside the scope; + the writer's ``__exit__`` derives the opposite-space matrices once and + synchronizes once. Only one writer scope may be active per view at a + time. View-level getters + (:meth:`~isaaclab.sim.views.BaseFrameView.get_world_poses` etc.) raise + :class:`RuntimeError` while a writer scope is active. + +* Added the two concrete tag classes + :class:`~isaaclab.sim.views.FrameViewWorldSpaceWriter` and + :class:`~isaaclab.sim.views.FrameViewLocalSpaceWriter` returned by + :meth:`~isaaclab.sim.views.BaseFrameView.xform_world_space_writer` / + :meth:`~isaaclab.sim.views.BaseFrameView.xform_local_space_writer`. + +Notes +^^^^^ + +* :meth:`~isaaclab.sim.views.BaseFrameView.set_world_poses`, + :meth:`~isaaclab.sim.views.BaseFrameView.set_local_poses`, and + :meth:`~isaaclab.sim.views.BaseFrameView.set_scales` remain supported as + convenience helpers (not deprecated). Each opens a single-statement writer + scope internally, so updating poses and scales through separate calls derives + the opposite-space matrices and synchronizes twice. For best performance, + update both inside one scope: + ``with view.xform_world_space_writer() as w: w.set_poses(...); w.set_scales(...)`` + (or :meth:`~isaaclab.sim.views.BaseFrameView.xform_local_space_writer`). + The bundled examples use the writer scope to get this benefit; callers may + keep the convenience helpers when code simplicity matters more than shaving + a redundant derive/sync. diff --git a/source/isaaclab/isaaclab/sensors/camera/camera.py b/source/isaaclab/isaaclab/sensors/camera/camera.py index d13a560600a5..881be4c7c9bf 100644 --- a/source/isaaclab/isaaclab/sensors/camera/camera.py +++ b/source/isaaclab/isaaclab/sensors/camera/camera.py @@ -374,7 +374,8 @@ def set_world_poses( orientations = convert_camera_frame_orientation_convention(orientations, origin=convention, target="opengl") ori_wp = wp.from_torch(orientations.contiguous(), dtype=wp.vec4f) idx_wp = self._resolve_env_ids_wp(env_ids) - self._view.set_world_poses(pos_wp, ori_wp, idx_wp) + with self._view.xform_world_space_writer() as writer: + writer.set_poses(pos_wp, ori_wp, idx_wp) def set_world_poses_from_view( self, eyes: torch.Tensor, targets: torch.Tensor, env_ids: Sequence[int] | None = None @@ -432,11 +433,12 @@ def set_world_poses_from_view( env_ids_torch = env_ids_torch.index_select(0, valid_indices) orientations = quat_from_matrix(rotation_matrix) idx_wp = wp.from_torch(env_ids_torch.contiguous(), dtype=wp.int32) - self._view.set_world_poses( - wp.from_torch(eyes.contiguous(), dtype=wp.vec3f), - wp.from_torch(orientations.contiguous(), dtype=wp.vec4f), - idx_wp, - ) + with self._view.xform_world_space_writer() as writer: + writer.set_poses( + wp.from_torch(eyes.contiguous(), dtype=wp.vec3f), + wp.from_torch(orientations.contiguous(), dtype=wp.vec4f), + idx_wp, + ) """ Operations diff --git a/source/isaaclab/isaaclab/sim/views/__init__.pyi b/source/isaaclab/isaaclab/sim/views/__init__.pyi index d578f85d6ada..734e925c19bb 100644 --- a/source/isaaclab/isaaclab/sim/views/__init__.pyi +++ b/source/isaaclab/isaaclab/sim/views/__init__.pyi @@ -7,6 +7,9 @@ __all__ = [ "BaseFrameView", "UsdFrameView", "FrameView", + "FrameViewSpaceWriterBase", + "FrameViewWorldSpaceWriter", + "FrameViewLocalSpaceWriter", # Deprecated alias "XformPrimView", ] @@ -14,5 +17,6 @@ __all__ = [ from .base_frame_view import BaseFrameView from .usd_frame_view import UsdFrameView from .frame_view import FrameView +from .xform_space_writer import FrameViewSpaceWriterBase, FrameViewWorldSpaceWriter, FrameViewLocalSpaceWriter # Deprecated alias from .xform_prim_view import XformPrimView diff --git a/source/isaaclab/isaaclab/sim/views/base_frame_view.py b/source/isaaclab/isaaclab/sim/views/base_frame_view.py index d58cb0fbb534..79c672b7cfd0 100644 --- a/source/isaaclab/isaaclab/sim/views/base_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/base_frame_view.py @@ -8,23 +8,48 @@ from __future__ import annotations import abc +from typing import TYPE_CHECKING import warp as wp from isaaclab.utils.warp import ProxyArray +if TYPE_CHECKING: + from .xform_space_writer import FrameViewLocalSpaceWriter, FrameViewSpaceWriterBase, FrameViewWorldSpaceWriter + class BaseFrameView(abc.ABC): - """Abstract interface for reading and writing world-space transforms of multiple prims. + """Abstract interface for reading and writing transforms of multiple prims. Backend-specific implementations (USD/Fabric, Newton GPU state, etc.) subclass this to provide efficient batched pose queries. The factory :class:`~isaaclab.sim.views.FrameView` selects the correct implementation at runtime based on the active physics backend. - All pose getters return :class:`~isaaclab.utils.warp.ProxyArray`. Setters accept ``wp.array``. + All getters return :class:`~isaaclab.utils.warp.ProxyArray`. All writes go + through the writer-scope API -- :meth:`xform_world_space_writer` or + :meth:`xform_local_space_writer`: + + .. code-block:: python + + with view.xform_world_space_writer() as writer: + writer.set_poses(positions=p, orientations=o) + writer.set_scales(scales=s) + # Derived-space matrices are recomputed and the writer scope is closed. + + Only one writer scope may be active per view at a time. While a writer + scope is active, the view-level getters + (:meth:`get_world_poses`, :meth:`get_local_poses`, + :meth:`get_world_scales`, :meth:`get_local_scales`) raise + :class:`RuntimeError` -- use the writer's :meth:`~FrameViewSpaceWriterBase.get_poses` + or :meth:`~FrameViewSpaceWriterBase.get_scales` inside the scope, or exit the + scope first. """ + # Class-level default; instance-level value is set by the writer's + # __enter__ / __exit__ to track the active scope on this view. + _active_writer: FrameViewSpaceWriterBase | None = None + @property @abc.abstractmethod def count(self) -> int: @@ -37,7 +62,77 @@ def device(self) -> str: """Device where arrays are allocated (``"cpu"`` or ``"cuda:0"``).""" ... + # ------------------------------------------------------------------ + # Write scope -- recommended API for all transform writes. + # ------------------------------------------------------------------ + + def xform_world_space_writer(self) -> FrameViewWorldSpaceWriter: + """Open a world-space write scope on this view (recommended write API). + + Inside the scope, :meth:`~FrameViewSpaceWriterBase.set_poses` / + :meth:`~FrameViewSpaceWriterBase.set_scales` write world-space values. + + Returns: + A :class:`~isaaclab.sim.views.FrameViewWorldSpaceWriter` context manager. + + Raises: + RuntimeError: On ``__enter__``, if another writer is already active + on this view. + + Example: + .. code-block:: python + + with view.xform_world_space_writer() as w: + w.set_poses(positions=p, orientations=o) + w.set_scales(scales=s) + """ + return self._make_world_space_writer() + + def xform_local_space_writer(self) -> FrameViewLocalSpaceWriter: + """Open a local-space write scope on this view (recommended write API). + + Inside the scope, :meth:`~FrameViewSpaceWriterBase.set_poses` / + :meth:`~FrameViewSpaceWriterBase.set_scales` write local-space values. + + Returns: + A :class:`~isaaclab.sim.views.FrameViewLocalSpaceWriter` context manager. + + Raises: + RuntimeError: On ``__enter__``, if another writer is already active + on this view. + + Example: + .. code-block:: python + + with view.xform_local_space_writer() as w: + w.set_poses(translations=t, orientations=o) + w.set_scales(scales=s) + """ + return self._make_local_space_writer() + + @abc.abstractmethod + def _make_world_space_writer(self) -> FrameViewWorldSpaceWriter: + """Backend hook: return a fresh :class:`FrameViewWorldSpaceWriter` for this view.""" + ... + @abc.abstractmethod + def _make_local_space_writer(self) -> FrameViewLocalSpaceWriter: + """Backend hook: return a fresh :class:`FrameViewLocalSpaceWriter` for this view.""" + ... + + def _assert_no_active_writer(self, method_name: str) -> None: + """Raise :class:`RuntimeError` if a writer scope is currently active on this view.""" + if self._active_writer is not None: + raise RuntimeError( + f"{type(self).__name__}.{method_name}() is not allowed while a writer " + f"scope is active ({type(self._active_writer).__name__}). Use the writer's " + f"get_poses / get_scales inside the scope, or exit the scope first." + ) + + # ------------------------------------------------------------------ + # Public getters -- guarded; delegate to backend ``_*_impl`` hooks. + # ------------------------------------------------------------------ + def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: """Get world-space positions and orientations for prims in the view. @@ -46,12 +141,110 @@ def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, Returns: A tuple ``(positions, orientations)`` of :class:`~isaaclab.utils.warp.ProxyArray` - wrappers. Use ``.warp`` for the underlying ``wp.array`` or ``.torch`` for a - cached zero-copy ``torch.Tensor`` view. + wrappers. + + Raises: + RuntimeError: If a writer scope is active on this view. + """ + self._assert_no_active_writer("get_world_poses") + return self._get_world_poses_impl(indices) + + def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: + """Get local-space translations and orientations for prims in the view. + + Args: + indices: Subset of prims to query. ``None`` means all prims. + + Returns: + A tuple ``(translations, orientations)`` of :class:`~isaaclab.utils.warp.ProxyArray` + wrappers. + + Raises: + RuntimeError: If a writer scope is active on this view. + """ + self._assert_no_active_writer("get_local_poses") + return self._get_local_poses_impl(indices) + + def get_local_scales(self, indices: wp.array | None = None) -> ProxyArray: + """Get local-space scales for prims in the view. + + Args: + indices: Subset of prims to query. ``None`` means all prims. + + Returns: + A :class:`~isaaclab.utils.warp.ProxyArray` of shape ``(M, 3)``. + + Raises: + RuntimeError: If a writer scope is active on this view. + """ + self._assert_no_active_writer("get_local_scales") + return self._get_local_scales_impl(indices) + + def get_world_scales(self, indices: wp.array | None = None) -> ProxyArray: + """Get world-space (composed) scales for prims in the view. + + Returns the effective scale in world space (``parent_scale * local_scale``). + + .. note:: + Scale extraction uses TRS (Translation-Rotation-Scale) decomposition, + which assumes no shear/skew in the transform matrix. If a prim's + world transform contains shear, the extracted scale values will be + approximate. + + Args: + indices: Subset of prims to query. ``None`` means all prims. + + Returns: + A :class:`~isaaclab.utils.warp.ProxyArray` of shape ``(M, 3)``. + + Raises: + RuntimeError: If a writer scope is active on this view. """ + self._assert_no_active_writer("get_world_scales") + return self._get_world_scales_impl(indices) + + # ------------------------------------------------------------------ + # Backend hooks for the public getters above. + # ------------------------------------------------------------------ + + @abc.abstractmethod + def _get_world_poses_impl(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: + """Backend implementation of :meth:`get_world_poses`.""" ... @abc.abstractmethod + def _get_local_poses_impl(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: + """Backend implementation of :meth:`get_local_poses`.""" + ... + + @abc.abstractmethod + def _get_local_scales_impl(self, indices: wp.array | None = None) -> ProxyArray: + """Backend implementation of :meth:`get_local_scales`.""" + ... + + @abc.abstractmethod + def _get_world_scales_impl(self, indices: wp.array | None = None) -> ProxyArray: + """Backend implementation of :meth:`get_world_scales`.""" + ... + + # ------------------------------------------------------------------ + # Convenience pose/scale setters -- route through the writer scope. + # + # These are kept as first-class convenience APIs (not deprecated). Each + # call opens its own single-statement writer scope internally, so writing + # poses and scales through two separate calls performs the opposite-space + # recompute + synchronize twice. For the best performance when updating + # both poses and scales together, open one writer scope and issue both + # writes inside it:: + # + # with view.xform_world_space_writer() as w: + # w.set_poses(...) + # w.set_scales(...) + # + # Prefer the writer scope in hot loops; prefer these helpers when code + # simplicity matters more than shaving a redundant derive/sync. + # ------------------------------------------------------------------ + def set_world_poses( self, positions: wp.array | None = None, @@ -60,28 +253,19 @@ def set_world_poses( ) -> None: """Set world-space positions and/or orientations for prims in the view. + This convenience method opens a single-statement writer scope + internally. To update poses and scales together without paying the + opposite-space derive/sync twice, prefer + ``with view.xform_world_space_writer() as w: w.set_poses(...); w.set_scales(...)``. + Args: positions: World-space positions ``(M, 3)``. ``None`` leaves positions unchanged. orientations: World-space quaternions ``(M, 4)``. ``None`` leaves orientations unchanged. indices: Subset of prims to update. ``None`` means all prims. """ - ... + with self.xform_world_space_writer() as writer: + writer.set_poses(positions, orientations, indices) - @abc.abstractmethod - def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: - """Get local-space positions and orientations for prims in the view. - - Args: - indices: Subset of prims to query. ``None`` means all prims. - - Returns: - A tuple ``(translations, orientations)`` of :class:`~isaaclab.utils.warp.ProxyArray` - wrappers. Use ``.warp`` for the underlying ``wp.array`` or ``.torch`` for a - cached zero-copy ``torch.Tensor`` view. - """ - ... - - @abc.abstractmethod def set_local_poses( self, translations: wp.array | None = None, @@ -90,31 +274,66 @@ def set_local_poses( ) -> None: """Set local-space translations and/or orientations for prims in the view. + This convenience method opens a single-statement writer scope + internally. To update poses and scales together without paying the + opposite-space derive/sync twice, prefer + ``with view.xform_local_space_writer() as w: w.set_poses(...); w.set_scales(...)``. + Args: translations: Local-space translations ``(M, 3)``. ``None`` leaves translations unchanged. orientations: Local-space quaternions ``(M, 4)``. ``None`` leaves orientations unchanged. indices: Subset of prims to update. ``None`` means all prims. """ - ... + with self.xform_local_space_writer() as writer: + writer.set_poses(translations, orientations, indices) - @abc.abstractmethod - def get_scales(self, indices: wp.array | None = None) -> wp.array: + # ------------------------------------------------------------------ + # Scale getter/setter convenience helpers. + # ------------------------------------------------------------------ + + def get_scales(self, indices: wp.array | None = None) -> ProxyArray: """Get scales for prims in the view. + .. note:: + Prefer the explicit :meth:`get_local_scales` or + :meth:`get_world_scales` when the space matters. This method + delegates to :meth:`_get_scales_impl`, which preserves each + backend's legacy space (world for Fabric, local for USD). + Args: indices: Subset of prims to query. ``None`` means all prims. Returns: - A ``wp.array`` of shape ``(M, 3)``. + A ``ProxyArray`` of shape ``(M, 3)``. + + Raises: + RuntimeError: If a writer scope is active on this view. """ - ... + self._assert_no_active_writer("get_scales") + return self._get_scales_impl(indices) - @abc.abstractmethod def set_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: """Set scales for prims in the view. + This convenience method delegates to :meth:`_set_scales_impl`, which + opens the backend's legacy space (world for Fabric, local for USD) and + calls ``writer.set_scales``. To update poses and scales together + without paying the opposite-space derive/sync twice, prefer + ``with view.xform_world_space_writer() as w: w.set_poses(...); w.set_scales(...)`` + (or :meth:`xform_local_space_writer`). + Args: scales: Scales ``(M, 3)`` as ``wp.array``. indices: Subset of prims to update. ``None`` means all prims. """ + self._set_scales_impl(scales, indices) + + @abc.abstractmethod + def _get_scales_impl(self, indices: wp.array | None = None) -> ProxyArray: + """Backend-specific implementation for :meth:`get_scales`.""" + ... + + @abc.abstractmethod + def _set_scales_impl(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Backend-specific implementation for :meth:`set_scales`.""" ... diff --git a/source/isaaclab/isaaclab/sim/views/usd_frame_view.py b/source/isaaclab/isaaclab/sim/views/usd_frame_view.py index 88392d54b2a0..22707ceb8a73 100644 --- a/source/isaaclab/isaaclab/sim/views/usd_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/usd_frame_view.py @@ -17,6 +17,7 @@ from isaaclab.utils.warp import ProxyArray from .base_frame_view import BaseFrameView +from .xform_space_writer import FrameViewLocalSpaceWriter, FrameViewWorldSpaceWriter logger = logging.getLogger(__name__) @@ -35,7 +36,14 @@ class UsdFrameView(BaseFrameView): For GPU-accelerated Fabric operations, use the PhysX backend variant obtained via :class:`~isaaclab.sim.views.FrameView`. - Pose getters return :class:`~isaaclab.utils.warp.ProxyArray`. Setters accept ``wp.array``. + All writes go through the writer-scope API (:meth:`xform_world_space_writer` + / :meth:`xform_local_space_writer`). The + USD backend's writers are pass-throughs: each :meth:`set_poses` / + :meth:`set_scales` call directly modifies the prim's USD ``xformOp:*`` + attributes (no batching, no derivation step on exit) -- USD has no + separate world-matrix storage to keep in sync. + + Getters return :class:`~isaaclab.utils.warp.ProxyArray`. .. note:: **Transform Requirements:** @@ -126,24 +134,70 @@ def prim_paths(self) -> list[str]: return self._prim_paths # ------------------------------------------------------------------ - # Setters + # Writer factory hooks (pass-through writers; USD has no derived state) + # ------------------------------------------------------------------ + + def _make_world_space_writer(self) -> FrameViewWorldSpaceWriter: + return _UsdWorldSpaceWriter(self) + + def _make_local_space_writer(self) -> FrameViewLocalSpaceWriter: + return _UsdLocalSpaceWriter(self) + + # ------------------------------------------------------------------ + # Visibility (USD-only, no writer scope) # ------------------------------------------------------------------ - def set_world_poses( + def set_visibility(self, visibility: torch.Tensor, indices: wp.array | None = None): + """Set visibility for prims in the view. + + Args: + visibility: Visibility as a boolean tensor of shape ``(M,)``. + indices: Indices of prims to set visibility for. Defaults to None (all prims). + """ + indices_list = self._resolve_indices(indices) + + if visibility.shape != (len(indices_list),): + raise ValueError(f"Expected visibility shape ({len(indices_list)},), got {visibility.shape}.") + + with Sdf.ChangeBlock(): + for idx, prim_idx in enumerate(indices_list): + imageable = UsdGeom.Imageable(self._prims[prim_idx]) + if visibility[idx]: + imageable.MakeVisible() + else: + imageable.MakeInvisible() + + def get_visibility(self, indices: wp.array | None = None) -> torch.Tensor: + """Get visibility for prims in the view. + + Args: + indices: Indices of prims to get visibility for. Defaults to None (all prims). + + Returns: + A tensor of shape ``(M,)`` containing the visibility of each prim (bool). + """ + indices_list = self._resolve_indices(indices) + + visibility = torch.zeros(len(indices_list), dtype=torch.bool, device=self._device) + for idx, prim_idx in enumerate(indices_list): + imageable = UsdGeom.Imageable(self._prims[prim_idx]) + visibility[idx] = imageable.ComputeVisibility() != UsdGeom.Tokens.invisible + return visibility + + # ------------------------------------------------------------------ + # Backend hooks: pose / scale writes (called by writers). + # ------------------------------------------------------------------ + + def _apply_world_pose_write( self, positions: wp.array | None = None, orientations: wp.array | None = None, indices: wp.array | None = None, - ): - """Set world-space poses for prims in the view. + ) -> None: + """Apply a world-space pose write directly to USD xform ops. Converts the desired world pose to local-space relative to each prim's - parent before writing to USD xform ops. - - Args: - positions: World-space positions of shape ``(M, 3)``. - orientations: World-space quaternions ``(w, x, y, z)`` of shape ``(M, 4)``. - indices: Indices of prims to set poses for. Defaults to None (all prims). + parent before writing. """ indices_list = self._resolve_indices(indices) @@ -187,19 +241,13 @@ def set_world_poses( if local_quat is not None: prim.GetAttribute("xformOp:orient").Set(local_quat) - def set_local_poses( + def _apply_local_pose_write( self, translations: wp.array | None = None, orientations: wp.array | None = None, indices: wp.array | None = None, - ): - """Set local-space poses for prims in the view. - - Args: - translations: Local-space translations of shape ``(M, 3)``. - orientations: Local-space quaternions ``(w, x, y, z)`` of shape ``(M, 4)``. - indices: Indices of prims to set poses for. Defaults to None (all prims). - """ + ) -> None: + """Apply a local-space pose write directly to USD xform ops.""" indices_list = self._resolve_indices(indices) translations_array = Vt.Vec3dArray.FromNumpy(self._to_numpy(translations)) if translations is not None else None @@ -213,13 +261,8 @@ def set_local_poses( if orientations_array is not None: prim.GetAttribute("xformOp:orient").Set(orientations_array[idx]) - def set_scales(self, scales: wp.array, indices: wp.array | None = None): - """Set scales for prims in the view. - - Args: - scales: Scales of shape ``(M, 3)``. - indices: Indices of prims to set scales for. Defaults to None (all prims). - """ + def _apply_local_scale_write(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Apply a local-space scale write (``xformOp:scale``).""" indices_list = self._resolve_indices(indices) scales_array = Vt.Vec3dArray.FromNumpy(self._to_numpy(scales)) @@ -228,41 +271,41 @@ def set_scales(self, scales: wp.array, indices: wp.array | None = None): prim = self._prims[prim_idx] prim.GetAttribute("xformOp:scale").Set(scales_array[idx]) - def set_visibility(self, visibility: torch.Tensor, indices: wp.array | None = None): - """Set visibility for prims in the view. + def _apply_world_scale_write(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Apply a world-space scale write. - Args: - visibility: Visibility as a boolean tensor of shape ``(M,)``. - indices: Indices of prims to set visibility for. Defaults to None (all prims). + Computes ``local_scale = world_scale / parent_world_scale`` and writes + to ``xformOp:scale``. """ indices_list = self._resolve_indices(indices) - - if visibility.shape != (len(indices_list),): - raise ValueError(f"Expected visibility shape ({len(indices_list)},), got {visibility.shape}.") + scales_np = self._to_numpy(scales) + xf_cache = UsdGeom.XformCache(Usd.TimeCode.Default()) with Sdf.ChangeBlock(): for idx, prim_idx in enumerate(indices_list): - imageable = UsdGeom.Imageable(self._prims[prim_idx]) - if visibility[idx]: - imageable.MakeVisible() + prim = self._prims[prim_idx] + parent = prim.GetParent() + if parent and parent.IsValid() and parent.GetPath() != Sdf.Path.absoluteRootPath: + parent_world = xf_cache.GetLocalToWorldTransform(parent) + parent_scale = Gf.Vec3d( + Gf.Vec3d(parent_world[0][0], parent_world[0][1], parent_world[0][2]).GetLength(), + Gf.Vec3d(parent_world[1][0], parent_world[1][1], parent_world[1][2]).GetLength(), + Gf.Vec3d(parent_world[2][0], parent_world[2][1], parent_world[2][2]).GetLength(), + ) else: - imageable.MakeInvisible() + parent_scale = Gf.Vec3d(1.0, 1.0, 1.0) + local_scale = Gf.Vec3d( + float(scales_np[idx][0] / parent_scale[0]), + float(scales_np[idx][1] / parent_scale[1]), + float(scales_np[idx][2] / parent_scale[2]), + ) + prim.GetAttribute("xformOp:scale").Set(local_scale) # ------------------------------------------------------------------ - # Getters + # Backend hooks: pose / scale reads. # ------------------------------------------------------------------ - def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: - """Get world-space poses for prims in the view. - - Args: - indices: Indices of prims to get poses for. Defaults to None (all prims). - - Returns: - A tuple ``(positions, orientations)`` of :class:`~isaaclab.utils.warp.ProxyArray` - wrappers. Use ``.warp`` for the underlying ``wp.array`` or ``.torch`` for a - cached zero-copy ``torch.Tensor`` view. - """ + def _get_world_poses_impl(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: indices_list = self._resolve_indices(indices) positions = Vt.Vec3dArray(len(indices_list)) @@ -280,17 +323,7 @@ def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, quat_wp = wp.array(np.array(orientations, dtype=np.float32), dtype=wp.float32, device=self._device) return ProxyArray(pos_wp), ProxyArray(quat_wp) - def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: - """Get local-space poses for prims in the view. - - Args: - indices: Indices of prims to get poses for. Defaults to None (all prims). - - Returns: - A tuple ``(translations, orientations)`` of :class:`~isaaclab.utils.warp.ProxyArray` - wrappers. Use ``.warp`` for the underlying ``wp.array`` or ``.torch`` for a - cached zero-copy ``torch.Tensor`` view. - """ + def _get_local_poses_impl(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: indices_list = self._resolve_indices(indices) translations = Vt.Vec3dArray(len(indices_list)) @@ -308,15 +341,7 @@ def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, quat_wp = wp.array(np.array(orientations, dtype=np.float32), dtype=wp.float32, device=self._device) return ProxyArray(pos_wp), ProxyArray(quat_wp) - def get_scales(self, indices: wp.array | None = None) -> wp.array: - """Get scales for prims in the view. - - Args: - indices: Indices of prims to get scales for. Defaults to None (all prims). - - Returns: - A ``wp.array`` of shape ``(M, 3)``. - """ + def _get_local_scales_impl(self, indices: wp.array | None = None) -> ProxyArray: indices_list = self._resolve_indices(indices) scales = Vt.Vec3dArray(len(indices_list)) @@ -324,24 +349,34 @@ def get_scales(self, indices: wp.array | None = None) -> wp.array: prim = self._prims[prim_idx] scales[idx] = prim.GetAttribute("xformOp:scale").Get() - return wp.array(np.array(scales, dtype=np.float32), dtype=wp.float32, device=self._device) - - def get_visibility(self, indices: wp.array | None = None) -> torch.Tensor: - """Get visibility for prims in the view. - - Args: - indices: Indices of prims to get visibility for. Defaults to None (all prims). + return ProxyArray(wp.array(np.array(scales, dtype=np.float32), dtype=wp.float32, device=self._device)) - Returns: - A tensor of shape ``(M,)`` containing the visibility of each prim (bool). - """ + def _get_world_scales_impl(self, indices: wp.array | None = None) -> ProxyArray: indices_list = self._resolve_indices(indices) + xf_cache = UsdGeom.XformCache(Usd.TimeCode.Default()) - visibility = torch.zeros(len(indices_list), dtype=torch.bool, device=self._device) + scales = np.empty((len(indices_list), 3), dtype=np.float32) for idx, prim_idx in enumerate(indices_list): - imageable = UsdGeom.Imageable(self._prims[prim_idx]) - visibility[idx] = imageable.ComputeVisibility() != UsdGeom.Tokens.invisible - return visibility + prim = self._prims[prim_idx] + world_mtx = xf_cache.GetLocalToWorldTransform(prim) + scales[idx, 0] = Gf.Vec3d(world_mtx[0][0], world_mtx[0][1], world_mtx[0][2]).GetLength() + scales[idx, 1] = Gf.Vec3d(world_mtx[1][0], world_mtx[1][1], world_mtx[1][2]).GetLength() + scales[idx, 2] = Gf.Vec3d(world_mtx[2][0], world_mtx[2][1], world_mtx[2][2]).GetLength() + + return ProxyArray(wp.array(scales, dtype=wp.float32, device=self._device)) + + # ------------------------------------------------------------------ + # Deprecated get_scales / set_scales hooks + # ------------------------------------------------------------------ + + def _get_scales_impl(self, indices: wp.array | None = None) -> ProxyArray: + """USD legacy: get_scales returns local scales.""" + return self._get_local_scales_impl(indices) + + def _set_scales_impl(self, scales: wp.array, indices: wp.array | None = None) -> None: + """USD legacy: set_scales writes local scales via a one-shot writer scope.""" + with self.xform_local_space_writer() as writer: + writer.set_scales(scales, indices) # ------------------------------------------------------------------ # Helpers @@ -359,3 +394,44 @@ def _to_numpy(data: wp.array | torch.Tensor) -> np.ndarray: if isinstance(data, wp.array): return data.numpy() return data.cpu().numpy() + + +# ---------------------------------------------------------------------- +# Pass-through writer classes +# ---------------------------------------------------------------------- + + +class _UsdWorldSpaceWriter(FrameViewWorldSpaceWriter): + """USD world-space writer: pass-through to backend ``_apply_*`` hooks. + + USD has no separate world-matrix storage to keep in sync; ``__exit__`` + is a no-op beyond releasing the single-writer lock. + """ + + def set_poses(self, positions=None, orientations=None, indices=None) -> None: + self._view._apply_world_pose_write(positions, orientations, indices) # type: ignore[attr-defined] + + def set_scales(self, scales, indices=None) -> None: + self._view._apply_world_scale_write(scales, indices) # type: ignore[attr-defined] + + def get_poses(self, indices=None) -> tuple[ProxyArray, ProxyArray]: + return self._view._get_world_poses_impl(indices) # type: ignore[attr-defined] + + def get_scales(self, indices=None) -> ProxyArray: + return self._view._get_world_scales_impl(indices) # type: ignore[attr-defined] + + +class _UsdLocalSpaceWriter(FrameViewLocalSpaceWriter): + """USD local-space writer: pass-through to backend ``_apply_*`` hooks.""" + + def set_poses(self, positions=None, orientations=None, indices=None) -> None: + self._view._apply_local_pose_write(positions, orientations, indices) # type: ignore[attr-defined] + + def set_scales(self, scales, indices=None) -> None: + self._view._apply_local_scale_write(scales, indices) # type: ignore[attr-defined] + + def get_poses(self, indices=None) -> tuple[ProxyArray, ProxyArray]: + return self._view._get_local_poses_impl(indices) # type: ignore[attr-defined] + + def get_scales(self, indices=None) -> ProxyArray: + return self._view._get_local_scales_impl(indices) # type: ignore[attr-defined] diff --git a/source/isaaclab/isaaclab/sim/views/xform_space_writer.py b/source/isaaclab/isaaclab/sim/views/xform_space_writer.py new file mode 100644 index 000000000000..2d657b031995 --- /dev/null +++ b/source/isaaclab/isaaclab/sim/views/xform_space_writer.py @@ -0,0 +1,145 @@ +# 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 + +"""Context-managed transform writers for :class:`~isaaclab.sim.views.BaseFrameView`. + +This module defines the recommended write API for FrameView poses and scales: + +.. code-block:: python + + with view.xform_world_space_writer() as writer: + writer.set_poses(positions=p, orientations=o) + writer.set_scales(scales=s) + # ... any number of writes ... + # On exit the writer derives the opposite-space matrices once, + # synchronizes once, and restores any saved Fabric tracking state. + +Only one writer may be active per view at a time. While a writer scope is +active on a view, view-level getters (``view.get_world_poses``, +``view.get_local_poses``, ``view.get_world_scales``, +``view.get_local_scales``) raise :class:`RuntimeError` -- use the writer's own +:meth:`~FrameViewSpaceWriterBase.get_poses` / :meth:`~FrameViewSpaceWriterBase.get_scales` +inside the scope, or exit the scope first. + +**Do not advance the simulation or render from inside a scope.** The scope +runs as synchronous Python code, so no ``sim.step()`` / ``world.render()`` / +``SimulationApp.update()`` is allowed inside the ``with`` block. Until the +scope exits, the backend's matrices may be mid-write (some prims updated, +others not; the opposite-space derive has not yet run) and rendering against +that state would read torn data. Keep scopes short and step the +simulation outside them. +""" + +from __future__ import annotations + +import abc +from typing import TYPE_CHECKING + +import warp as wp + +from isaaclab.utils.warp import ProxyArray + +if TYPE_CHECKING: + from .base_frame_view import BaseFrameView + + +class FrameViewSpaceWriterBase(abc.ABC): + """Abstract context-managed writer for a single transform space. + + Subclasses are returned by :meth:`BaseFrameView.xform_world_space_writer` / + :meth:`BaseFrameView.xform_local_space_writer`; they + are not constructed directly. The class is intentionally minimal -- the + pose/scale semantics depend on the writer's space (world or local), which + is conveyed by the concrete tag class :class:`FrameViewWorldSpaceWriter` or + :class:`FrameViewLocalSpaceWriter`. + + The scope runs as synchronous Python code: no simulation step and no + render tick can run while it is open, and the caller must not advance + either from inside the ``with`` block. See the module docstring for the + full contract. + """ + + def __init__(self, view: BaseFrameView): + self._view = view + + @abc.abstractmethod + def set_poses( + self, + positions: wp.array | None = None, + orientations: wp.array | None = None, + indices: wp.array | None = None, + ) -> None: + """Set positions and/or orientations in this writer's space. + + Args: + positions: Positions ``(M, 3)``. ``None`` leaves positions unchanged. + orientations: Quaternions ``(M, 4)`` in ``(x, y, z, w)``. + ``None`` leaves orientations unchanged. + indices: Subset of prims to update. ``None`` means all prims. + """ + ... + + @abc.abstractmethod + def set_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Set scales in this writer's space. + + Args: + scales: Scales ``(M, 3)`` as ``wp.array``. + indices: Subset of prims to update. ``None`` means all prims. + """ + ... + + @abc.abstractmethod + def get_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: + """Return ``(positions, orientations)`` in this writer's space. + + Reflects any in-scope writes that have already been queued on the + underlying device stream. + """ + ... + + @abc.abstractmethod + def get_scales(self, indices: wp.array | None = None) -> ProxyArray: + """Return scales in this writer's space.""" + ... + + def __enter__(self) -> FrameViewSpaceWriterBase: + if self._view._active_writer is not None: + raise RuntimeError( + f"{type(self._view).__name__} already has an active writer scope " + f"({type(self._view._active_writer).__name__}). Exit the existing scope before " + "opening a new one." + ) + self._view._active_writer = self + self._enter_impl() + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + try: + self._exit_impl(exc_type, exc_val, exc_tb) + finally: + self._view._active_writer = None + + def _enter_impl(self) -> None: + """Backend hook called after the single-active-writer lock is claimed.""" + + def _exit_impl(self, exc_type, exc_val, exc_tb) -> None: + """Backend hook called before the single-active-writer lock is released.""" + + +class FrameViewWorldSpaceWriter(FrameViewSpaceWriterBase): + """Writer whose :meth:`set_poses` / :meth:`set_scales` write world-space values. + + On context exit the opposite-space (``local``) matrices are derived from + the just-written world matrices in a single Warp kernel launch. + """ + + +class FrameViewLocalSpaceWriter(FrameViewSpaceWriterBase): + """Writer whose :meth:`set_poses` / :meth:`set_scales` write local-space values. + + On context exit the opposite-space (``world``) matrices are derived from + the just-written local matrices in a single Warp kernel launch. + """ diff --git a/source/isaaclab/isaaclab/utils/warp/fabric.py b/source/isaaclab/isaaclab/utils/warp/fabric.py index a48f773f4991..e0519d98c338 100644 --- a/source/isaaclab/isaaclab/utils/warp/fabric.py +++ b/source/isaaclab/isaaclab/utils/warp/fabric.py @@ -18,12 +18,14 @@ if TYPE_CHECKING: FabricArrayUInt32 = Any FabricArrayMat44d = Any + IndexedFabricArrayMat44d = Any ArrayUInt32 = Any ArrayUInt32_1d = Any ArrayFloat32_2d = Any else: FabricArrayUInt32 = wp.fabricarray(dtype=wp.uint32) FabricArrayMat44d = wp.fabricarray(dtype=wp.mat44d) + IndexedFabricArrayMat44d = wp.indexedfabricarray(dtype=wp.mat44d) ArrayUInt32 = wp.array(ndim=1, dtype=wp.uint32) ArrayUInt32_1d = wp.array(dtype=wp.uint32) ArrayFloat32_2d = wp.array(ndim=2, dtype=wp.float32) @@ -163,6 +165,180 @@ def compose_fabric_transformation_matrix_from_warp_arrays( ) +@wp.kernel(enable_backward=False) +def decompose_indexed_fabric_transforms( + fabric_matrices: IndexedFabricArrayMat44d, + array_positions: ArrayFloat32_2d, + array_orientations: ArrayFloat32_2d, + array_scales: ArrayFloat32_2d, + indices: ArrayUInt32, +): + """Decompose indexed Fabric transformation matrices into position, orientation, and scale. + + Like :func:`decompose_fabric_transformation_matrix_to_warp_arrays` but operates on a + :class:`wp.indexedfabricarray` that already encodes the view-to-fabric mapping, removing + the need for a separate ``mapping`` array. + + Args: + fabric_matrices: Indexed fabric array containing 4x4 transformation matrices. + array_positions: Output array for positions [m], shape (N, 3). + array_orientations: Output array for quaternions in xyzw format, shape (N, 4). + array_scales: Output array for scales, shape (N, 3). + indices: View indices to process (subset selection). + """ + output_index = wp.tid() + view_index = indices[output_index] + + position, rotation, scale = _decompose_transformation_matrix(wp.mat44f(fabric_matrices[view_index])) + + if array_positions.shape[0] > 0: + array_positions[output_index, 0] = position[0] + array_positions[output_index, 1] = position[1] + array_positions[output_index, 2] = position[2] + if array_orientations.shape[0] > 0: + array_orientations[output_index, 0] = rotation[0] + array_orientations[output_index, 1] = rotation[1] + array_orientations[output_index, 2] = rotation[2] + array_orientations[output_index, 3] = rotation[3] + if array_scales.shape[0] > 0: + array_scales[output_index, 0] = scale[0] + array_scales[output_index, 1] = scale[1] + array_scales[output_index, 2] = scale[2] + + +@wp.kernel(enable_backward=False) +def compose_indexed_fabric_transforms( + fabric_matrices: IndexedFabricArrayMat44d, + array_positions: ArrayFloat32_2d, + array_orientations: ArrayFloat32_2d, + array_scales: ArrayFloat32_2d, + broadcast_positions: bool, + broadcast_orientations: bool, + broadcast_scales: bool, + indices: ArrayUInt32, +): + """Compose indexed Fabric transformation matrices from position, orientation, and scale. + + Like :func:`compose_fabric_transformation_matrix_from_warp_arrays` but operates on a + :class:`wp.indexedfabricarray` that already encodes the view-to-fabric mapping, removing + the need for a separate ``mapping`` array. + + Args: + fabric_matrices: Indexed fabric array containing 4x4 transformation matrices to update. + array_positions: Input array for positions [m], shape (N, 3). + array_orientations: Input array for quaternions in xyzw format, shape (N, 4). + array_scales: Input array for scales, shape (N, 3). + broadcast_positions: If True, use first position for all prims. + broadcast_orientations: If True, use first orientation for all prims. + broadcast_scales: If True, use first scale for all prims. + indices: View indices to process (subset selection). + """ + i = wp.tid() + view_index = indices[i] + position, rotation, scale = _decompose_transformation_matrix(wp.mat44f(fabric_matrices[view_index])) + + if array_positions.shape[0] > 0: + if broadcast_positions: + index = 0 + else: + index = i + position[0] = array_positions[index, 0] + position[1] = array_positions[index, 1] + position[2] = array_positions[index, 2] + if array_orientations.shape[0] > 0: + if broadcast_orientations: + index = 0 + else: + index = i + rotation[0] = array_orientations[index, 0] + rotation[1] = array_orientations[index, 1] + rotation[2] = array_orientations[index, 2] + rotation[3] = array_orientations[index, 3] + if array_scales.shape[0] > 0: + if broadcast_scales: + index = 0 + else: + index = i + scale[0] = array_scales[index, 0] + scale[1] = array_scales[index, 1] + scale[2] = array_scales[index, 2] + + fabric_matrices[view_index] = wp.mat44d( # type: ignore[arg-type] + wp.transpose(wp.transform_compose(position, rotation, scale)) # type: ignore[arg-type] + ) + + +@wp.kernel(enable_backward=False) +def update_indexed_local_matrix_from_world( + child_world_matrices: IndexedFabricArrayMat44d, + parent_world_matrices: IndexedFabricArrayMat44d, + child_local_matrices: IndexedFabricArrayMat44d, + indices: ArrayUInt32, +): + """Recompute child localMatrix from (parent worldMatrix, child worldMatrix). + + Computes ``child_local = inv(parent_world) * child_world`` per prim and writes the + result back to the child's :data:`omni:fabric:localMatrix` so that subsequent + ``get_local_poses`` calls see consistent values after a world-pose write. + + All three indexed arrays are expected to be indexed by the same per-view indices + (i.e. ``view_to_child_fabric``, ``view_to_parent_fabric``, ``view_to_child_fabric``) + so the kernel only needs the view-side indices. + + Storage convention: Fabric matrices are stored as the transpose of the standard + column-major math convention. Math is ``local = inv(parent) * world``; under + the transpose identity ``(A * B)^T = B^T * A^T`` (and ``inv(A^T) = inv(A)^T``) + that is equivalent to storage-side ``local^T = world^T * inv(parent^T)``, so we + can compute it directly on the stored matrices without explicit transposes. + + Args: + child_world_matrices: Indexed fabric array of child world matrices (read). + parent_world_matrices: Indexed fabric array of parent world matrices (read). + child_local_matrices: Indexed fabric array of child local matrices (written). + indices: View indices to process. + """ + i = wp.tid() + view_index = indices[i] + child_world = wp.mat44f(child_world_matrices[view_index]) + parent_world = wp.mat44f(parent_world_matrices[view_index]) + child_local_matrices[view_index] = wp.mat44d( # type: ignore[arg-type] + child_world * wp.inverse(parent_world) + ) + + +@wp.kernel(enable_backward=False) +def update_indexed_world_matrix_from_local( + child_local_matrices: IndexedFabricArrayMat44d, + parent_world_matrices: IndexedFabricArrayMat44d, + child_world_matrices: IndexedFabricArrayMat44d, + indices: ArrayUInt32, +): + """Recompute child worldMatrix from (parent worldMatrix, child localMatrix). + + Computes ``child_world = parent_world * child_local`` per prim and writes the + result back to the child's :data:`omni:fabric:worldMatrix`. Used after a + ``set_local_poses`` write so that subsequent ``get_world_poses`` calls see + consistent values. Mirror of :func:`update_indexed_local_matrix_from_world`. + + Args: + child_local_matrices: Indexed fabric array of child local matrices (read). + parent_world_matrices: Indexed fabric array of parent world matrices (read). + child_world_matrices: Indexed fabric array of child world matrices (written). + indices: View indices to process. + + Storage convention: same as :func:`update_indexed_local_matrix_from_world`. + Math is ``world = parent * local``; under the transpose identity that becomes + storage-side ``world^T = local^T * parent^T``, no explicit transposes needed. + """ + i = wp.tid() + view_index = indices[i] + child_local = wp.mat44f(child_local_matrices[view_index]) + parent_world = wp.mat44f(parent_world_matrices[view_index]) + child_world_matrices[view_index] = wp.mat44d( # type: ignore[arg-type] + child_local * parent_world + ) + + @wp.func def _decompose_transformation_matrix(m: Any): # -> tuple[wp.vec3f, wp.quatf, wp.vec3f] """Decompose a 4x4 transformation matrix into position, orientation, and scale. diff --git a/source/isaaclab/test/sim/frame_view_contract_utils.py b/source/isaaclab/test/sim/frame_view_contract_utils.py index 4336f94bd3d7..7ff53025415e 100644 --- a/source/isaaclab/test/sim/frame_view_contract_utils.py +++ b/source/isaaclab/test/sim/frame_view_contract_utils.py @@ -193,7 +193,8 @@ def test_set_world_roundtrip(device, view_factory): try: new_pos = _wp_vec3f([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]], device=device) new_quat = _wp_vec4f([[0.0, 0.0, 0.7071068, 0.7071068], [0.0, 0.0, 0.0, 1.0]], device=device) - bundle.view.set_world_poses(new_pos, new_quat) + with bundle.view.xform_world_space_writer() as w: + w.set_poses(new_pos, new_quat) ret_pos, ret_quat = bundle.view.get_world_poses() torch.testing.assert_close(_t(ret_pos), _t(new_pos), atol=ATOL, rtol=0) @@ -209,7 +210,8 @@ def test_set_local_roundtrip(device, view_factory): try: new_pos = _wp_vec3f([[0.5, 0.3, 0.1], [0.2, 0.7, 0.4]], device=device) new_quat = _wp_vec4f([[0.0, 0.0, 0.0, 1.0]] * 2, device=device) - bundle.view.set_local_poses(new_pos, new_quat) + with bundle.view.xform_local_space_writer() as w: + w.set_poses(new_pos, new_quat) ret_pos, ret_quat = bundle.view.get_local_poses() torch.testing.assert_close(_t(ret_pos), _t(new_pos), atol=ATOL, rtol=0) @@ -224,10 +226,11 @@ def test_set_world_does_not_move_parent(device, view_factory): bundle = view_factory(num_envs=2, device=device) try: parent_before = bundle.get_parent_pos(2, device).clone() - bundle.view.set_world_poses( - _wp_vec3f([[99.0, 99.0, 99.0], [88.0, 88.0, 88.0]], device=device), - _wp_vec4f([[0.0, 0.0, 0.0, 1.0]] * 2, device=device), - ) + with bundle.view.xform_world_space_writer() as w: + w.set_poses( + _wp_vec3f([[99.0, 99.0, 99.0], [88.0, 88.0, 88.0]], device=device), + _wp_vec4f([[0.0, 0.0, 0.0, 1.0]] * 2, device=device), + ) parent_after = bundle.get_parent_pos(2, device) torch.testing.assert_close(parent_after, parent_before, atol=0, rtol=0) @@ -241,10 +244,11 @@ def test_set_local_does_not_move_parent(device, view_factory): bundle = view_factory(num_envs=2, device=device) try: parent_before = bundle.get_parent_pos(2, device).clone() - bundle.view.set_local_poses( - _wp_vec3f([[0.5, 0.5, 0.5], [1.0, 1.0, 1.0]], device=device), - _wp_vec4f([[0.0, 0.0, 0.0, 1.0]] * 2, device=device), - ) + with bundle.view.xform_local_space_writer() as w: + w.set_poses( + _wp_vec3f([[0.5, 0.5, 0.5], [1.0, 1.0, 1.0]], device=device), + _wp_vec4f([[0.0, 0.0, 0.0, 1.0]] * 2, device=device), + ) parent_after = bundle.get_parent_pos(2, device) torch.testing.assert_close(parent_after, parent_before, atol=0, rtol=0) @@ -264,10 +268,11 @@ def test_set_world_updates_local(device, view_factory): desired_offset = torch.tensor([[0.3, 0.7, 0.2], [0.8, 0.1, 0.6]], device=device) new_world = parent_pos + desired_offset - bundle.view.set_world_poses( - _wp_vec3f(new_world.tolist(), device=device), - _wp_vec4f([[0.0, 0.0, 0.0, 1.0]] * 2, device=device), - ) + with bundle.view.xform_world_space_writer() as w: + w.set_poses( + _wp_vec3f(new_world.tolist(), device=device), + _wp_vec4f([[0.0, 0.0, 0.0, 1.0]] * 2, device=device), + ) local_pos = _t(bundle.view.get_local_poses()[0]) torch.testing.assert_close(local_pos, desired_offset, atol=ATOL, rtol=0) @@ -285,10 +290,11 @@ def test_set_local_updates_world(device, view_factory): try: parent_pos = bundle.get_parent_pos(2, device) new_offset = torch.tensor([[0.4, 0.9, 0.15], [0.6, 0.2, 0.85]], device=device) - bundle.view.set_local_poses( - _wp_vec3f(new_offset.tolist(), device=device), - _wp_vec4f([[0.0, 0.0, 0.0, 1.0]] * 2, device=device), - ) + with bundle.view.xform_local_space_writer() as w: + w.set_poses( + _wp_vec3f(new_offset.tolist(), device=device), + _wp_vec4f([[0.0, 0.0, 0.0, 1.0]] * 2, device=device), + ) world_pos = _t(bundle.view.get_world_poses()[0]) torch.testing.assert_close(world_pos, parent_pos + new_offset, atol=ATOL, rtol=0) @@ -303,7 +309,8 @@ def test_set_world_partial_position_only(device, view_factory): try: _, orig_quat = bundle.view.get_world_poses() new_pos = _wp_vec3f([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], device=device) - bundle.view.set_world_poses(positions=new_pos) + with bundle.view.xform_world_space_writer() as w: + w.set_poses(positions=new_pos) ret_pos, ret_quat = bundle.view.get_world_poses() torch.testing.assert_close(_t(ret_pos), _t(new_pos), atol=ATOL, rtol=0) @@ -319,7 +326,8 @@ def test_set_world_partial_orientation_only(device, view_factory): try: orig_pos, _ = bundle.view.get_world_poses() new_quat = _wp_vec4f([[0.0, 0.0, 0.7071068, 0.7071068], [0.7071068, 0.0, 0.0, 0.7071068]], device=device) - bundle.view.set_world_poses(orientations=new_quat) + with bundle.view.xform_world_space_writer() as w: + w.set_poses(orientations=new_quat) ret_pos, ret_quat = bundle.view.get_world_poses() torch.testing.assert_close(_t(ret_pos), _t(orig_pos), atol=ATOL, rtol=0) @@ -335,7 +343,8 @@ def test_set_local_partial_position_only(device, view_factory): try: _, orig_quat = bundle.view.get_local_poses() new_pos = _wp_vec3f([[0.2, 0.3, 0.4], [0.5, 0.6, 0.7]], device=device) - bundle.view.set_local_poses(translations=new_pos) + with bundle.view.xform_local_space_writer() as w: + w.set_poses(positions=new_pos) ret_pos, ret_quat = bundle.view.get_local_poses() torch.testing.assert_close(_t(ret_pos), _t(new_pos), atol=ATOL, rtol=0) @@ -352,7 +361,8 @@ def test_set_world_indexed_only_affects_subset(device, view_factory): orig_pos = _t(bundle.view.get_world_poses()[0]).clone() indices = wp.array([1, 3], dtype=wp.int32, device=device) new_pos = _wp_vec3f([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]], device=device) - bundle.view.set_world_poses(positions=new_pos, indices=indices) + with bundle.view.xform_world_space_writer() as w: + w.set_poses(positions=new_pos, indices=indices) updated = _t(bundle.view.get_world_poses()[0]) torch.testing.assert_close(updated[0], orig_pos[0], atol=0, rtol=0) @@ -365,7 +375,7 @@ def test_set_world_indexed_only_affects_subset(device, view_factory): @pytest.mark.parametrize("device", ["cpu", "cuda:0"]) def test_return_types_are_torcharray(device, view_factory): - """Public API contract — every backend returns ProxyArray from the pose getters.""" + """Public API contract — every backend returns ProxyArray from the pose and scale getters.""" bundle = view_factory(num_envs=2, device=device) try: pos_full, quat_full = bundle.view.get_world_poses() @@ -400,5 +410,137 @@ def test_return_types_are_torcharray(device, view_factory): assert isinstance(lquat_idx, ProxyArray), ( f"get_local_poses(indices)[1] must be ProxyArray, got {type(lquat_idx).__name__}" ) + + world_scales_full = bundle.view.get_world_scales() + assert isinstance(world_scales_full, ProxyArray), ( + f"get_world_scales() must be ProxyArray, got {type(world_scales_full).__name__}" + ) + world_scales_idx = bundle.view.get_world_scales(indices) + assert isinstance(world_scales_idx, ProxyArray), ( + f"get_world_scales(indices) must be ProxyArray, got {type(world_scales_idx).__name__}" + ) + + local_scales_full = bundle.view.get_local_scales() + assert isinstance(local_scales_full, ProxyArray), ( + f"get_local_scales() must be ProxyArray, got {type(local_scales_full).__name__}" + ) + local_scales_idx = bundle.view.get_local_scales(indices) + assert isinstance(local_scales_idx, ProxyArray), ( + f"get_local_scales(indices) must be ProxyArray, got {type(local_scales_idx).__name__}" + ) + + scales_full = bundle.view.get_scales() + assert isinstance(scales_full, ProxyArray), f"get_scales() must be ProxyArray, got {type(scales_full).__name__}" + scales_idx = bundle.view.get_scales(indices) + assert isinstance(scales_idx, ProxyArray), ( + f"get_scales(indices) must be ProxyArray, got {type(scales_idx).__name__}" + ) + finally: + bundle.teardown() + + +# ================================================================== +# Contract: Scales +# ================================================================== + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_local_scales_default_identity(device, view_factory): + """Local scales are (1, 1, 1) by default (no authored scale transforms).""" + bundle = view_factory(num_envs=2, device=device) + try: + scales = _t(bundle.view.get_local_scales()) + expected = torch.ones(2, 3, device=device) + torch.testing.assert_close(scales, expected, atol=ATOL, rtol=0) + finally: + bundle.teardown() + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_world_scales_default_identity(device, view_factory): + """World scales are (1, 1, 1) by default (no authored scale transforms).""" + bundle = view_factory(num_envs=2, device=device) + try: + scales = _t(bundle.view.get_world_scales()) + expected = torch.ones(2, 3, device=device) + torch.testing.assert_close(scales, expected, atol=ATOL, rtol=0) + finally: + bundle.teardown() + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_local_scales_roundtrip(device, view_factory): + """Writing scales through the local-space writer roundtrips via ``get_local_scales``.""" + bundle = view_factory(num_envs=2, device=device) + try: + new_scales = _wp_vec3f([[2.0, 3.0, 4.0], [0.5, 1.5, 2.5]], device=device) + with bundle.view.xform_local_space_writer() as w: + w.set_scales(new_scales) + + ret_scales = _t(bundle.view.get_local_scales()) + torch.testing.assert_close(ret_scales, _t(new_scales), atol=ATOL, rtol=0) + finally: + bundle.teardown() + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_world_scales_roundtrip(device, view_factory): + """Writing scales through the world-space writer roundtrips via ``get_world_scales``.""" + bundle = view_factory(num_envs=2, device=device) + try: + new_scales = _wp_vec3f([[2.0, 3.0, 4.0], [0.5, 1.5, 2.5]], device=device) + with bundle.view.xform_world_space_writer() as w: + w.set_scales(new_scales) + + ret_scales = _t(bundle.view.get_world_scales()) + torch.testing.assert_close(ret_scales, _t(new_scales), atol=ATOL, rtol=0) + finally: + bundle.teardown() + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_local_scales_do_not_affect_local_poses(device, view_factory): + """Changing scales does not change local pose translations/orientations.""" + bundle = view_factory(num_envs=2, device=device) + try: + local_pos_before = _t(bundle.view.get_local_poses()[0]).clone() + local_ori_before = _t(bundle.view.get_local_poses()[1]).clone() + + new_scales = _wp_vec3f([[3.0, 3.0, 3.0], [5.0, 5.0, 5.0]], device=device) + with bundle.view.xform_local_space_writer() as w: + w.set_scales(new_scales) + + local_pos_after = _t(bundle.view.get_local_poses()[0]) + local_ori_after = _t(bundle.view.get_local_poses()[1]) + + torch.testing.assert_close(local_pos_after, local_pos_before, atol=ATOL, rtol=0) + torch.testing.assert_close(local_ori_after, local_ori_before, atol=ATOL, rtol=0) + finally: + bundle.teardown() + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_scale_getters_return_proxyarray(device, view_factory): + """Public API contract -- scale getters return ProxyArray.""" + bundle = view_factory(num_envs=2, device=device) + try: + local_scales = bundle.view.get_local_scales() + assert isinstance(local_scales, ProxyArray), ( + f"get_local_scales() must return ProxyArray, got {type(local_scales).__name__}" + ) + world_scales = bundle.view.get_world_scales() + assert isinstance(world_scales, ProxyArray), ( + f"get_world_scales() must return ProxyArray, got {type(world_scales).__name__}" + ) + + indices = wp.array([0], dtype=wp.int32, device=bundle.view.device) + local_indexed = bundle.view.get_local_scales(indices) + assert isinstance(local_indexed, ProxyArray), ( + f"get_local_scales(indices) must return ProxyArray, got {type(local_indexed).__name__}" + ) + world_indexed = bundle.view.get_world_scales(indices) + assert isinstance(world_indexed, ProxyArray), ( + f"get_world_scales(indices) must return ProxyArray, got {type(world_indexed).__name__}" + ) finally: bundle.teardown() diff --git a/source/isaaclab/test/sim/test_views_xform_prim.py b/source/isaaclab/test/sim/test_views_xform_prim.py index 64cd86a7466f..61b281ca5e9e 100644 --- a/source/isaaclab/test/sim/test_views_xform_prim.py +++ b/source/isaaclab/test/sim/test_views_xform_prim.py @@ -226,8 +226,10 @@ def test_nested_hierarchy_world_poses(device): frames_view = FrameView("/World/Frame_.*", device=device) targets_view = FrameView("/World/Frame_.*/Target", device=device) - frames_view.set_local_poses(translations=torch.tensor(frame_positions, device=device)) - targets_view.set_local_poses(translations=torch.tensor(target_positions, device=device)) + with frames_view.xform_local_space_writer() as w: + w.set_poses(positions=torch.tensor(frame_positions, device=device)) + with targets_view.xform_local_space_writer() as w: + w.set_poses(positions=torch.tensor(target_positions, device=device)) world_pos = targets_view.get_world_poses()[0].torch expected = torch.tensor( @@ -237,6 +239,67 @@ def test_nested_hierarchy_world_poses(device): torch.testing.assert_close(world_pos, expected, atol=1e-5, rtol=0) +# ================================================================== +# USD-only: Cross-space scale conversion under a scaled parent +# ================================================================== +# +# These exercise the USD-specific world<->local scale math that the shared +# contract suite cannot cover: the contract fixtures only expose a unit-scale +# parent, and Newton has no independent local scale (local == world), so the +# parent-aware conversions below are not universal invariants. OvPhysxFrameView +# inherits this behavior by delegating to UsdFrameView. + + +def _make_scaled_parent_child_view(device, parent_scale, child_scale=None): + """Build a 1-prim view with a scaled parent (and optional authored child scale).""" + stage = sim_utils.get_current_stage() + sim_utils.create_prim("/World/Parent_0", "Xform", translation=PARENT_POS, scale=parent_scale, stage=stage) + child_kwargs = {} if child_scale is None else {"scale": child_scale} + sim_utils.create_prim("/World/Parent_0/Child", "Xform", translation=CHILD_OFFSET, stage=stage, **child_kwargs) + return FrameView("/World/Parent_.*/Child", device=device) + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_world_scale_composes_with_parent_scale(device): + """Under a scaled parent, ``get_world_scales`` returns ``parent_scale * local_scale``. + + Writes the child's local scale via the local-space writer and verifies + that reading the world scale composes with the parent's scale. + """ + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + view = _make_scaled_parent_child_view(device, parent_scale=(2.0, 1.0, 1.0)) + local_scales = wp.array([wp.vec3f(3.0, 1.0, 1.0)], dtype=wp.vec3f, device=device) + with view.xform_local_space_writer() as w: + w.set_scales(local_scales) + + world_scales = view.get_world_scales().torch + expected = torch.tensor([[6.0, 1.0, 1.0]], dtype=torch.float32, device=device) + torch.testing.assert_close(world_scales, expected, atol=1e-5, rtol=0) + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_local_scale_inverts_parent_when_writing_world_scale(device): + """Writing a world scale derives ``local = world / parent_scale`` under a scaled parent. + + Writes the child's world scale via the world-space writer and verifies + that the derived local scale is the world scale divided by the + parent's scale. + """ + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + view = _make_scaled_parent_child_view(device, parent_scale=(2.0, 1.0, 1.0)) + world_scales = wp.array([wp.vec3f(6.0, 1.0, 1.0)], dtype=wp.vec3f, device=device) + with view.xform_world_space_writer() as w: + w.set_scales(world_scales) + + local_scales = view.get_local_scales().torch + expected = torch.tensor([[3.0, 1.0, 1.0]], dtype=torch.float32, device=device) + torch.testing.assert_close(local_scales, expected, atol=1e-5, rtol=0) + + # ================================================================== # USD-only: Comparison with Isaac Sim # ================================================================== @@ -291,7 +354,8 @@ def test_with_franka_robots(device): new_pos = torch.tensor([[10.0, 10.0, 0.0], [-40.0, -40.0, 0.0]], device=device) new_quat = torch.tensor([[0.0, 0.0, 0.7071068, 0.7071068], [0.0, 0.0, -0.7071068, 0.7071068]], device=device) - view.set_world_poses(positions=new_pos, orientations=new_quat) + with view.xform_world_space_writer() as w: + w.set_poses(positions=new_pos, orientations=new_quat) ret_pos = view.get_world_poses()[0].torch torch.testing.assert_close(ret_pos, new_pos, atol=1e-5, rtol=0) diff --git a/source/isaaclab/test/terrains/check_terrain_importer.py b/source/isaaclab/test/terrains/check_terrain_importer.py index a8229023f90c..901cfc420be3 100644 --- a/source/isaaclab/test/terrains/check_terrain_importer.py +++ b/source/isaaclab/test/terrains/check_terrain_importer.py @@ -159,7 +159,8 @@ def main(): ball_initial_positions = terrain_importer.env_origins.clone() ball_initial_positions[:, 2] += 5.0 # set initial poses (writes to USD before simulation) - xform_view.set_world_poses(positions=ball_initial_positions) + with xform_view.xform_world_space_writer() as w: + w.set_poses(positions=ball_initial_positions) # Play simulator sim.reset() diff --git a/source/isaaclab/test/terrains/test_terrain_importer.py b/source/isaaclab/test/terrains/test_terrain_importer.py index 5234df4cae51..96375cab0758 100644 --- a/source/isaaclab/test/terrains/test_terrain_importer.py +++ b/source/isaaclab/test/terrains/test_terrain_importer.py @@ -316,4 +316,5 @@ def _populate_scene(sim: SimulationContext, num_balls: int = 2048, geom_sphere: ball_initial_positions[:, 2] += 5.0 # set initial poses # note: setting here writes to USD :) - ball_view.set_world_poses(positions=wp.from_torch(ball_initial_positions)) + with ball_view.xform_world_space_writer() as w: + w.set_poses(positions=wp.from_torch(ball_initial_positions)) diff --git a/source/isaaclab_mimic/changelog.d/xform-space-writer.skip b/source/isaaclab_mimic/changelog.d/xform-space-writer.skip new file mode 100644 index 000000000000..8556272d818a --- /dev/null +++ b/source/isaaclab_mimic/changelog.d/xform-space-writer.skip @@ -0,0 +1 @@ +no user-facing change: internal migration of one set_world_poses call site to the new FrameViewSpaceWriter context API 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 4ba068fc8f56..00a29bc04afc 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py +++ b/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py @@ -126,7 +126,8 @@ def set_pose(self, pose: torch.Tensor): xform_prim = self._get_xform_view() position = pose[..., :3] orientation = pose[..., 3:] - xform_prim.set_world_poses(wp.from_torch(position.contiguous()), wp.from_torch(orientation.contiguous()), None) + with xform_prim.xform_world_space_writer() as writer: + writer.set_poses(wp.from_torch(position.contiguous()), wp.from_torch(orientation.contiguous()), None) class RelativePose(HasPose): diff --git a/source/isaaclab_newton/changelog.d/fabric-local-poses.rst b/source/isaaclab_newton/changelog.d/fabric-local-poses.rst new file mode 100644 index 000000000000..c5c0b0f3c7d2 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/fabric-local-poses.rst @@ -0,0 +1,19 @@ +Added +^^^^^ + +* Added :meth:`~isaaclab_newton.sim.views.NewtonSiteFrameView.get_local_scales` + and :meth:`~isaaclab_newton.sim.views.NewtonSiteFrameView.get_world_scales` + for reading transform (xform) scales. Scale writes go through the writer + scope (see the ``xform-space-writer`` fragment). These transform-scale + APIs are intentionally separate from Newton collision shape geometry + sizes. + +Deprecated +^^^^^^^^^^ + +* Deprecated :meth:`~isaaclab_newton.sim.views.NewtonSiteFrameView.get_scales` + and :meth:`~isaaclab_newton.sim.views.NewtonSiteFrameView.set_scales` in favor + of the explicit transform-scale getters ``get_world_scales`` / + ``get_local_scales`` (and the writer scope's ``set_scales``). The + deprecated methods still work but emit a ``DeprecationWarning`` and + preserve Newton's legacy collision shape geometry-scale behavior. diff --git a/source/isaaclab_newton/changelog.d/xform-space-writer.rst b/source/isaaclab_newton/changelog.d/xform-space-writer.rst new file mode 100644 index 000000000000..e04d5325e6aa --- /dev/null +++ b/source/isaaclab_newton/changelog.d/xform-space-writer.rst @@ -0,0 +1,13 @@ +Changed +^^^^^^^ + +* :class:`~isaaclab_newton.sim.views.NewtonSiteFrameView` now ships + pass-through ``FrameViewWorldSpaceWriter`` / ``FrameViewLocalSpaceWriter`` + implementations so writes follow the new + :meth:`~isaaclab.sim.views.BaseFrameView.xform_world_space_writer` / + :meth:`~isaaclab.sim.views.BaseFrameView.xform_local_space_writer` context API. + ``set_world_poses`` / ``set_local_poses`` shims still work (one-time + ``DeprecationWarning`` per class). The legacy ``set_scales`` / + ``get_scales`` paths continue to operate on Newton collision-shape + geometry sizes -- they are not routed through the writer because the + writer's ``set_scales`` writes the transform-scale state. diff --git a/source/isaaclab_newton/isaaclab_newton/sim/views/newton_site_frame_view.py b/source/isaaclab_newton/isaaclab_newton/sim/views/newton_site_frame_view.py index c079b2a86b09..bbb29c745ea1 100644 --- a/source/isaaclab_newton/isaaclab_newton/sim/views/newton_site_frame_view.py +++ b/source/isaaclab_newton/isaaclab_newton/sim/views/newton_site_frame_view.py @@ -14,9 +14,10 @@ from pxr import UsdPhysics import isaaclab.sim as sim_utils -from isaaclab.cloner.cloner_utils import get_suffix, iter_clone_plan_matches +from isaaclab.cloner.cloner_utils import get_suffix, iter_clone_plan_matches, split_clone_template from isaaclab.physics import PhysicsEvent from isaaclab.sim.views.base_frame_view import BaseFrameView +from isaaclab.sim.views.xform_space_writer import FrameViewLocalSpaceWriter, FrameViewWorldSpaceWriter from isaaclab.utils.string import resolve_matching_names from isaaclab.utils.warp import ProxyArray @@ -104,7 +105,7 @@ def _write_site_local_from_local_poses( @wp.kernel -def _gather_scales( +def _gather_shape_scales( shape_scale: wp.array(dtype=wp.vec3f), shape_body: wp.array(dtype=wp.int32), site_body: wp.array(dtype=wp.int32), @@ -112,7 +113,7 @@ def _gather_scales( num_shapes: wp.int32, out_scales: wp.array(dtype=wp.vec3f), ): - """Gather per-site scales from collision shapes on the same body.""" + """Gather legacy per-site geometry scales from collision shapes on the same body.""" i = wp.tid() si = indices[i] bid = site_body[si] @@ -126,7 +127,7 @@ def _gather_scales( @wp.kernel -def _scatter_scales( +def _scatter_shape_scales( site_body: wp.array(dtype=wp.int32), indices: wp.array(dtype=wp.int32), new_scales: wp.array(dtype=wp.vec3f), @@ -134,7 +135,7 @@ def _scatter_scales( num_shapes: wp.int32, shape_scale: wp.array(dtype=wp.vec3f), ): - """Scatter per-site scales to collision shapes on the same body.""" + """Scatter legacy per-site geometry scales to collision shapes on the same body.""" i = wp.tid() si = indices[i] bid = site_body[si] @@ -143,6 +144,28 @@ def _scatter_scales( shape_scale[s] = new_scales[i] +@wp.kernel +def _gather_xform_scales( + site_xform_scale: wp.array(dtype=wp.vec3f), + indices: wp.array(dtype=wp.int32), + out_scales: wp.array(dtype=wp.vec3f), +): + """Gather per-site xform scales.""" + i = wp.tid() + out_scales[i] = site_xform_scale[indices[i]] + + +@wp.kernel +def _scatter_xform_scales( + indices: wp.array(dtype=wp.int32), + new_scales: wp.array(dtype=wp.vec3f), + site_xform_scale: wp.array(dtype=wp.vec3f), +): + """Scatter per-site xform scales.""" + i = wp.tid() + site_xform_scale[indices[i]] = new_scales[i] + + class NewtonSiteFrameView(BaseFrameView): """Batched Newton site view for non-physics frames. @@ -178,43 +201,51 @@ def __init__( stage = sim_utils.get_current_stage() if stage is None else stage self._site_specs = self._resolve_site_specs(stage, validate_xform_ops) self._site_labels: list[str] = [] + self._site_label_scales: list[tuple[float, float, float]] = [] self._site_body: wp.array | None = None self._site_local: wp.array | None = None + self._site_xform_scale: wp.array | None = None self._site_indices: wp.array | None = None self._pos_buf: wp.array | None = None self._quat_buf: wp.array | None = None self._local_pos_buf: wp.array | None = None self._local_quat_buf: wp.array | None = None + self._scale_buf: wp.array | None = None self._pos_ta: ProxyArray | None = None self._quat_ta: ProxyArray | None = None self._local_pos_ta: ProxyArray | None = None self._local_quat_ta: ProxyArray | None = None + self._scale_ta: ProxyArray | None = None self._count = 0 model = NewtonManager.get_model() if model is not None: self._initialize_from_specs(model) else: - for body_patterns, xform, per_world, _env_ids in self._site_specs: + for body_patterns, xform, scale, per_world, _env_ids in self._site_specs: if body_patterns is None: self._site_labels.append(NewtonManager.cl_register_site(None, xform, per_world=per_world)) + self._site_label_scales.append(scale) else: for body_pattern in body_patterns: self._site_labels.append(NewtonManager.cl_register_site(body_pattern, xform)) + self._site_label_scales.append(scale) self._physics_ready_handle = NewtonManager.register_callback( self._on_physics_ready, PhysicsEvent.PHYSICS_READY, name=f"site_view_{self._prim_path}" ) def _resolve_site_specs( self, stage, validate_xform_ops: bool - ) -> list[tuple[tuple[str, ...] | None, wp.transform, bool, tuple[int, ...] | None]]: + ) -> list[tuple[tuple[str, ...] | None, wp.transform, tuple[float, float, float], bool, tuple[int, ...] | None]]: """Resolve source prims into Newton site registration specs.""" plan = sim_utils.SimulationContext.instance().get_clone_plan() model = NewtonManager.get_model() body_labels = list(model.body_label) if model is not None else () shape_labels = list(model.shape_label) if model is not None else () use_clone_body_pattern = model is None - specs: list[tuple[tuple[str, ...] | None, wp.transform, bool, tuple[int, ...] | None]] = [] + specs: list[ + tuple[tuple[str, ...] | None, wp.transform, tuple[float, float, float], bool, tuple[int, ...] | None] + ] = [] for path_expr in self._prim_paths: if resolve_matching_names(path_expr, body_labels, raise_when_no_match=False)[1]: @@ -268,8 +299,8 @@ def _resolve_source_prim( env_ids: tuple[int, ...] | None, use_clone_body_pattern: bool, stage, - ) -> tuple[tuple[str, ...] | None, wp.transform, bool, tuple[int, ...] | None]: - """Resolve one source prim into body patterns and a local frame.""" + ) -> tuple[tuple[str, ...] | None, wp.transform, tuple[float, float, float], bool, tuple[int, ...] | None]: + """Resolve one source prim into body patterns, local frame, and xform scale.""" prim_path = prim.GetPath().pathString if prim.HasAPI(UsdPhysics.RigidBodyAPI) or prim.HasAPI(UsdPhysics.ArticulationRootAPI): raise ValueError( @@ -281,6 +312,13 @@ def _resolve_source_prim( if not sim_utils.validate_standard_xform_ops(prim): raise ValueError(f"FrameView prim '{prim_path}' does not have standard xform ops.") + scale_attr = prim.GetAttribute("xformOp:scale") + scale = ( + tuple(float(v) for v in scale_attr.Get()) + if scale_attr and scale_attr.HasAuthoredValue() + else (1.0, 1.0, 1.0) + ) + body_prim = prim.GetParent() while body_prim and body_prim.IsValid(): if body_prim.HasAPI(UsdPhysics.RigidBodyAPI) or body_prim.HasAPI(UsdPhysics.ArticulationRootAPI): @@ -300,7 +338,7 @@ def _resolve_source_prim( raise RuntimeError( f"FrameView destination root '{destination_root}' does not end with '{suffix}'." ) - return (destination_root[: -len(suffix)],), wp.transform(pos, quat), False, env_ids + return (destination_root[: -len(suffix)],), wp.transform(pos, quat), scale, False, env_ids body_patterns = [] for env_id in env_ids: destination_root = destination_template.format(env_id) @@ -309,7 +347,7 @@ def _resolve_source_prim( f"FrameView destination root '{destination_root}' does not end with '{suffix}'." ) body_patterns.append(destination_root[: -len(suffix)]) - return tuple(body_patterns), wp.transform(pos, quat), False, env_ids + return tuple(body_patterns), wp.transform(pos, quat), scale, False, env_ids else: raise RuntimeError(f"FrameView source body '{body_path}' is not under '{source_root}'.") if use_clone_body_pattern: @@ -318,18 +356,18 @@ def _resolve_source_prim( body_patterns = tuple(destination_template.format(env_id) + suffix for env_id in env_ids) else: body_patterns = (body_path,) - return body_patterns, wp.transform(pos, quat), False, env_ids + return body_patterns, wp.transform(pos, quat), scale, False, env_ids body_prim = body_prim.GetParent() ref_path = source_root if source_root is not None and destination_template is not None: - instance_template = destination_template.partition("{}")[0] + "{}" - source_suffix = get_suffix(source_root, instance_template) + template_prefix, _ = split_clone_template(destination_template) + source_suffix = get_suffix(source_root, template_prefix + "{}") if source_suffix is not None: ref_path = source_root[: -len(source_suffix)] if source_suffix else source_root ref_prim = stage.GetPrimAtPath(ref_path) if ref_path is not None else None pos, quat = sim_utils.resolve_prim_pose(prim, ref_prim if ref_prim and ref_prim.IsValid() else None) - return None, wp.transform(pos, quat), source_root is not None, env_ids + return None, wp.transform(pos, quat), scale, source_root is not None, env_ids def _on_physics_ready(self, _event) -> None: """Callback invoked when the Newton model becomes available.""" @@ -342,8 +380,9 @@ def _initialize_from_site_map(self, model) -> None: xform_t = wp.to_torch(model.shape_transform) site_bodies: list[int] = [] site_locals: list[list[float]] = [] + site_scales: list[tuple[float, float, float]] = [] - for site_label in self._site_labels: + for site_label, scale in zip(self._site_labels, self._site_label_scales, strict=True): global_idx, per_world = site_map[site_label] site_indices = ( [global_idx] if per_world is None else [site_idx for sites in per_world for site_idx in sites] @@ -351,16 +390,18 @@ def _initialize_from_site_map(self, model) -> None: for site_idx in site_indices: site_bodies.append(int(body_t[site_idx].item())) site_locals.append([float(v) for v in xform_t[site_idx].tolist()]) + site_scales.append(scale) - self._create_buffers(site_bodies, site_locals) + self._create_buffers(site_bodies, site_locals, site_scales) def _initialize_from_specs(self, model) -> None: """Initialize arrays directly from resolved specs and Newton body labels.""" body_labels = list(model.body_label) site_bodies: list[int] = [] site_locals: list[list[float]] = [] + site_scales: list[tuple[float, float, float]] = [] - for body_patterns, xform, per_world, env_ids in self._site_specs: + for body_patterns, xform, scale, per_world, env_ids in self._site_specs: if body_patterns is None: if per_world: if NewtonManager._world_xforms is None: @@ -370,9 +411,11 @@ def _initialize_from_specs(self, model) -> None: world_xform = NewtonManager._world_xforms[world_id] site_bodies.append(WORLD_BODY_INDEX) site_locals.append([float(v) for v in wp.transform_multiply(world_xform, xform)]) + site_scales.append(scale) else: site_bodies.append(WORLD_BODY_INDEX) site_locals.append([float(v) for v in xform]) + site_scales.append(scale) continue for body_pattern in body_patterns: @@ -385,24 +428,33 @@ def _initialize_from_specs(self, model) -> None: for body_idx in matched_indices: site_bodies.append(body_idx) site_locals.append([float(v) for v in xform]) + site_scales.append(scale) - self._create_buffers(site_bodies, site_locals) + self._create_buffers(site_bodies, site_locals, site_scales) - def _create_buffers(self, site_bodies: list[int], site_locals: list[list[float]]) -> None: + def _create_buffers( + self, + site_bodies: list[int], + site_locals: list[list[float]], + site_scales: list[tuple[float, float, float]], + ) -> None: """Allocate view buffers from body indices and local transforms.""" self._count = len(site_bodies) device = self._device self._site_body = wp.array(site_bodies, dtype=wp.int32, device=device) self._site_local = wp.array([wp.transform(*x) for x in site_locals], dtype=wp.transformf, device=device) + self._site_xform_scale = wp.array([wp.vec3f(*scale) for scale in site_scales], dtype=wp.vec3f, device=device) self._site_indices = wp.array(list(range(self._count)), dtype=wp.int32, device=device) self._pos_buf = wp.zeros(self._count, dtype=wp.vec3f, device=device) self._quat_buf = wp.zeros(self._count, dtype=wp.vec4f, device=device) self._local_pos_buf = wp.zeros(self._count, dtype=wp.vec3f, device=device) self._local_quat_buf = wp.zeros(self._count, dtype=wp.vec4f, device=device) + self._scale_buf = wp.zeros(self._count, dtype=wp.vec3f, device=device) self._pos_ta = ProxyArray(self._pos_buf) self._quat_ta = ProxyArray(self._quat_buf) self._local_pos_ta = ProxyArray(self._local_pos_buf) self._local_quat_ta = ProxyArray(self._local_quat_buf) + self._scale_ta = ProxyArray(self._site_xform_scale) @property def prims(self) -> list: @@ -422,7 +474,21 @@ def device(self) -> str: """Device where arrays are allocated.""" return self._device - def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: + # ------------------------------------------------------------------ + # Writer factory hooks (pass-through; Newton has no separate Fabric storage) + # ------------------------------------------------------------------ + + def _make_world_space_writer(self) -> FrameViewWorldSpaceWriter: + return _NewtonWorldSpaceWriter(self) + + def _make_local_space_writer(self) -> FrameViewLocalSpaceWriter: + return _NewtonLocalSpaceWriter(self) + + # ------------------------------------------------------------------ + # Backend hooks + # ------------------------------------------------------------------ + + def _get_world_poses_impl(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: """Get world-space positions and orientations.""" state = NewtonManager.get_state_0() site_indices = self._site_indices if indices is None else indices @@ -441,7 +507,7 @@ def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, return self._pos_ta, self._quat_ta return ProxyArray(pos_buf), ProxyArray(quat_buf) - def set_world_poses( + def _apply_world_pose_write( self, positions: wp.array | None = None, orientations: wp.array | None = None, @@ -453,7 +519,7 @@ def set_world_poses( state = NewtonManager.get_state_0() if positions is None or orientations is None: - cur_pos_ta, cur_quat_ta = self.get_world_poses(indices) + cur_pos_ta, cur_quat_ta = self._get_world_poses_impl(indices) if positions is None: positions = cur_pos_ta.warp if orientations is None: @@ -468,7 +534,7 @@ def set_world_poses( device=self._device, ) - def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: + def _get_local_poses_impl(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: """Get body-local positions and orientations.""" site_indices = self._site_indices if indices is None else indices n = self.count if indices is None else len(indices) @@ -486,7 +552,7 @@ def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, return self._local_pos_ta, self._local_quat_ta return ProxyArray(pos_buf), ProxyArray(quat_buf) - def set_local_poses( + def _apply_local_pose_write( self, translations: wp.array | None = None, orientations: wp.array | None = None, @@ -497,7 +563,7 @@ def set_local_poses( return if translations is None or orientations is None: - cur_pos_ta, cur_quat_ta = self.get_local_poses(indices) + cur_pos_ta, cur_quat_ta = self._get_local_poses_impl(indices) if translations is None: translations = cur_pos_ta.warp if orientations is None: @@ -512,31 +578,137 @@ def set_local_poses( device=self._device, ) - def get_scales(self, indices: wp.array | None = None) -> wp.array: - """Get per-site scales by reading from the first collision shape on the same body.""" + # ------------------------------------------------------------------ + # Scales + # ------------------------------------------------------------------ + + def _get_world_scales_impl(self, indices: wp.array | None = None) -> ProxyArray: + """Get per-site world xform scales. + + These are transform scales, matching the USD FrameView scale API. They + are intentionally separate from Newton collision shape geometry sizes. + """ + if indices is None: + return self._scale_ta + n = len(indices) + out = wp.zeros(n, dtype=wp.vec3f, device=self._device) + wp.launch( + _gather_xform_scales, + dim=n, + inputs=[self._site_xform_scale, indices], + outputs=[out], + device=self._device, + ) + return ProxyArray(out) + + def _get_local_scales_impl(self, indices: wp.array | None = None) -> ProxyArray: + """Get per-site local xform scales. + + These are transform scales, matching the USD FrameView scale API. They + are intentionally separate from Newton collision shape geometry sizes. + """ + return self._get_world_scales_impl(indices) + + def _apply_world_scale_write(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Set per-site world xform scales. + + These update transform scale state only; use deprecated ``set_scales`` if + legacy Newton collision shape geometry-scale behavior is required. + """ + if indices is None: + indices = self._site_indices + n = self.count if indices is self._site_indices else len(indices) + wp.launch( + _scatter_xform_scales, + dim=n, + inputs=[indices, scales, self._site_xform_scale], + device=self._device, + ) + + def _apply_local_scale_write(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Set per-site local xform scales. + + These update transform scale state only; use deprecated ``set_scales`` if + legacy Newton collision shape geometry-scale behavior is required. + """ + self._apply_world_scale_write(scales, indices) + + def _get_legacy_shape_scales(self, indices: wp.array | None = None) -> ProxyArray: + """Get Newton legacy geometry scales from collision shapes.""" model = NewtonManager.get_model() num_shapes = model.shape_count site_indices = self._site_indices if indices is None else indices n = self.count if indices is None else len(indices) out = wp.zeros(n, dtype=wp.vec3f, device=self._device) wp.launch( - _gather_scales, + _gather_shape_scales, dim=n, inputs=[model.shape_scale, model.shape_body, self._site_body, site_indices, num_shapes], outputs=[out], device=self._device, ) - return out + return ProxyArray(out) - def set_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: - """Set per-site scales by writing to all collision shapes on the same body.""" + def _set_legacy_shape_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Set Newton legacy geometry scales on collision shapes.""" model = NewtonManager.get_model() num_shapes = model.shape_count site_indices = self._site_indices if indices is None else indices n = self.count if indices is None else len(indices) wp.launch( - _scatter_scales, + _scatter_shape_scales, dim=n, inputs=[self._site_body, site_indices, scales, model.shape_body, num_shapes, model.shape_scale], device=self._device, ) + + def _get_scales_impl(self, indices: wp.array | None = None) -> ProxyArray: + """Newton legacy: get_scales returns collision shape geometry scales.""" + return self._get_legacy_shape_scales(indices) + + def _set_scales_impl(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Newton legacy: deprecated set_scales writes collision shape geometry scales. + + Newton's legacy ``set_scales`` path is *not* routed through the + :class:`FrameViewSpaceWriterBase` API because it targets a different state + (collision-shape geometry sizes) than the transform-scale state that + the writer's :meth:`~FrameViewSpaceWriterBase.set_scales` operates on. + """ + self._set_legacy_shape_scales(scales, indices) + + +# ---------------------------------------------------------------------- +# Pass-through writer classes +# ---------------------------------------------------------------------- + + +class _NewtonWorldSpaceWriter(FrameViewWorldSpaceWriter): + """Newton world-space writer: pass-through to backend ``_apply_*`` hooks.""" + + def set_poses(self, positions=None, orientations=None, indices=None) -> None: + self._view._apply_world_pose_write(positions, orientations, indices) # type: ignore[attr-defined] + + def set_scales(self, scales, indices=None) -> None: + self._view._apply_world_scale_write(scales, indices) # type: ignore[attr-defined] + + def get_poses(self, indices=None) -> tuple[ProxyArray, ProxyArray]: + return self._view._get_world_poses_impl(indices) # type: ignore[attr-defined] + + def get_scales(self, indices=None) -> ProxyArray: + return self._view._get_world_scales_impl(indices) # type: ignore[attr-defined] + + +class _NewtonLocalSpaceWriter(FrameViewLocalSpaceWriter): + """Newton local-space writer: pass-through to backend ``_apply_*`` hooks.""" + + def set_poses(self, positions=None, orientations=None, indices=None) -> None: + self._view._apply_local_pose_write(positions, orientations, indices) # type: ignore[attr-defined] + + def set_scales(self, scales, indices=None) -> None: + self._view._apply_local_scale_write(scales, indices) # type: ignore[attr-defined] + + def get_poses(self, indices=None) -> tuple[ProxyArray, ProxyArray]: + return self._view._get_local_poses_impl(indices) # type: ignore[attr-defined] + + def get_scales(self, indices=None) -> ProxyArray: + return self._view._get_local_scales_impl(indices) # type: ignore[attr-defined] diff --git a/source/isaaclab_newton/test/sim/test_views_xform_prim_newton.py b/source/isaaclab_newton/test/sim/test_views_xform_prim_newton.py index d114a1da2a80..174f0684fb55 100644 --- a/source/isaaclab_newton/test/sim/test_views_xform_prim_newton.py +++ b/source/isaaclab_newton/test/sim/test_views_xform_prim_newton.py @@ -215,7 +215,8 @@ def test_world_attached_set_world_roundtrip(device): new_pos = _wp_vec3f([[10.0, 20.0, 30.0]], device=device) new_quat = _wp_vec4f([[0.0, 0.0, 0.0, 1.0]], device=device) - view.set_world_poses(new_pos, new_quat) + with view.xform_world_space_writer() as w: + w.set_poses(new_pos, new_quat) ret_pos, ret_quat = view.get_world_poses() torch.testing.assert_close(ret_pos.torch, wp.to_torch(new_pos), atol=1e-5, rtol=0) diff --git a/source/isaaclab_ovphysx/changelog.d/fabric-local-poses.rst b/source/isaaclab_ovphysx/changelog.d/fabric-local-poses.rst new file mode 100644 index 000000000000..fdf33f0b7d94 --- /dev/null +++ b/source/isaaclab_ovphysx/changelog.d/fabric-local-poses.rst @@ -0,0 +1,19 @@ +Added +^^^^^ + +* Added :meth:`~isaaclab_ovphysx.sim.views.OvPhysxFrameView.get_local_scales` + and :meth:`~isaaclab_ovphysx.sim.views.OvPhysxFrameView.get_world_scales`, + which delegate to the internal :class:`~isaaclab.sim.views.UsdFrameView`. + Scale writes go through the writer scope (see the ``xform-space-writer`` + fragment). + +Deprecated +^^^^^^^^^^ + +* Deprecated :meth:`~isaaclab_ovphysx.sim.views.OvPhysxFrameView.get_scales` and + :meth:`~isaaclab_ovphysx.sim.views.OvPhysxFrameView.set_scales`. For reads, + use the explicit ``get_local_scales`` (operates on ``xformOp:scale``) or + ``get_world_scales``. For writes, use the writer scope's + ``set_scales``. The deprecated methods still work but emit a + ``DeprecationWarning`` and default to local scales, preserving prior + behavior. diff --git a/source/isaaclab_ovphysx/changelog.d/xform-space-writer.rst b/source/isaaclab_ovphysx/changelog.d/xform-space-writer.rst new file mode 100644 index 000000000000..3561232a311e --- /dev/null +++ b/source/isaaclab_ovphysx/changelog.d/xform-space-writer.rst @@ -0,0 +1,13 @@ +Changed +^^^^^^^ + +* :class:`~isaaclab_ovphysx.sim.views.OvPhysxFrameView` now ships + pass-through ``FrameViewWorldSpaceWriter`` / ``FrameViewLocalSpaceWriter`` + implementations so writes follow the new + :meth:`~isaaclab.sim.views.BaseFrameView.xform_world_space_writer` / + :meth:`~isaaclab.sim.views.BaseFrameView.xform_local_space_writer` context API. + ``set_world_poses`` / ``set_local_poses`` shims still work (one-time + ``DeprecationWarning`` per class). Scale writes inside the writer scope + delegate to the internal :class:`~isaaclab.sim.views.UsdFrameView` and + land in the USD stage (no propagation to OVPhysX-side collision-shape + scales). diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/sim/views/ovphysx_frame_view.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/sim/views/ovphysx_frame_view.py index 1eda486a1205..ff7912f20f60 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/sim/views/ovphysx_frame_view.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/sim/views/ovphysx_frame_view.py @@ -19,6 +19,7 @@ from isaaclab.physics import PhysicsEvent from isaaclab.sim.views.base_frame_view import BaseFrameView from isaaclab.sim.views.usd_frame_view import UsdFrameView +from isaaclab.sim.views.xform_space_writer import FrameViewLocalSpaceWriter, FrameViewWorldSpaceWriter from isaaclab.utils.warp import ProxyArray from isaaclab_ovphysx.physics import OvPhysxManager @@ -284,7 +285,7 @@ class OvPhysxFrameView(BaseFrameView): Scales and visibility delegate to an internal :class:`UsdFrameView` (lazy-constructed on first call). - Pose getters return :class:`~isaaclab.utils.warp.ProxyArray`. Setters + Getters return :class:`~isaaclab.utils.warp.ProxyArray`. Setters accept ``wp.array``. Limitations (v1): @@ -357,6 +358,7 @@ def _initialize_impl(self, physx: Any) -> None: ``env_0`` replaced by the row's env_id. """ from isaaclab_ovphysx import tensor_types as TT # noqa: PLC0415 + from isaaclab_ovphysx.sim.views.ovphysx_view import OvPhysxView # noqa: PLC0415 xform_cache = UsdGeom.XformCache(Usd.TimeCode.Default()) identity_xform7 = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0] @@ -406,8 +408,11 @@ def _initialize_impl(self, physx: Any) -> None: # 4. Create the RIGID_BODY_POSE binding (or operate in world-only mode). if patterns: pattern = patterns[0] - self._pose_binding = physx.create_tensor_binding(pattern=pattern, tensor_type=TT.RIGID_BODY_POSE) - if self._pose_binding.shape[0] == 0: + self._root_view = OvPhysxView(physx, pattern=pattern, device=self._device) + # ``try_binding_for`` returns None for a zero-match binding (the view rejects a + # 0-count binding); surface that as the explicit zero-bodies error below. + self._pose_binding = self._root_view.try_binding_for(TT.RIGID_BODY_POSE) + if self._pose_binding is None: raise RuntimeError( f"OvPhysxFrameView: RIGID_BODY_POSE binding for pattern {pattern!r} matched zero bodies." ) @@ -415,6 +420,7 @@ def _initialize_impl(self, physx: Any) -> None: binding_paths: list[str] = list(self._pose_binding.prim_paths) else: # All prims resolved as world-attached: no binding needed; kernels only hit the -1 branch. + self._root_view = None self._pose_binding = None self._pose_buf = wp.zeros((1, 7), dtype=wp.float32, device=self._device) binding_paths = [] @@ -613,7 +619,7 @@ def _current_body_q(self) -> wp.array: buffer ``[num_bodies]``. """ if self._pose_binding is not None: - self._pose_binding.read(self._pose_buf) + self._root_view.read_into("rigid_body_pose", self._pose_buf) return self._pose_buf.view(wp.transformf) # ------------------------------------------------------------------ @@ -624,17 +630,22 @@ def _current_body_q(self) -> wp.array: # World poses # ------------------------------------------------------------------ - def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: - """Get world-space positions and orientations. + # ------------------------------------------------------------------ + # Writer factory hooks (pass-through; OvPhysX has no separate Fabric storage) + # ------------------------------------------------------------------ - Args: - indices: Subset of sites to query. ``None`` means all sites. + def _make_world_space_writer(self) -> FrameViewWorldSpaceWriter: + return _OvPhysxWorldSpaceWriter(self) - Returns: - A tuple ``(positions, orientations)`` of :class:`~isaaclab.utils.warp.ProxyArray` - wrappers. Use ``.warp`` for the underlying ``wp.array`` or ``.torch`` for a - cached zero-copy ``torch.Tensor`` view. - """ + def _make_local_space_writer(self) -> FrameViewLocalSpaceWriter: + return _OvPhysxLocalSpaceWriter(self) + + # ------------------------------------------------------------------ + # Backend hooks + # ------------------------------------------------------------------ + + def _get_world_poses_impl(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: + """Get world-space positions and orientations.""" self._require_initialized() body_q = self._current_body_q() @@ -660,31 +671,20 @@ def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ) return self._pos_ta, self._quat_ta - def set_world_poses( + def _apply_world_pose_write( self, positions: wp.array | None = None, orientations: wp.array | None = None, indices: wp.array | None = None, ) -> None: - """Set world-space positions and/or orientations. - - Updates ``site_local`` so that ``body_q[body] * site_local`` yields the - desired world pose. Does **not** modify ``body_q``. - - Args: - positions: Desired world positions ``(M, 3)`` [m]. ``None`` leaves - positions unchanged. - orientations: Desired world quaternions ``(M, 4)`` as - ``(qx, qy, qz, qw)``. ``None`` leaves orientations unchanged. - indices: Subset of sites to update. ``None`` means all sites. - """ + """Set world-space positions and/or orientations.""" if positions is None and orientations is None: return self._require_initialized() body_q = self._current_body_q() if positions is None or orientations is None: - cur_pos_ta, cur_quat_ta = self.get_world_poses(indices) + cur_pos_ta, cur_quat_ta = self._get_world_poses_impl(indices) if positions is None: positions = cur_pos_ta.warp if orientations is None: @@ -709,18 +709,8 @@ def set_world_poses( # Local poses (parent-relative) # ------------------------------------------------------------------ - def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: - """Get parent-relative positions and orientations. - - Computes ``inv(parent_world) * prim_world`` for each site. - - Args: - indices: Subset of sites to query. ``None`` means all sites. - - Returns: - A tuple ``(translations, orientations)`` of :class:`~isaaclab.utils.warp.ProxyArray` - wrappers. - """ + def _get_local_poses_impl(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: + """Get parent-relative positions and orientations.""" self._require_initialized() body_q = self._current_body_q() @@ -759,30 +749,20 @@ def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ) return self._local_pos_ta, self._local_quat_ta - def set_local_poses( + def _apply_local_pose_write( self, translations: wp.array | None = None, orientations: wp.array | None = None, indices: wp.array | None = None, ) -> None: - """Set parent-relative translations and/or orientations. - - Updates ``site_local`` only; does **not** modify ``body_q``. - - Args: - translations: Desired parent-relative translations ``(M, 3)`` [m]. - ``None`` leaves translations unchanged. - orientations: Desired parent-relative quaternions ``(M, 4)`` as - ``(qx, qy, qz, qw)``. ``None`` leaves orientations unchanged. - indices: Subset of sites to update. ``None`` means all sites. - """ + """Set parent-relative translations and/or orientations.""" if translations is None and orientations is None: return self._require_initialized() body_q = self._current_body_q() if translations is None or orientations is None: - cur_pos_ta, cur_quat_ta = self.get_local_poses(indices) + cur_pos_ta, cur_quat_ta = self._get_local_poses_impl(indices) if translations is None: translations = cur_pos_ta.warp if orientations is None: @@ -834,39 +814,43 @@ def _ensure_usd_view(self) -> UsdFrameView: ) return self._usd_view - def get_scales(self, indices: wp.array | None = None) -> wp.array: - """Get prim scales from the USD stage's ``xformOp:scale`` attribute. + def _get_local_scales_impl(self, indices: wp.array | None = None) -> ProxyArray: + """Get local-space scales (xformOp:scale) via the USD view. .. note:: This reads the *static* USD authored value, not a live physics-state value. OVPhysX does not maintain a per-shape ``shape_scale`` array equivalent to Newton's ``model.shape_scale``, so sim-driven scale - updates are not reflected here. For sites under ``clone_usd=False`` - envs without authored USD prims, the read returns the env_0 - template's scale via the lazy internal :class:`UsdFrameView`. - - Args: - indices: Subset of sites to query. ``None`` means all sites. - - Returns: - ``wp.array`` of shape ``(M, 3)``. + updates are not reflected here. """ - return self._ensure_usd_view().get_scales(indices) + return self._ensure_usd_view()._get_local_scales_impl(indices) - def set_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: - """Set prim scales by writing the USD ``xformOp:scale`` attribute. + def _get_world_scales_impl(self, indices: wp.array | None = None) -> ProxyArray: + """Get world-space (composed) scales via the USD view.""" + return self._ensure_usd_view()._get_world_scales_impl(indices) + + def _apply_local_scale_write(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Set local-space scales (xformOp:scale) via the USD view. .. note:: The write lands in the USD stage but does *not* propagate to any OVPhysX-side collision-shape scale. PhysX is unaffected; this is a - stage-only annotation. Use :class:`~isaaclab_ovphysx.assets.RigidObject` - APIs if you need to change physics-effective shape sizes. - - Args: - scales: Scales ``(M, 3)`` as ``wp.array``. - indices: Subset of sites to update. ``None`` means all sites. + stage-only annotation. """ - self._ensure_usd_view().set_scales(scales, indices) + self._ensure_usd_view()._apply_local_scale_write(scales, indices) + + def _apply_world_scale_write(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Set world-space scales via the USD view.""" + self._ensure_usd_view()._apply_world_scale_write(scales, indices) + + def _get_scales_impl(self, indices=None): + """OvPhysX legacy: deprecated get_scales returns local scales.""" + return self._get_local_scales_impl(indices) + + def _set_scales_impl(self, scales, indices=None): + """OvPhysX legacy: deprecated set_scales writes local scales via a one-shot writer scope.""" + with self.xform_local_space_writer() as writer: + writer.set_scales(scales, indices) def get_visibility(self, indices: wp.array | None = None): """Get visibility for prims in the view (USD-backed). @@ -888,3 +872,40 @@ def _gf_matrix_to_xform7(mat: Gf.Matrix4d) -> list[float]: q = mat.ExtractRotationQuat() imag = q.GetImaginary() return [float(t[0]), float(t[1]), float(t[2]), float(imag[0]), float(imag[1]), float(imag[2]), float(q.GetReal())] + + +# ---------------------------------------------------------------------- +# Pass-through writer classes +# ---------------------------------------------------------------------- + + +class _OvPhysxWorldSpaceWriter(FrameViewWorldSpaceWriter): + """OvPhysX world-space writer: pass-through to backend ``_apply_*`` hooks.""" + + def set_poses(self, positions=None, orientations=None, indices=None) -> None: + self._view._apply_world_pose_write(positions, orientations, indices) # type: ignore[attr-defined] + + def set_scales(self, scales, indices=None) -> None: + self._view._apply_world_scale_write(scales, indices) # type: ignore[attr-defined] + + def get_poses(self, indices=None) -> tuple[ProxyArray, ProxyArray]: + return self._view._get_world_poses_impl(indices) # type: ignore[attr-defined] + + def get_scales(self, indices=None) -> ProxyArray: + return self._view._get_world_scales_impl(indices) # type: ignore[attr-defined] + + +class _OvPhysxLocalSpaceWriter(FrameViewLocalSpaceWriter): + """OvPhysX local-space writer: pass-through to backend ``_apply_*`` hooks.""" + + def set_poses(self, positions=None, orientations=None, indices=None) -> None: + self._view._apply_local_pose_write(positions, orientations, indices) # type: ignore[attr-defined] + + def set_scales(self, scales, indices=None) -> None: + self._view._apply_local_scale_write(scales, indices) # type: ignore[attr-defined] + + def get_poses(self, indices=None) -> tuple[ProxyArray, ProxyArray]: + return self._view._get_local_poses_impl(indices) # type: ignore[attr-defined] + + def get_scales(self, indices=None) -> ProxyArray: + return self._view._get_local_scales_impl(indices) # type: ignore[attr-defined] diff --git a/source/isaaclab_physx/changelog.d/fabric-local-poses.rst b/source/isaaclab_physx/changelog.d/fabric-local-poses.rst new file mode 100644 index 000000000000..94c48cf83690 --- /dev/null +++ b/source/isaaclab_physx/changelog.d/fabric-local-poses.rst @@ -0,0 +1,36 @@ +Added +^^^^^ + +* Added :func:`~isaaclab.utils.warp.fabric.decompose_indexed_fabric_transforms` + and :func:`~isaaclab.utils.warp.fabric.compose_indexed_fabric_transforms` + Warp kernels. They mirror the existing + ``decompose_fabric_transformation_matrix_to_warp_arrays`` / + ``compose_fabric_transformation_matrix_from_warp_arrays`` kernels but + operate on :class:`wp.indexedfabricarray`, so the view-to-fabric mapping + is baked into the array and the kernel just dereferences + ``ifa[view_index]`` instead of taking a separate ``mapping`` argument. + +* Added :func:`~isaaclab.utils.warp.fabric.update_indexed_local_matrix_from_world` + and :func:`~isaaclab.utils.warp.fabric.update_indexed_world_matrix_from_local` + Warp kernels that propagate ``local = world * inv(parent)`` and + ``world = local * parent`` directly on Fabric storage matrices. + +* Added Fabric-accelerated local-pose read/write paths to + :class:`~isaaclab_physx.sim.views.FabricFrameView`. Local-pose + operations now use :class:`wp.indexedfabricarray` to read and write + ``omni:fabric:localMatrix`` directly on the GPU, propagating between + parent world matrices and child local/world matrices via Warp kernels + without round-tripping through USD. + +* Added topology-change recovery via automatic ``PrepareForReuse`` detection + and per-selection index rebuild. + +Deprecated +^^^^^^^^^^ + +* Deprecated ``get_scales`` / ``set_scales`` on ``FabricFrameView``. For + reads, use the explicit ``get_local_scales`` (operates on + ``localMatrix``) or ``get_world_scales`` (composed world-space scale). + For writes, use the writer scope's ``set_scales``. The deprecated + methods still work but emit a ``DeprecationWarning``; ``FabricFrameView`` + defaults to world (preserving prior behavior). diff --git a/source/isaaclab_physx/changelog.d/xform-space-writer.rst b/source/isaaclab_physx/changelog.d/xform-space-writer.rst new file mode 100644 index 000000000000..c87325bf69a5 --- /dev/null +++ b/source/isaaclab_physx/changelog.d/xform-space-writer.rst @@ -0,0 +1,36 @@ +Changed +^^^^^^^ + +* :class:`~isaaclab_physx.sim.views.FabricFrameView` now writes Fabric + ``omni:fabric:worldMatrix`` and ``omni:fabric:localMatrix`` through the + new context-managed + :class:`~isaaclab.sim.views.FrameViewSpaceWriterBase` scope. Each scope: + + - eagerly writes both the primary matrix (world or local, per the + chosen space) and derives the opposite-space matrix in a single Warp + kernel on ``__exit__``; + - calls ``wp.synchronize()`` once on ``__exit__``; + - pauses :meth:`IFabricHierarchy.track_local_xform_changes` and + :meth:`track_world_xform_changes` while the scope is active and + restores their prior state on exit, so Fabric Hierarchy's + ``update_world_xforms()`` on the next tick has no recorded changes + to replay for these prims. The Fabric Scene Delegate (FSD) reads + ``omni:fabric:worldMatrix`` from Fabric storage directly and + observes the writes. + - runs the opposite-space derive + ``wp.synchronize()`` on exit even + when the scope unwinds via exception (including ``KeyboardInterrupt`` + in interactive notebooks), as a best-effort to keep ``worldMatrix`` + and ``localMatrix`` mutually consistent prim-by-prim. The partial + write itself is not rolled back -- callers needing transactional + semantics should snapshot the matrices themselves before entering + the scope. + + Two persistent selections back the two access modes: ``_sel_ro`` + (``worldMatrix=RO, localMatrix=RO``, steady state) and ``_sel_rw`` + (``worldMatrix=RW, localMatrix=RW``, used inside a writer scope). + Both are built once during ``_initialize_fabric`` and kept for the + view's lifetime; the writer flips a single ``_is_rw`` flag on + enter/exit and neither selection is rebuilt on the flip. The RO + steady state tells Fabric Hierarchy's next ``update_world_xforms()`` + tick that no attribute is user-authored, so it leaves the pair + alone. 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 a0e79bae7fd1..068d4c5b875f 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 @@ -12,12 +12,12 @@ import torch import warp as wp -from pxr import Usd +from pxr import Gf, Usd, UsdGeom -import isaaclab.sim as sim_utils from isaaclab.app.settings_manager import SettingsManager from isaaclab.sim.views.base_frame_view import BaseFrameView from isaaclab.sim.views.usd_frame_view import UsdFrameView +from isaaclab.sim.views.xform_space_writer import FrameViewLocalSpaceWriter, FrameViewWorldSpaceWriter from isaaclab.utils.warp import ProxyArray from isaaclab.utils.warp import fabric as fabric_utils @@ -44,25 +44,125 @@ class FabricFrameView(BaseFrameView): """FrameView with Fabric GPU acceleration for the PhysX backend. Uses composition: holds a :class:`UsdFrameView` internally for USD - fallback and non-accelerated operations (local poses, visibility, scales - when Fabric is disabled). - - When Fabric is enabled, world-pose and scale operations use Warp kernels - operating on ``omni:fabric:worldMatrix``. Fabric acceleration runs on - the same CUDA device the view was constructed with — ``cuda:0``, - ``cuda:1``, or any other available CUDA index — so this view is safe - to use from distributed-training workers pinned to non-primary GPUs. - All other operations delegate to the internal USD view. + fallback and non-accelerated operations (visibility, and all pose/scale + operations when Fabric is disabled). - After every Fabric write (``set_world_poses``, ``set_scales``), - :meth:`PrepareForReuse` is called on the ``PrimSelection`` to notify - the FSD renderer that Fabric data has changed and to detect topology - changes that require rebuilding internal mappings. Read operations - do not call PrepareForReuse to avoid unnecessary renderer invalidation. + When Fabric is enabled, world-pose, local-pose, and scale operations run + on the GPU via Warp kernels that read and write + ``omni:fabric:worldMatrix`` and ``omni:fabric:localMatrix`` directly. + All other operations delegate to the internal USD view. - Pose getters return :class:`~isaaclab.utils.warp.ProxyArray`. Setters accept ``wp.array``. + All writes go through the writer-scope API + (:meth:`xform_world_space_writer` / :meth:`xform_local_space_writer`, + recommended) or the + convenience :meth:`set_world_poses` / :meth:`set_local_poses` / etc. helpers + inherited from :class:`BaseFrameView`. + + Behavior (Fabric path): + + * **Leaf-prim assumption.** This view manages a flat set of sibling prims + (e.g. all cameras under ``/World/Env_*/Camera``). It does NOT propagate + transforms to child prims. If a managed prim has children whose world + matrices depend on the parent, those children must be updated via a + separate view, a physics step, or ``IFabricHierarchy.update_world_xforms``. + * **No write-back to USD.** Fabric writes update only + ``omni:fabric:worldMatrix`` / ``omni:fabric:localMatrix``; the prim's + USD ``xformOp:*`` attributes are unchanged. Downstream consumers that + read the prim's USD attributes after a Fabric write will see stale + values until the next USD-side sync. + * **Eager dual-write inside a writer scope (no dirty tracking).** + When a writer scope is open, all writes go to the primary attribute + (``worldMatrix`` for the world writer, ``localMatrix`` for the local + writer). On scope exit, a single Warp kernel derives the opposite + attribute and a single ``wp.synchronize()`` runs. After the scope + exits, both Fabric matrices are self-consistent. Getters launch + their own decompose kernel and ``wp.synchronize()`` before returning, + so a returned :class:`ProxyArray` is always immediately readable from + either GPU or host code (no caller-side sync required). + + The opposite-space derive runs even when the scope unwinds via + exception (including ``KeyboardInterrupt`` in interactive notebooks), + as a best-effort to keep ``worldMatrix`` and ``localMatrix`` + mutually consistent on whatever partial-write state Fabric holds. + The partial write itself is not rolled back -- if you need + transactional all-or-nothing semantics, snapshot the matrices + yourself before entering the scope. + * **Fabric Hierarchy listeners are paused while a writer scope is + active.** On enter, the writer calls + :meth:`IFabricHierarchy.track_local_xform_changes(False)` / + :meth:`track_world_xform_changes(False)` (saving the prior state). + Fabric itself is just a flat attribute store; the plugin that + keeps ``omni:fabric:worldMatrix`` and ``omni:fabric:localMatrix`` + mutually consistent across the prim hierarchy is + :class:`usdrt.hierarchy.IFabricHierarchy` (a.k.a. Fabric + Hierarchy). Its change tracking is pull-based: a per-attribute + listener records writes into a private changelog, and the plugin + drains and processes that changelog on the next call to + ``IFabricHierarchy::update_world_xforms()`` (typically from the + render path). "Tracking off" just stops the listener from + recording new entries -- writes still land in Fabric storage; + they are simply invisible to the next ``update_world_xforms()`` + call. + + That is exactly what we want. Inside the scope we write one + space (world or local) and, at scope exit, derive the other in a + single batched kernel so both matrices are mutually consistent. + If tracking were left on, our writes would be queued in the + changelog and the next ``update_world_xforms()`` tick would + process them -- choosing a canonical direction (e.g. "user + authored local, recompute world from it") and potentially + overwriting one half of our just-consistent pair. With tracking + paused for the duration of the scope, the changelog stays empty + for these prims and the next tick is a no-op for them. + + ``__exit__`` restores the prior tracking state (so we do not + re-enable listeners the caller had previously paused). The + Fabric Scene Delegate (FSD) reads ``omni:fabric:worldMatrix`` + directly from Fabric storage on the render path; it observes our + final writes unchanged. + + Note: the scope is synchronous Python code, so no simulation step + and no render tick can run while it is open -- callers must not + advance the simulation from inside the scope (see + :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. + + Pose getters return :class:`~isaaclab.utils.warp.ProxyArray`; the + convenience :meth:`set_world_poses` / :meth:`set_local_poses` helpers accept + :class:`wp.array`. Inside a writer scope, the writer's + :meth:`~FrameViewSpaceWriterBase.set_poses` / :meth:`~FrameViewSpaceWriterBase.set_scales` + accept :class:`wp.array`. """ + _WORLD_MATRIX_NAME = "omni:fabric:worldMatrix" + _LOCAL_MATRIX_NAME = "omni:fabric:localMatrix" + def __init__( self, prim_path: str, @@ -76,7 +176,7 @@ def __init__( Args: prim_path: USD prim-path pattern to match. device: Device for Warp arrays. Either ``"cpu"`` or any CUDA - device string (``"cuda:0"``, ``"cuda:1"``, …); Fabric + device string (``"cuda:0"``, ``"cuda:1"``, ...); Fabric acceleration is supported on every CUDA index. validate_xform_ops: Whether to validate prim xform-ops. stage: USD stage; defaults to the current sim context's stage. @@ -90,18 +190,40 @@ def __init__( settings = SettingsManager.instance() self._use_fabric = bool(settings.get("/physics/fabricEnabled", False)) - # TODO(pv): Misleading abstraction — FabricFrameView can fall back to USD internally; + + # TODO(pv): Misleading abstraction -- FabricFrameView can fall back to USD internally; # the concrete class should be determined by the factory instead. (PR #5673 pv/fabric-view-no-fallback) - # TODO(pv): Fuse set_world_poses/set_scales into single kernel launch (PR #5674 pv/fabric-fused-compose) self._fabric_initialized = False - self._fabric_usd_sync_done = False - self._fabric_selection = None - self._fabric_to_view: wp.array | None = None - self._view_to_fabric: wp.array | None = None - self._default_view_indices: wp.array | None = None + self._stage = None self._fabric_hierarchy = None - self._view_index_attr = f"isaaclab:view_index:{abs(hash(self))}" + + # 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. + self._sel_ro = None + self._sel_rw = None + self._is_rw: bool = False + + # View-side indices array (shared across both bundles). + 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 + + # Sentinel passed to compose/decompose kernels for unused slots. + self._fabric_empty_2d_array_sentinel: wp.array | None = None # ------------------------------------------------------------------ # Delegated properties @@ -135,59 +257,29 @@ def set_visibility(self, visibility, indices=None): self._usd_view.set_visibility(visibility, indices) # ------------------------------------------------------------------ - # World poses — Fabric-accelerated or USD fallback + # Writer factory hooks # ------------------------------------------------------------------ - def set_world_poses(self, positions=None, orientations=None, indices=None): + def _make_world_space_writer(self) -> FrameViewWorldSpaceWriter: if not self._use_fabric: - self._usd_view.set_world_poses(positions, orientations, indices) - return + return _FabricFallbackWorldWriter(self) + return _FabricWorldSpaceWriter(self) - if not self._fabric_initialized: - self._initialize_fabric() - - self._prepare_for_reuse() - - indices_wp = self._resolve_indices_wp(indices) - count = indices_wp.shape[0] - - dummy = wp.zeros((0, 3), dtype=wp.float32, device=self._device) - positions_wp = _to_float32_2d(positions) if positions is not None else dummy - orientations_wp = ( - _to_float32_2d(orientations) - if orientations is not None - else wp.zeros((0, 4), dtype=wp.float32, device=self._device) - ) - - wp.launch( - kernel=fabric_utils.compose_fabric_transformation_matrix_from_warp_arrays, - dim=count, - inputs=[ - self._fabric_world_matrices, - positions_wp, - orientations_wp, - dummy, - False, - False, - False, - indices_wp, - self._view_to_fabric, - ], - device=self._fabric_device, - ) - wp.synchronize() + def _make_local_space_writer(self) -> FrameViewLocalSpaceWriter: + if not self._use_fabric: + return _FabricFallbackLocalWriter(self) + return _FabricLocalSpaceWriter(self) - self._fabric_hierarchy.update_world_xforms() - self._fabric_usd_sync_done = True + # ------------------------------------------------------------------ + # Getter hooks -- read directly from Fabric (no lazy sync) + # ------------------------------------------------------------------ - def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: + 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(indices) + return self._usd_view._get_world_poses_impl(indices) if not self._fabric_initialized: self._initialize_fabric() - if not self._fabric_usd_sync_done: - self._sync_fabric_from_usd_once() indices_wp = self._resolve_indices_wp(indices) count = indices_wp.shape[0] @@ -201,85 +293,85 @@ def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, orientations_wp = wp.zeros((count, 4), dtype=wp.float32, device=self._device) wp.launch( - kernel=fabric_utils.decompose_fabric_transformation_matrix_to_warp_arrays, + kernel=fabric_utils.decompose_indexed_fabric_transforms, dim=count, inputs=[ - self._fabric_world_matrices, + self._get_world_ifa(), positions_wp, orientations_wp, - self._fabric_dummy_buffer, + self._fabric_empty_2d_array_sentinel, indices_wp, - self._view_to_fabric, ], - device=self._fabric_device, + device=self._device, ) + # Sync before returning regardless of caching path: the cached buffers + # are reused (an in-flight kernel from a prior call must finish before + # the new write is observable) and the fresh-allocation path must also + # complete before the caller can read the returned ProxyArray via any + # host-visible accessor without observing zeros. + wp.synchronize() if use_cached: - wp.synchronize() return self._fabric_positions_ta, self._fabric_orientations_ta return ProxyArray(positions_wp), ProxyArray(orientations_wp) - # ------------------------------------------------------------------ - # Local poses — USD fallback (Fabric only accelerates world poses) - # ------------------------------------------------------------------ - - def set_local_poses(self, translations=None, orientations=None, indices=None): - self._usd_view.set_local_poses(translations, orientations, indices) - - def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: - return self._usd_view.get_local_poses(indices) - - # ------------------------------------------------------------------ - # Scales — Fabric-accelerated or USD fallback - # ------------------------------------------------------------------ - - def set_scales(self, scales, indices=None): + def _get_local_poses_impl(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: if not self._use_fabric: - self._usd_view.set_scales(scales, indices) - return + return self._usd_view._get_local_poses_impl(indices) if not self._fabric_initialized: self._initialize_fabric() - self._prepare_for_reuse() - indices_wp = self._resolve_indices_wp(indices) count = indices_wp.shape[0] - dummy3 = wp.zeros((0, 3), dtype=wp.float32, device=self._device) - dummy4 = wp.zeros((0, 4), dtype=wp.float32, device=self._device) - scales_wp = _to_float32_2d(scales) + use_cached = indices is None or indices == slice(None) + if use_cached: + translations_wp = self._fabric_local_translations_buf + orientations_wp = self._fabric_local_orientations_buf + else: + translations_wp = wp.zeros((count, 3), dtype=wp.float32, device=self._device) + orientations_wp = wp.zeros((count, 4), dtype=wp.float32, device=self._device) wp.launch( - kernel=fabric_utils.compose_fabric_transformation_matrix_from_warp_arrays, + kernel=fabric_utils.decompose_indexed_fabric_transforms, dim=count, inputs=[ - self._fabric_world_matrices, - dummy3, - dummy4, - scales_wp, - False, - False, - False, + self._get_local_ifa(), + translations_wp, + orientations_wp, + self._fabric_empty_2d_array_sentinel, indices_wp, - self._view_to_fabric, ], - device=self._fabric_device, + device=self._device, ) + + # See note in _get_world_poses_impl: sync regardless of caching path. wp.synchronize() + if use_cached: + return self._fabric_local_translations_ta, self._fabric_local_orientations_ta + return ProxyArray(translations_wp), ProxyArray(orientations_wp) - self._fabric_hierarchy.update_world_xforms() - self._fabric_usd_sync_done = True + def _get_world_scales_impl(self, indices=None) -> ProxyArray: + if not self._use_fabric: + return self._usd_view._get_world_scales_impl(indices) + + if not self._fabric_initialized: + self._initialize_fabric() - def get_scales(self, indices=None): + return self._decompose_scales(self._get_world_ifa(), indices) + + def _get_local_scales_impl(self, indices=None) -> ProxyArray: if not self._use_fabric: - return self._usd_view.get_scales(indices) + return self._usd_view._get_local_scales_impl(indices) if not self._fabric_initialized: self._initialize_fabric() - if not self._fabric_usd_sync_done: - self._sync_fabric_from_usd_once() + return self._decompose_scales(self._get_local_ifa(), indices) + + def _decompose_scales(self, ro_array, indices) -> ProxyArray: + """Shared scale-decompose path for world / local getters.""" indices_wp = self._resolve_indices_wp(indices) count = indices_wp.shape[0] @@ -290,168 +382,598 @@ def get_scales(self, indices=None): scales_wp = wp.zeros((count, 3), dtype=wp.float32, device=self._device) wp.launch( - kernel=fabric_utils.decompose_fabric_transformation_matrix_to_warp_arrays, + kernel=fabric_utils.decompose_indexed_fabric_transforms, dim=count, inputs=[ - self._fabric_world_matrices, - self._fabric_dummy_buffer, - self._fabric_dummy_buffer, + ro_array, + self._fabric_empty_2d_array_sentinel, + self._fabric_empty_2d_array_sentinel, scales_wp, indices_wp, - self._view_to_fabric, ], - device=self._fabric_device, + device=self._device, ) + # See note in _get_world_poses_impl: sync regardless of caching path. + wp.synchronize() if use_cached: - wp.synchronize() - return scales_wp + return self._fabric_scales_ta + return ProxyArray(scales_wp) # ------------------------------------------------------------------ - # Internal — PrepareForReuse (renderer notification + topology tracking) + # Deprecated get_scales / set_scales hooks # ------------------------------------------------------------------ - def _prepare_for_reuse(self) -> None: - """Call PrepareForReuse on the PrimSelection to notify the renderer. + def _get_scales_impl(self, indices=None) -> ProxyArray: + """Fabric: get_scales returns world-space scales (legacy behavior).""" + return self._get_world_scales_impl(indices) - PrepareForReuse serves two purposes: + def _set_scales_impl(self, scales, indices=None) -> None: + """Fabric: set_scales writes world-space scales via a one-shot writer scope.""" + with self.xform_world_space_writer() as writer: + writer.set_scales(scales, indices) - 1. **Renderer notification**: Tells FSD/Storm that Fabric data has - been (or will be) modified, so the next rendered frame reflects - the updated transforms. - 2. **Topology change detection**: Returns True when Fabric's - internal memory layout changed (e.g., prims added/removed). - In that case, view-to-fabric index mappings and fabricarrays - must be rebuilt. - """ - if self._fabric_selection is None: - return + # ------------------------------------------------------------------ + # Internal -- helpers shared by writers + initialization + # ------------------------------------------------------------------ - topology_changed = self._fabric_selection.PrepareForReuse() - if topology_changed: - logger.debug("Fabric topology changed — rebuilding view-to-fabric index mapping.") - self._rebuild_fabric_arrays() + def _to_float32_2d_or_empty(self, data): + return self._fabric_empty_2d_array_sentinel if data is None else _to_float32_2d(data) - def _rebuild_fabric_arrays(self) -> None: - """Rebuild fabricarray and view↔fabric mappings after a topology change. + def _recompute_local_from_world_all(self) -> None: + """Derive ``localMatrix = inv(parent) * worldMatrix`` for every prim in the view. - Note: Only index mappings and fabricarrays are rebuilt. Position/orientation/scale - buffers are *not* resized because ``self.count`` is derived from the USD prim-path - pattern (via ``_usd_view.count``) and does not change when Fabric rearranges its - internal memory layout. The assertion below guards this invariant. + Called from :class:`_FabricWorldSpaceWriter` ``__exit__`` to keep the + (world, local) pair self-consistent after a world-space write. + Storage convention: see + :func:`isaaclab.utils.warp.fabric.update_indexed_local_matrix_from_world`. """ - assert self.count == self._default_view_indices.shape[0], ( - f"Prim count changed ({self.count} vs {self._default_view_indices.shape[0]}). " - "Fabric topology change added/removed tracked prims — full re-initialization required." + wp.launch( + kernel=fabric_utils.update_indexed_local_matrix_from_world, + dim=self.count, + inputs=[ + self._get_world_ifa(), + self._get_parent_world_ifa(), + self._get_local_ifa(), + self._view_indices, + ], + device=self._device, ) - self._view_to_fabric = wp.zeros((self.count,), dtype=wp.uint32, device=self._fabric_device) - self._fabric_to_view = wp.fabricarray(self._fabric_selection, self._view_index_attr) + def _recompute_world_from_local_all(self) -> None: + """Derive ``worldMatrix = parent * localMatrix`` for every prim in the view. + + Called from :class:`_FabricLocalSpaceWriter` ``__exit__`` and from + :meth:`_sync_fabric_from_usd_initial` after seeding local matrices. + Storage convention: see + :func:`isaaclab.utils.warp.fabric.update_indexed_world_matrix_from_local`. + """ wp.launch( - kernel=fabric_utils.set_view_to_fabric_array, - dim=self._fabric_to_view.shape[0], - inputs=[self._fabric_to_view, self._view_to_fabric], - device=self._fabric_device, + kernel=fabric_utils.update_indexed_world_matrix_from_local, + dim=self.count, + inputs=[ + self._get_local_ifa(), + self._get_parent_world_ifa(), + self._get_world_ifa(), + self._view_indices, + ], + device=self._device, ) - wp.synchronize() - self._fabric_world_matrices = wp.fabricarray(self._fabric_selection, "omni:fabric:worldMatrix") + # ------------------------------------------------------------------ + # Internal -- selection accessors with on-demand index rebuild + # ------------------------------------------------------------------ + + 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_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_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 _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, + ) + + 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, + ) + + # ------------------------------------------------------------------ + # Internal -- index computation + # ------------------------------------------------------------------ + + 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)) + + def _compute_parent_fabric_indices(self, selection) -> wp.array: + """View-side indices that map each managed prim's parent into ``selection``.""" + + 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 + + return self._compute_fabric_indices_for(selection, [parent_path(p) for p in self.prim_paths]) + + 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 _resolve_indices_wp(self, indices: wp.array | None) -> wp.array: + """Resolve view indices as a Warp uint32 array.""" + if indices is None or indices == slice(None): + 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 # ------------------------------------------------------------------ - # Internal — Fabric initialization + # Internal -- Fabric initialization # ------------------------------------------------------------------ def _initialize_fabric(self) -> None: - """Initialize Fabric batch infrastructure for GPU-accelerated pose queries.""" + """One-time Fabric setup: hierarchy handle, attribute population, selections, indexed arrays.""" import usdrt # noqa: PLC0415 from usdrt import Rt # noqa: PLC0415 - stage_id = sim_utils.get_current_stage_id() - fabric_stage = usdrt.Usd.Stage.Attach(stage_id) + from isaaclab.sim.utils import get_current_stage_id # noqa: PLC0415 - for i in range(self.count): - rt_prim = fabric_stage.GetPrimAtPath(self.prim_paths[i]) - rt_xformable = Rt.Xformable(rt_prim) + # Attach usdrt stage and create hierarchy handle. + stage_id = get_current_stage_id() + self._stage = usdrt.Usd.Stage.Attach(stage_id) + fabric_id = self._stage.GetFabricId() + self._fabric_id = fabric_id.id + self._fabric_hierarchy = usdrt.hierarchy.IFabricHierarchy().get_fabric_hierarchy( + fabric_id, self._stage.GetStageIdAsStageId() + ) - has_attr = ( - rt_xformable.HasFabricHierarchyWorldMatrixAttr() - if hasattr(rt_xformable, "HasFabricHierarchyWorldMatrixAttr") - else False - ) - if not has_attr: + # 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) + rt_prim = self._stage.GetPrimAtPath(path) + if not rt_prim.IsValid(): + continue + rt_xformable = Rt.Xformable(rt_prim) rt_xformable.CreateFabricHierarchyWorldMatrixAttr() + rt_xformable.CreateFabricHierarchyLocalMatrixAttr() + rt_xformable.SetLocalXformFromUsd() + rt_xformable.SetWorldXformFromUsd() + + # 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``. + matrix = usdrt.Sdf.ValueTypeNames.Matrix4d + ro = usdrt.Usd.Access.Read + rw = usdrt.Usd.Access.ReadWrite + 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._view_indices = wp.array(list(range(self.count)), dtype=wp.uint32, device=self._device) + self._rebuild_ro_arrays() + self._rebuild_rw_arrays() + + # Pre-allocated reusable output buffers (world + local + scales). + self._fabric_positions_buf = wp.zeros((self.count, 3), dtype=wp.float32, device=self._device) + self._fabric_orientations_buf = wp.zeros((self.count, 4), dtype=wp.float32, device=self._device) + self._fabric_scales_buf = wp.zeros((self.count, 3), dtype=wp.float32, device=self._device) + self._fabric_local_translations_buf = wp.zeros((self.count, 3), dtype=wp.float32, device=self._device) + self._fabric_local_orientations_buf = wp.zeros((self.count, 4), dtype=wp.float32, device=self._device) + self._fabric_empty_2d_array_sentinel = wp.zeros((0, 0), dtype=wp.float32, device=self._device) - rt_xformable.SetWorldXformFromUsd() + self._fabric_positions_ta = ProxyArray(self._fabric_positions_buf) + self._fabric_orientations_ta = ProxyArray(self._fabric_orientations_buf) + self._fabric_scales_ta = ProxyArray(self._fabric_scales_buf) + self._fabric_local_translations_ta = ProxyArray(self._fabric_local_translations_buf) + self._fabric_local_orientations_ta = ProxyArray(self._fabric_local_orientations_buf) - rt_prim.CreateAttribute(self._view_index_attr, usdrt.Sdf.ValueTypeNames.UInt, custom=True) - rt_prim.GetAttribute(self._view_index_attr).Set(i) + self._fabric_initialized = True - self._fabric_hierarchy = usdrt.hierarchy.IFabricHierarchy().get_fabric_hierarchy( - fabric_stage.GetFabricId(), fabric_stage.GetStageIdAsStageId() + # 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. + self._is_rw = True + try: + self._sync_fabric_from_usd_initial() + finally: + self._is_rw = False + + def _sync_fabric_from_usd_initial(self) -> None: + """Populate Fabric world+local matrices for children and parents from USD. + + Performed once during ``_initialize_fabric``. Without this step Fabric's + matrices are identity for stages that haven't been rendered yet, and our + getters (which read from Fabric) would return wrong values. + """ + # --- Children: compose child localMatrix from USD-authored local transforms. + scales_wp = _to_float32_2d(self._usd_view.get_local_scales().warp) + local_pos_ta, local_ori_ta = self._usd_view.get_local_poses() + wp.launch( + kernel=fabric_utils.compose_indexed_fabric_transforms, + dim=self.count, + inputs=[ + self._local_ifa_rw, # explicit RW: 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), + False, + False, + False, + self._view_indices, + ], + device=self._device, ) - self._fabric_hierarchy.update_world_xforms() - self._default_view_indices = wp.zeros((self.count,), dtype=wp.uint32, device=self._device) + # --- Parents (one entry per unique parent path) --- + unique_parent_paths = list(dict.fromkeys(p.rsplit("/", 1)[0] for p in self.prim_paths)) + if unique_parent_paths: + from isaaclab.sim.utils import get_current_stage # noqa: PLC0415 + + usd_stage = get_current_stage() + xform_cache = UsdGeom.XformCache(Usd.TimeCode.Default()) + world_pos_rows: list[list[float]] = [] + world_ori_rows: list[list[float]] = [] + world_scale_rows: list[list[float]] = [] + decomposer = Gf.Transform() + warned_shear = False + for path in unique_parent_paths: + prim = usd_stage.GetPrimAtPath(path) + tf = xform_cache.GetLocalToWorldTransform(prim) + decomposer.SetMatrix(tf) + s = decomposer.GetScale() + if not warned_shear: + row0 = Gf.Vec3d(tf[0][0], tf[0][1], tf[0][2]).GetNormalized() + row1 = Gf.Vec3d(tf[1][0], tf[1][1], tf[1][2]).GetNormalized() + row2 = Gf.Vec3d(tf[2][0], tf[2][1], tf[2][2]).GetNormalized() + if ( + abs(Gf.Dot(row0, row1)) > 1e-3 + or abs(Gf.Dot(row0, row2)) > 1e-3 + or abs(Gf.Dot(row1, row2)) > 1e-3 + ): + warned_shear = True + logger.warning( + "FabricFrameView: parent prim '%s' has a sheared/skewed world " + "transform. TRS decomposition (used by scale getters and world<->local " + "propagation) does not support shear -- extracted scales and rotations " + "will be approximate. Avoid shear in parent transforms for correct results.", + path, + ) + tf.Orthonormalize() + t = tf.ExtractTranslation() + q = tf.ExtractRotationQuat() + img, real = q.GetImaginary(), q.GetReal() + world_pos_rows.append([float(t[0]), float(t[1]), float(t[2])]) + world_ori_rows.append([float(img[0]), float(img[1]), float(img[2]), float(real)]) + world_scale_rows.append([float(s[0]), float(s[1]), float(s[2])]) + parent_view_indices = wp.array(list(range(len(unique_parent_paths))), dtype=wp.uint32, device=self._device) + 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) + 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), + ) + wp.launch( + kernel=fabric_utils.compose_indexed_fabric_transforms, + dim=len(unique_parent_paths), + inputs=[ + parent_world_rw, + parent_pos_wp, + parent_ori_wp, + parent_scale_wp, + False, + False, + False, + parent_view_indices, + ], + device=self._device, + ) + wp.synchronize() + + # After seeding local matrices from USD, recompute world matrices so + # the view starts with consistent state. + 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 +# ---------------------------------------------------------------------- + + +class _FabricWriterMixin: + """Common ``__enter__`` / ``__exit__`` for the Fabric world / local writers. + + 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). + + 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 + steady-state reads) and restores hierarchy-tracking state. + + **Exception safety.** If the scope unwinds because of an exception + (including ``KeyboardInterrupt`` from an interactive notebook), the + opposite-space derive still runs so that ``worldMatrix`` and + ``localMatrix`` are mutually consistent prim-by-prim on whatever + partial-write state Fabric currently holds. The partial write itself + is **not** rolled back -- some prims may carry the new value and the + rest the pre-scope value -- so callers needing transactional + all-or-nothing semantics should snapshot matrices themselves before + entering the scope. If the recovery launch itself fails (typically + because the original exception came from a poisoned CUDA stream), the + failure is logged and the original exception propagates; the view + should then be recreated. + """ + + def _enter_impl(self) -> None: + view: FabricFrameView = self._view # type: ignore[assignment] + if not view._fabric_initialized: + view._initialize_fabric() + self._wrote_anything = False + h = view._fabric_hierarchy + self._was_tracking_local = h.tracking_local_xform_changes + self._was_tracking_world = h.tracking_world_xform_changes + if self._was_tracking_local: + h.track_local_xform_changes(False) + if self._was_tracking_world: + h.track_world_xform_changes(False) + view._is_rw = True + + def _exit_impl(self, exc_type, exc_val, exc_tb) -> None: + view: FabricFrameView = self._view # type: ignore[assignment] + try: + if self._wrote_anything: + try: + self._derive_opposite() + wp.synchronize() + except Exception as recovery_exc: + # Recovery itself failed (e.g. original exception came + # from a device error and the CUDA stream is poisoned). + # If we got here on the happy path, re-raise. If we are + # already unwinding an exception, log and let the original + # propagate -- masking it would hide the actual root cause. + if exc_type is None: + raise + logger.error( + "FabricFrameView writer scope: best-effort opposite-space sync " + "failed during exception handling: %s. World/local matrices may " + "be inconsistent prim-by-prim; recreate the view to recover.", + recovery_exc, + ) + finally: + # Flip back to RO before restoring hierarchy tracking so any + # subsequent updateWorldXforms tick sees a fully-RO selection. + view._is_rw = False + h = view._fabric_hierarchy + if self._was_tracking_world: + h.track_world_xform_changes(True) + if self._was_tracking_local: + h.track_local_xform_changes(True) + + def _derive_opposite(self) -> None: + raise NotImplementedError + + +class _FabricWorldSpaceWriter(_FabricWriterMixin, FrameViewWorldSpaceWriter): + """World-space writer for :class:`FabricFrameView`. + + Writes flow through ``_world_ifa_rw`` (the RW-bundle worldMatrix array); + on exit ``localMatrix`` is derived from the just-written ``worldMatrix`` + via :func:`update_indexed_local_matrix_from_world`. + """ + + def _derive_opposite(self) -> None: + self._view._recompute_local_from_world_all() # type: ignore[attr-defined] + + def set_poses(self, positions=None, orientations=None, indices=None) -> None: + view: FabricFrameView = self._view # type: ignore[assignment] + indices_wp = view._resolve_indices_wp(indices) + positions_wp = view._to_float32_2d_or_empty(positions) + orientations_wp = view._to_float32_2d_or_empty(orientations) wp.launch( - kernel=fabric_utils.arange_k, dim=self.count, inputs=[self._default_view_indices], device=self._device + kernel=fabric_utils.compose_indexed_fabric_transforms, + dim=indices_wp.shape[0], + inputs=[ + view._get_world_ifa(), + positions_wp, + orientations_wp, + view._fabric_empty_2d_array_sentinel, + False, + False, + False, + indices_wp, + ], + device=view._device, ) - wp.synchronize() + self._wrote_anything = True - self._fabric_selection = fabric_stage.SelectPrims( - require_attrs=[ - (usdrt.Sdf.ValueTypeNames.UInt, self._view_index_attr, usdrt.Usd.Access.Read), - (usdrt.Sdf.ValueTypeNames.Matrix4d, "omni:fabric:worldMatrix", usdrt.Usd.Access.ReadWrite), + def set_scales(self, scales, indices=None) -> None: + view: FabricFrameView = self._view # type: ignore[assignment] + indices_wp = view._resolve_indices_wp(indices) + scales_wp = view._to_float32_2d_or_empty(scales) + wp.launch( + kernel=fabric_utils.compose_indexed_fabric_transforms, + dim=indices_wp.shape[0], + inputs=[ + view._get_world_ifa(), + view._fabric_empty_2d_array_sentinel, + view._fabric_empty_2d_array_sentinel, + scales_wp, + False, + False, + False, + indices_wp, ], - device=self._device, + device=view._device, ) + self._wrote_anything = True + + def get_poses(self, indices=None) -> tuple[ProxyArray, ProxyArray]: + return self._view._get_world_poses_impl(indices) # type: ignore[attr-defined] - self._view_to_fabric = wp.zeros((self.count,), dtype=wp.uint32, device=self._device) - self._fabric_to_view = wp.fabricarray(self._fabric_selection, self._view_index_attr) + def get_scales(self, indices=None) -> ProxyArray: + return self._view._get_world_scales_impl(indices) # type: ignore[attr-defined] + +class _FabricLocalSpaceWriter(_FabricWriterMixin, FrameViewLocalSpaceWriter): + """Local-space writer for :class:`FabricFrameView`. + + Writes flow through ``_local_ifa_rw`` (the RW-bundle localMatrix array); + on exit ``worldMatrix`` is derived from the just-written ``localMatrix`` + via :func:`update_indexed_world_matrix_from_local`. + """ + + def _derive_opposite(self) -> None: + self._view._recompute_world_from_local_all() # type: ignore[attr-defined] + + def set_poses(self, positions=None, orientations=None, indices=None) -> None: + view: FabricFrameView = self._view # type: ignore[assignment] + indices_wp = view._resolve_indices_wp(indices) + translations_wp = view._to_float32_2d_or_empty(positions) + orientations_wp = view._to_float32_2d_or_empty(orientations) wp.launch( - kernel=fabric_utils.set_view_to_fabric_array, - dim=self._fabric_to_view.shape[0], - inputs=[self._fabric_to_view, self._view_to_fabric], - device=self._device, + kernel=fabric_utils.compose_indexed_fabric_transforms, + dim=indices_wp.shape[0], + inputs=[ + view._get_local_ifa(), + translations_wp, + orientations_wp, + view._fabric_empty_2d_array_sentinel, + False, + False, + False, + indices_wp, + ], + device=view._device, ) - wp.synchronize() + self._wrote_anything = True - self._fabric_positions_buf = wp.zeros((self.count, 3), dtype=wp.float32, device=self._device) - self._fabric_orientations_buf = wp.zeros((self.count, 4), dtype=wp.float32, device=self._device) - self._fabric_positions_ta = ProxyArray(self._fabric_positions_buf) - self._fabric_orientations_ta = ProxyArray(self._fabric_orientations_buf) - self._fabric_scales_buf = wp.zeros((self.count, 3), dtype=wp.float32, device=self._device) - self._fabric_dummy_buffer = wp.zeros((0, 3), dtype=wp.float32, device=self._device) - self._fabric_world_matrices = wp.fabricarray(self._fabric_selection, "omni:fabric:worldMatrix") - self._fabric_stage = fabric_stage - self._fabric_device = self._device + def set_scales(self, scales, indices=None) -> None: + view: FabricFrameView = self._view # type: ignore[assignment] + indices_wp = view._resolve_indices_wp(indices) + scales_wp = view._to_float32_2d_or_empty(scales) + wp.launch( + kernel=fabric_utils.compose_indexed_fabric_transforms, + dim=indices_wp.shape[0], + inputs=[ + view._get_local_ifa(), + view._fabric_empty_2d_array_sentinel, + view._fabric_empty_2d_array_sentinel, + scales_wp, + False, + False, + False, + indices_wp, + ], + device=view._device, + ) + self._wrote_anything = True - self._fabric_initialized = True - self._fabric_usd_sync_done = False + def get_poses(self, indices=None) -> tuple[ProxyArray, ProxyArray]: + return self._view._get_local_poses_impl(indices) # type: ignore[attr-defined] - def _sync_fabric_from_usd_once(self) -> None: - """Sync Fabric world matrices from USD once, on the first read. + def get_scales(self, indices=None) -> ProxyArray: + return self._view._get_local_scales_impl(indices) # type: ignore[attr-defined] - ``set_world_poses`` and ``set_scales`` each set ``_fabric_usd_sync_done`` - themselves, so no explicit flag assignment is needed here. - """ - if not self._fabric_initialized: - self._initialize_fabric() - positions_usd_ta, orientations_usd_ta = self._usd_view.get_world_poses() - positions_usd = positions_usd_ta.warp - orientations_usd = orientations_usd_ta.warp - scales_usd = self._usd_view.get_scales() +class _FabricFallbackWorldWriter(FrameViewWorldSpaceWriter): + """Fallback world-space writer used when Fabric is disabled. - self.set_world_poses(positions_usd, orientations_usd) - self.set_scales(scales_usd) + Delegates set/get calls to the internal :class:`UsdFrameView`'s backend + hooks directly. No batching, no listener pausing -- there's no Fabric to + confuse. + """ - def _resolve_indices_wp(self, indices: wp.array | None) -> wp.array: - """Resolve view indices as a Warp uint32 array.""" - if indices is None or indices == slice(None): - if self._default_view_indices is None: - raise RuntimeError("Fabric indices are not initialized.") - return self._default_view_indices - if indices.dtype != wp.uint32: - return wp.array(indices.numpy().astype("uint32"), dtype=wp.uint32, device=self._device) - return indices + def set_poses(self, positions=None, orientations=None, indices=None) -> None: + self._view._usd_view._apply_world_pose_write(positions, orientations, indices) # type: ignore[attr-defined] + + def set_scales(self, scales, indices=None) -> None: + self._view._usd_view._apply_world_scale_write(scales, indices) # type: ignore[attr-defined] + + def get_poses(self, indices=None) -> tuple[ProxyArray, ProxyArray]: + return self._view._usd_view._get_world_poses_impl(indices) # type: ignore[attr-defined] + + def get_scales(self, indices=None) -> ProxyArray: + return self._view._usd_view._get_world_scales_impl(indices) # type: ignore[attr-defined] + + +class _FabricFallbackLocalWriter(FrameViewLocalSpaceWriter): + """Fallback local-space writer used when Fabric is disabled.""" + + def set_poses(self, positions=None, orientations=None, indices=None) -> None: + self._view._usd_view._apply_local_pose_write(positions, orientations, indices) # type: ignore[attr-defined] + + def set_scales(self, scales, indices=None) -> None: + self._view._usd_view._apply_local_scale_write(scales, indices) # type: ignore[attr-defined] + + def get_poses(self, indices=None) -> tuple[ProxyArray, ProxyArray]: + return self._view._usd_view._get_local_poses_impl(indices) # type: ignore[attr-defined] + + def get_scales(self, indices=None) -> ProxyArray: + return self._view._usd_view._get_local_scales_impl(indices) # type: ignore[attr-defined] 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 3cfe70095fd3..c023ab40ba86 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 @@ -24,7 +24,7 @@ import torch # noqa: E402 import warp as wp # noqa: E402 from frame_view_contract_utils import * # noqa: F401, F403, E402 -from frame_view_contract_utils import CHILD_OFFSET, ViewBundle, test_set_world_updates_local # noqa: E402 +from frame_view_contract_utils import CHILD_OFFSET, ViewBundle # noqa: E402, F401 from isaaclab_physx.sim.views import FabricFrameView as FrameView # noqa: E402 from pxr import Gf, UsdGeom # noqa: E402 @@ -57,7 +57,7 @@ def _skip_if_unavailable(device: str): # a misconfigured multi-GPU runner is already caught there. Failing here would # only break the standard single-GPU CI runners that legitimately can't run # ``cuda:1+`` tests. - pytest.skip(f"{device} not available (device_count={n}) — multi-GPU test skipped") + pytest.skip(f"{device} not available (device_count={n}) -- multi-GPU test skipped") # ------------------------------------------------------------------ @@ -118,28 +118,11 @@ def factory(num_envs: int, device: str) -> ViewBundle: # ------------------------------------------------------------------ -# Override shared contract test with expected failure for Fabric. -# FabricFrameView.set_world_poses writes to Fabric worldMatrix only; the local -# pose (read via USD) does not reflect the change because there is no -# Fabric → USD writeback for local poses. This is tracked as Issue #5 -# (localMatrix: set_local_poses falls back to USD). +# Override: ensure the shared contract test runs without xfail now that +# get_local_poses computes local from Fabric world matrices. # ------------------------------------------------------------------ - - -@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) -@pytest.mark.xfail( - reason=( - "Issue #5: FabricFrameView.set_world_poses writes to Fabric worldMatrix only. " - "get_local_poses reads from stale USD because there is no Fabric→USD " - "writeback for local poses." - ), - strict=True, -) -def test_set_world_updates_local(device, view_factory): # noqa: F811 - """Override the shared test to mark it as expected failure.""" - from frame_view_contract_utils import test_set_world_updates_local as _impl # noqa: PLC0415 - - _impl(device, view_factory) +# (No override needed -- the shared test_set_world_updates_local from +# frame_view_contract_utils is imported via wildcard and will run as-is.) # ------------------------------------------------------------------ @@ -174,10 +157,11 @@ def test_fabric_set_world_does_not_write_back_to_usd(device, view_factory): usd_t_before = usd_tf_before.ExtractTranslation() orig_usd_pos = torch.tensor([float(usd_t_before[0]), float(usd_t_before[1]), float(usd_t_before[2])]) - # Write to Fabric — move to (99, 99, 99) + # Write to Fabric -- move to (99, 99, 99) new_pos = wp.zeros((1, 3), dtype=wp.float32, device=device) wp.launch(kernel=_fill_position, dim=1, inputs=[new_pos, 99.0, 99.0, 99.0], device=device) - view.set_world_poses(positions=new_pos) + with view.xform_world_space_writer() as w: + w.set_poses(positions=new_pos) # Verify Fabric has the new position fab_pos, _ = view.get_world_poses() @@ -187,7 +171,7 @@ def test_fabric_set_world_does_not_write_back_to_usd(device, view_factory): ) # Verify USD still has the ORIGINAL position (no writeback). Equality, not - # approximate — USD should literally not have moved, so any drift would + # approximate -- USD should literally not have moved, so any drift would # indicate a residual writeback path. xform_cache_after = UsdGeom.XformCache() usd_tf_after = xform_cache_after.GetLocalToWorldTransform(prim) @@ -200,55 +184,457 @@ def test_fabric_set_world_does_not_write_back_to_usd(device, view_factory): @pytest.mark.parametrize("device", ["cpu", "cuda:0"]) -def test_fabric_rebuild_after_topology_change(device, view_factory, monkeypatch): - """Forcing the topology-changed branch on a write triggers - :meth:`_rebuild_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 monkeypatch ``_prepare_for_reuse`` on the instance to always - take the rebuild branch and verify the view remains usable. +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. """ bundle = view_factory(2, device) view = bundle.view - # First write — initializes Fabric and binds _fabric_selection. + # First write -- initializes Fabric. initial = wp.zeros((2, 3), dtype=wp.float32, device=device) wp.launch(kernel=_fill_position, dim=2, inputs=[initial, 1.0, 2.0, 3.0], device=device) - view.set_world_poses(positions=initial) - - rebuild_calls = [] - real_rebuild = view._rebuild_fabric_arrays - - def spy_rebuild(): - rebuild_calls.append(True) - real_rebuild() + with view.xform_world_space_writer() as w: + w.set_poses(positions=initial) - def force_topology_changed(): - if view._fabric_selection is not None: - view._fabric_selection.PrepareForReuse() - spy_rebuild() + # 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() - monkeypatch.setattr(view, "_prepare_for_reuse", force_topology_changed) - - # Trigger another write — goes through the forced topology-change branch. + # Trigger another write through the rebuilt arrays. new = wp.zeros((2, 3), dtype=wp.float32, device=device) wp.launch(kernel=_fill_position, dim=2, inputs=[new, 4.0, 5.0, 6.0], device=device) - view.set_world_poses(positions=new) - - assert rebuild_calls, "Forced topology-change branch did not invoke _rebuild_fabric_arrays" + with view.xform_world_space_writer() as w: + w.set_poses(positions=new) - # Read back — proves the rebuilt _view_to_fabric and _fabric_world_matrices - # are still consistent. ret_pos, _ = view.get_world_poses() pos_torch = torch.as_tensor(ret_pos, device=device) expected = torch.tensor([[4.0, 5.0, 6.0], [4.0, 5.0, 6.0]], device=device) - assert torch.allclose(pos_torch, expected, atol=1e-7), f"Read after rebuild failed on {device}: {pos_torch}" + # 1e-5 ≈ 20 ULP at magnitudes ~4-6; absorbs float32 SRT compose/decompose drift. + assert torch.allclose(pos_torch, expected, atol=1e-5), f"Read after rebuild failed on {device}: {pos_torch}" + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_writer_scope_exception_recovers_state(device, view_factory): + """An exception raised inside a writer scope must still: + + 1. Restore ``IFabricHierarchy.track_*_xform_changes`` to their pre-scope state. + 2. Flip the view's ``_is_rw`` flag back to ``False``. + 3. Leave ``worldMatrix`` and ``localMatrix`` mutually consistent prim-by-prim + on whatever partial-write state Fabric currently holds (best-effort). + + Simulates an interactive-notebook scenario where a user-code exception + fires after a ``set_poses`` call but before the scope closes. + """ + bundle = view_factory(2, device) + view = bundle.view + view.get_world_poses() # ensure Fabric is initialized + + # Snapshot pre-scope tracking state. + h = view._fabric_hierarchy + was_tracking_local = h.tracking_local_xform_changes + was_tracking_world = h.tracking_world_xform_changes + + positions = wp.zeros((2, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=2, inputs=[positions, 7.0, 8.0, 9.0], device=device) + + with pytest.raises(RuntimeError, match="user-code failure"): + with view.xform_world_space_writer() as w: + w.set_poses(positions=positions) + raise RuntimeError("user-code failure") + + # Tracking state restored. + assert h.tracking_local_xform_changes == was_tracking_local + assert h.tracking_world_xform_changes == was_tracking_world + # _is_rw flipped back. + assert view._is_rw is False + # World/local mutually consistent: re-reading both spaces succeeds and the + # world positions reflect the partial write we made before the exception. + world_pos, _ = view.get_world_poses() + local_pos, _ = view.get_local_poses() + world_t = torch.as_tensor(world_pos, device=device) + local_t = torch.as_tensor(local_pos, device=device) + expected_world = torch.tensor([[7.0, 8.0, 9.0], [7.0, 8.0, 9.0]], device=device) + assert torch.allclose(world_t, expected_world, atol=1e-5), f"world not updated on {device}: {world_t}" + # Parents at PARENT_POS, so local = world - parent_world (parents identity-rotated). + expected_local = expected_world - torch.tensor(PARENT_POS, device=device) + assert torch.allclose(local_t, expected_local, atol=1e-5), ( + f"local not derived from partial world write on {device}: got {local_t}, expected {expected_local}" + ) + + # The view is still usable for further writer scopes after recovery. + next_pos = wp.zeros((2, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=2, inputs=[next_pos, 10.0, 11.0, 12.0], device=device) + with view.xform_world_space_writer() as w: + w.set_poses(positions=next_pos) + follow_up, _ = view.get_world_poses() + follow_up_t = torch.as_tensor(follow_up, device=device) + assert torch.allclose(follow_up_t, torch.tensor([[10.0, 11.0, 12.0]] * 2, device=device), atol=1e-5) + + +@pytest.mark.parametrize("device", ["cuda:0"]) +def test_prepare_for_reuse_detects_topology_change(device, view_factory): + """Each persistent ``PrimSelection`` exposes ``PrepareForReuse`` and returns a + bool. When the underlying Fabric topology is unchanged it returns False. + """ + bundle = view_factory(1, device) + view = bundle.view + view.get_world_poses() # trigger Fabric init + + assert view._sel_ro is not None, "RO selection not initialized" + assert view._sel_rw is not None, "RW selection not initialized" + for selection in (view._sel_ro, view._sel_rw): + result = selection.PrepareForReuse() + assert isinstance(result, bool), f"PrepareForReuse should return bool, got {type(result)}" + assert not result, "PrepareForReuse should return False when no topology change" + + +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]) + world_attr = rt_prim.GetAttribute(view._WORLD_MATRIX_NAME) + matrix = world_attr.Get() + translation = matrix.ExtractTranslation() + return torch.tensor( + [[float(translation[0]), float(translation[1]), float(translation[2])]], + dtype=torch.float32, + device=view._device, + ) + + +def _read_fabric_world_matrix_scale(view, prim_index=0): + """Read cached Fabric worldMatrix scale directly, without FrameView getter sync.""" + import usdrt # noqa: PLC0415 + + rt_prim = view._stage.GetPrimAtPath(view.prim_paths[prim_index]) + world_attr = rt_prim.GetAttribute(view._WORLD_MATRIX_NAME) + matrix = world_attr.Get() + scale = usdrt.Gf.Transform(matrix).GetScale() + return torch.tensor( + [[float(scale[0]), float(scale[1]), float(scale[2])]], + dtype=torch.float32, + device=view._device, + ) + + +def _read_fabric_local_matrix_translation(view, prim_index=0): + """Read cached Fabric localMatrix directly, without FrameView getter sync.""" + rt_prim = view._stage.GetPrimAtPath(view.prim_paths[prim_index]) + local_attr = rt_prim.GetAttribute(view._LOCAL_MATRIX_NAME) + matrix = local_attr.Get() + translation = matrix.ExtractTranslation() + return torch.tensor( + [[float(translation[0]), float(translation[1]), float(translation[2])]], + dtype=torch.float32, + device=view._device, + ) + + +@pytest.mark.parametrize("device", ["cuda:0"]) +def test_set_local_via_fabric_path(device, view_factory): + """Exercise the Fabric-native set_local_poses path. + + Ensures set_local_poses computes child_world = parent_world * local + entirely within Fabric (not falling back to USD) by first triggering + the Fabric sync via get_world_poses. + """ + bundle = view_factory(num_envs=1, device=device) + view = bundle.view + + # Trigger lazy `_initialize_fabric()` so subsequent calls take the Fabric path. + view.get_world_poses() + + # Now write via the writer scope (Fabric path). + new_local_pos = wp.zeros((1, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=1, inputs=[new_local_pos, 1.0, 2.0, 3.0], device=device) + ori = torch.tensor([[0.0, 0.0, 0.0, 1.0]], dtype=torch.float32, device=device) + new_local_ori = wp.from_torch(ori) + + with view.xform_local_space_writer() as w: + w.set_poses(positions=new_local_pos, orientations=new_local_ori) + + # Verify: world = parent(0,0,1) + local(1,2,3) = (1,2,4) + world_pos, _ = view.get_world_poses() + expected = torch.tensor([[1.0, 2.0, 4.0]], dtype=torch.float32, device=device) + torch.testing.assert_close(torch.as_tensor(world_pos, device=device), expected, atol=1e-4, rtol=0) + + # Verify get_local_poses returns the local offset + local_pos, _ = view.get_local_poses() + expected_local = torch.tensor([[1.0, 2.0, 3.0]], dtype=torch.float32, device=device) + torch.testing.assert_close(torch.as_tensor(local_pos, device=device), expected_local, atol=1e-4, rtol=0) + + +@pytest.mark.parametrize("device", ["cuda:0"]) +def test_get_scales_fabric_path(device, view_factory): + """Exercise the Fabric-native get_world_scales path.""" + bundle = view_factory(num_envs=1, device=device) + view = bundle.view + + # Trigger lazy `_initialize_fabric()` so the get_world_scales call below uses Fabric. + view.get_world_poses() + + scales = view.get_world_scales() + scales_t = scales.torch + # Default scale should be (1, 1, 1) + expected = torch.tensor([[1.0, 1.0, 1.0]], dtype=torch.float32, device=device) + torch.testing.assert_close(scales_t, expected, atol=1e-4, rtol=0) + + +@pytest.mark.parametrize("device", ["cuda:0"]) +def test_local_scales_roundtrip(device, view_factory): + """Writing scales through the local-space writer roundtrips via ``localMatrix``.""" + bundle = view_factory(num_envs=2, device=device) + view = bundle.view + + # Force Fabric init + view.get_world_poses() + + new_scales = wp.zeros((2, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=2, inputs=[new_scales, 2.0, 3.0, 4.0], device=device) + with view.xform_local_space_writer() as w: + w.set_scales(new_scales) + + ret_scales = view.get_local_scales() + scales_torch = ret_scales.torch + expected = torch.tensor([[2.0, 3.0, 4.0], [2.0, 3.0, 4.0]], device=device) + torch.testing.assert_close(scales_torch, expected, atol=1e-5, rtol=0) + + +@pytest.mark.parametrize("device", ["cuda:0"]) +def test_world_scales_roundtrip(device, view_factory): + """Writing scales through the world-space writer roundtrips via ``worldMatrix``.""" + bundle = view_factory(num_envs=2, device=device) + view = bundle.view + + # Force Fabric init + view.get_world_poses() + + new_scales = wp.zeros((2, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=2, inputs=[new_scales, 5.0, 6.0, 7.0], device=device) + with view.xform_world_space_writer() as w: + w.set_scales(new_scales) + + ret_scales = view.get_world_scales() + scales_torch = ret_scales.torch + expected = torch.tensor([[5.0, 6.0, 7.0], [5.0, 6.0, 7.0]], device=device) + torch.testing.assert_close(scales_torch, expected, atol=1e-5, rtol=0) + + +# ------------------------------------------------------------------ +# Transpose-convention verification: world ↔ local kernels rely on the +# identity ``(A·B)ᵀ = Bᵀ·Aᵀ`` to drop explicit transposes when operating +# on Fabric's column-transposed matrix storage. The translation-only +# parents used by the standard fixture cannot distinguish the right +# convention from the wrong one -- the rotation block is identity and +# equals its own transpose. These tests use a parent rotated 90° around +# Z so that an incorrect storage convention would produce a clearly +# wrong child pose. +# ------------------------------------------------------------------ + + +# Parent at (0, 0, 1) rotated +90° around Z (so the parent X axis points +# along world +Y). Quaternion components in (x, y, z, w) order. +_ROTATED_PARENT_POS = (0.0, 0.0, 1.0) +_ROTATED_PARENT_QUAT_XYZW = (0.0, 0.0, 0.70710678, 0.70710678) + + +def _build_rotated_parent_view(device: str) -> "FrameView": + """Build a 1-env FabricFrameView whose parent is rotated 90° around Z.""" + stage = sim_utils.get_current_stage() + sim_utils.create_prim( + "/World/Parent_0", + "Xform", + translation=_ROTATED_PARENT_POS, + orientation=_ROTATED_PARENT_QUAT_XYZW, + stage=stage, + ) + sim_utils.create_prim("/World/Parent_0/Child", "Camera", translation=(0.0, 0.0, 0.0), stage=stage) + 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() # force Fabric init and USD→Fabric seed + return view + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_set_local_then_get_world_with_rotated_parent(device): + """Verify ``update_indexed_world_matrix_from_local`` under non-identity parent rotation. + + With parent rotated +90° around Z, a child local translation of (1, 0, 0) + must produce world translation (0, 1, 1) -- parent_pos + R · local. If the + transpose convention in the kernel were wrong, the rotation would flip + direction and the world position would land at (0, -1, 1) instead. + """ + _skip_if_unavailable(device) + view = _build_rotated_parent_view(device) + + new_local = wp.zeros((1, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=1, inputs=[new_local, 1.0, 0.0, 0.0], device=device) + identity_quat = wp.from_torch(torch.tensor([[0.0, 0.0, 0.0, 1.0]], dtype=torch.float32, device=device)) + with view.xform_local_space_writer() as w: + w.set_poses(positions=new_local, orientations=identity_quat) + + 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) + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_set_world_then_get_local_with_rotated_parent(device): + """Verify ``update_indexed_local_matrix_from_world`` under non-identity parent rotation. + + With parent rotated +90° around Z and at (0, 0, 1), writing child world + translation (5, 0, 2) must yield child local translation Rᵀ · (5, 0, 1) = + (0, -5, 1). A wrong transpose convention would invert the rotation in the + wrong direction and produce (0, 5, 1) instead. + """ + _skip_if_unavailable(device) + view = _build_rotated_parent_view(device) + + new_world = wp.zeros((1, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=1, inputs=[new_world, 5.0, 0.0, 2.0], device=device) + with view.xform_world_space_writer() as w: + w.set_poses(positions=new_world) + + 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) + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_initial_seed_with_scaled_parent(device): + """Verify the initial USD→Fabric seed handles non-unit scales correctly. + + Sets up a parent with world scale (2, 1, 1) and a child with local scale + (3, 1, 1) at local translation (1, 0, 0). Expected world-space values for + the child: + + * world scale = parent_scale * child_local_scale = (6, 1, 1) + * world position = parent_pos + parent_scale * child_local_pos + = (0, 0, 1) + (2 * 1, 0, 0) = (2, 0, 1) + + If the parent's worldMatrix is seeded with a hardcoded unit scale, + ``get_scales`` returns (3, 1, 1) instead of (6, 1, 1) and ``get_world_poses`` + returns (1, 0, 1) instead of (2, 0, 1). If the child's localMatrix is + seeded without scale, the derived world scale collapses to (2, 1, 1). + This test catches both regressions. + """ + _skip_if_unavailable(device) + stage = sim_utils.get_current_stage() + sim_utils.create_prim("/World/Parent_0", "Xform", translation=(0.0, 0.0, 1.0), scale=(2.0, 1.0, 1.0), stage=stage) + sim_utils.create_prim( + "/World/Parent_0/Child", + "Camera", + translation=(1.0, 0.0, 0.0), + scale=(3.0, 1.0, 1.0), + stage=stage, + ) + sim_utils.SimulationContext(sim_utils.SimulationCfg(dt=0.01, device=device, use_fabric=True)) + view = FrameView("/World/Parent_.*/Child", device=device) + + world_pos, _ = view.get_world_poses() + torch.testing.assert_close( + torch.as_tensor(world_pos, device=device), + torch.tensor([[2.0, 0.0, 1.0]], dtype=torch.float32, device=device), + atol=1e-5, + rtol=0, + ) + + scales = view.get_world_scales().torch + torch.testing.assert_close( + scales, + torch.tensor([[6.0, 1.0, 1.0]], dtype=torch.float32, device=device), + atol=1e-5, + rtol=0, + ) # ------------------------------------------------------------------ -# Multi-GPU tests (cuda:1) — skipped automatically on single-GPU workstations +# Multi-view per stage: per-view single-writer isolation +# ------------------------------------------------------------------ + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_multi_view_writer_isolation(device): + """Two ``FabricFrameView`` instances on the same stage have independent writer scopes. + + Each view's ``_active_writer`` is per-instance, so view B's read or writer + scope must not interfere with view A's pending writes. Verifies that + writes through one view do not corrupt the other view's poses, and that + each view can hold its own writer scope concurrently with reads on the + other view. + """ + _skip_if_unavailable(device) + stage = sim_utils.get_current_stage() + + sim_utils.create_prim("/World/EnvA_0", "Xform", translation=(0.0, 0.0, 1.0), stage=stage) + sim_utils.create_prim("/World/EnvA_0/ChildA", "Camera", translation=(0.1, 0.0, 0.0), stage=stage) + sim_utils.create_prim("/World/EnvB_0", "Xform", translation=(0.0, 0.0, 2.0), stage=stage) + sim_utils.create_prim("/World/EnvB_0/ChildB", "Camera", translation=(0.2, 0.0, 0.0), stage=stage) + + sim_utils.SimulationContext(sim_utils.SimulationCfg(dt=0.01, device=device, use_fabric=True)) + view_a = FrameView("/World/EnvA_.*/ChildA", device=device) + view_b = FrameView("/World/EnvB_.*/ChildB", device=device) + + expected_a0 = torch.tensor([[0.1, 0.0, 1.0]], dtype=torch.float32, device=device) + expected_b0 = torch.tensor([[0.2, 0.0, 2.0]], dtype=torch.float32, device=device) + torch.testing.assert_close( + torch.as_tensor(view_a.get_world_poses()[0], device=device), expected_a0, atol=1e-5, rtol=0 + ) + torch.testing.assert_close( + torch.as_tensor(view_b.get_world_poses()[0], device=device), expected_b0, atol=1e-5, rtol=0 + ) + + # Write a new local pose on view A only via a writer scope. + new_local_a = wp.zeros((1, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=1, inputs=[new_local_a, 1.0, 0.0, 0.0], device=device) + identity_quat = wp.from_torch(torch.tensor([[0.0, 0.0, 0.0, 1.0]], dtype=torch.float32, device=device)) + with view_a.xform_local_space_writer() as w: + w.set_poses(positions=new_local_a, orientations=identity_quat) + + # View B remains undisturbed. + torch.testing.assert_close( + torch.as_tensor(view_b.get_world_poses()[0], device=device), expected_b0, atol=1e-5, rtol=0 + ) + + # View A's world reflects the new local. + expected_a1 = torch.tensor([[1.0, 0.0, 1.0]], dtype=torch.float32, device=device) + torch.testing.assert_close( + torch.as_tensor(view_a.get_world_poses()[0], device=device), expected_a1, atol=1e-5, rtol=0 + ) + + # Write a new local pose on view B; view A unaffected. + new_local_b = wp.zeros((1, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=1, inputs=[new_local_b, 3.0, 0.0, 0.0], device=device) + with view_b.xform_local_space_writer() as w: + w.set_poses(positions=new_local_b, orientations=identity_quat) + torch.testing.assert_close( + torch.as_tensor(view_a.get_world_poses()[0], device=device), expected_a1, atol=1e-5, rtol=0 + ) + expected_b1 = torch.tensor([[3.0, 0.0, 2.0]], dtype=torch.float32, device=device) + torch.testing.assert_close( + torch.as_tensor(view_b.get_world_poses()[0], device=device), expected_b1, atol=1e-5, rtol=0 + ) + + # Single-active-writer is per-view: opening a writer on A leaves B free. + with view_a.xform_world_space_writer(): + assert view_a._active_writer is not None + assert view_b._active_writer is None + # B can still open its own writer concurrently. + with view_b.xform_world_space_writer(): + assert view_b._active_writer is not None + assert view_a._active_writer is None + assert view_b._active_writer is None + + +# ------------------------------------------------------------------ +# Multi-GPU tests (cuda:1) -- skipped automatically on single-GPU workstations # ------------------------------------------------------------------ @@ -268,7 +654,8 @@ def test_fabric_cuda1_world_pose_roundtrip(device, view_factory): new_pos = wp.zeros((2, 3), dtype=wp.float32, device=device) wp.launch(kernel=_fill_position, dim=2, inputs=[new_pos, 10.0, 20.0, 30.0], device=device) - view.set_world_poses(positions=new_pos) + with view.xform_world_space_writer() as w: + w.set_poses(positions=new_pos) ret_pos, _ = view.get_world_poses() pos_torch = torch.as_tensor(ret_pos, device=device) @@ -298,9 +685,10 @@ def test_fabric_cuda1_no_usd_writeback(device, view_factory): new_pos = wp.zeros((1, 3), dtype=wp.float32, device=device) wp.launch(kernel=_fill_position, dim=1, inputs=[new_pos, 99.0, 99.0, 99.0], device=device) - view.set_world_poses(positions=new_pos) + with view.xform_world_space_writer() as w: + w.set_poses(positions=new_pos) - # USD must not have moved at all — equality, not approximate. + # USD must not have moved at all -- equality, not approximate. t_after = UsdGeom.XformCache().GetLocalToWorldTransform(prim).ExtractTranslation() usd_pos_after = torch.tensor([float(t_after[0]), float(t_after[1]), float(t_after[2])]) assert torch.allclose(usd_pos_after, orig_usd_pos, atol=0.0), ( @@ -314,20 +702,323 @@ def test_fabric_cuda1_no_usd_writeback(device, view_factory): ) @pytest.mark.parametrize("device", ["cuda:1"]) def test_fabric_cuda1_scales_roundtrip(device, view_factory): - """set_scales -> get_scales roundtrip works on cuda:1. + """World-space scale writes roundtrip via ``worldMatrix`` on ``cuda:1``. - Both write paths (``set_world_poses`` and ``set_scales``) call - ``_prepare_for_reuse`` and launch on ``self._device``; this test covers - the scales path on the non-primary CUDA device. + Covers the scales path on the non-primary CUDA device: the writer + launches its compose kernel on ``self._device`` and the indexed-fabric + arrays are rebuilt on demand from the corresponding selection. """ bundle = view_factory(2, device) view = bundle.view new_scales = wp.zeros((2, 3), dtype=wp.float32, device=device) wp.launch(kernel=_fill_position, dim=2, inputs=[new_scales, 2.0, 3.0, 4.0], device=device) - view.set_scales(new_scales) + with view.xform_world_space_writer() as w: + w.set_scales(new_scales) - ret_scales = view.get_scales() - scales_torch = torch.as_tensor(ret_scales, device=device) + ret_scales = view.get_world_scales() + scales_torch = ret_scales.torch expected = torch.tensor([[2.0, 3.0, 4.0], [2.0, 3.0, 4.0]], device=device) assert torch.allclose(scales_torch, expected, atol=1e-7), f"Scales roundtrip failed on {device}: {scales_torch}" + + +# ------------------------------------------------------------------ +# Sequential writer scopes (interleaved world / local writes via two scopes) +# ------------------------------------------------------------------ + + +def _build_two_child_view(device: str) -> "FrameView": + """Build a 2-env FabricFrameView with rotated parent for interleave tests. + + Parent at (0, 0, 1) rotated 90° around Z. Two child prims at identity local. + """ + _skip_if_unavailable(device) + stage = sim_utils.get_current_stage() + for i in range(2): + sim_utils.create_prim( + f"/World/Parent_{i}", + "Xform", + translation=_ROTATED_PARENT_POS, + orientation=_ROTATED_PARENT_QUAT_XYZW, + stage=stage, + ) + sim_utils.create_prim(f"/World/Parent_{i}/Child", "Camera", translation=(0.0, 0.0, 0.0), stage=stage) + 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() # force init + return view + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_sequential_world_then_local_scopes_partial_indices(device): + """A world writer scope (idx 0), then a local writer scope (idx 1). Both correct.""" + view = _build_two_child_view(device) + + new_world_pos = wp.zeros((1, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=1, inputs=[new_world_pos, 5.0, 0.0, 2.0], device=device) + idx0 = wp.from_torch(torch.tensor([0], dtype=torch.int32, device=device)) + with view.xform_world_space_writer() as w: + w.set_poses(positions=new_world_pos, indices=idx0) + + new_local_pos = wp.zeros((1, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=1, inputs=[new_local_pos, 1.0, 0.0, 0.0], device=device) + identity_quat = wp.from_torch(torch.tensor([[0.0, 0.0, 0.0, 1.0]], dtype=torch.float32, device=device)) + idx1 = wp.from_torch(torch.tensor([1], dtype=torch.int32, device=device)) + with view.xform_local_space_writer() as w: + w.set_poses(positions=new_local_pos, orientations=identity_quat, indices=idx1) + + # Verify index 0's world pose is still (5, 0, 2) -- index 1's local-scope write + # derives world from the just-written local; index 0 was outside the derived + # set on entry but the world-scope already wrote it and the second scope + # re-derives world from local for all prims (including idx 0). After the + # second scope, idx 0's local was derived from its world (= (0, -5, 1)), + # so re-deriving world = parent * local lands back on (5, 0, 2). + world_pos, _ = view.get_world_poses(indices=idx0) + torch.testing.assert_close( + torch.as_tensor(world_pos, device=device), + torch.tensor([[5.0, 0.0, 2.0]], dtype=torch.float32, device=device), + atol=1e-5, + rtol=0, + ) + + # Index 0's local (derived from its world): + # local = Rᵀ · (child_world_pos - parent_pos) = Rz(-90)·(5, 0, 1) = (0, -5, 1) + local_pos_0, _ = view.get_local_poses(indices=idx0) + torch.testing.assert_close( + torch.as_tensor(local_pos_0, device=device), + torch.tensor([[0.0, -5.0, 1.0]], dtype=torch.float32, device=device), + atol=1e-5, + rtol=0, + ) + + # Index 1's world (derived from local): + # world = parent_world * local = Rz(90)·(1, 0, 0) + parent_pos = (0, 1, 0) + (0, 0, 1) = (0, 1, 1) + world_pos_1, _ = view.get_world_poses(indices=idx1) + torch.testing.assert_close( + torch.as_tensor(world_pos_1, device=device), + torch.tensor([[0.0, 1.0, 1.0]], dtype=torch.float32, device=device), + atol=1e-5, + rtol=0, + ) + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_sequential_local_then_world_scopes_partial_indices(device): + """A local writer scope (idx 0), then a world writer scope (idx 1). Both correct.""" + view = _build_two_child_view(device) + + new_local_pos = wp.zeros((1, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=1, inputs=[new_local_pos, 2.0, 3.0, 0.0], device=device) + identity_quat = wp.from_torch(torch.tensor([[0.0, 0.0, 0.0, 1.0]], dtype=torch.float32, device=device)) + idx0 = wp.from_torch(torch.tensor([0], dtype=torch.int32, device=device)) + with view.xform_local_space_writer() as w: + w.set_poses(positions=new_local_pos, orientations=identity_quat, indices=idx0) + + new_world_pos = wp.zeros((1, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=1, inputs=[new_world_pos, 10.0, 20.0, 30.0], device=device) + idx1 = wp.from_torch(torch.tensor([1], dtype=torch.int32, device=device)) + with view.xform_world_space_writer() as w: + w.set_poses(positions=new_world_pos, indices=idx1) + + # Index 0's world (derived from local): + # world = Rz(90)·(2, 3, 0) + (0, 0, 1) = (-3, 2, 0) + (0, 0, 1) = (-3, 2, 1) + world_pos_0, _ = view.get_world_poses(indices=idx0) + torch.testing.assert_close( + torch.as_tensor(world_pos_0, device=device), + torch.tensor([[-3.0, 2.0, 1.0]], dtype=torch.float32, device=device), + atol=1e-5, + rtol=0, + ) + + # Index 1's world is still (10, 20, 30). + world_pos_1, _ = view.get_world_poses(indices=idx1) + torch.testing.assert_close( + torch.as_tensor(world_pos_1, device=device), + torch.tensor([[10.0, 20.0, 30.0]], dtype=torch.float32, device=device), + atol=1e-5, + rtol=0, + ) + + # Index 1's local (derived from world): + # local = Rᵀ·(world - parent) = Rz(-90)·(10, 20, 29) = (20, -10, 29) + local_pos_1, _ = view.get_local_poses(indices=idx1) + torch.testing.assert_close( + torch.as_tensor(local_pos_1, device=device), + torch.tensor([[20.0, -10.0, 29.0]], dtype=torch.float32, device=device), + atol=1e-5, + rtol=0, + ) + + +# ------------------------------------------------------------------ +# FrameViewSpaceWriterBase contract tests +# ------------------------------------------------------------------ + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_world_writer_writes_world_and_derives_local(device, view_factory): + """A world writer's set_poses + set_scales updates cached Fabric world AND derives local on exit.""" + bundle = view_factory(num_envs=1, device=device) + view = bundle.view + view.get_world_poses() # trigger Fabric init + + new_pos = wp.zeros((1, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=1, inputs=[new_pos, 1.0, 2.0, 4.0], device=device) + new_scales = wp.zeros((1, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=1, inputs=[new_scales, 2.0, 3.0, 4.0], device=device) + + with view.xform_world_space_writer() as w: + w.set_poses(positions=new_pos) + w.set_scales(new_scales) + + # Cached world reflects the writes (parent is at (0, 0, 1) so child world pos + # is whatever we wrote). + expected_world = torch.tensor([[1.0, 2.0, 4.0]], dtype=torch.float32, device=device) + cached_world = _read_fabric_world_matrix_translation(view) + torch.testing.assert_close(cached_world, expected_world, atol=1e-5, rtol=0) + + expected_scale = torch.tensor([[2.0, 3.0, 4.0]], dtype=torch.float32, device=device) + cached_scale = _read_fabric_world_matrix_scale(view) + torch.testing.assert_close(cached_scale, expected_scale, atol=1e-5, rtol=0) + + # Local derived: with parent at (0, 0, 1) (identity rotation, unit scale), + # local translation = world - parent = (1, 2, 3). + expected_local = torch.tensor([[1.0, 2.0, 3.0]], dtype=torch.float32, device=device) + cached_local = _read_fabric_local_matrix_translation(view) + torch.testing.assert_close(cached_local, expected_local, atol=1e-5, rtol=0) + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_local_writer_writes_local_and_derives_world(device, view_factory): + """A local writer's set_poses updates cached Fabric local AND derives world on exit.""" + bundle = view_factory(num_envs=1, device=device) + view = bundle.view + view.get_world_poses() + + new_pos = wp.zeros((1, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=1, inputs=[new_pos, 1.0, 2.0, 3.0], device=device) + identity_quat = wp.from_torch(torch.tensor([[0.0, 0.0, 0.0, 1.0]], dtype=torch.float32, device=device)) + + with view.xform_local_space_writer() as w: + w.set_poses(positions=new_pos, orientations=identity_quat) + + expected_local = torch.tensor([[1.0, 2.0, 3.0]], dtype=torch.float32, device=device) + cached_local = _read_fabric_local_matrix_translation(view) + torch.testing.assert_close(cached_local, expected_local, atol=1e-5, rtol=0) + + # World derived: with parent at (0, 0, 1), world = parent + local = (1, 2, 4). + expected_world = torch.tensor([[1.0, 2.0, 4.0]], dtype=torch.float32, device=device) + cached_world = _read_fabric_world_matrix_translation(view) + torch.testing.assert_close(cached_world, expected_world, atol=1e-5, rtol=0) + + +@pytest.mark.parametrize("device", ["cuda:0"]) +def test_writer_single_derivation_per_scope(device, view_factory, monkeypatch): + """Multiple set_* calls inside one scope produce exactly one derive-kernel launch.""" + bundle = view_factory(num_envs=1, device=device) + view = bundle.view + view.get_world_poses() + + calls = 0 + original = view._recompute_local_from_world_all + + def counted(): + nonlocal calls + calls += 1 + original() + + monkeypatch.setattr(view, "_recompute_local_from_world_all", counted) + + new_pos = wp.zeros((1, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=1, inputs=[new_pos, 1.0, 2.0, 4.0], device=device) + new_scales = wp.zeros((1, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=1, inputs=[new_scales, 2.0, 3.0, 4.0], device=device) + + with view.xform_world_space_writer() as w: + w.set_poses(positions=new_pos) + w.set_scales(new_scales) + assert calls == 0 # no derive yet + + assert calls == 1 # exactly one derive on exit + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_writer_single_active_invariant(device, view_factory): + """Only one writer scope may be active per view at a time.""" + bundle = view_factory(num_envs=1, device=device) + view = bundle.view + view.get_world_poses() + + with view.xform_world_space_writer(): + with pytest.raises(RuntimeError, match="already has an active writer"): + view.xform_world_space_writer().__enter__() + with pytest.raises(RuntimeError, match="already has an active writer"): + view.xform_local_space_writer().__enter__() + # After the outer scope exits, the lock is released and a new scope succeeds. + with view.xform_local_space_writer(): + pass + + +@pytest.mark.parametrize("device", ["cuda:0"]) +def test_writer_restores_hierarchy_change_tracking(device, view_factory): + """``__exit__`` restores the prior ``track_*_xform_changes`` state (don't re-enable paused listeners).""" + bundle = view_factory(num_envs=1, device=device) + view = bundle.view + view.get_world_poses() + h = view._fabric_hierarchy + + # Case 1: pre-paused local stays paused after exit. + h.track_local_xform_changes(False) + assert not h.tracking_local_xform_changes + with view.xform_world_space_writer(): + pass + assert not h.tracking_local_xform_changes, "writer must not re-enable a pre-paused local listener" + + # Case 2: pre-enabled local stays enabled after exit. + h.track_local_xform_changes(True) + with view.xform_world_space_writer(): + pass + assert h.tracking_local_xform_changes, "writer must restore the pre-enabled local listener" + + +@pytest.mark.parametrize("device", ["cuda:0"]) +def test_writer_empty_scope_does_no_derivation(device, view_factory, monkeypatch): + """Entering and exiting a writer scope without any ``set_*`` call must not launch the derive kernel.""" + bundle = view_factory(num_envs=1, device=device) + view = bundle.view + view.get_world_poses() + + calls = 0 + original = view._recompute_local_from_world_all + + def counted(): + nonlocal calls + calls += 1 + original() + + monkeypatch.setattr(view, "_recompute_local_from_world_all", counted) + + with view.xform_world_space_writer(): + pass + + assert calls == 0 + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_view_getter_inside_scope_raises(device, view_factory): + """View-level getters raise ``RuntimeError`` while a writer scope is active.""" + bundle = view_factory(num_envs=1, device=device) + view = bundle.view + view.get_world_poses() + + with view.xform_world_space_writer(): + with pytest.raises(RuntimeError, match="while a writer scope is active"): + view.get_world_poses() + with pytest.raises(RuntimeError, match="while a writer scope is active"): + view.get_local_poses() + with pytest.raises(RuntimeError, match="while a writer scope is active"): + view.get_world_scales() + with pytest.raises(RuntimeError, match="while a writer scope is active"): + view.get_local_scales() + # After the scope exits, view getters work again. + view.get_world_poses() From 63e4005d20d73e2994fb5d9ed71597654d682e31 Mon Sep 17 00:00:00 2001 From: pv-nvidia <197907000+pv-nvidia@users.noreply.github.com> Date: Fri, 7 Aug 2026 08:49:21 +0200 Subject: [PATCH 3/4] perf: scale FabricFrameView selections to the view, not the stage (#6805) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Camera world-pose resolution stalled at high environment counts on the PhysX backend, badly enough to time out benchmarks. There are two independent bottlenecks; the PR also adds a lifecycle hook that came out of review. whole stage Cameras read their poses from Fabric. To do that the frame view needs to know where each camera's data sits in Fabric memory, so it builds a lookup table from "camera number" to "Fabric slot". Building that table was very slow. The view selected prims by requiring the Fabric world and local matrix attributes — but every xformable in the stage carries those, so the selection matched the **entire stage** (~1.1M prims at 8192 environments), not the view's handful of cameras. Finding its own prims in that list then meant building a Python dictionary over all of it, on the host: ```python path_to_idx = {str(p): i for i, p in enumerate(selection.GetPaths())} ``` That ran twice per rebuild (once for the children, once for their parents), and rebuilt whenever `PrepareForReuse()` reported a bucket change — in practice on every environment reset. The cost of that dictionary scales with the size of the whole scene, not with the number of cameras: it converts every path in the stage to a Python string and inserts it. Measured directly, one full rebuild takes 10 ms on a bare stage but **556 ms once 100k unrelated prims are present** (table below), and it grows linearly from there — seconds of host work per reset at the ~1.1M prims of an 8192-environment scene. **Fix:** author a private per-view `uint` index attribute on each managed prim (and on each unique parent), holding that prim's view index. Every selection requires the matching index attribute, so **selections resolve to exactly the view's prims**. The view-to-Fabric slot mapping is then rebuilt by inverting that attribute in a single Warp kernel launch: ```python fabric_slots[view_indices[fabric_slot]] = fabric_slot ``` That is `O(view)` device work with no host-side path resolution and no Python objects created per frame. There is also **no cache**, so there is nothing to invalidate. The table is rebuilt from live Fabric data on every access, which means a bucket reorder can never leave a stale mapping behind. Selections are still checked with `GetCount()`, which is exact here: the index attribute belongs to one view, so the count can only change if one of that view's prims actually disappeared — and then the view raises a clear error instead of silently reading the wrong prim. Attribute names embed a process-wide monotonic uid, so a dead view's leftover attributes can never satisfy a live view's selection — and since review, they are removed outright when the view is released (section 3). Parent reads get their own read-only selection, keeping the child RO/RW flip semantics from #5677 intact. ```python for i, path in enumerate(input_paths): mapping[i] = paths.index(path) # O(N) scan, per path ``` **Fix:** build the reverse dict once (first occurrence wins, matching `list.index` semantics), then resolve each path in `O(1)`. Measured over 200k paths, preserving first-occurrence semantics costs nothing versus an unguarded dict comprehension, so the fix carries no behaviour change. review) The per-view index attributes from fix 1 were initially never removed, so views recreated over the same prims on a long-lived stage accumulated attribute pairs, widening those prims' Fabric buckets. Measured over 8 successive views on one stage (512 prims), attributes on a single prim went 1 → 8 and first-access rose monotonically from 120 ms to 173 ms. **Fix:** `FabricFrameView.close()` removes the view's tags (measured at 4.7 µs per prim, ~77 ms for an 8192-env view) and is safe to call more than once. `BaseFrameView` gains a no-op `close()` so callers can close any backend uniformly. Every in-repo callsite now closes its view: `Camera` when invalidated **and when dropped** (`__del__`), `prim_world_positions` after each per-env view it creates in a loop, and `SceneAsset` before replacing its cached view. Views dropped without `close()` are cleaned up best-effort from `__del__`, following the shutdown-safe idiom the env classes already use (`sys` bound as a default argument; nothing runs during interpreter finalization, where Fabric dies with the process anyway), and a warning names the view so the missing `close()` call can be fixed. The stall was invisible in Tracy. Comparing captures before and after the commit that introduced it, Kit's own per-frame work is unchanged — `App Update` totals 65.1 s before and 62.3 s after, over ~120 frames. All the lost time sits in the gaps *between* frames, where the main thread is blocked in Python and Kit records nothing. A `py-spy` capture found the cause: `_compute_fabric_indices_for` accounted for **11.86% of samples (23.2 s of a 195 s run)**, reached through `camera.reset()` → `get_world_poses()` → the selection rebuild. Every frame on that chain carries essentially the same share, so the whole cost is that one function. Fixes nvbug 6535498. - Bug fix (non-breaking change which fixes an issue) - New feature (non-breaking change which adds functionality) — `FrameView.close()` Task `Isaac-Lift-KukaAllegro-Camera`, 8192 environments, `presets=physx,isaacsim_rtx_renderer,duo_camera`. Commit `69888c34e471` (which introduced the problem) against its parent `e5f99d320338`: | | Parent | With the problem | | --- | ---: | ---: | | Mean environment step | 699 ms | 7193 ms | | Slow steps (over 3 s) | 1 of 100 | 28 of 100, averaging 18.9 s | | GPU utilization | 23.4% | 4.7% | GPU and CPU utilization both drop during the slow steps, which is what a stall looks like — the pipeline is waiting, not doing extra work. With this PR, the same task at 8192 environments steps at **mean 1.554 s, peak 1.663 s** — peak within 7% of mean, i.e. the stalls are gone. (Measured on an RTX A6000, so the absolute step time is not comparable to the table above, which was captured on different hardware; the stall signature is the comparison that matters.) L40, `cuda:0`, 1024-prim view (1024 children + 1024 parents). "Filler" is xformable prims **not** in the view but carrying Fabric matrices — i.e. the rest of a real scene. | filler prims | `develop` selection | this PR | | ---: | ---: | ---: | | 0 | 2,048 | 1,024 | | 10,000 | 12,048 | 1,024 | | 100,000 | 102,048 | **1,024** | `develop` scales with the stage; this PR is pinned to the view. | filler prims | `develop` child | `develop` parent | `develop` full RO rebuild | PR child | PR parent | | ---: | ---: | ---: | ---: | ---: | ---: | | 0 | 4.67 | 5.40 | 10.24 | 0.137 | 0.191 | | 10,000 | 24.05 | 24.69 | 62.47 | 0.139 | 0.193 | | 100,000 | 274.21 | 270.47 | **556.41** | **0.137** | **0.190** | `develop` grows linearly with stage size — **556 ms per rebuild** at 100k filler prims. This PR is flat at **~0.14 ms**, a **~2000×** reduction, and unchanged from 0 to 100k filler prims. First access at 100k filler: 1824 ms → 383 ms. | N | `develop` | this PR | speedup | | ---: | ---: | ---: | ---: | | 1,000 | 0.0038 s | 0.0001 s | 27× | | 5,000 | 0.0862 s | 0.0007 s | 120× | | 20,000 | 1.3241 s | 0.0030 s | 434× | | 50,000 | 8.5507 s | 0.0092 s | **929×** | | 200,000 | (not run) | 0.0467 s | — | Outputs verified identical at every N. At the ~200k rigid bodies of an 8192-environment scene the old path takes minutes. | Suite | Result | | --- | --- | | `isaaclab_physx/test/sim/test_views_xform_prim_fabric.py` | 79 passed, 4 skipped (`cpu` + `cuda:0` on an L40) | | `isaaclab/test/sim/test_views_xform_prim.py` (USD contract) | 63 passed | | `isaaclab_newton/test/sim/test_views_xform_prim_newton.py` | 58 passed | | `isaaclab/test/sensors/test_camera.py` (main consumer, incl. the new `close()` call) | 35 passed, 0 missing-`close()` warnings | | `isaaclab/test/utils/warp/test_proxy_array.py` | 81 passed | New tests: - `test_selections_match_only_the_view_prims` asserts each selection matches exactly the prims the view manages. Verified to fail without the fix — with the selections unscoped it reports `matched 8 prims, expected 4`, because the child selections pick up the parents too. - `test_close_removes_index_attributes` asserts the tag-scoped selection drops to zero after `close()` and that a second `close()` is a no-op. - `test_garbage_collection_removes_index_attributes_and_warns` asserts a view dropped without `close()` still removes its tags on collection and logs the warning. `test_fabric_rebuild_after_topology_change` was updated to drive the new refresh paths instead of the removed `_rebuild_{ro,rw}_arrays`. The 4 skips are environmental: empty device parameter sets, and Fabric hierarchy bindings that are unavailable in a headless experience. **Coverage gap:** `test_physx_scene_data_backend.py` and `test_ovphysx_scene_data_backend.py` both skip at collection on the machine used here, so the `create_mapping` change is exercised only by the standalone benchmark above, not by the test suite. Worth a look in CI, where those backends are available. Not applicable — no visual change. The benchmarks above cover the behaviour change. **Review feedback addressed so far:** the index-attribute accumulation is fixed by `close()` (section 3); the opposite-space recompute no longer refreshes the same selection twice on writer-scope exit; `create_mapping` keeps `list.index`'s first-occurrence semantics for duplicate paths (measured free at 200k paths); and docstrings/changelogs were trimmed to the observable contract. **The `isaaclab`/`isaaclab_physx` changelog fragments are still named `fix-physx-newton-camera-pose-scaling.rst`**, inherited from #6554 — nothing here touches Newton or scaling. The new `isaaclab_mimic` fragment uses the branch slug. Happy to rename if reviewers prefer. **A handful of tests still drop views without `close()`** and now log the new warning during runs; harmless, cleanup to follow. EDIT: fixed — the only remaining warning is the one `test_garbage_collection_removes_index_attributes_and_warns` asserts on. `Camera` previously closed its view only on invalidation, so a dropped camera still went through `__del__`; it now closes in `Camera.__del__` too. `test_camera.py` went from 3 warnings to 0. `SceneDataProvider` fix unchanged (credited with a `Co-authored-by` line) and replaces its Fabric frame view fix. cached the result, reusing it while the number of selected prims stayed the same. Two problems with that. The cache key is unsafe: an equal prim count does not mean the prims or their order stayed the same, so after a bucket reorder or a same-count membership change the cached indices point at the wrong prims and cameras silently read another prim's transform. And it treats the symptom — the underlying operation is still proportional to the size of the whole scene, just performed less often. Making the selection small removes the need for a cache at all, so both problems go away. - [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 added a changelog fragment under `source//changelog.d/` for every touched package (do **not** edit `CHANGELOG.rst` or bump `extension.toml` — CI handles that) - [x] I have added my name to the `CONTRIBUTORS.md` or my name already exists there --------- Co-authored-by: pv-nvidia <197907000+pv-nvidia@users.noreply.github.com> Co-authored-by: yts-nv Co-authored-by: Kelly Guo (cherry picked from commit 480370cdc5b08f5ca8ca733ecf670e2bcc83c0ee) --- .../fix-physx-newton-camera-pose-scaling.rst | 13 + .../isaaclab/envs/utils/camera_view.py | 11 +- .../scene_data/scene_data_provider.py | 12 +- .../isaaclab/sensors/camera/camera.py | 12 +- .../isaaclab/sim/views/base_frame_view.py | 11 + source/isaaclab/isaaclab/utils/warp/fabric.py | 25 + source/isaaclab/test/sensors/test_camera.py | 11 + .../fix-fabric-frameview-stall.rst | 5 + .../locomanipulation_sdg/scene_utils.py | 5 +- .../fix-physx-newton-camera-pose-scaling.rst | 14 + .../sim/views/fabric_frame_view.py | 438 ++++++++++++------ .../test/sim/test_views_xform_prim_fabric.py | 126 ++++- 12 files changed, 503 insertions(+), 180 deletions(-) create mode 100644 source/isaaclab/changelog.d/fix-physx-newton-camera-pose-scaling.rst create mode 100644 source/isaaclab_mimic/changelog.d/fix-fabric-frameview-stall.rst create mode 100644 source/isaaclab_physx/changelog.d/fix-physx-newton-camera-pose-scaling.rst 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 c39fe9c906a8..1447d63577c9 100644 --- a/source/isaaclab/isaaclab/envs/utils/camera_view.py +++ b/source/isaaclab/isaaclab/envs/utils/camera_view.py @@ -242,10 +242,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 ba4c051920e1..f1b6f5592018 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 re from collections import deque from typing import TYPE_CHECKING, Any @@ -212,10 +211,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 881be4c7c9bf..395e8c5495d2 100644 --- a/source/isaaclab/isaaclab/sensors/camera/camera.py +++ b/source/isaaclab/isaaclab/sensors/camera/camera.py @@ -207,11 +207,17 @@ def __init__(self, cfg: CameraCfg): # Renderer and render data — assigned in _initialize_impl. self._renderer: BaseRenderer | None = None self._render_data = None + # Frame view — assigned in _initialize_impl. + self._view: FrameView | None = None def __del__(self): """Unsubscribes from callbacks and cleans up renderer resources.""" # unsubscribe callbacks super().__del__() + # release the frame view's backend state + if self._view is not None: + self._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) @@ -809,5 +815,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/test/sensors/test_camera.py b/source/isaaclab/test/sensors/test_camera.py index 9b41899765b7..5f242984b576 100644 --- a/source/isaaclab/test/sensors/test_camera.py +++ b/source/isaaclab/test/sensors/test_camera.py @@ -1241,6 +1241,17 @@ def test_camera_pose_update_reflected_in_render(setup_camera_device, device): del camera +def test_camera_invalidate_before_initialize(setup_sim_camera): + """Invalidation on a camera that never initialized does not raise.""" + _, camera_cfg, _ = setup_sim_camera + camera = Camera(camera_cfg.replace(prim_path="/World/NeverInitialized", spawn=None)) + try: + assert camera._view is None + camera._invalidate_initialize_callback(None) + finally: + del camera + + def _populate_scene(): """Add prims to the scene.""" # Ground-plane 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 00a29bc04afc..76e34b5403f3 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py +++ b/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py @@ -107,8 +107,11 @@ def _get_xform_view(self) -> FrameView: if xform_prim.count == 0: # The view was created before environment cloning; rebuild it now that prims exist. # FabricFrameView composes UsdFrameView; the template prim_path lives on the inner USD view. + # Close the stale view so its backend state is released now, not on garbage collection. inner = getattr(xform_prim, "_usd_view", xform_prim) - xform_prim = FrameView(inner._prim_path, device=xform_prim.device) + prim_path, device = inner._prim_path, xform_prim.device + xform_prim.close() + xform_prim = FrameView(prim_path, device=device) self.scene.extras[self.entity_name] = xform_prim return xform_prim 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 068d4c5b875f..8caa1b4fa454 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). @@ -127,31 +150,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 @@ -163,6 +178,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, @@ -198,33 +219,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 # ------------------------------------------------------------------ @@ -273,7 +349,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) @@ -428,13 +503,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, @@ -448,91 +524,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.""" @@ -540,16 +634,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 from usdrt import Rt # noqa: PLC0415 @@ -564,39 +661,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) @@ -615,8 +745,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() @@ -637,7 +767,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), @@ -650,8 +780,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() @@ -694,9 +826,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, @@ -720,24 +868,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 @@ -750,11 +880,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 @@ -823,7 +953,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`. """ @@ -884,7 +1014,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 c023ab40ba86..c62f343211d6 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 os 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", ["cpu", "cuda:0"]) 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) @@ -296,6 +303,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", ["cpu", "cuda:0"]) +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]) @@ -483,6 +574,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"]) @@ -505,6 +597,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"]) @@ -553,6 +646,7 @@ def test_initial_seed_with_scaled_parent(device): atol=1e-5, rtol=0, ) + view.close() # ------------------------------------------------------------------ @@ -631,6 +725,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() # ------------------------------------------------------------------ @@ -800,6 +896,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"]) @@ -848,6 +945,7 @@ def test_sequential_local_then_world_scopes_partial_indices(device): atol=1e-5, rtol=0, ) + view.close() # ------------------------------------------------------------------ From 89adc465dbb482a82a298633bb1f4bc06a1089c3 Mon Sep 17 00:00:00 2001 From: pv-nvidia <197907000+pv-nvidia@users.noreply.github.com> Date: Tue, 25 Aug 2026 10:59:01 +0000 Subject: [PATCH 4/4] fix: use the OVPhysX tensor binding API available on release/3.0.0-beta2 The #5677 cherry-pick brought in OvPhysxFrameView code calling OvPhysxView (try_binding_for / read_into). That class arrives with #6224, which is not on this release branch, so the module import raised ModuleNotFoundError as soon as an OVPhysX frame view was initialized. Restore the direct physx.create_tensor_binding / binding.read calls this branch provides. --- .../isaaclab_ovphysx/sim/views/ovphysx_frame_view.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/sim/views/ovphysx_frame_view.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/sim/views/ovphysx_frame_view.py index ff7912f20f60..4eb60809ad6e 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/sim/views/ovphysx_frame_view.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/sim/views/ovphysx_frame_view.py @@ -358,7 +358,6 @@ def _initialize_impl(self, physx: Any) -> None: ``env_0`` replaced by the row's env_id. """ from isaaclab_ovphysx import tensor_types as TT # noqa: PLC0415 - from isaaclab_ovphysx.sim.views.ovphysx_view import OvPhysxView # noqa: PLC0415 xform_cache = UsdGeom.XformCache(Usd.TimeCode.Default()) identity_xform7 = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0] @@ -408,11 +407,8 @@ def _initialize_impl(self, physx: Any) -> None: # 4. Create the RIGID_BODY_POSE binding (or operate in world-only mode). if patterns: pattern = patterns[0] - self._root_view = OvPhysxView(physx, pattern=pattern, device=self._device) - # ``try_binding_for`` returns None for a zero-match binding (the view rejects a - # 0-count binding); surface that as the explicit zero-bodies error below. - self._pose_binding = self._root_view.try_binding_for(TT.RIGID_BODY_POSE) - if self._pose_binding is None: + self._pose_binding = physx.create_tensor_binding(pattern=pattern, tensor_type=TT.RIGID_BODY_POSE) + if self._pose_binding.shape[0] == 0: raise RuntimeError( f"OvPhysxFrameView: RIGID_BODY_POSE binding for pattern {pattern!r} matched zero bodies." ) @@ -420,7 +416,6 @@ def _initialize_impl(self, physx: Any) -> None: binding_paths: list[str] = list(self._pose_binding.prim_paths) else: # All prims resolved as world-attached: no binding needed; kernels only hit the -1 branch. - self._root_view = None self._pose_binding = None self._pose_buf = wp.zeros((1, 7), dtype=wp.float32, device=self._device) binding_paths = [] @@ -619,7 +614,7 @@ def _current_body_q(self) -> wp.array: buffer ``[num_bodies]``. """ if self._pose_binding is not None: - self._root_view.read_into("rigid_body_pose", self._pose_buf) + self._pose_binding.read(self._pose_buf) return self._pose_buf.view(wp.transformf) # ------------------------------------------------------------------