From cefd225ab2976ce29160d49de2648bfad9161a54 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Mon, 18 May 2026 08:27:17 +0000 Subject: [PATCH 01/54] feat: add indexed fabric transform kernels for local/world matrix propagation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add Warp kernels that operate on wp.indexedfabricarray for direct local↔world matrix propagation without round-tripping through USD: - decompose_indexed_fabric_transforms: extract pos/quat/scale from ifa - compose_indexed_fabric_transforms: write pos/quat/scale into ifa - update_indexed_local_matrix_from_world: local = inv(parent) * world - update_indexed_world_matrix_from_local: world = parent * local Also refactor existing kernels to use wp.where (branchless) instead of if/else for broadcast index selection. These kernels are the foundation for Fabric-accelerated get/set_local_poses in FabricFrameView. --- docs/source/api/lab/isaaclab.utils.rst | 13 ++ .../changelog.d/indexed-fabric-kernels.rst | 24 +++ source/isaaclab/isaaclab/utils/warp/fabric.py | 189 +++++++++++++++-- .../test/utils/warp/test_fabric_kernels.py | 193 ++++++++++++++++++ 4 files changed, 407 insertions(+), 12 deletions(-) create mode 100644 source/isaaclab/changelog.d/indexed-fabric-kernels.rst create mode 100644 source/isaaclab/test/utils/warp/test_fabric_kernels.py diff --git a/docs/source/api/lab/isaaclab.utils.rst b/docs/source/api/lab/isaaclab.utils.rst index 5b352152e0b5..f236ebcb6a15 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`. Will be 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/source/isaaclab/changelog.d/indexed-fabric-kernels.rst b/source/isaaclab/changelog.d/indexed-fabric-kernels.rst new file mode 100644 index 000000000000..0f68b4c60e9d --- /dev/null +++ b/source/isaaclab/changelog.d/indexed-fabric-kernels.rst @@ -0,0 +1,24 @@ +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 (no + explicit transposes). Will be used by + :class:`~isaaclab_physx.sim.views.FabricFrameView` to keep child world and + local matrices consistent across writes without round-tripping through USD. + +Changed +^^^^^^^ + +* Replaced ``if/else`` branching with ``wp.where`` in existing Fabric + compose/decompose kernels for branchless GPU execution. diff --git a/source/isaaclab/isaaclab/utils/warp/fabric.py b/source/isaaclab/isaaclab/utils/warp/fabric.py index a48f773f4991..6f9963f290a1 100644 --- a/source/isaaclab/isaaclab/utils/warp/fabric.py +++ b/source/isaaclab/isaaclab/utils/warp/fabric.py @@ -15,15 +15,28 @@ import warp as wp +__all__ = [ + "arange_k", + "compose_fabric_transformation_matrix_from_warp_arrays", + "compose_indexed_fabric_transforms", + "decompose_fabric_transformation_matrix_to_warp_arrays", + "decompose_indexed_fabric_transforms", + "set_view_to_fabric_array", + "update_indexed_local_matrix_from_world", + "update_indexed_world_matrix_from_local", +] + 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) @@ -130,29 +143,20 @@ def compose_fabric_transformation_matrix_from_warp_arrays( position, rotation, scale = _decompose_transformation_matrix(wp.mat44f(fabric_matrices[fabric_index])) # update position (check if array has elements, not just if it exists) if array_positions.shape[0] > 0: - if broadcast_positions: - index = 0 - else: - index = i + index = wp.where(broadcast_positions, 0, i) position[0] = array_positions[index, 0] position[1] = array_positions[index, 1] position[2] = array_positions[index, 2] # update orientation (convert from wxyz to xyzw for Warp) if array_orientations.shape[0] > 0: - if broadcast_orientations: - index = 0 - else: - index = i + index = wp.where(broadcast_orientations, 0, i) rotation[0] = array_orientations[index, 0] # x rotation[1] = array_orientations[index, 1] # y rotation[2] = array_orientations[index, 2] # z rotation[3] = array_orientations[index, 3] # w # update scale if array_scales.shape[0] > 0: - if broadcast_scales: - index = 0 - else: - index = i + index = wp.where(broadcast_scales, 0, i) scale[0] = array_scales[index, 0] scale[1] = array_scales[index, 1] scale[2] = array_scales[index, 2] @@ -163,6 +167,167 @@ 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: + index = wp.where(broadcast_positions, 0, 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: + index = wp.where(broadcast_orientations, 0, 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: + index = wp.where(broadcast_scales, 0, 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(child_world * wp.inverse(parent_world)) # type: ignore[arg-type] + + +@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(child_local * parent_world) # type: ignore[arg-type] + + @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/utils/warp/test_fabric_kernels.py b/source/isaaclab/test/utils/warp/test_fabric_kernels.py new file mode 100644 index 000000000000..13bad50da500 --- /dev/null +++ b/source/isaaclab/test/utils/warp/test_fabric_kernels.py @@ -0,0 +1,193 @@ +# 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 + +"""Unit tests for the Warp fabric transform kernels. + +Tests the decompose/compose math via helper kernels that operate on +regular wp.array (no Fabric/USDRT runtime required). +""" + +import numpy as np +import warp as wp + +wp.init() + +from isaaclab.utils.warp.fabric import _decompose_transformation_matrix # noqa: E402 + + +@wp.kernel(enable_backward=False) +def _test_decompose_kernel( + matrices: wp.array(dtype=wp.mat44f), + out_positions: wp.array(dtype=wp.vec3f), + out_rotations: wp.array(dtype=wp.vec4f), + out_scales: wp.array(dtype=wp.vec3f), +): + """Decompose a batch of 4x4 matrices into pos/quat/scale.""" + i = wp.tid() + pos, rot, scale = _decompose_transformation_matrix(matrices[i]) + out_positions[i] = pos + out_rotations[i] = wp.vec4f(rot[0], rot[1], rot[2], rot[3]) + out_scales[i] = scale + + +@wp.kernel(enable_backward=False) +def _test_compose_kernel( + positions: wp.array(dtype=wp.vec3f), + rotations: wp.array(dtype=wp.vec4f), + scales: wp.array(dtype=wp.vec3f), + out_matrices: wp.array(dtype=wp.mat44f), +): + """Compose a batch of pos/quat/scale into 4x4 matrices.""" + i = wp.tid() + pos = positions[i] + rot = wp.quatf(rotations[i][0], rotations[i][1], rotations[i][2], rotations[i][3]) + scale = scales[i] + out_matrices[i] = wp.transpose(wp.transform_compose(pos, rot, scale)) + + +class TestDecomposeCompose: + """Round-trip tests for decompose ↔ compose transform math.""" + + def test_identity_matrix(self): + """Identity matrix decomposes to pos=0, quat=identity, scale=1.""" + mat = np.eye(4, dtype=np.float32).reshape(1, 4, 4) + matrices = wp.array(mat, dtype=wp.mat44f, device="cpu") + out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") + out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") + out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") + + wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") + wp.synchronize() + + pos = out_pos.numpy() + rot = out_rot.numpy() + scale = out_scale.numpy() + + np.testing.assert_allclose(pos[0], [0, 0, 0], atol=1e-6) + np.testing.assert_allclose(scale[0], [1, 1, 1], atol=1e-6) + # Identity quaternion: either (0,0,0,1) or (0,0,0,-1) + assert abs(abs(rot[0, 3]) - 1.0) < 1e-5 + + def test_translation_only(self): + """Matrix with only translation decomposes correctly.""" + mat = np.eye(4, dtype=np.float32) + mat[3, 0] = 1.0 # row-major: translation in last row + mat[3, 1] = 2.0 + mat[3, 2] = 3.0 + matrices = wp.array(mat.reshape(1, 4, 4), dtype=wp.mat44f, device="cpu") + out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") + out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") + out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") + + wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") + wp.synchronize() + + np.testing.assert_allclose(out_pos.numpy()[0], [1, 2, 3], atol=1e-6) + np.testing.assert_allclose(out_scale.numpy()[0], [1, 1, 1], atol=1e-6) + + def test_uniform_scale(self): + """Matrix with uniform scale decomposes correctly.""" + mat = np.eye(4, dtype=np.float32) + mat[0, 0] = 2.0 + mat[1, 1] = 2.0 + mat[2, 2] = 2.0 + matrices = wp.array(mat.reshape(1, 4, 4), dtype=wp.mat44f, device="cpu") + out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") + out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") + out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") + + wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") + wp.synchronize() + + np.testing.assert_allclose(out_scale.numpy()[0], [2, 2, 2], atol=1e-6) + + def test_round_trip(self): + """Compose then decompose recovers original pos/quat/scale.""" + # Known transform: translate (5,6,7), rotate 90° about Z, scale (1,2,3) + pos = np.array([[5.0, 6.0, 7.0]], dtype=np.float32) + # 90° about Z in xyzw: (0, 0, sin(45°), cos(45°)) + s45 = np.sin(np.pi / 4) + c45 = np.cos(np.pi / 4) + rot = np.array([[0.0, 0.0, s45, c45]], dtype=np.float32) + scale = np.array([[1.0, 2.0, 3.0]], dtype=np.float32) + + wp_pos = wp.array(pos, dtype=wp.vec3f, device="cpu") + wp_rot = wp.array(rot, dtype=wp.vec4f, device="cpu") + wp_scale = wp.array(scale, dtype=wp.vec3f, device="cpu") + out_mat = wp.zeros(1, dtype=wp.mat44f, device="cpu") + + # Compose + wp.launch(_test_compose_kernel, dim=1, inputs=[wp_pos, wp_rot, wp_scale, out_mat], device="cpu") + wp.synchronize() + + # Decompose + out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") + out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") + out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") + + wp.launch(_test_decompose_kernel, dim=1, inputs=[out_mat, out_pos, out_rot, out_scale], device="cpu") + wp.synchronize() + + np.testing.assert_allclose(out_pos.numpy()[0], pos[0], atol=1e-5) + np.testing.assert_allclose(out_scale.numpy()[0], scale[0], atol=1e-5) + # Quaternion sign ambiguity + r_out = out_rot.numpy()[0] + r_exp = rot[0] + dot = np.dot(r_out, r_exp) + np.testing.assert_allclose(abs(dot), 1.0, atol=1e-5) + + def test_non_uniform_scale_round_trip(self): + """Non-uniform scale round-trips correctly.""" + pos = np.array([[0.0, 0.0, 0.0]], dtype=np.float32) + rot = np.array([[0.0, 0.0, 0.0, 1.0]], dtype=np.float32) # identity + scale = np.array([[0.5, 2.0, 3.0]], dtype=np.float32) + + wp_pos = wp.array(pos, dtype=wp.vec3f, device="cpu") + wp_rot = wp.array(rot, dtype=wp.vec4f, device="cpu") + wp_scale = wp.array(scale, dtype=wp.vec3f, device="cpu") + out_mat = wp.zeros(1, dtype=wp.mat44f, device="cpu") + + wp.launch(_test_compose_kernel, dim=1, inputs=[wp_pos, wp_rot, wp_scale, out_mat], device="cpu") + wp.synchronize() + + out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") + out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") + out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") + + wp.launch(_test_decompose_kernel, dim=1, inputs=[out_mat, out_pos, out_rot, out_scale], device="cpu") + wp.synchronize() + + np.testing.assert_allclose(out_scale.numpy()[0], scale[0], atol=1e-5) + + +class TestKernelSignatures: + """Verify all exported kernels are importable and are Warp Kernels.""" + + def test_all_kernels_importable(self): + """All public kernels listed in __all__ should be importable and be Warp Kernels.""" + from isaaclab.utils.warp import fabric as fabric_utils + + expected_kernels = [ + "arange_k", + "compose_fabric_transformation_matrix_from_warp_arrays", + "compose_indexed_fabric_transforms", + "decompose_fabric_transformation_matrix_to_warp_arrays", + "decompose_indexed_fabric_transforms", + "set_view_to_fabric_array", + "update_indexed_local_matrix_from_world", + "update_indexed_world_matrix_from_local", + ] + + for name in expected_kernels: + obj = getattr(fabric_utils, name, None) + assert obj is not None, f"{name} not found in fabric_utils" + assert isinstance(obj, wp.Kernel), f"{name} should be a wp.Kernel, got {type(obj)}" + + def test_module_exports_match_all(self): + """__all__ should list every public kernel.""" + from isaaclab.utils.warp import fabric as fabric_utils + + for name in fabric_utils.__all__: + assert hasattr(fabric_utils, name), f"__all__ lists '{name}' but it's not defined" From 6341d1d42b23d0f35621822daad4b0197841e532 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Thu, 21 May 2026 15:06:59 +0000 Subject: [PATCH 02/54] test: convert fabric kernel tests to plain functions (no classes) --- .../test/utils/warp/test_fabric_kernels.py | 297 +++++++++--------- 1 file changed, 153 insertions(+), 144 deletions(-) diff --git a/source/isaaclab/test/utils/warp/test_fabric_kernels.py b/source/isaaclab/test/utils/warp/test_fabric_kernels.py index 13bad50da500..c2d3e6d37381 100644 --- a/source/isaaclab/test/utils/warp/test_fabric_kernels.py +++ b/source/isaaclab/test/utils/warp/test_fabric_kernels.py @@ -47,147 +47,156 @@ def _test_compose_kernel( out_matrices[i] = wp.transpose(wp.transform_compose(pos, rot, scale)) -class TestDecomposeCompose: - """Round-trip tests for decompose ↔ compose transform math.""" - - def test_identity_matrix(self): - """Identity matrix decomposes to pos=0, quat=identity, scale=1.""" - mat = np.eye(4, dtype=np.float32).reshape(1, 4, 4) - matrices = wp.array(mat, dtype=wp.mat44f, device="cpu") - out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") - out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") - out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") - - wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") - wp.synchronize() - - pos = out_pos.numpy() - rot = out_rot.numpy() - scale = out_scale.numpy() - - np.testing.assert_allclose(pos[0], [0, 0, 0], atol=1e-6) - np.testing.assert_allclose(scale[0], [1, 1, 1], atol=1e-6) - # Identity quaternion: either (0,0,0,1) or (0,0,0,-1) - assert abs(abs(rot[0, 3]) - 1.0) < 1e-5 - - def test_translation_only(self): - """Matrix with only translation decomposes correctly.""" - mat = np.eye(4, dtype=np.float32) - mat[3, 0] = 1.0 # row-major: translation in last row - mat[3, 1] = 2.0 - mat[3, 2] = 3.0 - matrices = wp.array(mat.reshape(1, 4, 4), dtype=wp.mat44f, device="cpu") - out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") - out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") - out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") - - wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") - wp.synchronize() - - np.testing.assert_allclose(out_pos.numpy()[0], [1, 2, 3], atol=1e-6) - np.testing.assert_allclose(out_scale.numpy()[0], [1, 1, 1], atol=1e-6) - - def test_uniform_scale(self): - """Matrix with uniform scale decomposes correctly.""" - mat = np.eye(4, dtype=np.float32) - mat[0, 0] = 2.0 - mat[1, 1] = 2.0 - mat[2, 2] = 2.0 - matrices = wp.array(mat.reshape(1, 4, 4), dtype=wp.mat44f, device="cpu") - out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") - out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") - out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") - - wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") - wp.synchronize() - - np.testing.assert_allclose(out_scale.numpy()[0], [2, 2, 2], atol=1e-6) - - def test_round_trip(self): - """Compose then decompose recovers original pos/quat/scale.""" - # Known transform: translate (5,6,7), rotate 90° about Z, scale (1,2,3) - pos = np.array([[5.0, 6.0, 7.0]], dtype=np.float32) - # 90° about Z in xyzw: (0, 0, sin(45°), cos(45°)) - s45 = np.sin(np.pi / 4) - c45 = np.cos(np.pi / 4) - rot = np.array([[0.0, 0.0, s45, c45]], dtype=np.float32) - scale = np.array([[1.0, 2.0, 3.0]], dtype=np.float32) - - wp_pos = wp.array(pos, dtype=wp.vec3f, device="cpu") - wp_rot = wp.array(rot, dtype=wp.vec4f, device="cpu") - wp_scale = wp.array(scale, dtype=wp.vec3f, device="cpu") - out_mat = wp.zeros(1, dtype=wp.mat44f, device="cpu") - - # Compose - wp.launch(_test_compose_kernel, dim=1, inputs=[wp_pos, wp_rot, wp_scale, out_mat], device="cpu") - wp.synchronize() - - # Decompose - out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") - out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") - out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") - - wp.launch(_test_decompose_kernel, dim=1, inputs=[out_mat, out_pos, out_rot, out_scale], device="cpu") - wp.synchronize() - - np.testing.assert_allclose(out_pos.numpy()[0], pos[0], atol=1e-5) - np.testing.assert_allclose(out_scale.numpy()[0], scale[0], atol=1e-5) - # Quaternion sign ambiguity - r_out = out_rot.numpy()[0] - r_exp = rot[0] - dot = np.dot(r_out, r_exp) - np.testing.assert_allclose(abs(dot), 1.0, atol=1e-5) - - def test_non_uniform_scale_round_trip(self): - """Non-uniform scale round-trips correctly.""" - pos = np.array([[0.0, 0.0, 0.0]], dtype=np.float32) - rot = np.array([[0.0, 0.0, 0.0, 1.0]], dtype=np.float32) # identity - scale = np.array([[0.5, 2.0, 3.0]], dtype=np.float32) - - wp_pos = wp.array(pos, dtype=wp.vec3f, device="cpu") - wp_rot = wp.array(rot, dtype=wp.vec4f, device="cpu") - wp_scale = wp.array(scale, dtype=wp.vec3f, device="cpu") - out_mat = wp.zeros(1, dtype=wp.mat44f, device="cpu") - - wp.launch(_test_compose_kernel, dim=1, inputs=[wp_pos, wp_rot, wp_scale, out_mat], device="cpu") - wp.synchronize() - - out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") - out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") - out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") - - wp.launch(_test_decompose_kernel, dim=1, inputs=[out_mat, out_pos, out_rot, out_scale], device="cpu") - wp.synchronize() - - np.testing.assert_allclose(out_scale.numpy()[0], scale[0], atol=1e-5) - - -class TestKernelSignatures: - """Verify all exported kernels are importable and are Warp Kernels.""" - - def test_all_kernels_importable(self): - """All public kernels listed in __all__ should be importable and be Warp Kernels.""" - from isaaclab.utils.warp import fabric as fabric_utils - - expected_kernels = [ - "arange_k", - "compose_fabric_transformation_matrix_from_warp_arrays", - "compose_indexed_fabric_transforms", - "decompose_fabric_transformation_matrix_to_warp_arrays", - "decompose_indexed_fabric_transforms", - "set_view_to_fabric_array", - "update_indexed_local_matrix_from_world", - "update_indexed_world_matrix_from_local", - ] - - for name in expected_kernels: - obj = getattr(fabric_utils, name, None) - assert obj is not None, f"{name} not found in fabric_utils" - assert isinstance(obj, wp.Kernel), f"{name} should be a wp.Kernel, got {type(obj)}" - - def test_module_exports_match_all(self): - """__all__ should list every public kernel.""" - from isaaclab.utils.warp import fabric as fabric_utils - - for name in fabric_utils.__all__: - assert hasattr(fabric_utils, name), f"__all__ lists '{name}' but it's not defined" +# ------------------------------------------------------------------ +# Decompose / Compose round-trip tests +# ------------------------------------------------------------------ + + +def test_identity_matrix(): + """Identity matrix decomposes to pos=0, quat=identity, scale=1.""" + mat = np.eye(4, dtype=np.float32).reshape(1, 4, 4) + matrices = wp.array(mat, dtype=wp.mat44f, device="cpu") + out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") + out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") + out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") + + wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") + wp.synchronize() + + pos = out_pos.numpy() + rot = out_rot.numpy() + scale = out_scale.numpy() + + np.testing.assert_allclose(pos[0], [0, 0, 0], atol=1e-6) + np.testing.assert_allclose(scale[0], [1, 1, 1], atol=1e-6) + # Identity quaternion: either (0,0,0,1) or (0,0,0,-1) + assert abs(abs(rot[0, 3]) - 1.0) < 1e-5 + + +def test_translation_only(): + """Matrix with only translation decomposes correctly.""" + mat = np.eye(4, dtype=np.float32) + mat[3, 0] = 1.0 # row-major: translation in last row + mat[3, 1] = 2.0 + mat[3, 2] = 3.0 + matrices = wp.array(mat.reshape(1, 4, 4), dtype=wp.mat44f, device="cpu") + out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") + out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") + out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") + + wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") + wp.synchronize() + + np.testing.assert_allclose(out_pos.numpy()[0], [1, 2, 3], atol=1e-6) + np.testing.assert_allclose(out_scale.numpy()[0], [1, 1, 1], atol=1e-6) + + +def test_uniform_scale(): + """Matrix with uniform scale decomposes correctly.""" + mat = np.eye(4, dtype=np.float32) + mat[0, 0] = 2.0 + mat[1, 1] = 2.0 + mat[2, 2] = 2.0 + matrices = wp.array(mat.reshape(1, 4, 4), dtype=wp.mat44f, device="cpu") + out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") + out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") + out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") + + wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") + wp.synchronize() + + np.testing.assert_allclose(out_scale.numpy()[0], [2, 2, 2], atol=1e-6) + + +def test_round_trip(): + """Compose then decompose recovers original pos/quat/scale.""" + # Known transform: translate (5,6,7), rotate 90deg about Z, scale (1,2,3) + pos = np.array([[5.0, 6.0, 7.0]], dtype=np.float32) + # 90deg about Z in xyzw: (0, 0, sin(45), cos(45)) + s45 = np.sin(np.pi / 4) + c45 = np.cos(np.pi / 4) + rot = np.array([[0.0, 0.0, s45, c45]], dtype=np.float32) + scale = np.array([[1.0, 2.0, 3.0]], dtype=np.float32) + + wp_pos = wp.array(pos, dtype=wp.vec3f, device="cpu") + wp_rot = wp.array(rot, dtype=wp.vec4f, device="cpu") + wp_scale = wp.array(scale, dtype=wp.vec3f, device="cpu") + out_mat = wp.zeros(1, dtype=wp.mat44f, device="cpu") + + # Compose + wp.launch(_test_compose_kernel, dim=1, inputs=[wp_pos, wp_rot, wp_scale, out_mat], device="cpu") + wp.synchronize() + + # Decompose + out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") + out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") + out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") + + wp.launch(_test_decompose_kernel, dim=1, inputs=[out_mat, out_pos, out_rot, out_scale], device="cpu") + wp.synchronize() + + np.testing.assert_allclose(out_pos.numpy()[0], pos[0], atol=1e-5) + np.testing.assert_allclose(out_scale.numpy()[0], scale[0], atol=1e-5) + # Quaternion sign ambiguity + r_out = out_rot.numpy()[0] + r_exp = rot[0] + dot = np.dot(r_out, r_exp) + np.testing.assert_allclose(abs(dot), 1.0, atol=1e-5) + + +def test_non_uniform_scale_round_trip(): + """Non-uniform scale round-trips correctly.""" + pos = np.array([[0.0, 0.0, 0.0]], dtype=np.float32) + rot = np.array([[0.0, 0.0, 0.0, 1.0]], dtype=np.float32) # identity + scale = np.array([[0.5, 2.0, 3.0]], dtype=np.float32) + + wp_pos = wp.array(pos, dtype=wp.vec3f, device="cpu") + wp_rot = wp.array(rot, dtype=wp.vec4f, device="cpu") + wp_scale = wp.array(scale, dtype=wp.vec3f, device="cpu") + out_mat = wp.zeros(1, dtype=wp.mat44f, device="cpu") + + wp.launch(_test_compose_kernel, dim=1, inputs=[wp_pos, wp_rot, wp_scale, out_mat], device="cpu") + wp.synchronize() + + out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") + out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") + out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") + + wp.launch(_test_decompose_kernel, dim=1, inputs=[out_mat, out_pos, out_rot, out_scale], device="cpu") + wp.synchronize() + + np.testing.assert_allclose(out_scale.numpy()[0], scale[0], atol=1e-5) + + +# ------------------------------------------------------------------ +# Kernel signature / importability tests +# ------------------------------------------------------------------ + + +def test_all_kernels_importable(): + """All public kernels listed in __all__ should be importable and be Warp Kernels.""" + from isaaclab.utils.warp import fabric as fabric_utils + + expected_kernels = [ + "arange_k", + "compose_fabric_transformation_matrix_from_warp_arrays", + "compose_indexed_fabric_transforms", + "decompose_fabric_transformation_matrix_to_warp_arrays", + "decompose_indexed_fabric_transforms", + "set_view_to_fabric_array", + "update_indexed_local_matrix_from_world", + "update_indexed_world_matrix_from_local", + ] + + for name in expected_kernels: + obj = getattr(fabric_utils, name, None) + assert obj is not None, f"{name} not found in fabric_utils" + assert isinstance(obj, wp.Kernel), f"{name} should be a wp.Kernel, got {type(obj)}" + + +def test_module_exports_match_all(): + """__all__ should list every public kernel.""" + from isaaclab.utils.warp import fabric as fabric_utils + + for name in fabric_utils.__all__: + assert hasattr(fabric_utils, name), f"__all__ lists '{name}' but it's not defined" From 59bdf955c7059e2d1dd82ad7a8617a9c707252e3 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Thu, 21 May 2026 16:01:45 +0000 Subject: [PATCH 03/54] test: rewrite fabric kernel tests with wp.array (no Fabric runtime deps) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Replace indexedfabricarray hacks with plain wp.array(dtype=wp.mat44d) test kernels - Test kernels mirror production math: local^T = world^T * inv(parent^T) - Add 5 tests for world↔local transforms including non-orthogonal/sheared cases - All tests run on CPU without USDRT/Fabric runtime - Convert existing class-based tests to plain functions (IsaacLab guideline) --- .../test/utils/warp/test_fabric_kernels.py | 208 +++++++++++++++++- 1 file changed, 199 insertions(+), 9 deletions(-) diff --git a/source/isaaclab/test/utils/warp/test_fabric_kernels.py b/source/isaaclab/test/utils/warp/test_fabric_kernels.py index c2d3e6d37381..3c17ff2c0bb2 100644 --- a/source/isaaclab/test/utils/warp/test_fabric_kernels.py +++ b/source/isaaclab/test/utils/warp/test_fabric_kernels.py @@ -5,8 +5,11 @@ """Unit tests for the Warp fabric transform kernels. -Tests the decompose/compose math via helper kernels that operate on -regular wp.array (no Fabric/USDRT runtime required). +Tests the shared @wp.func math (decompose/compose and matrix inverse/multiply) +through plain wp.array kernels — no Fabric/USDRT runtime required. + +The production fabric kernels are thin adapters over the same math; testing the +math in isolation avoids coupling tests to Fabric container internals. """ import numpy as np @@ -17,6 +20,11 @@ from isaaclab.utils.warp.fabric import _decompose_transformation_matrix # noqa: E402 +# ------------------------------------------------------------------ +# Test kernels (wp.array wrappers around the same math as production) +# ------------------------------------------------------------------ + + @wp.kernel(enable_backward=False) def _test_decompose_kernel( matrices: wp.array(dtype=wp.mat44f), @@ -47,6 +55,57 @@ def _test_compose_kernel( out_matrices[i] = wp.transpose(wp.transform_compose(pos, rot, scale)) +@wp.kernel(enable_backward=False) +def _test_local_from_world_kernel( + child_world: wp.array(dtype=wp.mat44d), + parent_world: wp.array(dtype=wp.mat44d), + out_local: wp.array(dtype=wp.mat44d), +): + """Same math as update_indexed_local_matrix_from_world: local^T = world^T * inv(parent^T). + + Casts to mat44f for compute (matching production precision), writes back as mat44d. + """ + i = wp.tid() + cw = wp.mat44f(child_world[i]) + pw = wp.mat44f(parent_world[i]) + out_local[i] = wp.mat44d(cw * wp.inverse(pw)) + + +@wp.kernel(enable_backward=False) +def _test_world_from_local_kernel( + child_local: wp.array(dtype=wp.mat44d), + parent_world: wp.array(dtype=wp.mat44d), + out_world: wp.array(dtype=wp.mat44d), +): + """Same math as update_indexed_world_matrix_from_local: world^T = local^T * parent^T.""" + i = wp.tid() + cl = wp.mat44f(child_local[i]) + pw = wp.mat44f(parent_world[i]) + out_world[i] = wp.mat44d(cl * pw) + + +# ------------------------------------------------------------------ +# Helpers +# ------------------------------------------------------------------ + + +def _make_transform_matrix(pos, rot_quat_xyzw, scale): + """Build a 4x4 Fabric-transposed transform from pos/quat/scale. + + Returns numpy (4,4) float64 in the transposed storage convention (row-major with + translation in the last row) that Fabric uses. + """ + from scipy.spatial.transform import Rotation + + r = Rotation.from_quat(rot_quat_xyzw).as_matrix().astype(np.float64) + rs = r * np.array(scale, dtype=np.float64) + m = np.eye(4, dtype=np.float64) + m[:3, :3] = rs + m[:3, 3] = pos + # Transpose for Fabric storage convention + return m.T + + # ------------------------------------------------------------------ # Decompose / Compose round-trip tests # ------------------------------------------------------------------ @@ -76,7 +135,7 @@ def test_identity_matrix(): def test_translation_only(): """Matrix with only translation decomposes correctly.""" mat = np.eye(4, dtype=np.float32) - mat[3, 0] = 1.0 # row-major: translation in last row + mat[3, 0] = 1.0 # row-major transposed: translation in last row mat[3, 1] = 2.0 mat[3, 2] = 3.0 matrices = wp.array(mat.reshape(1, 4, 4), dtype=wp.mat44f, device="cpu") @@ -110,9 +169,7 @@ def test_uniform_scale(): def test_round_trip(): """Compose then decompose recovers original pos/quat/scale.""" - # Known transform: translate (5,6,7), rotate 90deg about Z, scale (1,2,3) pos = np.array([[5.0, 6.0, 7.0]], dtype=np.float32) - # 90deg about Z in xyzw: (0, 0, sin(45), cos(45)) s45 = np.sin(np.pi / 4) c45 = np.cos(np.pi / 4) rot = np.array([[0.0, 0.0, s45, c45]], dtype=np.float32) @@ -123,11 +180,9 @@ def test_round_trip(): wp_scale = wp.array(scale, dtype=wp.vec3f, device="cpu") out_mat = wp.zeros(1, dtype=wp.mat44f, device="cpu") - # Compose wp.launch(_test_compose_kernel, dim=1, inputs=[wp_pos, wp_rot, wp_scale, out_mat], device="cpu") wp.synchronize() - # Decompose out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") @@ -137,7 +192,6 @@ def test_round_trip(): np.testing.assert_allclose(out_pos.numpy()[0], pos[0], atol=1e-5) np.testing.assert_allclose(out_scale.numpy()[0], scale[0], atol=1e-5) - # Quaternion sign ambiguity r_out = out_rot.numpy()[0] r_exp = rot[0] dot = np.dot(r_out, r_exp) @@ -147,7 +201,7 @@ def test_round_trip(): def test_non_uniform_scale_round_trip(): """Non-uniform scale round-trips correctly.""" pos = np.array([[0.0, 0.0, 0.0]], dtype=np.float32) - rot = np.array([[0.0, 0.0, 0.0, 1.0]], dtype=np.float32) # identity + rot = np.array([[0.0, 0.0, 0.0, 1.0]], dtype=np.float32) scale = np.array([[0.5, 2.0, 3.0]], dtype=np.float32) wp_pos = wp.array(pos, dtype=wp.vec3f, device="cpu") @@ -168,6 +222,142 @@ def test_non_uniform_scale_round_trip(): np.testing.assert_allclose(out_scale.numpy()[0], scale[0], atol=1e-5) +# ------------------------------------------------------------------ +# World ↔ Local matrix tests +# +# These test the same math the production fabric kernels use: +# local^T = world^T * inv(parent^T) +# world^T = local^T * parent^T +# +# Under the transposed storage convention, this is equivalent to: +# local = inv(parent) * world +# world = parent * local +# ------------------------------------------------------------------ + + +def test_local_from_world_identity_parent(): + """With identity parent, local should equal world.""" + child_world_T = _make_transform_matrix([3, -1, 7], [0, 0, 0.3826834, 0.9238795], [1.5, 2.0, 0.5]) + parent_world_T = _make_transform_matrix([0, 0, 0], [0, 0, 0, 1], [1, 1, 1]) + + cw = wp.array(child_world_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") + pw = wp.array(parent_world_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") + out = wp.zeros(1, dtype=wp.mat44d, device="cpu") + + wp.launch(_test_local_from_world_kernel, dim=1, inputs=[cw, pw, out], device="cpu") + wp.synchronize() + + np.testing.assert_allclose(out.numpy()[0], child_world_T, atol=1e-5) + + +def test_local_from_world_non_orthogonal(): + """Non-uniform scale + rotation in parent produces non-orthogonal local matrix. + + Verifies: local^T @ parent^T == child_world^T (reconstruction check). + """ + parent_world_T = _make_transform_matrix([10, -5, 2], [0, 0.2588190, 0, 0.9659258], [2.0, 0.5, 3.0]) + child_world_T = _make_transform_matrix([1, 2, 3], [0.5, 0, 0, 0.8660254], [1.0, 1.5, 0.8]) + + cw = wp.array(child_world_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") + pw = wp.array(parent_world_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") + out = wp.zeros(1, dtype=wp.mat44d, device="cpu") + + wp.launch(_test_local_from_world_kernel, dim=1, inputs=[cw, pw, out], device="cpu") + wp.synchronize() + + # Verify reconstruction: local^T @ parent^T should equal child_world^T + local_T = out.numpy()[0] + reconstructed = local_T @ parent_world_T + np.testing.assert_allclose(reconstructed, child_world_T, atol=1e-5) + + +def test_world_from_local_non_orthogonal(): + """Recompose world from local with non-orthogonal parent.""" + parent_world_T = _make_transform_matrix([10, -5, 2], [0, 0.2588190, 0, 0.9659258], [2.0, 0.5, 3.0]) + child_world_T = _make_transform_matrix([1, 2, 3], [0.5, 0, 0, 0.8660254], [1.0, 1.5, 0.8]) + + # Ground-truth local + child_local_T = child_world_T @ np.linalg.inv(parent_world_T) + + cl = wp.array(child_local_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") + pw = wp.array(parent_world_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") + out = wp.zeros(1, dtype=wp.mat44d, device="cpu") + + wp.launch(_test_world_from_local_kernel, dim=1, inputs=[cl, pw, out], device="cpu") + wp.synchronize() + + np.testing.assert_allclose(out.numpy()[0], child_world_T, atol=1e-5) + + +def test_world_local_round_trip_non_orthogonal_batch(): + """Batch of 4 prims with non-orthogonal transforms: world->local->world round-trip.""" + from scipy.spatial.transform import Rotation + + rng = np.random.default_rng(42) + n = 4 + + parent_scales = rng.uniform(0.3, 3.0, size=(n, 3)).astype(np.float64) + child_scales = rng.uniform(0.3, 3.0, size=(n, 3)).astype(np.float64) + parent_rots = Rotation.random(n, random_state=42).as_quat().astype(np.float64) + child_rots = Rotation.random(n, random_state=99).as_quat().astype(np.float64) + parent_positions = rng.uniform(-10, 10, size=(n, 3)).astype(np.float64) + child_positions = rng.uniform(-10, 10, size=(n, 3)).astype(np.float64) + + parent_world_Ts = np.stack( + [_make_transform_matrix(parent_positions[i], parent_rots[i], parent_scales[i]) for i in range(n)] + ) + child_world_Ts = np.stack( + [_make_transform_matrix(child_positions[i], child_rots[i], child_scales[i]) for i in range(n)] + ) + + # world -> local + cw = wp.array(child_world_Ts, dtype=wp.mat44d, device="cpu") + pw = wp.array(parent_world_Ts, dtype=wp.mat44d, device="cpu") + local_out = wp.zeros(n, dtype=wp.mat44d, device="cpu") + + wp.launch(_test_local_from_world_kernel, dim=n, inputs=[cw, pw, local_out], device="cpu") + wp.synchronize() + + # local -> world (round-trip) + cl = wp.array(local_out.numpy(), dtype=wp.mat44d, device="cpu") + pw2 = wp.array(parent_world_Ts, dtype=wp.mat44d, device="cpu") + world_out = wp.zeros(n, dtype=wp.mat44d, device="cpu") + + wp.launch(_test_world_from_local_kernel, dim=n, inputs=[cl, pw2, world_out], device="cpu") + wp.synchronize() + + np.testing.assert_allclose(world_out.numpy(), child_world_Ts, atol=1e-4) + + +def test_local_from_world_sheared_parent(): + """Parent with extreme non-uniform scale (10:1 ratio) creating significant shear. + + Verifies both correctness and that the resulting local matrix is genuinely + non-orthogonal (the whole point of testing with sheared transforms). + """ + parent_world_T = _make_transform_matrix([0, 0, 0], [0, 0, 0.3826834, 0.9238795], [10.0, 1.0, 1.0]) + child_world_T = _make_transform_matrix([5, 5, 0], [0, 0, 0, 1], [1.0, 1.0, 1.0]) + + cw = wp.array(child_world_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") + pw = wp.array(parent_world_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") + out = wp.zeros(1, dtype=wp.mat44d, device="cpu") + + wp.launch(_test_local_from_world_kernel, dim=1, inputs=[cw, pw, out], device="cpu") + wp.synchronize() + + # Verify reconstruction + local_T = out.numpy()[0] + reconstructed = local_T @ parent_world_T + np.testing.assert_allclose(reconstructed, child_world_T, atol=1e-4) + + # Verify the local matrix upper-3x3 is NOT orthogonal + upper3x3 = local_T[:3, :3] + gram = upper3x3 @ upper3x3.T + assert not np.allclose(gram, np.eye(3), atol=0.1), ( + "Local matrix upper-3x3 should NOT be orthogonal with 10:1 sheared parent" + ) + + # ------------------------------------------------------------------ # Kernel signature / importability tests # ------------------------------------------------------------------ From d6a65aff4aa3655775b2ecc9588b25842942ab53 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Thu, 21 May 2026 16:12:43 +0000 Subject: [PATCH 04/54] =?UTF-8?q?refactor:=20extract=20world=E2=86=94local?= =?UTF-8?q?=20math=20into=20@wp.func=20for=20testability?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add _local_from_world() and _world_from_local() as @wp.func - Production fabric kernels delegate to these shared funcs - Test kernels import and call the same funcs via wp.array adapters - Zero copy-pasted math: tests exercise identical code paths as production --- source/isaaclab/isaaclab/utils/warp/fabric.py | 20 ++++++++++++++-- .../test/utils/warp/test_fabric_kernels.py | 23 ++++++++----------- 2 files changed, 28 insertions(+), 15 deletions(-) diff --git a/source/isaaclab/isaaclab/utils/warp/fabric.py b/source/isaaclab/isaaclab/utils/warp/fabric.py index 6f9963f290a1..aad13d4daab3 100644 --- a/source/isaaclab/isaaclab/utils/warp/fabric.py +++ b/source/isaaclab/isaaclab/utils/warp/fabric.py @@ -261,6 +261,18 @@ def compose_indexed_fabric_transforms( ) +@wp.func +def _local_from_world(child_world: wp.mat44f, parent_world: wp.mat44f) -> wp.mat44f: + """Compute local^T = world^T * inv(parent^T) on transposed storage matrices.""" + return child_world * wp.inverse(parent_world) + + +@wp.func +def _world_from_local(child_local: wp.mat44f, parent_world: wp.mat44f) -> wp.mat44f: + """Compute world^T = local^T * parent^T on transposed storage matrices.""" + return child_local * parent_world + + @wp.kernel(enable_backward=False) def update_indexed_local_matrix_from_world( child_world_matrices: IndexedFabricArrayMat44d, @@ -294,7 +306,9 @@ def update_indexed_local_matrix_from_world( 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(child_world * wp.inverse(parent_world)) # type: ignore[arg-type] + child_local_matrices[view_index] = wp.mat44d( # type: ignore[arg-type] + _local_from_world(child_world, parent_world) + ) @wp.kernel(enable_backward=False) @@ -325,7 +339,9 @@ def update_indexed_world_matrix_from_local( 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(child_local * parent_world) # type: ignore[arg-type] + child_world_matrices[view_index] = wp.mat44d( # type: ignore[arg-type] + _world_from_local(child_local, parent_world) + ) @wp.func diff --git a/source/isaaclab/test/utils/warp/test_fabric_kernels.py b/source/isaaclab/test/utils/warp/test_fabric_kernels.py index 3c17ff2c0bb2..4d303dc36cb6 100644 --- a/source/isaaclab/test/utils/warp/test_fabric_kernels.py +++ b/source/isaaclab/test/utils/warp/test_fabric_kernels.py @@ -17,11 +17,15 @@ wp.init() -from isaaclab.utils.warp.fabric import _decompose_transformation_matrix # noqa: E402 +from isaaclab.utils.warp.fabric import ( # noqa: E402 + _decompose_transformation_matrix, + _local_from_world, + _world_from_local, +) # ------------------------------------------------------------------ -# Test kernels (wp.array wrappers around the same math as production) +# Test kernels — thin wp.array wrappers that delegate to production @wp.func # ------------------------------------------------------------------ @@ -61,14 +65,9 @@ def _test_local_from_world_kernel( parent_world: wp.array(dtype=wp.mat44d), out_local: wp.array(dtype=wp.mat44d), ): - """Same math as update_indexed_local_matrix_from_world: local^T = world^T * inv(parent^T). - - Casts to mat44f for compute (matching production precision), writes back as mat44d. - """ + """wp.array adapter for _local_from_world — same func as production fabric kernel.""" i = wp.tid() - cw = wp.mat44f(child_world[i]) - pw = wp.mat44f(parent_world[i]) - out_local[i] = wp.mat44d(cw * wp.inverse(pw)) + out_local[i] = wp.mat44d(_local_from_world(wp.mat44f(child_world[i]), wp.mat44f(parent_world[i]))) @wp.kernel(enable_backward=False) @@ -77,11 +76,9 @@ def _test_world_from_local_kernel( parent_world: wp.array(dtype=wp.mat44d), out_world: wp.array(dtype=wp.mat44d), ): - """Same math as update_indexed_world_matrix_from_local: world^T = local^T * parent^T.""" + """wp.array adapter for _world_from_local — same func as production fabric kernel.""" i = wp.tid() - cl = wp.mat44f(child_local[i]) - pw = wp.mat44f(parent_world[i]) - out_world[i] = wp.mat44d(cl * pw) + out_world[i] = wp.mat44d(_world_from_local(wp.mat44f(child_local[i]), wp.mat44f(parent_world[i]))) # ------------------------------------------------------------------ From 0c52abcea75c1d06c907409db60b018f86869e74 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Thu, 21 May 2026 16:57:50 +0000 Subject: [PATCH 05/54] =?UTF-8?q?refactor:=20extract=20world=E2=86=94local?= =?UTF-8?q?=20math=20into=20@wp.func=20for=20testability?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add _local_from_world_transposed() and _world_from_local_transposed() as @wp.func (names make explicit they operate on transposed storage matrices) - Production fabric kernels delegate to these shared funcs - Test kernels import and call the same funcs via wp.array adapters - Remove __all__ (not an IsaacLab convention) - Remove importability/export tests (unnecessary) - Remove batch test (Warp handles batching) - Add inline det != 0 assertions to world↔local tests --- source/isaaclab/isaaclab/utils/warp/fabric.py | 23 +- .../test/utils/warp/test_fabric_kernels.py | 309 +++--------------- 2 files changed, 51 insertions(+), 281 deletions(-) diff --git a/source/isaaclab/isaaclab/utils/warp/fabric.py b/source/isaaclab/isaaclab/utils/warp/fabric.py index aad13d4daab3..c5a86a58b292 100644 --- a/source/isaaclab/isaaclab/utils/warp/fabric.py +++ b/source/isaaclab/isaaclab/utils/warp/fabric.py @@ -15,17 +15,6 @@ import warp as wp -__all__ = [ - "arange_k", - "compose_fabric_transformation_matrix_from_warp_arrays", - "compose_indexed_fabric_transforms", - "decompose_fabric_transformation_matrix_to_warp_arrays", - "decompose_indexed_fabric_transforms", - "set_view_to_fabric_array", - "update_indexed_local_matrix_from_world", - "update_indexed_world_matrix_from_local", -] - if TYPE_CHECKING: FabricArrayUInt32 = Any FabricArrayMat44d = Any @@ -262,15 +251,15 @@ def compose_indexed_fabric_transforms( @wp.func -def _local_from_world(child_world: wp.mat44f, parent_world: wp.mat44f) -> wp.mat44f: +def _local_from_world_transposed(child_world_T: wp.mat44f, parent_world_T: wp.mat44f) -> wp.mat44f: """Compute local^T = world^T * inv(parent^T) on transposed storage matrices.""" - return child_world * wp.inverse(parent_world) + return child_world_T * wp.inverse(parent_world_T) @wp.func -def _world_from_local(child_local: wp.mat44f, parent_world: wp.mat44f) -> wp.mat44f: +def _world_from_local_transposed(child_local_T: wp.mat44f, parent_world_T: wp.mat44f) -> wp.mat44f: """Compute world^T = local^T * parent^T on transposed storage matrices.""" - return child_local * parent_world + return child_local_T * parent_world_T @wp.kernel(enable_backward=False) @@ -307,7 +296,7 @@ def update_indexed_local_matrix_from_world( 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] - _local_from_world(child_world, parent_world) + _local_from_world_transposed(child_world, parent_world) ) @@ -340,7 +329,7 @@ def update_indexed_world_matrix_from_local( 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] - _world_from_local(child_local, parent_world) + _world_from_local_transposed(child_local, parent_world) ) diff --git a/source/isaaclab/test/utils/warp/test_fabric_kernels.py b/source/isaaclab/test/utils/warp/test_fabric_kernels.py index 4d303dc36cb6..9153838e4366 100644 --- a/source/isaaclab/test/utils/warp/test_fabric_kernels.py +++ b/source/isaaclab/test/utils/warp/test_fabric_kernels.py @@ -19,8 +19,8 @@ from isaaclab.utils.warp.fabric import ( # noqa: E402 _decompose_transformation_matrix, - _local_from_world, - _world_from_local, + _local_from_world_transposed, + _world_from_local_transposed, ) @@ -33,41 +33,26 @@ def _test_decompose_kernel( matrices: wp.array(dtype=wp.mat44f), out_positions: wp.array(dtype=wp.vec3f), - out_rotations: wp.array(dtype=wp.vec4f), + out_rotations: wp.array(dtype=wp.quatf), out_scales: wp.array(dtype=wp.vec3f), ): """Decompose a batch of 4x4 matrices into pos/quat/scale.""" i = wp.tid() pos, rot, scale = _decompose_transformation_matrix(matrices[i]) out_positions[i] = pos - out_rotations[i] = wp.vec4f(rot[0], rot[1], rot[2], rot[3]) + out_rotations[i] = rot out_scales[i] = scale -@wp.kernel(enable_backward=False) -def _test_compose_kernel( - positions: wp.array(dtype=wp.vec3f), - rotations: wp.array(dtype=wp.vec4f), - scales: wp.array(dtype=wp.vec3f), - out_matrices: wp.array(dtype=wp.mat44f), -): - """Compose a batch of pos/quat/scale into 4x4 matrices.""" - i = wp.tid() - pos = positions[i] - rot = wp.quatf(rotations[i][0], rotations[i][1], rotations[i][2], rotations[i][3]) - scale = scales[i] - out_matrices[i] = wp.transpose(wp.transform_compose(pos, rot, scale)) - - @wp.kernel(enable_backward=False) def _test_local_from_world_kernel( child_world: wp.array(dtype=wp.mat44d), parent_world: wp.array(dtype=wp.mat44d), out_local: wp.array(dtype=wp.mat44d), ): - """wp.array adapter for _local_from_world — same func as production fabric kernel.""" + """wp.array adapter for _local_from_world_transposed — same func as production fabric kernel.""" i = wp.tid() - out_local[i] = wp.mat44d(_local_from_world(wp.mat44f(child_world[i]), wp.mat44f(parent_world[i]))) + out_local[i] = wp.mat44d(_local_from_world_transposed(wp.mat44f(child_world[i]), wp.mat44f(parent_world[i]))) @wp.kernel(enable_backward=False) @@ -76,9 +61,9 @@ def _test_world_from_local_kernel( parent_world: wp.array(dtype=wp.mat44d), out_world: wp.array(dtype=wp.mat44d), ): - """wp.array adapter for _world_from_local — same func as production fabric kernel.""" + """wp.array adapter for _world_from_local_transposed — same func as production fabric kernel.""" i = wp.tid() - out_world[i] = wp.mat44d(_world_from_local(wp.mat44f(child_local[i]), wp.mat44f(parent_world[i]))) + out_world[i] = wp.mat44d(_world_from_local_transposed(wp.mat44f(child_local[i]), wp.mat44f(parent_world[i]))) # ------------------------------------------------------------------ @@ -91,6 +76,9 @@ def _make_transform_matrix(pos, rot_quat_xyzw, scale): Returns numpy (4,4) float64 in the transposed storage convention (row-major with translation in the last row) that Fabric uses. + + Raises: + AssertionError: If the resulting matrix is singular (e.g. zero scale component). """ from scipy.spatial.transform import Rotation @@ -100,7 +88,10 @@ def _make_transform_matrix(pos, rot_quat_xyzw, scale): m[:3, :3] = rs m[:3, 3] = pos # Transpose for Fabric storage convention - return m.T + result = m.T + det = np.linalg.det(result) + assert abs(det) > 1e-6, f"Singular matrix: det={det:.2e}, scale={scale}" + return result # ------------------------------------------------------------------ @@ -108,117 +99,30 @@ def _make_transform_matrix(pos, rot_quat_xyzw, scale): # ------------------------------------------------------------------ -def test_identity_matrix(): - """Identity matrix decomposes to pos=0, quat=identity, scale=1.""" - mat = np.eye(4, dtype=np.float32).reshape(1, 4, 4) - matrices = wp.array(mat, dtype=wp.mat44f, device="cpu") - out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") - out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") - out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") - - wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") - wp.synchronize() - - pos = out_pos.numpy() - rot = out_rot.numpy() - scale = out_scale.numpy() - - np.testing.assert_allclose(pos[0], [0, 0, 0], atol=1e-6) - np.testing.assert_allclose(scale[0], [1, 1, 1], atol=1e-6) - # Identity quaternion: either (0,0,0,1) or (0,0,0,-1) - assert abs(abs(rot[0, 3]) - 1.0) < 1e-5 - - -def test_translation_only(): - """Matrix with only translation decomposes correctly.""" - mat = np.eye(4, dtype=np.float32) - mat[3, 0] = 1.0 # row-major transposed: translation in last row - mat[3, 1] = 2.0 - mat[3, 2] = 3.0 - matrices = wp.array(mat.reshape(1, 4, 4), dtype=wp.mat44f, device="cpu") - out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") - out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") - out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") - - wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") - wp.synchronize() - - np.testing.assert_allclose(out_pos.numpy()[0], [1, 2, 3], atol=1e-6) - np.testing.assert_allclose(out_scale.numpy()[0], [1, 1, 1], atol=1e-6) - - -def test_uniform_scale(): - """Matrix with uniform scale decomposes correctly.""" - mat = np.eye(4, dtype=np.float32) - mat[0, 0] = 2.0 - mat[1, 1] = 2.0 - mat[2, 2] = 2.0 - matrices = wp.array(mat.reshape(1, 4, 4), dtype=wp.mat44f, device="cpu") - out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") - out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") - out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") - - wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") - wp.synchronize() - - np.testing.assert_allclose(out_scale.numpy()[0], [2, 2, 2], atol=1e-6) - - -def test_round_trip(): - """Compose then decompose recovers original pos/quat/scale.""" - pos = np.array([[5.0, 6.0, 7.0]], dtype=np.float32) +def test_decompose_round_trip(): + """Decompose a matrix with translation, rotation, and non-uniform scale; verify round-trip.""" + pos = np.array([5.0, -3.0, 7.0]) s45 = np.sin(np.pi / 4) c45 = np.cos(np.pi / 4) - rot = np.array([[0.0, 0.0, s45, c45]], dtype=np.float32) - scale = np.array([[1.0, 2.0, 3.0]], dtype=np.float32) - - wp_pos = wp.array(pos, dtype=wp.vec3f, device="cpu") - wp_rot = wp.array(rot, dtype=wp.vec4f, device="cpu") - wp_scale = wp.array(scale, dtype=wp.vec3f, device="cpu") - out_mat = wp.zeros(1, dtype=wp.mat44f, device="cpu") + quat_xyzw = np.array([0.0, 0.0, s45, c45]) # 45° Z rotation + scale = np.array([1.5, 0.8, 3.0]) - wp.launch(_test_compose_kernel, dim=1, inputs=[wp_pos, wp_rot, wp_scale, out_mat], device="cpu") - wp.synchronize() + mat = _make_transform_matrix(pos, quat_xyzw, scale).astype(np.float32) + matrices = wp.array(mat.reshape(1, 4, 4), dtype=wp.mat44f, device="cpu") out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") - out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") + out_rot = wp.zeros(1, dtype=wp.quatf, device="cpu") out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") - wp.launch(_test_decompose_kernel, dim=1, inputs=[out_mat, out_pos, out_rot, out_scale], device="cpu") + wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") wp.synchronize() - np.testing.assert_allclose(out_pos.numpy()[0], pos[0], atol=1e-5) - np.testing.assert_allclose(out_scale.numpy()[0], scale[0], atol=1e-5) - r_out = out_rot.numpy()[0] - r_exp = rot[0] - dot = np.dot(r_out, r_exp) + np.testing.assert_allclose(out_pos.numpy()[0], pos, atol=1e-5) + np.testing.assert_allclose(out_scale.numpy()[0], scale, atol=1e-5) + dot = np.dot(out_rot.numpy()[0], quat_xyzw) np.testing.assert_allclose(abs(dot), 1.0, atol=1e-5) -def test_non_uniform_scale_round_trip(): - """Non-uniform scale round-trips correctly.""" - pos = np.array([[0.0, 0.0, 0.0]], dtype=np.float32) - rot = np.array([[0.0, 0.0, 0.0, 1.0]], dtype=np.float32) - scale = np.array([[0.5, 2.0, 3.0]], dtype=np.float32) - - wp_pos = wp.array(pos, dtype=wp.vec3f, device="cpu") - wp_rot = wp.array(rot, dtype=wp.vec4f, device="cpu") - wp_scale = wp.array(scale, dtype=wp.vec3f, device="cpu") - out_mat = wp.zeros(1, dtype=wp.mat44f, device="cpu") - - wp.launch(_test_compose_kernel, dim=1, inputs=[wp_pos, wp_rot, wp_scale, out_mat], device="cpu") - wp.synchronize() - - out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") - out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") - out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") - - wp.launch(_test_decompose_kernel, dim=1, inputs=[out_mat, out_pos, out_rot, out_scale], device="cpu") - wp.synchronize() - - np.testing.assert_allclose(out_scale.numpy()[0], scale[0], atol=1e-5) - - # ------------------------------------------------------------------ # World ↔ Local matrix tests # @@ -226,164 +130,41 @@ def test_non_uniform_scale_round_trip(): # local^T = world^T * inv(parent^T) # world^T = local^T * parent^T # -# Under the transposed storage convention, this is equivalent to: -# local = inv(parent) * world -# world = parent * local +# Both parent and child have rotation, translation, and non-uniform scale +# (producing sheared/non-orthogonal upper-3x3 blocks). # ------------------------------------------------------------------ +# Shared test data: parent with 10:1 non-uniform scale + 45° Z rotation + translation +_PARENT_WORLD_T = _make_transform_matrix([10, -5, 2], [0, 0, 0.3826834, 0.9238795], [4.0, 0.5, 2.0]) +_CHILD_WORLD_T = _make_transform_matrix([1, 2, 3], [0.2588190, 0, 0, 0.9659258], [1.5, 0.8, 3.0]) -def test_local_from_world_identity_parent(): - """With identity parent, local should equal world.""" - child_world_T = _make_transform_matrix([3, -1, 7], [0, 0, 0.3826834, 0.9238795], [1.5, 2.0, 0.5]) - parent_world_T = _make_transform_matrix([0, 0, 0], [0, 0, 0, 1], [1, 1, 1]) - cw = wp.array(child_world_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") - pw = wp.array(parent_world_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") +def test_local_from_world_transposed(): + """local^T = world^T * inv(parent^T) — verified by reconstruction.""" + cw = wp.array(_CHILD_WORLD_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") + pw = wp.array(_PARENT_WORLD_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") out = wp.zeros(1, dtype=wp.mat44d, device="cpu") wp.launch(_test_local_from_world_kernel, dim=1, inputs=[cw, pw, out], device="cpu") wp.synchronize() - np.testing.assert_allclose(out.numpy()[0], child_world_T, atol=1e-5) - - -def test_local_from_world_non_orthogonal(): - """Non-uniform scale + rotation in parent produces non-orthogonal local matrix. - - Verifies: local^T @ parent^T == child_world^T (reconstruction check). - """ - parent_world_T = _make_transform_matrix([10, -5, 2], [0, 0.2588190, 0, 0.9659258], [2.0, 0.5, 3.0]) - child_world_T = _make_transform_matrix([1, 2, 3], [0.5, 0, 0, 0.8660254], [1.0, 1.5, 0.8]) - - cw = wp.array(child_world_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") - pw = wp.array(parent_world_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") - out = wp.zeros(1, dtype=wp.mat44d, device="cpu") - - wp.launch(_test_local_from_world_kernel, dim=1, inputs=[cw, pw, out], device="cpu") - wp.synchronize() - - # Verify reconstruction: local^T @ parent^T should equal child_world^T + # Reconstruction: local^T @ parent^T must equal child_world^T local_T = out.numpy()[0] - reconstructed = local_T @ parent_world_T - np.testing.assert_allclose(reconstructed, child_world_T, atol=1e-5) - + reconstructed = local_T @ _PARENT_WORLD_T + np.testing.assert_allclose(reconstructed, _CHILD_WORLD_T, atol=1e-5) -def test_world_from_local_non_orthogonal(): - """Recompose world from local with non-orthogonal parent.""" - parent_world_T = _make_transform_matrix([10, -5, 2], [0, 0.2588190, 0, 0.9659258], [2.0, 0.5, 3.0]) - child_world_T = _make_transform_matrix([1, 2, 3], [0.5, 0, 0, 0.8660254], [1.0, 1.5, 0.8]) - # Ground-truth local - child_local_T = child_world_T @ np.linalg.inv(parent_world_T) +def test_world_from_local_transposed(): + """world^T = local^T * parent^T — verified against known child world.""" + # Ground-truth local computed via numpy + child_local_T = _CHILD_WORLD_T @ np.linalg.inv(_PARENT_WORLD_T) cl = wp.array(child_local_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") - pw = wp.array(parent_world_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") + pw = wp.array(_PARENT_WORLD_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") out = wp.zeros(1, dtype=wp.mat44d, device="cpu") wp.launch(_test_world_from_local_kernel, dim=1, inputs=[cl, pw, out], device="cpu") wp.synchronize() - np.testing.assert_allclose(out.numpy()[0], child_world_T, atol=1e-5) - - -def test_world_local_round_trip_non_orthogonal_batch(): - """Batch of 4 prims with non-orthogonal transforms: world->local->world round-trip.""" - from scipy.spatial.transform import Rotation - - rng = np.random.default_rng(42) - n = 4 - - parent_scales = rng.uniform(0.3, 3.0, size=(n, 3)).astype(np.float64) - child_scales = rng.uniform(0.3, 3.0, size=(n, 3)).astype(np.float64) - parent_rots = Rotation.random(n, random_state=42).as_quat().astype(np.float64) - child_rots = Rotation.random(n, random_state=99).as_quat().astype(np.float64) - parent_positions = rng.uniform(-10, 10, size=(n, 3)).astype(np.float64) - child_positions = rng.uniform(-10, 10, size=(n, 3)).astype(np.float64) - - parent_world_Ts = np.stack( - [_make_transform_matrix(parent_positions[i], parent_rots[i], parent_scales[i]) for i in range(n)] - ) - child_world_Ts = np.stack( - [_make_transform_matrix(child_positions[i], child_rots[i], child_scales[i]) for i in range(n)] - ) - - # world -> local - cw = wp.array(child_world_Ts, dtype=wp.mat44d, device="cpu") - pw = wp.array(parent_world_Ts, dtype=wp.mat44d, device="cpu") - local_out = wp.zeros(n, dtype=wp.mat44d, device="cpu") - - wp.launch(_test_local_from_world_kernel, dim=n, inputs=[cw, pw, local_out], device="cpu") - wp.synchronize() - - # local -> world (round-trip) - cl = wp.array(local_out.numpy(), dtype=wp.mat44d, device="cpu") - pw2 = wp.array(parent_world_Ts, dtype=wp.mat44d, device="cpu") - world_out = wp.zeros(n, dtype=wp.mat44d, device="cpu") - - wp.launch(_test_world_from_local_kernel, dim=n, inputs=[cl, pw2, world_out], device="cpu") - wp.synchronize() - - np.testing.assert_allclose(world_out.numpy(), child_world_Ts, atol=1e-4) - - -def test_local_from_world_sheared_parent(): - """Parent with extreme non-uniform scale (10:1 ratio) creating significant shear. - - Verifies both correctness and that the resulting local matrix is genuinely - non-orthogonal (the whole point of testing with sheared transforms). - """ - parent_world_T = _make_transform_matrix([0, 0, 0], [0, 0, 0.3826834, 0.9238795], [10.0, 1.0, 1.0]) - child_world_T = _make_transform_matrix([5, 5, 0], [0, 0, 0, 1], [1.0, 1.0, 1.0]) - - cw = wp.array(child_world_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") - pw = wp.array(parent_world_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") - out = wp.zeros(1, dtype=wp.mat44d, device="cpu") - - wp.launch(_test_local_from_world_kernel, dim=1, inputs=[cw, pw, out], device="cpu") - wp.synchronize() - - # Verify reconstruction - local_T = out.numpy()[0] - reconstructed = local_T @ parent_world_T - np.testing.assert_allclose(reconstructed, child_world_T, atol=1e-4) - - # Verify the local matrix upper-3x3 is NOT orthogonal - upper3x3 = local_T[:3, :3] - gram = upper3x3 @ upper3x3.T - assert not np.allclose(gram, np.eye(3), atol=0.1), ( - "Local matrix upper-3x3 should NOT be orthogonal with 10:1 sheared parent" - ) - - -# ------------------------------------------------------------------ -# Kernel signature / importability tests -# ------------------------------------------------------------------ - - -def test_all_kernels_importable(): - """All public kernels listed in __all__ should be importable and be Warp Kernels.""" - from isaaclab.utils.warp import fabric as fabric_utils - - expected_kernels = [ - "arange_k", - "compose_fabric_transformation_matrix_from_warp_arrays", - "compose_indexed_fabric_transforms", - "decompose_fabric_transformation_matrix_to_warp_arrays", - "decompose_indexed_fabric_transforms", - "set_view_to_fabric_array", - "update_indexed_local_matrix_from_world", - "update_indexed_world_matrix_from_local", - ] - - for name in expected_kernels: - obj = getattr(fabric_utils, name, None) - assert obj is not None, f"{name} not found in fabric_utils" - assert isinstance(obj, wp.Kernel), f"{name} should be a wp.Kernel, got {type(obj)}" - - -def test_module_exports_match_all(): - """__all__ should list every public kernel.""" - from isaaclab.utils.warp import fabric as fabric_utils + np.testing.assert_allclose(out.numpy()[0], _CHILD_WORLD_T, atol=1e-5) - for name in fabric_utils.__all__: - assert hasattr(fabric_utils, name), f"__all__ lists '{name}' but it's not defined" From 936ae76fb56f3a0325ecc64a6244e2857c2cc381 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Fri, 22 May 2026 14:10:38 +0000 Subject: [PATCH 06/54] refactor: inline @wp.func helpers, add __all__, sync with full-stack MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Remove _local_from_world_transposed / _world_from_local_transposed @wp.func helpers (one-liner math, abstraction added no value) - Add __all__ export list for public API surface - Sync test file with full-stack state - Changelog: 'Will be used by' → 'Used by' --- .../changelog.d/indexed-fabric-kernels.rst | 2 +- source/isaaclab/isaaclab/utils/warp/fabric.py | 31 +- .../test/utils/warp/test_fabric_kernels.py | 299 ++++++++++-------- 3 files changed, 175 insertions(+), 157 deletions(-) diff --git a/source/isaaclab/changelog.d/indexed-fabric-kernels.rst b/source/isaaclab/changelog.d/indexed-fabric-kernels.rst index 0f68b4c60e9d..281881bb7972 100644 --- a/source/isaaclab/changelog.d/indexed-fabric-kernels.rst +++ b/source/isaaclab/changelog.d/indexed-fabric-kernels.rst @@ -13,7 +13,7 @@ Added 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 (no - explicit transposes). Will be used by + explicit transposes). Used by :class:`~isaaclab_physx.sim.views.FabricFrameView` to keep child world and local matrices consistent across writes without round-tripping through USD. diff --git a/source/isaaclab/isaaclab/utils/warp/fabric.py b/source/isaaclab/isaaclab/utils/warp/fabric.py index c5a86a58b292..6f9963f290a1 100644 --- a/source/isaaclab/isaaclab/utils/warp/fabric.py +++ b/source/isaaclab/isaaclab/utils/warp/fabric.py @@ -15,6 +15,17 @@ import warp as wp +__all__ = [ + "arange_k", + "compose_fabric_transformation_matrix_from_warp_arrays", + "compose_indexed_fabric_transforms", + "decompose_fabric_transformation_matrix_to_warp_arrays", + "decompose_indexed_fabric_transforms", + "set_view_to_fabric_array", + "update_indexed_local_matrix_from_world", + "update_indexed_world_matrix_from_local", +] + if TYPE_CHECKING: FabricArrayUInt32 = Any FabricArrayMat44d = Any @@ -250,18 +261,6 @@ def compose_indexed_fabric_transforms( ) -@wp.func -def _local_from_world_transposed(child_world_T: wp.mat44f, parent_world_T: wp.mat44f) -> wp.mat44f: - """Compute local^T = world^T * inv(parent^T) on transposed storage matrices.""" - return child_world_T * wp.inverse(parent_world_T) - - -@wp.func -def _world_from_local_transposed(child_local_T: wp.mat44f, parent_world_T: wp.mat44f) -> wp.mat44f: - """Compute world^T = local^T * parent^T on transposed storage matrices.""" - return child_local_T * parent_world_T - - @wp.kernel(enable_backward=False) def update_indexed_local_matrix_from_world( child_world_matrices: IndexedFabricArrayMat44d, @@ -295,9 +294,7 @@ def update_indexed_local_matrix_from_world( 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] - _local_from_world_transposed(child_world, parent_world) - ) + child_local_matrices[view_index] = wp.mat44d(child_world * wp.inverse(parent_world)) # type: ignore[arg-type] @wp.kernel(enable_backward=False) @@ -328,9 +325,7 @@ def update_indexed_world_matrix_from_local( 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] - _world_from_local_transposed(child_local, parent_world) - ) + child_world_matrices[view_index] = wp.mat44d(child_local * parent_world) # type: ignore[arg-type] @wp.func diff --git a/source/isaaclab/test/utils/warp/test_fabric_kernels.py b/source/isaaclab/test/utils/warp/test_fabric_kernels.py index 9153838e4366..13bad50da500 100644 --- a/source/isaaclab/test/utils/warp/test_fabric_kernels.py +++ b/source/isaaclab/test/utils/warp/test_fabric_kernels.py @@ -5,11 +5,8 @@ """Unit tests for the Warp fabric transform kernels. -Tests the shared @wp.func math (decompose/compose and matrix inverse/multiply) -through plain wp.array kernels — no Fabric/USDRT runtime required. - -The production fabric kernels are thin adapters over the same math; testing the -math in isolation avoids coupling tests to Fabric container internals. +Tests the decompose/compose math via helper kernels that operate on +regular wp.array (no Fabric/USDRT runtime required). """ import numpy as np @@ -17,154 +14,180 @@ wp.init() -from isaaclab.utils.warp.fabric import ( # noqa: E402 - _decompose_transformation_matrix, - _local_from_world_transposed, - _world_from_local_transposed, -) - - -# ------------------------------------------------------------------ -# Test kernels — thin wp.array wrappers that delegate to production @wp.func -# ------------------------------------------------------------------ +from isaaclab.utils.warp.fabric import _decompose_transformation_matrix # noqa: E402 @wp.kernel(enable_backward=False) def _test_decompose_kernel( matrices: wp.array(dtype=wp.mat44f), out_positions: wp.array(dtype=wp.vec3f), - out_rotations: wp.array(dtype=wp.quatf), + out_rotations: wp.array(dtype=wp.vec4f), out_scales: wp.array(dtype=wp.vec3f), ): """Decompose a batch of 4x4 matrices into pos/quat/scale.""" i = wp.tid() pos, rot, scale = _decompose_transformation_matrix(matrices[i]) out_positions[i] = pos - out_rotations[i] = rot + out_rotations[i] = wp.vec4f(rot[0], rot[1], rot[2], rot[3]) out_scales[i] = scale @wp.kernel(enable_backward=False) -def _test_local_from_world_kernel( - child_world: wp.array(dtype=wp.mat44d), - parent_world: wp.array(dtype=wp.mat44d), - out_local: wp.array(dtype=wp.mat44d), -): - """wp.array adapter for _local_from_world_transposed — same func as production fabric kernel.""" - i = wp.tid() - out_local[i] = wp.mat44d(_local_from_world_transposed(wp.mat44f(child_world[i]), wp.mat44f(parent_world[i]))) - - -@wp.kernel(enable_backward=False) -def _test_world_from_local_kernel( - child_local: wp.array(dtype=wp.mat44d), - parent_world: wp.array(dtype=wp.mat44d), - out_world: wp.array(dtype=wp.mat44d), +def _test_compose_kernel( + positions: wp.array(dtype=wp.vec3f), + rotations: wp.array(dtype=wp.vec4f), + scales: wp.array(dtype=wp.vec3f), + out_matrices: wp.array(dtype=wp.mat44f), ): - """wp.array adapter for _world_from_local_transposed — same func as production fabric kernel.""" + """Compose a batch of pos/quat/scale into 4x4 matrices.""" i = wp.tid() - out_world[i] = wp.mat44d(_world_from_local_transposed(wp.mat44f(child_local[i]), wp.mat44f(parent_world[i]))) - - -# ------------------------------------------------------------------ -# Helpers -# ------------------------------------------------------------------ - - -def _make_transform_matrix(pos, rot_quat_xyzw, scale): - """Build a 4x4 Fabric-transposed transform from pos/quat/scale. - - Returns numpy (4,4) float64 in the transposed storage convention (row-major with - translation in the last row) that Fabric uses. - - Raises: - AssertionError: If the resulting matrix is singular (e.g. zero scale component). - """ - from scipy.spatial.transform import Rotation - - r = Rotation.from_quat(rot_quat_xyzw).as_matrix().astype(np.float64) - rs = r * np.array(scale, dtype=np.float64) - m = np.eye(4, dtype=np.float64) - m[:3, :3] = rs - m[:3, 3] = pos - # Transpose for Fabric storage convention - result = m.T - det = np.linalg.det(result) - assert abs(det) > 1e-6, f"Singular matrix: det={det:.2e}, scale={scale}" - return result - - -# ------------------------------------------------------------------ -# Decompose / Compose round-trip tests -# ------------------------------------------------------------------ - - -def test_decompose_round_trip(): - """Decompose a matrix with translation, rotation, and non-uniform scale; verify round-trip.""" - pos = np.array([5.0, -3.0, 7.0]) - s45 = np.sin(np.pi / 4) - c45 = np.cos(np.pi / 4) - quat_xyzw = np.array([0.0, 0.0, s45, c45]) # 45° Z rotation - scale = np.array([1.5, 0.8, 3.0]) - - mat = _make_transform_matrix(pos, quat_xyzw, scale).astype(np.float32) - matrices = wp.array(mat.reshape(1, 4, 4), dtype=wp.mat44f, device="cpu") - - out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") - out_rot = wp.zeros(1, dtype=wp.quatf, device="cpu") - out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") - - wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") - wp.synchronize() - - np.testing.assert_allclose(out_pos.numpy()[0], pos, atol=1e-5) - np.testing.assert_allclose(out_scale.numpy()[0], scale, atol=1e-5) - dot = np.dot(out_rot.numpy()[0], quat_xyzw) - np.testing.assert_allclose(abs(dot), 1.0, atol=1e-5) - - -# ------------------------------------------------------------------ -# World ↔ Local matrix tests -# -# These test the same math the production fabric kernels use: -# local^T = world^T * inv(parent^T) -# world^T = local^T * parent^T -# -# Both parent and child have rotation, translation, and non-uniform scale -# (producing sheared/non-orthogonal upper-3x3 blocks). -# ------------------------------------------------------------------ - -# Shared test data: parent with 10:1 non-uniform scale + 45° Z rotation + translation -_PARENT_WORLD_T = _make_transform_matrix([10, -5, 2], [0, 0, 0.3826834, 0.9238795], [4.0, 0.5, 2.0]) -_CHILD_WORLD_T = _make_transform_matrix([1, 2, 3], [0.2588190, 0, 0, 0.9659258], [1.5, 0.8, 3.0]) - - -def test_local_from_world_transposed(): - """local^T = world^T * inv(parent^T) — verified by reconstruction.""" - cw = wp.array(_CHILD_WORLD_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") - pw = wp.array(_PARENT_WORLD_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") - out = wp.zeros(1, dtype=wp.mat44d, device="cpu") - - wp.launch(_test_local_from_world_kernel, dim=1, inputs=[cw, pw, out], device="cpu") - wp.synchronize() - - # Reconstruction: local^T @ parent^T must equal child_world^T - local_T = out.numpy()[0] - reconstructed = local_T @ _PARENT_WORLD_T - np.testing.assert_allclose(reconstructed, _CHILD_WORLD_T, atol=1e-5) - - -def test_world_from_local_transposed(): - """world^T = local^T * parent^T — verified against known child world.""" - # Ground-truth local computed via numpy - child_local_T = _CHILD_WORLD_T @ np.linalg.inv(_PARENT_WORLD_T) - - cl = wp.array(child_local_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") - pw = wp.array(_PARENT_WORLD_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") - out = wp.zeros(1, dtype=wp.mat44d, device="cpu") - - wp.launch(_test_world_from_local_kernel, dim=1, inputs=[cl, pw, out], device="cpu") - wp.synchronize() - - np.testing.assert_allclose(out.numpy()[0], _CHILD_WORLD_T, atol=1e-5) - + pos = positions[i] + rot = wp.quatf(rotations[i][0], rotations[i][1], rotations[i][2], rotations[i][3]) + scale = scales[i] + out_matrices[i] = wp.transpose(wp.transform_compose(pos, rot, scale)) + + +class TestDecomposeCompose: + """Round-trip tests for decompose ↔ compose transform math.""" + + def test_identity_matrix(self): + """Identity matrix decomposes to pos=0, quat=identity, scale=1.""" + mat = np.eye(4, dtype=np.float32).reshape(1, 4, 4) + matrices = wp.array(mat, dtype=wp.mat44f, device="cpu") + out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") + out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") + out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") + + wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") + wp.synchronize() + + pos = out_pos.numpy() + rot = out_rot.numpy() + scale = out_scale.numpy() + + np.testing.assert_allclose(pos[0], [0, 0, 0], atol=1e-6) + np.testing.assert_allclose(scale[0], [1, 1, 1], atol=1e-6) + # Identity quaternion: either (0,0,0,1) or (0,0,0,-1) + assert abs(abs(rot[0, 3]) - 1.0) < 1e-5 + + def test_translation_only(self): + """Matrix with only translation decomposes correctly.""" + mat = np.eye(4, dtype=np.float32) + mat[3, 0] = 1.0 # row-major: translation in last row + mat[3, 1] = 2.0 + mat[3, 2] = 3.0 + matrices = wp.array(mat.reshape(1, 4, 4), dtype=wp.mat44f, device="cpu") + out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") + out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") + out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") + + wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") + wp.synchronize() + + np.testing.assert_allclose(out_pos.numpy()[0], [1, 2, 3], atol=1e-6) + np.testing.assert_allclose(out_scale.numpy()[0], [1, 1, 1], atol=1e-6) + + def test_uniform_scale(self): + """Matrix with uniform scale decomposes correctly.""" + mat = np.eye(4, dtype=np.float32) + mat[0, 0] = 2.0 + mat[1, 1] = 2.0 + mat[2, 2] = 2.0 + matrices = wp.array(mat.reshape(1, 4, 4), dtype=wp.mat44f, device="cpu") + out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") + out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") + out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") + + wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") + wp.synchronize() + + np.testing.assert_allclose(out_scale.numpy()[0], [2, 2, 2], atol=1e-6) + + def test_round_trip(self): + """Compose then decompose recovers original pos/quat/scale.""" + # Known transform: translate (5,6,7), rotate 90° about Z, scale (1,2,3) + pos = np.array([[5.0, 6.0, 7.0]], dtype=np.float32) + # 90° about Z in xyzw: (0, 0, sin(45°), cos(45°)) + s45 = np.sin(np.pi / 4) + c45 = np.cos(np.pi / 4) + rot = np.array([[0.0, 0.0, s45, c45]], dtype=np.float32) + scale = np.array([[1.0, 2.0, 3.0]], dtype=np.float32) + + wp_pos = wp.array(pos, dtype=wp.vec3f, device="cpu") + wp_rot = wp.array(rot, dtype=wp.vec4f, device="cpu") + wp_scale = wp.array(scale, dtype=wp.vec3f, device="cpu") + out_mat = wp.zeros(1, dtype=wp.mat44f, device="cpu") + + # Compose + wp.launch(_test_compose_kernel, dim=1, inputs=[wp_pos, wp_rot, wp_scale, out_mat], device="cpu") + wp.synchronize() + + # Decompose + out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") + out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") + out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") + + wp.launch(_test_decompose_kernel, dim=1, inputs=[out_mat, out_pos, out_rot, out_scale], device="cpu") + wp.synchronize() + + np.testing.assert_allclose(out_pos.numpy()[0], pos[0], atol=1e-5) + np.testing.assert_allclose(out_scale.numpy()[0], scale[0], atol=1e-5) + # Quaternion sign ambiguity + r_out = out_rot.numpy()[0] + r_exp = rot[0] + dot = np.dot(r_out, r_exp) + np.testing.assert_allclose(abs(dot), 1.0, atol=1e-5) + + def test_non_uniform_scale_round_trip(self): + """Non-uniform scale round-trips correctly.""" + pos = np.array([[0.0, 0.0, 0.0]], dtype=np.float32) + rot = np.array([[0.0, 0.0, 0.0, 1.0]], dtype=np.float32) # identity + scale = np.array([[0.5, 2.0, 3.0]], dtype=np.float32) + + wp_pos = wp.array(pos, dtype=wp.vec3f, device="cpu") + wp_rot = wp.array(rot, dtype=wp.vec4f, device="cpu") + wp_scale = wp.array(scale, dtype=wp.vec3f, device="cpu") + out_mat = wp.zeros(1, dtype=wp.mat44f, device="cpu") + + wp.launch(_test_compose_kernel, dim=1, inputs=[wp_pos, wp_rot, wp_scale, out_mat], device="cpu") + wp.synchronize() + + out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") + out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") + out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") + + wp.launch(_test_decompose_kernel, dim=1, inputs=[out_mat, out_pos, out_rot, out_scale], device="cpu") + wp.synchronize() + + np.testing.assert_allclose(out_scale.numpy()[0], scale[0], atol=1e-5) + + +class TestKernelSignatures: + """Verify all exported kernels are importable and are Warp Kernels.""" + + def test_all_kernels_importable(self): + """All public kernels listed in __all__ should be importable and be Warp Kernels.""" + from isaaclab.utils.warp import fabric as fabric_utils + + expected_kernels = [ + "arange_k", + "compose_fabric_transformation_matrix_from_warp_arrays", + "compose_indexed_fabric_transforms", + "decompose_fabric_transformation_matrix_to_warp_arrays", + "decompose_indexed_fabric_transforms", + "set_view_to_fabric_array", + "update_indexed_local_matrix_from_world", + "update_indexed_world_matrix_from_local", + ] + + for name in expected_kernels: + obj = getattr(fabric_utils, name, None) + assert obj is not None, f"{name} not found in fabric_utils" + assert isinstance(obj, wp.Kernel), f"{name} should be a wp.Kernel, got {type(obj)}" + + def test_module_exports_match_all(self): + """__all__ should list every public kernel.""" + from isaaclab.utils.warp import fabric as fabric_utils + + for name in fabric_utils.__all__: + assert hasattr(fabric_utils, name), f"__all__ lists '{name}' but it's not defined" From d1275259b2fdb3f0e50c7ad8841868e95bbca0d6 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Fri, 22 May 2026 14:13:57 +0000 Subject: [PATCH 07/54] Revert "refactor: inline @wp.func helpers, add __all__, sync with full-stack" This reverts commit 735132e0d896e3c8182152abeb372c4dfe113870. --- .../changelog.d/indexed-fabric-kernels.rst | 2 +- source/isaaclab/isaaclab/utils/warp/fabric.py | 31 +- .../test/utils/warp/test_fabric_kernels.py | 299 ++++++++---------- 3 files changed, 157 insertions(+), 175 deletions(-) diff --git a/source/isaaclab/changelog.d/indexed-fabric-kernels.rst b/source/isaaclab/changelog.d/indexed-fabric-kernels.rst index 281881bb7972..0f68b4c60e9d 100644 --- a/source/isaaclab/changelog.d/indexed-fabric-kernels.rst +++ b/source/isaaclab/changelog.d/indexed-fabric-kernels.rst @@ -13,7 +13,7 @@ Added 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 (no - explicit transposes). Used by + explicit transposes). Will be used by :class:`~isaaclab_physx.sim.views.FabricFrameView` to keep child world and local matrices consistent across writes without round-tripping through USD. diff --git a/source/isaaclab/isaaclab/utils/warp/fabric.py b/source/isaaclab/isaaclab/utils/warp/fabric.py index 6f9963f290a1..c5a86a58b292 100644 --- a/source/isaaclab/isaaclab/utils/warp/fabric.py +++ b/source/isaaclab/isaaclab/utils/warp/fabric.py @@ -15,17 +15,6 @@ import warp as wp -__all__ = [ - "arange_k", - "compose_fabric_transformation_matrix_from_warp_arrays", - "compose_indexed_fabric_transforms", - "decompose_fabric_transformation_matrix_to_warp_arrays", - "decompose_indexed_fabric_transforms", - "set_view_to_fabric_array", - "update_indexed_local_matrix_from_world", - "update_indexed_world_matrix_from_local", -] - if TYPE_CHECKING: FabricArrayUInt32 = Any FabricArrayMat44d = Any @@ -261,6 +250,18 @@ def compose_indexed_fabric_transforms( ) +@wp.func +def _local_from_world_transposed(child_world_T: wp.mat44f, parent_world_T: wp.mat44f) -> wp.mat44f: + """Compute local^T = world^T * inv(parent^T) on transposed storage matrices.""" + return child_world_T * wp.inverse(parent_world_T) + + +@wp.func +def _world_from_local_transposed(child_local_T: wp.mat44f, parent_world_T: wp.mat44f) -> wp.mat44f: + """Compute world^T = local^T * parent^T on transposed storage matrices.""" + return child_local_T * parent_world_T + + @wp.kernel(enable_backward=False) def update_indexed_local_matrix_from_world( child_world_matrices: IndexedFabricArrayMat44d, @@ -294,7 +295,9 @@ def update_indexed_local_matrix_from_world( 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(child_world * wp.inverse(parent_world)) # type: ignore[arg-type] + child_local_matrices[view_index] = wp.mat44d( # type: ignore[arg-type] + _local_from_world_transposed(child_world, parent_world) + ) @wp.kernel(enable_backward=False) @@ -325,7 +328,9 @@ def update_indexed_world_matrix_from_local( 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(child_local * parent_world) # type: ignore[arg-type] + child_world_matrices[view_index] = wp.mat44d( # type: ignore[arg-type] + _world_from_local_transposed(child_local, parent_world) + ) @wp.func diff --git a/source/isaaclab/test/utils/warp/test_fabric_kernels.py b/source/isaaclab/test/utils/warp/test_fabric_kernels.py index 13bad50da500..9153838e4366 100644 --- a/source/isaaclab/test/utils/warp/test_fabric_kernels.py +++ b/source/isaaclab/test/utils/warp/test_fabric_kernels.py @@ -5,8 +5,11 @@ """Unit tests for the Warp fabric transform kernels. -Tests the decompose/compose math via helper kernels that operate on -regular wp.array (no Fabric/USDRT runtime required). +Tests the shared @wp.func math (decompose/compose and matrix inverse/multiply) +through plain wp.array kernels — no Fabric/USDRT runtime required. + +The production fabric kernels are thin adapters over the same math; testing the +math in isolation avoids coupling tests to Fabric container internals. """ import numpy as np @@ -14,180 +17,154 @@ wp.init() -from isaaclab.utils.warp.fabric import _decompose_transformation_matrix # noqa: E402 +from isaaclab.utils.warp.fabric import ( # noqa: E402 + _decompose_transformation_matrix, + _local_from_world_transposed, + _world_from_local_transposed, +) + + +# ------------------------------------------------------------------ +# Test kernels — thin wp.array wrappers that delegate to production @wp.func +# ------------------------------------------------------------------ @wp.kernel(enable_backward=False) def _test_decompose_kernel( matrices: wp.array(dtype=wp.mat44f), out_positions: wp.array(dtype=wp.vec3f), - out_rotations: wp.array(dtype=wp.vec4f), + out_rotations: wp.array(dtype=wp.quatf), out_scales: wp.array(dtype=wp.vec3f), ): """Decompose a batch of 4x4 matrices into pos/quat/scale.""" i = wp.tid() pos, rot, scale = _decompose_transformation_matrix(matrices[i]) out_positions[i] = pos - out_rotations[i] = wp.vec4f(rot[0], rot[1], rot[2], rot[3]) + out_rotations[i] = rot out_scales[i] = scale @wp.kernel(enable_backward=False) -def _test_compose_kernel( - positions: wp.array(dtype=wp.vec3f), - rotations: wp.array(dtype=wp.vec4f), - scales: wp.array(dtype=wp.vec3f), - out_matrices: wp.array(dtype=wp.mat44f), +def _test_local_from_world_kernel( + child_world: wp.array(dtype=wp.mat44d), + parent_world: wp.array(dtype=wp.mat44d), + out_local: wp.array(dtype=wp.mat44d), +): + """wp.array adapter for _local_from_world_transposed — same func as production fabric kernel.""" + i = wp.tid() + out_local[i] = wp.mat44d(_local_from_world_transposed(wp.mat44f(child_world[i]), wp.mat44f(parent_world[i]))) + + +@wp.kernel(enable_backward=False) +def _test_world_from_local_kernel( + child_local: wp.array(dtype=wp.mat44d), + parent_world: wp.array(dtype=wp.mat44d), + out_world: wp.array(dtype=wp.mat44d), ): - """Compose a batch of pos/quat/scale into 4x4 matrices.""" + """wp.array adapter for _world_from_local_transposed — same func as production fabric kernel.""" i = wp.tid() - pos = positions[i] - rot = wp.quatf(rotations[i][0], rotations[i][1], rotations[i][2], rotations[i][3]) - scale = scales[i] - out_matrices[i] = wp.transpose(wp.transform_compose(pos, rot, scale)) - - -class TestDecomposeCompose: - """Round-trip tests for decompose ↔ compose transform math.""" - - def test_identity_matrix(self): - """Identity matrix decomposes to pos=0, quat=identity, scale=1.""" - mat = np.eye(4, dtype=np.float32).reshape(1, 4, 4) - matrices = wp.array(mat, dtype=wp.mat44f, device="cpu") - out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") - out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") - out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") - - wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") - wp.synchronize() - - pos = out_pos.numpy() - rot = out_rot.numpy() - scale = out_scale.numpy() - - np.testing.assert_allclose(pos[0], [0, 0, 0], atol=1e-6) - np.testing.assert_allclose(scale[0], [1, 1, 1], atol=1e-6) - # Identity quaternion: either (0,0,0,1) or (0,0,0,-1) - assert abs(abs(rot[0, 3]) - 1.0) < 1e-5 - - def test_translation_only(self): - """Matrix with only translation decomposes correctly.""" - mat = np.eye(4, dtype=np.float32) - mat[3, 0] = 1.0 # row-major: translation in last row - mat[3, 1] = 2.0 - mat[3, 2] = 3.0 - matrices = wp.array(mat.reshape(1, 4, 4), dtype=wp.mat44f, device="cpu") - out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") - out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") - out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") - - wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") - wp.synchronize() - - np.testing.assert_allclose(out_pos.numpy()[0], [1, 2, 3], atol=1e-6) - np.testing.assert_allclose(out_scale.numpy()[0], [1, 1, 1], atol=1e-6) - - def test_uniform_scale(self): - """Matrix with uniform scale decomposes correctly.""" - mat = np.eye(4, dtype=np.float32) - mat[0, 0] = 2.0 - mat[1, 1] = 2.0 - mat[2, 2] = 2.0 - matrices = wp.array(mat.reshape(1, 4, 4), dtype=wp.mat44f, device="cpu") - out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") - out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") - out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") - - wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") - wp.synchronize() - - np.testing.assert_allclose(out_scale.numpy()[0], [2, 2, 2], atol=1e-6) - - def test_round_trip(self): - """Compose then decompose recovers original pos/quat/scale.""" - # Known transform: translate (5,6,7), rotate 90° about Z, scale (1,2,3) - pos = np.array([[5.0, 6.0, 7.0]], dtype=np.float32) - # 90° about Z in xyzw: (0, 0, sin(45°), cos(45°)) - s45 = np.sin(np.pi / 4) - c45 = np.cos(np.pi / 4) - rot = np.array([[0.0, 0.0, s45, c45]], dtype=np.float32) - scale = np.array([[1.0, 2.0, 3.0]], dtype=np.float32) - - wp_pos = wp.array(pos, dtype=wp.vec3f, device="cpu") - wp_rot = wp.array(rot, dtype=wp.vec4f, device="cpu") - wp_scale = wp.array(scale, dtype=wp.vec3f, device="cpu") - out_mat = wp.zeros(1, dtype=wp.mat44f, device="cpu") - - # Compose - wp.launch(_test_compose_kernel, dim=1, inputs=[wp_pos, wp_rot, wp_scale, out_mat], device="cpu") - wp.synchronize() - - # Decompose - out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") - out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") - out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") - - wp.launch(_test_decompose_kernel, dim=1, inputs=[out_mat, out_pos, out_rot, out_scale], device="cpu") - wp.synchronize() - - np.testing.assert_allclose(out_pos.numpy()[0], pos[0], atol=1e-5) - np.testing.assert_allclose(out_scale.numpy()[0], scale[0], atol=1e-5) - # Quaternion sign ambiguity - r_out = out_rot.numpy()[0] - r_exp = rot[0] - dot = np.dot(r_out, r_exp) - np.testing.assert_allclose(abs(dot), 1.0, atol=1e-5) - - def test_non_uniform_scale_round_trip(self): - """Non-uniform scale round-trips correctly.""" - pos = np.array([[0.0, 0.0, 0.0]], dtype=np.float32) - rot = np.array([[0.0, 0.0, 0.0, 1.0]], dtype=np.float32) # identity - scale = np.array([[0.5, 2.0, 3.0]], dtype=np.float32) - - wp_pos = wp.array(pos, dtype=wp.vec3f, device="cpu") - wp_rot = wp.array(rot, dtype=wp.vec4f, device="cpu") - wp_scale = wp.array(scale, dtype=wp.vec3f, device="cpu") - out_mat = wp.zeros(1, dtype=wp.mat44f, device="cpu") - - wp.launch(_test_compose_kernel, dim=1, inputs=[wp_pos, wp_rot, wp_scale, out_mat], device="cpu") - wp.synchronize() - - out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") - out_rot = wp.zeros(1, dtype=wp.vec4f, device="cpu") - out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") - - wp.launch(_test_decompose_kernel, dim=1, inputs=[out_mat, out_pos, out_rot, out_scale], device="cpu") - wp.synchronize() - - np.testing.assert_allclose(out_scale.numpy()[0], scale[0], atol=1e-5) - - -class TestKernelSignatures: - """Verify all exported kernels are importable and are Warp Kernels.""" - - def test_all_kernels_importable(self): - """All public kernels listed in __all__ should be importable and be Warp Kernels.""" - from isaaclab.utils.warp import fabric as fabric_utils - - expected_kernels = [ - "arange_k", - "compose_fabric_transformation_matrix_from_warp_arrays", - "compose_indexed_fabric_transforms", - "decompose_fabric_transformation_matrix_to_warp_arrays", - "decompose_indexed_fabric_transforms", - "set_view_to_fabric_array", - "update_indexed_local_matrix_from_world", - "update_indexed_world_matrix_from_local", - ] - - for name in expected_kernels: - obj = getattr(fabric_utils, name, None) - assert obj is not None, f"{name} not found in fabric_utils" - assert isinstance(obj, wp.Kernel), f"{name} should be a wp.Kernel, got {type(obj)}" - - def test_module_exports_match_all(self): - """__all__ should list every public kernel.""" - from isaaclab.utils.warp import fabric as fabric_utils - - for name in fabric_utils.__all__: - assert hasattr(fabric_utils, name), f"__all__ lists '{name}' but it's not defined" + out_world[i] = wp.mat44d(_world_from_local_transposed(wp.mat44f(child_local[i]), wp.mat44f(parent_world[i]))) + + +# ------------------------------------------------------------------ +# Helpers +# ------------------------------------------------------------------ + + +def _make_transform_matrix(pos, rot_quat_xyzw, scale): + """Build a 4x4 Fabric-transposed transform from pos/quat/scale. + + Returns numpy (4,4) float64 in the transposed storage convention (row-major with + translation in the last row) that Fabric uses. + + Raises: + AssertionError: If the resulting matrix is singular (e.g. zero scale component). + """ + from scipy.spatial.transform import Rotation + + r = Rotation.from_quat(rot_quat_xyzw).as_matrix().astype(np.float64) + rs = r * np.array(scale, dtype=np.float64) + m = np.eye(4, dtype=np.float64) + m[:3, :3] = rs + m[:3, 3] = pos + # Transpose for Fabric storage convention + result = m.T + det = np.linalg.det(result) + assert abs(det) > 1e-6, f"Singular matrix: det={det:.2e}, scale={scale}" + return result + + +# ------------------------------------------------------------------ +# Decompose / Compose round-trip tests +# ------------------------------------------------------------------ + + +def test_decompose_round_trip(): + """Decompose a matrix with translation, rotation, and non-uniform scale; verify round-trip.""" + pos = np.array([5.0, -3.0, 7.0]) + s45 = np.sin(np.pi / 4) + c45 = np.cos(np.pi / 4) + quat_xyzw = np.array([0.0, 0.0, s45, c45]) # 45° Z rotation + scale = np.array([1.5, 0.8, 3.0]) + + mat = _make_transform_matrix(pos, quat_xyzw, scale).astype(np.float32) + matrices = wp.array(mat.reshape(1, 4, 4), dtype=wp.mat44f, device="cpu") + + out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") + out_rot = wp.zeros(1, dtype=wp.quatf, device="cpu") + out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") + + wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") + wp.synchronize() + + np.testing.assert_allclose(out_pos.numpy()[0], pos, atol=1e-5) + np.testing.assert_allclose(out_scale.numpy()[0], scale, atol=1e-5) + dot = np.dot(out_rot.numpy()[0], quat_xyzw) + np.testing.assert_allclose(abs(dot), 1.0, atol=1e-5) + + +# ------------------------------------------------------------------ +# World ↔ Local matrix tests +# +# These test the same math the production fabric kernels use: +# local^T = world^T * inv(parent^T) +# world^T = local^T * parent^T +# +# Both parent and child have rotation, translation, and non-uniform scale +# (producing sheared/non-orthogonal upper-3x3 blocks). +# ------------------------------------------------------------------ + +# Shared test data: parent with 10:1 non-uniform scale + 45° Z rotation + translation +_PARENT_WORLD_T = _make_transform_matrix([10, -5, 2], [0, 0, 0.3826834, 0.9238795], [4.0, 0.5, 2.0]) +_CHILD_WORLD_T = _make_transform_matrix([1, 2, 3], [0.2588190, 0, 0, 0.9659258], [1.5, 0.8, 3.0]) + + +def test_local_from_world_transposed(): + """local^T = world^T * inv(parent^T) — verified by reconstruction.""" + cw = wp.array(_CHILD_WORLD_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") + pw = wp.array(_PARENT_WORLD_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") + out = wp.zeros(1, dtype=wp.mat44d, device="cpu") + + wp.launch(_test_local_from_world_kernel, dim=1, inputs=[cw, pw, out], device="cpu") + wp.synchronize() + + # Reconstruction: local^T @ parent^T must equal child_world^T + local_T = out.numpy()[0] + reconstructed = local_T @ _PARENT_WORLD_T + np.testing.assert_allclose(reconstructed, _CHILD_WORLD_T, atol=1e-5) + + +def test_world_from_local_transposed(): + """world^T = local^T * parent^T — verified against known child world.""" + # Ground-truth local computed via numpy + child_local_T = _CHILD_WORLD_T @ np.linalg.inv(_PARENT_WORLD_T) + + cl = wp.array(child_local_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") + pw = wp.array(_PARENT_WORLD_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") + out = wp.zeros(1, dtype=wp.mat44d, device="cpu") + + wp.launch(_test_world_from_local_kernel, dim=1, inputs=[cl, pw, out], device="cpu") + wp.synchronize() + + np.testing.assert_allclose(out.numpy()[0], _CHILD_WORLD_T, atol=1e-5) + From e795df9be2d102edb06b9038e17ac9bc987c183f Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Fri, 22 May 2026 14:14:05 +0000 Subject: [PATCH 08/54] =?UTF-8?q?docs:=20update=20changelog=20wording=20('?= =?UTF-8?q?Will=20be=20used=20by'=20=E2=86=92=20'Used=20by')?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- source/isaaclab/changelog.d/indexed-fabric-kernels.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/isaaclab/changelog.d/indexed-fabric-kernels.rst b/source/isaaclab/changelog.d/indexed-fabric-kernels.rst index 0f68b4c60e9d..281881bb7972 100644 --- a/source/isaaclab/changelog.d/indexed-fabric-kernels.rst +++ b/source/isaaclab/changelog.d/indexed-fabric-kernels.rst @@ -13,7 +13,7 @@ Added 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 (no - explicit transposes). Will be used by + explicit transposes). Used by :class:`~isaaclab_physx.sim.views.FabricFrameView` to keep child world and local matrices consistent across writes without round-tripping through USD. From 0af88a64795566c3187ad4876704f5aba893c653 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Sat, 23 May 2026 11:03:02 +0000 Subject: [PATCH 09/54] =?UTF-8?q?docs:=20fix=20changelog=20wording=20(Used?= =?UTF-8?q?=20by=20=E2=86=92=20Will=20be=20used=20by)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- source/isaaclab/changelog.d/indexed-fabric-kernels.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/isaaclab/changelog.d/indexed-fabric-kernels.rst b/source/isaaclab/changelog.d/indexed-fabric-kernels.rst index 281881bb7972..0f68b4c60e9d 100644 --- a/source/isaaclab/changelog.d/indexed-fabric-kernels.rst +++ b/source/isaaclab/changelog.d/indexed-fabric-kernels.rst @@ -13,7 +13,7 @@ Added 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 (no - explicit transposes). Used by + explicit transposes). Will be used by :class:`~isaaclab_physx.sim.views.FabricFrameView` to keep child world and local matrices consistent across writes without round-tripping through USD. From b4672525d9500bf8b01d378f61f21db3d60d060c Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Sat, 23 May 2026 11:10:28 +0000 Subject: [PATCH 10/54] revert: remove wp.where refactor (unrelated to this PR) --- .../changelog.d/indexed-fabric-kernels.rst | 6 ---- source/isaaclab/isaaclab/utils/warp/fabric.py | 30 +++++++++++++++---- 2 files changed, 24 insertions(+), 12 deletions(-) diff --git a/source/isaaclab/changelog.d/indexed-fabric-kernels.rst b/source/isaaclab/changelog.d/indexed-fabric-kernels.rst index 0f68b4c60e9d..baaf33ceb8a3 100644 --- a/source/isaaclab/changelog.d/indexed-fabric-kernels.rst +++ b/source/isaaclab/changelog.d/indexed-fabric-kernels.rst @@ -16,9 +16,3 @@ Added explicit transposes). Will be used by :class:`~isaaclab_physx.sim.views.FabricFrameView` to keep child world and local matrices consistent across writes without round-tripping through USD. - -Changed -^^^^^^^ - -* Replaced ``if/else`` branching with ``wp.where`` in existing Fabric - compose/decompose kernels for branchless GPU execution. diff --git a/source/isaaclab/isaaclab/utils/warp/fabric.py b/source/isaaclab/isaaclab/utils/warp/fabric.py index c5a86a58b292..f665323cbe34 100644 --- a/source/isaaclab/isaaclab/utils/warp/fabric.py +++ b/source/isaaclab/isaaclab/utils/warp/fabric.py @@ -132,20 +132,29 @@ def compose_fabric_transformation_matrix_from_warp_arrays( position, rotation, scale = _decompose_transformation_matrix(wp.mat44f(fabric_matrices[fabric_index])) # update position (check if array has elements, not just if it exists) if array_positions.shape[0] > 0: - index = wp.where(broadcast_positions, 0, i) + 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] # update orientation (convert from wxyz to xyzw for Warp) if array_orientations.shape[0] > 0: - index = wp.where(broadcast_orientations, 0, i) + if broadcast_orientations: + index = 0 + else: + index = i rotation[0] = array_orientations[index, 0] # x rotation[1] = array_orientations[index, 1] # y rotation[2] = array_orientations[index, 2] # z rotation[3] = array_orientations[index, 3] # w # update scale if array_scales.shape[0] > 0: - index = wp.where(broadcast_scales, 0, i) + 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] @@ -229,18 +238,27 @@ def compose_indexed_fabric_transforms( position, rotation, scale = _decompose_transformation_matrix(wp.mat44f(fabric_matrices[view_index])) if array_positions.shape[0] > 0: - index = wp.where(broadcast_positions, 0, i) + 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: - index = wp.where(broadcast_orientations, 0, i) + 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: - index = wp.where(broadcast_scales, 0, i) + 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] From db069811c2e88a87e0f2e445805b7b88c832b761 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Sat, 23 May 2026 17:21:27 +0000 Subject: [PATCH 11/54] test: replace scipy with warp builtins in fabric kernel test helper --- .../test/utils/warp/test_fabric_kernels.py | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/source/isaaclab/test/utils/warp/test_fabric_kernels.py b/source/isaaclab/test/utils/warp/test_fabric_kernels.py index 9153838e4366..de48afadab24 100644 --- a/source/isaaclab/test/utils/warp/test_fabric_kernels.py +++ b/source/isaaclab/test/utils/warp/test_fabric_kernels.py @@ -23,7 +23,6 @@ _world_from_local_transposed, ) - # ------------------------------------------------------------------ # Test kernels — thin wp.array wrappers that delegate to production @wp.func # ------------------------------------------------------------------ @@ -80,15 +79,11 @@ def _make_transform_matrix(pos, rot_quat_xyzw, scale): Raises: AssertionError: If the resulting matrix is singular (e.g. zero scale component). """ - from scipy.spatial.transform import Rotation - - r = Rotation.from_quat(rot_quat_xyzw).as_matrix().astype(np.float64) - rs = r * np.array(scale, dtype=np.float64) - m = np.eye(4, dtype=np.float64) - m[:3, :3] = rs - m[:3, 3] = pos - # Transpose for Fabric storage convention - result = m.T + p = wp.vec3f(*pos) + q = wp.quatf(*rot_quat_xyzw) + s = wp.vec3f(*scale) + m = wp.transpose(wp.transform_compose(p, q, s)) + result = np.array(m).reshape(4, 4).astype(np.float64) det = np.linalg.det(result) assert abs(det) > 1e-6, f"Singular matrix: det={det:.2e}, scale={scale}" return result @@ -167,4 +162,3 @@ def test_world_from_local_transposed(): wp.synchronize() np.testing.assert_allclose(out.numpy()[0], _CHILD_WORLD_T, atol=1e-5) - From a3176fade02b81040aad31995fd940d8d12518f2 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Sat, 23 May 2026 18:12:50 +0000 Subject: [PATCH 12/54] docs: remove dead ThreeDWorld link (domain squatted) --- docs/source/setup/ecosystem.rst | 2 -- 1 file changed, 2 deletions(-) diff --git a/docs/source/setup/ecosystem.rst b/docs/source/setup/ecosystem.rst index 4eac6d13cb53..718cc18b5b49 100644 --- a/docs/source/setup/ecosystem.rst +++ b/docs/source/setup/ecosystem.rst @@ -1,7 +1,6 @@ .. _isaac-lab-ecosystem: Isaac Lab Ecosystem -=================== Isaac Lab is a modular, extensible framework for robot learning built on top of `Isaac Sim`_ and `Newton`_. It provides a unified interface for the most common workflows in robotics research — @@ -209,7 +208,6 @@ contributing, please reach out to us. .. _AirSim: https://microsoft.github.io/AirSim/ .. _DoorGym: https://github.com/PSVL/DoorGym/ .. _ManiSkill: https://github.com/haosulab/ManiSkill -.. _ThreeDWorld: https://github.com/threedworld-mit/tdw .. _RoboSuite: https://github.com/ARISE-Initiative/robosuite .. _MuJoCo: https://mujoco.org/ .. _MuJoCo Playground: https://playground.mujoco.org/ From 4dc3b5cefa26e95f22c1ecfa6d0735724a496c3f Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Mon, 25 May 2026 17:24:21 +0000 Subject: [PATCH 13/54] feat: Fabric-accelerated get/set_local_poses via indexedfabricarray --- .../benchmarks/benchmark_xform_prim_view.py | 10 +- .../changelog.d/fabric-local-poses.rst | 21 + .../sim/views/fabric_frame_view.py | 733 +++++++++++++----- .../test/sim/test_views_xform_prim_fabric.py | 550 +++++++++++-- 4 files changed, 1072 insertions(+), 242 deletions(-) create mode 100644 source/isaaclab_physx/changelog.d/fabric-local-poses.rst diff --git a/scripts/benchmarks/benchmark_xform_prim_view.py b/scripts/benchmarks/benchmark_xform_prim_view.py index fee3b9642c79..76b5fab35863 100644 --- a/scripts/benchmarks/benchmark_xform_prim_view.py +++ b/scripts/benchmarks/benchmark_xform_prim_view.py @@ -139,7 +139,11 @@ def benchmark_frame_view( # noqa: C901 is_newton = api == "isaaclab-newton-site" 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 -------------------------------------------------------- @@ -162,7 +166,7 @@ def to_torch(a): # -- set_world_poses ----------------------------------------------- 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() @@ -198,7 +202,7 @@ 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() 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..5ad67e31d4e5 --- /dev/null +++ b/source/isaaclab_physx/changelog.d/fabric-local-poses.rst @@ -0,0 +1,21 @@ +Added +^^^^^ + +* Added Fabric-accelerated ``get_local_poses`` / ``set_local_poses`` to + :class:`~isaaclab_physx.sim.views.FabricFrameView`. + + Local-pose operations now use ``wp.indexedfabricarray`` to read/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 lazy per-view dirty tracking: ``set_local_poses`` marks the world + matrix dirty and vice-versa, triggering automatic re-propagation only on + the next read (no eager kernel launches on the write path). + +* Added interleave detection: interleaving ``set_world_poses`` and + ``set_local_poses`` on the same view within a frame flushes the stale + direction automatically and emits a one-time performance warning. + +* Added topology-change recovery via automatic ``PrepareForReuse`` detection + and per-selection index rebuild. 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 3d3c0a3b0d9e..86be4a5bd3ca 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,14 +7,14 @@ from __future__ import annotations +import enum import logging 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 @@ -24,6 +24,16 @@ logger = logging.getLogger(__name__) +class _DirtyFlag(enum.Enum): + """Which matrix direction is stale and needs recomputation on the next read.""" + + NONE = 0 + #: World matrices are stale (a prior ``set_local_poses`` wrote new locals). + WORLD = 1 + #: Local matrices are stale (a prior ``set_world_poses``/``set_scales`` wrote new worlds). + LOCAL = 2 + + def _to_float32_2d(a: wp.array | torch.Tensor) -> wp.array | torch.Tensor: """Ensure array is compatible with Fabric kernels (2-D float32). @@ -44,25 +54,61 @@ 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``. + 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. + * **World ↔ local consistency (lazy).** Getters are lazy: after + ``set_world_poses`` (or ``set_scales``), local matrices are only + recomputed when ``get_local_poses`` is called; after ``set_local_poses``, + world matrices are only recomputed when ``get_world_poses`` is called. + Both directions stay in sync without round-tripping through USD. + * **Dirty-flag invariant.** The ``_dirty`` enum is one of ``NONE``, + ``WORLD``, or ``LOCAL`` -- mutually exclusive by construction. + ``set_world_poses`` / ``set_scales`` sets ``_dirty = LOCAL``; + ``set_local_poses`` sets ``_dirty = WORLD``. + If the user interleaves both setters on the same view within a single + frame, the second setter flushes the first's stale data before writing. + This is correct but incurs an extra kernel launch -- a one-time warning + is logged when this happens. + * **Topology-adaptive.** Fabric topology changes are detected on each + access; the view rebuilds its internal mapping automatically and no + manual refresh is required. Steady-state overhead is negligible. + + Performance note: + The fast path assumes the user calls **either** ``set_world_poses`` + **or** ``set_local_poses`` exclusively within a frame (not both). + In that case, setters are O(1) kernel launches with no synchronization + overhead beyond the single ``wp.synchronize()``; getters lazily flush + the opposite direction only when actually needed. + + Interleaving both setters on different index subsets within the same + frame is supported and correct, but triggers an extra flush kernel + per transition. A warning is emitted once per view instance. + + Pose getters return :class:`~isaaclab.utils.warp.ProxyArray`; setters + accept :class:`wp.array`. """ + _WORLD_MATRIX_NAME = "omni:fabric:worldMatrix" + _LOCAL_MATRIX_NAME = "omni:fabric:localMatrix" + def __init__( self, prim_path: str, @@ -90,18 +136,34 @@ 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))}" + # Tracks which matrix direction is stale. Mutually exclusive by construction. + # Per-view (not per-stage) so concurrent views on the same stage don't interfere. + self._dirty: _DirtyFlag = _DirtyFlag.NONE + self._warned_interleaved_set: bool = False + + # Selection (single RW covering both world + local matrix). + self._sel = None + + # Index arrays (view-side indices and view->fabric mapping). + self._view_indices: wp.array | None = None + self._fabric_indices: wp.array | None = None + + # Indexed fabric arrays. + self._world_ifa = None + self._local_ifa = None + self._parent_world_ifa = None + + # Sentinel passed to compose/decompose kernels for unused slots. + # Kernels gate per-row access on ``shape[0] > 0``, so (0, 0) suffices. + self._fabric_empty_2d_array_sentinel: wp.array | None = None # ------------------------------------------------------------------ # Delegated properties @@ -146,48 +208,59 @@ def set_world_poses(self, positions=None, orientations=None, indices=None): if not self._fabric_initialized: self._initialize_fabric() - self._prepare_for_reuse() + # If a prior set_local_poses left worlds stale, flush them now. + if self._dirty == _DirtyFlag.WORLD and not self._warned_interleaved_set: + self._warned_interleaved_set = True + logger.warning( + "FabricFrameView: set_world_poses called while world matrices are stale from a " + "prior set_local_poses. Flushing stale worlds first. " + "For best performance, avoid interleaving set_world_poses and set_local_poses " + "on the same view within a single frame -- use one or the other exclusively." + ) - indices_wp = self._resolve_indices_wp(indices) - count = indices_wp.shape[0] + self._sync_world_from_local_if_dirty() - 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) - ) + indices_wp = self._resolve_indices_wp(indices) + positions_wp = self._to_float32_2d_or_empty(positions) + orientations_wp = self._to_float32_2d_or_empty(orientations) wp.launch( - kernel=fabric_utils.compose_fabric_transformation_matrix_from_warp_arrays, - dim=count, + kernel=fabric_utils.compose_indexed_fabric_transforms, + dim=indices_wp.shape[0], inputs=[ - self._fabric_world_matrices, + self._get_world_array(), positions_wp, orientations_wp, - dummy, + self._fabric_empty_2d_array_sentinel, False, False, False, indices_wp, - self._view_to_fabric, ], - device=self._fabric_device, + device=self._device, ) wp.synchronize() - self._fabric_hierarchy.update_world_xforms() - self._fabric_usd_sync_done = True + # World was just written -- mark local poses as stale so the next + # get_local_poses recomputes them lazily. + self._dirty = _DirtyFlag.LOCAL def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: + """Return (positions, orientations) in world frame. + + .. warning:: + When *indices* is None (all prims), the returned arrays are **shared + pre-allocated buffers** that are overwritten on the next call. Do not + hold references across calls -- copy if persistence is needed. + """ if not self._use_fabric: return self._usd_view.get_world_poses(indices) if not self._fabric_initialized: self._initialize_fabric() - if not self._fabric_usd_sync_done: - self._sync_fabric_from_usd_once() + + # If a prior set_local_poses left worldMatrix stale, propagate local -> world first. + self._sync_world_from_local_if_dirty() indices_wp = self._resolve_indices_wp(indices) count = indices_wp.shape[0] @@ -201,17 +274,16 @@ 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_array(), 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, ) if use_cached: @@ -220,20 +292,109 @@ def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, return ProxyArray(positions_wp), ProxyArray(orientations_wp) # ------------------------------------------------------------------ - # Local poses — USD fallback (Fabric only accelerates world poses) + # Local poses # ------------------------------------------------------------------ def set_local_poses(self, translations=None, orientations=None, indices=None): - self._usd_view.set_local_poses(translations, orientations, indices) + if not self._use_fabric: + self._usd_view.set_local_poses(translations, orientations, indices) + return + + if not self._fabric_initialized: + self._initialize_fabric() + + # If a prior set_world_poses left locals stale, flush them now before we + # overwrite a (possibly different) subset of local matrices. + if self._dirty == _DirtyFlag.LOCAL and not self._warned_interleaved_set: + self._warned_interleaved_set = True + logger.warning( + "FabricFrameView: set_local_poses called while local matrices are stale from a " + "prior set_world_poses/set_scales. Flushing stale locals first. " + "For best performance, avoid interleaving set_world_poses and set_local_poses " + "on the same view within a single frame -- use one or the other exclusively." + ) + + self._sync_local_from_world_if_dirty() + + indices_wp = self._resolve_indices_wp(indices) + translations_wp = self._to_float32_2d_or_empty(translations) + orientations_wp = self._to_float32_2d_or_empty(orientations) + + wp.launch( + kernel=fabric_utils.compose_indexed_fabric_transforms, + dim=indices_wp.shape[0], + inputs=[ + self._get_local_array(), + translations_wp, + orientations_wp, + self._fabric_empty_2d_array_sentinel, + False, + False, + False, + indices_wp, + ], + device=self._device, + ) + wp.synchronize() + + # Mark this view's worlds stale so the next world read recomputes them. + self._dirty = _DirtyFlag.WORLD def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: - return self._usd_view.get_local_poses(indices) + """Return (translations, orientations) in parent-local frame. + + .. warning:: + When *indices* is None (all prims), the returned arrays are **shared + pre-allocated buffers** that are overwritten on the next call. Do not + hold references across calls -- copy if persistence is needed. + """ + if not self._use_fabric: + return self._usd_view.get_local_poses(indices) + + if not self._fabric_initialized: + self._initialize_fabric() + + # If a prior set_world_poses/set_scales left localMatrix stale, recompute. + self._sync_local_from_world_if_dirty() + + indices_wp = self._resolve_indices_wp(indices) + count = indices_wp.shape[0] + + 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.decompose_indexed_fabric_transforms, + dim=count, + inputs=[ + self._get_local_array(), + translations_wp, + orientations_wp, + self._fabric_empty_2d_array_sentinel, + indices_wp, + ], + device=self._device, + ) + + if use_cached: + wp.synchronize() + return self._fabric_local_translations_ta, self._fabric_local_orientations_ta + return ProxyArray(translations_wp), ProxyArray(orientations_wp) # ------------------------------------------------------------------ - # Scales — Fabric-accelerated or USD fallback + # Scales # ------------------------------------------------------------------ def set_scales(self, scales, indices=None): + # TODO(pv): This decomposes/recomposes the *world* matrix, so the scale + # value is the composed (parent × local) scale. UsdFrameView.set_scales + # writes xformOp:scale which is purely local. The two diverge when the + # parent has non-identity scale. Fix: operate on localMatrix instead. if not self._use_fabric: self._usd_view.set_scales(scales, indices) return @@ -241,44 +402,49 @@ def set_scales(self, scales, indices=None): if not self._fabric_initialized: self._initialize_fabric() - self._prepare_for_reuse() + # Sync world matrices first if local writes are pending. + self._sync_world_from_local_if_dirty() 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) + scales_wp = self._to_float32_2d_or_empty(scales) wp.launch( - kernel=fabric_utils.compose_fabric_transformation_matrix_from_warp_arrays, - dim=count, + kernel=fabric_utils.compose_indexed_fabric_transforms, + dim=indices_wp.shape[0], inputs=[ - self._fabric_world_matrices, - dummy3, - dummy4, + self._get_world_array(), + self._fabric_empty_2d_array_sentinel, + self._fabric_empty_2d_array_sentinel, scales_wp, False, False, False, indices_wp, - self._view_to_fabric, ], - device=self._fabric_device, + device=self._device, ) wp.synchronize() - self._fabric_hierarchy.update_world_xforms() - self._fabric_usd_sync_done = True + # World was just written -- mark local poses as stale. + self._dirty = _DirtyFlag.LOCAL def get_scales(self, indices: wp.array | None = None) -> ProxyArray: + """Return per-prim (sx, sy, sz) scales extracted from world matrix. + + .. warning:: + When *indices* is None (all prims), the returned array is a **shared + pre-allocated buffer** that is overwritten on the next call. Do not + hold references across calls -- copy if persistence is needed. + """ + # TODO(pv): Same world-vs-local divergence as set_scales -- see note above. if not self._use_fabric: return self._usd_view.get_scales(indices) if not self._fabric_initialized: self._initialize_fabric() - if not self._fabric_usd_sync_done: - self._sync_fabric_from_usd_once() + + # Sync world matrices first if local writes are pending. + self._sync_world_from_local_if_dirty() indices_wp = self._resolve_indices_wp(indices) count = indices_wp.shape[0] @@ -290,17 +456,16 @@ def get_scales(self, indices: wp.array | None = None) -> ProxyArray: 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, + self._get_world_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, ) if use_cached: @@ -308,150 +473,352 @@ def get_scales(self, indices: wp.array | None = None) -> ProxyArray: return ProxyArray(scales_wp) # ------------------------------------------------------------------ - # Internal — PrepareForReuse (renderer notification + topology tracking) + # Internal -- sync helpers # ------------------------------------------------------------------ - def _prepare_for_reuse(self) -> None: - """Call PrepareForReuse on the PrimSelection to notify the renderer. + def _to_float32_2d_or_empty(self, data): + return self._fabric_empty_2d_array_sentinel if data is None else _to_float32_2d(data) - PrepareForReuse serves two purposes: - - 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: + def _sync_world_from_local_if_dirty(self) -> None: + """If a prior local write left world matrices stale, recompute them.""" + if self._dirty != _DirtyFlag.WORLD: return + self._recompute_world_from_local() + self._dirty = _DirtyFlag.NONE - 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 _rebuild_fabric_arrays(self) -> None: - """Rebuild fabricarray and view↔fabric mappings after a topology change. + def _recompute_world_from_local(self) -> None: + """Recompute world matrices: child_world = parent_world * child_local. - 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. + We deliberately do NOT call ``IFabricHierarchy.update_world_xforms()`` -- + in practice that re-reads USD's authored xformOps and overwrites the Fabric + local+world matrices we just authored. Instead we fire a Warp kernel that + does the multiply per child, leaving the Fabric-side localMatrix untouched. """ - 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." + self._refresh_if_needed() + wp.launch( + kernel=fabric_utils.update_indexed_world_matrix_from_local, + dim=self.count, + inputs=[ + self._local_ifa, + self._parent_world_ifa, + self._world_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) + wp.synchronize() + def _sync_local_from_world(self, indices_wp: wp.array) -> None: + """Recompute child localMatrix from (parent worldMatrix, child worldMatrix). + + Called after ``set_world_poses`` so that subsequent ``get_local_poses`` returns + values consistent with the just-written world poses. + """ + self._refresh_if_needed() 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_local_matrix_from_world, + dim=indices_wp.shape[0], + inputs=[ + self._world_ifa, + self._parent_world_ifa, + self._local_ifa, + indices_wp, + ], + device=self._device, ) wp.synchronize() - self._fabric_world_matrices = wp.fabricarray(self._fabric_selection, "omni:fabric:worldMatrix") + def _sync_local_from_world_if_dirty(self) -> None: + """If a prior world write left local matrices stale, recompute them lazily.""" + if self._dirty != _DirtyFlag.LOCAL: + return + self._sync_local_from_world(self._view_indices) + self._dirty = _DirtyFlag.NONE # ------------------------------------------------------------------ - # Internal — Fabric initialization + # Internal -- selection accessors with on-demand index rebuild # ------------------------------------------------------------------ - def _initialize_fabric(self) -> None: - """Initialize Fabric batch infrastructure for GPU-accelerated pose queries.""" - import usdrt # noqa: PLC0415 - from usdrt import Rt # noqa: PLC0415 + def _refresh_if_needed(self): + """Rebuild indexed arrays if the selection's prim set changed.""" + if self._sel.PrepareForReuse() or self._world_ifa is None: + self._fabric_indices = self._compute_fabric_indices(self._sel) + self._world_ifa = self._build_indexed_array(self._sel, self._WORLD_MATRIX_NAME, self._fabric_indices) + self._local_ifa = self._build_indexed_array(self._sel, self._LOCAL_MATRIX_NAME, self._fabric_indices) + self._parent_world_ifa = self._build_parent_indexed_array(self._sel) - stage_id = sim_utils.get_current_stage_id() - fabric_stage = usdrt.Usd.Stage.Attach(stage_id) + def _get_world_array(self): + self._refresh_if_needed() + return self._world_ifa - for i in range(self.count): - rt_prim = fabric_stage.GetPrimAtPath(self.prim_paths[i]) - rt_xformable = Rt.Xformable(rt_prim) + def _get_local_array(self): + self._refresh_if_needed() + return self._local_ifa - has_attr = ( - rt_xformable.HasFabricHierarchyWorldMatrixAttr() - if hasattr(rt_xformable, "HasFabricHierarchyWorldMatrixAttr") - else False - ) - if not has_attr: - rt_xformable.CreateFabricHierarchyWorldMatrixAttr() + def _get_parent_world_array(self): + self._refresh_if_needed() + return self._parent_world_ifa - rt_xformable.SetWorldXformFromUsd() + # ------------------------------------------------------------------ + # Internal -- index computation + # ------------------------------------------------------------------ - rt_prim.CreateAttribute(self._view_index_attr, usdrt.Sdf.ValueTypeNames.UInt, custom=True) - rt_prim.GetAttribute(self._view_index_attr).Set(i) + def _compute_fabric_indices(self, selection) -> wp.array: + fabric_paths = selection.GetPaths() + path_to_fabric_idx: dict[str, int] = {str(p): i for i, p in enumerate(fabric_paths)} + indices: list[int] = [] + for prim_path in self.prim_paths: + fabric_idx = path_to_fabric_idx.get(prim_path) + if fabric_idx is None: + raise RuntimeError( + f"Prim '{prim_path}' not found in Fabric selection. Ensure the hierarchy has been populated." + ) + indices.append(fabric_idx) + return wp.array(indices, dtype=wp.int32, device=self._device) + + def _compute_parent_fabric_indices(self, selection) -> wp.array: + """For each child in this view, look up the parent prim's fabric index.""" + fabric_paths = selection.GetPaths() + path_to_fabric_idx: dict[str, int] = {str(p): i for i, p in enumerate(fabric_paths)} + indices: list[int] = [] + for prim_path in self.prim_paths: + parent_path = prim_path.rsplit("/", 1)[0] + if parent_path == "": + 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." + ) + fabric_idx = path_to_fabric_idx.get(parent_path) + if fabric_idx is None: + raise RuntimeError( + f"Parent prim '{parent_path}' (for child '{prim_path}') not found in Fabric selection. " + "Ensure parents have Fabric world+local matrices populated." + ) + indices.append(fabric_idx) + return wp.array(indices, dtype=wp.int32, 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 _build_parent_indexed_array(self, selection) -> wp.indexedfabricarray: + self._parent_fabric_indices = self._compute_parent_fabric_indices(selection) + fa = wp.fabricarray(selection, self._WORLD_MATRIX_NAME) + return wp.indexedfabricarray(fa=fa, indices=self._parent_fabric_indices) - self._fabric_hierarchy = usdrt.hierarchy.IFabricHierarchy().get_fabric_hierarchy( - fabric_stage.GetFabricId(), fabric_stage.GetStageIdAsStageId() - ) - self._fabric_hierarchy.update_world_xforms() + 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 - self._default_view_indices = wp.zeros((self.count,), dtype=wp.uint32, device=self._device) - wp.launch( - kernel=fabric_utils.arange_k, dim=self.count, inputs=[self._default_view_indices], device=self._device - ) - wp.synchronize() + # ------------------------------------------------------------------ + # Internal -- Fabric initialization + # ------------------------------------------------------------------ - 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), - ], - device=self._device, - ) + def _initialize_fabric(self) -> None: + """One-time Fabric setup: hierarchy handle, attribute population, selections, indexed arrays.""" + import usdrt # noqa: PLC0415 + from usdrt import Rt # noqa: PLC0415 - 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) + from isaaclab.sim.utils import get_current_stage_id # noqa: PLC0415 - 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, + # 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() ) - wp.synchronize() + # Ensure each child prim AND its parent have BOTH Fabric world and local matrix + # attributes. Our ``trans_ro`` selection requires both, so prims missing either + # would silently be excluded. ``Create*Attr`` calls are idempotent. + # + # ``SetWorldXformFromUsd`` writes Fabric's worldMatrix from USD's accumulated + # local-to-world transform (so it picks up the parent chain). + # ``SetLocalXformFromUsd`` writes Fabric's localMatrix from USD's authored + # xformOps on this prim only. Calling both gives Fabric a consistent + # (worldMatrix, localMatrix) pair for each prim before we touch the hierarchy. + 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() + + # Single RW selection covering both matrices. + # TODO: Benchmark RO vs RW selection split -- separate RO selections could reduce + # lock contention under concurrent Fabric access, but current usage is single-threaded. + matrix = usdrt.Sdf.ValueTypeNames.Matrix4d + rw = usdrt.Usd.Access.ReadWrite + wm_rw = (matrix, self._WORLD_MATRIX_NAME, rw) + lm_rw = (matrix, self._LOCAL_MATRIX_NAME, rw) + self._sel = self._stage.SelectPrims(require_attrs=[wm_rw, lm_rw], device=self._device, want_paths=True) + + # Build the view-side indices array (just [0..count-1]) and a + # view->fabric mapping (selections do not guarantee a shared path ordering). + self._view_indices = wp.array(list(range(self.count)), dtype=wp.uint32, device=self._device) + self._fabric_indices = self._compute_fabric_indices(self._sel) + + # Indexed fabric arrays per attribute. + self._world_ifa = self._build_indexed_array(self._sel, self._WORLD_MATRIX_NAME, self._fabric_indices) + self._local_ifa = self._build_indexed_array(self._sel, self._LOCAL_MATRIX_NAME, self._fabric_indices) + self._parent_world_ifa = self._build_parent_indexed_array(self._sel) + + # 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) + 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 + self._fabric_local_translations_ta = ProxyArray(self._fabric_local_translations_buf) + self._fabric_local_orientations_ta = ProxyArray(self._fabric_local_orientations_buf) self._fabric_initialized = True - self._fabric_usd_sync_done = False - def _sync_fabric_from_usd_once(self) -> None: - """Sync Fabric world matrices from USD once, on the first read. + # Seed Fabric matrices from USD authoritatively. ``SetWorldXformFromUsd`` / + # ``SetLocalXformFromUsd`` are no-ops on freshly authored stages that haven't + # been rendered yet; we instead read through the USD view (children) and + # ``UsdGeom.XformCache`` (parents) and write via the same compose kernel that + # ``set_world_poses`` uses. + self._sync_fabric_from_usd_initial() - ``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() + def _sync_fabric_from_usd_initial(self) -> None: + """Populate Fabric world+local matrices for children and parents from USD. - 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().warp + 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 --- + pos_ta, ori_ta = self._usd_view.get_world_poses() + scales_obj = self._usd_view.get_scales() + scales_wp = ( + scales_obj.warp + if hasattr(scales_obj, "warp") + else scales_obj + if isinstance(scales_obj, wp.array) + else self._fabric_empty_2d_array_sentinel + ) + local_pos_ta, local_ori_ta = self._usd_view.get_local_poses() + # Compose into child worldMatrix. + wp.launch( + kernel=fabric_utils.compose_indexed_fabric_transforms, + dim=self.count, + inputs=[ + self._world_ifa, + _to_float32_2d(pos_ta.warp), + _to_float32_2d(ori_ta.warp), + _to_float32_2d(scales_wp), + False, + False, + False, + self._view_indices, + ], + device=self._device, + ) + # Compose into child localMatrix. Pass the locally-authored scale so + # that a subsequent ``_sync_world_from_local_if_dirty`` produces the + # right world-space scale (``world = parent_world * local`` carries + # ``local``'s scale through the multiply). + wp.launch( + kernel=fabric_utils.compose_indexed_fabric_transforms, + dim=self.count, + inputs=[ + self._local_ifa, + _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.set_world_poses(positions_usd, orientations_usd) - self.set_scales(scales_usd) + # --- 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() + for path in unique_parent_paths: + prim = usd_stage.GetPrimAtPath(path) + tf = xform_cache.GetLocalToWorldTransform(prim) + # Extract scale before ``Orthonormalize`` strips it from the rows. + decomposer.SetMatrix(tf) + s = decomposer.GetScale() + 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) + # Compose worldMatrix for parents (use a one-shot indexed array against + # ``world_sel_rw`` keyed on the unique parent paths). + parent_world_rw = wp.indexedfabricarray( + fa=wp.fabricarray(self._sel, self._WORLD_MATRIX_NAME), + indices=self._compute_fabric_indices_for(self._sel, 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() - 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 + # After seeding local matrices from USD, recompute world matrices so + # the view starts with consistent state (child_world = parent_world * child_local). + self._recompute_world_from_local() + + def _compute_fabric_indices_for(self, selection, paths: list[str]) -> wp.array: + """Path-dict lookup helper used to build one-shot indexed arrays for a custom path set.""" + fabric_paths = selection.GetPaths() + path_to_idx = {str(p): i for i, p in enumerate(fabric_paths)} + indices: list[int] = [] + for path in paths: + idx = path_to_idx.get(path) + if idx is None: + raise RuntimeError(f"Path '{path}' not found in Fabric selection.") + indices.append(idx) + return wp.array(indices, dtype=wp.int32, device=self._device) 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..441a8bd0ac84 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 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,7 +157,7 @@ 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) @@ -187,7 +170,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 +183,328 @@ 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() - - def force_topology_changed(): - if view._fabric_selection is not None: - view._fabric_selection.PrepareForReuse() - spy_rebuild() - - monkeypatch.setattr(view, "_prepare_for_reuse", force_topology_changed) + # Simulate topology change: force rebuild of the selection's indexed arrays. + view._refresh_if_needed() - # 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" - - # 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", ["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 is not None, "selection not initialized" + result = view._sel.PrepareForReuse() + assert isinstance(result, bool), f"PrepareForReuse should return bool, got {type(result)}" + assert not result, "PrepareForReuse should return False when no topology change" + + +@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 set_local_poses should take the 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) + + view.set_local_poses(translations=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_scales path.""" + bundle = view_factory(num_envs=1, device=device) + view = bundle.view + + # Trigger lazy `_initialize_fabric()` so the get_scales call below uses Fabric. + view.get_world_poses() + + scales = view.get_scales() + scales_t = torch.as_tensor(scales, device=device) + # 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) # ------------------------------------------------------------------ -# Multi-GPU tests (cuda:1) — skipped automatically on single-GPU workstations +# 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)) + view.set_local_poses(translations=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) + view.set_world_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, after ``_sync_world_from_local_if_dirty`` the 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 = torch.as_tensor(view.get_scales(), device=device) + torch.testing.assert_close( + scales, + torch.tensor([[6.0, 1.0, 1.0]], dtype=torch.float32, device=device), + atol=1e-5, + rtol=0, + ) + + +# ------------------------------------------------------------------ +# Multi-view per stage: per-view dirty-flag isolation +# ------------------------------------------------------------------ + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_multi_view_per_view_dirty_isolation(device): + """Two ``FabricFrameView`` instances on the same stage must not clear each other's + pending local→world sync. + + Background: an earlier implementation stored the world-dirty flag at the class + level keyed by ``stage_id``. With two views on the same stage, view B reading + worlds would clear the flag set by view A's ``set_local_poses``, leaving A's + world matrices silently stale because A's per-view sync kernel never fired. + + This test sets up two views over disjoint child prims (under different parent + sub-trees of the same stage), interleaves their writes and reads, and verifies: + + * view A's ``set_local_poses`` only dirties view A + * view B's ``get_world_poses`` does not clear view A's flag + * after both views' world reads, each one's worlds reflect its own latest local + * neither view's reads/writes corrupt the other view's poses + """ + _skip_if_unavailable(device) + stage = sim_utils.get_current_stage() + + # Two disjoint sub-trees under the same stage. Use different parent names so + # the regex patterns for the two views don't accidentally overlap. + 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) + + # Initial reads -- triggers Fabric init + the seed-time ``_dirty = WORLD`` + # path on both views, then clears it. + 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 + ) + assert view_a._dirty.name == "NONE" + assert view_b._dirty.name == "NONE" + + # Write a new local pose on view A only. + 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)) + view_a.set_local_poses(translations=new_local_a, orientations=identity_quat) + + # Only view A should be dirty. Critical: a per-stage flag would have dirtied + # both views (or neither) at this point. + assert view_a._dirty.name == "WORLD", "set_local_poses should mark its own view dirty" + assert view_b._dirty.name == "NONE", "set_local_poses on view A must not dirty view B" + + # Read worlds from view B FIRST. With a per-stage flag, B's + # ``_sync_world_from_local_if_dirty`` would fire and clear the flag, leaving A + # stale. With the per-view flag, B's read is a no-op sync-wise. + torch.testing.assert_close( + torch.as_tensor(view_b.get_world_poses()[0], device=device), expected_b0, atol=1e-5, rtol=0 + ) + assert view_b._dirty.name == "NONE" + assert view_a._dirty.name == "WORLD", "view B's world read must not clear view A's dirty flag" + + # Now read view A's worlds -- sync fires, 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 + ) + assert view_a._dirty.name == "NONE" + + # Symmetric pass: write on B, ensure A is undisturbed. + 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) + view_b.set_local_poses(translations=new_local_b, orientations=identity_quat) + assert view_a._dirty.name == "NONE" + assert view_b._dirty.name == "WORLD" + + # A's worlds must still read back the post-set-local value from above; no + # cross-view stomp on the world matrix. + 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 + ) + assert view_a._dirty.name == "NONE" + assert view_b._dirty.name == "NONE" + + +# ------------------------------------------------------------------ +# Multi-GPU tests (cuda:1) -- skipped automatically on single-GPU workstations # ------------------------------------------------------------------ @@ -300,7 +556,7 @@ def test_fabric_cuda1_no_usd_writeback(device, view_factory): 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) - # 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), ( @@ -331,3 +587,185 @@ def test_fabric_cuda1_scales_roundtrip(device, view_factory): scales_torch = torch.as_tensor(ret_scales, device=device) 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}" + + +# ------------------------------------------------------------------ +# Interleaved set_world_poses / set_local_poses tests +# ------------------------------------------------------------------ + + +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_interleaved_set_world_then_set_local_partial_indices(device): + """set_world_poses on index 0, then set_local_poses on index 1 -- both must be correct. + + This exercises the dirty-flag flush: after set_world_poses marks _dirty == LOCAL, + set_local_poses must flush stale locals before writing index 1, ensuring index 0's + local is correctly derived from its new world pose. + """ + view = _build_two_child_view(device) + + # Step 1: set world pose on index 0 only + 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)) + view.set_world_poses(positions=new_world_pos, indices=idx0) + + # Step 2: set local pose on index 1 only + 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)) + view.set_local_poses(translations=new_local_pos, orientations=identity_quat, indices=idx1) + + # Verify index 0's world pose is still (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, + ) + + # Verify index 0's local pose (derived from 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, + ) + + # Verify index 1's world pose (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_interleaved_set_local_then_set_world_partial_indices(device): + """set_local_poses on index 0, then set_world_poses on index 1 -- both must be correct. + + The reverse direction of the above: after set_local_poses marks _dirty = WORLD, + set_world_poses must flush stale worlds before writing index 1. + """ + view = _build_two_child_view(device) + + # Step 1: set local pose on index 0 only + 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)) + view.set_local_poses(translations=new_local_pos, orientations=identity_quat, indices=idx0) + + # Step 2: set world pose on index 1 only + 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)) + view.set_world_poses(positions=new_world_pos, indices=idx1) + + # Verify index 0's world pose (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, + ) + + # Verify index 1's world pose 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, + ) + + # Verify 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, + ) + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_interleaved_set_emits_warning(device, caplog): + """Interleaving set_world_poses and set_local_poses logs a one-time warning.""" + view = _build_two_child_view(device) + + # First set_world_poses -- no warning (first user setter) + new_world = wp.zeros((2, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=2, inputs=[new_world, 1.0, 2.0, 3.0], device=device) + view.set_world_poses(positions=new_world) + + # Now set_local_poses -- should trigger warning about interleaving + new_local = wp.zeros((2, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=2, inputs=[new_local, 0.0, 0.0, 0.0], device=device) + identity_quat = wp.from_torch(torch.tensor([[0.0, 0.0, 0.0, 1.0]] * 2, dtype=torch.float32, device=device)) + + with caplog.at_level("WARNING", logger="isaaclab_physx.sim.views.fabric_frame_view"): + caplog.clear() + view.set_local_poses(translations=new_local, orientations=identity_quat) + + assert any("interleaving" in r.message.lower() for r in caplog.records), ( + f"Expected interleave warning, got: {[r.message for r in caplog.records]}" + ) + + # Second interleave -- warning should NOT repeat (one-time only) + caplog.clear() + view.set_world_poses(positions=new_world) + assert not any("interleaving" in r.message.lower() for r in caplog.records), ( + "Warning should only fire once per view instance" + ) + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_no_warning_when_using_single_setter(device, caplog): + """Calling only set_world_poses (or only set_local_poses) should never warn.""" + view = _build_two_child_view(device) + + new_world = wp.zeros((2, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=2, inputs=[new_world, 1.0, 2.0, 3.0], device=device) + + with caplog.at_level("WARNING", logger="isaaclab_physx.sim.views.fabric_frame_view"): + caplog.clear() + view.set_world_poses(positions=new_world) + view.set_world_poses(positions=new_world) + view.set_world_poses(positions=new_world) + + assert not any("interleaving" in r.message.lower() for r in caplog.records), ( + f"Unexpected interleave warning with single setter: {[r.message for r in caplog.records]}" + ) From 77fbb0ed11da0be3d76f9bea82438d3ec563c780 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Mon, 25 May 2026 21:38:00 +0000 Subject: [PATCH 14/54] refactor: remove kernel unit tests, consolidate changelogs for combined MR - Remove test_fabric_kernels.py: @wp.func math is fully exercised by the 57 integration tests in test_views_xform_prim_fabric.py - Merge indexed-fabric-kernels.rst into fabric-local-poses.rst (single changelog) - @wp.func helpers remain module-private (used by production kernels) --- .../changelog.d/indexed-fabric-kernels.rst | 18 -- .../test/utils/warp/test_fabric_kernels.py | 164 ------------------ .../changelog.d/fabric-local-poses.rst | 14 ++ 3 files changed, 14 insertions(+), 182 deletions(-) delete mode 100644 source/isaaclab/changelog.d/indexed-fabric-kernels.rst delete mode 100644 source/isaaclab/test/utils/warp/test_fabric_kernels.py diff --git a/source/isaaclab/changelog.d/indexed-fabric-kernels.rst b/source/isaaclab/changelog.d/indexed-fabric-kernels.rst deleted file mode 100644 index baaf33ceb8a3..000000000000 --- a/source/isaaclab/changelog.d/indexed-fabric-kernels.rst +++ /dev/null @@ -1,18 +0,0 @@ -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 (no - explicit transposes). Will be used by - :class:`~isaaclab_physx.sim.views.FabricFrameView` to keep child world and - local matrices consistent across writes without round-tripping through USD. diff --git a/source/isaaclab/test/utils/warp/test_fabric_kernels.py b/source/isaaclab/test/utils/warp/test_fabric_kernels.py deleted file mode 100644 index de48afadab24..000000000000 --- a/source/isaaclab/test/utils/warp/test_fabric_kernels.py +++ /dev/null @@ -1,164 +0,0 @@ -# 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 - -"""Unit tests for the Warp fabric transform kernels. - -Tests the shared @wp.func math (decompose/compose and matrix inverse/multiply) -through plain wp.array kernels — no Fabric/USDRT runtime required. - -The production fabric kernels are thin adapters over the same math; testing the -math in isolation avoids coupling tests to Fabric container internals. -""" - -import numpy as np -import warp as wp - -wp.init() - -from isaaclab.utils.warp.fabric import ( # noqa: E402 - _decompose_transformation_matrix, - _local_from_world_transposed, - _world_from_local_transposed, -) - -# ------------------------------------------------------------------ -# Test kernels — thin wp.array wrappers that delegate to production @wp.func -# ------------------------------------------------------------------ - - -@wp.kernel(enable_backward=False) -def _test_decompose_kernel( - matrices: wp.array(dtype=wp.mat44f), - out_positions: wp.array(dtype=wp.vec3f), - out_rotations: wp.array(dtype=wp.quatf), - out_scales: wp.array(dtype=wp.vec3f), -): - """Decompose a batch of 4x4 matrices into pos/quat/scale.""" - i = wp.tid() - pos, rot, scale = _decompose_transformation_matrix(matrices[i]) - out_positions[i] = pos - out_rotations[i] = rot - out_scales[i] = scale - - -@wp.kernel(enable_backward=False) -def _test_local_from_world_kernel( - child_world: wp.array(dtype=wp.mat44d), - parent_world: wp.array(dtype=wp.mat44d), - out_local: wp.array(dtype=wp.mat44d), -): - """wp.array adapter for _local_from_world_transposed — same func as production fabric kernel.""" - i = wp.tid() - out_local[i] = wp.mat44d(_local_from_world_transposed(wp.mat44f(child_world[i]), wp.mat44f(parent_world[i]))) - - -@wp.kernel(enable_backward=False) -def _test_world_from_local_kernel( - child_local: wp.array(dtype=wp.mat44d), - parent_world: wp.array(dtype=wp.mat44d), - out_world: wp.array(dtype=wp.mat44d), -): - """wp.array adapter for _world_from_local_transposed — same func as production fabric kernel.""" - i = wp.tid() - out_world[i] = wp.mat44d(_world_from_local_transposed(wp.mat44f(child_local[i]), wp.mat44f(parent_world[i]))) - - -# ------------------------------------------------------------------ -# Helpers -# ------------------------------------------------------------------ - - -def _make_transform_matrix(pos, rot_quat_xyzw, scale): - """Build a 4x4 Fabric-transposed transform from pos/quat/scale. - - Returns numpy (4,4) float64 in the transposed storage convention (row-major with - translation in the last row) that Fabric uses. - - Raises: - AssertionError: If the resulting matrix is singular (e.g. zero scale component). - """ - p = wp.vec3f(*pos) - q = wp.quatf(*rot_quat_xyzw) - s = wp.vec3f(*scale) - m = wp.transpose(wp.transform_compose(p, q, s)) - result = np.array(m).reshape(4, 4).astype(np.float64) - det = np.linalg.det(result) - assert abs(det) > 1e-6, f"Singular matrix: det={det:.2e}, scale={scale}" - return result - - -# ------------------------------------------------------------------ -# Decompose / Compose round-trip tests -# ------------------------------------------------------------------ - - -def test_decompose_round_trip(): - """Decompose a matrix with translation, rotation, and non-uniform scale; verify round-trip.""" - pos = np.array([5.0, -3.0, 7.0]) - s45 = np.sin(np.pi / 4) - c45 = np.cos(np.pi / 4) - quat_xyzw = np.array([0.0, 0.0, s45, c45]) # 45° Z rotation - scale = np.array([1.5, 0.8, 3.0]) - - mat = _make_transform_matrix(pos, quat_xyzw, scale).astype(np.float32) - matrices = wp.array(mat.reshape(1, 4, 4), dtype=wp.mat44f, device="cpu") - - out_pos = wp.zeros(1, dtype=wp.vec3f, device="cpu") - out_rot = wp.zeros(1, dtype=wp.quatf, device="cpu") - out_scale = wp.zeros(1, dtype=wp.vec3f, device="cpu") - - wp.launch(_test_decompose_kernel, dim=1, inputs=[matrices, out_pos, out_rot, out_scale], device="cpu") - wp.synchronize() - - np.testing.assert_allclose(out_pos.numpy()[0], pos, atol=1e-5) - np.testing.assert_allclose(out_scale.numpy()[0], scale, atol=1e-5) - dot = np.dot(out_rot.numpy()[0], quat_xyzw) - np.testing.assert_allclose(abs(dot), 1.0, atol=1e-5) - - -# ------------------------------------------------------------------ -# World ↔ Local matrix tests -# -# These test the same math the production fabric kernels use: -# local^T = world^T * inv(parent^T) -# world^T = local^T * parent^T -# -# Both parent and child have rotation, translation, and non-uniform scale -# (producing sheared/non-orthogonal upper-3x3 blocks). -# ------------------------------------------------------------------ - -# Shared test data: parent with 10:1 non-uniform scale + 45° Z rotation + translation -_PARENT_WORLD_T = _make_transform_matrix([10, -5, 2], [0, 0, 0.3826834, 0.9238795], [4.0, 0.5, 2.0]) -_CHILD_WORLD_T = _make_transform_matrix([1, 2, 3], [0.2588190, 0, 0, 0.9659258], [1.5, 0.8, 3.0]) - - -def test_local_from_world_transposed(): - """local^T = world^T * inv(parent^T) — verified by reconstruction.""" - cw = wp.array(_CHILD_WORLD_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") - pw = wp.array(_PARENT_WORLD_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") - out = wp.zeros(1, dtype=wp.mat44d, device="cpu") - - wp.launch(_test_local_from_world_kernel, dim=1, inputs=[cw, pw, out], device="cpu") - wp.synchronize() - - # Reconstruction: local^T @ parent^T must equal child_world^T - local_T = out.numpy()[0] - reconstructed = local_T @ _PARENT_WORLD_T - np.testing.assert_allclose(reconstructed, _CHILD_WORLD_T, atol=1e-5) - - -def test_world_from_local_transposed(): - """world^T = local^T * parent^T — verified against known child world.""" - # Ground-truth local computed via numpy - child_local_T = _CHILD_WORLD_T @ np.linalg.inv(_PARENT_WORLD_T) - - cl = wp.array(child_local_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") - pw = wp.array(_PARENT_WORLD_T.reshape(1, 4, 4), dtype=wp.mat44d, device="cpu") - out = wp.zeros(1, dtype=wp.mat44d, device="cpu") - - wp.launch(_test_world_from_local_kernel, dim=1, inputs=[cl, pw, out], device="cpu") - wp.synchronize() - - np.testing.assert_allclose(out.numpy()[0], _CHILD_WORLD_T, atol=1e-5) diff --git a/source/isaaclab_physx/changelog.d/fabric-local-poses.rst b/source/isaaclab_physx/changelog.d/fabric-local-poses.rst index 5ad67e31d4e5..61b59c22398b 100644 --- a/source/isaaclab_physx/changelog.d/fabric-local-poses.rst +++ b/source/isaaclab_physx/changelog.d/fabric-local-poses.rst @@ -1,6 +1,20 @@ 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 ``get_local_poses`` / ``set_local_poses`` to :class:`~isaaclab_physx.sim.views.FabricFrameView`. From fee3e75222eb4a061694c9c62a1c8673a5abb791 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Tue, 26 May 2026 08:58:38 +0000 Subject: [PATCH 15/54] refactor: inline @wp.func helpers into production kernels _local_from_world_transposed and _world_from_local_transposed were one-line wrappers. Inline the math directly into the kernel bodies (child_world * wp.inverse(parent_world) and child_local * parent_world) to reduce indirection. --- source/isaaclab/isaaclab/utils/warp/fabric.py | 16 ++-------------- 1 file changed, 2 insertions(+), 14 deletions(-) diff --git a/source/isaaclab/isaaclab/utils/warp/fabric.py b/source/isaaclab/isaaclab/utils/warp/fabric.py index f665323cbe34..e0519d98c338 100644 --- a/source/isaaclab/isaaclab/utils/warp/fabric.py +++ b/source/isaaclab/isaaclab/utils/warp/fabric.py @@ -268,18 +268,6 @@ def compose_indexed_fabric_transforms( ) -@wp.func -def _local_from_world_transposed(child_world_T: wp.mat44f, parent_world_T: wp.mat44f) -> wp.mat44f: - """Compute local^T = world^T * inv(parent^T) on transposed storage matrices.""" - return child_world_T * wp.inverse(parent_world_T) - - -@wp.func -def _world_from_local_transposed(child_local_T: wp.mat44f, parent_world_T: wp.mat44f) -> wp.mat44f: - """Compute world^T = local^T * parent^T on transposed storage matrices.""" - return child_local_T * parent_world_T - - @wp.kernel(enable_backward=False) def update_indexed_local_matrix_from_world( child_world_matrices: IndexedFabricArrayMat44d, @@ -314,7 +302,7 @@ def update_indexed_local_matrix_from_world( 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] - _local_from_world_transposed(child_world, parent_world) + child_world * wp.inverse(parent_world) ) @@ -347,7 +335,7 @@ def update_indexed_world_matrix_from_local( 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] - _world_from_local_transposed(child_local, parent_world) + child_local * parent_world ) From 3765299116543eddb275d657893f330043ab6be6 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Tue, 26 May 2026 10:06:40 +0000 Subject: [PATCH 16/54] feat: deprecate set/get_scales, add explicit local/world scale API Introduce set_local_scales/get_local_scales and set_world_scales/ get_world_scales on BaseFrameView and all implementations: - UsdFrameView: local scales read/write xformOp:scale; world scales decompose/compose from world transform matrix - FabricFrameView: local scales operate on localMatrix (marks world dirty); world scales operate on worldMatrix (marks local dirty) - NewtonSiteFrameView: both map to shape_scale (Newton treats scale as composed) - OvPhysxFrameView: delegates to UsdFrameView The deprecated set_scales/get_scales still work but emit DeprecationWarning. USD/OvPhysX/Newton default to prior behavior; Fabric defaults to world scales (matching its pre-existing semantics). --- .../isaaclab/sim/views/base_frame_view.py | 92 +- .../isaaclab/sim/views/usd_frame_view.py | 76 +- .../sim/views/newton_site_frame_view.py | 1156 ++++++++++++----- .../sim/views/ovphysx_frame_view.py | 30 +- .../changelog.d/fabric-local-poses.rst | 11 + .../sim/views/fabric_frame_view.py | 100 +- .../test/sim/test_views_xform_prim_fabric.py | 60 +- 7 files changed, 1151 insertions(+), 374 deletions(-) diff --git a/source/isaaclab/isaaclab/sim/views/base_frame_view.py b/source/isaaclab/isaaclab/sim/views/base_frame_view.py index 656108f24d2c..8095b3963692 100644 --- a/source/isaaclab/isaaclab/sim/views/base_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/base_frame_view.py @@ -8,6 +8,7 @@ from __future__ import annotations import abc +import warnings import warp as wp @@ -98,8 +99,11 @@ def set_local_poses( ... @abc.abstractmethod - def get_scales(self, indices: wp.array | None = None) -> ProxyArray: - """Get scales for prims in the view. + def get_local_scales(self, indices: wp.array | None = None) -> wp.array: + """Get local-space scales for prims in the view. + + Returns the per-prim scale as stored in local space (``xformOp:scale`` + for USD, decomposition of ``localMatrix`` for Fabric). Args: indices: Subset of prims to query. ``None`` means all prims. @@ -110,11 +114,95 @@ def get_scales(self, indices: wp.array | None = None) -> ProxyArray: ... @abc.abstractmethod + def set_local_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Set local-space scales for prims in the view. + + Writes the per-prim scale in local space. For the Fabric backend this + marks the world matrix as dirty (will be re-propagated on next read). + + Args: + scales: Scales ``(M, 3)`` as ``wp.array``. + indices: Subset of prims to update. ``None`` means all prims. + """ + ... + + @abc.abstractmethod + def get_world_scales(self, indices: wp.array | None = None) -> wp.array: + """Get world-space (composed) scales for prims in the view. + + Returns the effective scale in world space (``parent_scale * local_scale``). + + Args: + indices: Subset of prims to query. ``None`` means all prims. + + Returns: + A ``wp.array`` of shape ``(M, 3)``. + """ + ... + + @abc.abstractmethod + def set_world_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Set world-space (composed) scales for prims in the view. + + Writes the effective scale in world space. For the Fabric backend this + marks the local matrix as dirty (will be re-propagated on next read). + + Args: + scales: Scales ``(M, 3)`` as ``wp.array``. + indices: Subset of prims to update. ``None`` means all prims. + """ + ... + + # ------------------------------------------------------------------ + # Deprecated -- use get/set_local_scales or get/set_world_scales + # ------------------------------------------------------------------ + + def get_scales(self, indices: wp.array | None = None) -> wp.array: + """Get scales for prims in the view. + + .. deprecated:: + Use :meth:`get_local_scales` or :meth:`get_world_scales` instead. + This method calls ``get_local_scales`` for USD backends and + ``get_world_scales`` for Fabric backends (preserving legacy behavior). + + Args: + indices: Subset of prims to query. ``None`` means all prims. + + Returns: + A ``wp.array`` of shape ``(M, 3)``. + """ + warnings.warn( + "get_scales() is deprecated. Use get_local_scales() or get_world_scales() instead.", + DeprecationWarning, + stacklevel=2, + ) + return self._get_scales_default(indices) + def set_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: """Set scales for prims in the view. + .. deprecated:: + Use :meth:`set_local_scales` or :meth:`set_world_scales` instead. + This method calls ``set_local_scales`` for USD backends and + ``set_world_scales`` for Fabric backends (preserving legacy behavior). + Args: scales: Scales ``(M, 3)`` as ``wp.array``. indices: Subset of prims to update. ``None`` means all prims. """ + warnings.warn( + "set_scales() is deprecated. Use set_local_scales() or set_world_scales() instead.", + DeprecationWarning, + stacklevel=2, + ) + self._set_scales_default(scales, indices) + + @abc.abstractmethod + def _get_scales_default(self, indices: wp.array | None = None) -> wp.array: + """Backend-specific default for deprecated get_scales().""" + ... + + @abc.abstractmethod + def _set_scales_default(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Backend-specific default for deprecated 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 31331221e46f..166f7c36288f 100644 --- a/source/isaaclab/isaaclab/sim/views/usd_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/usd_frame_view.py @@ -213,8 +213,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. + def set_local_scales(self, scales: wp.array, indices: wp.array | None = None): + """Set local-space scales (xformOp:scale) for prims in the view. Args: scales: Scales of shape ``(M, 3)``. @@ -228,6 +228,49 @@ 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_world_scales(self, scales: wp.array, indices: wp.array | None = None): + """Set world-space (composed) scales for prims in the view. + + Computes ``local_scale = world_scale / parent_world_scale`` and writes + to ``xformOp:scale``. + + Args: + scales: World-space scales of shape ``(M, 3)``. + indices: Indices of prims to set scales for. Defaults to None (all prims). + """ + indices_list = self._resolve_indices(indices) + scales_np = self._to_numpy(scales) + xf_cache = UsdGeom.XformCache(Usd.TimeCode.Default()) + + with Sdf.ChangeBlock(): + for idx, prim_idx in enumerate(indices_list): + prim = self._prims[prim_idx] + parent = prim.GetParent() + if parent and parent.IsValid(): + parent_world = xf_cache.GetLocalToWorldTransform(parent) + parent_scale = Gf.Vec3d(*[parent_world.GetRow(i)[:3] for i in range(3)]) + parent_scale = Gf.Vec3d( + Gf.Vec3d(*parent_world.GetRow(0)[:3]).GetLength(), + Gf.Vec3d(*parent_world.GetRow(1)[:3]).GetLength(), + Gf.Vec3d(*parent_world.GetRow(2)[:3]).GetLength(), + ) + else: + parent_scale = Gf.Vec3d(1.0, 1.0, 1.0) + local_scale = Gf.Vec3d( + scales_np[idx][0] / parent_scale[0], + scales_np[idx][1] / parent_scale[1], + scales_np[idx][2] / parent_scale[2], + ) + prim.GetAttribute("xformOp:scale").Set(local_scale) + + def _get_scales_default(self, indices: wp.array | None = None) -> wp.array: + """USD default: get_scales returns local scales.""" + return self.get_local_scales(indices) + + def _set_scales_default(self, scales: wp.array, indices: wp.array | None = None) -> None: + """USD default: set_scales writes local scales.""" + self.set_local_scales(scales, indices) + def set_visibility(self, visibility: torch.Tensor, indices: wp.array | None = None): """Set visibility for prims in the view. @@ -308,8 +351,8 @@ 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) -> ProxyArray: - """Get scales for prims in the view. + def get_local_scales(self, indices: wp.array | None = None) -> wp.array: + """Get local-space scales (xformOp:scale) for prims in the view. Args: indices: Indices of prims to get scales for. Defaults to None (all prims). @@ -327,6 +370,31 @@ def get_scales(self, indices: wp.array | None = None) -> ProxyArray: scales_wp = wp.array(np.array(scales, dtype=np.float32), dtype=wp.float32, device=self._device) return ProxyArray(scales_wp) + def get_world_scales(self, indices: wp.array | None = None) -> wp.array: + """Get world-space (composed) scales for prims in the view. + + Computes the effective world-space scale by extracting column lengths + from the world transform matrix. + + Args: + indices: Indices of prims to get scales for. Defaults to None (all prims). + + Returns: + A ``wp.array`` of shape ``(M, 3)``. + """ + indices_list = self._resolve_indices(indices) + xf_cache = UsdGeom.XformCache(Usd.TimeCode.Default()) + + scales = np.empty((len(indices_list), 3), dtype=np.float32) + for idx, prim_idx in enumerate(indices_list): + prim = self._prims[prim_idx] + world_mtx = xf_cache.GetLocalToWorldTransform(prim) + scales[idx, 0] = Gf.Vec3d(*world_mtx.GetRow(0)[:3]).GetLength() + scales[idx, 1] = Gf.Vec3d(*world_mtx.GetRow(1)[:3]).GetLength() + scales[idx, 2] = Gf.Vec3d(*world_mtx.GetRow(2)[:3]).GetLength() + + return wp.array(scales, dtype=wp.float32, device=self._device) + def get_visibility(self, indices: wp.array | None = None) -> torch.Tensor: """Get visibility for prims in the view. 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 c7c31ae36928..35d5be997c56 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 @@ -3,7 +3,7 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Newton-backed FrameView using Newton body labels and injected sites.""" +"""Newton-backed FrameView — Warp-native, GPU-resident pose queries.""" from __future__ import annotations @@ -11,13 +11,11 @@ import warp as wp -from pxr import UsdPhysics +from pxr import Gf, Usd, UsdGeom import isaaclab.sim as sim_utils -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.utils.string import resolve_matching_names from isaaclab.utils.warp import ProxyArray from isaaclab_newton.physics.newton_manager import NewtonManager @@ -27,8 +25,47 @@ WORLD_BODY_INDEX = -1 +# ------------------------------------------------------------------ +# Warp kernels +# ------------------------------------------------------------------ + + @wp.kernel def _compute_site_world_transforms( + body_q: wp.array(dtype=wp.transformf), + site_body: wp.array(dtype=wp.int32), + site_local: wp.array(dtype=wp.transformf), + out_pos: wp.array(dtype=wp.vec3f), + out_quat: wp.array(dtype=wp.vec4f), +): + """Compute world-space transforms for every site in the view. + + For each site *i*, computes ``world = body_q[site_body[i]] * site_local[i]`` + and splits the result into position and quaternion outputs. When + ``site_body[i] == -1`` the site is world-attached and ``site_local[i]`` is + returned directly. + + Args: + body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. + site_body: Per-site body index (flat model-level), shape ``[num_sites]``. + A value of ``-1`` indicates a world-attached site. + site_local: Per-site local offset relative to its parent body, shape ``[num_sites]``. + out_pos: Output world positions [m], shape ``[num_sites]``. + out_quat: Output world orientations as ``(qx, qy, qz, qw)``, shape ``[num_sites]``. + """ + i = wp.tid() + bid = site_body[i] + if bid == -1: + world = site_local[i] + else: + world = wp.transform_multiply(body_q[bid], site_local[i]) + out_pos[i] = wp.transform_get_translation(world) + q = wp.transform_get_rotation(world) + out_quat[i] = wp.vec4f(q[0], q[1], q[2], q[3]) + + +@wp.kernel +def _compute_site_world_transforms_indexed( body_q: wp.array(dtype=wp.transformf), site_body: wp.array(dtype=wp.int32), site_local: wp.array(dtype=wp.transformf), @@ -36,11 +73,24 @@ def _compute_site_world_transforms( out_pos: wp.array(dtype=wp.vec3f), out_quat: wp.array(dtype=wp.vec4f), ): - """Compute world-space transforms for selected sites.""" + """Indexed variant of :func:`_compute_site_world_transforms`. + + Only computes world transforms for the subset of sites selected by + ``indices``. Thread *i* reads ``indices[i]`` to obtain the site index, + then writes the result to ``out_pos[i]`` / ``out_quat[i]``. + + Args: + body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. + site_body: Per-site body index (flat model-level), shape ``[num_sites]``. + site_local: Per-site local offset relative to its parent body, shape ``[num_sites]``. + indices: Site indices to query, shape ``[M]``. + out_pos: Output world positions [m], shape ``[M]``. + out_quat: Output world orientations as ``(qx, qy, qz, qw)``, shape ``[M]``. + """ i = wp.tid() si = indices[i] bid = site_body[si] - if bid == WORLD_BODY_INDEX: + if bid == -1: world = site_local[si] else: world = wp.transform_multiply(body_q[bid], site_local[si]) @@ -50,23 +100,169 @@ def _compute_site_world_transforms( @wp.kernel -def _gather_site_local_transforms( - site_local: wp.array(dtype=wp.transformf), +def _gather_scales( + shape_scale: wp.array(dtype=wp.vec3f), + shape_body: wp.array(dtype=wp.int32), + site_body: wp.array(dtype=wp.int32), + num_shapes: wp.int32, + out_scales: wp.array(dtype=wp.vec3f), +): + """Gather per-site scales from collision shapes on the same body. + + For each site *i*, linearly scans all shapes to find the first one whose + ``shape_body`` matches ``site_body[i]`` and copies its scale. Falls back + to ``(1, 1, 1)`` if no shape is found on that body. + + Args: + shape_scale: Per-shape scale vectors from the Newton model, shape ``[num_shapes]``. + shape_body: Per-shape parent body index, shape ``[num_shapes]``. + site_body: Per-site body index, shape ``[num_sites]``. + num_shapes: Total number of shapes in the model. + out_scales: Output scale per site, shape ``[num_sites]``. + """ + i = wp.tid() + bid = site_body[i] + found = int(0) + for s in range(num_shapes): + if shape_body[s] == bid and found == 0: + out_scales[i] = shape_scale[s] + found = 1 + if found == 0: + out_scales[i] = wp.vec3f(1.0, 1.0, 1.0) + + +@wp.kernel +def _gather_scales_indexed( + shape_scale: wp.array(dtype=wp.vec3f), + shape_body: wp.array(dtype=wp.int32), + site_body: wp.array(dtype=wp.int32), indices: wp.array(dtype=wp.int32), - out_pos: wp.array(dtype=wp.vec3f), - out_quat: wp.array(dtype=wp.vec4f), + num_shapes: wp.int32, + out_scales: wp.array(dtype=wp.vec3f), ): - """Gather local transforms for selected sites.""" + """Indexed variant of :func:`_gather_scales`. + + Args: + shape_scale: Per-shape scale vectors from the Newton model, shape ``[num_shapes]``. + shape_body: Per-shape parent body index, shape ``[num_shapes]``. + site_body: Per-site body index, shape ``[num_sites]``. + indices: Site indices to query, shape ``[M]``. + num_shapes: Total number of shapes in the model. + out_scales: Output scale per queried site, shape ``[M]``. + """ i = wp.tid() si = indices[i] - local_tf = site_local[si] - out_pos[i] = wp.transform_get_translation(local_tf) - q = wp.transform_get_rotation(local_tf) - out_quat[i] = wp.vec4f(q[0], q[1], q[2], q[3]) + bid = site_body[si] + found = int(0) + for s in range(num_shapes): + if shape_body[s] == bid and found == 0: + out_scales[i] = shape_scale[s] + found = 1 + if found == 0: + out_scales[i] = wp.vec3f(1.0, 1.0, 1.0) + + +@wp.kernel +def _scatter_scales( + site_body: wp.array(dtype=wp.int32), + new_scales: wp.array(dtype=wp.vec3f), + shape_body: wp.array(dtype=wp.int32), + num_shapes: wp.int32, + shape_scale: wp.array(dtype=wp.vec3f), +): + """Scatter per-site scales to all collision shapes on the same body. + + For each site *i*, writes ``new_scales[i]`` to every shape whose + ``shape_body`` matches ``site_body[i]``. Multiple shapes on the same + body all receive the same scale. + + Args: + site_body: Per-site body index, shape ``[num_sites]``. + new_scales: New scale to apply per site, shape ``[num_sites]``. + shape_body: Per-shape parent body index, shape ``[num_shapes]``. + num_shapes: Total number of shapes in the model. + shape_scale: Per-shape scale vectors to write into (modified in-place), + shape ``[num_shapes]``. + """ + i = wp.tid() + bid = site_body[i] + for s in range(num_shapes): + if shape_body[s] == bid: + shape_scale[s] = new_scales[i] + + +@wp.kernel +def _scatter_scales_indexed( + site_body: wp.array(dtype=wp.int32), + indices: wp.array(dtype=wp.int32), + new_scales: wp.array(dtype=wp.vec3f), + shape_body: wp.array(dtype=wp.int32), + num_shapes: wp.int32, + shape_scale: wp.array(dtype=wp.vec3f), +): + """Indexed variant of :func:`_scatter_scales`. + + Args: + site_body: Per-site body index, shape ``[num_sites]``. + indices: Site indices to update, shape ``[M]``. + new_scales: New scale to apply per selected site, shape ``[M]``. + shape_body: Per-shape parent body index, shape ``[num_shapes]``. + num_shapes: Total number of shapes in the model. + shape_scale: Per-shape scale vectors to write into (modified in-place), + shape ``[num_shapes]``. + """ + i = wp.tid() + si = indices[i] + bid = site_body[si] + for s in range(num_shapes): + if shape_body[s] == bid: + shape_scale[s] = new_scales[i] + + +# ------------------------------------------------------------------ +# World-pose site_local write kernels +# ------------------------------------------------------------------ @wp.kernel def _write_site_local_from_world_poses( + body_q: wp.array(dtype=wp.transformf), + site_body: wp.array(dtype=wp.int32), + world_pos: wp.array(dtype=wp.vec3f), + world_quat: wp.array(dtype=wp.vec4f), + site_local: wp.array(dtype=wp.transformf), +): + """Update site local offsets so that the sites reach desired world poses. + + For each site *i*, computes + ``site_local[i] = inv(body_q[site_body[i]]) * desired_world`` so that + a subsequent ``body_q[bid] * site_local[i]`` yields the requested world + pose. For world-attached sites (``site_body[i] == -1``) the desired world + transform is written directly into ``site_local[i]``. + + Does **not** modify ``body_q``. + + Args: + body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. + site_body: Per-site body index (flat model-level), shape ``[num_sites]``. + world_pos: Desired world positions [m], shape ``[num_sites]``. + world_quat: Desired world orientations as ``(qx, qy, qz, qw)``, shape ``[num_sites]``. + site_local: Per-site local offset (modified in-place), shape ``[num_sites]``. + """ + i = wp.tid() + w_pos = world_pos[i] + w_q = world_quat[i] + desired_world = wp.transform(w_pos, wp.quatf(w_q[0], w_q[1], w_q[2], w_q[3])) + + bid = site_body[i] + if bid == -1: + site_local[i] = desired_world + else: + site_local[i] = wp.transform_multiply(wp.transform_inverse(body_q[bid]), desired_world) + + +@wp.kernel +def _write_site_local_from_world_poses_indexed( body_q: wp.array(dtype=wp.transformf), site_body: wp.array(dtype=wp.int32), indices: wp.array(dtype=wp.int32), @@ -74,7 +270,16 @@ def _write_site_local_from_world_poses( world_quat: wp.array(dtype=wp.vec4f), site_local: wp.array(dtype=wp.transformf), ): - """Update local offsets so selected sites reach desired world poses.""" + """Indexed variant of :func:`_write_site_local_from_world_poses`. + + Args: + body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. + site_body: Per-site body index (flat model-level), shape ``[num_sites]``. + indices: Site indices to update, shape ``[M]``. + world_pos: Desired world positions [m], shape ``[M]``. + world_quat: Desired world orientations as ``(qx, qy, qz, qw)``, shape ``[M]``. + site_local: Per-site local offset (modified in-place), shape ``[num_sites]``. + """ i = wp.tid() si = indices[i] w_pos = world_pos[i] @@ -82,364 +287,435 @@ def _write_site_local_from_world_poses( desired_world = wp.transform(w_pos, wp.quatf(w_q[0], w_q[1], w_q[2], w_q[3])) bid = site_body[si] - if bid == WORLD_BODY_INDEX: + if bid == -1: site_local[si] = desired_world else: site_local[si] = wp.transform_multiply(wp.transform_inverse(body_q[bid]), desired_world) +# ------------------------------------------------------------------ +# Local-pose Warp kernels +# ------------------------------------------------------------------ + + @wp.kernel -def _write_site_local_from_local_poses( - indices: wp.array(dtype=wp.int32), - local_pos: wp.array(dtype=wp.vec3f), - local_quat: wp.array(dtype=wp.vec4f), +def _compute_site_local_transforms( + body_q: wp.array(dtype=wp.transformf), + site_body: wp.array(dtype=wp.int32), site_local: wp.array(dtype=wp.transformf), + parent_site_body: wp.array(dtype=wp.int32), + parent_site_local: wp.array(dtype=wp.transformf), + out_pos: wp.array(dtype=wp.vec3f), + out_quat: wp.array(dtype=wp.vec4f), ): - """Update local offsets for selected sites.""" + """Compute parent-relative transforms for every site in the view. + + For each site *i*, computes the world pose of both the site and its USD + parent, then returns ``inv(parent_world) * prim_world``. When + ``site_body[i] == -1`` the site is world-attached and ``site_local[i]`` + is used as the world transform directly. The same convention applies to + the parent arrays. + + Args: + body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. + site_body: Per-site body index (flat model-level), shape ``[num_sites]``. + site_local: Per-site local offset relative to its parent body, shape ``[num_sites]``. + parent_site_body: Per-site USD-parent body index, shape ``[num_sites]``. + parent_site_local: Per-site USD-parent local offset, shape ``[num_sites]``. + out_pos: Output parent-relative positions [m], shape ``[num_sites]``. + out_quat: Output parent-relative orientations as ``(qx, qy, qz, qw)``, + shape ``[num_sites]``. + """ i = wp.tid() - si = indices[i] - l_pos = local_pos[i] - l_q = local_quat[i] - site_local[si] = wp.transform(l_pos, wp.quatf(l_q[0], l_q[1], l_q[2], l_q[3])) + prim_bid = site_body[i] + if prim_bid == -1: + prim_world = site_local[i] + else: + prim_world = wp.transform_multiply(body_q[prim_bid], site_local[i]) + + parent_bid = parent_site_body[i] + if parent_bid == -1: + parent_world = parent_site_local[i] + else: + parent_world = wp.transform_multiply(body_q[parent_bid], parent_site_local[i]) + + local_tf = wp.transform_multiply(wp.transform_inverse(parent_world), prim_world) + out_pos[i] = wp.transform_get_translation(local_tf) + q = wp.transform_get_rotation(local_tf) + out_quat[i] = wp.vec4f(q[0], q[1], q[2], q[3]) @wp.kernel -def _gather_scales( - shape_scale: wp.array(dtype=wp.vec3f), - shape_body: wp.array(dtype=wp.int32), +def _compute_site_local_transforms_indexed( + body_q: wp.array(dtype=wp.transformf), site_body: wp.array(dtype=wp.int32), + site_local: wp.array(dtype=wp.transformf), + parent_site_body: wp.array(dtype=wp.int32), + parent_site_local: wp.array(dtype=wp.transformf), indices: wp.array(dtype=wp.int32), - num_shapes: wp.int32, - out_scales: wp.array(dtype=wp.vec3f), + out_pos: wp.array(dtype=wp.vec3f), + out_quat: wp.array(dtype=wp.vec4f), ): - """Gather per-site scales from collision shapes on the same body.""" + """Indexed variant of :func:`_compute_site_local_transforms`. + + Args: + body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. + site_body: Per-site body index (flat model-level), shape ``[num_sites]``. + site_local: Per-site local offset relative to its parent body, shape ``[num_sites]``. + parent_site_body: Per-site USD-parent body index, shape ``[num_sites]``. + parent_site_local: Per-site USD-parent local offset, shape ``[num_sites]``. + indices: Site indices to query, shape ``[M]``. + out_pos: Output parent-relative positions [m], shape ``[M]``. + out_quat: Output parent-relative orientations as ``(qx, qy, qz, qw)``, + shape ``[M]``. + """ i = wp.tid() si = indices[i] - bid = site_body[si] - found = int(0) - for s in range(num_shapes): - if shape_body[s] == bid and found == 0: - out_scales[i] = shape_scale[s] - found = 1 - if found == 0: - out_scales[i] = wp.vec3f(1.0, 1.0, 1.0) + prim_bid = site_body[si] + if prim_bid == -1: + prim_world = site_local[si] + else: + prim_world = wp.transform_multiply(body_q[prim_bid], site_local[si]) + + parent_bid = parent_site_body[si] + if parent_bid == -1: + parent_world = parent_site_local[si] + else: + parent_world = wp.transform_multiply(body_q[parent_bid], parent_site_local[si]) + + local_tf = wp.transform_multiply(wp.transform_inverse(parent_world), prim_world) + out_pos[i] = wp.transform_get_translation(local_tf) + q = wp.transform_get_rotation(local_tf) + out_quat[i] = wp.vec4f(q[0], q[1], q[2], q[3]) @wp.kernel -def _scatter_scales( +def _write_site_local_from_local_poses( + body_q: wp.array(dtype=wp.transformf), site_body: wp.array(dtype=wp.int32), + parent_site_body: wp.array(dtype=wp.int32), + parent_site_local: wp.array(dtype=wp.transformf), + local_pos: wp.array(dtype=wp.vec3f), + local_quat: wp.array(dtype=wp.vec4f), + site_local: wp.array(dtype=wp.transformf), +): + """Update site local offsets so that sites reach desired parent-relative poses. + + For each site *i*, reconstructs the desired world pose as + ``parent_world * desired_local``, then solves for the body-relative offset: + ``site_local[i] = inv(body_q[bid]) * desired_world``. For world-attached + sites (``site_body[i] == -1``) the world transform is written directly. + + Does **not** modify ``body_q``. + + Args: + body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. + site_body: Per-site body index (flat model-level), shape ``[num_sites]``. + parent_site_body: Per-site USD-parent body index, shape ``[num_sites]``. + parent_site_local: Per-site USD-parent local offset, shape ``[num_sites]``. + local_pos: Desired parent-relative positions [m], shape ``[num_sites]``. + local_quat: Desired parent-relative orientations as ``(qx, qy, qz, qw)``, + shape ``[num_sites]``. + site_local: Per-site local offset (modified in-place), shape ``[num_sites]``. + """ + i = wp.tid() + parent_bid = parent_site_body[i] + if parent_bid == -1: + parent_world = parent_site_local[i] + else: + parent_world = wp.transform_multiply(body_q[parent_bid], parent_site_local[i]) + + l_pos = local_pos[i] + l_q = local_quat[i] + local_tf = wp.transform(l_pos, wp.quatf(l_q[0], l_q[1], l_q[2], l_q[3])) + desired_world = wp.transform_multiply(parent_world, local_tf) + + bid = site_body[i] + if bid == -1: + site_local[i] = desired_world + else: + site_local[i] = wp.transform_multiply(wp.transform_inverse(body_q[bid]), desired_world) + + +@wp.kernel +def _write_site_local_from_local_poses_indexed( + body_q: wp.array(dtype=wp.transformf), + site_body: wp.array(dtype=wp.int32), + parent_site_body: wp.array(dtype=wp.int32), + parent_site_local: wp.array(dtype=wp.transformf), indices: wp.array(dtype=wp.int32), - new_scales: wp.array(dtype=wp.vec3f), - shape_body: wp.array(dtype=wp.int32), - num_shapes: wp.int32, - shape_scale: wp.array(dtype=wp.vec3f), + local_pos: wp.array(dtype=wp.vec3f), + local_quat: wp.array(dtype=wp.vec4f), + site_local: wp.array(dtype=wp.transformf), ): - """Scatter per-site scales to collision shapes on the same body.""" + """Indexed variant of :func:`_write_site_local_from_local_poses`. + + Args: + body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. + site_body: Per-site body index (flat model-level), shape ``[num_sites]``. + parent_site_body: Per-site USD-parent body index, shape ``[num_sites]``. + parent_site_local: Per-site USD-parent local offset, shape ``[num_sites]``. + indices: Site indices to update, shape ``[M]``. + local_pos: Desired parent-relative positions [m], shape ``[M]``. + local_quat: Desired parent-relative orientations as ``(qx, qy, qz, qw)``, + shape ``[M]``. + site_local: Per-site local offset (modified in-place), shape ``[num_sites]``. + """ i = wp.tid() si = indices[i] + parent_bid = parent_site_body[si] + if parent_bid == -1: + parent_world = parent_site_local[si] + else: + parent_world = wp.transform_multiply(body_q[parent_bid], parent_site_local[si]) + + l_pos = local_pos[i] + l_q = local_quat[i] + local_tf = wp.transform(l_pos, wp.quatf(l_q[0], l_q[1], l_q[2], l_q[3])) + desired_world = wp.transform_multiply(parent_world, local_tf) + bid = site_body[si] - for s in range(num_shapes): - if shape_body[s] == bid: - shape_scale[s] = new_scales[i] + if bid == -1: + site_local[si] = desired_world + else: + site_local[si] = wp.transform_multiply(wp.transform_inverse(body_q[bid]), desired_world) -class NewtonSiteFrameView(BaseFrameView): - """Batched Newton site view for non-physics frames. +# ------------------------------------------------------------------ +# View class +# ------------------------------------------------------------------ + - The public construction contract matches the generic :class:`FrameView`: - callers provide a prim expression and the backend resolves the source prim - into Newton body-local or world-local sites. +class NewtonSiteFrameView(BaseFrameView): + """Batched prim view for non-physics prims tracked as sites on Newton bodies. + + Each matched USD prim must be a **non-physics** prim (camera, sensor, + Xform marker, etc.) that sits as a child of a Newton rigid body in the + USD hierarchy. The prim path must **not** resolve directly to a physics + body or collision shape -- those are owned by Newton and should be + accessed through :class:`~isaaclab_newton.assets.Articulation` or + :class:`~isaaclab_newton.assets.RigidObject` instead. + + At init time each prim is resolved to a ``(body_index, site_local)`` + pair via ancestor walk: the nearest ancestor that appears in + ``model.body_label`` becomes the attachment body, and the relative USD + transform becomes the site offset. If no body ancestor exists the prim + is attached to the world frame (``body_index = -1``). + + World poses are computed on GPU as + ``body_q[body_index] * site_local`` via a Warp kernel. Both + ``set_world_poses`` and ``set_local_poses`` update ``site_local`` -- + neither touches ``body_q``. + + Pose getters return :class:`~isaaclab.utils.warp.ProxyArray`. Setters accept ``wp.array``. + + Raises: + ValueError: If any matched prim resolves to a Newton physics body + or collision shape. """ - def __init__( - self, - prim_path: str | list[str], - device: str = "cpu", - validate_xform_ops: bool = True, - stage: object | None = None, - **kwargs, - ): - """Initialize the Newton site frame view. + def __init__(self, prim_path: str, device: str = "cpu", stage: Usd.Stage | None = None, **kwargs): + """Initialize the Newton site-based frame view. + + Resolves all USD prims matching ``prim_path`` and, for each one, walks + the USD ancestor hierarchy to find the nearest Newton rigid body. The + relative transform between the prim and its ancestor body becomes the + site's local offset. + + If the Newton model is already finalized the view initializes + immediately; otherwise initialization is deferred to a + :attr:`PhysicsEvent.PHYSICS_READY` callback. Args: - prim_path: User-facing frame path pattern, or list of patterns. - device: Warp device for GPU arrays. - validate_xform_ops: Whether to validate source USD xform ops. - stage: USD stage that contains the source prims. - **kwargs: Unused. + prim_path: USD prim path pattern (may contain regex). + device: Warp device for GPU arrays (e.g. ``"cuda:0"``). + stage: USD stage to search. Defaults to the current stage. + **kwargs: Unused; accepted for interface compatibility with other + :class:`~isaaclab.sim.views.BaseFrameView` backends. """ - del kwargs - - self._prim_paths = [prim_path] if isinstance(prim_path, str) else list(prim_path) - self._prim_path = prim_path if isinstance(prim_path, str) else ", ".join(self._prim_paths) + self._prim_path = prim_path self._device = device - self._prims = [] 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_body: wp.array | None = None - self._site_local: 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._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._count = 0 + self._prims: list[Usd.Prim] = sim_utils.find_matching_prims(prim_path, stage=stage) model = NewtonManager.get_model() if model is not None: - self._initialize_from_specs(model) + self._initialize_impl(model) else: - for body_patterns, xform, 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)) - else: - for body_pattern in body_patterns: - self._site_labels.append(NewtonManager.cl_register_site(body_pattern, xform)) 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]]: - """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]] = [] - - for path_expr in self._prim_paths: - if resolve_matching_names(path_expr, body_labels, raise_when_no_match=False)[1]: - raise ValueError( - f"FrameView prim '{path_expr}' is a Newton physics body. " - "FrameView should only be used for non-physics frames." - ) - if resolve_matching_names(path_expr, shape_labels, raise_when_no_match=False)[1]: - raise ValueError( - f"FrameView prim '{path_expr}' is a Newton collision shape. " - "FrameView should only be used for non-physics frames." - ) - matches = tuple(iter_clone_plan_matches(plan, path_expr)) if plan is not None else () - if matches: - for source_root, destination_template, source_path, env_ids in matches: - source_prim = None - if not any(token in source_path for token in "*[]()+?|\\"): - source_prim = stage.GetPrimAtPath(source_path) - if source_prim is None or not source_prim.IsValid(): - source_prim = sim_utils.find_first_matching_prim(source_path, stage) - if source_prim is None or not source_prim.IsValid(): - raise RuntimeError(f"FrameView '{path_expr}' could not resolve source prim '{source_path}'.") - specs.append( - self._resolve_source_prim( - source_prim, - validate_xform_ops, - source_root, - destination_template, - env_ids, - use_clone_body_pattern, - stage, - ) - ) - continue - - prim = sim_utils.find_first_matching_prim(path_expr, stage) - if prim is None or not prim.IsValid(): - raise RuntimeError(f"FrameView '{path_expr}' could not resolve a source prim.") - specs.append( - self._resolve_source_prim(prim, validate_xform_ops, None, None, None, use_clone_body_pattern, stage) - ) - - return specs - - def _resolve_source_prim( - self, - prim, - validate_xform_ops: bool, - source_root: str | None, - destination_template: str | None, - 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.""" - prim_path = prim.GetPath().pathString - if prim.HasAPI(UsdPhysics.RigidBodyAPI) or prim.HasAPI(UsdPhysics.ArticulationRootAPI): - raise ValueError( - f"FrameView prim '{prim_path}' is a Newton physics body. " - "FrameView should only be used for non-physics frames." + self._on_physics_ready, PhysicsEvent.PHYSICS_READY, name=f"site_view_{prim_path}" ) - if validate_xform_ops: - sim_utils.standardize_xform_ops(prim) - if not sim_utils.validate_standard_xform_ops(prim): - raise ValueError(f"FrameView prim '{prim_path}' does not have standard xform ops.") - - body_prim = prim.GetParent() - while body_prim and body_prim.IsValid(): - if body_prim.HasAPI(UsdPhysics.RigidBodyAPI) or body_prim.HasAPI(UsdPhysics.ArticulationRootAPI): - pos, quat = sim_utils.resolve_prim_pose(prim, body_prim) - body_path = body_prim.GetPath().pathString - if source_root is not None and destination_template is not None: - assert env_ids is not None - if body_path == source_root: - suffix = "" - elif body_path.startswith(source_root + "/"): - suffix = body_path[len(source_root) :] - elif source_root.startswith(body_path + "/"): - suffix = source_root[len(body_path) :] - if use_clone_body_pattern: - destination_root = destination_template.format(".*") - if not destination_root.endswith(suffix): - 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 - body_patterns = [] - for env_id in env_ids: - destination_root = destination_template.format(env_id) - if not destination_root.endswith(suffix): - raise RuntimeError( - 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 - else: - raise RuntimeError(f"FrameView source body '{body_path}' is not under '{source_root}'.") - if use_clone_body_pattern: - body_patterns = (destination_template.format(".*") + suffix,) - else: - 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 - body_prim = body_prim.GetParent() - - ref_path = source_root - if source_root is not None and destination_template is not None: - 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 def _on_physics_ready(self, _event) -> None: """Callback invoked when the Newton model becomes available.""" - self._initialize_from_site_map(NewtonManager.get_model()) + self._initialize_impl(NewtonManager.get_model()) - def _initialize_from_site_map(self, model) -> None: - """Initialize arrays from injected Newton sites.""" - site_map = NewtonManager._cl_site_index_map - body_t = wp.to_torch(model.shape_body) - xform_t = wp.to_torch(model.shape_transform) - site_bodies: list[int] = [] - site_locals: list[list[float]] = [] - - for site_label in self._site_labels: - 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] - ) - 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()]) + def _initialize_impl(self, model) -> None: + """Resolve USD prims to Newton body indices and allocate GPU buffers.""" + body_labels = list(model.body_label) + body_label_set = set(body_labels) + body_label_to_idx = {path: idx for idx, path in enumerate(body_labels)} + shape_label_set = set(model.shape_label) - self._create_buffers(site_bodies, site_locals) + xform_cache = UsdGeom.XformCache(Usd.TimeCode.Default()) - 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]] = [] + parent_bodies: list[int] = [] + parent_locals: list[list[float]] = [] + + identity_xform = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0] + resolve_cache: dict[str, tuple[int, list[float]]] = {} + + for prim in self._prims: + pp = prim.GetPath().pathString + if pp in body_label_set: + raise ValueError( + f"FrameView prim '{pp}' is a Newton physics body. " + "FrameView should only be used for non-physics prims (cameras, sensors, Xform markers). " + "Use Articulation or RigidObject APIs to control physics bodies." + ) + if pp in shape_label_set: + raise ValueError( + f"FrameView prim '{pp}' is a Newton collision shape. " + "FrameView should only be used for non-physics prims (cameras, sensors, Xform markers). " + "Use Articulation or RigidObject APIs to control collision shapes." + ) - for body_patterns, xform, per_world, env_ids in self._site_specs: - if body_patterns is None: - if per_world: - if NewtonManager._world_xforms is None: - raise RuntimeError(f"FrameView '{self._prim_path}' needs Newton cloned-world transforms.") - world_ids = range(len(NewtonManager._world_xforms)) if env_ids is None else env_ids - for world_id in world_ids: - 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)]) + body_idx, local_xform = self._resolve_ancestor_body(prim, body_label_to_idx, xform_cache) + site_bodies.append(body_idx) + site_locals.append(local_xform) + + parent = prim.GetParent() + if not parent or not parent.IsValid() or parent.GetPath().pathString == "/": + parent_bodies.append(WORLD_BODY_INDEX) + parent_locals.append(identity_xform) + else: + parent_path = parent.GetPath().pathString + if parent_path in resolve_cache: + pb_idx, pb_local = resolve_cache[parent_path] + elif parent_path in body_label_to_idx: + pb_idx = body_label_to_idx[parent_path] + pb_local = identity_xform + resolve_cache[parent_path] = (pb_idx, pb_local) else: - site_bodies.append(WORLD_BODY_INDEX) - site_locals.append([float(v) for v in xform]) - continue - - for body_pattern in body_patterns: - matched_indices, _ = resolve_matching_names(body_pattern, body_labels, raise_when_no_match=False) - if not matched_indices: - raise ValueError( - f"FrameView '{self._prim_path}' body pattern '{body_pattern}' matched no Newton bodies." - ) - - for body_idx in matched_indices: - site_bodies.append(body_idx) - site_locals.append([float(v) for v in xform]) - - self._create_buffers(site_bodies, site_locals) - - def _create_buffers(self, site_bodies: list[int], site_locals: list[list[float]]) -> None: - """Allocate view buffers from body indices and local transforms.""" - self._count = len(site_bodies) + pb_idx, pb_local = self._resolve_ancestor_body(parent, body_label_to_idx, xform_cache) + resolve_cache[parent_path] = (pb_idx, pb_local) + parent_bodies.append(pb_idx) + parent_locals.append(pb_local) + 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_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._site_local = wp.array( + [wp.transform(*x) for x in site_locals], + dtype=wp.transformf, + device=device, + ) + self._parent_site_body = wp.array(parent_bodies, dtype=wp.int32, device=device) + self._parent_site_local = wp.array( + [wp.transform(*x) for x in parent_locals], + dtype=wp.transformf, + 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._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) - @property - def prims(self) -> list: - """List of USD prims being managed by this view. + @staticmethod + def _resolve_ancestor_body( + prim: Usd.Prim, + body_label_to_idx: dict[str, int], + xform_cache: UsdGeom.XformCache, + ) -> tuple[int, list[float]]: + """Walk USD ancestors to find the nearest Newton body and compute the relative local transform. - Newton site views do not retain USD prim handles. + Args: + prim: The USD prim to resolve. + body_label_to_idx: Dict mapping body prim paths to their Newton body indices. + xform_cache: USD xform cache for efficient transform lookups. + + Returns: + A tuple ``(body_index, local_xform_7)`` where *local_xform_7* is + ``[tx, ty, tz, qx, qy, qz, qw]``. If no body ancestor exists, + ``body_index`` is :data:`WORLD_BODY_INDEX` and the local transform + is the prim's world transform. """ + prim_world_tf = xform_cache.GetLocalToWorldTransform(prim) + prim_world_tf.Orthonormalize() + + ancestor = prim.GetParent() + while ancestor and ancestor.IsValid() and ancestor.GetPath().pathString != "/": + ancestor_path = ancestor.GetPath().pathString + body_idx = body_label_to_idx.get(ancestor_path) + if body_idx is not None: + ancestor_world_tf = xform_cache.GetLocalToWorldTransform(ancestor) + ancestor_world_tf.Orthonormalize() + local_tf = prim_world_tf * ancestor_world_tf.GetInverse() + return body_idx, _gf_matrix_to_xform7(local_tf) + ancestor = ancestor.GetParent() + + return WORLD_BODY_INDEX, _gf_matrix_to_xform7(prim_world_tf) + + @property + def prims(self) -> list: + """List of USD prims being managed by this view.""" return self._prims @property def count(self) -> int: - """Number of frames in this view.""" - return self._count + """Number of prims in this view.""" + return len(self._prims) @property def device(self) -> str: - """Device where arrays are allocated.""" + """Device where arrays are allocated (cpu or cuda).""" return self._device + # ------------------------------------------------------------------ + # World poses + # ------------------------------------------------------------------ + def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: - """Get world-space positions and orientations.""" + """Get world-space positions and orientations. + + Args: + indices: Subset of sites to query. ``None`` means all sites. + + 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. + """ state = NewtonManager.get_state_0() - site_indices = self._site_indices if indices is None else indices - n = self.count if indices is None else len(indices) - pos_buf = self._pos_buf if indices is None else wp.zeros(n, dtype=wp.vec3f, device=self._device) - quat_buf = self._quat_buf if indices is None else wp.zeros(n, dtype=wp.vec4f, device=self._device) + + if indices is not None: + n = len(indices) + pos_buf = wp.zeros(n, dtype=wp.vec3f, device=self._device) + quat_buf = wp.zeros(n, dtype=wp.vec4f, device=self._device) + wp.launch( + _compute_site_world_transforms_indexed, + dim=n, + inputs=[state.body_q, self._site_body, self._site_local, indices], + outputs=[pos_buf, quat_buf], + device=self._device, + ) + return ProxyArray(pos_buf), ProxyArray(quat_buf) wp.launch( _compute_site_world_transforms, - dim=n, - inputs=[state.body_q, self._site_body, self._site_local, site_indices], - outputs=[pos_buf, quat_buf], + dim=self.count, + inputs=[state.body_q, self._site_body, self._site_local], + outputs=[self._pos_buf, self._quat_buf], device=self._device, ) - if indices is None: - return self._pos_ta, self._quat_ta - return ProxyArray(pos_buf), ProxyArray(quat_buf) + return self._pos_ta, self._quat_ta def set_world_poses( self, @@ -447,11 +723,24 @@ def set_world_poses( orientations: wp.array | None = None, indices: wp.array | None = None, ) -> None: - """Set world-space positions and/or orientations.""" + """Set world-space positions and/or orientations. + + Updates the internal ``site_local`` offsets so that + ``body_q[body] * new_site_local`` yields the desired world pose. + Does **not** modify ``body_q``. + + Args: + positions: Desired world positions ``(M, 3)``. ``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. + """ if positions is None and orientations is None: return state = NewtonManager.get_state_0() + if positions is None or orientations is None: cur_pos_ta, cur_quat_ta = self.get_world_poses(indices) if positions is None: @@ -459,32 +748,74 @@ def set_world_poses( if orientations is None: orientations = cur_quat_ta.warp - site_indices = self._site_indices if indices is None else indices - n = self.count if indices is None else len(indices) - wp.launch( - _write_site_local_from_world_poses, - dim=n, - inputs=[state.body_q, self._site_body, site_indices, positions, orientations, self._site_local], - device=self._device, - ) + if indices is not None: + wp.launch( + _write_site_local_from_world_poses_indexed, + dim=len(indices), + inputs=[state.body_q, self._site_body, indices, positions, orientations, self._site_local], + device=self._device, + ) + else: + wp.launch( + _write_site_local_from_world_poses, + dim=self.count, + inputs=[state.body_q, self._site_body, positions, orientations, self._site_local], + device=self._device, + ) + + # ------------------------------------------------------------------ + # Local poses (parent-relative) + # ------------------------------------------------------------------ def get_local_poses(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) - pos_buf = self._local_pos_buf if indices is None else wp.zeros(n, dtype=wp.vec3f, device=self._device) - quat_buf = self._local_quat_buf if indices is None else wp.zeros(n, dtype=wp.vec4f, device=self._device) + """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. Use ``.warp`` for the underlying ``wp.array`` or ``.torch`` for a + cached zero-copy ``torch.Tensor`` view. + """ + state = NewtonManager.get_state_0() + + if indices is not None: + n = len(indices) + pos_buf = wp.zeros(n, dtype=wp.vec3f, device=self._device) + quat_buf = wp.zeros(n, dtype=wp.vec4f, device=self._device) + wp.launch( + _compute_site_local_transforms_indexed, + dim=n, + inputs=[ + state.body_q, + self._site_body, + self._site_local, + self._parent_site_body, + self._parent_site_local, + indices, + ], + outputs=[pos_buf, quat_buf], + device=self._device, + ) + return ProxyArray(pos_buf), ProxyArray(quat_buf) wp.launch( - _gather_site_local_transforms, - dim=n, - inputs=[self._site_local, site_indices], - outputs=[pos_buf, quat_buf], + _compute_site_local_transforms, + dim=self.count, + inputs=[ + state.body_q, + self._site_body, + self._site_local, + self._parent_site_body, + self._parent_site_local, + ], + outputs=[self._local_pos_buf, self._local_quat_buf], device=self._device, ) - if indices is None: - return self._local_pos_ta, self._local_quat_ta - return ProxyArray(pos_buf), ProxyArray(quat_buf) + return self._local_pos_ta, self._local_quat_ta def set_local_poses( self, @@ -492,10 +823,24 @@ def set_local_poses( orientations: wp.array | None = None, indices: wp.array | None = None, ) -> None: - """Set body-local translations and/or orientations.""" + """Set parent-relative translations and/or orientations. + + Updates the internal ``site_local`` offsets so that + ``inv(parent_world) * (body_q[bid] * site_local)`` yields the desired + local pose. Does **not** modify ``body_q``. + + Args: + translations: Desired parent-relative translations ``(M, 3)``. + ``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. + """ if translations is None and orientations is None: return + state = NewtonManager.get_state_0() + if translations is None or orientations is None: cur_pos_ta, cur_quat_ta = self.get_local_poses(indices) if translations is None: @@ -503,40 +848,153 @@ def set_local_poses( if orientations is None: orientations = cur_quat_ta.warp - site_indices = self._site_indices if indices is None else indices - n = self.count if indices is None else len(indices) - wp.launch( - _write_site_local_from_local_poses, - dim=n, - inputs=[site_indices, translations, orientations, self._site_local], - device=self._device, - ) + if indices is not None: + wp.launch( + _write_site_local_from_local_poses_indexed, + dim=len(indices), + inputs=[ + state.body_q, + self._site_body, + self._parent_site_body, + self._parent_site_local, + indices, + translations, + orientations, + self._site_local, + ], + device=self._device, + ) + else: + wp.launch( + _write_site_local_from_local_poses, + dim=self.count, + inputs=[ + state.body_q, + self._site_body, + self._parent_site_body, + self._parent_site_local, + translations, + orientations, + self._site_local, + ], + device=self._device, + ) + + # ------------------------------------------------------------------ + # Scales + # ------------------------------------------------------------------ - def get_scales(self, indices: wp.array | None = None) -> ProxyArray: - """Get per-site scales by reading from the first collision shape on the same body.""" + def get_local_scales(self, indices: wp.array | None = None) -> wp.array: + """Get local-space scales. + + .. note:: + Newton stores shape_scale as the effective (composed) scale. + get_local_scales returns the same value since Newton does not + decompose parent/child scale independently. + """ + return self._get_shape_scales(indices) + + def get_world_scales(self, indices: wp.array | None = None) -> wp.array: + """Get world-space (composed) scales from Newton shape_scale.""" + return self._get_shape_scales(indices) + + def set_local_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Set local-space scales. + + .. note:: + Newton stores shape_scale as the effective (composed) scale. + set_local_scales writes directly to shape_scale. + """ + self._set_shape_scales(scales, indices) + + def set_world_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Set world-space (composed) scales to Newton shape_scale.""" + self._set_shape_scales(scales, indices) + + def _get_scales_default(self, indices=None): + """Newton default: get_scales returns shape_scale (world-like).""" + return self.get_world_scales(indices) + + def _set_scales_default(self, scales, indices=None): + """Newton default: set_scales writes shape_scale (world-like).""" + self.set_world_scales(scales, indices) + + 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. + + .. deprecated:: + Use :meth:`get_local_scales` or :meth:`get_world_scales` instead. + + Args: + indices: Subset of sites to query. ``None`` means all sites. + + Returns: + A ``wp.array`` of shape ``(M, 3)``. + """ + return self._get_shape_scales(indices) + + 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. + + .. deprecated:: + Use :meth:`set_local_scales` or :meth:`set_world_scales` instead. + + Args: + scales: New scales ``(M, 3)`` as ``wp.array``. + indices: Subset of sites to update. ``None`` means all sites. + """ + self._set_shape_scales(scales, indices) + + def _get_shape_scales(self, indices: wp.array | None = None) -> wp.array: + """Internal: read shape_scale from Newton model.""" 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, - dim=n, - inputs=[model.shape_scale, model.shape_body, self._site_body, site_indices, num_shapes], - outputs=[out], - device=self._device, - ) - 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.""" + if indices is not None: + n = len(indices) + out = wp.zeros(n, dtype=wp.vec3f, device=self._device) + wp.launch( + _gather_scales_indexed, + dim=n, + inputs=[model.shape_scale, model.shape_body, self._site_body, indices, num_shapes], + outputs=[out], + device=self._device, + ) + else: + out = wp.zeros(self.count, dtype=wp.vec3f, device=self._device) + wp.launch( + _gather_scales, + dim=self.count, + inputs=[model.shape_scale, model.shape_body, self._site_body, num_shapes], + outputs=[out], + device=self._device, + ) + return out + + def _set_shape_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Internal: write shape_scale to Newton model.""" 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, - dim=n, - inputs=[self._site_body, site_indices, scales, model.shape_body, num_shapes, model.shape_scale], - device=self._device, - ) + + if indices is not None: + wp.launch( + _scatter_scales_indexed, + dim=len(indices), + inputs=[self._site_body, indices, scales, model.shape_body, num_shapes, model.shape_scale], + device=self._device, + ) + else: + wp.launch( + _scatter_scales, + dim=self.count, + inputs=[self._site_body, scales, model.shape_body, num_shapes, model.shape_scale], + device=self._device, + ) + + +def _gf_matrix_to_xform7(mat: Gf.Matrix4d) -> list[float]: + """Convert a ``Gf.Matrix4d`` to ``[tx, ty, tz, qx, qy, qz, qw]``.""" + t = mat.ExtractTranslation() + 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())] 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 879ecddf385f..eff35de9759d 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 @@ -837,6 +837,9 @@ def _ensure_usd_view(self) -> UsdFrameView: def get_scales(self, indices: wp.array | None = None) -> ProxyArray: """Get prim scales from the USD stage's ``xformOp:scale`` attribute. + .. deprecated:: + Use :meth:`get_local_scales` or :meth:`get_world_scales` instead. + .. note:: This reads the *static* USD authored value, not a live physics-state value. OVPhysX does not maintain a per-shape ``shape_scale`` array @@ -851,11 +854,34 @@ def get_scales(self, indices: wp.array | None = None) -> ProxyArray: Returns: A :class:`~isaaclab.utils.warp.ProxyArray` of shape ``(M, 3)``. """ - return self._ensure_usd_view().get_scales(indices) + return self.get_local_scales(indices) + + def get_world_scales(self, indices: wp.array | None = None) -> wp.array: + """Get world-space (composed) scales via the USD view.""" + return self._ensure_usd_view().get_world_scales(indices) + + def set_local_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Set local-space scales via the USD view.""" + self._ensure_usd_view().set_local_scales(scales, indices) + + def set_world_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Set world-space scales via the USD view.""" + self._ensure_usd_view().set_world_scales(scales, indices) + + def _get_scales_default(self, indices=None): + """OvPhysX default: get_scales returns local scales (same as USD).""" + return self.get_local_scales(indices) + + def _set_scales_default(self, scales, indices=None): + """OvPhysX default: set_scales writes local scales (same as USD).""" + self.set_local_scales(scales, indices) def set_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: """Set prim scales by writing the USD ``xformOp:scale`` attribute. + .. deprecated:: + Use :meth:`set_local_scales` or :meth:`set_world_scales` instead. + .. 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 @@ -866,7 +892,7 @@ def set_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: scales: Scales ``(M, 3)`` as ``wp.array``. indices: Subset of sites to update. ``None`` means all sites. """ - self._ensure_usd_view().set_scales(scales, indices) + self.set_local_scales(scales, indices) def get_visibility(self, indices: wp.array | None = None): """Get visibility for prims in the view (USD-backed). diff --git a/source/isaaclab_physx/changelog.d/fabric-local-poses.rst b/source/isaaclab_physx/changelog.d/fabric-local-poses.rst index 61b59c22398b..546d395cb78e 100644 --- a/source/isaaclab_physx/changelog.d/fabric-local-poses.rst +++ b/source/isaaclab_physx/changelog.d/fabric-local-poses.rst @@ -33,3 +33,14 @@ Added * Added topology-change recovery via automatic ``PrepareForReuse`` detection and per-selection index rebuild. + +Deprecated +^^^^^^^^^^ + +* Deprecated ``get_scales`` / ``set_scales`` on all ``BaseFrameView`` subclasses. + Use the new explicit ``get_local_scales`` / ``set_local_scales`` (operates on + ``xformOp:scale`` / ``localMatrix``) or ``get_world_scales`` / + ``set_world_scales`` (operates on composed world-space scale) instead. + The deprecated methods still work but emit a ``DeprecationWarning``; + ``UsdFrameView`` defaults to local, ``FabricFrameView`` defaults to world + (preserving prior behavior). 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 86be4a5bd3ca..0dba645e5334 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 @@ -390,13 +390,10 @@ def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, # Scales # ------------------------------------------------------------------ - def set_scales(self, scales, indices=None): - # TODO(pv): This decomposes/recomposes the *world* matrix, so the scale - # value is the composed (parent × local) scale. UsdFrameView.set_scales - # writes xformOp:scale which is purely local. The two diverge when the - # parent has non-identity scale. Fix: operate on localMatrix instead. + def set_world_scales(self, scales, indices=None): + """Set world-space (composed) scales by decomposing/recomposing worldMatrix.""" if not self._use_fabric: - self._usd_view.set_scales(scales, indices) + self._usd_view.set_world_scales(scales, indices) return if not self._fabric_initialized: @@ -428,7 +425,7 @@ def set_scales(self, scales, indices=None): # World was just written -- mark local poses as stale. self._dirty = _DirtyFlag.LOCAL - def get_scales(self, indices: wp.array | None = None) -> ProxyArray: + def get_world_scales(self, indices=None): """Return per-prim (sx, sy, sz) scales extracted from world matrix. .. warning:: @@ -436,9 +433,8 @@ def get_scales(self, indices: wp.array | None = None) -> ProxyArray: pre-allocated buffer** that is overwritten on the next call. Do not hold references across calls -- copy if persistence is needed. """ - # TODO(pv): Same world-vs-local divergence as set_scales -- see note above. if not self._use_fabric: - return self._usd_view.get_scales(indices) + return self._usd_view.get_world_scales(indices) if not self._fabric_initialized: self._initialize_fabric() @@ -472,6 +468,92 @@ def get_scales(self, indices: wp.array | None = None) -> ProxyArray: wp.synchronize() return ProxyArray(scales_wp) + def set_local_scales(self, scales, indices=None): + """Set local-space scales by decomposing/recomposing localMatrix.""" + if not self._use_fabric: + self._usd_view.set_local_scales(scales, indices) + return + + if not self._fabric_initialized: + self._initialize_fabric() + + # Sync local matrices first if world writes are pending. + self._sync_local_from_world_if_dirty() + + indices_wp = self._resolve_indices_wp(indices) + scales_wp = self._to_float32_2d_or_empty(scales) + + wp.launch( + kernel=fabric_utils.compose_indexed_fabric_transforms, + dim=indices_wp.shape[0], + inputs=[ + self._get_local_array(), + self._fabric_empty_2d_array_sentinel, + self._fabric_empty_2d_array_sentinel, + scales_wp, + False, + False, + False, + indices_wp, + ], + device=self._device, + ) + wp.synchronize() + + # Local was just written -- mark world poses as stale. + self._dirty = _DirtyFlag.WORLD + + def get_local_scales(self, indices=None): + """Return per-prim (sx, sy, sz) scales extracted from local matrix. + + .. warning:: + When *indices* is None (all prims), the returned array is a **shared + pre-allocated buffer** that is overwritten on the next call. Do not + hold references across calls -- copy if persistence is needed. + """ + if not self._use_fabric: + return self._usd_view.get_local_scales(indices) + + if not self._fabric_initialized: + self._initialize_fabric() + + # Sync local matrices first if world writes are pending. + self._sync_local_from_world_if_dirty() + + indices_wp = self._resolve_indices_wp(indices) + count = indices_wp.shape[0] + + use_cached = indices is None or indices == slice(None) + if use_cached: + scales_wp = self._fabric_scales_buf + else: + scales_wp = wp.zeros((count, 3), dtype=wp.float32, device=self._device) + + wp.launch( + kernel=fabric_utils.decompose_indexed_fabric_transforms, + dim=count, + inputs=[ + self._get_local_array(), + self._fabric_empty_2d_array_sentinel, + self._fabric_empty_2d_array_sentinel, + scales_wp, + indices_wp, + ], + device=self._device, + ) + + if use_cached: + wp.synchronize() + return scales_wp + + def _get_scales_default(self, indices=None): + """Fabric default: get_scales returns world scales (backwards compat).""" + return self.get_world_scales(indices) + + def _set_scales_default(self, scales, indices=None): + """Fabric default: set_scales writes world scales (backwards compat).""" + self.set_world_scales(scales, indices) + # ------------------------------------------------------------------ # Internal -- sync helpers # ------------------------------------------------------------------ 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 441a8bd0ac84..d3f21caa0e05 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 @@ -266,20 +266,64 @@ def test_set_local_via_fabric_path(device, view_factory): @pytest.mark.parametrize("device", ["cuda:0"]) def test_get_scales_fabric_path(device, view_factory): - """Exercise the Fabric-native get_scales path.""" + """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_scales call below uses Fabric. + # Trigger lazy `_initialize_fabric()` so the get_world_scales call below uses Fabric. view.get_world_poses() - scales = view.get_scales() + scales = view.get_world_scales() scales_t = torch.as_tensor(scales, device=device) # 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): + """set_local_scales -> get_local_scales roundtrip 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) + view.set_local_scales(new_scales) + + # Should have dirtied world + assert view._dirty.name == "WORLD" + + ret_scales = view.get_local_scales() + scales_torch = torch.as_tensor(ret_scales, device=device) + 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): + """set_world_scales -> get_world_scales roundtrip 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) + view.set_world_scales(new_scales) + + # Should have dirtied local + assert view._dirty.name == "LOCAL" + + ret_scales = view.get_world_scales() + scales_torch = torch.as_tensor(ret_scales, device=device) + 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 @@ -397,7 +441,7 @@ def test_initial_seed_with_scaled_parent(device): rtol=0, ) - scales = torch.as_tensor(view.get_scales(), device=device) + scales = torch.as_tensor(view.get_world_scales(), device=device) torch.testing.assert_close( scales, torch.tensor([[6.0, 1.0, 1.0]], dtype=torch.float32, device=device), @@ -570,9 +614,9 @@ 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. + """set_world_scales -> get_world_scales roundtrip works on cuda:1. - Both write paths (``set_world_poses`` and ``set_scales``) call + Both write paths (``set_world_poses`` and ``set_world_scales``) call ``_prepare_for_reuse`` and launch on ``self._device``; this test covers the scales path on the non-primary CUDA device. """ @@ -581,9 +625,9 @@ def test_fabric_cuda1_scales_roundtrip(device, view_factory): 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) + view.set_world_scales(new_scales) - ret_scales = view.get_scales() + ret_scales = view.get_world_scales() scales_torch = torch.as_tensor(ret_scales, device=device) 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}" From cf89d7f535c92fc0b3c38de96a25c5f78109b20f Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Tue, 26 May 2026 16:03:28 +0000 Subject: [PATCH 17/54] fix: UsdFrameView.set/get_world_scales Gf.Vec3d type errors - Use row-major indexing (world_mtx[row][col]) instead of GetRow() which returns a Vec4 that can't be unpacked into Vec3d - Cast numpy float32 to Python float for Gf.Vec3d constructor - Replace internal get_scales() call with get_local_scales() in FabricFrameView._initialize_fabric to avoid deprecation warning --- .../isaaclab/sim/views/usd_frame_view.py | 19 +++++++++---------- .../sim/views/fabric_frame_view.py | 2 +- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/source/isaaclab/isaaclab/sim/views/usd_frame_view.py b/source/isaaclab/isaaclab/sim/views/usd_frame_view.py index 166f7c36288f..dd83cfc1f2c9 100644 --- a/source/isaaclab/isaaclab/sim/views/usd_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/usd_frame_view.py @@ -248,18 +248,17 @@ def set_world_scales(self, scales: wp.array, indices: wp.array | None = None): parent = prim.GetParent() if parent and parent.IsValid(): parent_world = xf_cache.GetLocalToWorldTransform(parent) - parent_scale = Gf.Vec3d(*[parent_world.GetRow(i)[:3] for i in range(3)]) parent_scale = Gf.Vec3d( - Gf.Vec3d(*parent_world.GetRow(0)[:3]).GetLength(), - Gf.Vec3d(*parent_world.GetRow(1)[:3]).GetLength(), - Gf.Vec3d(*parent_world.GetRow(2)[:3]).GetLength(), + 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: parent_scale = Gf.Vec3d(1.0, 1.0, 1.0) local_scale = Gf.Vec3d( - scales_np[idx][0] / parent_scale[0], - scales_np[idx][1] / parent_scale[1], - scales_np[idx][2] / parent_scale[2], + 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) @@ -389,9 +388,9 @@ def get_world_scales(self, indices: wp.array | None = None) -> wp.array: for idx, prim_idx in enumerate(indices_list): prim = self._prims[prim_idx] world_mtx = xf_cache.GetLocalToWorldTransform(prim) - scales[idx, 0] = Gf.Vec3d(*world_mtx.GetRow(0)[:3]).GetLength() - scales[idx, 1] = Gf.Vec3d(*world_mtx.GetRow(1)[:3]).GetLength() - scales[idx, 2] = Gf.Vec3d(*world_mtx.GetRow(2)[:3]).GetLength() + 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 wp.array(scales, dtype=wp.float32, device=self._device) 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 0dba645e5334..73fcae410aed 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 @@ -793,7 +793,7 @@ def _sync_fabric_from_usd_initial(self) -> None: """ # --- Children --- pos_ta, ori_ta = self._usd_view.get_world_poses() - scales_obj = self._usd_view.get_scales() + scales_obj = self._usd_view.get_local_scales() scales_wp = ( scales_obj.warp if hasattr(scales_obj, "warp") From 8e00c5acb9094964eb0cba9f433d5456f8c082b2 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Tue, 26 May 2026 16:19:35 +0000 Subject: [PATCH 18/54] cleanup: remove redundant get/set_scales overrides from subclasses BaseFrameView.get_scales/set_scales are now concrete methods that emit DeprecationWarning and delegate to _get_scales_default/_set_scales_default. Subclasses that override them bypass the warning. Remove the overrides from OvPhysxFrameView and NewtonSiteFrameView so users get a consistent deprecation warning regardless of backend. Move the OvPhysxFrameView notes into get_local_scales (the actual impl). --- .../sim/views/newton_site_frame_view.py | 26 ------------ .../sim/views/ovphysx_frame_view.py | 40 +++++++------------ 2 files changed, 15 insertions(+), 51 deletions(-) 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 35d5be997c56..587ff6ea5b9b 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 @@ -919,32 +919,6 @@ def _set_scales_default(self, scales, indices=None): """Newton default: set_scales writes shape_scale (world-like).""" self.set_world_scales(scales, indices) - 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. - - .. deprecated:: - Use :meth:`get_local_scales` or :meth:`get_world_scales` instead. - - Args: - indices: Subset of sites to query. ``None`` means all sites. - - Returns: - A ``wp.array`` of shape ``(M, 3)``. - """ - return self._get_shape_scales(indices) - - 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. - - .. deprecated:: - Use :meth:`set_local_scales` or :meth:`set_world_scales` instead. - - Args: - scales: New scales ``(M, 3)`` as ``wp.array``. - indices: Subset of sites to update. ``None`` means all sites. - """ - self._set_shape_scales(scales, indices) - def _get_shape_scales(self, indices: wp.array | None = None) -> wp.array: """Internal: read shape_scale from Newton model.""" model = NewtonManager.get_model() 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 eff35de9759d..b6c765c424e6 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 @@ -834,11 +834,8 @@ def _ensure_usd_view(self) -> UsdFrameView: ) return self._usd_view - def get_scales(self, indices: wp.array | None = None) -> ProxyArray: - """Get prim scales from the USD stage's ``xformOp:scale`` attribute. - - .. deprecated:: - Use :meth:`get_local_scales` or :meth:`get_world_scales` instead. + def get_local_scales(self, indices: wp.array | None = None) -> wp.array: + """Get local-space scales (xformOp:scale) via the USD view. .. note:: This reads the *static* USD authored value, not a live physics-state @@ -854,14 +851,25 @@ def get_scales(self, indices: wp.array | None = None) -> ProxyArray: Returns: A :class:`~isaaclab.utils.warp.ProxyArray` of shape ``(M, 3)``. """ - return self.get_local_scales(indices) + return self._ensure_usd_view().get_local_scales(indices) def get_world_scales(self, indices: wp.array | None = None) -> wp.array: """Get world-space (composed) scales via the USD view.""" return self._ensure_usd_view().get_world_scales(indices) def set_local_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: - """Set local-space scales via the USD view.""" + """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. + """ self._ensure_usd_view().set_local_scales(scales, indices) def set_world_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: @@ -876,24 +884,6 @@ def _set_scales_default(self, scales, indices=None): """OvPhysX default: set_scales writes local scales (same as USD).""" self.set_local_scales(scales, indices) - def set_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: - """Set prim scales by writing the USD ``xformOp:scale`` attribute. - - .. deprecated:: - Use :meth:`set_local_scales` or :meth:`set_world_scales` instead. - - .. 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. - """ - self.set_local_scales(scales, indices) - def get_visibility(self, indices: wp.array | None = None): """Get visibility for prims in the view (USD-backed). From a304d88db0bc1b743225cc4a5943c33b17a76e22 Mon Sep 17 00:00:00 2001 From: pv-nvidia Date: Sat, 30 May 2026 06:37:44 +0000 Subject: [PATCH 19/54] fix: review improvements to scale API refactoring - Remove redundant child world matrix composition in _sync_fabric_from_usd_initial: the world matrix is immediately recomputed by _recompute_world_from_local() at the end of the method via child_world = child_local * parent_world. Eliminating the dead compose kernel saves one GPU kernel launch during initialization. - Fix UsdFrameView.get_world_scales docstring: correctly describes extracting 'row lengths' (USD uses a row-vector convention where basis vectors are stored in rows), not 'column lengths'. - Add pseudoroot check to UsdFrameView.set_world_scales: mirrors the guard in set_world_poses for consistency. The pseudoroot's identity transform means this is not a bug (parent_scale would be (1,1,1) regardless), but explicit guards prevent surprises with unusual stage structures. - Document that FabricFrameView.get_local_scales and get_world_scales share the same pre-allocated buffer (_fabric_scales_buf). Callers interleaving both getters without copying will see overwritten values. --- .../isaaclab/sim/views/usd_frame_view.py | 7 ++-- .../sim/views/fabric_frame_view.py | 36 ++++++------------- 2 files changed, 15 insertions(+), 28 deletions(-) diff --git a/source/isaaclab/isaaclab/sim/views/usd_frame_view.py b/source/isaaclab/isaaclab/sim/views/usd_frame_view.py index dd83cfc1f2c9..30402e181d19 100644 --- a/source/isaaclab/isaaclab/sim/views/usd_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/usd_frame_view.py @@ -246,7 +246,7 @@ def set_world_scales(self, scales: wp.array, indices: wp.array | None = None): for idx, prim_idx in enumerate(indices_list): prim = self._prims[prim_idx] parent = prim.GetParent() - if parent and parent.IsValid(): + 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(), @@ -372,8 +372,9 @@ def get_local_scales(self, indices: wp.array | None = None) -> wp.array: def get_world_scales(self, indices: wp.array | None = None) -> wp.array: """Get world-space (composed) scales for prims in the view. - Computes the effective world-space scale by extracting column lengths - from the world transform matrix. + Computes the effective world-space scale by extracting row lengths + from the world transform matrix (USD uses a row-vector convention + where each row of the 3x3 sub-matrix is a basis vector). Args: indices: Indices of prims to get scales for. Defaults to None (all prims). 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 73fcae410aed..3ba306d5cea6 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 @@ -430,8 +430,9 @@ def get_world_scales(self, indices=None): .. warning:: When *indices* is None (all prims), the returned array is a **shared - pre-allocated buffer** that is overwritten on the next call. Do not - hold references across calls -- copy if persistence is needed. + pre-allocated buffer** (shared with :meth:`get_local_scales`) that is + overwritten on the next call. Do not hold references across calls -- + copy if persistence is needed. """ if not self._use_fabric: return self._usd_view.get_world_scales(indices) @@ -508,8 +509,9 @@ def get_local_scales(self, indices=None): .. warning:: When *indices* is None (all prims), the returned array is a **shared - pre-allocated buffer** that is overwritten on the next call. Do not - hold references across calls -- copy if persistence is needed. + pre-allocated buffer** (shared with :meth:`get_world_scales`) that is + overwritten on the next call. Do not hold references across calls -- + copy if persistence is needed. """ if not self._use_fabric: return self._usd_view.get_local_scales(indices) @@ -792,7 +794,11 @@ def _sync_fabric_from_usd_initial(self) -> None: getters (which read from Fabric) would return wrong values. """ # --- Children --- - pos_ta, ori_ta = self._usd_view.get_world_poses() + # Compose child localMatrix from USD-authored local transforms. + # The child world matrix is NOT composed here -- it will be computed + # by ``_recompute_world_from_local()`` at the end of this method as + # ``child_world = child_local * parent_world``, which naturally + # composes scales through the matrix multiplication. scales_obj = self._usd_view.get_local_scales() scales_wp = ( scales_obj.warp @@ -802,26 +808,6 @@ def _sync_fabric_from_usd_initial(self) -> None: else self._fabric_empty_2d_array_sentinel ) local_pos_ta, local_ori_ta = self._usd_view.get_local_poses() - # Compose into child worldMatrix. - wp.launch( - kernel=fabric_utils.compose_indexed_fabric_transforms, - dim=self.count, - inputs=[ - self._world_ifa, - _to_float32_2d(pos_ta.warp), - _to_float32_2d(ori_ta.warp), - _to_float32_2d(scales_wp), - False, - False, - False, - self._view_indices, - ], - device=self._device, - ) - # Compose into child localMatrix. Pass the locally-authored scale so - # that a subsequent ``_sync_world_from_local_if_dirty`` produces the - # right world-space scale (``world = parent_world * local`` carries - # ``local``'s scale through the multiply). wp.launch( kernel=fabric_utils.compose_indexed_fabric_transforms, dim=self.count, From 769e4dfc8f642dd34ae645da6ab2176083e76541 Mon Sep 17 00:00:00 2001 From: pv-nvidia Date: Sat, 30 May 2026 10:09:51 +0000 Subject: [PATCH 20/54] feat: return ProxyArray from scale getters, add scale contract tests, shear warning - Change get_local_scales/get_world_scales return type from wp.array to ProxyArray across all backends (BaseFrameView, UsdFrameView, FabricFrameView, NewtonSiteFrameView, OvPhysxFrameView). - The deprecated get_scales() still returns wp.array (via .warp unwrap in _get_scales_default) to preserve backward compatibility. - Add scale contract tests to frame_view_contract_utils.py: - test_local_scales_default_identity: verify (1,1,1) default - test_world_scales_default_identity: verify (1,1,1) default - test_set_local_scales_roundtrip: write/read consistency - test_set_world_scales_roundtrip: write/read consistency - test_local_scales_do_not_affect_local_poses: scale changes preserve T/R - test_scale_getters_return_proxyarray: type contract check - Add shear/skew detection in FabricFrameView._sync_fabric_from_usd_initial: logs a one-time warning if any parent world transform has non-orthogonal rows (indicating shear), since TRS decomposition cannot represent shear. - Document the shear limitation in BaseFrameView.get_world_scales docstring. - Update existing Fabric scale tests to use .torch instead of torch.as_tensor (which would trigger ProxyArray deprecation bridge). - Simplify _sync_fabric_from_usd_initial scale extraction to directly unwrap ProxyArray.warp (no longer needs defensive hasattr check). --- .../isaaclab/sim/views/base_frame_view.py | 15 ++- .../isaaclab/sim/views/usd_frame_view.py | 15 ++- .../test/sim/frame_view_contract_utils.py | 104 ++++++++++++++++++ .../sim/views/newton_site_frame_view.py | 10 +- .../sim/views/ovphysx_frame_view.py | 8 +- .../sim/views/fabric_frame_view.py | 35 ++++-- .../test/sim/test_views_xform_prim_fabric.py | 10 +- 7 files changed, 161 insertions(+), 36 deletions(-) diff --git a/source/isaaclab/isaaclab/sim/views/base_frame_view.py b/source/isaaclab/isaaclab/sim/views/base_frame_view.py index 8095b3963692..201303b0a618 100644 --- a/source/isaaclab/isaaclab/sim/views/base_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/base_frame_view.py @@ -99,7 +99,7 @@ def set_local_poses( ... @abc.abstractmethod - def get_local_scales(self, indices: wp.array | None = None) -> wp.array: + def get_local_scales(self, indices: wp.array | None = None) -> ProxyArray: """Get local-space scales for prims in the view. Returns the per-prim scale as stored in local space (``xformOp:scale`` @@ -109,7 +109,7 @@ def get_local_scales(self, indices: wp.array | None = None) -> wp.array: indices: Subset of prims to query. ``None`` means all prims. Returns: - A :class:`~isaaclab.utils.warp.ProxyArray` of shape ``(M, 3)``. + A :class:`~isaaclab.utils.warp.ProxyArray` wrapping a ``wp.array`` of shape ``(M, 3)``. """ ... @@ -127,16 +127,23 @@ def set_local_scales(self, scales: wp.array, indices: wp.array | None = None) -> ... @abc.abstractmethod - def get_world_scales(self, indices: wp.array | None = None) -> wp.array: + 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. A warning is emitted at initialization time when sheared + parent transforms are detected. + Args: indices: Subset of prims to query. ``None`` means all prims. Returns: - A ``wp.array`` of shape ``(M, 3)``. + A :class:`~isaaclab.utils.warp.ProxyArray` wrapping a ``wp.array`` of shape ``(M, 3)``. """ ... diff --git a/source/isaaclab/isaaclab/sim/views/usd_frame_view.py b/source/isaaclab/isaaclab/sim/views/usd_frame_view.py index 30402e181d19..67dcc9a2300e 100644 --- a/source/isaaclab/isaaclab/sim/views/usd_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/usd_frame_view.py @@ -264,7 +264,7 @@ def set_world_scales(self, scales: wp.array, indices: wp.array | None = None): def _get_scales_default(self, indices: wp.array | None = None) -> wp.array: """USD default: get_scales returns local scales.""" - return self.get_local_scales(indices) + return self.get_local_scales(indices).warp def _set_scales_default(self, scales: wp.array, indices: wp.array | None = None) -> None: """USD default: set_scales writes local scales.""" @@ -350,14 +350,14 @@ 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_local_scales(self, indices: wp.array | None = None) -> wp.array: + def get_local_scales(self, indices: wp.array | None = None) -> ProxyArray: """Get local-space scales (xformOp:scale) for prims in the view. Args: indices: Indices of prims to get scales for. Defaults to None (all prims). Returns: - A :class:`~isaaclab.utils.warp.ProxyArray` of shape ``(M, 3)``. + A :class:`~isaaclab.utils.warp.ProxyArray` wrapping a ``wp.array`` of shape ``(M, 3)``. """ indices_list = self._resolve_indices(indices) @@ -366,10 +366,9 @@ def get_local_scales(self, indices: wp.array | None = None) -> wp.array: prim = self._prims[prim_idx] scales[idx] = prim.GetAttribute("xformOp:scale").Get() - scales_wp = wp.array(np.array(scales, dtype=np.float32), dtype=wp.float32, device=self._device) - return ProxyArray(scales_wp) + return ProxyArray(wp.array(np.array(scales, dtype=np.float32), dtype=wp.float32, device=self._device)) - def get_world_scales(self, indices: wp.array | None = None) -> wp.array: + def get_world_scales(self, indices: wp.array | None = None) -> ProxyArray: """Get world-space (composed) scales for prims in the view. Computes the effective world-space scale by extracting row lengths @@ -380,7 +379,7 @@ def get_world_scales(self, indices: wp.array | None = None) -> wp.array: indices: Indices of prims to get scales for. Defaults to None (all prims). Returns: - A ``wp.array`` of shape ``(M, 3)``. + A :class:`~isaaclab.utils.warp.ProxyArray` wrapping a ``wp.array`` of shape ``(M, 3)``. """ indices_list = self._resolve_indices(indices) xf_cache = UsdGeom.XformCache(Usd.TimeCode.Default()) @@ -393,7 +392,7 @@ def get_world_scales(self, indices: wp.array | None = None) -> wp.array: 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 wp.array(scales, dtype=wp.float32, device=self._device) + return ProxyArray(wp.array(scales, dtype=wp.float32, device=self._device)) def get_visibility(self, indices: wp.array | None = None) -> torch.Tensor: """Get visibility for prims in the view. diff --git a/source/isaaclab/test/sim/frame_view_contract_utils.py b/source/isaaclab/test/sim/frame_view_contract_utils.py index 8cd73c02b6f6..92e4eaafe31f 100644 --- a/source/isaaclab/test/sim/frame_view_contract_utils.py +++ b/source/isaaclab/test/sim/frame_view_contract_utils.py @@ -409,3 +409,107 @@ def test_return_types_are_torcharray(device, view_factory): ) 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_set_local_scales_roundtrip(device, view_factory): + """set_local_scales -> get_local_scales returns the same values.""" + 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) + bundle.view.set_local_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_set_world_scales_roundtrip(device, view_factory): + """set_world_scales -> get_world_scales returns the same values.""" + 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) + bundle.view.set_world_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) + bundle.view.set_local_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_newton/isaaclab_newton/sim/views/newton_site_frame_view.py b/source/isaaclab_newton/isaaclab_newton/sim/views/newton_site_frame_view.py index 587ff6ea5b9b..28a22743515b 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 @@ -884,7 +884,7 @@ def set_local_poses( # Scales # ------------------------------------------------------------------ - def get_local_scales(self, indices: wp.array | None = None) -> wp.array: + def get_local_scales(self, indices: wp.array | None = None) -> ProxyArray: """Get local-space scales. .. note:: @@ -892,11 +892,11 @@ def get_local_scales(self, indices: wp.array | None = None) -> wp.array: get_local_scales returns the same value since Newton does not decompose parent/child scale independently. """ - return self._get_shape_scales(indices) + return ProxyArray(self._get_shape_scales(indices)) - def get_world_scales(self, indices: wp.array | None = None) -> wp.array: + def get_world_scales(self, indices: wp.array | None = None) -> ProxyArray: """Get world-space (composed) scales from Newton shape_scale.""" - return self._get_shape_scales(indices) + return ProxyArray(self._get_shape_scales(indices)) def set_local_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: """Set local-space scales. @@ -913,7 +913,7 @@ def set_world_scales(self, scales: wp.array, indices: wp.array | None = None) -> def _get_scales_default(self, indices=None): """Newton default: get_scales returns shape_scale (world-like).""" - return self.get_world_scales(indices) + return self.get_world_scales(indices).warp def _set_scales_default(self, scales, indices=None): """Newton default: set_scales writes shape_scale (world-like).""" 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 b6c765c424e6..e51d417ec802 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 @@ -834,7 +834,7 @@ def _ensure_usd_view(self) -> UsdFrameView: ) return self._usd_view - def get_local_scales(self, indices: wp.array | None = None) -> wp.array: + def get_local_scales(self, indices: wp.array | None = None) -> ProxyArray: """Get local-space scales (xformOp:scale) via the USD view. .. note:: @@ -849,11 +849,11 @@ def get_local_scales(self, indices: wp.array | None = None) -> wp.array: indices: Subset of sites to query. ``None`` means all sites. Returns: - A :class:`~isaaclab.utils.warp.ProxyArray` of shape ``(M, 3)``. + A :class:`~isaaclab.utils.warp.ProxyArray` wrapping a ``wp.array`` of shape ``(M, 3)``. """ return self._ensure_usd_view().get_local_scales(indices) - def get_world_scales(self, indices: wp.array | None = None) -> wp.array: + def get_world_scales(self, indices: wp.array | None = None) -> ProxyArray: """Get world-space (composed) scales via the USD view.""" return self._ensure_usd_view().get_world_scales(indices) @@ -878,7 +878,7 @@ def set_world_scales(self, scales: wp.array, indices: wp.array | None = None) -> def _get_scales_default(self, indices=None): """OvPhysX default: get_scales returns local scales (same as USD).""" - return self.get_local_scales(indices) + return self.get_local_scales(indices).warp def _set_scales_default(self, scales, indices=None): """OvPhysX default: set_scales writes local scales (same as USD).""" 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 3ba306d5cea6..6a4a9f731730 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 @@ -467,6 +467,7 @@ def get_world_scales(self, indices=None): if use_cached: wp.synchronize() + return self._fabric_scales_ta return ProxyArray(scales_wp) def set_local_scales(self, scales, indices=None): @@ -546,11 +547,12 @@ def get_local_scales(self, indices=None): if use_cached: wp.synchronize() - return scales_wp + return self._fabric_scales_ta + return ProxyArray(scales_wp) def _get_scales_default(self, indices=None): """Fabric default: get_scales returns world scales (backwards compat).""" - return self.get_world_scales(indices) + return self.get_world_scales(indices).warp def _set_scales_default(self, scales, indices=None): """Fabric default: set_scales writes world scales (backwards compat).""" @@ -774,6 +776,7 @@ def _initialize_fabric(self) -> None: 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) @@ -799,14 +802,7 @@ def _sync_fabric_from_usd_initial(self) -> None: # by ``_recompute_world_from_local()`` at the end of this method as # ``child_world = child_local * parent_world``, which naturally # composes scales through the matrix multiplication. - scales_obj = self._usd_view.get_local_scales() - scales_wp = ( - scales_obj.warp - if hasattr(scales_obj, "warp") - else scales_obj - if isinstance(scales_obj, wp.array) - else self._fabric_empty_2d_array_sentinel - ) + 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, @@ -835,12 +831,31 @@ def _sync_fabric_from_usd_initial(self) -> None: 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) # Extract scale before ``Orthonormalize`` strips it from the rows. decomposer.SetMatrix(tf) s = decomposer.GetScale() + # Check for shear/skew: after removing scale, rows should be orthogonal. + 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() 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 d3f21caa0e05..72845ed3121b 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 @@ -274,7 +274,7 @@ def test_get_scales_fabric_path(device, view_factory): view.get_world_poses() scales = view.get_world_scales() - scales_t = torch.as_tensor(scales, device=device) + 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) @@ -297,7 +297,7 @@ def test_local_scales_roundtrip(device, view_factory): assert view._dirty.name == "WORLD" ret_scales = view.get_local_scales() - scales_torch = torch.as_tensor(ret_scales, device=device) + 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) @@ -319,7 +319,7 @@ def test_world_scales_roundtrip(device, view_factory): assert view._dirty.name == "LOCAL" ret_scales = view.get_world_scales() - scales_torch = torch.as_tensor(ret_scales, device=device) + 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) @@ -441,7 +441,7 @@ def test_initial_seed_with_scaled_parent(device): rtol=0, ) - scales = torch.as_tensor(view.get_world_scales(), device=device) + scales = view.get_world_scales().torch torch.testing.assert_close( scales, torch.tensor([[6.0, 1.0, 1.0]], dtype=torch.float32, device=device), @@ -628,7 +628,7 @@ def test_fabric_cuda1_scales_roundtrip(device, view_factory): view.set_world_scales(new_scales) ret_scales = view.get_world_scales() - scales_torch = torch.as_tensor(ret_scales, device=device) + 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}" From 9c060389e9d6b43719825ad33063248309b4397e Mon Sep 17 00:00:00 2001 From: pv-nvidia Date: Mon, 1 Jun 2026 13:43:13 +0000 Subject: [PATCH 21/54] cleanup: remove derived-class details from BaseFrameView docstrings Base class docstrings should describe the abstract contract without mentioning implementation specifics (Fabric dirty flags, USD xformOps, etc.). Move those details to the respective subclass docstrings where they belong. --- .../isaaclab/sim/views/base_frame_view.py | 20 +++++-------------- 1 file changed, 5 insertions(+), 15 deletions(-) diff --git a/source/isaaclab/isaaclab/sim/views/base_frame_view.py b/source/isaaclab/isaaclab/sim/views/base_frame_view.py index 201303b0a618..3e6a5154fae8 100644 --- a/source/isaaclab/isaaclab/sim/views/base_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/base_frame_view.py @@ -102,9 +102,6 @@ def set_local_poses( def get_local_scales(self, indices: wp.array | None = None) -> ProxyArray: """Get local-space scales for prims in the view. - Returns the per-prim scale as stored in local space (``xformOp:scale`` - for USD, decomposition of ``localMatrix`` for Fabric). - Args: indices: Subset of prims to query. ``None`` means all prims. @@ -117,9 +114,6 @@ def get_local_scales(self, indices: wp.array | None = None) -> ProxyArray: def set_local_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: """Set local-space scales for prims in the view. - Writes the per-prim scale in local space. For the Fabric backend this - marks the world matrix as dirty (will be re-propagated on next read). - Args: scales: Scales ``(M, 3)`` as ``wp.array``. indices: Subset of prims to update. ``None`` means all prims. @@ -136,8 +130,7 @@ def get_world_scales(self, indices: wp.array | None = None) -> ProxyArray: 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. A warning is emitted at initialization time when sheared - parent transforms are detected. + approximate. Args: indices: Subset of prims to query. ``None`` means all prims. @@ -151,9 +144,6 @@ def get_world_scales(self, indices: wp.array | None = None) -> ProxyArray: def set_world_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: """Set world-space (composed) scales for prims in the view. - Writes the effective scale in world space. For the Fabric backend this - marks the local matrix as dirty (will be re-propagated on next read). - Args: scales: Scales ``(M, 3)`` as ``wp.array``. indices: Subset of prims to update. ``None`` means all prims. @@ -169,8 +159,8 @@ def get_scales(self, indices: wp.array | None = None) -> wp.array: .. deprecated:: Use :meth:`get_local_scales` or :meth:`get_world_scales` instead. - This method calls ``get_local_scales`` for USD backends and - ``get_world_scales`` for Fabric backends (preserving legacy behavior). + This method delegates to :meth:`_get_scales_default` which preserves + each backend's legacy behavior. Args: indices: Subset of prims to query. ``None`` means all prims. @@ -190,8 +180,8 @@ def set_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: .. deprecated:: Use :meth:`set_local_scales` or :meth:`set_world_scales` instead. - This method calls ``set_local_scales`` for USD backends and - ``set_world_scales`` for Fabric backends (preserving legacy behavior). + This method delegates to :meth:`_set_scales_default` which preserves + each backend's legacy behavior. Args: scales: Scales ``(M, 3)`` as ``wp.array``. From 4c1e020e2fb898c51cecdc0f3617394753d890c8 Mon Sep 17 00:00:00 2001 From: pv-nvidia Date: Mon, 1 Jun 2026 13:48:27 +0000 Subject: [PATCH 22/54] cleanup: rename _get/set_scales_default to _impl, warn only once - Rename _get_scales_default/_set_scales_default to _get_scales_impl/_set_scales_impl across all backends. The 'default' name leaked implementation reasoning into the interface. - Emit the DeprecationWarning only once (class-level flag) instead of on every call. Avoids flooding logs when legacy code calls get_scales/ set_scales in a tight loop. --- .../isaaclab/sim/views/base_frame_view.py | 43 +++++++++++-------- .../isaaclab/sim/views/usd_frame_view.py | 4 +- .../sim/views/newton_site_frame_view.py | 4 +- .../sim/views/ovphysx_frame_view.py | 4 +- .../sim/views/fabric_frame_view.py | 4 +- 5 files changed, 33 insertions(+), 26 deletions(-) diff --git a/source/isaaclab/isaaclab/sim/views/base_frame_view.py b/source/isaaclab/isaaclab/sim/views/base_frame_view.py index 3e6a5154fae8..fb936bcd1f89 100644 --- a/source/isaaclab/isaaclab/sim/views/base_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/base_frame_view.py @@ -154,12 +154,15 @@ def set_world_scales(self, scales: wp.array, indices: wp.array | None = None) -> # Deprecated -- use get/set_local_scales or get/set_world_scales # ------------------------------------------------------------------ + _get_scales_deprecated_warned: bool = False + _set_scales_deprecated_warned: bool = False + def get_scales(self, indices: wp.array | None = None) -> wp.array: """Get scales for prims in the view. .. deprecated:: Use :meth:`get_local_scales` or :meth:`get_world_scales` instead. - This method delegates to :meth:`_get_scales_default` which preserves + This method delegates to :meth:`_get_scales_impl` which preserves each backend's legacy behavior. Args: @@ -168,38 +171,42 @@ def get_scales(self, indices: wp.array | None = None) -> wp.array: Returns: A ``wp.array`` of shape ``(M, 3)``. """ - warnings.warn( - "get_scales() is deprecated. Use get_local_scales() or get_world_scales() instead.", - DeprecationWarning, - stacklevel=2, - ) - return self._get_scales_default(indices) + if not BaseFrameView._get_scales_deprecated_warned: + BaseFrameView._get_scales_deprecated_warned = True + warnings.warn( + "get_scales() is deprecated. Use get_local_scales() or get_world_scales() instead.", + DeprecationWarning, + stacklevel=2, + ) + return self._get_scales_impl(indices) def set_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: """Set scales for prims in the view. .. deprecated:: Use :meth:`set_local_scales` or :meth:`set_world_scales` instead. - This method delegates to :meth:`_set_scales_default` which preserves + This method delegates to :meth:`_set_scales_impl` which preserves each backend's legacy behavior. Args: scales: Scales ``(M, 3)`` as ``wp.array``. indices: Subset of prims to update. ``None`` means all prims. """ - warnings.warn( - "set_scales() is deprecated. Use set_local_scales() or set_world_scales() instead.", - DeprecationWarning, - stacklevel=2, - ) - self._set_scales_default(scales, indices) + if not BaseFrameView._set_scales_deprecated_warned: + BaseFrameView._set_scales_deprecated_warned = True + warnings.warn( + "set_scales() is deprecated. Use set_local_scales() or set_world_scales() instead.", + DeprecationWarning, + stacklevel=2, + ) + self._set_scales_impl(scales, indices) @abc.abstractmethod - def _get_scales_default(self, indices: wp.array | None = None) -> wp.array: - """Backend-specific default for deprecated get_scales().""" + def _get_scales_impl(self, indices: wp.array | None = None) -> wp.array: + """Backend-specific implementation for deprecated get_scales().""" ... @abc.abstractmethod - def _set_scales_default(self, scales: wp.array, indices: wp.array | None = None) -> None: - """Backend-specific default for deprecated set_scales().""" + def _set_scales_impl(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Backend-specific implementation for deprecated 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 67dcc9a2300e..24c73ae036b0 100644 --- a/source/isaaclab/isaaclab/sim/views/usd_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/usd_frame_view.py @@ -262,11 +262,11 @@ def set_world_scales(self, scales: wp.array, indices: wp.array | None = None): ) prim.GetAttribute("xformOp:scale").Set(local_scale) - def _get_scales_default(self, indices: wp.array | None = None) -> wp.array: + def _get_scales_impl(self, indices: wp.array | None = None) -> wp.array: """USD default: get_scales returns local scales.""" return self.get_local_scales(indices).warp - def _set_scales_default(self, scales: wp.array, indices: wp.array | None = None) -> None: + def _set_scales_impl(self, scales: wp.array, indices: wp.array | None = None) -> None: """USD default: set_scales writes local scales.""" self.set_local_scales(scales, indices) 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 28a22743515b..1331d9d682f0 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 @@ -911,11 +911,11 @@ def set_world_scales(self, scales: wp.array, indices: wp.array | None = None) -> """Set world-space (composed) scales to Newton shape_scale.""" self._set_shape_scales(scales, indices) - def _get_scales_default(self, indices=None): + def _get_scales_impl(self, indices=None): """Newton default: get_scales returns shape_scale (world-like).""" return self.get_world_scales(indices).warp - def _set_scales_default(self, scales, indices=None): + def _set_scales_impl(self, scales, indices=None): """Newton default: set_scales writes shape_scale (world-like).""" self.set_world_scales(scales, indices) 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 e51d417ec802..ac0c591b40f6 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 @@ -876,11 +876,11 @@ def set_world_scales(self, scales: wp.array, indices: wp.array | None = None) -> """Set world-space scales via the USD view.""" self._ensure_usd_view().set_world_scales(scales, indices) - def _get_scales_default(self, indices=None): + def _get_scales_impl(self, indices=None): """OvPhysX default: get_scales returns local scales (same as USD).""" return self.get_local_scales(indices).warp - def _set_scales_default(self, scales, indices=None): + def _set_scales_impl(self, scales, indices=None): """OvPhysX default: set_scales writes local scales (same as USD).""" self.set_local_scales(scales, indices) 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 6a4a9f731730..56ebf673abba 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 @@ -550,11 +550,11 @@ def get_local_scales(self, indices=None): return self._fabric_scales_ta return ProxyArray(scales_wp) - def _get_scales_default(self, indices=None): + def _get_scales_impl(self, indices=None): """Fabric default: get_scales returns world scales (backwards compat).""" return self.get_world_scales(indices).warp - def _set_scales_default(self, scales, indices=None): + def _set_scales_impl(self, scales, indices=None): """Fabric default: set_scales writes world scales (backwards compat).""" self.set_world_scales(scales, indices) From 05e0ed67f04e5af4e14d4fb76b3ad97bc8b6851d Mon Sep 17 00:00:00 2001 From: pv-nvidia Date: Mon, 1 Jun 2026 14:49:06 +0000 Subject: [PATCH 23/54] revert: restore docs/source/setup/ecosystem.rst to upstream state The ThreeDWorld link removal was unrelated to this PR and caused a rebase conflict. Revert to upstream's version. --- docs/source/setup/ecosystem.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/source/setup/ecosystem.rst b/docs/source/setup/ecosystem.rst index 718cc18b5b49..4eac6d13cb53 100644 --- a/docs/source/setup/ecosystem.rst +++ b/docs/source/setup/ecosystem.rst @@ -1,6 +1,7 @@ .. _isaac-lab-ecosystem: Isaac Lab Ecosystem +=================== Isaac Lab is a modular, extensible framework for robot learning built on top of `Isaac Sim`_ and `Newton`_. It provides a unified interface for the most common workflows in robotics research — @@ -208,6 +209,7 @@ contributing, please reach out to us. .. _AirSim: https://microsoft.github.io/AirSim/ .. _DoorGym: https://github.com/PSVL/DoorGym/ .. _ManiSkill: https://github.com/haosulab/ManiSkill +.. _ThreeDWorld: https://github.com/threedworld-mit/tdw .. _RoboSuite: https://github.com/ARISE-Initiative/robosuite .. _MuJoCo: https://mujoco.org/ .. _MuJoCo Playground: https://playground.mujoco.org/ From 5d73a3e511576c759c398e65a60eebb6afcf35ac Mon Sep 17 00:00:00 2001 From: pv-nvidia Date: Mon, 1 Jun 2026 14:51:30 +0000 Subject: [PATCH 24/54] cleanup: update stale set_scales/get_scales references in FabricFrameView Replace all remaining references to the deprecated set_scales/get_scales with the new explicit set_world_scales/get_world_scales (and local variants) in comments, docstrings, and log messages. --- .../sim/views/fabric_frame_view.py | 25 ++++++++++--------- 1 file changed, 13 insertions(+), 12 deletions(-) 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 56ebf673abba..29302376f5ea 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 @@ -30,7 +30,7 @@ class _DirtyFlag(enum.Enum): NONE = 0 #: World matrices are stale (a prior ``set_local_poses`` wrote new locals). WORLD = 1 - #: Local matrices are stale (a prior ``set_world_poses``/``set_scales`` wrote new worlds). + #: Local matrices are stale (a prior ``set_world_poses``/``set_world_scales`` wrote new worlds). LOCAL = 2 @@ -75,14 +75,15 @@ class FabricFrameView(BaseFrameView): read the prim's USD attributes after a Fabric write will see stale values until the next USD-side sync. * **World ↔ local consistency (lazy).** Getters are lazy: after - ``set_world_poses`` (or ``set_scales``), local matrices are only - recomputed when ``get_local_poses`` is called; after ``set_local_poses``, - world matrices are only recomputed when ``get_world_poses`` is called. - Both directions stay in sync without round-tripping through USD. + ``set_world_poses`` or ``set_world_scales``, local matrices are only + recomputed when ``get_local_poses`` (or ``get_local_scales``) is called; + after ``set_local_poses`` or ``set_local_scales``, world matrices are + only recomputed when ``get_world_poses`` (or ``get_world_scales``) is + called. Both directions stay in sync without round-tripping through USD. * **Dirty-flag invariant.** The ``_dirty`` enum is one of ``NONE``, ``WORLD``, or ``LOCAL`` -- mutually exclusive by construction. - ``set_world_poses`` / ``set_scales`` sets ``_dirty = LOCAL``; - ``set_local_poses`` sets ``_dirty = WORLD``. + ``set_world_poses`` / ``set_world_scales`` sets ``_dirty = LOCAL``; + ``set_local_poses`` / ``set_local_scales`` sets ``_dirty = WORLD``. If the user interleaves both setters on the same view within a single frame, the second setter flushes the first's stale data before writing. This is correct but incurs an extra kernel launch -- a one-time warning @@ -139,7 +140,7 @@ def __init__( # 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) + # TODO(pv): Fuse set_world_poses/set_world_scales into single kernel launch (PR #5674 pv/fabric-fused-compose) self._fabric_initialized = False self._stage = None @@ -309,7 +310,7 @@ def set_local_poses(self, translations=None, orientations=None, indices=None): self._warned_interleaved_set = True logger.warning( "FabricFrameView: set_local_poses called while local matrices are stale from a " - "prior set_world_poses/set_scales. Flushing stale locals first. " + "prior set_world_poses/set_world_scales. Flushing stale locals first. " "For best performance, avoid interleaving set_world_poses and set_local_poses " "on the same view within a single frame -- use one or the other exclusively." ) @@ -354,7 +355,7 @@ def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, if not self._fabric_initialized: self._initialize_fabric() - # If a prior set_world_poses/set_scales left localMatrix stale, recompute. + # If a prior set_world_poses/set_world_scales left localMatrix stale, recompute. self._sync_local_from_world_if_dirty() indices_wp = self._resolve_indices_wp(indices) @@ -551,11 +552,11 @@ def get_local_scales(self, indices=None): return ProxyArray(scales_wp) def _get_scales_impl(self, indices=None): - """Fabric default: get_scales returns world scales (backwards compat).""" + """Fabric: deprecated get_scales delegates to get_world_scales.""" return self.get_world_scales(indices).warp def _set_scales_impl(self, scales, indices=None): - """Fabric default: set_scales writes world scales (backwards compat).""" + """Fabric: deprecated set_scales delegates to set_world_scales.""" self.set_world_scales(scales, indices) # ------------------------------------------------------------------ From 3f1820f8323200951f5883ad6d82cdbcc41f046f Mon Sep 17 00:00:00 2001 From: pv-nvidia Date: Tue, 2 Jun 2026 15:14:45 +0000 Subject: [PATCH 25/54] fix: reset newton_site_frame_view.py to upstream + minimal scale additions The previous rebase accidentally included a complete rewrite of the Newton file (from a different branch). Reset to upstream/develop and add only the 6 scale methods required by the new BaseFrameView ABC. --- .../sim/views/newton_site_frame_view.py | 1132 ++++++----------- 1 file changed, 364 insertions(+), 768 deletions(-) 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 1331d9d682f0..33d34a89dd8e 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 @@ -3,7 +3,7 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Newton-backed FrameView — Warp-native, GPU-resident pose queries.""" +"""Newton-backed FrameView using Newton body labels and injected sites.""" from __future__ import annotations @@ -11,11 +11,13 @@ import warp as wp -from pxr import Gf, Usd, UsdGeom +from pxr import UsdPhysics import isaaclab.sim as sim_utils +from isaaclab.cloner.cloner_utils import iter_clone_plan_matches from isaaclab.physics import PhysicsEvent from isaaclab.sim.views.base_frame_view import BaseFrameView +from isaaclab.utils.string import resolve_matching_names from isaaclab.utils.warp import ProxyArray from isaaclab_newton.physics.newton_manager import NewtonManager @@ -25,47 +27,8 @@ WORLD_BODY_INDEX = -1 -# ------------------------------------------------------------------ -# Warp kernels -# ------------------------------------------------------------------ - - @wp.kernel def _compute_site_world_transforms( - body_q: wp.array(dtype=wp.transformf), - site_body: wp.array(dtype=wp.int32), - site_local: wp.array(dtype=wp.transformf), - out_pos: wp.array(dtype=wp.vec3f), - out_quat: wp.array(dtype=wp.vec4f), -): - """Compute world-space transforms for every site in the view. - - For each site *i*, computes ``world = body_q[site_body[i]] * site_local[i]`` - and splits the result into position and quaternion outputs. When - ``site_body[i] == -1`` the site is world-attached and ``site_local[i]`` is - returned directly. - - Args: - body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. - site_body: Per-site body index (flat model-level), shape ``[num_sites]``. - A value of ``-1`` indicates a world-attached site. - site_local: Per-site local offset relative to its parent body, shape ``[num_sites]``. - out_pos: Output world positions [m], shape ``[num_sites]``. - out_quat: Output world orientations as ``(qx, qy, qz, qw)``, shape ``[num_sites]``. - """ - i = wp.tid() - bid = site_body[i] - if bid == -1: - world = site_local[i] - else: - world = wp.transform_multiply(body_q[bid], site_local[i]) - out_pos[i] = wp.transform_get_translation(world) - q = wp.transform_get_rotation(world) - out_quat[i] = wp.vec4f(q[0], q[1], q[2], q[3]) - - -@wp.kernel -def _compute_site_world_transforms_indexed( body_q: wp.array(dtype=wp.transformf), site_body: wp.array(dtype=wp.int32), site_local: wp.array(dtype=wp.transformf), @@ -73,24 +36,11 @@ def _compute_site_world_transforms_indexed( out_pos: wp.array(dtype=wp.vec3f), out_quat: wp.array(dtype=wp.vec4f), ): - """Indexed variant of :func:`_compute_site_world_transforms`. - - Only computes world transforms for the subset of sites selected by - ``indices``. Thread *i* reads ``indices[i]`` to obtain the site index, - then writes the result to ``out_pos[i]`` / ``out_quat[i]``. - - Args: - body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. - site_body: Per-site body index (flat model-level), shape ``[num_sites]``. - site_local: Per-site local offset relative to its parent body, shape ``[num_sites]``. - indices: Site indices to query, shape ``[M]``. - out_pos: Output world positions [m], shape ``[M]``. - out_quat: Output world orientations as ``(qx, qy, qz, qw)``, shape ``[M]``. - """ + """Compute world-space transforms for selected sites.""" i = wp.tid() si = indices[i] bid = site_body[si] - if bid == -1: + if bid == WORLD_BODY_INDEX: world = site_local[si] else: world = wp.transform_multiply(body_q[bid], site_local[si]) @@ -100,169 +50,23 @@ def _compute_site_world_transforms_indexed( @wp.kernel -def _gather_scales( - shape_scale: wp.array(dtype=wp.vec3f), - shape_body: wp.array(dtype=wp.int32), - site_body: wp.array(dtype=wp.int32), - num_shapes: wp.int32, - out_scales: wp.array(dtype=wp.vec3f), -): - """Gather per-site scales from collision shapes on the same body. - - For each site *i*, linearly scans all shapes to find the first one whose - ``shape_body`` matches ``site_body[i]`` and copies its scale. Falls back - to ``(1, 1, 1)`` if no shape is found on that body. - - Args: - shape_scale: Per-shape scale vectors from the Newton model, shape ``[num_shapes]``. - shape_body: Per-shape parent body index, shape ``[num_shapes]``. - site_body: Per-site body index, shape ``[num_sites]``. - num_shapes: Total number of shapes in the model. - out_scales: Output scale per site, shape ``[num_sites]``. - """ - i = wp.tid() - bid = site_body[i] - found = int(0) - for s in range(num_shapes): - if shape_body[s] == bid and found == 0: - out_scales[i] = shape_scale[s] - found = 1 - if found == 0: - out_scales[i] = wp.vec3f(1.0, 1.0, 1.0) - - -@wp.kernel -def _gather_scales_indexed( - shape_scale: wp.array(dtype=wp.vec3f), - shape_body: wp.array(dtype=wp.int32), - site_body: wp.array(dtype=wp.int32), - indices: wp.array(dtype=wp.int32), - num_shapes: wp.int32, - out_scales: wp.array(dtype=wp.vec3f), -): - """Indexed variant of :func:`_gather_scales`. - - Args: - shape_scale: Per-shape scale vectors from the Newton model, shape ``[num_shapes]``. - shape_body: Per-shape parent body index, shape ``[num_shapes]``. - site_body: Per-site body index, shape ``[num_sites]``. - indices: Site indices to query, shape ``[M]``. - num_shapes: Total number of shapes in the model. - out_scales: Output scale per queried site, shape ``[M]``. - """ - i = wp.tid() - si = indices[i] - bid = site_body[si] - found = int(0) - for s in range(num_shapes): - if shape_body[s] == bid and found == 0: - out_scales[i] = shape_scale[s] - found = 1 - if found == 0: - out_scales[i] = wp.vec3f(1.0, 1.0, 1.0) - - -@wp.kernel -def _scatter_scales( - site_body: wp.array(dtype=wp.int32), - new_scales: wp.array(dtype=wp.vec3f), - shape_body: wp.array(dtype=wp.int32), - num_shapes: wp.int32, - shape_scale: wp.array(dtype=wp.vec3f), -): - """Scatter per-site scales to all collision shapes on the same body. - - For each site *i*, writes ``new_scales[i]`` to every shape whose - ``shape_body`` matches ``site_body[i]``. Multiple shapes on the same - body all receive the same scale. - - Args: - site_body: Per-site body index, shape ``[num_sites]``. - new_scales: New scale to apply per site, shape ``[num_sites]``. - shape_body: Per-shape parent body index, shape ``[num_shapes]``. - num_shapes: Total number of shapes in the model. - shape_scale: Per-shape scale vectors to write into (modified in-place), - shape ``[num_shapes]``. - """ - i = wp.tid() - bid = site_body[i] - for s in range(num_shapes): - if shape_body[s] == bid: - shape_scale[s] = new_scales[i] - - -@wp.kernel -def _scatter_scales_indexed( - site_body: wp.array(dtype=wp.int32), +def _gather_site_local_transforms( + site_local: wp.array(dtype=wp.transformf), indices: wp.array(dtype=wp.int32), - new_scales: wp.array(dtype=wp.vec3f), - shape_body: wp.array(dtype=wp.int32), - num_shapes: wp.int32, - shape_scale: wp.array(dtype=wp.vec3f), + out_pos: wp.array(dtype=wp.vec3f), + out_quat: wp.array(dtype=wp.vec4f), ): - """Indexed variant of :func:`_scatter_scales`. - - Args: - site_body: Per-site body index, shape ``[num_sites]``. - indices: Site indices to update, shape ``[M]``. - new_scales: New scale to apply per selected site, shape ``[M]``. - shape_body: Per-shape parent body index, shape ``[num_shapes]``. - num_shapes: Total number of shapes in the model. - shape_scale: Per-shape scale vectors to write into (modified in-place), - shape ``[num_shapes]``. - """ + """Gather local transforms for selected sites.""" i = wp.tid() si = indices[i] - bid = site_body[si] - for s in range(num_shapes): - if shape_body[s] == bid: - shape_scale[s] = new_scales[i] - - -# ------------------------------------------------------------------ -# World-pose site_local write kernels -# ------------------------------------------------------------------ + local_tf = site_local[si] + out_pos[i] = wp.transform_get_translation(local_tf) + q = wp.transform_get_rotation(local_tf) + out_quat[i] = wp.vec4f(q[0], q[1], q[2], q[3]) @wp.kernel def _write_site_local_from_world_poses( - body_q: wp.array(dtype=wp.transformf), - site_body: wp.array(dtype=wp.int32), - world_pos: wp.array(dtype=wp.vec3f), - world_quat: wp.array(dtype=wp.vec4f), - site_local: wp.array(dtype=wp.transformf), -): - """Update site local offsets so that the sites reach desired world poses. - - For each site *i*, computes - ``site_local[i] = inv(body_q[site_body[i]]) * desired_world`` so that - a subsequent ``body_q[bid] * site_local[i]`` yields the requested world - pose. For world-attached sites (``site_body[i] == -1``) the desired world - transform is written directly into ``site_local[i]``. - - Does **not** modify ``body_q``. - - Args: - body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. - site_body: Per-site body index (flat model-level), shape ``[num_sites]``. - world_pos: Desired world positions [m], shape ``[num_sites]``. - world_quat: Desired world orientations as ``(qx, qy, qz, qw)``, shape ``[num_sites]``. - site_local: Per-site local offset (modified in-place), shape ``[num_sites]``. - """ - i = wp.tid() - w_pos = world_pos[i] - w_q = world_quat[i] - desired_world = wp.transform(w_pos, wp.quatf(w_q[0], w_q[1], w_q[2], w_q[3])) - - bid = site_body[i] - if bid == -1: - site_local[i] = desired_world - else: - site_local[i] = wp.transform_multiply(wp.transform_inverse(body_q[bid]), desired_world) - - -@wp.kernel -def _write_site_local_from_world_poses_indexed( body_q: wp.array(dtype=wp.transformf), site_body: wp.array(dtype=wp.int32), indices: wp.array(dtype=wp.int32), @@ -270,16 +74,7 @@ def _write_site_local_from_world_poses_indexed( world_quat: wp.array(dtype=wp.vec4f), site_local: wp.array(dtype=wp.transformf), ): - """Indexed variant of :func:`_write_site_local_from_world_poses`. - - Args: - body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. - site_body: Per-site body index (flat model-level), shape ``[num_sites]``. - indices: Site indices to update, shape ``[M]``. - world_pos: Desired world positions [m], shape ``[M]``. - world_quat: Desired world orientations as ``(qx, qy, qz, qw)``, shape ``[M]``. - site_local: Per-site local offset (modified in-place), shape ``[num_sites]``. - """ + """Update local offsets so selected sites reach desired world poses.""" i = wp.tid() si = indices[i] w_pos = world_pos[i] @@ -287,435 +82,358 @@ def _write_site_local_from_world_poses_indexed( desired_world = wp.transform(w_pos, wp.quatf(w_q[0], w_q[1], w_q[2], w_q[3])) bid = site_body[si] - if bid == -1: + if bid == WORLD_BODY_INDEX: site_local[si] = desired_world else: site_local[si] = wp.transform_multiply(wp.transform_inverse(body_q[bid]), desired_world) -# ------------------------------------------------------------------ -# Local-pose Warp kernels -# ------------------------------------------------------------------ - - @wp.kernel -def _compute_site_local_transforms( - body_q: wp.array(dtype=wp.transformf), - site_body: wp.array(dtype=wp.int32), +def _write_site_local_from_local_poses( + indices: wp.array(dtype=wp.int32), + local_pos: wp.array(dtype=wp.vec3f), + local_quat: wp.array(dtype=wp.vec4f), site_local: wp.array(dtype=wp.transformf), - parent_site_body: wp.array(dtype=wp.int32), - parent_site_local: wp.array(dtype=wp.transformf), - out_pos: wp.array(dtype=wp.vec3f), - out_quat: wp.array(dtype=wp.vec4f), ): - """Compute parent-relative transforms for every site in the view. - - For each site *i*, computes the world pose of both the site and its USD - parent, then returns ``inv(parent_world) * prim_world``. When - ``site_body[i] == -1`` the site is world-attached and ``site_local[i]`` - is used as the world transform directly. The same convention applies to - the parent arrays. - - Args: - body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. - site_body: Per-site body index (flat model-level), shape ``[num_sites]``. - site_local: Per-site local offset relative to its parent body, shape ``[num_sites]``. - parent_site_body: Per-site USD-parent body index, shape ``[num_sites]``. - parent_site_local: Per-site USD-parent local offset, shape ``[num_sites]``. - out_pos: Output parent-relative positions [m], shape ``[num_sites]``. - out_quat: Output parent-relative orientations as ``(qx, qy, qz, qw)``, - shape ``[num_sites]``. - """ + """Update local offsets for selected sites.""" i = wp.tid() - prim_bid = site_body[i] - if prim_bid == -1: - prim_world = site_local[i] - else: - prim_world = wp.transform_multiply(body_q[prim_bid], site_local[i]) - - parent_bid = parent_site_body[i] - if parent_bid == -1: - parent_world = parent_site_local[i] - else: - parent_world = wp.transform_multiply(body_q[parent_bid], parent_site_local[i]) - - local_tf = wp.transform_multiply(wp.transform_inverse(parent_world), prim_world) - out_pos[i] = wp.transform_get_translation(local_tf) - q = wp.transform_get_rotation(local_tf) - out_quat[i] = wp.vec4f(q[0], q[1], q[2], q[3]) + si = indices[i] + l_pos = local_pos[i] + l_q = local_quat[i] + site_local[si] = wp.transform(l_pos, wp.quatf(l_q[0], l_q[1], l_q[2], l_q[3])) @wp.kernel -def _compute_site_local_transforms_indexed( - body_q: wp.array(dtype=wp.transformf), +def _gather_scales( + shape_scale: wp.array(dtype=wp.vec3f), + shape_body: wp.array(dtype=wp.int32), site_body: wp.array(dtype=wp.int32), - site_local: wp.array(dtype=wp.transformf), - parent_site_body: wp.array(dtype=wp.int32), - parent_site_local: wp.array(dtype=wp.transformf), indices: wp.array(dtype=wp.int32), - out_pos: wp.array(dtype=wp.vec3f), - out_quat: wp.array(dtype=wp.vec4f), + num_shapes: wp.int32, + out_scales: wp.array(dtype=wp.vec3f), ): - """Indexed variant of :func:`_compute_site_local_transforms`. - - Args: - body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. - site_body: Per-site body index (flat model-level), shape ``[num_sites]``. - site_local: Per-site local offset relative to its parent body, shape ``[num_sites]``. - parent_site_body: Per-site USD-parent body index, shape ``[num_sites]``. - parent_site_local: Per-site USD-parent local offset, shape ``[num_sites]``. - indices: Site indices to query, shape ``[M]``. - out_pos: Output parent-relative positions [m], shape ``[M]``. - out_quat: Output parent-relative orientations as ``(qx, qy, qz, qw)``, - shape ``[M]``. - """ + """Gather per-site scales from collision shapes on the same body.""" i = wp.tid() si = indices[i] - prim_bid = site_body[si] - if prim_bid == -1: - prim_world = site_local[si] - else: - prim_world = wp.transform_multiply(body_q[prim_bid], site_local[si]) - - parent_bid = parent_site_body[si] - if parent_bid == -1: - parent_world = parent_site_local[si] - else: - parent_world = wp.transform_multiply(body_q[parent_bid], parent_site_local[si]) - - local_tf = wp.transform_multiply(wp.transform_inverse(parent_world), prim_world) - out_pos[i] = wp.transform_get_translation(local_tf) - q = wp.transform_get_rotation(local_tf) - out_quat[i] = wp.vec4f(q[0], q[1], q[2], q[3]) - - -@wp.kernel -def _write_site_local_from_local_poses( - body_q: wp.array(dtype=wp.transformf), - site_body: wp.array(dtype=wp.int32), - parent_site_body: wp.array(dtype=wp.int32), - parent_site_local: wp.array(dtype=wp.transformf), - local_pos: wp.array(dtype=wp.vec3f), - local_quat: wp.array(dtype=wp.vec4f), - site_local: wp.array(dtype=wp.transformf), -): - """Update site local offsets so that sites reach desired parent-relative poses. - - For each site *i*, reconstructs the desired world pose as - ``parent_world * desired_local``, then solves for the body-relative offset: - ``site_local[i] = inv(body_q[bid]) * desired_world``. For world-attached - sites (``site_body[i] == -1``) the world transform is written directly. - - Does **not** modify ``body_q``. - - Args: - body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. - site_body: Per-site body index (flat model-level), shape ``[num_sites]``. - parent_site_body: Per-site USD-parent body index, shape ``[num_sites]``. - parent_site_local: Per-site USD-parent local offset, shape ``[num_sites]``. - local_pos: Desired parent-relative positions [m], shape ``[num_sites]``. - local_quat: Desired parent-relative orientations as ``(qx, qy, qz, qw)``, - shape ``[num_sites]``. - site_local: Per-site local offset (modified in-place), shape ``[num_sites]``. - """ - i = wp.tid() - parent_bid = parent_site_body[i] - if parent_bid == -1: - parent_world = parent_site_local[i] - else: - parent_world = wp.transform_multiply(body_q[parent_bid], parent_site_local[i]) - - l_pos = local_pos[i] - l_q = local_quat[i] - local_tf = wp.transform(l_pos, wp.quatf(l_q[0], l_q[1], l_q[2], l_q[3])) - desired_world = wp.transform_multiply(parent_world, local_tf) - - bid = site_body[i] - if bid == -1: - site_local[i] = desired_world - else: - site_local[i] = wp.transform_multiply(wp.transform_inverse(body_q[bid]), desired_world) + bid = site_body[si] + found = int(0) + for s in range(num_shapes): + if shape_body[s] == bid and found == 0: + out_scales[i] = shape_scale[s] + found = 1 + if found == 0: + out_scales[i] = wp.vec3f(1.0, 1.0, 1.0) @wp.kernel -def _write_site_local_from_local_poses_indexed( - body_q: wp.array(dtype=wp.transformf), +def _scatter_scales( site_body: wp.array(dtype=wp.int32), - parent_site_body: wp.array(dtype=wp.int32), - parent_site_local: wp.array(dtype=wp.transformf), indices: wp.array(dtype=wp.int32), - local_pos: wp.array(dtype=wp.vec3f), - local_quat: wp.array(dtype=wp.vec4f), - site_local: wp.array(dtype=wp.transformf), + new_scales: wp.array(dtype=wp.vec3f), + shape_body: wp.array(dtype=wp.int32), + num_shapes: wp.int32, + shape_scale: wp.array(dtype=wp.vec3f), ): - """Indexed variant of :func:`_write_site_local_from_local_poses`. - - Args: - body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. - site_body: Per-site body index (flat model-level), shape ``[num_sites]``. - parent_site_body: Per-site USD-parent body index, shape ``[num_sites]``. - parent_site_local: Per-site USD-parent local offset, shape ``[num_sites]``. - indices: Site indices to update, shape ``[M]``. - local_pos: Desired parent-relative positions [m], shape ``[M]``. - local_quat: Desired parent-relative orientations as ``(qx, qy, qz, qw)``, - shape ``[M]``. - site_local: Per-site local offset (modified in-place), shape ``[num_sites]``. - """ + """Scatter per-site scales to collision shapes on the same body.""" i = wp.tid() si = indices[i] - parent_bid = parent_site_body[si] - if parent_bid == -1: - parent_world = parent_site_local[si] - else: - parent_world = wp.transform_multiply(body_q[parent_bid], parent_site_local[si]) - - l_pos = local_pos[i] - l_q = local_quat[i] - local_tf = wp.transform(l_pos, wp.quatf(l_q[0], l_q[1], l_q[2], l_q[3])) - desired_world = wp.transform_multiply(parent_world, local_tf) - bid = site_body[si] - if bid == -1: - site_local[si] = desired_world - else: - site_local[si] = wp.transform_multiply(wp.transform_inverse(body_q[bid]), desired_world) - - -# ------------------------------------------------------------------ -# View class -# ------------------------------------------------------------------ + for s in range(num_shapes): + if shape_body[s] == bid: + shape_scale[s] = new_scales[i] class NewtonSiteFrameView(BaseFrameView): - """Batched prim view for non-physics prims tracked as sites on Newton bodies. - - Each matched USD prim must be a **non-physics** prim (camera, sensor, - Xform marker, etc.) that sits as a child of a Newton rigid body in the - USD hierarchy. The prim path must **not** resolve directly to a physics - body or collision shape -- those are owned by Newton and should be - accessed through :class:`~isaaclab_newton.assets.Articulation` or - :class:`~isaaclab_newton.assets.RigidObject` instead. - - At init time each prim is resolved to a ``(body_index, site_local)`` - pair via ancestor walk: the nearest ancestor that appears in - ``model.body_label`` becomes the attachment body, and the relative USD - transform becomes the site offset. If no body ancestor exists the prim - is attached to the world frame (``body_index = -1``). - - World poses are computed on GPU as - ``body_q[body_index] * site_local`` via a Warp kernel. Both - ``set_world_poses`` and ``set_local_poses`` update ``site_local`` -- - neither touches ``body_q``. - - Pose getters return :class:`~isaaclab.utils.warp.ProxyArray`. Setters accept ``wp.array``. - - Raises: - ValueError: If any matched prim resolves to a Newton physics body - or collision shape. - """ + """Batched Newton site view for non-physics frames. - def __init__(self, prim_path: str, device: str = "cpu", stage: Usd.Stage | None = None, **kwargs): - """Initialize the Newton site-based frame view. - - Resolves all USD prims matching ``prim_path`` and, for each one, walks - the USD ancestor hierarchy to find the nearest Newton rigid body. The - relative transform between the prim and its ancestor body becomes the - site's local offset. + The public construction contract matches the generic :class:`FrameView`: + callers provide a prim expression and the backend resolves the source prim + into Newton body-local or world-local sites. + """ - If the Newton model is already finalized the view initializes - immediately; otherwise initialization is deferred to a - :attr:`PhysicsEvent.PHYSICS_READY` callback. + def __init__( + self, + prim_path: str | list[str], + device: str = "cpu", + validate_xform_ops: bool = True, + stage: object | None = None, + **kwargs, + ): + """Initialize the Newton site frame view. Args: - prim_path: USD prim path pattern (may contain regex). - device: Warp device for GPU arrays (e.g. ``"cuda:0"``). - stage: USD stage to search. Defaults to the current stage. - **kwargs: Unused; accepted for interface compatibility with other - :class:`~isaaclab.sim.views.BaseFrameView` backends. + prim_path: User-facing frame path pattern, or list of patterns. + device: Warp device for GPU arrays. + validate_xform_ops: Whether to validate source USD xform ops. + stage: USD stage that contains the source prims. + **kwargs: Unused. """ - self._prim_path = prim_path + del kwargs + + self._prim_paths = [prim_path] if isinstance(prim_path, str) else list(prim_path) + self._prim_path = prim_path if isinstance(prim_path, str) else ", ".join(self._prim_paths) self._device = device + self._prims = [] stage = sim_utils.get_current_stage() if stage is None else stage - self._prims: list[Usd.Prim] = sim_utils.find_matching_prims(prim_path, stage=stage) + self._site_specs = self._resolve_site_specs(stage, validate_xform_ops) + self._site_labels: list[str] = [] + self._site_body: wp.array | None = None + self._site_local: 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._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._count = 0 model = NewtonManager.get_model() if model is not None: - self._initialize_impl(model) + self._initialize_from_specs(model) else: + for body_patterns, xform, 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)) + else: + for body_pattern in body_patterns: + self._site_labels.append(NewtonManager.cl_register_site(body_pattern, xform)) self._physics_ready_handle = NewtonManager.register_callback( - self._on_physics_ready, PhysicsEvent.PHYSICS_READY, name=f"site_view_{prim_path}" + self._on_physics_ready, PhysicsEvent.PHYSICS_READY, name=f"site_view_{self._prim_path}" ) - def _on_physics_ready(self, _event) -> None: - """Callback invoked when the Newton model becomes available.""" - self._initialize_impl(NewtonManager.get_model()) + def _resolve_site_specs( + self, stage, validate_xform_ops: bool + ) -> list[tuple[tuple[str, ...] | None, wp.transform, 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]] = [] - def _initialize_impl(self, model) -> None: - """Resolve USD prims to Newton body indices and allocate GPU buffers.""" - body_labels = list(model.body_label) - body_label_set = set(body_labels) - body_label_to_idx = {path: idx for idx, path in enumerate(body_labels)} - shape_label_set = set(model.shape_label) + for path_expr in self._prim_paths: + if resolve_matching_names(path_expr, body_labels, raise_when_no_match=False)[1]: + raise ValueError( + f"FrameView prim '{path_expr}' is a Newton physics body. " + "FrameView should only be used for non-physics frames." + ) + if resolve_matching_names(path_expr, shape_labels, raise_when_no_match=False)[1]: + raise ValueError( + f"FrameView prim '{path_expr}' is a Newton collision shape. " + "FrameView should only be used for non-physics frames." + ) + matches = tuple(iter_clone_plan_matches(plan, path_expr)) if plan is not None else () + if matches: + for source_root, destination_template, source_path, env_ids in matches: + source_prim = None + if not any(token in source_path for token in "*[]()+?|\\"): + source_prim = stage.GetPrimAtPath(source_path) + if source_prim is None or not source_prim.IsValid(): + source_prim = sim_utils.find_first_matching_prim(source_path, stage) + if source_prim is None or not source_prim.IsValid(): + raise RuntimeError(f"FrameView '{path_expr}' could not resolve source prim '{source_path}'.") + specs.append( + self._resolve_source_prim( + source_prim, + validate_xform_ops, + source_root, + destination_template, + env_ids, + use_clone_body_pattern, + stage, + ) + ) + continue + + prim = sim_utils.find_first_matching_prim(path_expr, stage) + if prim is None or not prim.IsValid(): + raise RuntimeError(f"FrameView '{path_expr}' could not resolve a source prim.") + specs.append( + self._resolve_source_prim(prim, validate_xform_ops, None, None, None, use_clone_body_pattern, stage) + ) + + return specs - xform_cache = UsdGeom.XformCache(Usd.TimeCode.Default()) + def _resolve_source_prim( + self, + prim, + validate_xform_ops: bool, + source_root: str | None, + destination_template: str | None, + 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.""" + prim_path = prim.GetPath().pathString + if prim.HasAPI(UsdPhysics.RigidBodyAPI) or prim.HasAPI(UsdPhysics.ArticulationRootAPI): + raise ValueError( + f"FrameView prim '{prim_path}' is a Newton physics body. " + "FrameView should only be used for non-physics frames." + ) + if validate_xform_ops: + sim_utils.standardize_xform_ops(prim) + if not sim_utils.validate_standard_xform_ops(prim): + raise ValueError(f"FrameView prim '{prim_path}' does not have standard xform ops.") + + body_prim = prim.GetParent() + while body_prim and body_prim.IsValid(): + if body_prim.HasAPI(UsdPhysics.RigidBodyAPI) or body_prim.HasAPI(UsdPhysics.ArticulationRootAPI): + pos, quat = sim_utils.resolve_prim_pose(prim, body_prim) + body_path = body_prim.GetPath().pathString + if source_root is not None and destination_template is not None: + assert env_ids is not None + if body_path == source_root: + suffix = "" + elif body_path.startswith(source_root + "/"): + suffix = body_path[len(source_root) :] + elif source_root.startswith(body_path + "/"): + suffix = source_root[len(body_path) :] + if use_clone_body_pattern: + destination_root = destination_template.format(".*") + if not destination_root.endswith(suffix): + 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 + body_patterns = [] + for env_id in env_ids: + destination_root = destination_template.format(env_id) + if not destination_root.endswith(suffix): + raise RuntimeError( + 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 + else: + raise RuntimeError(f"FrameView source body '{body_path}' is not under '{source_root}'.") + if use_clone_body_pattern: + body_patterns = (destination_template.format(".*") + suffix,) + else: + 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 + body_prim = body_prim.GetParent() + ref_prim = stage.GetPrimAtPath(source_root) if source_root 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 + + def _on_physics_ready(self, _event) -> None: + """Callback invoked when the Newton model becomes available.""" + self._initialize_from_site_map(NewtonManager.get_model()) + + def _initialize_from_site_map(self, model) -> None: + """Initialize arrays from injected Newton sites.""" + site_map = NewtonManager._cl_site_index_map + body_t = wp.to_torch(model.shape_body) + xform_t = wp.to_torch(model.shape_transform) site_bodies: list[int] = [] site_locals: list[list[float]] = [] - parent_bodies: list[int] = [] - parent_locals: list[list[float]] = [] - identity_xform = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0] - resolve_cache: dict[str, tuple[int, list[float]]] = {} + for site_label in self._site_labels: + 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] + ) + 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()]) - for prim in self._prims: - pp = prim.GetPath().pathString - if pp in body_label_set: - raise ValueError( - f"FrameView prim '{pp}' is a Newton physics body. " - "FrameView should only be used for non-physics prims (cameras, sensors, Xform markers). " - "Use Articulation or RigidObject APIs to control physics bodies." - ) - if pp in shape_label_set: - raise ValueError( - f"FrameView prim '{pp}' is a Newton collision shape. " - "FrameView should only be used for non-physics prims (cameras, sensors, Xform markers). " - "Use Articulation or RigidObject APIs to control collision shapes." - ) + self._create_buffers(site_bodies, site_locals) - body_idx, local_xform = self._resolve_ancestor_body(prim, body_label_to_idx, xform_cache) - site_bodies.append(body_idx) - site_locals.append(local_xform) - - parent = prim.GetParent() - if not parent or not parent.IsValid() or parent.GetPath().pathString == "/": - parent_bodies.append(WORLD_BODY_INDEX) - parent_locals.append(identity_xform) - else: - parent_path = parent.GetPath().pathString - if parent_path in resolve_cache: - pb_idx, pb_local = resolve_cache[parent_path] - elif parent_path in body_label_to_idx: - pb_idx = body_label_to_idx[parent_path] - pb_local = identity_xform - resolve_cache[parent_path] = (pb_idx, pb_local) - else: - pb_idx, pb_local = self._resolve_ancestor_body(parent, body_label_to_idx, xform_cache) - resolve_cache[parent_path] = (pb_idx, pb_local) - parent_bodies.append(pb_idx) - parent_locals.append(pb_local) + 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]] = [] + for body_patterns, xform, per_world, env_ids in self._site_specs: + if body_patterns is None: + if per_world: + if NewtonManager._world_xforms is None: + raise RuntimeError(f"FrameView '{self._prim_path}' needs Newton cloned-world transforms.") + world_ids = range(len(NewtonManager._world_xforms)) if env_ids is None else env_ids + for world_id in world_ids: + 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)]) + else: + site_bodies.append(WORLD_BODY_INDEX) + site_locals.append([float(v) for v in xform]) + continue + + for body_pattern in body_patterns: + matched_indices, _ = resolve_matching_names(body_pattern, body_labels, raise_when_no_match=False) + if not matched_indices: + raise ValueError( + f"FrameView '{self._prim_path}' body pattern '{body_pattern}' matched no Newton bodies." + ) + + for body_idx in matched_indices: + site_bodies.append(body_idx) + site_locals.append([float(v) for v in xform]) + + self._create_buffers(site_bodies, site_locals) + + def _create_buffers(self, site_bodies: list[int], site_locals: list[list[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._parent_site_body = wp.array(parent_bodies, dtype=wp.int32, device=device) - self._parent_site_local = wp.array( - [wp.transform(*x) for x in parent_locals], - dtype=wp.transformf, - 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._site_local = wp.array([wp.transform(*x) for x in site_locals], dtype=wp.transformf, 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._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) - @staticmethod - def _resolve_ancestor_body( - prim: Usd.Prim, - body_label_to_idx: dict[str, int], - xform_cache: UsdGeom.XformCache, - ) -> tuple[int, list[float]]: - """Walk USD ancestors to find the nearest Newton body and compute the relative local transform. - - Args: - prim: The USD prim to resolve. - body_label_to_idx: Dict mapping body prim paths to their Newton body indices. - xform_cache: USD xform cache for efficient transform lookups. - - Returns: - A tuple ``(body_index, local_xform_7)`` where *local_xform_7* is - ``[tx, ty, tz, qx, qy, qz, qw]``. If no body ancestor exists, - ``body_index`` is :data:`WORLD_BODY_INDEX` and the local transform - is the prim's world transform. - """ - prim_world_tf = xform_cache.GetLocalToWorldTransform(prim) - prim_world_tf.Orthonormalize() - - ancestor = prim.GetParent() - while ancestor and ancestor.IsValid() and ancestor.GetPath().pathString != "/": - ancestor_path = ancestor.GetPath().pathString - body_idx = body_label_to_idx.get(ancestor_path) - if body_idx is not None: - ancestor_world_tf = xform_cache.GetLocalToWorldTransform(ancestor) - ancestor_world_tf.Orthonormalize() - local_tf = prim_world_tf * ancestor_world_tf.GetInverse() - return body_idx, _gf_matrix_to_xform7(local_tf) - ancestor = ancestor.GetParent() - - return WORLD_BODY_INDEX, _gf_matrix_to_xform7(prim_world_tf) - @property def prims(self) -> list: - """List of USD prims being managed by this view.""" + """List of USD prims being managed by this view. + + Newton site views do not retain USD prim handles. + """ return self._prims @property def count(self) -> int: - """Number of prims in this view.""" - return len(self._prims) + """Number of frames in this view.""" + return self._count @property def device(self) -> str: - """Device where arrays are allocated (cpu or cuda).""" + """Device where arrays are allocated.""" return self._device - # ------------------------------------------------------------------ - # World poses - # ------------------------------------------------------------------ - def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: - """Get world-space positions and orientations. - - Args: - indices: Subset of sites to query. ``None`` means all sites. - - 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. - """ + """Get world-space positions and orientations.""" state = NewtonManager.get_state_0() - - if indices is not None: - n = len(indices) - pos_buf = wp.zeros(n, dtype=wp.vec3f, device=self._device) - quat_buf = wp.zeros(n, dtype=wp.vec4f, device=self._device) - wp.launch( - _compute_site_world_transforms_indexed, - dim=n, - inputs=[state.body_q, self._site_body, self._site_local, indices], - outputs=[pos_buf, quat_buf], - device=self._device, - ) - return ProxyArray(pos_buf), ProxyArray(quat_buf) + site_indices = self._site_indices if indices is None else indices + n = self.count if indices is None else len(indices) + pos_buf = self._pos_buf if indices is None else wp.zeros(n, dtype=wp.vec3f, device=self._device) + quat_buf = self._quat_buf if indices is None else wp.zeros(n, dtype=wp.vec4f, device=self._device) wp.launch( _compute_site_world_transforms, - dim=self.count, - inputs=[state.body_q, self._site_body, self._site_local], - outputs=[self._pos_buf, self._quat_buf], + dim=n, + inputs=[state.body_q, self._site_body, self._site_local, site_indices], + outputs=[pos_buf, quat_buf], device=self._device, ) - return self._pos_ta, self._quat_ta + if indices is None: + return self._pos_ta, self._quat_ta + return ProxyArray(pos_buf), ProxyArray(quat_buf) def set_world_poses( self, @@ -723,24 +441,11 @@ def set_world_poses( orientations: wp.array | None = None, indices: wp.array | None = None, ) -> None: - """Set world-space positions and/or orientations. - - Updates the internal ``site_local`` offsets so that - ``body_q[body] * new_site_local`` yields the desired world pose. - Does **not** modify ``body_q``. - - Args: - positions: Desired world positions ``(M, 3)``. ``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 state = NewtonManager.get_state_0() - if positions is None or orientations is None: cur_pos_ta, cur_quat_ta = self.get_world_poses(indices) if positions is None: @@ -748,74 +453,32 @@ def set_world_poses( if orientations is None: orientations = cur_quat_ta.warp - if indices is not None: - wp.launch( - _write_site_local_from_world_poses_indexed, - dim=len(indices), - inputs=[state.body_q, self._site_body, indices, positions, orientations, self._site_local], - device=self._device, - ) - else: - wp.launch( - _write_site_local_from_world_poses, - dim=self.count, - inputs=[state.body_q, self._site_body, positions, orientations, self._site_local], - device=self._device, - ) - - # ------------------------------------------------------------------ - # Local poses (parent-relative) - # ------------------------------------------------------------------ + site_indices = self._site_indices if indices is None else indices + n = self.count if indices is None else len(indices) + wp.launch( + _write_site_local_from_world_poses, + dim=n, + inputs=[state.body_q, self._site_body, site_indices, positions, orientations, self._site_local], + device=self._device, + ) 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. Use ``.warp`` for the underlying ``wp.array`` or ``.torch`` for a - cached zero-copy ``torch.Tensor`` view. - """ - state = NewtonManager.get_state_0() - - if indices is not None: - n = len(indices) - pos_buf = wp.zeros(n, dtype=wp.vec3f, device=self._device) - quat_buf = wp.zeros(n, dtype=wp.vec4f, device=self._device) - wp.launch( - _compute_site_local_transforms_indexed, - dim=n, - inputs=[ - state.body_q, - self._site_body, - self._site_local, - self._parent_site_body, - self._parent_site_local, - indices, - ], - outputs=[pos_buf, quat_buf], - device=self._device, - ) - return ProxyArray(pos_buf), ProxyArray(quat_buf) + """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) + pos_buf = self._local_pos_buf if indices is None else wp.zeros(n, dtype=wp.vec3f, device=self._device) + quat_buf = self._local_quat_buf if indices is None else wp.zeros(n, dtype=wp.vec4f, device=self._device) wp.launch( - _compute_site_local_transforms, - dim=self.count, - inputs=[ - state.body_q, - self._site_body, - self._site_local, - self._parent_site_body, - self._parent_site_local, - ], - outputs=[self._local_pos_buf, self._local_quat_buf], + _gather_site_local_transforms, + dim=n, + inputs=[self._site_local, site_indices], + outputs=[pos_buf, quat_buf], device=self._device, ) - return self._local_pos_ta, self._local_quat_ta + if indices is None: + return self._local_pos_ta, self._local_quat_ta + return ProxyArray(pos_buf), ProxyArray(quat_buf) def set_local_poses( self, @@ -823,24 +486,10 @@ def set_local_poses( orientations: wp.array | None = None, indices: wp.array | None = None, ) -> None: - """Set parent-relative translations and/or orientations. - - Updates the internal ``site_local`` offsets so that - ``inv(parent_world) * (body_q[bid] * site_local)`` yields the desired - local pose. Does **not** modify ``body_q``. - - Args: - translations: Desired parent-relative translations ``(M, 3)``. - ``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 body-local translations and/or orientations.""" if translations is None and orientations is None: return - state = NewtonManager.get_state_0() - if translations is None or orientations is None: cur_pos_ta, cur_quat_ta = self.get_local_poses(indices) if translations is None: @@ -848,127 +497,74 @@ def set_local_poses( if orientations is None: orientations = cur_quat_ta.warp - if indices is not None: - wp.launch( - _write_site_local_from_local_poses_indexed, - dim=len(indices), - inputs=[ - state.body_q, - self._site_body, - self._parent_site_body, - self._parent_site_local, - indices, - translations, - orientations, - self._site_local, - ], - device=self._device, - ) - else: - wp.launch( - _write_site_local_from_local_poses, - dim=self.count, - inputs=[ - state.body_q, - self._site_body, - self._parent_site_body, - self._parent_site_local, - translations, - orientations, - self._site_local, - ], - device=self._device, - ) + site_indices = self._site_indices if indices is None else indices + n = self.count if indices is None else len(indices) + wp.launch( + _write_site_local_from_local_poses, + dim=n, + inputs=[site_indices, translations, orientations, self._site_local], + device=self._device, + ) # ------------------------------------------------------------------ # Scales # ------------------------------------------------------------------ - def get_local_scales(self, indices: wp.array | None = None) -> ProxyArray: - """Get local-space scales. + def get_world_scales(self, indices: wp.array | None = None) -> ProxyArray: + """Get per-site world scales from the first collision shape on the same body. - .. note:: - Newton stores shape_scale as the effective (composed) scale. - get_local_scales returns the same value since Newton does not - decompose parent/child scale independently. + Newton's ``shape_scale`` is an absolute (world-space) quantity. """ - return ProxyArray(self._get_shape_scales(indices)) - - def get_world_scales(self, indices: wp.array | None = None) -> ProxyArray: - """Get world-space (composed) scales from Newton shape_scale.""" - return ProxyArray(self._get_shape_scales(indices)) + 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, + dim=n, + inputs=[model.shape_scale, model.shape_body, self._site_body, site_indices, num_shapes], + outputs=[out], + device=self._device, + ) + return ProxyArray(out) - def set_local_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: - """Set local-space scales. + def get_local_scales(self, indices: wp.array | None = None) -> ProxyArray: + """Get per-site local scales. - .. note:: - Newton stores shape_scale as the effective (composed) scale. - set_local_scales writes directly to shape_scale. + Newton does not distinguish local from world scale (``shape_scale`` is + absolute). Returns the same value as :meth:`get_world_scales`. """ - self._set_shape_scales(scales, indices) + return self.get_world_scales(indices) def set_world_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: - """Set world-space (composed) scales to Newton shape_scale.""" - self._set_shape_scales(scales, indices) + """Set per-site world scales by writing to all collision shapes on the same body. - def _get_scales_impl(self, indices=None): - """Newton default: get_scales returns shape_scale (world-like).""" - return self.get_world_scales(indices).warp - - def _set_scales_impl(self, scales, indices=None): - """Newton default: set_scales writes shape_scale (world-like).""" - self.set_world_scales(scales, indices) - - def _get_shape_scales(self, indices: wp.array | None = None) -> wp.array: - """Internal: read shape_scale from Newton model.""" + Newton's ``shape_scale`` is an absolute (world-space) quantity. + """ 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, + dim=n, + inputs=[self._site_body, site_indices, scales, model.shape_body, num_shapes, model.shape_scale], + device=self._device, + ) - if indices is not None: - n = len(indices) - out = wp.zeros(n, dtype=wp.vec3f, device=self._device) - wp.launch( - _gather_scales_indexed, - dim=n, - inputs=[model.shape_scale, model.shape_body, self._site_body, indices, num_shapes], - outputs=[out], - device=self._device, - ) - else: - out = wp.zeros(self.count, dtype=wp.vec3f, device=self._device) - wp.launch( - _gather_scales, - dim=self.count, - inputs=[model.shape_scale, model.shape_body, self._site_body, num_shapes], - outputs=[out], - device=self._device, - ) - return out - - def _set_shape_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: - """Internal: write shape_scale to Newton model.""" - model = NewtonManager.get_model() - num_shapes = model.shape_count + def set_local_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Set per-site local scales. - if indices is not None: - wp.launch( - _scatter_scales_indexed, - dim=len(indices), - inputs=[self._site_body, indices, scales, model.shape_body, num_shapes, model.shape_scale], - device=self._device, - ) - else: - wp.launch( - _scatter_scales, - dim=self.count, - inputs=[self._site_body, scales, model.shape_body, num_shapes, model.shape_scale], - device=self._device, - ) + Newton does not distinguish local from world scale (``shape_scale`` is + absolute). Delegates to :meth:`set_world_scales`. + """ + self.set_world_scales(scales, indices) + def _get_scales_impl(self, indices: wp.array | None = None) -> wp.array: + """Newton legacy: get_scales returns world scales.""" + return self.get_world_scales(indices).warp -def _gf_matrix_to_xform7(mat: Gf.Matrix4d) -> list[float]: - """Convert a ``Gf.Matrix4d`` to ``[tx, ty, tz, qx, qy, qz, qw]``.""" - t = mat.ExtractTranslation() - 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())] + def _set_scales_impl(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Newton legacy: set_scales writes world scales.""" + self.set_world_scales(scales, indices) From 3c6de78a28874708315984d38d586134704dd50f Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Tue, 2 Jun 2026 19:21:16 +0200 Subject: [PATCH 26/54] feat: add tests for world and local scale conversions under scaled parent Implemented new tests to validate the behavior of world and local scale conversions in a hierarchy with a scaled parent. The tests ensure that setting local scales correctly computes world scales and vice versa, addressing USD-specific scale math not covered by existing tests. --- .../changelog.d/fabric-local-poses.rst | 29 +++++++++++ .../test/sim/test_views_xform_prim.py | 50 +++++++++++++++++++ .../changelog.d/fabric-local-poses.rst | 19 +++++++ .../changelog.d/fabric-local-poses.rst | 19 +++++++ 4 files changed, 117 insertions(+) create mode 100644 source/isaaclab/changelog.d/fabric-local-poses.rst create mode 100644 source/isaaclab_newton/changelog.d/fabric-local-poses.rst create mode 100644 source/isaaclab_ovphysx/changelog.d/fabric-local-poses.rst 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..4705c62d5b34 --- /dev/null +++ b/source/isaaclab/changelog.d/fabric-local-poses.rst @@ -0,0 +1,29 @@ +Added +^^^^^ + +* Added explicit local/world scale methods + :meth:`~isaaclab.sim.views.BaseFrameView.get_local_scales`, + :meth:`~isaaclab.sim.views.BaseFrameView.set_local_scales`, + :meth:`~isaaclab.sim.views.BaseFrameView.get_world_scales`, and + :meth:`~isaaclab.sim.views.BaseFrameView.set_world_scales` to the FrameView + API, implemented for :class:`~isaaclab.sim.views.UsdFrameView`. Scale getters + now return :class:`~isaaclab.utils.warp.ProxyArray`. + +* 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``). + +Deprecated +^^^^^^^^^^ + +* Deprecated :meth:`~isaaclab.sim.views.BaseFrameView.get_scales` and + :meth:`~isaaclab.sim.views.BaseFrameView.set_scales` in favor of the explicit + ``get_local_scales`` / ``set_local_scales`` (operates on ``xformOp:scale``) or + ``get_world_scales`` / ``set_world_scales`` (operates on composed world-space + scale). The deprecated methods still work but emit a ``DeprecationWarning``; + :class:`~isaaclab.sim.views.UsdFrameView` preserves prior behavior by + defaulting to local scales. diff --git a/source/isaaclab/test/sim/test_views_xform_prim.py b/source/isaaclab/test/sim/test_views_xform_prim.py index 64cd86a7466f..52cc7a05fc80 100644 --- a/source/isaaclab/test/sim/test_views_xform_prim.py +++ b/source/isaaclab/test/sim/test_views_xform_prim.py @@ -237,6 +237,56 @@ 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_set_local_scales_then_get_world_scales(device): + """Under a scaled parent, world scale == parent_scale * local_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) + view.set_local_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_set_world_scales_then_get_local_scales(device): + """Under a scaled parent, set_world_scales writes local = world / parent_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) + view.set_world_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 # ================================================================== 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..373c3b589f52 --- /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`, + :meth:`~isaaclab_newton.sim.views.NewtonSiteFrameView.set_local_scales`, + :meth:`~isaaclab_newton.sim.views.NewtonSiteFrameView.get_world_scales`, and + :meth:`~isaaclab_newton.sim.views.NewtonSiteFrameView.set_world_scales`. + Newton's ``shape_scale`` is an absolute (world-space) quantity, so the local + methods return the same value as the world methods. Scale getters now return + :class:`~isaaclab.utils.warp.ProxyArray`. + +Deprecated +^^^^^^^^^^ + +* Deprecated :meth:`~isaaclab_newton.sim.views.NewtonSiteFrameView.get_scales` + and :meth:`~isaaclab_newton.sim.views.NewtonSiteFrameView.set_scales` in favor + of the explicit ``get_world_scales`` / ``set_world_scales`` (or their local + equivalents). The deprecated methods still work but emit a + ``DeprecationWarning`` and default to world scales, preserving prior behavior. 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..9ca040119b62 --- /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`, + :meth:`~isaaclab_ovphysx.sim.views.OvPhysxFrameView.set_local_scales`, + :meth:`~isaaclab_ovphysx.sim.views.OvPhysxFrameView.get_world_scales`, and + :meth:`~isaaclab_ovphysx.sim.views.OvPhysxFrameView.set_world_scales`, which + delegate to the internal :class:`~isaaclab.sim.views.UsdFrameView`. Scale + getters now return :class:`~isaaclab.utils.warp.ProxyArray`. + +Deprecated +^^^^^^^^^^ + +* Deprecated :meth:`~isaaclab_ovphysx.sim.views.OvPhysxFrameView.get_scales` and + :meth:`~isaaclab_ovphysx.sim.views.OvPhysxFrameView.set_scales` in favor of the + explicit ``get_local_scales`` / ``set_local_scales`` (operates on + ``xformOp:scale``) or ``get_world_scales`` / ``set_world_scales``. The + deprecated methods still work but emit a ``DeprecationWarning`` and default to + local scales, preserving prior behavior. From 025eea74042f32cc81b833ba8d6293c25464e18d Mon Sep 17 00:00:00 2001 From: pv-nvidia <{"message":"Not Found","documentation_url":"https://docs.github.com/rest/users/emails#list-email-addresses-for-the-authenticated-user","status":"404"}> Date: Wed, 3 Jun 2026 23:32:18 +0000 Subject: [PATCH 27/54] test: fix FrameView scale contracts --- .../test/sim/frame_view_contract_utils.py | 24 +++++++++++++++---- .../test/sim/test_views_xform_prim_newton.py | 2 +- 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/source/isaaclab/test/sim/frame_view_contract_utils.py b/source/isaaclab/test/sim/frame_view_contract_utils.py index 92e4eaafe31f..cc9161c45072 100644 --- a/source/isaaclab/test/sim/frame_view_contract_utils.py +++ b/source/isaaclab/test/sim/frame_view_contract_utils.py @@ -401,12 +401,28 @@ def test_return_types_are_torcharray(device, view_factory): 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__}" + assert isinstance(scales_full, wp.array), f"get_scales() must be wp.array, 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__}" - ) + assert isinstance(scales_idx, wp.array), f"get_scales(indices) must be wp.array, got {type(scales_idx).__name__}" finally: bundle.teardown() 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..53f884e99b4d 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 @@ -40,7 +40,7 @@ class _SceneCfg(InteractiveSceneCfg): cube: RigidObjectCfg = RigidObjectCfg( prim_path="{ENV_REGEX_NS}/Cube", spawn=sim_utils.CuboidCfg( - size=(0.2, 0.2, 0.2), + size=(2.0, 2.0, 2.0), rigid_props=sim_utils.RigidBodyPropertiesCfg(), mass_props=sim_utils.MassPropertiesCfg(mass=1.0), collision_props=sim_utils.CollisionPropertiesCfg(), From 7e82bd1b4923faeff7e677741e5dea30decfc589 Mon Sep 17 00:00:00 2001 From: pv-nvidia <{"message":"Not Found","documentation_url":"https://docs.github.com/rest/users/emails#list-email-addresses-for-the-authenticated-user","status":"404"}> Date: Thu, 4 Jun 2026 04:49:54 +0000 Subject: [PATCH 28/54] style: format FrameView scale contract test --- source/isaaclab/test/sim/frame_view_contract_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/source/isaaclab/test/sim/frame_view_contract_utils.py b/source/isaaclab/test/sim/frame_view_contract_utils.py index cc9161c45072..d378cee63ae5 100644 --- a/source/isaaclab/test/sim/frame_view_contract_utils.py +++ b/source/isaaclab/test/sim/frame_view_contract_utils.py @@ -422,7 +422,9 @@ def test_return_types_are_torcharray(device, view_factory): scales_full = bundle.view.get_scales() assert isinstance(scales_full, wp.array), f"get_scales() must be wp.array, got {type(scales_full).__name__}" scales_idx = bundle.view.get_scales(indices) - assert isinstance(scales_idx, wp.array), f"get_scales(indices) must be wp.array, got {type(scales_idx).__name__}" + assert isinstance(scales_idx, wp.array), ( + f"get_scales(indices) must be wp.array, got {type(scales_idx).__name__}" + ) finally: bundle.teardown() From fc9f12c5748d16b8320d8068baf1a96d74910318 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Thu, 4 Jun 2026 09:41:23 +0000 Subject: [PATCH 29/54] bench: include FrameView scale operations --- .../benchmarks/benchmark_xform_prim_view.py | 69 +++++++++++++++++++ 1 file changed, 69 insertions(+) diff --git a/scripts/benchmarks/benchmark_xform_prim_view.py b/scripts/benchmarks/benchmark_xform_prim_view.py index 76b5fab35863..ca3c53a18b0a 100644 --- a/scripts/benchmarks/benchmark_xform_prim_view.py +++ b/scripts/benchmarks/benchmark_xform_prim_view.py @@ -148,6 +148,7 @@ def to_torch(a): try: # -- Warmup -------------------------------------------------------- xform_view.get_world_poses() + xform_view.get_world_scales() # -- get_world_poses ----------------------------------------------- if is_newton: @@ -221,6 +222,70 @@ def to_torch(a): computed_results["local_translations_after_set"] = to_torch(ta).clone() computed_results["local_orientations_after_set"] = to_torch(ola).clone() + # -- get_world_scales ---------------------------------------------- + if is_newton: + torch.cuda.synchronize() + start_time = time.perf_counter() + for _ in range(num_iterations): + world_scales = xform_view.get_world_scales() + if is_newton: + torch.cuda.synchronize() + 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 + + if is_newton: + torch.cuda.synchronize() + start_time = time.perf_counter() + for _ in range(num_iterations): + xform_view.set_world_scales(new_world_scales) + if is_newton: + torch.cuda.synchronize() + 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 ---------------------------------------------- + if is_newton: + torch.cuda.synchronize() + start_time = time.perf_counter() + for _ in range(num_iterations): + local_scales = xform_view.get_local_scales() + if is_newton: + torch.cuda.synchronize() + 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: + 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 + + if is_newton: + torch.cuda.synchronize() + start_time = time.perf_counter() + for _ in range(num_iterations): + xform_view.set_local_scales(new_local_scales) + if is_newton: + torch.cuda.synchronize() + 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) -------------------------------------- if is_newton: torch.cuda.synchronize() @@ -277,6 +342,10 @@ def print_results(results_dict: dict[str, dict[str, float]], num_prims: int, num ("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"), ] From 1566c637b6be9c3f3a74f1301a1599ae1fb958c0 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Thu, 4 Jun 2026 19:54:17 +0000 Subject: [PATCH 30/54] fix: preserve ProxyArray get_scales contract --- source/isaaclab/isaaclab/sim/views/base_frame_view.py | 6 +++--- source/isaaclab/isaaclab/sim/views/usd_frame_view.py | 6 +++--- source/isaaclab/test/sim/frame_view_contract_utils.py | 6 +++--- .../isaaclab_newton/sim/views/newton_site_frame_view.py | 4 ++-- .../isaaclab_ovphysx/sim/views/ovphysx_frame_view.py | 4 ++-- .../isaaclab_physx/sim/views/fabric_frame_view.py | 2 +- 6 files changed, 14 insertions(+), 14 deletions(-) diff --git a/source/isaaclab/isaaclab/sim/views/base_frame_view.py b/source/isaaclab/isaaclab/sim/views/base_frame_view.py index fb936bcd1f89..ee5e313d7f3c 100644 --- a/source/isaaclab/isaaclab/sim/views/base_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/base_frame_view.py @@ -157,7 +157,7 @@ def set_world_scales(self, scales: wp.array, indices: wp.array | None = None) -> _get_scales_deprecated_warned: bool = False _set_scales_deprecated_warned: bool = False - def get_scales(self, indices: wp.array | None = None) -> wp.array: + def get_scales(self, indices: wp.array | None = None) -> ProxyArray: """Get scales for prims in the view. .. deprecated:: @@ -169,7 +169,7 @@ def get_scales(self, indices: wp.array | None = None) -> wp.array: indices: Subset of prims to query. ``None`` means all prims. Returns: - A ``wp.array`` of shape ``(M, 3)``. + A ``ProxyArray`` of shape ``(M, 3)``. """ if not BaseFrameView._get_scales_deprecated_warned: BaseFrameView._get_scales_deprecated_warned = True @@ -202,7 +202,7 @@ def set_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: self._set_scales_impl(scales, indices) @abc.abstractmethod - def _get_scales_impl(self, indices: wp.array | None = None) -> wp.array: + def _get_scales_impl(self, indices: wp.array | None = None) -> ProxyArray: """Backend-specific implementation for deprecated get_scales().""" ... diff --git a/source/isaaclab/isaaclab/sim/views/usd_frame_view.py b/source/isaaclab/isaaclab/sim/views/usd_frame_view.py index 24c73ae036b0..3927b4690a62 100644 --- a/source/isaaclab/isaaclab/sim/views/usd_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/usd_frame_view.py @@ -262,9 +262,9 @@ def set_world_scales(self, scales: wp.array, indices: wp.array | None = None): ) prim.GetAttribute("xformOp:scale").Set(local_scale) - def _get_scales_impl(self, indices: wp.array | None = None) -> wp.array: - """USD default: get_scales returns local scales.""" - return self.get_local_scales(indices).warp + def _get_scales_impl(self, indices: wp.array | None = None) -> ProxyArray: + """USD legacy: get_scales returns local scales.""" + return self.get_local_scales(indices) def _set_scales_impl(self, scales: wp.array, indices: wp.array | None = None) -> None: """USD default: set_scales writes local scales.""" diff --git a/source/isaaclab/test/sim/frame_view_contract_utils.py b/source/isaaclab/test/sim/frame_view_contract_utils.py index d378cee63ae5..b27335547d72 100644 --- a/source/isaaclab/test/sim/frame_view_contract_utils.py +++ b/source/isaaclab/test/sim/frame_view_contract_utils.py @@ -420,10 +420,10 @@ def test_return_types_are_torcharray(device, view_factory): ) scales_full = bundle.view.get_scales() - assert isinstance(scales_full, wp.array), f"get_scales() must be wp.array, got {type(scales_full).__name__}" + 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, wp.array), ( - f"get_scales(indices) must be wp.array, got {type(scales_idx).__name__}" + assert isinstance(scales_idx, ProxyArray), ( + f"get_scales(indices) must be ProxyArray, got {type(scales_idx).__name__}" ) finally: bundle.teardown() 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 33d34a89dd8e..2c016d1c7e9f 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 @@ -561,9 +561,9 @@ def set_local_scales(self, scales: wp.array, indices: wp.array | None = None) -> """ self.set_world_scales(scales, indices) - def _get_scales_impl(self, indices: wp.array | None = None) -> wp.array: + def _get_scales_impl(self, indices: wp.array | None = None) -> ProxyArray: """Newton legacy: get_scales returns world scales.""" - return self.get_world_scales(indices).warp + return self.get_world_scales(indices) def _set_scales_impl(self, scales: wp.array, indices: wp.array | None = None) -> None: """Newton legacy: set_scales writes world 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 ac0c591b40f6..acf8abe5822b 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 @@ -877,8 +877,8 @@ def set_world_scales(self, scales: wp.array, indices: wp.array | None = None) -> self._ensure_usd_view().set_world_scales(scales, indices) def _get_scales_impl(self, indices=None): - """OvPhysX default: get_scales returns local scales (same as USD).""" - return self.get_local_scales(indices).warp + """OvPhysX legacy: get_scales returns local scales (same as USD).""" + return self.get_local_scales(indices) def _set_scales_impl(self, scales, indices=None): """OvPhysX default: set_scales writes local scales (same as USD).""" 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 29302376f5ea..8bfb5224abce 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 @@ -553,7 +553,7 @@ def get_local_scales(self, indices=None): def _get_scales_impl(self, indices=None): """Fabric: deprecated get_scales delegates to get_world_scales.""" - return self.get_world_scales(indices).warp + return self.get_world_scales(indices) def _set_scales_impl(self, scales, indices=None): """Fabric: deprecated set_scales delegates to set_world_scales.""" From ead581b5e04a82e45d448b3b271f1912e4d8d365 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Thu, 4 Jun 2026 20:01:26 +0000 Subject: [PATCH 31/54] docs: remove ProxyArray scale return attribution --- source/isaaclab/changelog.d/fabric-local-poses.rst | 3 +-- source/isaaclab_newton/changelog.d/fabric-local-poses.rst | 3 +-- source/isaaclab_ovphysx/changelog.d/fabric-local-poses.rst | 3 +-- 3 files changed, 3 insertions(+), 6 deletions(-) diff --git a/source/isaaclab/changelog.d/fabric-local-poses.rst b/source/isaaclab/changelog.d/fabric-local-poses.rst index 4705c62d5b34..1ca61cfda97b 100644 --- a/source/isaaclab/changelog.d/fabric-local-poses.rst +++ b/source/isaaclab/changelog.d/fabric-local-poses.rst @@ -6,8 +6,7 @@ Added :meth:`~isaaclab.sim.views.BaseFrameView.set_local_scales`, :meth:`~isaaclab.sim.views.BaseFrameView.get_world_scales`, and :meth:`~isaaclab.sim.views.BaseFrameView.set_world_scales` to the FrameView - API, implemented for :class:`~isaaclab.sim.views.UsdFrameView`. Scale getters - now return :class:`~isaaclab.utils.warp.ProxyArray`. + API, implemented for :class:`~isaaclab.sim.views.UsdFrameView`. * Added :func:`~isaaclab.utils.warp.fabric.decompose_indexed_fabric_transforms`, :func:`~isaaclab.utils.warp.fabric.compose_indexed_fabric_transforms`, diff --git a/source/isaaclab_newton/changelog.d/fabric-local-poses.rst b/source/isaaclab_newton/changelog.d/fabric-local-poses.rst index 373c3b589f52..314c137f4743 100644 --- a/source/isaaclab_newton/changelog.d/fabric-local-poses.rst +++ b/source/isaaclab_newton/changelog.d/fabric-local-poses.rst @@ -6,8 +6,7 @@ Added :meth:`~isaaclab_newton.sim.views.NewtonSiteFrameView.get_world_scales`, and :meth:`~isaaclab_newton.sim.views.NewtonSiteFrameView.set_world_scales`. Newton's ``shape_scale`` is an absolute (world-space) quantity, so the local - methods return the same value as the world methods. Scale getters now return - :class:`~isaaclab.utils.warp.ProxyArray`. + methods return the same value as the world methods. Deprecated ^^^^^^^^^^ diff --git a/source/isaaclab_ovphysx/changelog.d/fabric-local-poses.rst b/source/isaaclab_ovphysx/changelog.d/fabric-local-poses.rst index 9ca040119b62..eb8caa5d254f 100644 --- a/source/isaaclab_ovphysx/changelog.d/fabric-local-poses.rst +++ b/source/isaaclab_ovphysx/changelog.d/fabric-local-poses.rst @@ -5,8 +5,7 @@ Added :meth:`~isaaclab_ovphysx.sim.views.OvPhysxFrameView.set_local_scales`, :meth:`~isaaclab_ovphysx.sim.views.OvPhysxFrameView.get_world_scales`, and :meth:`~isaaclab_ovphysx.sim.views.OvPhysxFrameView.set_world_scales`, which - delegate to the internal :class:`~isaaclab.sim.views.UsdFrameView`. Scale - getters now return :class:`~isaaclab.utils.warp.ProxyArray`. + delegate to the internal :class:`~isaaclab.sim.views.UsdFrameView`. Deprecated ^^^^^^^^^^ From 57fd1268e3b23a6529e2ca07cdbb1430bbd7f31c Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Thu, 4 Jun 2026 20:21:31 +0000 Subject: [PATCH 32/54] fix: preserve Newton legacy scale semantics --- .../changelog.d/fabric-local-poses.rst | 13 +- .../sim/views/newton_site_frame_view.py | 166 +++++++++++++----- .../test/sim/test_views_xform_prim_newton.py | 2 +- 3 files changed, 130 insertions(+), 51 deletions(-) diff --git a/source/isaaclab_newton/changelog.d/fabric-local-poses.rst b/source/isaaclab_newton/changelog.d/fabric-local-poses.rst index 314c137f4743..4e18af38eb2b 100644 --- a/source/isaaclab_newton/changelog.d/fabric-local-poses.rst +++ b/source/isaaclab_newton/changelog.d/fabric-local-poses.rst @@ -4,15 +4,16 @@ Added * Added :meth:`~isaaclab_newton.sim.views.NewtonSiteFrameView.get_local_scales`, :meth:`~isaaclab_newton.sim.views.NewtonSiteFrameView.set_local_scales`, :meth:`~isaaclab_newton.sim.views.NewtonSiteFrameView.get_world_scales`, and - :meth:`~isaaclab_newton.sim.views.NewtonSiteFrameView.set_world_scales`. - Newton's ``shape_scale`` is an absolute (world-space) quantity, so the local - methods return the same value as the world methods. + :meth:`~isaaclab_newton.sim.views.NewtonSiteFrameView.set_world_scales` for + transform (xform) scales. These explicit 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 ``get_world_scales`` / ``set_world_scales`` (or their local - equivalents). The deprecated methods still work but emit a - ``DeprecationWarning`` and default to world scales, preserving prior behavior. + of the explicit xform-scale ``get_world_scales`` / ``set_world_scales`` (or + their local equivalents). 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/isaaclab_newton/sim/views/newton_site_frame_view.py b/source/isaaclab_newton/isaaclab_newton/sim/views/newton_site_frame_view.py index 2c016d1c7e9f..3cab904c1ce0 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 @@ -104,7 +104,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 +112,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 +126,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 +134,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 +143,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 +200,49 @@ 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 +296,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 +309,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 +335,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 +344,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,12 +353,12 @@ 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_prim = stage.GetPrimAtPath(source_root) if source_root 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.""" @@ -336,8 +371,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] @@ -345,16 +381,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: @@ -364,9 +402,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: @@ -379,24 +419,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: @@ -511,60 +560,89 @@ def set_local_poses( # ------------------------------------------------------------------ def get_world_scales(self, indices: wp.array | None = None) -> ProxyArray: - """Get per-site world scales from the first collision shape on the same body. + """Get per-site world xform scales. - Newton's ``shape_scale`` is an absolute (world-space) quantity. + These are transform scales, matching the USD FrameView scale API. They + are intentionally separate from Newton collision shape geometry sizes. """ - 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) + if indices is None: + return self._scale_ta + n = len(indices) out = wp.zeros(n, dtype=wp.vec3f, device=self._device) wp.launch( - _gather_scales, + _gather_xform_scales, dim=n, - inputs=[model.shape_scale, model.shape_body, self._site_body, site_indices, num_shapes], + inputs=[self._site_xform_scale, indices], outputs=[out], device=self._device, ) return ProxyArray(out) def get_local_scales(self, indices: wp.array | None = None) -> ProxyArray: - """Get per-site local scales. + """Get per-site local xform scales. - Newton does not distinguish local from world scale (``shape_scale`` is - absolute). Returns the same value as :meth:`get_world_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(indices) def set_world_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: - """Set per-site world scales by writing to all collision shapes on the same body. + """Set per-site world xform scales. - Newton's ``shape_scale`` is an absolute (world-space) quantity. + These update transform scale state only; use deprecated ``set_scales`` if + legacy Newton collision shape geometry-scale behavior is required. """ - 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) + if indices is None: + indices = self._site_indices + n = self.count if indices is self._site_indices else len(indices) wp.launch( - _scatter_scales, + _scatter_xform_scales, dim=n, - inputs=[self._site_body, site_indices, scales, model.shape_body, num_shapes, model.shape_scale], + inputs=[indices, scales, self._site_xform_scale], device=self._device, ) def set_local_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: - """Set per-site local scales. + """Set per-site local xform scales. - Newton does not distinguish local from world scale (``shape_scale`` is - absolute). Delegates to :meth:`set_world_scales`. + These update transform scale state only; use deprecated ``set_scales`` if + legacy Newton collision shape geometry-scale behavior is required. """ self.set_world_scales(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_shape_scales, + dim=n, + inputs=[model.shape_scale, model.shape_body, self._site_body, site_indices, num_shapes], + outputs=[out], + device=self._device, + ) + return ProxyArray(out) + + 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_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 world scales.""" - return self.get_world_scales(indices) + """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: set_scales writes world scales.""" - self.set_world_scales(scales, indices) + """Newton legacy: set_scales writes collision shape geometry scales.""" + self._set_legacy_shape_scales(scales, indices) 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 53f884e99b4d..d114a1da2a80 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 @@ -40,7 +40,7 @@ class _SceneCfg(InteractiveSceneCfg): cube: RigidObjectCfg = RigidObjectCfg( prim_path="{ENV_REGEX_NS}/Cube", spawn=sim_utils.CuboidCfg( - size=(2.0, 2.0, 2.0), + size=(0.2, 0.2, 0.2), rigid_props=sim_utils.RigidBodyPropertiesCfg(), mass_props=sim_utils.MassPropertiesCfg(mass=1.0), collision_props=sim_utils.CollisionPropertiesCfg(), From 96195d72274b29873b820276e7cb5951ae9c9b21 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Thu, 4 Jun 2026 20:59:33 +0000 Subject: [PATCH 33/54] style: format Newton scale type annotation --- .../isaaclab_newton/sim/views/newton_site_frame_view.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 3cab904c1ce0..b7d2918e01c7 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 @@ -242,7 +242,9 @@ def _resolve_site_specs( 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, tuple[float, float, float], 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]: From c109f6a0e73d010256a399b3d82a4114b0b4acdf Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Thu, 4 Jun 2026 23:32:10 +0000 Subject: [PATCH 34/54] fix: flush Fabric world matrices after local writes --- .../isaaclab/sim/views/base_frame_view.py | 16 ++ .../sim/views/fabric_frame_view.py | 55 +++++-- .../test/sim/test_views_xform_prim_fabric.py | 151 ++++++++++++++++-- 3 files changed, 202 insertions(+), 20 deletions(-) diff --git a/source/isaaclab/isaaclab/sim/views/base_frame_view.py b/source/isaaclab/isaaclab/sim/views/base_frame_view.py index ee5e313d7f3c..0146c6328b4b 100644 --- a/source/isaaclab/isaaclab/sim/views/base_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/base_frame_view.py @@ -8,6 +8,7 @@ from __future__ import annotations import abc +import contextlib import warnings import warp as wp @@ -38,6 +39,21 @@ def device(self) -> str: """Device where arrays are allocated (``"cpu"`` or ``"cuda:0"``).""" ... + + @contextlib.contextmanager + def change_block(self): + """Batch multiple transform writes into one logical change. + + Backends may use this context to defer expensive derived-state updates + until the outermost block exits. The default implementation is a no-op + so callers can use it with every FrameView backend. + """ + yield self + + def changeBlock(self): + """Alias for :meth:`change_block`.""" + return self.change_block() + @abc.abstractmethod def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: """Get world-space positions and orientations for prims in the view. 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 8bfb5224abce..fc6b919706ed 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,6 +7,7 @@ from __future__ import annotations +import contextlib import enum import logging @@ -74,16 +75,18 @@ class FabricFrameView(BaseFrameView): 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. - * **World ↔ local consistency (lazy).** Getters are lazy: after - ``set_world_poses`` or ``set_world_scales``, local matrices are only - recomputed when ``get_local_poses`` (or ``get_local_scales``) is called; - after ``set_local_poses`` or ``set_local_scales``, world matrices are - only recomputed when ``get_world_poses`` (or ``get_world_scales``) is - called. Both directions stay in sync without round-tripping through USD. + * **World ↔ local consistency.** Getters are lazy after world writes: + ``set_world_poses`` or ``set_world_scales`` recompute local matrices only + when ``get_local_poses`` (or ``get_local_scales``) is called. Local + writes update renderer-facing ``omni:fabric:worldMatrix`` eagerly so + Fabric render delegates see the new pose/scale without requiring an Isaac + Lab world getter. Wrap multiple local writes in :meth:`change_block` to + defer this local→world recompute until the outermost block exits. * **Dirty-flag invariant.** The ``_dirty`` enum is one of ``NONE``, ``WORLD``, or ``LOCAL`` -- mutually exclusive by construction. ``set_world_poses`` / ``set_world_scales`` sets ``_dirty = LOCAL``; - ``set_local_poses`` / ``set_local_scales`` sets ``_dirty = WORLD``. + ``set_local_poses`` / ``set_local_scales`` flushes local→world + immediately unless a :meth:`change_block` is active. If the user interleaves both setters on the same view within a single frame, the second setter flushes the first's stale data before writing. This is correct but incurs an extra kernel launch -- a one-time warning @@ -149,6 +152,7 @@ def __init__( # Per-view (not per-stage) so concurrent views on the same stage don't interfere. self._dirty: _DirtyFlag = _DirtyFlag.NONE self._warned_interleaved_set: bool = False + self._change_block_depth: int = 0 # Selection (single RW covering both world + local matrix). self._sel = None @@ -179,6 +183,33 @@ def device(self) -> str: """Device where arrays are allocated (cpu or cuda).""" return self._device + @contextlib.contextmanager + def change_block(self): + """Batch Fabric writes and flush derived matrices on outermost exit. + + Local pose/scale setters normally update cached Fabric world matrices + immediately for renderer/FSD consumers. Use this block when applying + several local edits to the same view so the local→world recompute runs + once after the final edit. + """ + self._change_block_depth += 1 + try: + yield self + finally: + self._change_block_depth -= 1 + if self._change_block_depth == 0: + self._flush_deferred_local_writes() + + def changeBlock(self): + """Alias for :meth:`change_block`.""" + return self.change_block() + + def _flush_deferred_local_writes(self) -> None: + """Flush pending local→world recompute unless still inside a change block.""" + if self._change_block_depth > 0: + return + self._sync_world_from_local_if_dirty() + @property def prims(self) -> list: return self._usd_view.prims @@ -338,8 +369,11 @@ def set_local_poses(self, translations=None, orientations=None, indices=None): ) wp.synchronize() - # Mark this view's worlds stale so the next world read recomputes them. + # Keep renderer-facing world matrices current. Inside change_block(), + # defer the recompute until the outermost block exits so pose+scale + # edits only pay one local→world kernel launch. self._dirty = _DirtyFlag.WORLD + self._flush_deferred_local_writes() def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: """Return (translations, orientations) in parent-local frame. @@ -503,8 +537,11 @@ def set_local_scales(self, scales, indices=None): ) wp.synchronize() - # Local was just written -- mark world poses as stale. + # Keep renderer-facing world matrices current. Inside change_block(), + # defer the recompute until the outermost block exits so pose+scale + # edits only pay one local→world kernel launch. self._dirty = _DirtyFlag.WORLD + self._flush_deferred_local_writes() def get_local_scales(self, indices=None): """Return per-prim (sx, sy, sz) scales extracted from local matrix. 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 72845ed3121b..d56c9cef59ac 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 @@ -231,6 +231,136 @@ def test_prepare_for_reuse_detects_topology_change(device, view_factory): 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.""" + 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() + 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, + ) + + +@pytest.mark.parametrize("device", ["cuda:0"]) +def test_set_local_poses_updates_renderer_facing_fabric_world_matrix(device, view_factory): + """Local pose writes must update cached Fabric worldMatrix immediately. + + The FSD renderer reads Fabric's cached ``omni:fabric:worldMatrix`` directly; + it does not call ``FrameView.get_world_poses()`` to trigger Isaac Lab's lazy + local→world sync. This test intentionally reads the Fabric attribute + directly after ``set_local_poses`` and before any world getter call. + """ + bundle = view_factory(num_envs=1, device=device) + view = bundle.view + + # Initialize Fabric and clear the initial dirty state. + view.get_world_poses() + assert view._dirty.name == "NONE" + + 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) + new_local_ori = wp.from_torch(torch.tensor([[0.0, 0.0, 0.0, 1.0]], dtype=torch.float32, device=device)) + + view.set_local_poses(translations=new_local_pos, orientations=new_local_ori) + + # Parent is at (0, 0, 1), so renderer-facing cached worldMatrix should + # already contain world translation (1, 2, 4) without get_world_poses(). + 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_set_local_scales_updates_renderer_facing_fabric_world_matrix(device, view_factory): + """Local scale writes must update cached Fabric worldMatrix immediately. + + This is the scale analogue of the local-pose renderer/FSD contract: read + cached ``omni:fabric:worldMatrix`` directly after ``set_local_scales`` and + before any FrameView world getter can repair stale state. + """ + bundle = view_factory(num_envs=1, device=device) + view = bundle.view + + # Initialize Fabric and clear the initial dirty state. + view.get_world_poses() + assert view._dirty.name == "NONE" + + 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) + + view.set_local_scales(new_scales) + + 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) + + +@pytest.mark.parametrize("device", ["cuda:0"]) +def test_change_block_batches_local_world_matrix_update(device, view_factory, monkeypatch): + """Local pose+scale writes inside change_block flush worldMatrix once on exit.""" + bundle = view_factory(num_envs=1, device=device) + view = bundle.view + + view.get_world_poses() + assert view._dirty.name == "NONE" + + calls = 0 + original_recompute = view._recompute_world_from_local + + def counted_recompute(): + nonlocal calls + calls += 1 + original_recompute() + + monkeypatch.setattr(view, "_recompute_world_from_local", counted_recompute) + + 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) + new_local_ori = wp.from_torch(torch.tensor([[0.0, 0.0, 0.0, 1.0]], dtype=torch.float32, 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.change_block(): + view.set_local_poses(translations=new_local_pos, orientations=new_local_ori) + assert view._dirty.name == "WORLD" + assert calls == 0 + + view.set_local_scales(new_scales) + assert view._dirty.name == "WORLD" + assert calls == 0 + + assert calls == 1 + assert view._dirty.name == "NONE" + + 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) + + @pytest.mark.parametrize("device", ["cuda:0"]) def test_set_local_via_fabric_path(device, view_factory): """Exercise the Fabric-native set_local_poses path. @@ -293,8 +423,8 @@ def test_local_scales_roundtrip(device, view_factory): wp.launch(kernel=_fill_position, dim=2, inputs=[new_scales, 2.0, 3.0, 4.0], device=device) view.set_local_scales(new_scales) - # Should have dirtied world - assert view._dirty.name == "WORLD" + # Local writes eagerly update renderer-facing world matrices outside change_block(). + assert view._dirty.name == "NONE" ret_scales = view.get_local_scales() scales_torch = ret_scales.torch @@ -506,21 +636,20 @@ def test_multi_view_per_view_dirty_isolation(device): identity_quat = wp.from_torch(torch.tensor([[0.0, 0.0, 0.0, 1.0]], dtype=torch.float32, device=device)) view_a.set_local_poses(translations=new_local_a, orientations=identity_quat) - # Only view A should be dirty. Critical: a per-stage flag would have dirtied - # both views (or neither) at this point. - assert view_a._dirty.name == "WORLD", "set_local_poses should mark its own view dirty" + # Local writes eagerly update renderer-facing world matrices outside change_block(), + # and they must not dirty another view on the same stage. + assert view_a._dirty.name == "NONE", "set_local_poses should flush its own view outside change_block" assert view_b._dirty.name == "NONE", "set_local_poses on view A must not dirty view B" - # Read worlds from view B FIRST. With a per-stage flag, B's - # ``_sync_world_from_local_if_dirty`` would fire and clear the flag, leaving A - # stale. With the per-view flag, B's read is a no-op sync-wise. + # Read worlds from view B FIRST. This must not affect view A's already-flushed + # world matrices. torch.testing.assert_close( torch.as_tensor(view_b.get_world_poses()[0], device=device), expected_b0, atol=1e-5, rtol=0 ) assert view_b._dirty.name == "NONE" - assert view_a._dirty.name == "WORLD", "view B's world read must not clear view A's dirty flag" + assert view_a._dirty.name == "NONE" - # Now read view A's worlds -- sync fires, world reflects the new local. + # Now read view A's worlds -- world already 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 @@ -532,7 +661,7 @@ def test_multi_view_per_view_dirty_isolation(device): wp.launch(kernel=_fill_position, dim=1, inputs=[new_local_b, 3.0, 0.0, 0.0], device=device) view_b.set_local_poses(translations=new_local_b, orientations=identity_quat) assert view_a._dirty.name == "NONE" - assert view_b._dirty.name == "WORLD" + assert view_b._dirty.name == "NONE" # A's worlds must still read back the post-set-local value from above; no # cross-view stomp on the world matrix. From 2bc2c64d134ed27e4af0f056e998775df621f96e Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Thu, 4 Jun 2026 23:33:11 +0000 Subject: [PATCH 35/54] style: apply pre-commit formatting --- source/isaaclab/isaaclab/sim/views/base_frame_view.py | 1 - source/isaaclab_physx/test/sim/test_views_xform_prim_fabric.py | 2 -- 2 files changed, 3 deletions(-) diff --git a/source/isaaclab/isaaclab/sim/views/base_frame_view.py b/source/isaaclab/isaaclab/sim/views/base_frame_view.py index 0146c6328b4b..0f9ee66281da 100644 --- a/source/isaaclab/isaaclab/sim/views/base_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/base_frame_view.py @@ -39,7 +39,6 @@ def device(self) -> str: """Device where arrays are allocated (``"cpu"`` or ``"cuda:0"``).""" ... - @contextlib.contextmanager def change_block(self): """Batch multiple transform writes into one logical change. 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 d56c9cef59ac..f81881ba0540 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 @@ -233,8 +233,6 @@ def test_prepare_for_reuse_detects_topology_change(device, view_factory): def _read_fabric_world_matrix_translation(view, prim_index=0): """Read cached Fabric worldMatrix 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() From 2966561f59a02c8fcc575176ea0708bb22111438 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Fri, 5 Jun 2026 00:00:31 +0000 Subject: [PATCH 36/54] feat: configure Fabric change-block matrix flushes --- .../isaaclab/sim/views/base_frame_view.py | 25 +++- .../sim/views/fabric_frame_view.py | 79 +++++++++--- .../test/sim/test_views_xform_prim_fabric.py | 120 ++++++++++++++++++ 3 files changed, 204 insertions(+), 20 deletions(-) diff --git a/source/isaaclab/isaaclab/sim/views/base_frame_view.py b/source/isaaclab/isaaclab/sim/views/base_frame_view.py index 0f9ee66281da..31ea771a4132 100644 --- a/source/isaaclab/isaaclab/sim/views/base_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/base_frame_view.py @@ -40,18 +40,35 @@ def device(self) -> str: ... @contextlib.contextmanager - def change_block(self): + def change_block(self, update_world_matrices: bool = True, update_local_matrices: bool = False): """Batch multiple transform writes into one logical change. Backends may use this context to defer expensive derived-state updates until the outermost block exits. The default implementation is a no-op so callers can use it with every FrameView backend. + + Args: + update_world_matrices: Whether derived world matrices should be + updated before the outermost block exits. Backends may ignore + this option when they do not maintain separate cached world + matrices. + update_local_matrices: Whether derived local matrices should be + updated before the outermost block exits. Backends may ignore + this option when they do not maintain separate cached local + matrices. """ yield self - def changeBlock(self): - """Alias for :meth:`change_block`.""" - return self.change_block() + def changeBlock( + self, + updateWorldMatrices: bool = True, + updateLocalMatrices: bool = False, + ): + """CamelCase alias for :meth:`change_block`.""" + return self.change_block( + update_world_matrices=updateWorldMatrices, + update_local_matrices=updateLocalMatrices, + ) @abc.abstractmethod def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: 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 fc6b919706ed..8aa1c788ea8f 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 @@ -80,13 +80,18 @@ class FabricFrameView(BaseFrameView): when ``get_local_poses`` (or ``get_local_scales``) is called. Local writes update renderer-facing ``omni:fabric:worldMatrix`` eagerly so Fabric render delegates see the new pose/scale without requiring an Isaac - Lab world getter. Wrap multiple local writes in :meth:`change_block` to - defer this local→world recompute until the outermost block exits. + Lab world getter. The default is intentionally asymmetric: render + delegates consume world matrices, while no known downstream render path + consumes local matrices directly. Wrap multiple writes in + :meth:`change_block` to defer and configure derived-matrix updates. * **Dirty-flag invariant.** The ``_dirty`` enum is one of ``NONE``, ``WORLD``, or ``LOCAL`` -- mutually exclusive by construction. ``set_world_poses`` / ``set_world_scales`` sets ``_dirty = LOCAL``; ``set_local_poses`` / ``set_local_scales`` flushes local→world - immediately unless a :meth:`change_block` is active. + immediately unless a :meth:`change_block` is active. By default, + ``change_block`` flushes pending world matrices on exit but leaves local + matrices lazy; callers can override this with + ``update_world_matrices`` / ``update_local_matrices``. If the user interleaves both setters on the same view within a single frame, the second setter flushes the first's stale data before writing. This is correct but incurs an extra kernel launch -- a one-time warning @@ -153,6 +158,8 @@ def __init__( self._dirty: _DirtyFlag = _DirtyFlag.NONE self._warned_interleaved_set: bool = False self._change_block_depth: int = 0 + self._change_block_update_world_matrices: bool = True + self._change_block_update_local_matrices: bool = False # Selection (single RW covering both world + local matrix). self._sel = None @@ -184,31 +191,71 @@ def device(self) -> str: return self._device @contextlib.contextmanager - def change_block(self): - """Batch Fabric writes and flush derived matrices on outermost exit. + def change_block(self, update_world_matrices: bool = True, update_local_matrices: bool = False): + """Batch Fabric writes and configure derived-matrix flushes. - Local pose/scale setters normally update cached Fabric world matrices - immediately for renderer/FSD consumers. Use this block when applying - several local edits to the same view so the local→world recompute runs - once after the final edit. + Args: + update_world_matrices: When True (default), pending local writes are + propagated to cached renderer-facing ``omni:fabric:worldMatrix`` + before the outermost block exits. Set False when the caller + knows no downstream consumer will read world matrices before an + explicit FrameView world getter or later sync. + update_local_matrices: When True, pending world writes are propagated + to ``omni:fabric:localMatrix`` before the outermost block exits. + The default is False because Fabric render delegates consume + cached world matrices, and no known downstream render path reads + local matrices directly. + + Nested blocks inherit the outermost block's policy. This lets callers + reliably suppress derived-matrix updates around helper code that may + also use ``change_block()`` with default arguments. """ + if self._change_block_depth == 0: + self._change_block_update_world_matrices = update_world_matrices + self._change_block_update_local_matrices = update_local_matrices + self._change_block_depth += 1 try: yield self finally: self._change_block_depth -= 1 if self._change_block_depth == 0: - self._flush_deferred_local_writes() + update_world = self._change_block_update_world_matrices + update_local = self._change_block_update_local_matrices + self._change_block_update_world_matrices = True + self._change_block_update_local_matrices = False + self._flush_deferred_writes( + update_world_matrices=update_world, + update_local_matrices=update_local, + ) + # Nested blocks inherit the outermost block's flush policy. - def changeBlock(self): - """Alias for :meth:`change_block`.""" - return self.change_block() + def changeBlock( + self, + updateWorldMatrices: bool = True, + updateLocalMatrices: bool = False, + ): + """CamelCase alias for :meth:`change_block`.""" + return self.change_block( + update_world_matrices=updateWorldMatrices, + update_local_matrices=updateLocalMatrices, + ) - def _flush_deferred_local_writes(self) -> None: - """Flush pending local→world recompute unless still inside a change block.""" + def _flush_deferred_writes(self, update_world_matrices: bool = True, update_local_matrices: bool = False) -> None: + """Flush pending derived matrices unless still inside a change block.""" if self._change_block_depth > 0: return - self._sync_world_from_local_if_dirty() + if update_world_matrices: + self._sync_world_from_local_if_dirty() + if update_local_matrices: + self._sync_local_from_world_if_dirty() + + def _flush_deferred_local_writes(self) -> None: + """Flush local→world recompute using the active change-block policy.""" + self._flush_deferred_writes( + update_world_matrices=self._change_block_update_world_matrices, + update_local_matrices=False, + ) @property def prims(self) -> list: 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 f81881ba0540..ccc4e396e40e 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 @@ -259,6 +259,19 @@ def _read_fabric_world_matrix_scale(view, prim_index=0): ) +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_poses_updates_renderer_facing_fabric_world_matrix(device, view_factory): """Local pose writes must update cached Fabric worldMatrix immediately. @@ -359,6 +372,113 @@ def counted_recompute(): torch.testing.assert_close(cached_scale, expected_scale, atol=1e-5, rtol=0) +@pytest.mark.parametrize("device", ["cuda:0"]) +def test_change_block_can_defer_world_matrix_update(device, view_factory, monkeypatch): + """Callers can suppress local→world flushing when no world consumer needs it.""" + bundle = view_factory(num_envs=1, device=device) + view = bundle.view + + view.get_world_poses() + assert view._dirty.name == "NONE" + initial_cached_world = _read_fabric_world_matrix_translation(view) + + calls = 0 + original_recompute = view._recompute_world_from_local + + def counted_recompute(): + nonlocal calls + calls += 1 + original_recompute() + + monkeypatch.setattr(view, "_recompute_world_from_local", counted_recompute) + + 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) + + with view.change_block(update_world_matrices=False): + view.set_local_poses(translations=new_local_pos) + + assert calls == 0 + assert view._dirty.name == "WORLD" + cached_world = _read_fabric_world_matrix_translation(view) + torch.testing.assert_close(cached_world, initial_cached_world, atol=1e-5, rtol=0) + + expected_world = torch.tensor([[1.0, 2.0, 4.0]], dtype=torch.float32, device=device) + world_pos, _ = view.get_world_poses() + assert calls == 1 + torch.testing.assert_close(torch.as_tensor(world_pos, device=device), expected_world, atol=1e-5, rtol=0) + + +@pytest.mark.parametrize("device", ["cuda:0"]) +def test_change_block_can_flush_local_matrix_update(device, view_factory, monkeypatch): + """Callers can request world→local flushing when a localMatrix consumer needs it.""" + bundle = view_factory(num_envs=1, device=device) + view = bundle.view + + view.get_world_poses() + assert view._dirty.name == "NONE" + + calls = 0 + original_sync_local = view._sync_local_from_world + + def counted_sync_local(indices_wp): + nonlocal calls + calls += 1 + original_sync_local(indices_wp) + + monkeypatch.setattr(view, "_sync_local_from_world", counted_sync_local) + + new_world_pos = wp.zeros((1, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=1, inputs=[new_world_pos, 1.0, 2.0, 4.0], device=device) + + with view.change_block(update_local_matrices=True): + view.set_world_poses(positions=new_world_pos) + assert view._dirty.name == "LOCAL" + assert calls == 0 + + assert calls == 1 + assert view._dirty.name == "NONE" + + # Parent is at (0, 0, 1), so local translation should be world - parent. + 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", ["cuda:0"]) +def test_change_block_leaves_local_matrix_lazy_by_default(device, view_factory, monkeypatch): + """Default block policy updates world matrices, but not local matrices.""" + bundle = view_factory(num_envs=1, device=device) + view = bundle.view + + view.get_world_poses() + assert view._dirty.name == "NONE" + + calls = 0 + original_sync_local = view._sync_local_from_world + + def counted_sync_local(indices_wp): + nonlocal calls + calls += 1 + original_sync_local(indices_wp) + + monkeypatch.setattr(view, "_sync_local_from_world", counted_sync_local) + + new_world_pos = wp.zeros((1, 3), dtype=wp.float32, device=device) + wp.launch(kernel=_fill_position, dim=1, inputs=[new_world_pos, 1.0, 2.0, 4.0], device=device) + + with view.change_block(): + view.set_world_poses(positions=new_world_pos) + + assert calls == 0 + assert view._dirty.name == "LOCAL" + + expected_local = torch.tensor([[1.0, 2.0, 3.0]], dtype=torch.float32, device=device) + local_pos, _ = view.get_local_poses() + assert calls == 1 + torch.testing.assert_close(torch.as_tensor(local_pos, device=device), expected_local, atol=1e-5, rtol=0) + + @pytest.mark.parametrize("device", ["cuda:0"]) def test_set_local_via_fabric_path(device, view_factory): """Exercise the Fabric-native set_local_poses path. From ce9ab26ecf744a7275b73305f6c90b472a66b9e2 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Fri, 5 Jun 2026 15:27:34 +0000 Subject: [PATCH 37/54] Restore three-selection RO/RW design, drop eager world<->local flushes Restores PR #5728's three-selection layout (`_trans_sel_ro`, `_world_sel_rw`, `_local_sel_rw`) with asymmetric Fabric access flags on `worldMatrix` and `localMatrix`. Those flags are what protect the user's write from being clobbered by Kit's per-tick `IFabricHierarchy.update_world_xforms`: * On `set_world_poses` (via `_world_sel_rw`, `localMatrix=RO`), Fabric does not recompute world from local -- the user's worldMatrix write survives until the renderer reads it. * On `set_local_poses` (via `_local_sel_rw`, `worldMatrix=RO`), Fabric recomputes world from the new local on the next tick -- the renderer reads the correct world. A single combined `worldMatrix=RW, localMatrix=RW` selection (the recent design on this branch) removed that protection. Fabric saw both attributes as user-authored and fell back to the hierarchy's canonical direction (local -> world), recomputing world from a stale local and silently overwriting the user's world write. That was the failure mode behind the `test_output_equal_to_usdcamera` regression and any other Camera + RTX path that drives world poses through Fabric. With the RO/RW protection back in place, the eager world<->local flushes introduced by commits "fix: flush Fabric world matrices after local writes" and the follow-up "set_world_poses eager local sync" are no longer needed and are removed. The `change_block` context manager and its companion helpers existed only to batch those eager flushes; with the flushes gone, the API has nothing to defer and is removed from both `BaseFrameView` and `FabricFrameView`. Class docstring now spells out the load-bearing role of the RO/RW layout so a future refactor doesn't reintroduce the single-selection shape. Tests: * Removed `test_set_local_*_updates_renderer_facing_fabric_world_matrix` (asserted an eager-update contract that the lazy design deliberately does not hold; correctness across the next render tick is provided by the RO/RW protection, not by an extra Warp kernel). * Removed the four `test_change_block_*` tests; the API is gone. * Inverted `test_interleaved_set_emits_no_warning` back to `test_interleaved_set_emits_warning`, restored `_dirty == LOCAL` assertions in `test_world_scales_roundtrip` and the symmetric `WORLD` assertion in `test_local_scales_roundtrip`, and updated `test_multi_view_per_view_dirty_isolation` to expect the lazy cross-view behavior. * Adapted `test_prepare_for_reuse_detects_topology_change` and `test_fabric_rebuild_after_topology_change` to poll/rebuild all three selections. Verified: * `pytest test_ray_caster_camera.py::test_output_equal_to_usdcamera` passes. * `pytest test_ray_caster_camera.py::test_output_equal_to_usd_camera_when_intrinsics_set` 4/4 pass. * `pytest test_views_xform_prim_fabric.py` 71 passed, 3 skipped (cuda:1). * `./isaaclab.sh -f` clean. Net diff -198 lines. --- .../isaaclab/sim/views/base_frame_view.py | 32 -- .../sim/views/fabric_frame_view.py | 357 +++++++++--------- .../test/sim/test_views_xform_prim_fabric.py | 273 +++----------- 3 files changed, 232 insertions(+), 430 deletions(-) diff --git a/source/isaaclab/isaaclab/sim/views/base_frame_view.py b/source/isaaclab/isaaclab/sim/views/base_frame_view.py index 31ea771a4132..ee5e313d7f3c 100644 --- a/source/isaaclab/isaaclab/sim/views/base_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/base_frame_view.py @@ -8,7 +8,6 @@ from __future__ import annotations import abc -import contextlib import warnings import warp as wp @@ -39,37 +38,6 @@ def device(self) -> str: """Device where arrays are allocated (``"cpu"`` or ``"cuda:0"``).""" ... - @contextlib.contextmanager - def change_block(self, update_world_matrices: bool = True, update_local_matrices: bool = False): - """Batch multiple transform writes into one logical change. - - Backends may use this context to defer expensive derived-state updates - until the outermost block exits. The default implementation is a no-op - so callers can use it with every FrameView backend. - - Args: - update_world_matrices: Whether derived world matrices should be - updated before the outermost block exits. Backends may ignore - this option when they do not maintain separate cached world - matrices. - update_local_matrices: Whether derived local matrices should be - updated before the outermost block exits. Backends may ignore - this option when they do not maintain separate cached local - matrices. - """ - yield self - - def changeBlock( - self, - updateWorldMatrices: bool = True, - updateLocalMatrices: bool = False, - ): - """CamelCase alias for :meth:`change_block`.""" - return self.change_block( - update_world_matrices=updateWorldMatrices, - update_local_matrices=updateLocalMatrices, - ) - @abc.abstractmethod def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: """Get world-space positions and orientations for prims in the view. 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 8aa1c788ea8f..87f98f0c51b4 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,6 @@ from __future__ import annotations -import contextlib import enum import logging @@ -75,41 +74,72 @@ class FabricFrameView(BaseFrameView): 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. - * **World ↔ local consistency.** Getters are lazy after world writes: - ``set_world_poses`` or ``set_world_scales`` recompute local matrices only - when ``get_local_poses`` (or ``get_local_scales``) is called. Local - writes update renderer-facing ``omni:fabric:worldMatrix`` eagerly so - Fabric render delegates see the new pose/scale without requiring an Isaac - Lab world getter. The default is intentionally asymmetric: render - delegates consume world matrices, while no known downstream render path - consumes local matrices directly. Wrap multiple writes in - :meth:`change_block` to defer and configure derived-matrix updates. + * **World <-> local consistency (lazy, RO/RW-protected).** After + ``set_world_poses`` / ``set_world_scales``, local matrices are + recomputed lazily when ``get_local_poses`` / ``get_local_scales`` is + called; after ``set_local_poses`` / ``set_local_scales``, world + matrices are recomputed lazily when ``get_world_poses`` / + ``get_world_scales`` is called. Both directions stay in sync without + round-tripping through USD. The renderer/FSD reads cached + ``omni:fabric:worldMatrix`` *between* user writes and the next + ``IFabricHierarchy.update_world_xforms`` tick (which Kit runs as part + of the render path). Correctness across that tick depends on the + RO/RW selection layout below -- a single combined ``ReadWrite`` + selection is **not** safe; see the next bullet. + * **Three selections with asymmetric RO/RW access (load-bearing).** + The view holds three persistent Fabric selections: + + .. code-block:: text + + _trans_sel_ro : worldMatrix=RO, localMatrix=RO (reads + sync helpers) + _world_sel_rw : worldMatrix=RW, localMatrix=RO (set_world_poses / _scales) + _local_sel_rw : worldMatrix=RO, localMatrix=RW (set_local_poses / _scales) + + The access flags tell Fabric which attribute the user is *authoring* + for a given operation. When ``IFabricHierarchy.update_world_xforms`` + runs (as part of Kit's render tick), it respects those flags and does + not recompute an attribute marked RO on the most-recently-written + selection. Concretely: + + - After ``set_world_poses`` (via ``_world_sel_rw``): Fabric sees + ``localMatrix=RO`` -> does not recompute world from local -> the + user's worldMatrix write survives until the renderer reads it. + - After ``set_local_poses`` (via ``_local_sel_rw``): Fabric sees + ``localMatrix=RW`` -> recomputes world from the new local on the + next tick -> the renderer sees the correct world. + + Using a single combined ``worldMatrix=RW, localMatrix=RW`` selection + removes this protection. Fabric then sees both attributes as + user-authored and falls back to the hierarchy's canonical direction + (local -> world), recomputing world from a stale local and clobbering + the user's worldMatrix write. This was the failure mode that broke + the Camera + RTX renderer path when the camera prim sat in a + hierarchy traversed during the next render tick. * **Dirty-flag invariant.** The ``_dirty`` enum is one of ``NONE``, ``WORLD``, or ``LOCAL`` -- mutually exclusive by construction. ``set_world_poses`` / ``set_world_scales`` sets ``_dirty = LOCAL``; - ``set_local_poses`` / ``set_local_scales`` flushes local→world - immediately unless a :meth:`change_block` is active. By default, - ``change_block`` flushes pending world matrices on exit but leaves local - matrices lazy; callers can override this with - ``update_world_matrices`` / ``update_local_matrices``. + ``set_local_poses`` / ``set_local_scales`` sets ``_dirty = WORLD``. If the user interleaves both setters on the same view within a single frame, the second setter flushes the first's stale data before writing. - This is correct but incurs an extra kernel launch -- a one-time warning - is logged when this happens. + This is correct but incurs an extra kernel launch -- a one-time + warning is logged when this happens. * **Topology-adaptive.** Fabric topology changes are detected on each - access; the view rebuilds its internal mapping automatically and no - manual refresh is required. Steady-state overhead is negligible. + access via per-selection ``PrepareForReuse()`` polls; the affected + indexed arrays rebuild automatically and no manual refresh is required. + Steady-state overhead is negligible. Performance note: The fast path assumes the user calls **either** ``set_world_poses`` **or** ``set_local_poses`` exclusively within a frame (not both). - In that case, setters are O(1) kernel launches with no synchronization - overhead beyond the single ``wp.synchronize()``; getters lazily flush - the opposite direction only when actually needed. + In that case, setters are O(1) kernel launches with no + synchronization overhead beyond the single ``wp.synchronize()``; + getters lazily flush the opposite direction only when actually + needed. - Interleaving both setters on different index subsets within the same - frame is supported and correct, but triggers an extra flush kernel - per transition. A warning is emitted once per view instance. + Interleaving both setters on different index subsets within the + same frame is supported and correct, but triggers an extra flush + kernel per transition. A warning is emitted once per view + instance. Pose getters return :class:`~isaaclab.utils.warp.ProxyArray`; setters accept :class:`wp.array`. @@ -131,7 +161,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. @@ -157,21 +187,30 @@ def __init__( # Per-view (not per-stage) so concurrent views on the same stage don't interfere. self._dirty: _DirtyFlag = _DirtyFlag.NONE self._warned_interleaved_set: bool = False - self._change_block_depth: int = 0 - self._change_block_update_world_matrices: bool = True - self._change_block_update_local_matrices: bool = False - # Selection (single RW covering both world + local matrix). - self._sel = None - - # Index arrays (view-side indices and view->fabric mapping). + # Three persistent Fabric selections with asymmetric access flags. + # See the class docstring "Three selections with asymmetric RO/RW access" + # bullet for why this layout is required. + self._trans_sel_ro = None + self._world_sel_rw = None + self._local_sel_rw = None + + # Index arrays (view-side indices and per-selection view->fabric mappings). + # Each selection's ``GetPaths()`` ordering is independent, so view->fabric + # is cached per selection -- sharing would silently corrupt indexed arrays + # whose selection didn't fire ``PrepareForReuse`` on the same frame. self._view_indices: wp.array | None = None - self._fabric_indices: wp.array | None = None - - # Indexed fabric arrays. - self._world_ifa = None - self._local_ifa = None - self._parent_world_ifa = None + self._trans_ro_fabric_indices: wp.array | None = None + self._world_rw_fabric_indices: wp.array | None = None + self._local_rw_fabric_indices: wp.array | None = None + self._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._world_ifa_rw = None + self._local_ifa_rw = None + self._parent_world_ifa_ro = None # Sentinel passed to compose/decompose kernels for unused slots. # Kernels gate per-row access on ``shape[0] > 0``, so (0, 0) suffices. @@ -190,73 +229,6 @@ def device(self) -> str: """Device where arrays are allocated (cpu or cuda).""" return self._device - @contextlib.contextmanager - def change_block(self, update_world_matrices: bool = True, update_local_matrices: bool = False): - """Batch Fabric writes and configure derived-matrix flushes. - - Args: - update_world_matrices: When True (default), pending local writes are - propagated to cached renderer-facing ``omni:fabric:worldMatrix`` - before the outermost block exits. Set False when the caller - knows no downstream consumer will read world matrices before an - explicit FrameView world getter or later sync. - update_local_matrices: When True, pending world writes are propagated - to ``omni:fabric:localMatrix`` before the outermost block exits. - The default is False because Fabric render delegates consume - cached world matrices, and no known downstream render path reads - local matrices directly. - - Nested blocks inherit the outermost block's policy. This lets callers - reliably suppress derived-matrix updates around helper code that may - also use ``change_block()`` with default arguments. - """ - if self._change_block_depth == 0: - self._change_block_update_world_matrices = update_world_matrices - self._change_block_update_local_matrices = update_local_matrices - - self._change_block_depth += 1 - try: - yield self - finally: - self._change_block_depth -= 1 - if self._change_block_depth == 0: - update_world = self._change_block_update_world_matrices - update_local = self._change_block_update_local_matrices - self._change_block_update_world_matrices = True - self._change_block_update_local_matrices = False - self._flush_deferred_writes( - update_world_matrices=update_world, - update_local_matrices=update_local, - ) - # Nested blocks inherit the outermost block's flush policy. - - def changeBlock( - self, - updateWorldMatrices: bool = True, - updateLocalMatrices: bool = False, - ): - """CamelCase alias for :meth:`change_block`.""" - return self.change_block( - update_world_matrices=updateWorldMatrices, - update_local_matrices=updateLocalMatrices, - ) - - def _flush_deferred_writes(self, update_world_matrices: bool = True, update_local_matrices: bool = False) -> None: - """Flush pending derived matrices unless still inside a change block.""" - if self._change_block_depth > 0: - return - if update_world_matrices: - self._sync_world_from_local_if_dirty() - if update_local_matrices: - self._sync_local_from_world_if_dirty() - - def _flush_deferred_local_writes(self) -> None: - """Flush local→world recompute using the active change-block policy.""" - self._flush_deferred_writes( - update_world_matrices=self._change_block_update_world_matrices, - update_local_matrices=False, - ) - @property def prims(self) -> list: return self._usd_view.prims @@ -276,7 +248,7 @@ def set_visibility(self, visibility, indices=None): self._usd_view.set_visibility(visibility, indices) # ------------------------------------------------------------------ - # World poses — Fabric-accelerated or USD fallback + # World poses -- Fabric-accelerated or USD fallback # ------------------------------------------------------------------ def set_world_poses(self, positions=None, orientations=None, indices=None): @@ -307,7 +279,7 @@ def set_world_poses(self, positions=None, orientations=None, indices=None): kernel=fabric_utils.compose_indexed_fabric_transforms, dim=indices_wp.shape[0], inputs=[ - self._get_world_array(), + self._get_world_rw_array(), positions_wp, orientations_wp, self._fabric_empty_2d_array_sentinel, @@ -321,7 +293,10 @@ def set_world_poses(self, positions=None, orientations=None, indices=None): wp.synchronize() # World was just written -- mark local poses as stale so the next - # get_local_poses recomputes them lazily. + # get_local_poses recomputes them lazily. No eager local recompute is + # needed: the worldMatrix write is protected by ``_world_sel_rw`` having + # ``localMatrix=RO``, so the next ``update_world_xforms`` tick will not + # overwrite world from a stale local. self._dirty = _DirtyFlag.LOCAL def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: @@ -356,7 +331,7 @@ def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, kernel=fabric_utils.decompose_indexed_fabric_transforms, dim=count, inputs=[ - self._get_world_array(), + self._get_world_ro_array(), positions_wp, orientations_wp, self._fabric_empty_2d_array_sentinel, @@ -403,7 +378,7 @@ def set_local_poses(self, translations=None, orientations=None, indices=None): kernel=fabric_utils.compose_indexed_fabric_transforms, dim=indices_wp.shape[0], inputs=[ - self._get_local_array(), + self._get_local_rw_array(), translations_wp, orientations_wp, self._fabric_empty_2d_array_sentinel, @@ -416,11 +391,11 @@ def set_local_poses(self, translations=None, orientations=None, indices=None): ) wp.synchronize() - # Keep renderer-facing world matrices current. Inside change_block(), - # defer the recompute until the outermost block exits so pose+scale - # edits only pay one local→world kernel launch. + # Local was just written -- mark world matrices stale. No eager world + # recompute: the ``_local_sel_rw`` selection has ``worldMatrix=RO``, so + # Kit's next ``update_world_xforms`` tick will recompute world from the + # new local automatically, and the renderer will read the correct world. self._dirty = _DirtyFlag.WORLD - self._flush_deferred_local_writes() def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: """Return (translations, orientations) in parent-local frame. @@ -454,7 +429,7 @@ def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, kernel=fabric_utils.decompose_indexed_fabric_transforms, dim=count, inputs=[ - self._get_local_array(), + self._get_local_ro_array(), translations_wp, orientations_wp, self._fabric_empty_2d_array_sentinel, @@ -491,7 +466,7 @@ def set_world_scales(self, scales, indices=None): kernel=fabric_utils.compose_indexed_fabric_transforms, dim=indices_wp.shape[0], inputs=[ - self._get_world_array(), + self._get_world_rw_array(), self._fabric_empty_2d_array_sentinel, self._fabric_empty_2d_array_sentinel, scales_wp, @@ -504,7 +479,8 @@ def set_world_scales(self, scales, indices=None): ) wp.synchronize() - # World was just written -- mark local poses as stale. + # World was just written -- mark local poses as stale. See set_world_poses + # for why no eager local recompute is needed. self._dirty = _DirtyFlag.LOCAL def get_world_scales(self, indices=None): @@ -538,7 +514,7 @@ def get_world_scales(self, indices=None): kernel=fabric_utils.decompose_indexed_fabric_transforms, dim=count, inputs=[ - self._get_world_array(), + self._get_world_ro_array(), self._fabric_empty_2d_array_sentinel, self._fabric_empty_2d_array_sentinel, scales_wp, @@ -571,7 +547,7 @@ def set_local_scales(self, scales, indices=None): kernel=fabric_utils.compose_indexed_fabric_transforms, dim=indices_wp.shape[0], inputs=[ - self._get_local_array(), + self._get_local_rw_array(), self._fabric_empty_2d_array_sentinel, self._fabric_empty_2d_array_sentinel, scales_wp, @@ -584,11 +560,9 @@ def set_local_scales(self, scales, indices=None): ) wp.synchronize() - # Keep renderer-facing world matrices current. Inside change_block(), - # defer the recompute until the outermost block exits so pose+scale - # edits only pay one local→world kernel launch. + # Local was just written -- mark world matrices stale. See set_local_poses + # for why no eager world recompute is needed. self._dirty = _DirtyFlag.WORLD - self._flush_deferred_local_writes() def get_local_scales(self, indices=None): """Return per-prim (sx, sy, sz) scales extracted from local matrix. @@ -621,7 +595,7 @@ def get_local_scales(self, indices=None): kernel=fabric_utils.decompose_indexed_fabric_transforms, dim=count, inputs=[ - self._get_local_array(), + self._get_local_ro_array(), self._fabric_empty_2d_array_sentinel, self._fabric_empty_2d_array_sentinel, scales_wp, @@ -665,14 +639,17 @@ def _recompute_world_from_local(self) -> None: local+world matrices we just authored. Instead we fire a Warp kernel that does the multiply per child, leaving the Fabric-side localMatrix untouched. """ - self._refresh_if_needed() + # Refresh trans_sel_ro once, then read _local_ifa_ro and _parent_world_ifa_ro + # directly to avoid calling PrepareForReuse twice on the same selection. + if self._trans_sel_ro.PrepareForReuse() or self._parent_world_ifa_ro is None: + self._rebuild_trans_ro_arrays() wp.launch( kernel=fabric_utils.update_indexed_world_matrix_from_local, dim=self.count, inputs=[ - self._local_ifa, - self._parent_world_ifa, - self._world_ifa, + self._local_ifa_ro, + self._parent_world_ifa_ro, + self._get_world_rw_array(), self._view_indices, ], device=self._device, @@ -685,14 +662,16 @@ def _sync_local_from_world(self, indices_wp: wp.array) -> None: Called after ``set_world_poses`` so that subsequent ``get_local_poses`` returns values consistent with the just-written world poses. """ - self._refresh_if_needed() + # Refresh trans_sel_ro once; _world_ifa_ro and _parent_world_ifa_ro share it. + if self._trans_sel_ro.PrepareForReuse() or self._parent_world_ifa_ro is None: + self._rebuild_trans_ro_arrays() wp.launch( kernel=fabric_utils.update_indexed_local_matrix_from_world, dim=indices_wp.shape[0], inputs=[ - self._world_ifa, - self._parent_world_ifa, - self._local_ifa, + self._world_ifa_ro, + self._parent_world_ifa_ro, + self._get_local_rw_array(), indices_wp, ], device=self._device, @@ -710,25 +689,52 @@ def _sync_local_from_world_if_dirty(self) -> None: # Internal -- selection accessors with on-demand index rebuild # ------------------------------------------------------------------ - def _refresh_if_needed(self): - """Rebuild indexed arrays if the selection's prim set changed.""" - if self._sel.PrepareForReuse() or self._world_ifa is None: - self._fabric_indices = self._compute_fabric_indices(self._sel) - self._world_ifa = self._build_indexed_array(self._sel, self._WORLD_MATRIX_NAME, self._fabric_indices) - self._local_ifa = self._build_indexed_array(self._sel, self._LOCAL_MATRIX_NAME, self._fabric_indices) - self._parent_world_ifa = self._build_parent_indexed_array(self._sel) + def _get_world_ro_array(self): + if self._trans_sel_ro.PrepareForReuse(): + self._rebuild_trans_ro_arrays() + return self._world_ifa_ro + + def _get_local_ro_array(self): + if self._trans_sel_ro.PrepareForReuse(): + self._rebuild_trans_ro_arrays() + return self._local_ifa_ro + + def _get_world_rw_array(self): + if self._world_sel_rw.PrepareForReuse(): + self._world_rw_fabric_indices = self._compute_fabric_indices(self._world_sel_rw) + self._world_ifa_rw = self._build_indexed_array( + self._world_sel_rw, self._WORLD_MATRIX_NAME, self._world_rw_fabric_indices + ) + return self._world_ifa_rw + + def _get_local_rw_array(self): + if self._local_sel_rw.PrepareForReuse(): + self._local_rw_fabric_indices = self._compute_fabric_indices(self._local_sel_rw) + self._local_ifa_rw = self._build_indexed_array( + self._local_sel_rw, self._LOCAL_MATRIX_NAME, self._local_rw_fabric_indices + ) + return self._local_ifa_rw - def _get_world_array(self): - self._refresh_if_needed() - return self._world_ifa + def _get_parent_world_ro_array(self): + # Built and refreshed alongside the trans_ro selection (parents share that selection). + if self._parent_world_ifa_ro is None or self._trans_sel_ro.PrepareForReuse(): + self._rebuild_trans_ro_arrays() + return self._parent_world_ifa_ro - def _get_local_array(self): - self._refresh_if_needed() - return self._local_ifa + def _rebuild_trans_ro_arrays(self) -> None: + """Rebuild the trans_ro indices and the three indexed arrays that depend on them. - def _get_parent_world_array(self): - self._refresh_if_needed() - return self._parent_world_ifa + ``_world_ifa_ro``, ``_local_ifa_ro`` and ``_parent_world_ifa_ro`` are all + keyed off the ``trans_sel_ro`` path ordering, so they are refreshed together. + """ + self._trans_ro_fabric_indices = self._compute_fabric_indices(self._trans_sel_ro) + self._world_ifa_ro = self._build_indexed_array( + self._trans_sel_ro, self._WORLD_MATRIX_NAME, self._trans_ro_fabric_indices + ) + self._local_ifa_ro = self._build_indexed_array( + self._trans_sel_ro, self._LOCAL_MATRIX_NAME, self._trans_ro_fabric_indices + ) + self._parent_world_ifa_ro = self._build_parent_indexed_array(self._trans_sel_ro) # ------------------------------------------------------------------ # Internal -- index computation @@ -832,24 +838,41 @@ def _initialize_fabric(self) -> None: rt_xformable.SetLocalXformFromUsd() rt_xformable.SetWorldXformFromUsd() - # Single RW selection covering both matrices. - # TODO: Benchmark RO vs RW selection split -- separate RO selections could reduce - # lock contention under concurrent Fabric access, but current usage is single-threaded. + # Three persistent selections with asymmetric access flags. This layout + # is load-bearing for correctness against Kit's ``update_world_xforms`` + # hierarchy update; see the class docstring for the why. 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 = self._stage.SelectPrims(require_attrs=[wm_rw, lm_rw], device=self._device, want_paths=True) + self._trans_sel_ro = self._stage.SelectPrims(require_attrs=[wm_ro, lm_ro], device=self._device, want_paths=True) + self._world_sel_rw = self._stage.SelectPrims(require_attrs=[wm_rw, lm_ro], device=self._device, want_paths=True) + self._local_sel_rw = self._stage.SelectPrims(require_attrs=[wm_ro, lm_rw], device=self._device, want_paths=True) - # Build the view-side indices array (just [0..count-1]) and a + # Build the view-side indices array (just [0..count-1]) and a per-selection # view->fabric mapping (selections do not guarantee a shared path ordering). self._view_indices = wp.array(list(range(self.count)), dtype=wp.uint32, device=self._device) - self._fabric_indices = self._compute_fabric_indices(self._sel) + self._trans_ro_fabric_indices = self._compute_fabric_indices(self._trans_sel_ro) + self._world_rw_fabric_indices = self._compute_fabric_indices(self._world_sel_rw) + self._local_rw_fabric_indices = self._compute_fabric_indices(self._local_sel_rw) - # Indexed fabric arrays per attribute. - self._world_ifa = self._build_indexed_array(self._sel, self._WORLD_MATRIX_NAME, self._fabric_indices) - self._local_ifa = self._build_indexed_array(self._sel, self._LOCAL_MATRIX_NAME, self._fabric_indices) - self._parent_world_ifa = self._build_parent_indexed_array(self._sel) + # Indexed fabric arrays per (selection x attribute). + self._world_ifa_ro = self._build_indexed_array( + self._trans_sel_ro, self._WORLD_MATRIX_NAME, self._trans_ro_fabric_indices + ) + self._local_ifa_ro = self._build_indexed_array( + self._trans_sel_ro, self._LOCAL_MATRIX_NAME, self._trans_ro_fabric_indices + ) + self._world_ifa_rw = self._build_indexed_array( + self._world_sel_rw, self._WORLD_MATRIX_NAME, self._world_rw_fabric_indices + ) + self._local_ifa_rw = self._build_indexed_array( + self._local_sel_rw, self._LOCAL_MATRIX_NAME, self._local_rw_fabric_indices + ) + self._parent_world_ifa_ro = self._build_parent_indexed_array(self._trans_sel_ro) # Pre-allocated reusable output buffers (world + local + scales). self._fabric_positions_buf = wp.zeros((self.count, 3), dtype=wp.float32, device=self._device) @@ -883,17 +906,15 @@ def _sync_fabric_from_usd_initial(self) -> None: """ # --- Children --- # Compose child localMatrix from USD-authored local transforms. - # The child world matrix is NOT composed here -- it will be computed - # by ``_recompute_world_from_local()`` at the end of this method as - # ``child_world = child_local * parent_world``, which naturally - # composes scales through the matrix multiplication. + # The child world matrix is computed at the end via + # ``_recompute_world_from_local()`` as ``child_world = parent_world * child_local``. 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, + self._local_ifa_rw, _to_float32_2d(local_pos_ta.warp), _to_float32_2d(local_ori_ta.warp), _to_float32_2d(scales_wp), @@ -936,7 +957,7 @@ def _sync_fabric_from_usd_initial(self) -> None: warned_shear = True logger.warning( "FabricFrameView: parent prim '%s' has a sheared/skewed world " - "transform. TRS decomposition (used by scale getters and world↔local " + "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, @@ -953,10 +974,10 @@ def _sync_fabric_from_usd_initial(self) -> None: 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) # Compose worldMatrix for parents (use a one-shot indexed array against - # ``world_sel_rw`` keyed on the unique parent paths). + # ``_world_sel_rw`` keyed on the unique parent paths). parent_world_rw = wp.indexedfabricarray( - fa=wp.fabricarray(self._sel, self._WORLD_MATRIX_NAME), - indices=self._compute_fabric_indices_for(self._sel, unique_parent_paths), + fa=wp.fabricarray(self._world_sel_rw, self._WORLD_MATRIX_NAME), + indices=self._compute_fabric_indices_for(self._world_sel_rw, unique_parent_paths), ) wp.launch( kernel=fabric_utils.compose_indexed_fabric_transforms, 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 ccc4e396e40e..a08432c9366d 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 @@ -201,8 +201,17 @@ def test_fabric_rebuild_after_topology_change(device, view_factory): wp.launch(kernel=_fill_position, dim=2, inputs=[initial, 1.0, 2.0, 3.0], device=device) view.set_world_poses(positions=initial) - # Simulate topology change: force rebuild of the selection's indexed arrays. - view._refresh_if_needed() + # Simulate topology change: recompute per-selection fabric indices and rebuild + # every indexed array, mirroring the lazy paths in the ``_get_*_array`` accessors. + view._rebuild_trans_ro_arrays() + view._world_rw_fabric_indices = view._compute_fabric_indices(view._world_sel_rw) + view._world_ifa_rw = view._build_indexed_array( + view._world_sel_rw, view._WORLD_MATRIX_NAME, view._world_rw_fabric_indices + ) + view._local_rw_fabric_indices = view._compute_fabric_indices(view._local_sel_rw) + view._local_ifa_rw = view._build_indexed_array( + view._local_sel_rw, view._LOCAL_MATRIX_NAME, view._local_rw_fabric_indices + ) # Trigger another write through the rebuilt arrays. new = wp.zeros((2, 3), dtype=wp.float32, device=device) @@ -225,10 +234,11 @@ def test_prepare_for_reuse_detects_topology_change(device, view_factory): view = bundle.view view.get_world_poses() # trigger Fabric init - assert view._sel is not None, "selection not initialized" - result = view._sel.PrepareForReuse() - assert isinstance(result, bool), f"PrepareForReuse should return bool, got {type(result)}" - assert not result, "PrepareForReuse should return False when no topology change" + assert view._trans_sel_ro is not None, "trans_sel_ro selection not initialized" + for selection in (view._trans_sel_ro, view._world_sel_rw, view._local_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): @@ -272,213 +282,6 @@ def _read_fabric_local_matrix_translation(view, prim_index=0): ) -@pytest.mark.parametrize("device", ["cuda:0"]) -def test_set_local_poses_updates_renderer_facing_fabric_world_matrix(device, view_factory): - """Local pose writes must update cached Fabric worldMatrix immediately. - - The FSD renderer reads Fabric's cached ``omni:fabric:worldMatrix`` directly; - it does not call ``FrameView.get_world_poses()`` to trigger Isaac Lab's lazy - local→world sync. This test intentionally reads the Fabric attribute - directly after ``set_local_poses`` and before any world getter call. - """ - bundle = view_factory(num_envs=1, device=device) - view = bundle.view - - # Initialize Fabric and clear the initial dirty state. - view.get_world_poses() - assert view._dirty.name == "NONE" - - 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) - new_local_ori = wp.from_torch(torch.tensor([[0.0, 0.0, 0.0, 1.0]], dtype=torch.float32, device=device)) - - view.set_local_poses(translations=new_local_pos, orientations=new_local_ori) - - # Parent is at (0, 0, 1), so renderer-facing cached worldMatrix should - # already contain world translation (1, 2, 4) without get_world_poses(). - 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_set_local_scales_updates_renderer_facing_fabric_world_matrix(device, view_factory): - """Local scale writes must update cached Fabric worldMatrix immediately. - - This is the scale analogue of the local-pose renderer/FSD contract: read - cached ``omni:fabric:worldMatrix`` directly after ``set_local_scales`` and - before any FrameView world getter can repair stale state. - """ - bundle = view_factory(num_envs=1, device=device) - view = bundle.view - - # Initialize Fabric and clear the initial dirty state. - view.get_world_poses() - assert view._dirty.name == "NONE" - - 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) - - view.set_local_scales(new_scales) - - 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) - - -@pytest.mark.parametrize("device", ["cuda:0"]) -def test_change_block_batches_local_world_matrix_update(device, view_factory, monkeypatch): - """Local pose+scale writes inside change_block flush worldMatrix once on exit.""" - bundle = view_factory(num_envs=1, device=device) - view = bundle.view - - view.get_world_poses() - assert view._dirty.name == "NONE" - - calls = 0 - original_recompute = view._recompute_world_from_local - - def counted_recompute(): - nonlocal calls - calls += 1 - original_recompute() - - monkeypatch.setattr(view, "_recompute_world_from_local", counted_recompute) - - 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) - new_local_ori = wp.from_torch(torch.tensor([[0.0, 0.0, 0.0, 1.0]], dtype=torch.float32, 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.change_block(): - view.set_local_poses(translations=new_local_pos, orientations=new_local_ori) - assert view._dirty.name == "WORLD" - assert calls == 0 - - view.set_local_scales(new_scales) - assert view._dirty.name == "WORLD" - assert calls == 0 - - assert calls == 1 - assert view._dirty.name == "NONE" - - 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) - - -@pytest.mark.parametrize("device", ["cuda:0"]) -def test_change_block_can_defer_world_matrix_update(device, view_factory, monkeypatch): - """Callers can suppress local→world flushing when no world consumer needs it.""" - bundle = view_factory(num_envs=1, device=device) - view = bundle.view - - view.get_world_poses() - assert view._dirty.name == "NONE" - initial_cached_world = _read_fabric_world_matrix_translation(view) - - calls = 0 - original_recompute = view._recompute_world_from_local - - def counted_recompute(): - nonlocal calls - calls += 1 - original_recompute() - - monkeypatch.setattr(view, "_recompute_world_from_local", counted_recompute) - - 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) - - with view.change_block(update_world_matrices=False): - view.set_local_poses(translations=new_local_pos) - - assert calls == 0 - assert view._dirty.name == "WORLD" - cached_world = _read_fabric_world_matrix_translation(view) - torch.testing.assert_close(cached_world, initial_cached_world, atol=1e-5, rtol=0) - - expected_world = torch.tensor([[1.0, 2.0, 4.0]], dtype=torch.float32, device=device) - world_pos, _ = view.get_world_poses() - assert calls == 1 - torch.testing.assert_close(torch.as_tensor(world_pos, device=device), expected_world, atol=1e-5, rtol=0) - - -@pytest.mark.parametrize("device", ["cuda:0"]) -def test_change_block_can_flush_local_matrix_update(device, view_factory, monkeypatch): - """Callers can request world→local flushing when a localMatrix consumer needs it.""" - bundle = view_factory(num_envs=1, device=device) - view = bundle.view - - view.get_world_poses() - assert view._dirty.name == "NONE" - - calls = 0 - original_sync_local = view._sync_local_from_world - - def counted_sync_local(indices_wp): - nonlocal calls - calls += 1 - original_sync_local(indices_wp) - - monkeypatch.setattr(view, "_sync_local_from_world", counted_sync_local) - - new_world_pos = wp.zeros((1, 3), dtype=wp.float32, device=device) - wp.launch(kernel=_fill_position, dim=1, inputs=[new_world_pos, 1.0, 2.0, 4.0], device=device) - - with view.change_block(update_local_matrices=True): - view.set_world_poses(positions=new_world_pos) - assert view._dirty.name == "LOCAL" - assert calls == 0 - - assert calls == 1 - assert view._dirty.name == "NONE" - - # Parent is at (0, 0, 1), so local translation should be world - parent. - 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", ["cuda:0"]) -def test_change_block_leaves_local_matrix_lazy_by_default(device, view_factory, monkeypatch): - """Default block policy updates world matrices, but not local matrices.""" - bundle = view_factory(num_envs=1, device=device) - view = bundle.view - - view.get_world_poses() - assert view._dirty.name == "NONE" - - calls = 0 - original_sync_local = view._sync_local_from_world - - def counted_sync_local(indices_wp): - nonlocal calls - calls += 1 - original_sync_local(indices_wp) - - monkeypatch.setattr(view, "_sync_local_from_world", counted_sync_local) - - new_world_pos = wp.zeros((1, 3), dtype=wp.float32, device=device) - wp.launch(kernel=_fill_position, dim=1, inputs=[new_world_pos, 1.0, 2.0, 4.0], device=device) - - with view.change_block(): - view.set_world_poses(positions=new_world_pos) - - assert calls == 0 - assert view._dirty.name == "LOCAL" - - expected_local = torch.tensor([[1.0, 2.0, 3.0]], dtype=torch.float32, device=device) - local_pos, _ = view.get_local_poses() - assert calls == 1 - torch.testing.assert_close(torch.as_tensor(local_pos, device=device), expected_local, atol=1e-5, rtol=0) - - @pytest.mark.parametrize("device", ["cuda:0"]) def test_set_local_via_fabric_path(device, view_factory): """Exercise the Fabric-native set_local_poses path. @@ -541,8 +344,11 @@ def test_local_scales_roundtrip(device, view_factory): wp.launch(kernel=_fill_position, dim=2, inputs=[new_scales, 2.0, 3.0, 4.0], device=device) view.set_local_scales(new_scales) - # Local writes eagerly update renderer-facing world matrices outside change_block(). - assert view._dirty.name == "NONE" + # Local writes mark world matrices stale; the renderer-facing world is + # recomputed lazily by Kit's next ``update_world_xforms`` tick (the + # ``_local_sel_rw`` selection has worldMatrix=RO, so that recompute is + # safe and the lazy design is correct). + assert view._dirty.name == "WORLD" ret_scales = view.get_local_scales() scales_torch = ret_scales.torch @@ -563,7 +369,10 @@ def test_world_scales_roundtrip(device, view_factory): wp.launch(kernel=_fill_position, dim=2, inputs=[new_scales, 5.0, 6.0, 7.0], device=device) view.set_world_scales(new_scales) - # Should have dirtied local + # World writes mark local matrices stale; they are recomputed lazily on the + # next ``get_local_*`` call. The renderer-facing world write is protected + # by ``_world_sel_rw`` having localMatrix=RO, so Kit's next + # ``update_world_xforms`` tick will not clobber it from a stale local. assert view._dirty.name == "LOCAL" ret_scales = view.get_world_scales() @@ -754,20 +563,18 @@ def test_multi_view_per_view_dirty_isolation(device): identity_quat = wp.from_torch(torch.tensor([[0.0, 0.0, 0.0, 1.0]], dtype=torch.float32, device=device)) view_a.set_local_poses(translations=new_local_a, orientations=identity_quat) - # Local writes eagerly update renderer-facing world matrices outside change_block(), - # and they must not dirty another view on the same stage. - assert view_a._dirty.name == "NONE", "set_local_poses should flush its own view outside change_block" + # Local writes are lazy: view A is dirty (WORLD), view B must not be touched. + assert view_a._dirty.name == "WORLD", "set_local_poses must mark its own view dirty" assert view_b._dirty.name == "NONE", "set_local_poses on view A must not dirty view B" - # Read worlds from view B FIRST. This must not affect view A's already-flushed - # world matrices. + # Read worlds from view B FIRST. This must not clear view A's dirty flag. torch.testing.assert_close( torch.as_tensor(view_b.get_world_poses()[0], device=device), expected_b0, atol=1e-5, rtol=0 ) assert view_b._dirty.name == "NONE" - assert view_a._dirty.name == "NONE" + assert view_a._dirty.name == "WORLD", "view B's read must not clear view A's dirty flag" - # Now read view A's worlds -- world already reflects the new local. + # Now read view A's worlds -- triggers the per-view local->world flush. 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 @@ -779,10 +586,10 @@ def test_multi_view_per_view_dirty_isolation(device): wp.launch(kernel=_fill_position, dim=1, inputs=[new_local_b, 3.0, 0.0, 0.0], device=device) view_b.set_local_poses(translations=new_local_b, orientations=identity_quat) assert view_a._dirty.name == "NONE" - assert view_b._dirty.name == "NONE" + assert view_b._dirty.name == "WORLD" - # A's worlds must still read back the post-set-local value from above; no - # cross-view stomp on the world matrix. + # A's worlds must still read back the post-flush value from above; no cross-view + # stomp on the world matrix. torch.testing.assert_close( torch.as_tensor(view_a.get_world_poses()[0], device=device), expected_a1, atol=1e-5, rtol=0 ) @@ -1014,15 +821,21 @@ def test_interleaved_set_local_then_set_world_partial_indices(device): @pytest.mark.parametrize("device", ["cpu", "cuda:0"]) def test_interleaved_set_emits_warning(device, caplog): - """Interleaving set_world_poses and set_local_poses logs a one-time warning.""" + """Interleaving set_world_poses and set_local_poses logs a one-time warning. + + With the lazy design, a set_world_poses leaves localMatrix stale; a subsequent + set_local_poses on a different (possibly disjoint) index subset must flush the + stale local first to avoid silent inconsistency. That extra kernel launch is + a performance hazard and is flagged once per view instance. + """ view = _build_two_child_view(device) - # First set_world_poses -- no warning (first user setter) + # First set_world_poses -- no warning (first user setter). new_world = wp.zeros((2, 3), dtype=wp.float32, device=device) wp.launch(kernel=_fill_position, dim=2, inputs=[new_world, 1.0, 2.0, 3.0], device=device) view.set_world_poses(positions=new_world) - # Now set_local_poses -- should trigger warning about interleaving + # Now set_local_poses -- should trigger warning about interleaving. new_local = wp.zeros((2, 3), dtype=wp.float32, device=device) wp.launch(kernel=_fill_position, dim=2, inputs=[new_local, 0.0, 0.0, 0.0], device=device) identity_quat = wp.from_torch(torch.tensor([[0.0, 0.0, 0.0, 1.0]] * 2, dtype=torch.float32, device=device)) @@ -1035,7 +848,7 @@ def test_interleaved_set_emits_warning(device, caplog): f"Expected interleave warning, got: {[r.message for r in caplog.records]}" ) - # Second interleave -- warning should NOT repeat (one-time only) + # Second interleave -- warning should NOT repeat (one-time only). caplog.clear() view.set_world_poses(positions=new_world) assert not any("interleaving" in r.message.lower() for r in caplog.records), ( From 264baf5c0a6065a880d1c98f463715f49a2e7774 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Sat, 6 Jun 2026 10:34:47 +0000 Subject: [PATCH 38/54] docs: clarify Fabric getter synchronization --- .../sim/views/fabric_frame_view.py | 51 ++++++++++++++++++- 1 file changed, 50 insertions(+), 1 deletion(-) 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 87f98f0c51b4..9a3a966586ef 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 @@ -306,6 +306,12 @@ def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, When *indices* is None (all prims), the returned arrays are **shared pre-allocated buffers** that are overwritten on the next call. Do not hold references across calls -- copy if persistence is needed. + + When *indices* selects a subset, Fabric launches the decompose kernel + into freshly allocated Warp buffers and returns their + :class:`~isaaclab.utils.warp.ProxyArray` wrappers without blocking. + Callers that need host-visible values immediately must synchronize or + copy explicitly; GPU consumers can rely on normal Warp stream ordering. """ if not self._use_fabric: return self._usd_view.get_world_poses(indices) @@ -313,7 +319,12 @@ def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, if not self._fabric_initialized: self._initialize_fabric() - # If a prior set_local_poses left worldMatrix stale, propagate local -> world first. + # If a prior set_local_poses/set_local_scales left worldMatrix stale, + # propagate local -> world first. + # TODO(pv): This dirty bit tracks Isaac Lab writes only. If Kit/Fabric + # hierarchy update_world_xforms() has already satisfied the dirty world + # matrices during a render tick, we currently have no Fabric-side version + # stamp to observe that and clear the flag; conservatively recompute. self._sync_world_from_local_if_dirty() indices_wp = self._resolve_indices_wp(indices) @@ -404,6 +415,12 @@ def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, When *indices* is None (all prims), the returned arrays are **shared pre-allocated buffers** that are overwritten on the next call. Do not hold references across calls -- copy if persistence is needed. + + When *indices* selects a subset, Fabric launches the decompose kernel + into freshly allocated Warp buffers and returns their + :class:`~isaaclab.utils.warp.ProxyArray` wrappers without blocking. + Callers that need host-visible values immediately must synchronize or + copy explicitly; GPU consumers can rely on normal Warp stream ordering. """ if not self._use_fabric: return self._usd_view.get_local_poses(indices) @@ -412,6 +429,10 @@ def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, self._initialize_fabric() # If a prior set_world_poses/set_world_scales left localMatrix stale, recompute. + # TODO(pv): This dirty bit tracks Isaac Lab writes only. If a future + # Fabric hierarchy tick learns to materialize the corresponding local + # matrices before this access, we have no Fabric-side version stamp to + # observe that and clear the flag; conservatively recompute. self._sync_local_from_world_if_dirty() indices_wp = self._resolve_indices_wp(indices) @@ -457,6 +478,10 @@ def set_world_scales(self, scales, indices=None): self._initialize_fabric() # Sync world matrices first if local writes are pending. + # TODO(pv): This dirty bit tracks Isaac Lab writes only. If Kit/Fabric + # hierarchy update_world_xforms() has already satisfied the dirty world + # matrices during a render tick, we currently have no Fabric-side version + # stamp to observe that and clear the flag; conservatively recompute. self._sync_world_from_local_if_dirty() indices_wp = self._resolve_indices_wp(indices) @@ -491,6 +516,12 @@ def get_world_scales(self, indices=None): pre-allocated buffer** (shared with :meth:`get_local_scales`) that is overwritten on the next call. Do not hold references across calls -- copy if persistence is needed. + + When *indices* selects a subset, Fabric launches the decompose kernel + into a freshly allocated Warp buffer and returns its + :class:`~isaaclab.utils.warp.ProxyArray` wrapper without blocking. + Callers that need host-visible values immediately must synchronize or + copy explicitly; GPU consumers can rely on normal Warp stream ordering. """ if not self._use_fabric: return self._usd_view.get_world_scales(indices) @@ -499,6 +530,10 @@ def get_world_scales(self, indices=None): self._initialize_fabric() # Sync world matrices first if local writes are pending. + # TODO(pv): This dirty bit tracks Isaac Lab writes only. If Kit/Fabric + # hierarchy update_world_xforms() has already satisfied the dirty world + # matrices during a render tick, we currently have no Fabric-side version + # stamp to observe that and clear the flag; conservatively recompute. self._sync_world_from_local_if_dirty() indices_wp = self._resolve_indices_wp(indices) @@ -538,6 +573,10 @@ def set_local_scales(self, scales, indices=None): self._initialize_fabric() # Sync local matrices first if world writes are pending. + # TODO(pv): This dirty bit tracks Isaac Lab writes only. If a future + # Fabric hierarchy tick learns to materialize the corresponding local + # matrices before this access, we have no Fabric-side version stamp to + # observe that and clear the flag; conservatively recompute. self._sync_local_from_world_if_dirty() indices_wp = self._resolve_indices_wp(indices) @@ -572,6 +611,12 @@ def get_local_scales(self, indices=None): pre-allocated buffer** (shared with :meth:`get_world_scales`) that is overwritten on the next call. Do not hold references across calls -- copy if persistence is needed. + + When *indices* selects a subset, Fabric launches the decompose kernel + into a freshly allocated Warp buffer and returns its + :class:`~isaaclab.utils.warp.ProxyArray` wrapper without blocking. + Callers that need host-visible values immediately must synchronize or + copy explicitly; GPU consumers can rely on normal Warp stream ordering. """ if not self._use_fabric: return self._usd_view.get_local_scales(indices) @@ -580,6 +625,10 @@ def get_local_scales(self, indices=None): self._initialize_fabric() # Sync local matrices first if world writes are pending. + # TODO(pv): This dirty bit tracks Isaac Lab writes only. If a future + # Fabric hierarchy tick learns to materialize the corresponding local + # matrices before this access, we have no Fabric-side version stamp to + # observe that and clear the flag; conservatively recompute. self._sync_local_from_world_if_dirty() indices_wp = self._resolve_indices_wp(indices) From 1713bfcbdcc4f3a07ae4d412941c861c450282d1 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Tue, 9 Jun 2026 19:29:11 +0000 Subject: [PATCH 39/54] Introduce FrameViewSpaceWriter context API; drop dirty tracking Replaces the four FrameView pose/scale setters with a single context-managed writer scope on every backend. The writer batches multiple writes (poses + scales) inside one scope, derives the opposite-space matrix once on ``__exit__``, and synchronizes once. Only one writer scope may be active per view at a time; view-level getters raise ``RuntimeError`` while a writer scope is active. API summary: with view.xform_space_writer("world") as w: w.set_poses(positions=p, orientations=o) w.set_scales(scales=s) # Derived (local) matrices are recomputed, the scope releases. Public classes (in ``isaaclab.sim.views``): * ``FrameViewSpaceWriterBase`` -- abstract base * ``FrameViewWorldSpaceWriter`` -- world-space tag class * ``FrameViewLocalSpaceWriter`` -- local-space tag class Behind the scenes: * ``BaseFrameView`` gains ``xform_space_writer()``, ``_active_writer``, and the abstract factory hooks ``_make_world_space_writer`` / ``_make_local_space_writer``. Public getters become guarded wrappers around new ``_get_*_impl`` backend hooks; they raise ``RuntimeError`` when a writer scope is active. * ``FabricFrameView`` ships ``_FabricWorldSpaceWriter`` and ``_FabricLocalSpaceWriter`` that pause ``IFabricHierarchy.track_local_xform_changes`` / ``track_world_xform_changes`` for the scope's lifetime (saving and restoring the prior state so we never re-enable a listener the caller had paused). Eager dual-write inside the scope means Kit's per-tick ``updateWorldXforms`` does not redundantly recompute matrices we just wrote. The renderer's independent ``omni:fabric:worldMatrix`` listener still observes the writes. The lazy-dirty mechanism is gone: the ``_DirtyFlag`` enum, ``_dirty`` field, ``_warned_interleaved_set`` field, the ``_sync_*_if_dirty`` helpers, and the one-time "interleaved set_world_poses / set_local_poses" warning are all deleted. The three-selection RO/RW layout is kept as a defensive layer and for clarity of authoring intent. * ``UsdFrameView`` / ``NewtonSiteFrameView`` / ``OvPhysxFrameView`` ship pass-through writers (their ``set_poses`` / ``set_scales`` immediately delegate to the backend's ``_apply_*_write`` helpers; ``__exit__`` is a no-op beyond releasing the single-writer lock). Setter deprecation / removal: * ``set_world_poses`` and ``set_local_poses`` are kept as one-time-warning shims on ``BaseFrameView`` that route through the writer internally. Use ``view.xform_space_writer("world" | "local")`` and ``w.set_poses(...)``. * ``set_world_scales`` and ``set_local_scales`` were introduced in this release cycle without external users and are removed outright (no deprecation). Use ``w.set_scales(...)`` inside a writer scope. * The existing ``set_scales`` deprecation shim is kept and now opens the backend-appropriate writer scope internally (Fabric: world, USD/OvPhysx: local). Migration: All 81 call sites across ``source/``, ``scripts/``, and the test suites are migrated to the new API in this commit so the repo's own code base raises no deprecation warnings. External callers on ``set_world_poses`` / ``set_local_poses`` / ``set_scales`` keep working (one warning per class on first call). Verification: * ``pytest source/isaaclab_physx/test/sim/test_views_xform_prim_fabric.py`` -- 81 passed, 3 skipped (cuda:1 multi-GPU). * ``pytest source/isaaclab/test/sensors/test_ray_caster_camera.py::test_output_equal_to_usdcamera`` -- the original Camera + RTX regression that motivated this arc, passes. * ``pytest source/isaaclab/test/sensors/test_ray_caster_camera.py::test_output_equal_to_usd_camera_when_intrinsics_set`` -- 4 parametrizations pass. * ``./isaaclab.sh -f`` -- clean. New test coverage for the writer contract (in ``test_views_xform_prim_fabric.py``): * world / local writer derives opposite space on exit * exactly one derive-kernel launch per scope (monkeypatched counter) regardless of how many set_poses / set_scales calls are made inside the scope * single-active-writer invariant raises ``RuntimeError`` * exit restores prior ``track_local_xform_changes`` / ``track_world_xform_changes`` state (does not re-enable listeners the caller had paused) * invalid space string raises ``ValueError`` * empty scope does not launch the derive kernel * view-level getters raise inside an active scope --- .../benchmarks/benchmark_view_comparison.py | 3 +- .../benchmarks/benchmark_xform_prim_view.py | 15 +- .../changelog.d/xform-space-writer.rst | 38 + .../isaaclab/sensors/camera/camera.py | 14 +- .../isaaclab/isaaclab/sim/views/__init__.pyi | 4 + .../isaaclab/sim/views/base_frame_view.py | 265 ++++-- .../isaaclab/sim/views/usd_frame_view.py | 247 +++--- .../isaaclab/sim/views/xform_space_writer.py | 131 +++ .../test/sim/frame_view_contract_utils.py | 63 +- .../test/sim/test_views_xform_prim.py | 15 +- .../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/xform-space-writer.rst | 12 + .../sim/views/newton_site_frame_view.py | 84 +- .../test/sim/test_views_xform_prim_newton.py | 3 +- .../changelog.d/xform-space-writer.rst | 12 + .../sim/views/ovphysx_frame_view.py | 150 ++-- .../changelog.d/xform-space-writer.rst | 28 + .../sim/views/fabric_frame_view.py | 757 +++++++----------- .../test/sim/test_views_xform_prim_fabric.py | 390 +++++---- 22 files changed, 1329 insertions(+), 912 deletions(-) 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/xform-space-writer.rst create mode 100644 source/isaaclab_ovphysx/changelog.d/xform-space-writer.rst create mode 100644 source/isaaclab_physx/changelog.d/xform-space-writer.rst diff --git a/scripts/benchmarks/benchmark_view_comparison.py b/scripts/benchmarks/benchmark_view_comparison.py index aa5927e10b6a..8270be2cd225 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_space_writer("world") 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 ca3c53a18b0a..c3ac1228feee 100644 --- a/scripts/benchmarks/benchmark_xform_prim_view.py +++ b/scripts/benchmarks/benchmark_xform_prim_view.py @@ -177,7 +177,8 @@ def to_torch(a): torch.cuda.synchronize() start_time = time.perf_counter() for _ in range(num_iterations): - xform_view.set_world_poses(new_positions, orientations) + with xform_view.xform_space_writer("world") as w: + w.set_poses(new_positions, orientations) if is_newton: torch.cuda.synchronize() timing_results["set_world_poses"] = (time.perf_counter() - start_time) / num_iterations @@ -213,7 +214,8 @@ def to_torch(a): torch.cuda.synchronize() start_time = time.perf_counter() for _ in range(num_iterations): - xform_view.set_local_poses(new_translations, orientations_local) + with xform_view.xform_space_writer("local") as w: + w.set_poses(new_translations, orientations_local) if is_newton: torch.cuda.synchronize() timing_results["set_local_poses"] = (time.perf_counter() - start_time) / num_iterations @@ -247,7 +249,8 @@ def to_torch(a): torch.cuda.synchronize() start_time = time.perf_counter() for _ in range(num_iterations): - xform_view.set_world_scales(new_world_scales) + with xform_view.xform_space_writer("world") as w: + w.set_scales(new_world_scales) if is_newton: torch.cuda.synchronize() timing_results["set_world_scales"] = (time.perf_counter() - start_time) / num_iterations @@ -279,7 +282,8 @@ def to_torch(a): torch.cuda.synchronize() start_time = time.perf_counter() for _ in range(num_iterations): - xform_view.set_local_scales(new_local_scales) + with xform_view.xform_space_writer("local") as w: + w.set_scales(new_local_scales) if is_newton: torch.cuda.synchronize() timing_results["set_local_scales"] = (time.perf_counter() - start_time) / num_iterations @@ -302,7 +306,8 @@ def to_torch(a): torch.cuda.synchronize() start_time = time.perf_counter() for _ in range(num_iterations): - xform_view.set_world_poses(new_positions, orientations) + with xform_view.xform_space_writer("world") as w: + w.set_poses(new_positions, orientations) xform_view.get_world_poses() if is_newton: torch.cuda.synchronize() 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..3a3236e042cf --- /dev/null +++ b/source/isaaclab/changelog.d/xform-space-writer.rst @@ -0,0 +1,38 @@ +Added +^^^^^ + +* Added :class:`~isaaclab.sim.views.FrameViewSpaceWriterBase`, the new context-managed + write API for ``FrameView``-managed prim transforms. Open with + ``view.xform_space_writer("world" | "local")`` 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_space_writer`. + +Deprecated +^^^^^^^^^^ + +* Deprecated :meth:`~isaaclab.sim.views.BaseFrameView.set_world_poses` and + :meth:`~isaaclab.sim.views.BaseFrameView.set_local_poses`. Use + ``with view.xform_space_writer("world" | "local") as w: w.set_poses(...)`` + instead. The deprecated methods still work but emit a one-time + ``DeprecationWarning`` per class and open a single-statement writer scope + internally. + +Removed +^^^^^^^ + +* **Breaking:** Removed ``set_world_scales`` and ``set_local_scales`` + from :class:`~isaaclab.sim.views.BaseFrameView` (and all subclasses). + These were introduced in this release cycle without a stable downstream + user, so they are removed outright (no deprecation cycle). Use + ``with view.xform_space_writer("world" | "local") as w: w.set_scales(...)`` + instead. diff --git a/source/isaaclab/isaaclab/sensors/camera/camera.py b/source/isaaclab/isaaclab/sensors/camera/camera.py index 5ed97a3825f5..4269b1fd25c0 100644 --- a/source/isaaclab/isaaclab/sensors/camera/camera.py +++ b/source/isaaclab/isaaclab/sensors/camera/camera.py @@ -376,7 +376,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_space_writer("world") 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 @@ -434,11 +435,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_space_writer("world") 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 ee5e313d7f3c..cda64dd4f435 100644 --- a/source/isaaclab/isaaclab/sim/views/base_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/base_frame_view.py @@ -9,23 +9,47 @@ import abc import warnings +from typing import TYPE_CHECKING, Literal 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 getters return :class:`~isaaclab.utils.warp.ProxyArray`. Setters accept ``wp.array``. + All getters return :class:`~isaaclab.utils.warp.ProxyArray`. All writes go + through :meth:`xform_space_writer` -- the recommended API: + + .. code-block:: python + + with view.xform_space_writer("world") 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: @@ -38,67 +62,93 @@ def device(self) -> str: """Device where arrays are allocated (``"cpu"`` or ``"cuda:0"``).""" ... - @abc.abstractmethod - def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: - """Get world-space positions and orientations for prims in the view. + # ------------------------------------------------------------------ + # Write scope -- recommended API for all transform writes. + # ------------------------------------------------------------------ + + def xform_space_writer(self, space: Literal["world", "local"]) -> FrameViewSpaceWriterBase: + """Open a write scope on this view (recommended write API). Args: - indices: Subset of prims to query. ``None`` means all prims. + space: ``"world"`` or ``"local"``. 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. + An :class:`~isaaclab.sim.views.FrameViewWorldSpaceWriter` or + :class:`~isaaclab.sim.views.FrameViewLocalSpaceWriter` context manager. + + Raises: + ValueError: If *space* is neither ``"world"`` nor ``"local"``. + RuntimeError: On ``__enter__``, if another writer is already active + on this view. + + Example: + .. code-block:: python + + with view.xform_space_writer("world") as w: + w.set_poses(positions=p, orientations=o) + w.set_scales(scales=s) """ + if space == "world": + return self._make_world_space_writer() + if space == "local": + return self._make_local_space_writer() + raise ValueError(f"Invalid space {space!r}; expected 'world' or 'local'.") + + @abc.abstractmethod + def _make_world_space_writer(self) -> FrameViewWorldSpaceWriter: + """Backend hook: return a fresh :class:`FrameViewWorldSpaceWriter` for this view.""" ... @abc.abstractmethod - def set_world_poses( - self, - positions: wp.array | None = None, - orientations: wp.array | None = None, - indices: wp.array | None = None, - ) -> None: - """Set world-space positions and/or orientations for prims in the view. + 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 an xform_space_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. 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. + indices: Subset of prims to query. ``None`` means all prims. + + Returns: + A tuple ``(positions, 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_world_poses") + return self._get_world_poses_impl(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. + """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. Use ``.warp`` for the underlying ``wp.array`` or ``.torch`` for a - cached zero-copy ``torch.Tensor`` view. - """ - ... + wrappers. - @abc.abstractmethod - def set_local_poses( - self, - translations: wp.array | None = None, - orientations: wp.array | None = None, - indices: wp.array | None = None, - ) -> None: - """Set local-space translations and/or orientations for prims in the view. - - 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. + 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) - @abc.abstractmethod def get_local_scales(self, indices: wp.array | None = None) -> ProxyArray: """Get local-space scales for prims in the view. @@ -106,21 +156,14 @@ def get_local_scales(self, indices: wp.array | None = None) -> ProxyArray: indices: Subset of prims to query. ``None`` means all prims. Returns: - A :class:`~isaaclab.utils.warp.ProxyArray` wrapping a ``wp.array`` of shape ``(M, 3)``. - """ - ... - - @abc.abstractmethod - def set_local_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: - """Set local-space scales for prims in the view. + A :class:`~isaaclab.utils.warp.ProxyArray` of shape ``(M, 3)``. - Args: - scales: Scales ``(M, 3)`` as ``wp.array``. - indices: Subset of prims to update. ``None`` means all prims. + 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) - @abc.abstractmethod def get_world_scales(self, indices: wp.array | None = None) -> ProxyArray: """Get world-space (composed) scales for prims in the view. @@ -136,26 +179,106 @@ def get_world_scales(self, indices: wp.array | None = None) -> ProxyArray: indices: Subset of prims to query. ``None`` means all prims. Returns: - A :class:`~isaaclab.utils.warp.ProxyArray` wrapping a ``wp.array`` of shape ``(M, 3)``. + 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 set_world_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: - """Set world-space (composed) scales for prims in the view. + def _get_local_poses_impl(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: + """Backend implementation of :meth:`get_local_poses`.""" + ... - Args: - scales: Scales ``(M, 3)`` as ``wp.array``. - indices: Subset of prims to update. ``None`` means all prims. - """ + @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`.""" ... # ------------------------------------------------------------------ - # Deprecated -- use get/set_local_scales or get/set_world_scales + # Deprecated pose setters -- route through the writer scope. # ------------------------------------------------------------------ - _get_scales_deprecated_warned: bool = False + _set_world_poses_deprecated_warned: bool = False + _set_local_poses_deprecated_warned: bool = False _set_scales_deprecated_warned: bool = False + _get_scales_deprecated_warned: bool = False + + def set_world_poses( + self, + positions: wp.array | None = None, + orientations: wp.array | None = None, + indices: wp.array | None = None, + ) -> None: + """Set world-space positions and/or orientations for prims in the view. + + .. deprecated:: + Use ``with view.xform_space_writer("world") as w: w.set_poses(...)`` instead. + This method opens a single-statement writer scope internally. + + 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. + """ + if not BaseFrameView._set_world_poses_deprecated_warned: + BaseFrameView._set_world_poses_deprecated_warned = True + warnings.warn( + 'set_world_poses() is deprecated. Use \'with view.xform_space_writer("world") as w:' + " w.set_poses(...)' instead.", + DeprecationWarning, + stacklevel=2, + ) + with self.xform_space_writer("world") as writer: + writer.set_poses(positions, orientations, indices) + + def set_local_poses( + self, + translations: wp.array | None = None, + orientations: wp.array | None = None, + indices: wp.array | None = None, + ) -> None: + """Set local-space translations and/or orientations for prims in the view. + + .. deprecated:: + Use ``with view.xform_space_writer("local") as w: w.set_poses(...)`` instead. + This method opens a single-statement writer scope internally. + + 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. + """ + if not BaseFrameView._set_local_poses_deprecated_warned: + BaseFrameView._set_local_poses_deprecated_warned = True + warnings.warn( + 'set_local_poses() is deprecated. Use \'with view.xform_space_writer("local") as w:' + " w.set_poses(...)' instead.", + DeprecationWarning, + stacklevel=2, + ) + with self.xform_space_writer("local") as writer: + writer.set_poses(translations, orientations, indices) + + # ------------------------------------------------------------------ + # Deprecated -- use writer scope or get_local_scales / get_world_scales. + # ------------------------------------------------------------------ def get_scales(self, indices: wp.array | None = None) -> ProxyArray: """Get scales for prims in the view. @@ -170,6 +293,9 @@ def get_scales(self, indices: wp.array | None = None) -> ProxyArray: Returns: A ``ProxyArray`` of shape ``(M, 3)``. + + Raises: + RuntimeError: If a writer scope is active on this view. """ if not BaseFrameView._get_scales_deprecated_warned: BaseFrameView._get_scales_deprecated_warned = True @@ -178,15 +304,17 @@ def get_scales(self, indices: wp.array | None = None) -> ProxyArray: DeprecationWarning, stacklevel=2, ) + self._assert_no_active_writer("get_scales") return self._get_scales_impl(indices) def set_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: """Set scales for prims in the view. .. deprecated:: - Use :meth:`set_local_scales` or :meth:`set_world_scales` instead. - This method delegates to :meth:`_set_scales_impl` which preserves - each backend's legacy behavior. + Use ``with view.xform_space_writer("world" | "local") as w: w.set_scales(...)`` instead. + This 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``. Args: scales: Scales ``(M, 3)`` as ``wp.array``. @@ -195,7 +323,8 @@ def set_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: if not BaseFrameView._set_scales_deprecated_warned: BaseFrameView._set_scales_deprecated_warned = True warnings.warn( - "set_scales() is deprecated. Use set_local_scales() or set_world_scales() instead.", + 'set_scales() is deprecated. Use \'with view.xform_space_writer("world" or' + ' "local") as w: w.set_scales(...)\' instead.', DeprecationWarning, stacklevel=2, ) @@ -203,10 +332,10 @@ def set_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: @abc.abstractmethod def _get_scales_impl(self, indices: wp.array | None = None) -> ProxyArray: - """Backend-specific implementation for deprecated get_scales().""" + """Backend-specific implementation for deprecated :meth:`get_scales`.""" ... @abc.abstractmethod def _set_scales_impl(self, scales: wp.array, indices: wp.array | None = None) -> None: - """Backend-specific implementation for deprecated set_scales().""" + """Backend-specific implementation for deprecated :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 3927b4690a62..8cf6bbe816d0 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,13 @@ class UsdFrameView(BaseFrameView): For GPU-accelerated Fabric operations, use the PhysX backend variant obtained via :class:`~isaaclab.sim.views.FrameView`. - Getters return :class:`~isaaclab.utils.warp.ProxyArray`. Setters accept ``wp.array``. + All writes go through :meth:`xform_space_writer` (recommended). 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 +133,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_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}.") - def set_world_poses( + 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 +240,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 +260,8 @@ def set_local_poses( if orientations_array is not None: prim.GetAttribute("xformOp:orient").Set(orientations_array[idx]) - def set_local_scales(self, scales: wp.array, indices: wp.array | None = None): - """Set local-space scales (xformOp:scale) 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,15 +270,11 @@ def set_local_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_world_scales(self, scales: wp.array, indices: wp.array | None = None): - """Set world-space (composed) scales 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. Computes ``local_scale = world_scale / parent_world_scale`` and writes to ``xformOp:scale``. - - Args: - scales: World-space scales of shape ``(M, 3)``. - indices: Indices of prims to set scales for. Defaults to None (all prims). """ indices_list = self._resolve_indices(indices) scales_np = self._to_numpy(scales) @@ -262,49 +300,11 @@ def set_world_scales(self, scales: wp.array, indices: wp.array | None = None): ) prim.GetAttribute("xformOp:scale").Set(local_scale) - def _get_scales_impl(self, indices: wp.array | None = None) -> ProxyArray: - """USD legacy: get_scales returns local scales.""" - return self.get_local_scales(indices) - - def _set_scales_impl(self, scales: wp.array, indices: wp.array | None = None) -> None: - """USD default: set_scales writes local scales.""" - self.set_local_scales(scales, indices) - - 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() - # ------------------------------------------------------------------ - # 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)) @@ -322,17 +322,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)) @@ -350,15 +340,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_local_scales(self, indices: wp.array | None = None) -> ProxyArray: - """Get local-space scales (xformOp:scale) for prims in the view. - - Args: - indices: Indices of prims to get scales for. Defaults to None (all prims). - - Returns: - A :class:`~isaaclab.utils.warp.ProxyArray` wrapping 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)) @@ -368,19 +350,7 @@ def get_local_scales(self, indices: wp.array | None = None) -> ProxyArray: return ProxyArray(wp.array(np.array(scales, dtype=np.float32), dtype=wp.float32, device=self._device)) - def get_world_scales(self, indices: wp.array | None = None) -> ProxyArray: - """Get world-space (composed) scales for prims in the view. - - Computes the effective world-space scale by extracting row lengths - from the world transform matrix (USD uses a row-vector convention - where each row of the 3x3 sub-matrix is a basis vector). - - Args: - indices: Indices of prims to get scales for. Defaults to None (all prims). - - Returns: - A :class:`~isaaclab.utils.warp.ProxyArray` wrapping a ``wp.array`` of shape ``(M, 3)``. - """ + 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()) @@ -394,22 +364,18 @@ def get_world_scales(self, indices: wp.array | None = None) -> ProxyArray: return ProxyArray(wp.array(scales, 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). + # ------------------------------------------------------------------ + # Deprecated get_scales / set_scales hooks + # ------------------------------------------------------------------ - Returns: - A tensor of shape ``(M,)`` containing the visibility of each prim (bool). - """ - indices_list = self._resolve_indices(indices) + def _get_scales_impl(self, indices: wp.array | None = None) -> ProxyArray: + """USD legacy: deprecated get_scales returns local scales.""" + return self._get_local_scales_impl(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 + def _set_scales_impl(self, scales: wp.array, indices: wp.array | None = None) -> None: + """USD legacy: deprecated set_scales writes local scales via a one-shot writer scope.""" + with self.xform_space_writer("local") as writer: + writer.set_scales(scales, indices) # ------------------------------------------------------------------ # Helpers @@ -427,3 +393,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..4be224ef679c --- /dev/null +++ b/source/isaaclab/isaaclab/sim/views/xform_space_writer.py @@ -0,0 +1,131 @@ +# 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_space_writer("world") 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. +""" + +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_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`. + """ + + 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 xform_space_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/test/sim/frame_view_contract_utils.py b/source/isaaclab/test/sim/frame_view_contract_utils.py index b27335547d72..2741eb955ae5 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_space_writer("world") 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_space_writer("local") 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_space_writer("world") 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_space_writer("local") 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_space_writer("world") 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_space_writer("local") 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_space_writer("world") 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_space_writer("world") 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_space_writer("local") 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_space_writer("world") 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) @@ -464,7 +474,8 @@ def test_set_local_scales_roundtrip(device, view_factory): 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) - bundle.view.set_local_scales(new_scales) + with bundle.view.xform_space_writer("local") 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) @@ -478,7 +489,8 @@ def test_set_world_scales_roundtrip(device, view_factory): 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) - bundle.view.set_world_scales(new_scales) + with bundle.view.xform_space_writer("world") 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) @@ -495,7 +507,8 @@ def test_local_scales_do_not_affect_local_poses(device, view_factory): 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) - bundle.view.set_local_scales(new_scales) + with bundle.view.xform_space_writer("local") 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]) diff --git a/source/isaaclab/test/sim/test_views_xform_prim.py b/source/isaaclab/test/sim/test_views_xform_prim.py index 52cc7a05fc80..099fd127a66f 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_space_writer("local") as w: + w.set_poses(positions=torch.tensor(frame_positions, device=device)) + with targets_view.xform_space_writer("local") as w: + w.set_poses(positions=torch.tensor(target_positions, device=device)) world_pos = targets_view.get_world_poses()[0].torch expected = torch.tensor( @@ -265,7 +267,8 @@ def test_set_local_scales_then_get_world_scales(device): 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) - view.set_local_scales(local_scales) + with view.xform_space_writer("local") 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) @@ -280,7 +283,8 @@ def test_set_world_scales_then_get_local_scales(device): 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) - view.set_world_scales(world_scales) + with view.xform_space_writer("world") 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) @@ -341,7 +345,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_space_writer("world") 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..951bbe3a3e54 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_space_writer("world") 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..c712812bee60 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_space_writer("world") 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..b4fca471f981 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_space_writer("world") 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/xform-space-writer.rst b/source/isaaclab_newton/changelog.d/xform-space-writer.rst new file mode 100644 index 000000000000..85df80394b1f --- /dev/null +++ b/source/isaaclab_newton/changelog.d/xform-space-writer.rst @@ -0,0 +1,12 @@ +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_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 b7d2918e01c7..d28d27f34730 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 @@ -17,6 +17,7 @@ from isaaclab.cloner.cloner_utils import iter_clone_plan_matches 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 @@ -467,7 +468,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 @@ -486,7 +501,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, @@ -498,7 +513,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: @@ -513,7 +528,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) @@ -531,7 +546,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, @@ -542,7 +557,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: @@ -561,7 +576,7 @@ def set_local_poses( # Scales # ------------------------------------------------------------------ - def get_world_scales(self, indices: wp.array | None = None) -> ProxyArray: + 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 @@ -580,15 +595,15 @@ def get_world_scales(self, indices: wp.array | None = None) -> ProxyArray: ) return ProxyArray(out) - def get_local_scales(self, indices: wp.array | None = None) -> ProxyArray: + 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(indices) + return self._get_world_scales_impl(indices) - def set_world_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: + 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 @@ -604,13 +619,13 @@ def set_world_scales(self, scales: wp.array, indices: wp.array | None = None) -> device=self._device, ) - def set_local_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: + 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.set_world_scales(scales, indices) + 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.""" @@ -646,5 +661,48 @@ def _get_scales_impl(self, indices: wp.array | None = None) -> ProxyArray: return self._get_legacy_shape_scales(indices) def _set_scales_impl(self, scales: wp.array, indices: wp.array | None = None) -> None: - """Newton legacy: set_scales writes collision shape geometry scales.""" + """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..d743140c48b3 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_space_writer("world") 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/xform-space-writer.rst b/source/isaaclab_ovphysx/changelog.d/xform-space-writer.rst new file mode 100644 index 000000000000..16b170753e7a --- /dev/null +++ b/source/isaaclab_ovphysx/changelog.d/xform-space-writer.rst @@ -0,0 +1,12 @@ +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_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 acf8abe5822b..e293e14bbb7a 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 @@ -624,17 +625,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 +666,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 +704,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 +744,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,55 +809,43 @@ def _ensure_usd_view(self) -> UsdFrameView: ) return self._usd_view - def get_local_scales(self, indices: wp.array | None = None) -> ProxyArray: + 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: - A :class:`~isaaclab.utils.warp.ProxyArray` wrapping a ``wp.array`` of shape ``(M, 3)``. + updates are not reflected here. """ - return self._ensure_usd_view().get_local_scales(indices) + return self._ensure_usd_view()._get_local_scales_impl(indices) - def get_world_scales(self, indices: wp.array | None = None) -> ProxyArray: + 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(indices) + return self._ensure_usd_view()._get_world_scales_impl(indices) - def set_local_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: + 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_local_scales(scales, indices) + self._ensure_usd_view()._apply_local_scale_write(scales, indices) - def set_world_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: + 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().set_world_scales(scales, indices) + self._ensure_usd_view()._apply_world_scale_write(scales, indices) def _get_scales_impl(self, indices=None): - """OvPhysX legacy: get_scales returns local scales (same as USD).""" - return self.get_local_scales(indices) + """OvPhysX legacy: deprecated get_scales returns local scales.""" + return self._get_local_scales_impl(indices) def _set_scales_impl(self, scales, indices=None): - """OvPhysX default: set_scales writes local scales (same as USD).""" - self.set_local_scales(scales, indices) + """OvPhysX legacy: deprecated set_scales writes local scales via a one-shot writer scope.""" + with self.xform_space_writer("local") 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). @@ -904,3 +867,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/xform-space-writer.rst b/source/isaaclab_physx/changelog.d/xform-space-writer.rst new file mode 100644 index 000000000000..2ff88d32d5b4 --- /dev/null +++ b/source/isaaclab_physx/changelog.d/xform-space-writer.rst @@ -0,0 +1,28 @@ +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 Kit's per-tick + ``updateWorldXforms()`` does not redundantly recompute matrices the + user just wrote. The renderer's independent ``omni:fabric:worldMatrix`` + listener is unaffected and observes the writes. + + The lazy-dirty-flag mechanism (the ``_DirtyFlag`` enum, ``_dirty`` field, + ``_sync_*_if_dirty`` helpers, and the one-time + ``interleaved set_world_poses / set_local_poses`` warning) has been + removed -- the eager dual-write inside the scope makes all of that + unnecessary. + + The three-selection RO/RW layout (``_trans_sel_ro``, + ``_world_sel_rw``, ``_local_sel_rw``) is kept as a defensive layer and + for clarity of authoring intent. 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 9a3a966586ef..e0dcd055b181 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,6 @@ from __future__ import annotations -import enum import logging import torch @@ -18,22 +17,13 @@ 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 logger = logging.getLogger(__name__) -class _DirtyFlag(enum.Enum): - """Which matrix direction is stale and needs recomputation on the next read.""" - - NONE = 0 - #: World matrices are stale (a prior ``set_local_poses`` wrote new locals). - WORLD = 1 - #: Local matrices are stale (a prior ``set_world_poses``/``set_world_scales`` wrote new worlds). - LOCAL = 2 - - def _to_float32_2d(a: wp.array | torch.Tensor) -> wp.array | torch.Tensor: """Ensure array is compatible with Fabric kernels (2-D float32). @@ -62,6 +52,10 @@ class FabricFrameView(BaseFrameView): ``omni:fabric:worldMatrix`` and ``omni:fabric:localMatrix`` directly. All other operations delegate to the internal USD view. + All writes go through :meth:`xform_space_writer` (recommended) or the + deprecated :meth:`set_world_poses` / :meth:`set_local_poses` / etc. shims + inherited from :class:`BaseFrameView`. + Behavior (Fabric path): * **Leaf-prim assumption.** This view manages a flat set of sibling prims @@ -74,74 +68,45 @@ class FabricFrameView(BaseFrameView): 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. - * **World <-> local consistency (lazy, RO/RW-protected).** After - ``set_world_poses`` / ``set_world_scales``, local matrices are - recomputed lazily when ``get_local_poses`` / ``get_local_scales`` is - called; after ``set_local_poses`` / ``set_local_scales``, world - matrices are recomputed lazily when ``get_world_poses`` / - ``get_world_scales`` is called. Both directions stay in sync without - round-tripping through USD. The renderer/FSD reads cached - ``omni:fabric:worldMatrix`` *between* user writes and the next - ``IFabricHierarchy.update_world_xforms`` tick (which Kit runs as part - of the render path). Correctness across that tick depends on the - RO/RW selection layout below -- a single combined ``ReadWrite`` - selection is **not** safe; see the next bullet. - * **Three selections with asymmetric RO/RW access (load-bearing).** - The view holds three persistent Fabric selections: + * **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 read directly + from Fabric storage without any further synchronization. + * **Hierarchy listeners are paused while a writer scope is active.** + The writer's ``__enter__`` calls + :meth:`IFabricHierarchy.track_local_xform_changes(False)` / + :meth:`track_world_xform_changes(False)` (saving the prior state) so + that Kit's per-tick ``updateWorldXforms()`` does not redundantly + recompute matrices we just wrote. ``__exit__`` restores the prior + tracking state (so we do not re-enable listeners the caller had + previously paused). The renderer's own independent worldMatrix + listener is unaffected and still observes our writes. + * **Three selections with asymmetric RO/RW access.** Despite the + pause/restore above, we keep three selections as a defensive layer: .. code-block:: text - _trans_sel_ro : worldMatrix=RO, localMatrix=RO (reads + sync helpers) - _world_sel_rw : worldMatrix=RW, localMatrix=RO (set_world_poses / _scales) - _local_sel_rw : worldMatrix=RO, localMatrix=RW (set_local_poses / _scales) - - The access flags tell Fabric which attribute the user is *authoring* - for a given operation. When ``IFabricHierarchy.update_world_xforms`` - runs (as part of Kit's render tick), it respects those flags and does - not recompute an attribute marked RO on the most-recently-written - selection. Concretely: - - - After ``set_world_poses`` (via ``_world_sel_rw``): Fabric sees - ``localMatrix=RO`` -> does not recompute world from local -> the - user's worldMatrix write survives until the renderer reads it. - - After ``set_local_poses`` (via ``_local_sel_rw``): Fabric sees - ``localMatrix=RW`` -> recomputes world from the new local on the - next tick -> the renderer sees the correct world. - - Using a single combined ``worldMatrix=RW, localMatrix=RW`` selection - removes this protection. Fabric then sees both attributes as - user-authored and falls back to the hierarchy's canonical direction - (local -> world), recomputing world from a stale local and clobbering - the user's worldMatrix write. This was the failure mode that broke - the Camera + RTX renderer path when the camera prim sat in a - hierarchy traversed during the next render tick. - * **Dirty-flag invariant.** The ``_dirty`` enum is one of ``NONE``, - ``WORLD``, or ``LOCAL`` -- mutually exclusive by construction. - ``set_world_poses`` / ``set_world_scales`` sets ``_dirty = LOCAL``; - ``set_local_poses`` / ``set_local_scales`` sets ``_dirty = WORLD``. - If the user interleaves both setters on the same view within a single - frame, the second setter flushes the first's stale data before writing. - This is correct but incurs an extra kernel launch -- a one-time - warning is logged when this happens. + _trans_sel_ro : worldMatrix=RO, localMatrix=RO (reads) + _world_sel_rw : worldMatrix=RW, localMatrix=RO (world writer) + _local_sel_rw : worldMatrix=RO, localMatrix=RW (local writer) + + A combined ``ReadWrite(world, local)`` selection is unsafe even with + tracking pause -- if a refactor accidentally re-enables tracking, + Fabric would see both attributes as user-authored and fall back to the + hierarchy's canonical direction (local -> world), clobbering our world + write. The separate RO/RW layout makes the intended authoring + direction explicit. * **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. - Steady-state overhead is negligible. - - Performance note: - The fast path assumes the user calls **either** ``set_world_poses`` - **or** ``set_local_poses`` exclusively within a frame (not both). - In that case, setters are O(1) kernel launches with no - synchronization overhead beyond the single ``wp.synchronize()``; - getters lazily flush the opposite direction only when actually - needed. - - Interleaving both setters on different index subsets within the - same frame is supported and correct, but triggers an extra flush - kernel per transition. A warning is emitted once per view - instance. - - Pose getters return :class:`~isaaclab.utils.warp.ProxyArray`; setters + + Pose getters return :class:`~isaaclab.utils.warp.ProxyArray`; the + deprecated :meth:`set_world_poses` / :meth:`set_local_poses` shims accept + :class:`wp.array`. Inside a writer scope, the writer's + :meth:`~FrameViewSpaceWriterBase.set_poses` / :meth:`~FrameViewSpaceWriterBase.set_scales` accept :class:`wp.array`. """ @@ -178,27 +143,17 @@ def __init__( # 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_world_scales into single kernel launch (PR #5674 pv/fabric-fused-compose) self._fabric_initialized = False self._stage = None self._fabric_hierarchy = None - # Tracks which matrix direction is stale. Mutually exclusive by construction. - # Per-view (not per-stage) so concurrent views on the same stage don't interfere. - self._dirty: _DirtyFlag = _DirtyFlag.NONE - self._warned_interleaved_set: bool = False # Three persistent Fabric selections with asymmetric access flags. - # See the class docstring "Three selections with asymmetric RO/RW access" - # bullet for why this layout is required. self._trans_sel_ro = None self._world_sel_rw = None self._local_sel_rw = None # Index arrays (view-side indices and per-selection view->fabric mappings). - # Each selection's ``GetPaths()`` ordering is independent, so view->fabric - # is cached per selection -- sharing would silently corrupt indexed arrays - # whose selection didn't fire ``PrepareForReuse`` on the same frame. self._view_indices: wp.array | None = None self._trans_ro_fabric_indices: wp.array | None = None self._world_rw_fabric_indices: wp.array | None = None @@ -213,7 +168,6 @@ def __init__( self._parent_world_ifa_ro = None # Sentinel passed to compose/decompose kernels for unused slots. - # Kernels gate per-row access on ``shape[0] > 0``, so (0, 0) suffices. self._fabric_empty_2d_array_sentinel: wp.array | None = None # ------------------------------------------------------------------ @@ -248,85 +202,30 @@ 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 - - if not self._fabric_initialized: - self._initialize_fabric() + return _FabricFallbackWorldWriter(self) + return _FabricWorldSpaceWriter(self) - # If a prior set_local_poses left worlds stale, flush them now. - if self._dirty == _DirtyFlag.WORLD and not self._warned_interleaved_set: - self._warned_interleaved_set = True - logger.warning( - "FabricFrameView: set_world_poses called while world matrices are stale from a " - "prior set_local_poses. Flushing stale worlds first. " - "For best performance, avoid interleaving set_world_poses and set_local_poses " - "on the same view within a single frame -- use one or the other exclusively." - ) - - self._sync_world_from_local_if_dirty() - - indices_wp = self._resolve_indices_wp(indices) - positions_wp = self._to_float32_2d_or_empty(positions) - orientations_wp = self._to_float32_2d_or_empty(orientations) + def _make_local_space_writer(self) -> FrameViewLocalSpaceWriter: + if not self._use_fabric: + return _FabricFallbackLocalWriter(self) + return _FabricLocalSpaceWriter(self) - wp.launch( - kernel=fabric_utils.compose_indexed_fabric_transforms, - dim=indices_wp.shape[0], - inputs=[ - self._get_world_rw_array(), - positions_wp, - orientations_wp, - self._fabric_empty_2d_array_sentinel, - False, - False, - False, - indices_wp, - ], - device=self._device, - ) - wp.synchronize() + # ------------------------------------------------------------------ + # Getter hooks -- read directly from Fabric (no lazy sync) + # ------------------------------------------------------------------ - # World was just written -- mark local poses as stale so the next - # get_local_poses recomputes them lazily. No eager local recompute is - # needed: the worldMatrix write is protected by ``_world_sel_rw`` having - # ``localMatrix=RO``, so the next ``update_world_xforms`` tick will not - # overwrite world from a stale local. - self._dirty = _DirtyFlag.LOCAL - - def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: - """Return (positions, orientations) in world frame. - - .. warning:: - When *indices* is None (all prims), the returned arrays are **shared - pre-allocated buffers** that are overwritten on the next call. Do not - hold references across calls -- copy if persistence is needed. - - When *indices* selects a subset, Fabric launches the decompose kernel - into freshly allocated Warp buffers and returns their - :class:`~isaaclab.utils.warp.ProxyArray` wrappers without blocking. - Callers that need host-visible values immediately must synchronize or - copy explicitly; GPU consumers can rely on normal Warp stream ordering. - """ + 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 a prior set_local_poses/set_local_scales left worldMatrix stale, - # propagate local -> world first. - # TODO(pv): This dirty bit tracks Isaac Lab writes only. If Kit/Fabric - # hierarchy update_world_xforms() has already satisfied the dirty world - # matrices during a render tick, we currently have no Fabric-side version - # stamp to observe that and clear the flag; conservatively recompute. - self._sync_world_from_local_if_dirty() - indices_wp = self._resolve_indices_wp(indices) count = indices_wp.shape[0] @@ -356,85 +255,13 @@ def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, return self._fabric_positions_ta, self._fabric_orientations_ta return ProxyArray(positions_wp), ProxyArray(orientations_wp) - # ------------------------------------------------------------------ - # Local poses - # ------------------------------------------------------------------ - - def set_local_poses(self, translations=None, orientations=None, indices=None): - if not self._use_fabric: - self._usd_view.set_local_poses(translations, orientations, indices) - return - - if not self._fabric_initialized: - self._initialize_fabric() - - # If a prior set_world_poses left locals stale, flush them now before we - # overwrite a (possibly different) subset of local matrices. - if self._dirty == _DirtyFlag.LOCAL and not self._warned_interleaved_set: - self._warned_interleaved_set = True - logger.warning( - "FabricFrameView: set_local_poses called while local matrices are stale from a " - "prior set_world_poses/set_world_scales. Flushing stale locals first. " - "For best performance, avoid interleaving set_world_poses and set_local_poses " - "on the same view within a single frame -- use one or the other exclusively." - ) - - self._sync_local_from_world_if_dirty() - - indices_wp = self._resolve_indices_wp(indices) - translations_wp = self._to_float32_2d_or_empty(translations) - orientations_wp = self._to_float32_2d_or_empty(orientations) - - wp.launch( - kernel=fabric_utils.compose_indexed_fabric_transforms, - dim=indices_wp.shape[0], - inputs=[ - self._get_local_rw_array(), - translations_wp, - orientations_wp, - self._fabric_empty_2d_array_sentinel, - False, - False, - False, - indices_wp, - ], - device=self._device, - ) - wp.synchronize() - - # Local was just written -- mark world matrices stale. No eager world - # recompute: the ``_local_sel_rw`` selection has ``worldMatrix=RO``, so - # Kit's next ``update_world_xforms`` tick will recompute world from the - # new local automatically, and the renderer will read the correct world. - self._dirty = _DirtyFlag.WORLD - - def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: - """Return (translations, orientations) in parent-local frame. - - .. warning:: - When *indices* is None (all prims), the returned arrays are **shared - pre-allocated buffers** that are overwritten on the next call. Do not - hold references across calls -- copy if persistence is needed. - - When *indices* selects a subset, Fabric launches the decompose kernel - into freshly allocated Warp buffers and returns their - :class:`~isaaclab.utils.warp.ProxyArray` wrappers without blocking. - Callers that need host-visible values immediately must synchronize or - copy explicitly; GPU consumers can rely on normal Warp stream ordering. - """ + def _get_local_poses_impl(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: if not self._use_fabric: - return self._usd_view.get_local_poses(indices) + return self._usd_view._get_local_poses_impl(indices) if not self._fabric_initialized: self._initialize_fabric() - # If a prior set_world_poses/set_world_scales left localMatrix stale, recompute. - # TODO(pv): This dirty bit tracks Isaac Lab writes only. If a future - # Fabric hierarchy tick learns to materialize the corresponding local - # matrices before this access, we have no Fabric-side version stamp to - # observe that and clear the flag; conservatively recompute. - self._sync_local_from_world_if_dirty() - indices_wp = self._resolve_indices_wp(indices) count = indices_wp.shape[0] @@ -464,78 +291,26 @@ def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, return self._fabric_local_translations_ta, self._fabric_local_orientations_ta return ProxyArray(translations_wp), ProxyArray(orientations_wp) - # ------------------------------------------------------------------ - # Scales - # ------------------------------------------------------------------ - - def set_world_scales(self, scales, indices=None): - """Set world-space (composed) scales by decomposing/recomposing worldMatrix.""" + def _get_world_scales_impl(self, indices=None) -> ProxyArray: if not self._use_fabric: - self._usd_view.set_world_scales(scales, indices) - return + return self._usd_view._get_world_scales_impl(indices) if not self._fabric_initialized: self._initialize_fabric() - # Sync world matrices first if local writes are pending. - # TODO(pv): This dirty bit tracks Isaac Lab writes only. If Kit/Fabric - # hierarchy update_world_xforms() has already satisfied the dirty world - # matrices during a render tick, we currently have no Fabric-side version - # stamp to observe that and clear the flag; conservatively recompute. - self._sync_world_from_local_if_dirty() + return self._decompose_scales(self._get_world_ro_array(), indices) - indices_wp = self._resolve_indices_wp(indices) - scales_wp = self._to_float32_2d_or_empty(scales) - - wp.launch( - kernel=fabric_utils.compose_indexed_fabric_transforms, - dim=indices_wp.shape[0], - inputs=[ - self._get_world_rw_array(), - self._fabric_empty_2d_array_sentinel, - self._fabric_empty_2d_array_sentinel, - scales_wp, - False, - False, - False, - indices_wp, - ], - device=self._device, - ) - wp.synchronize() - - # World was just written -- mark local poses as stale. See set_world_poses - # for why no eager local recompute is needed. - self._dirty = _DirtyFlag.LOCAL - - def get_world_scales(self, indices=None): - """Return per-prim (sx, sy, sz) scales extracted from world matrix. - - .. warning:: - When *indices* is None (all prims), the returned array is a **shared - pre-allocated buffer** (shared with :meth:`get_local_scales`) that is - overwritten on the next call. Do not hold references across calls -- - copy if persistence is needed. - - When *indices* selects a subset, Fabric launches the decompose kernel - into a freshly allocated Warp buffer and returns its - :class:`~isaaclab.utils.warp.ProxyArray` wrapper without blocking. - Callers that need host-visible values immediately must synchronize or - copy explicitly; GPU consumers can rely on normal Warp stream ordering. - """ + def _get_local_scales_impl(self, indices=None) -> ProxyArray: if not self._use_fabric: - return self._usd_view.get_world_scales(indices) + return self._usd_view._get_local_scales_impl(indices) if not self._fabric_initialized: self._initialize_fabric() - # Sync world matrices first if local writes are pending. - # TODO(pv): This dirty bit tracks Isaac Lab writes only. If Kit/Fabric - # hierarchy update_world_xforms() has already satisfied the dirty world - # matrices during a render tick, we currently have no Fabric-side version - # stamp to observe that and clear the flag; conservatively recompute. - self._sync_world_from_local_if_dirty() + return self._decompose_scales(self._get_local_ro_array(), 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] @@ -549,7 +324,7 @@ def get_world_scales(self, indices=None): kernel=fabric_utils.decompose_indexed_fabric_transforms, dim=count, inputs=[ - self._get_world_ro_array(), + ro_array, self._fabric_empty_2d_array_sentinel, self._fabric_empty_2d_array_sentinel, scales_wp, @@ -563,176 +338,69 @@ def get_world_scales(self, indices=None): return self._fabric_scales_ta return ProxyArray(scales_wp) - def set_local_scales(self, scales, indices=None): - """Set local-space scales by decomposing/recomposing localMatrix.""" - if not self._use_fabric: - self._usd_view.set_local_scales(scales, indices) - return - - if not self._fabric_initialized: - self._initialize_fabric() - - # Sync local matrices first if world writes are pending. - # TODO(pv): This dirty bit tracks Isaac Lab writes only. If a future - # Fabric hierarchy tick learns to materialize the corresponding local - # matrices before this access, we have no Fabric-side version stamp to - # observe that and clear the flag; conservatively recompute. - self._sync_local_from_world_if_dirty() - - indices_wp = self._resolve_indices_wp(indices) - scales_wp = self._to_float32_2d_or_empty(scales) - - wp.launch( - kernel=fabric_utils.compose_indexed_fabric_transforms, - dim=indices_wp.shape[0], - inputs=[ - self._get_local_rw_array(), - self._fabric_empty_2d_array_sentinel, - self._fabric_empty_2d_array_sentinel, - scales_wp, - False, - False, - False, - indices_wp, - ], - device=self._device, - ) - wp.synchronize() - - # Local was just written -- mark world matrices stale. See set_local_poses - # for why no eager world recompute is needed. - self._dirty = _DirtyFlag.WORLD - - def get_local_scales(self, indices=None): - """Return per-prim (sx, sy, sz) scales extracted from local matrix. - - .. warning:: - When *indices* is None (all prims), the returned array is a **shared - pre-allocated buffer** (shared with :meth:`get_world_scales`) that is - overwritten on the next call. Do not hold references across calls -- - copy if persistence is needed. - - When *indices* selects a subset, Fabric launches the decompose kernel - into a freshly allocated Warp buffer and returns its - :class:`~isaaclab.utils.warp.ProxyArray` wrapper without blocking. - Callers that need host-visible values immediately must synchronize or - copy explicitly; GPU consumers can rely on normal Warp stream ordering. - """ - if not self._use_fabric: - return self._usd_view.get_local_scales(indices) - - if not self._fabric_initialized: - self._initialize_fabric() - - # Sync local matrices first if world writes are pending. - # TODO(pv): This dirty bit tracks Isaac Lab writes only. If a future - # Fabric hierarchy tick learns to materialize the corresponding local - # matrices before this access, we have no Fabric-side version stamp to - # observe that and clear the flag; conservatively recompute. - self._sync_local_from_world_if_dirty() - - indices_wp = self._resolve_indices_wp(indices) - count = indices_wp.shape[0] - - use_cached = indices is None or indices == slice(None) - if use_cached: - scales_wp = self._fabric_scales_buf - else: - scales_wp = wp.zeros((count, 3), dtype=wp.float32, device=self._device) - - wp.launch( - kernel=fabric_utils.decompose_indexed_fabric_transforms, - dim=count, - inputs=[ - self._get_local_ro_array(), - self._fabric_empty_2d_array_sentinel, - self._fabric_empty_2d_array_sentinel, - scales_wp, - indices_wp, - ], - device=self._device, - ) - - if use_cached: - wp.synchronize() - return self._fabric_scales_ta - return ProxyArray(scales_wp) + # ------------------------------------------------------------------ + # Deprecated get_scales / set_scales hooks + # ------------------------------------------------------------------ - def _get_scales_impl(self, indices=None): - """Fabric: deprecated get_scales delegates to get_world_scales.""" - return self.get_world_scales(indices) + def _get_scales_impl(self, indices=None) -> ProxyArray: + """Fabric: deprecated get_scales returns world-space scales (legacy behavior).""" + return self._get_world_scales_impl(indices) - def _set_scales_impl(self, scales, indices=None): - """Fabric: deprecated set_scales delegates to set_world_scales.""" - self.set_world_scales(scales, indices) + def _set_scales_impl(self, scales, indices=None) -> None: + """Fabric: deprecated set_scales writes world-space scales via a one-shot writer scope.""" + with self.xform_space_writer("world") as writer: + writer.set_scales(scales, indices) # ------------------------------------------------------------------ - # Internal -- sync helpers + # Internal -- helpers shared by writers + initialization # ------------------------------------------------------------------ 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 _sync_world_from_local_if_dirty(self) -> None: - """If a prior local write left world matrices stale, recompute them.""" - if self._dirty != _DirtyFlag.WORLD: - return - self._recompute_world_from_local() - self._dirty = _DirtyFlag.NONE - - def _recompute_world_from_local(self) -> None: - """Recompute world matrices: child_world = parent_world * child_local. + def _recompute_local_from_world_all(self) -> None: + """Derive ``localMatrix = inv(parent) * worldMatrix`` for every prim in the view. - We deliberately do NOT call ``IFabricHierarchy.update_world_xforms()`` -- - in practice that re-reads USD's authored xformOps and overwrites the Fabric - local+world matrices we just authored. Instead we fire a Warp kernel that - does the multiply per child, leaving the Fabric-side localMatrix untouched. + 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`. """ - # Refresh trans_sel_ro once, then read _local_ifa_ro and _parent_world_ifa_ro - # directly to avoid calling PrepareForReuse twice on the same selection. if self._trans_sel_ro.PrepareForReuse() or self._parent_world_ifa_ro is None: self._rebuild_trans_ro_arrays() wp.launch( - kernel=fabric_utils.update_indexed_world_matrix_from_local, + kernel=fabric_utils.update_indexed_local_matrix_from_world, dim=self.count, inputs=[ - self._local_ifa_ro, + self._world_ifa_ro, self._parent_world_ifa_ro, - self._get_world_rw_array(), + self._get_local_rw_array(), self._view_indices, ], device=self._device, ) - wp.synchronize() - def _sync_local_from_world(self, indices_wp: wp.array) -> None: - """Recompute child localMatrix from (parent worldMatrix, child worldMatrix). + def _recompute_world_from_local_all(self) -> None: + """Derive ``worldMatrix = parent * localMatrix`` for every prim in the view. - Called after ``set_world_poses`` so that subsequent ``get_local_poses`` returns - values consistent with the just-written world poses. + 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`. """ - # Refresh trans_sel_ro once; _world_ifa_ro and _parent_world_ifa_ro share it. if self._trans_sel_ro.PrepareForReuse() or self._parent_world_ifa_ro is None: self._rebuild_trans_ro_arrays() wp.launch( - kernel=fabric_utils.update_indexed_local_matrix_from_world, - dim=indices_wp.shape[0], + kernel=fabric_utils.update_indexed_world_matrix_from_local, + dim=self.count, inputs=[ - self._world_ifa_ro, + self._local_ifa_ro, self._parent_world_ifa_ro, - self._get_local_rw_array(), - indices_wp, + self._get_world_rw_array(), + self._view_indices, ], device=self._device, ) - wp.synchronize() - - def _sync_local_from_world_if_dirty(self) -> None: - """If a prior world write left local matrices stale, recompute them lazily.""" - if self._dirty != _DirtyFlag.LOCAL: - return - self._sync_local_from_world(self._view_indices) - self._dirty = _DirtyFlag.NONE # ------------------------------------------------------------------ # Internal -- selection accessors with on-demand index rebuild @@ -864,14 +532,7 @@ def _initialize_fabric(self) -> None: ) # Ensure each child prim AND its parent have BOTH Fabric world and local matrix - # attributes. Our ``trans_ro`` selection requires both, so prims missing either - # would silently be excluded. ``Create*Attr`` calls are idempotent. - # - # ``SetWorldXformFromUsd`` writes Fabric's worldMatrix from USD's accumulated - # local-to-world transform (so it picks up the parent chain). - # ``SetLocalXformFromUsd`` writes Fabric's localMatrix from USD's authored - # xformOps on this prim only. Calling both gives Fabric a consistent - # (worldMatrix, localMatrix) pair for each prim before we touch the hierarchy. + # 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]): @@ -887,9 +548,7 @@ def _initialize_fabric(self) -> None: rt_xformable.SetLocalXformFromUsd() rt_xformable.SetWorldXformFromUsd() - # Three persistent selections with asymmetric access flags. This layout - # is load-bearing for correctness against Kit's ``update_world_xforms`` - # hierarchy update; see the class docstring for the why. + # Three persistent selections with asymmetric access flags. matrix = usdrt.Sdf.ValueTypeNames.Matrix4d ro = usdrt.Usd.Access.Read rw = usdrt.Usd.Access.ReadWrite @@ -901,8 +560,7 @@ def _initialize_fabric(self) -> None: self._world_sel_rw = self._stage.SelectPrims(require_attrs=[wm_rw, lm_ro], device=self._device, want_paths=True) self._local_sel_rw = self._stage.SelectPrims(require_attrs=[wm_ro, lm_rw], device=self._device, want_paths=True) - # Build the view-side indices array (just [0..count-1]) and a per-selection - # view->fabric mapping (selections do not guarantee a shared path ordering). + # Build the view-side indices array and per-selection view->fabric mappings. self._view_indices = wp.array(list(range(self.count)), dtype=wp.uint32, device=self._device) self._trans_ro_fabric_indices = self._compute_fabric_indices(self._trans_sel_ro) self._world_rw_fabric_indices = self._compute_fabric_indices(self._world_sel_rw) @@ -939,11 +597,7 @@ def _initialize_fabric(self) -> None: self._fabric_initialized = True - # Seed Fabric matrices from USD authoritatively. ``SetWorldXformFromUsd`` / - # ``SetLocalXformFromUsd`` are no-ops on freshly authored stages that haven't - # been rendered yet; we instead read through the USD view (children) and - # ``UsdGeom.XformCache`` (parents) and write via the same compose kernel that - # ``set_world_poses`` uses. + # Seed Fabric matrices from USD authoritatively. self._sync_fabric_from_usd_initial() def _sync_fabric_from_usd_initial(self) -> None: @@ -953,10 +607,7 @@ def _sync_fabric_from_usd_initial(self) -> None: 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. - # The child world matrix is computed at the end via - # ``_recompute_world_from_local()`` as ``child_world = parent_world * child_local``. + # --- 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( @@ -990,10 +641,8 @@ def _sync_fabric_from_usd_initial(self) -> None: for path in unique_parent_paths: prim = usd_stage.GetPrimAtPath(path) tf = xform_cache.GetLocalToWorldTransform(prim) - # Extract scale before ``Orthonormalize`` strips it from the rows. decomposer.SetMatrix(tf) s = decomposer.GetScale() - # Check for shear/skew: after removing scale, rows should be orthogonal. 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() @@ -1022,8 +671,6 @@ 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) - # Compose worldMatrix for parents (use a one-shot indexed array against - # ``_world_sel_rw`` keyed on the unique parent paths). parent_world_rw = wp.indexedfabricarray( fa=wp.fabricarray(self._world_sel_rw, self._WORLD_MATRIX_NAME), indices=self._compute_fabric_indices_for(self._world_sel_rw, unique_parent_paths), @@ -1046,8 +693,9 @@ def _sync_fabric_from_usd_initial(self) -> None: wp.synchronize() # After seeding local matrices from USD, recompute world matrices so - # the view starts with consistent state (child_world = parent_world * child_local). - self._recompute_world_from_local() + # 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: """Path-dict lookup helper used to build one-shot indexed arrays for a custom path set.""" @@ -1060,3 +708,206 @@ def _compute_fabric_indices_for(self, selection, paths: list[str]) -> wp.array: raise RuntimeError(f"Path '{path}' not found in Fabric selection.") indices.append(idx) return wp.array(indices, dtype=wp.int32, device=self._device) + + +# ---------------------------------------------------------------------- +# Concrete writer classes for FabricFrameView +# ---------------------------------------------------------------------- + + +class _FabricWriterMixin: + """Common ``__enter__`` / ``__exit__`` for the Fabric world / local writers. + + Pauses ``track_local_xform_changes`` / ``track_world_xform_changes`` on + the Fabric hierarchy while the scope is active so Kit does not redundantly + recompute the matrices we just wrote, then restores the prior state on + exit. + """ + + 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) + + def _exit_impl(self, exc_type, exc_val, exc_tb) -> None: + view: FabricFrameView = self._view # type: ignore[assignment] + try: + if self._wrote_anything and exc_type is None: + self._derive_opposite() + wp.synchronize() + finally: + 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_sel_rw``; 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.compose_indexed_fabric_transforms, + dim=indices_wp.shape[0], + inputs=[ + view._get_world_rw_array(), + positions_wp, + orientations_wp, + view._fabric_empty_2d_array_sentinel, + False, + False, + False, + indices_wp, + ], + device=view._device, + ) + self._wrote_anything = True + + 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_rw_array(), + 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 + + 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 _FabricLocalSpaceWriter(_FabricWriterMixin, FrameViewLocalSpaceWriter): + """Local-space writer for :class:`FabricFrameView`. + + Writes flow through ``_local_sel_rw``; 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.compose_indexed_fabric_transforms, + dim=indices_wp.shape[0], + inputs=[ + view._get_local_rw_array(), + translations_wp, + orientations_wp, + view._fabric_empty_2d_array_sentinel, + False, + False, + False, + indices_wp, + ], + device=view._device, + ) + self._wrote_anything = True + + 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_rw_array(), + 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 + + 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] + + +class _FabricFallbackWorldWriter(FrameViewWorldSpaceWriter): + """Fallback world-space writer used when Fabric is disabled. + + 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 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 a08432c9366d..9bd4a8cce30b 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 @@ -160,7 +160,8 @@ def test_fabric_set_world_does_not_write_back_to_usd(device, view_factory): # 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_space_writer("world") as w: + w.set_poses(positions=new_pos) # Verify Fabric has the new position fab_pos, _ = view.get_world_poses() @@ -199,7 +200,8 @@ def test_fabric_rebuild_after_topology_change(device, view_factory): # 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) + with view.xform_space_writer("world") as w: + w.set_poses(positions=initial) # Simulate topology change: recompute per-selection fabric indices and rebuild # every indexed array, mirroring the lazy paths in the ``_get_*_array`` accessors. @@ -216,7 +218,8 @@ def test_fabric_rebuild_after_topology_change(device, view_factory): # 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) + with view.xform_space_writer("world") as w: + w.set_poses(positions=new) ret_pos, _ = view.get_world_poses() pos_torch = torch.as_tensor(ret_pos, device=device) @@ -296,13 +299,14 @@ def test_set_local_via_fabric_path(device, view_factory): # Trigger lazy `_initialize_fabric()` so subsequent calls take the Fabric path. view.get_world_poses() - # Now set_local_poses should take the Fabric path + # 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) - view.set_local_poses(translations=new_local_pos, orientations=new_local_ori) + with view.xform_space_writer("local") 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() @@ -342,13 +346,8 @@ def test_local_scales_roundtrip(device, view_factory): 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_local_scales(new_scales) - - # Local writes mark world matrices stale; the renderer-facing world is - # recomputed lazily by Kit's next ``update_world_xforms`` tick (the - # ``_local_sel_rw`` selection has worldMatrix=RO, so that recompute is - # safe and the lazy design is correct). - assert view._dirty.name == "WORLD" + with view.xform_space_writer("local") as w: + w.set_scales(new_scales) ret_scales = view.get_local_scales() scales_torch = ret_scales.torch @@ -367,13 +366,8 @@ def test_world_scales_roundtrip(device, view_factory): 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) - view.set_world_scales(new_scales) - - # World writes mark local matrices stale; they are recomputed lazily on the - # next ``get_local_*`` call. The renderer-facing world write is protected - # by ``_world_sel_rw`` having localMatrix=RO, so Kit's next - # ``update_world_xforms`` tick will not clobber it from a stale local. - assert view._dirty.name == "LOCAL" + with view.xform_space_writer("world") as w: + w.set_scales(new_scales) ret_scales = view.get_world_scales() scales_torch = ret_scales.torch @@ -431,7 +425,8 @@ def test_set_local_then_get_world_with_rotated_parent(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)) - view.set_local_poses(translations=new_local, orientations=identity_quat) + with view.xform_space_writer("local") 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) @@ -452,7 +447,8 @@ def test_set_world_then_get_local_with_rotated_parent(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) - view.set_world_poses(positions=new_world) + with view.xform_space_writer("world") 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) @@ -508,33 +504,23 @@ def test_initial_seed_with_scaled_parent(device): # ------------------------------------------------------------------ -# Multi-view per stage: per-view dirty-flag isolation +# Multi-view per stage: per-view single-writer isolation # ------------------------------------------------------------------ @pytest.mark.parametrize("device", ["cpu", "cuda:0"]) -def test_multi_view_per_view_dirty_isolation(device): - """Two ``FabricFrameView`` instances on the same stage must not clear each other's - pending local→world sync. - - Background: an earlier implementation stored the world-dirty flag at the class - level keyed by ``stage_id``. With two views on the same stage, view B reading - worlds would clear the flag set by view A's ``set_local_poses``, leaving A's - world matrices silently stale because A's per-view sync kernel never fired. - - This test sets up two views over disjoint child prims (under different parent - sub-trees of the same stage), interleaves their writes and reads, and verifies: - - * view A's ``set_local_poses`` only dirties view A - * view B's ``get_world_poses`` does not clear view A's flag - * after both views' world reads, each one's worlds reflect its own latest local - * neither view's reads/writes corrupt the other view's poses +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() - # Two disjoint sub-trees under the same stage. Use different parent names so - # the regex patterns for the two views don't accidentally overlap. 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) @@ -544,8 +530,6 @@ def test_multi_view_per_view_dirty_isolation(device): view_a = FrameView("/World/EnvA_.*/ChildA", device=device) view_b = FrameView("/World/EnvB_.*/ChildB", device=device) - # Initial reads -- triggers Fabric init + the seed-time ``_dirty = WORLD`` - # path on both views, then clears it. 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( @@ -554,42 +538,30 @@ def test_multi_view_per_view_dirty_isolation(device): torch.testing.assert_close( torch.as_tensor(view_b.get_world_poses()[0], device=device), expected_b0, atol=1e-5, rtol=0 ) - assert view_a._dirty.name == "NONE" - assert view_b._dirty.name == "NONE" - # Write a new local pose on view A only. + # 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)) - view_a.set_local_poses(translations=new_local_a, orientations=identity_quat) - - # Local writes are lazy: view A is dirty (WORLD), view B must not be touched. - assert view_a._dirty.name == "WORLD", "set_local_poses must mark its own view dirty" - assert view_b._dirty.name == "NONE", "set_local_poses on view A must not dirty view B" + with view_a.xform_space_writer("local") as w: + w.set_poses(positions=new_local_a, orientations=identity_quat) - # Read worlds from view B FIRST. This must not clear view A's dirty flag. + # 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 ) - assert view_b._dirty.name == "NONE" - assert view_a._dirty.name == "WORLD", "view B's read must not clear view A's dirty flag" - # Now read view A's worlds -- triggers the per-view local->world flush. + # 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 ) - assert view_a._dirty.name == "NONE" - # Symmetric pass: write on B, ensure A is undisturbed. + # 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) - view_b.set_local_poses(translations=new_local_b, orientations=identity_quat) - assert view_a._dirty.name == "NONE" - assert view_b._dirty.name == "WORLD" - - # A's worlds must still read back the post-flush value from above; no cross-view - # stomp on the world matrix. + with view_b.xform_space_writer("local") 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 ) @@ -597,8 +569,16 @@ def test_multi_view_per_view_dirty_isolation(device): torch.testing.assert_close( torch.as_tensor(view_b.get_world_poses()[0], device=device), expected_b1, atol=1e-5, rtol=0 ) - assert view_a._dirty.name == "NONE" - assert view_b._dirty.name == "NONE" + + # Single-active-writer is per-view: opening a writer on A leaves B free. + with view_a.xform_space_writer("world"): + 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_space_writer("world"): + assert view_b._active_writer is not None + assert view_a._active_writer is None + assert view_b._active_writer is None # ------------------------------------------------------------------ @@ -622,7 +602,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_space_writer("world") as w: + w.set_poses(positions=new_pos) ret_pos, _ = view.get_world_poses() pos_torch = torch.as_tensor(ret_pos, device=device) @@ -652,7 +633,8 @@ 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_space_writer("world") as w: + w.set_poses(positions=new_pos) # USD must not have moved at all -- equality, not approximate. t_after = UsdGeom.XformCache().GetLocalToWorldTransform(prim).ExtractTranslation() @@ -679,7 +661,8 @@ def test_fabric_cuda1_scales_roundtrip(device, view_factory): 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_world_scales(new_scales) + with view.xform_space_writer("world") as w: + w.set_scales(new_scales) ret_scales = view.get_world_scales() scales_torch = ret_scales.torch @@ -688,7 +671,7 @@ def test_fabric_cuda1_scales_roundtrip(device, view_factory): # ------------------------------------------------------------------ -# Interleaved set_world_poses / set_local_poses tests +# Sequential writer scopes (interleaved world / local writes via two scopes) # ------------------------------------------------------------------ @@ -715,29 +698,29 @@ def _build_two_child_view(device: str) -> "FrameView": @pytest.mark.parametrize("device", ["cpu", "cuda:0"]) -def test_interleaved_set_world_then_set_local_partial_indices(device): - """set_world_poses on index 0, then set_local_poses on index 1 -- both must be correct. - - This exercises the dirty-flag flush: after set_world_poses marks _dirty == LOCAL, - set_local_poses must flush stale locals before writing index 1, ensuring index 0's - local is correctly derived from its new world pose. - """ +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) - # Step 1: set world pose on index 0 only 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)) - view.set_world_poses(positions=new_world_pos, indices=idx0) + with view.xform_space_writer("world") as w: + w.set_poses(positions=new_world_pos, indices=idx0) - # Step 2: set local pose on index 1 only 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)) - view.set_local_poses(translations=new_local_pos, orientations=identity_quat, indices=idx1) - - # Verify index 0's world pose is still (5, 0, 2) + with view.xform_space_writer("local") 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), @@ -746,7 +729,7 @@ def test_interleaved_set_world_then_set_local_partial_indices(device): rtol=0, ) - # Verify index 0's local pose (derived from world): + # 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( @@ -756,7 +739,7 @@ def test_interleaved_set_world_then_set_local_partial_indices(device): rtol=0, ) - # Verify index 1's world pose (derived from local): + # 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( @@ -768,28 +751,24 @@ def test_interleaved_set_world_then_set_local_partial_indices(device): @pytest.mark.parametrize("device", ["cpu", "cuda:0"]) -def test_interleaved_set_local_then_set_world_partial_indices(device): - """set_local_poses on index 0, then set_world_poses on index 1 -- both must be correct. - - The reverse direction of the above: after set_local_poses marks _dirty = WORLD, - set_world_poses must flush stale worlds before writing index 1. - """ +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) - # Step 1: set local pose on index 0 only 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)) - view.set_local_poses(translations=new_local_pos, orientations=identity_quat, indices=idx0) + with view.xform_space_writer("local") as w: + w.set_poses(positions=new_local_pos, orientations=identity_quat, indices=idx0) - # Step 2: set world pose on index 1 only 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)) - view.set_world_poses(positions=new_world_pos, indices=idx1) + with view.xform_space_writer("world") as w: + w.set_poses(positions=new_world_pos, indices=idx1) - # Verify index 0's world pose (derived from local): + # 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( @@ -799,7 +778,7 @@ def test_interleaved_set_local_then_set_world_partial_indices(device): rtol=0, ) - # Verify index 1's world pose is still (10, 20, 30) + # 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), @@ -808,7 +787,7 @@ def test_interleaved_set_local_then_set_world_partial_indices(device): rtol=0, ) - # Verify index 1's local (derived from world): + # 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( @@ -819,57 +798,194 @@ def test_interleaved_set_local_then_set_world_partial_indices(device): ) +# ------------------------------------------------------------------ +# FrameViewSpaceWriterBase contract tests +# ------------------------------------------------------------------ + + @pytest.mark.parametrize("device", ["cpu", "cuda:0"]) -def test_interleaved_set_emits_warning(device, caplog): - """Interleaving set_world_poses and set_local_poses logs a one-time warning. +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 - With the lazy design, a set_world_poses leaves localMatrix stale; a subsequent - set_local_poses on a different (possibly disjoint) index subset must flush the - stale local first to avoid silent inconsistency. That extra kernel launch is - a performance hazard and is flagged once per view instance. - """ - view = _build_two_child_view(device) + 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_space_writer("world") 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) - # First set_world_poses -- no warning (first user setter). - new_world = wp.zeros((2, 3), dtype=wp.float32, device=device) - wp.launch(kernel=_fill_position, dim=2, inputs=[new_world, 1.0, 2.0, 3.0], device=device) - view.set_world_poses(positions=new_world) - # Now set_local_poses -- should trigger warning about interleaving. - new_local = wp.zeros((2, 3), dtype=wp.float32, device=device) - wp.launch(kernel=_fill_position, dim=2, inputs=[new_local, 0.0, 0.0, 0.0], device=device) - identity_quat = wp.from_torch(torch.tensor([[0.0, 0.0, 0.0, 1.0]] * 2, dtype=torch.float32, device=device)) +@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 caplog.at_level("WARNING", logger="isaaclab_physx.sim.views.fabric_frame_view"): - caplog.clear() - view.set_local_poses(translations=new_local, orientations=identity_quat) + with view.xform_space_writer("local") as w: + w.set_poses(positions=new_pos, orientations=identity_quat) - assert any("interleaving" in r.message.lower() for r in caplog.records), ( - f"Expected interleave warning, got: {[r.message for r in caplog.records]}" - ) + 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) - # Second interleave -- warning should NOT repeat (one-time only). - caplog.clear() - view.set_world_poses(positions=new_world) - assert not any("interleaving" in r.message.lower() for r in caplog.records), ( - "Warning should only fire once per view instance" - ) + # 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_space_writer("world") 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_no_warning_when_using_single_setter(device, caplog): - """Calling only set_world_poses (or only set_local_poses) should never warn.""" - view = _build_two_child_view(device) +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() - new_world = wp.zeros((2, 3), dtype=wp.float32, device=device) - wp.launch(kernel=_fill_position, dim=2, inputs=[new_world, 1.0, 2.0, 3.0], device=device) + with view.xform_space_writer("world"): + with pytest.raises(RuntimeError, match="already has an active xform_space_writer"): + view.xform_space_writer("world").__enter__() + with pytest.raises(RuntimeError, match="already has an active xform_space_writer"): + view.xform_space_writer("local").__enter__() + # After the outer scope exits, the lock is released and a new scope succeeds. + with view.xform_space_writer("local"): + pass - with caplog.at_level("WARNING", logger="isaaclab_physx.sim.views.fabric_frame_view"): - caplog.clear() - view.set_world_poses(positions=new_world) - view.set_world_poses(positions=new_world) - view.set_world_poses(positions=new_world) - assert not any("interleaving" in r.message.lower() for r in caplog.records), ( - f"Unexpected interleave warning with single setter: {[r.message for r in caplog.records]}" - ) +@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_space_writer("world"): + 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_space_writer("world"): + pass + assert h.tracking_local_xform_changes, "writer must restore the pre-enabled local listener" + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_writer_invalid_space_raises(device, view_factory): + """``xform_space_writer`` with an invalid space raises ``ValueError``.""" + bundle = view_factory(num_envs=1, device=device) + view = bundle.view + with pytest.raises(ValueError, match="Invalid space"): + view.xform_space_writer("foo") + + +@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_space_writer("world"): + 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_space_writer("world"): + with pytest.raises(RuntimeError, match="xform_space_writer"): + view.get_world_poses() + with pytest.raises(RuntimeError, match="xform_space_writer"): + view.get_local_poses() + with pytest.raises(RuntimeError, match="xform_space_writer"): + view.get_world_scales() + with pytest.raises(RuntimeError, match="xform_space_writer"): + view.get_local_scales() + # After the scope exits, view getters work again. + view.get_world_poses() + + +def test_set_world_scales_method_no_longer_exists(): + """``set_world_scales`` / ``set_local_scales`` were deleted from this PR's surface.""" + # The deprecated set_world_poses / set_local_poses shims remain (with warnings), + # but the never-shipped set_world_scales / set_local_scales were removed. + from isaaclab.sim.views import BaseFrameView # noqa: PLC0415 + + assert not hasattr(BaseFrameView, "set_world_scales") + assert not hasattr(BaseFrameView, "set_local_scales") From 0c47acb72b68a1b1bb8fb1f0b8732d780a71ddbc Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Mon, 22 Jun 2026 17:37:43 +0000 Subject: [PATCH 40/54] refactor: split xform_space_writer(space) into xform_world_space_writer() / xform_local_space_writer() Replace the single dispatch method that took a 'world'|'local' string with two explicit methods. More discoverable, better typed (each returns its concrete writer class), and removes the invalid-space ValueError path. Update all callsites (camera, USD/Fabric/Newton/OVPhysX backends, mimic), benchmarks, tests, and changelog fragments. Drop the now-meaningless test_writer_invalid_space_raises. --- .../benchmarks/benchmark_view_comparison.py | 2 +- .../benchmarks/benchmark_xform_prim_view.py | 10 +-- .../changelog.d/xform-space-writer.rst | 11 ++- .../isaaclab/sensors/camera/camera.py | 4 +- .../isaaclab/sim/views/base_frame_view.py | 75 ++++++++++------- .../isaaclab/sim/views/usd_frame_view.py | 5 +- .../isaaclab/sim/views/xform_space_writer.py | 7 +- .../test/sim/frame_view_contract_utils.py | 26 +++--- .../test/sim/test_views_xform_prim.py | 10 +-- .../test/terrains/check_terrain_importer.py | 2 +- .../test/terrains/test_terrain_importer.py | 2 +- .../locomanipulation_sdg/scene_utils.py | 2 +- .../changelog.d/xform-space-writer.rst | 3 +- .../test/sim/test_views_xform_prim_newton.py | 2 +- .../changelog.d/xform-space-writer.rst | 3 +- .../sim/views/ovphysx_frame_view.py | 2 +- .../sim/views/fabric_frame_view.py | 6 +- .../test/sim/test_views_xform_prim_fabric.py | 81 +++++++++---------- 18 files changed, 135 insertions(+), 118 deletions(-) diff --git a/scripts/benchmarks/benchmark_view_comparison.py b/scripts/benchmarks/benchmark_view_comparison.py index 8270be2cd225..2f19554d5ffc 100644 --- a/scripts/benchmarks/benchmark_view_comparison.py +++ b/scripts/benchmarks/benchmark_view_comparison.py @@ -284,7 +284,7 @@ def _run_pose_benchmarks( start_time = time.perf_counter() for _ in range(num_iterations): - with view.xform_space_writer("world") as w: + 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 diff --git a/scripts/benchmarks/benchmark_xform_prim_view.py b/scripts/benchmarks/benchmark_xform_prim_view.py index c3ac1228feee..ae6656e3d8d0 100644 --- a/scripts/benchmarks/benchmark_xform_prim_view.py +++ b/scripts/benchmarks/benchmark_xform_prim_view.py @@ -177,7 +177,7 @@ def to_torch(a): torch.cuda.synchronize() start_time = time.perf_counter() for _ in range(num_iterations): - with xform_view.xform_space_writer("world") as w: + with xform_view.xform_world_space_writer() as w: w.set_poses(new_positions, orientations) if is_newton: torch.cuda.synchronize() @@ -214,7 +214,7 @@ def to_torch(a): torch.cuda.synchronize() start_time = time.perf_counter() for _ in range(num_iterations): - with xform_view.xform_space_writer("local") as w: + with xform_view.xform_local_space_writer() as w: w.set_poses(new_translations, orientations_local) if is_newton: torch.cuda.synchronize() @@ -249,7 +249,7 @@ def to_torch(a): torch.cuda.synchronize() start_time = time.perf_counter() for _ in range(num_iterations): - with xform_view.xform_space_writer("world") as w: + with xform_view.xform_world_space_writer() as w: w.set_scales(new_world_scales) if is_newton: torch.cuda.synchronize() @@ -282,7 +282,7 @@ def to_torch(a): torch.cuda.synchronize() start_time = time.perf_counter() for _ in range(num_iterations): - with xform_view.xform_space_writer("local") as w: + with xform_view.xform_local_space_writer() as w: w.set_scales(new_local_scales) if is_newton: torch.cuda.synchronize() @@ -306,7 +306,7 @@ def to_torch(a): torch.cuda.synchronize() start_time = time.perf_counter() for _ in range(num_iterations): - with xform_view.xform_space_writer("world") as w: + with xform_view.xform_world_space_writer() as w: w.set_poses(new_positions, orientations) xform_view.get_world_poses() if is_newton: diff --git a/source/isaaclab/changelog.d/xform-space-writer.rst b/source/isaaclab/changelog.d/xform-space-writer.rst index 3a3236e042cf..41c9c97a55cc 100644 --- a/source/isaaclab/changelog.d/xform-space-writer.rst +++ b/source/isaaclab/changelog.d/xform-space-writer.rst @@ -3,7 +3,7 @@ Added * Added :class:`~isaaclab.sim.views.FrameViewSpaceWriterBase`, the new context-managed write API for ``FrameView``-managed prim transforms. Open with - ``view.xform_space_writer("world" | "local")`` and call + ``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 @@ -15,14 +15,16 @@ Added * 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_space_writer`. + :meth:`~isaaclab.sim.views.BaseFrameView.xform_world_space_writer` / + :meth:`~isaaclab.sim.views.BaseFrameView.xform_local_space_writer`. Deprecated ^^^^^^^^^^ * Deprecated :meth:`~isaaclab.sim.views.BaseFrameView.set_world_poses` and :meth:`~isaaclab.sim.views.BaseFrameView.set_local_poses`. Use - ``with view.xform_space_writer("world" | "local") as w: w.set_poses(...)`` + ``with view.xform_world_space_writer() as w: w.set_poses(...)`` (or + :meth:`~isaaclab.sim.views.BaseFrameView.xform_local_space_writer`) instead. The deprecated methods still work but emit a one-time ``DeprecationWarning`` per class and open a single-statement writer scope internally. @@ -34,5 +36,6 @@ Removed from :class:`~isaaclab.sim.views.BaseFrameView` (and all subclasses). These were introduced in this release cycle without a stable downstream user, so they are removed outright (no deprecation cycle). Use - ``with view.xform_space_writer("world" | "local") as w: w.set_scales(...)`` + ``with view.xform_world_space_writer() as w: w.set_scales(...)`` (or + :meth:`~isaaclab.sim.views.BaseFrameView.xform_local_space_writer`) instead. diff --git a/source/isaaclab/isaaclab/sensors/camera/camera.py b/source/isaaclab/isaaclab/sensors/camera/camera.py index 4269b1fd25c0..58cf4423eaf5 100644 --- a/source/isaaclab/isaaclab/sensors/camera/camera.py +++ b/source/isaaclab/isaaclab/sensors/camera/camera.py @@ -376,7 +376,7 @@ 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) - with self._view.xform_space_writer("world") as writer: + with self._view.xform_world_space_writer() as writer: writer.set_poses(pos_wp, ori_wp, idx_wp) def set_world_poses_from_view( @@ -435,7 +435,7 @@ 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) - with self._view.xform_space_writer("world") as writer: + 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), diff --git a/source/isaaclab/isaaclab/sim/views/base_frame_view.py b/source/isaaclab/isaaclab/sim/views/base_frame_view.py index cda64dd4f435..dc7cc324f53a 100644 --- a/source/isaaclab/isaaclab/sim/views/base_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/base_frame_view.py @@ -9,7 +9,7 @@ import abc import warnings -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING import warp as wp @@ -28,11 +28,12 @@ class BaseFrameView(abc.ABC): implementation at runtime based on the active physics backend. All getters return :class:`~isaaclab.utils.warp.ProxyArray`. All writes go - through :meth:`xform_space_writer` -- the recommended API: + through the writer-scope API -- :meth:`xform_world_space_writer` or + :meth:`xform_local_space_writer`: .. code-block:: python - with view.xform_space_writer("world") as writer: + 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. @@ -66,33 +67,49 @@ def device(self) -> str: # Write scope -- recommended API for all transform writes. # ------------------------------------------------------------------ - def xform_space_writer(self, space: Literal["world", "local"]) -> FrameViewSpaceWriterBase: - """Open a write scope on this view (recommended write API). + def xform_world_space_writer(self) -> FrameViewWorldSpaceWriter: + """Open a world-space write scope on this view (recommended write API). - Args: - space: ``"world"`` or ``"local"``. + Inside the scope, :meth:`~FrameViewSpaceWriterBase.set_poses` / + :meth:`~FrameViewSpaceWriterBase.set_scales` write world-space values. Returns: - An :class:`~isaaclab.sim.views.FrameViewWorldSpaceWriter` or - :class:`~isaaclab.sim.views.FrameViewLocalSpaceWriter` context manager. + A :class:`~isaaclab.sim.views.FrameViewWorldSpaceWriter` context manager. Raises: - ValueError: If *space* is neither ``"world"`` nor ``"local"``. RuntimeError: On ``__enter__``, if another writer is already active on this view. Example: .. code-block:: python - with view.xform_space_writer("world") as w: + with view.xform_world_space_writer() as w: w.set_poses(positions=p, orientations=o) w.set_scales(scales=s) """ - if space == "world": - return self._make_world_space_writer() - if space == "local": - return self._make_local_space_writer() - raise ValueError(f"Invalid space {space!r}; expected 'world' or 'local'.") + 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: @@ -108,7 +125,7 @@ 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 an xform_space_writer " + 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." ) @@ -229,7 +246,7 @@ def set_world_poses( """Set world-space positions and/or orientations for prims in the view. .. deprecated:: - Use ``with view.xform_space_writer("world") as w: w.set_poses(...)`` instead. + Use ``with view.xform_world_space_writer() as w: w.set_poses(...)`` instead. This method opens a single-statement writer scope internally. Args: @@ -240,12 +257,12 @@ def set_world_poses( if not BaseFrameView._set_world_poses_deprecated_warned: BaseFrameView._set_world_poses_deprecated_warned = True warnings.warn( - 'set_world_poses() is deprecated. Use \'with view.xform_space_writer("world") as w:' + "set_world_poses() is deprecated. Use 'with view.xform_world_space_writer() as w:" " w.set_poses(...)' instead.", DeprecationWarning, stacklevel=2, ) - with self.xform_space_writer("world") as writer: + with self.xform_world_space_writer() as writer: writer.set_poses(positions, orientations, indices) def set_local_poses( @@ -257,7 +274,7 @@ def set_local_poses( """Set local-space translations and/or orientations for prims in the view. .. deprecated:: - Use ``with view.xform_space_writer("local") as w: w.set_poses(...)`` instead. + Use ``with view.xform_local_space_writer() as w: w.set_poses(...)`` instead. This method opens a single-statement writer scope internally. Args: @@ -268,12 +285,12 @@ def set_local_poses( if not BaseFrameView._set_local_poses_deprecated_warned: BaseFrameView._set_local_poses_deprecated_warned = True warnings.warn( - 'set_local_poses() is deprecated. Use \'with view.xform_space_writer("local") as w:' + "set_local_poses() is deprecated. Use 'with view.xform_local_space_writer() as w:" " w.set_poses(...)' instead.", DeprecationWarning, stacklevel=2, ) - with self.xform_space_writer("local") as writer: + with self.xform_local_space_writer() as writer: writer.set_poses(translations, orientations, indices) # ------------------------------------------------------------------ @@ -311,10 +328,10 @@ def set_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: """Set scales for prims in the view. .. deprecated:: - Use ``with view.xform_space_writer("world" | "local") as w: w.set_scales(...)`` instead. - This 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``. + Use ``with view.xform_world_space_writer() as w: w.set_scales(...)`` (or + :meth:`xform_local_space_writer`) instead. This 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``. Args: scales: Scales ``(M, 3)`` as ``wp.array``. @@ -323,8 +340,8 @@ def set_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: if not BaseFrameView._set_scales_deprecated_warned: BaseFrameView._set_scales_deprecated_warned = True warnings.warn( - 'set_scales() is deprecated. Use \'with view.xform_space_writer("world" or' - ' "local") as w: w.set_scales(...)\' instead.', + "set_scales() is deprecated. Use 'with view.xform_world_space_writer() as w:" + " w.set_scales(...)' (or xform_local_space_writer()) instead.", DeprecationWarning, stacklevel=2, ) diff --git a/source/isaaclab/isaaclab/sim/views/usd_frame_view.py b/source/isaaclab/isaaclab/sim/views/usd_frame_view.py index 8cf6bbe816d0..7145d4e07b00 100644 --- a/source/isaaclab/isaaclab/sim/views/usd_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/usd_frame_view.py @@ -36,7 +36,8 @@ class UsdFrameView(BaseFrameView): For GPU-accelerated Fabric operations, use the PhysX backend variant obtained via :class:`~isaaclab.sim.views.FrameView`. - All writes go through :meth:`xform_space_writer` (recommended). The + 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 @@ -374,7 +375,7 @@ def _get_scales_impl(self, indices: wp.array | None = None) -> ProxyArray: def _set_scales_impl(self, scales: wp.array, indices: wp.array | None = None) -> None: """USD legacy: deprecated set_scales writes local scales via a one-shot writer scope.""" - with self.xform_space_writer("local") as writer: + with self.xform_local_space_writer() as writer: writer.set_scales(scales, indices) # ------------------------------------------------------------------ diff --git a/source/isaaclab/isaaclab/sim/views/xform_space_writer.py b/source/isaaclab/isaaclab/sim/views/xform_space_writer.py index 4be224ef679c..17c437294a0b 100644 --- a/source/isaaclab/isaaclab/sim/views/xform_space_writer.py +++ b/source/isaaclab/isaaclab/sim/views/xform_space_writer.py @@ -9,7 +9,7 @@ .. code-block:: python - with view.xform_space_writer("world") as writer: + with view.xform_world_space_writer() as writer: writer.set_poses(positions=p, orientations=o) writer.set_scales(scales=s) # ... any number of writes ... @@ -40,7 +40,8 @@ class FrameViewSpaceWriterBase(abc.ABC): """Abstract context-managed writer for a single transform space. - Subclasses are returned by :meth:`BaseFrameView.xform_space_writer`; they + 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 @@ -94,7 +95,7 @@ def get_scales(self, indices: wp.array | None = None) -> ProxyArray: def __enter__(self) -> FrameViewSpaceWriterBase: if self._view._active_writer is not None: raise RuntimeError( - f"{type(self._view).__name__} already has an active xform_space_writer scope " + 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." ) diff --git a/source/isaaclab/test/sim/frame_view_contract_utils.py b/source/isaaclab/test/sim/frame_view_contract_utils.py index 2741eb955ae5..e67330f36791 100644 --- a/source/isaaclab/test/sim/frame_view_contract_utils.py +++ b/source/isaaclab/test/sim/frame_view_contract_utils.py @@ -193,7 +193,7 @@ 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) - with bundle.view.xform_space_writer("world") as w: + 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() @@ -210,7 +210,7 @@ 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) - with bundle.view.xform_space_writer("local") as w: + 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() @@ -226,7 +226,7 @@ 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() - with bundle.view.xform_space_writer("world") as w: + 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), @@ -244,7 +244,7 @@ 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() - with bundle.view.xform_space_writer("local") as w: + 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), @@ -268,7 +268,7 @@ 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 - with bundle.view.xform_space_writer("world") as w: + 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), @@ -290,7 +290,7 @@ 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) - with bundle.view.xform_space_writer("local") as w: + 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), @@ -309,7 +309,7 @@ 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) - with bundle.view.xform_space_writer("world") as w: + with bundle.view.xform_world_space_writer() as w: w.set_poses(positions=new_pos) ret_pos, ret_quat = bundle.view.get_world_poses() @@ -326,7 +326,7 @@ 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) - with bundle.view.xform_space_writer("world") as w: + with bundle.view.xform_world_space_writer() as w: w.set_poses(orientations=new_quat) ret_pos, ret_quat = bundle.view.get_world_poses() @@ -343,7 +343,7 @@ 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) - with bundle.view.xform_space_writer("local") as w: + with bundle.view.xform_local_space_writer() as w: w.set_poses(positions=new_pos) ret_pos, ret_quat = bundle.view.get_local_poses() @@ -361,7 +361,7 @@ 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) - with bundle.view.xform_space_writer("world") as w: + 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]) @@ -474,7 +474,7 @@ def test_set_local_scales_roundtrip(device, view_factory): 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_space_writer("local") as w: + with bundle.view.xform_local_space_writer() as w: w.set_scales(new_scales) ret_scales = _t(bundle.view.get_local_scales()) @@ -489,7 +489,7 @@ def test_set_world_scales_roundtrip(device, view_factory): 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_space_writer("world") as w: + with bundle.view.xform_world_space_writer() as w: w.set_scales(new_scales) ret_scales = _t(bundle.view.get_world_scales()) @@ -507,7 +507,7 @@ def test_local_scales_do_not_affect_local_poses(device, view_factory): 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_space_writer("local") as w: + with bundle.view.xform_local_space_writer() as w: w.set_scales(new_scales) local_pos_after = _t(bundle.view.get_local_poses()[0]) diff --git a/source/isaaclab/test/sim/test_views_xform_prim.py b/source/isaaclab/test/sim/test_views_xform_prim.py index 099fd127a66f..2fa404f26a7f 100644 --- a/source/isaaclab/test/sim/test_views_xform_prim.py +++ b/source/isaaclab/test/sim/test_views_xform_prim.py @@ -226,9 +226,9 @@ def test_nested_hierarchy_world_poses(device): frames_view = FrameView("/World/Frame_.*", device=device) targets_view = FrameView("/World/Frame_.*/Target", device=device) - with frames_view.xform_space_writer("local") as w: + with frames_view.xform_local_space_writer() as w: w.set_poses(positions=torch.tensor(frame_positions, device=device)) - with targets_view.xform_space_writer("local") as w: + 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 @@ -267,7 +267,7 @@ def test_set_local_scales_then_get_world_scales(device): 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_space_writer("local") as w: + with view.xform_local_space_writer() as w: w.set_scales(local_scales) world_scales = view.get_world_scales().torch @@ -283,7 +283,7 @@ def test_set_world_scales_then_get_local_scales(device): 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_space_writer("world") as w: + with view.xform_world_space_writer() as w: w.set_scales(world_scales) local_scales = view.get_local_scales().torch @@ -345,7 +345,7 @@ 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) - with view.xform_space_writer("world") as w: + 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 diff --git a/source/isaaclab/test/terrains/check_terrain_importer.py b/source/isaaclab/test/terrains/check_terrain_importer.py index 951bbe3a3e54..901cfc420be3 100644 --- a/source/isaaclab/test/terrains/check_terrain_importer.py +++ b/source/isaaclab/test/terrains/check_terrain_importer.py @@ -159,7 +159,7 @@ def main(): ball_initial_positions = terrain_importer.env_origins.clone() ball_initial_positions[:, 2] += 5.0 # set initial poses (writes to USD before simulation) - with xform_view.xform_space_writer("world") as w: + with xform_view.xform_world_space_writer() as w: w.set_poses(positions=ball_initial_positions) # Play simulator diff --git a/source/isaaclab/test/terrains/test_terrain_importer.py b/source/isaaclab/test/terrains/test_terrain_importer.py index c712812bee60..96375cab0758 100644 --- a/source/isaaclab/test/terrains/test_terrain_importer.py +++ b/source/isaaclab/test/terrains/test_terrain_importer.py @@ -316,5 +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 :) - with ball_view.xform_space_writer("world") as w: + 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/isaaclab_mimic/locomanipulation_sdg/scene_utils.py b/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py index b4fca471f981..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,7 @@ def set_pose(self, pose: torch.Tensor): xform_prim = self._get_xform_view() position = pose[..., :3] orientation = pose[..., 3:] - with xform_prim.xform_space_writer("world") as writer: + with xform_prim.xform_world_space_writer() as writer: writer.set_poses(wp.from_torch(position.contiguous()), wp.from_torch(orientation.contiguous()), None) diff --git a/source/isaaclab_newton/changelog.d/xform-space-writer.rst b/source/isaaclab_newton/changelog.d/xform-space-writer.rst index 85df80394b1f..e04d5325e6aa 100644 --- a/source/isaaclab_newton/changelog.d/xform-space-writer.rst +++ b/source/isaaclab_newton/changelog.d/xform-space-writer.rst @@ -4,7 +4,8 @@ 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_space_writer` context API. + :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 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 d743140c48b3..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,7 @@ 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) - with view.xform_space_writer("world") as w: + with view.xform_world_space_writer() as w: w.set_poses(new_pos, new_quat) ret_pos, ret_quat = view.get_world_poses() diff --git a/source/isaaclab_ovphysx/changelog.d/xform-space-writer.rst b/source/isaaclab_ovphysx/changelog.d/xform-space-writer.rst index 16b170753e7a..3561232a311e 100644 --- a/source/isaaclab_ovphysx/changelog.d/xform-space-writer.rst +++ b/source/isaaclab_ovphysx/changelog.d/xform-space-writer.rst @@ -4,7 +4,8 @@ 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_space_writer` context API. + :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 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 e293e14bbb7a..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 @@ -844,7 +844,7 @@ def _get_scales_impl(self, indices=None): 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_space_writer("local") as writer: + with self.xform_local_space_writer() as writer: writer.set_scales(scales, indices) def get_visibility(self, indices: wp.array | None = None): 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 e0dcd055b181..8d4d64720349 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 @@ -52,7 +52,9 @@ class FabricFrameView(BaseFrameView): ``omni:fabric:worldMatrix`` and ``omni:fabric:localMatrix`` directly. All other operations delegate to the internal USD view. - All writes go through :meth:`xform_space_writer` (recommended) or the + All writes go through the writer-scope API + (:meth:`xform_world_space_writer` / :meth:`xform_local_space_writer`, + recommended) or the deprecated :meth:`set_world_poses` / :meth:`set_local_poses` / etc. shims inherited from :class:`BaseFrameView`. @@ -348,7 +350,7 @@ def _get_scales_impl(self, indices=None) -> ProxyArray: def _set_scales_impl(self, scales, indices=None) -> None: """Fabric: deprecated set_scales writes world-space scales via a one-shot writer scope.""" - with self.xform_space_writer("world") as writer: + with self.xform_world_space_writer() as writer: writer.set_scales(scales, indices) # ------------------------------------------------------------------ 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 9bd4a8cce30b..4b030abc20c6 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 @@ -160,7 +160,7 @@ def test_fabric_set_world_does_not_write_back_to_usd(device, view_factory): # 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) - with view.xform_space_writer("world") as w: + with view.xform_world_space_writer() as w: w.set_poses(positions=new_pos) # Verify Fabric has the new position @@ -200,7 +200,7 @@ def test_fabric_rebuild_after_topology_change(device, view_factory): # 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) - with view.xform_space_writer("world") as w: + with view.xform_world_space_writer() as w: w.set_poses(positions=initial) # Simulate topology change: recompute per-selection fabric indices and rebuild @@ -218,7 +218,7 @@ def test_fabric_rebuild_after_topology_change(device, view_factory): # 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) - with view.xform_space_writer("world") as w: + with view.xform_world_space_writer() as w: w.set_poses(positions=new) ret_pos, _ = view.get_world_poses() @@ -305,7 +305,7 @@ def test_set_local_via_fabric_path(device, view_factory): 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_space_writer("local") as w: + 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) @@ -346,7 +346,7 @@ def test_local_scales_roundtrip(device, view_factory): 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_space_writer("local") as w: + with view.xform_local_space_writer() as w: w.set_scales(new_scales) ret_scales = view.get_local_scales() @@ -366,7 +366,7 @@ def test_world_scales_roundtrip(device, view_factory): 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_space_writer("world") as w: + with view.xform_world_space_writer() as w: w.set_scales(new_scales) ret_scales = view.get_world_scales() @@ -425,7 +425,7 @@ def test_set_local_then_get_world_with_rotated_parent(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_space_writer("local") as w: + with view.xform_local_space_writer() as w: w.set_poses(positions=new_local, orientations=identity_quat) world_pos, _ = view.get_world_poses() @@ -447,7 +447,7 @@ def test_set_world_then_get_local_with_rotated_parent(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_space_writer("world") as w: + with view.xform_world_space_writer() as w: w.set_poses(positions=new_world) local_pos, _ = view.get_local_poses() @@ -543,7 +543,7 @@ def test_multi_view_writer_isolation(device): 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_space_writer("local") as w: + with view_a.xform_local_space_writer() as w: w.set_poses(positions=new_local_a, orientations=identity_quat) # View B remains undisturbed. @@ -560,7 +560,7 @@ def test_multi_view_writer_isolation(device): # 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_space_writer("local") as w: + 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 @@ -571,11 +571,11 @@ def test_multi_view_writer_isolation(device): ) # Single-active-writer is per-view: opening a writer on A leaves B free. - with view_a.xform_space_writer("world"): + 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_space_writer("world"): + 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 @@ -602,7 +602,7 @@ 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) - with view.xform_space_writer("world") as w: + with view.xform_world_space_writer() as w: w.set_poses(positions=new_pos) ret_pos, _ = view.get_world_poses() @@ -633,7 +633,7 @@ 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) - with view.xform_space_writer("world") as w: + with view.xform_world_space_writer() as w: w.set_poses(positions=new_pos) # USD must not have moved at all -- equality, not approximate. @@ -661,7 +661,7 @@ def test_fabric_cuda1_scales_roundtrip(device, view_factory): 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_space_writer("world") as w: + with view.xform_world_space_writer() as w: w.set_scales(new_scales) ret_scales = view.get_world_scales() @@ -705,14 +705,14 @@ def test_sequential_world_then_local_scopes_partial_indices(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_space_writer("world") as w: + 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_space_writer("local") as w: + 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 @@ -759,13 +759,13 @@ def test_sequential_local_then_world_scopes_partial_indices(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_space_writer("local") as w: + 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_space_writer("world") as w: + with view.xform_world_space_writer() as w: w.set_poses(positions=new_world_pos, indices=idx1) # Index 0's world (derived from local): @@ -815,7 +815,7 @@ def test_world_writer_writes_world_and_derives_local(device, view_factory): 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_space_writer("world") as w: + with view.xform_world_space_writer() as w: w.set_poses(positions=new_pos) w.set_scales(new_scales) @@ -847,7 +847,7 @@ def test_local_writer_writes_local_and_derives_world(device, view_factory): 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_space_writer("local") as w: + 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) @@ -882,7 +882,7 @@ def counted(): 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_space_writer("world") as w: + 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 @@ -897,13 +897,13 @@ def test_writer_single_active_invariant(device, view_factory): view = bundle.view view.get_world_poses() - with view.xform_space_writer("world"): - with pytest.raises(RuntimeError, match="already has an active xform_space_writer"): - view.xform_space_writer("world").__enter__() - with pytest.raises(RuntimeError, match="already has an active xform_space_writer"): - view.xform_space_writer("local").__enter__() + 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_space_writer("local"): + with view.xform_local_space_writer(): pass @@ -918,26 +918,17 @@ def test_writer_restores_hierarchy_change_tracking(device, view_factory): # 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_space_writer("world"): + 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_space_writer("world"): + 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", ["cpu", "cuda:0"]) -def test_writer_invalid_space_raises(device, view_factory): - """``xform_space_writer`` with an invalid space raises ``ValueError``.""" - bundle = view_factory(num_envs=1, device=device) - view = bundle.view - with pytest.raises(ValueError, match="Invalid space"): - view.xform_space_writer("foo") - - @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.""" @@ -955,7 +946,7 @@ def counted(): monkeypatch.setattr(view, "_recompute_local_from_world_all", counted) - with view.xform_space_writer("world"): + with view.xform_world_space_writer(): pass assert calls == 0 @@ -968,14 +959,14 @@ def test_view_getter_inside_scope_raises(device, view_factory): view = bundle.view view.get_world_poses() - with view.xform_space_writer("world"): - with pytest.raises(RuntimeError, match="xform_space_writer"): + 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="xform_space_writer"): + with pytest.raises(RuntimeError, match="while a writer scope is active"): view.get_local_poses() - with pytest.raises(RuntimeError, match="xform_space_writer"): + with pytest.raises(RuntimeError, match="while a writer scope is active"): view.get_world_scales() - with pytest.raises(RuntimeError, match="xform_space_writer"): + 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 eecb4c243a951984cac4b89893e4c34de3d458b3 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Tue, 23 Jun 2026 14:59:43 +0000 Subject: [PATCH 41/54] refactor: collapse Fabric selections from three to two Previously FabricFrameView held three persistent PrimSelection handles with asymmetric RO/RW flags (trans_sel_ro = both-RO for reads, world_sel_rw = worldMatrix-RW, local_sel_rw = localMatrix-RW). That layout was the load-bearing protection against Kit's per-tick updateWorldXforms() before the writer scope started pausing IFabricHierarchy tracking. With tracking-pause in place, the asymmetric flags are no longer load-bearing. Reduce to two persistent selections: _sel_ro : worldMatrix=RO, localMatrix=RO steady state _sel_rw : worldMatrix=RW, localMatrix=RW inside writer scope Each selection owns its own bundle of indexed-fabric arrays. The writer scope flips a single _is_rw flag on enter/exit; both bundles stay alive for the view's lifetime, so no selection is rebuilt on the flip. PrepareForReuse() topology polls run independently per bundle. Net: -10 lines of code, -1 selection, -1 set of fabric_indices. The two whitebox tests that reached into the old triple are updated to poll the new pair. All 79 fabric-suite tests pass and the test_output_equal_to_usdcamera camera+RTX regression that motivated this PR still passes. --- .../changelog.d/xform-space-writer.rst | 9 +- .../sim/views/fabric_frame_view.py | 244 +++++++++--------- .../test/sim/test_views_xform_prim_fabric.py | 20 +- 3 files changed, 133 insertions(+), 140 deletions(-) diff --git a/source/isaaclab_physx/changelog.d/xform-space-writer.rst b/source/isaaclab_physx/changelog.d/xform-space-writer.rst index 2ff88d32d5b4..2630339b0153 100644 --- a/source/isaaclab_physx/changelog.d/xform-space-writer.rst +++ b/source/isaaclab_physx/changelog.d/xform-space-writer.rst @@ -23,6 +23,9 @@ Changed removed -- the eager dual-write inside the scope makes all of that unnecessary. - The three-selection RO/RW layout (``_trans_sel_ro``, - ``_world_sel_rw``, ``_local_sel_rw``) is kept as a defensive layer and - for clarity of authoring intent. + The previous three-selection RO/RW layout has been collapsed to two + persistent selections, ``_sel_ro`` (worldMatrix=RO, localMatrix=RO) and + ``_sel_rw`` (worldMatrix=RW, localMatrix=RW). The writer scope flips an + ``_is_rw`` flag on enter/exit; both selection bundles are kept alive for + the view's lifetime, so no selection is rebuilt on flip. The hierarchy + tracking-pause above is the load-bearing renderer-clobber protection. 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 8d4d64720349..5e41ccaf66b9 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 @@ -86,21 +86,27 @@ class FabricFrameView(BaseFrameView): tracking state (so we do not re-enable listeners the caller had previously paused). The renderer's own independent worldMatrix listener is unaffected and still observes our writes. - * **Three selections with asymmetric RO/RW access.** Despite the - pause/restore above, we keep three selections as a defensive layer: + * **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 - _trans_sel_ro : worldMatrix=RO, localMatrix=RO (reads) - _world_sel_rw : worldMatrix=RW, localMatrix=RO (world writer) - _local_sel_rw : worldMatrix=RO, localMatrix=RW (local writer) - - A combined ``ReadWrite(world, local)`` selection is unsafe even with - tracking pause -- if a refactor accidentally re-enables tracking, - Fabric would see both attributes as user-authored and fall back to the - hierarchy's canonical direction (local -> world), clobbering our world - write. The separate RO/RW layout makes the intended authoring - direction explicit. + _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 renderer-clobber protection comes from the hierarchy + tracking-pause above; the RO/RW split tells Fabric which steady-state + attribute is user-authored (none, in the RO case) once the scope + closes and the matrices are mutually consistent. * **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. @@ -150,24 +156,29 @@ def __init__( self._stage = None self._fabric_hierarchy = None - # Three persistent Fabric selections with asymmetric access flags. - self._trans_sel_ro = None - self._world_sel_rw = None - self._local_sel_rw = 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. + self._sel_ro = None + self._sel_rw = None + self._is_rw: bool = False - # Index arrays (view-side indices and per-selection view->fabric mappings). + # View-side indices array (shared across both bundles). self._view_indices: wp.array | None = None - self._trans_ro_fabric_indices: wp.array | None = None - self._world_rw_fabric_indices: wp.array | None = None - self._local_rw_fabric_indices: wp.array | None = None - self._parent_fabric_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_ro = 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 @@ -243,7 +254,7 @@ def _get_world_poses_impl(self, indices: wp.array | None = None) -> tuple[ProxyA kernel=fabric_utils.decompose_indexed_fabric_transforms, dim=count, inputs=[ - self._get_world_ro_array(), + self._get_world_ifa(), positions_wp, orientations_wp, self._fabric_empty_2d_array_sentinel, @@ -279,7 +290,7 @@ def _get_local_poses_impl(self, indices: wp.array | None = None) -> tuple[ProxyA kernel=fabric_utils.decompose_indexed_fabric_transforms, dim=count, inputs=[ - self._get_local_ro_array(), + self._get_local_ifa(), translations_wp, orientations_wp, self._fabric_empty_2d_array_sentinel, @@ -300,7 +311,7 @@ def _get_world_scales_impl(self, indices=None) -> ProxyArray: if not self._fabric_initialized: self._initialize_fabric() - return self._decompose_scales(self._get_world_ro_array(), indices) + return self._decompose_scales(self._get_world_ifa(), indices) def _get_local_scales_impl(self, indices=None) -> ProxyArray: if not self._use_fabric: @@ -309,7 +320,7 @@ def _get_local_scales_impl(self, indices=None) -> ProxyArray: if not self._fabric_initialized: self._initialize_fabric() - return self._decompose_scales(self._get_local_ro_array(), indices) + 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.""" @@ -368,15 +379,13 @@ def _recompute_local_from_world_all(self) -> None: Storage convention: see :func:`isaaclab.utils.warp.fabric.update_indexed_local_matrix_from_world`. """ - if self._trans_sel_ro.PrepareForReuse() or self._parent_world_ifa_ro is None: - self._rebuild_trans_ro_arrays() wp.launch( kernel=fabric_utils.update_indexed_local_matrix_from_world, dim=self.count, inputs=[ - self._world_ifa_ro, - self._parent_world_ifa_ro, - self._get_local_rw_array(), + self._get_world_ifa(), + self._get_parent_world_ifa(), + self._get_local_ifa(), self._view_indices, ], device=self._device, @@ -390,15 +399,13 @@ def _recompute_world_from_local_all(self) -> None: Storage convention: see :func:`isaaclab.utils.warp.fabric.update_indexed_world_matrix_from_local`. """ - if self._trans_sel_ro.PrepareForReuse() or self._parent_world_ifa_ro is None: - self._rebuild_trans_ro_arrays() wp.launch( kernel=fabric_utils.update_indexed_world_matrix_from_local, dim=self.count, inputs=[ - self._local_ifa_ro, - self._parent_world_ifa_ro, - self._get_world_rw_array(), + self._get_local_ifa(), + self._get_parent_world_ifa(), + self._get_world_ifa(), self._view_indices, ], device=self._device, @@ -408,52 +415,48 @@ def _recompute_world_from_local_all(self) -> None: # Internal -- selection accessors with on-demand index rebuild # ------------------------------------------------------------------ - def _get_world_ro_array(self): - if self._trans_sel_ro.PrepareForReuse(): - self._rebuild_trans_ro_arrays() - return self._world_ifa_ro - - def _get_local_ro_array(self): - if self._trans_sel_ro.PrepareForReuse(): - self._rebuild_trans_ro_arrays() - return self._local_ifa_ro - - def _get_world_rw_array(self): - if self._world_sel_rw.PrepareForReuse(): - self._world_rw_fabric_indices = self._compute_fabric_indices(self._world_sel_rw) - self._world_ifa_rw = self._build_indexed_array( - self._world_sel_rw, self._WORLD_MATRIX_NAME, self._world_rw_fabric_indices - ) - return self._world_ifa_rw - - def _get_local_rw_array(self): - if self._local_sel_rw.PrepareForReuse(): - self._local_rw_fabric_indices = self._compute_fabric_indices(self._local_sel_rw) - self._local_ifa_rw = self._build_indexed_array( - self._local_sel_rw, self._LOCAL_MATRIX_NAME, self._local_rw_fabric_indices - ) - return self._local_ifa_rw + 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_parent_world_ro_array(self): - # Built and refreshed alongside the trans_ro selection (parents share that selection). - if self._parent_world_ifa_ro is None or self._trans_sel_ro.PrepareForReuse(): - self._rebuild_trans_ro_arrays() - return self._parent_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 _rebuild_trans_ro_arrays(self) -> None: - """Rebuild the trans_ro indices and the three indexed arrays that depend on them. + 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 - ``_world_ifa_ro``, ``_local_ifa_ro`` and ``_parent_world_ifa_ro`` are all - keyed off the ``trans_sel_ro`` path ordering, so they are refreshed together. - """ - self._trans_ro_fabric_indices = self._compute_fabric_indices(self._trans_sel_ro) - self._world_ifa_ro = self._build_indexed_array( - self._trans_sel_ro, self._WORLD_MATRIX_NAME, self._trans_ro_fabric_indices + 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, ) - self._local_ifa_ro = self._build_indexed_array( - self._trans_sel_ro, self._LOCAL_MATRIX_NAME, self._trans_ro_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, ) - self._parent_world_ifa_ro = self._build_parent_indexed_array(self._trans_sel_ro) # ------------------------------------------------------------------ # Internal -- index computation @@ -498,11 +501,6 @@ def _build_indexed_array(self, selection, attribute_name: str, fabric_indices: w fa = wp.fabricarray(selection, attribute_name) return wp.indexedfabricarray(fa=fa, indices=fabric_indices) - def _build_parent_indexed_array(self, selection) -> wp.indexedfabricarray: - self._parent_fabric_indices = self._compute_parent_fabric_indices(selection) - fa = wp.fabricarray(selection, self._WORLD_MATRIX_NAME) - return wp.indexedfabricarray(fa=fa, indices=self._parent_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): @@ -550,7 +548,9 @@ def _initialize_fabric(self) -> None: rt_xformable.SetLocalXformFromUsd() rt_xformable.SetWorldXformFromUsd() - # Three persistent selections with asymmetric access flags. + # 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 @@ -558,30 +558,12 @@ def _initialize_fabric(self) -> None: 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._trans_sel_ro = self._stage.SelectPrims(require_attrs=[wm_ro, lm_ro], device=self._device, want_paths=True) - self._world_sel_rw = self._stage.SelectPrims(require_attrs=[wm_rw, lm_ro], device=self._device, want_paths=True) - self._local_sel_rw = self._stage.SelectPrims(require_attrs=[wm_ro, lm_rw], device=self._device, want_paths=True) + 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) - # Build the view-side indices array and per-selection view->fabric mappings. self._view_indices = wp.array(list(range(self.count)), dtype=wp.uint32, device=self._device) - self._trans_ro_fabric_indices = self._compute_fabric_indices(self._trans_sel_ro) - self._world_rw_fabric_indices = self._compute_fabric_indices(self._world_sel_rw) - self._local_rw_fabric_indices = self._compute_fabric_indices(self._local_sel_rw) - - # Indexed fabric arrays per (selection x attribute). - self._world_ifa_ro = self._build_indexed_array( - self._trans_sel_ro, self._WORLD_MATRIX_NAME, self._trans_ro_fabric_indices - ) - self._local_ifa_ro = self._build_indexed_array( - self._trans_sel_ro, self._LOCAL_MATRIX_NAME, self._trans_ro_fabric_indices - ) - self._world_ifa_rw = self._build_indexed_array( - self._world_sel_rw, self._WORLD_MATRIX_NAME, self._world_rw_fabric_indices - ) - self._local_ifa_rw = self._build_indexed_array( - self._local_sel_rw, self._LOCAL_MATRIX_NAME, self._local_rw_fabric_indices - ) - self._parent_world_ifa_ro = self._build_parent_indexed_array(self._trans_sel_ro) + 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) @@ -599,8 +581,14 @@ def _initialize_fabric(self) -> None: self._fabric_initialized = True - # Seed Fabric matrices from USD authoritatively. - self._sync_fabric_from_usd_initial() + # 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. @@ -616,7 +604,7 @@ def _sync_fabric_from_usd_initial(self) -> None: kernel=fabric_utils.compose_indexed_fabric_transforms, dim=self.count, inputs=[ - self._local_ifa_rw, + 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), @@ -674,8 +662,8 @@ def _sync_fabric_from_usd_initial(self) -> None: 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._world_sel_rw, self._WORLD_MATRIX_NAME), - indices=self._compute_fabric_indices_for(self._world_sel_rw, unique_parent_paths), + 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, @@ -720,10 +708,14 @@ def _compute_fabric_indices_for(self, selection, paths: list[str]) -> wp.array: class _FabricWriterMixin: """Common ``__enter__`` / ``__exit__`` for the Fabric world / local writers. - Pauses ``track_local_xform_changes`` / ``track_world_xform_changes`` on - the Fabric hierarchy while the scope is active so Kit does not redundantly - recompute the matrices we just wrote, then restores the prior state on - exit. + 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: derives the opposite-space matrix, synchronizes, flips ``_is_rw`` + back to ``False`` (RO bundle for steady-state reads), restores + hierarchy-tracking state. """ def _enter_impl(self) -> None: @@ -738,6 +730,7 @@ def _enter_impl(self) -> None: 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] @@ -746,6 +739,9 @@ def _exit_impl(self, exc_type, exc_val, exc_tb) -> None: self._derive_opposite() wp.synchronize() 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) @@ -759,9 +755,9 @@ def _derive_opposite(self) -> None: class _FabricWorldSpaceWriter(_FabricWriterMixin, FrameViewWorldSpaceWriter): """World-space writer for :class:`FabricFrameView`. - Writes flow through ``_world_sel_rw``; on exit ``localMatrix`` is derived - from the just-written ``worldMatrix`` via - :func:`update_indexed_local_matrix_from_world`. + 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: @@ -776,7 +772,7 @@ def set_poses(self, positions=None, orientations=None, indices=None) -> None: kernel=fabric_utils.compose_indexed_fabric_transforms, dim=indices_wp.shape[0], inputs=[ - view._get_world_rw_array(), + view._get_world_ifa(), positions_wp, orientations_wp, view._fabric_empty_2d_array_sentinel, @@ -797,7 +793,7 @@ def set_scales(self, scales, indices=None) -> None: kernel=fabric_utils.compose_indexed_fabric_transforms, dim=indices_wp.shape[0], inputs=[ - view._get_world_rw_array(), + view._get_world_ifa(), view._fabric_empty_2d_array_sentinel, view._fabric_empty_2d_array_sentinel, scales_wp, @@ -820,9 +816,9 @@ def get_scales(self, indices=None) -> ProxyArray: class _FabricLocalSpaceWriter(_FabricWriterMixin, FrameViewLocalSpaceWriter): """Local-space writer for :class:`FabricFrameView`. - Writes flow through ``_local_sel_rw``; on exit ``worldMatrix`` is derived - from the just-written ``localMatrix`` via - :func:`update_indexed_world_matrix_from_local`. + 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: @@ -837,7 +833,7 @@ def set_poses(self, positions=None, orientations=None, indices=None) -> None: kernel=fabric_utils.compose_indexed_fabric_transforms, dim=indices_wp.shape[0], inputs=[ - view._get_local_rw_array(), + view._get_local_ifa(), translations_wp, orientations_wp, view._fabric_empty_2d_array_sentinel, @@ -858,7 +854,7 @@ def set_scales(self, scales, indices=None) -> None: kernel=fabric_utils.compose_indexed_fabric_transforms, dim=indices_wp.shape[0], inputs=[ - view._get_local_rw_array(), + view._get_local_ifa(), view._fabric_empty_2d_array_sentinel, view._fabric_empty_2d_array_sentinel, scales_wp, 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 4b030abc20c6..3786d0ae271d 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 @@ -203,17 +203,10 @@ 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: recompute per-selection fabric indices and rebuild - # every indexed array, mirroring the lazy paths in the ``_get_*_array`` accessors. - view._rebuild_trans_ro_arrays() - view._world_rw_fabric_indices = view._compute_fabric_indices(view._world_sel_rw) - view._world_ifa_rw = view._build_indexed_array( - view._world_sel_rw, view._WORLD_MATRIX_NAME, view._world_rw_fabric_indices - ) - view._local_rw_fabric_indices = view._compute_fabric_indices(view._local_sel_rw) - view._local_ifa_rw = view._build_indexed_array( - view._local_sel_rw, view._LOCAL_MATRIX_NAME, view._local_rw_fabric_indices - ) + # 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() # Trigger another write through the rebuilt arrays. new = wp.zeros((2, 3), dtype=wp.float32, device=device) @@ -237,8 +230,9 @@ def test_prepare_for_reuse_detects_topology_change(device, view_factory): view = bundle.view view.get_world_poses() # trigger Fabric init - assert view._trans_sel_ro is not None, "trans_sel_ro selection not initialized" - for selection in (view._trans_sel_ro, view._world_sel_rw, view._local_sel_rw): + 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" From 7ce0f0f631d8346803e0c4534732082bd769fc79 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Tue, 23 Jun 2026 15:04:53 +0000 Subject: [PATCH 42/54] docs: describe Fabric selection layout as end state in changelog --- .../changelog.d/xform-space-writer.rst | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/source/isaaclab_physx/changelog.d/xform-space-writer.rst b/source/isaaclab_physx/changelog.d/xform-space-writer.rst index 2630339b0153..cc8bb4800bcd 100644 --- a/source/isaaclab_physx/changelog.d/xform-space-writer.rst +++ b/source/isaaclab_physx/changelog.d/xform-space-writer.rst @@ -23,9 +23,12 @@ Changed removed -- the eager dual-write inside the scope makes all of that unnecessary. - The previous three-selection RO/RW layout has been collapsed to two - persistent selections, ``_sel_ro`` (worldMatrix=RO, localMatrix=RO) and - ``_sel_rw`` (worldMatrix=RW, localMatrix=RW). The writer scope flips an - ``_is_rw`` flag on enter/exit; both selection bundles are kept alive for - the view's lifetime, so no selection is rebuilt on flip. The hierarchy - tracking-pause above is the load-bearing renderer-clobber protection. + Two persistent Fabric selections are built once during + ``_initialize_fabric`` and kept for the view's lifetime: ``_sel_ro`` + (``worldMatrix=RO, localMatrix=RO``, steady state) and ``_sel_rw`` + (``worldMatrix=RW, localMatrix=RW``, used inside a writer scope). The + writer flips a single ``_is_rw`` flag on enter/exit; neither selection + is rebuilt on the flip. Renderer correctness comes from the hierarchy + tracking pause above; the RO/RW split tells Kit's per-tick + ``updateWorldXforms()`` which matrices the user is currently + authoring. From c2f7a7ffc9ccb4a9e611f05be35827a3d0a5a9d6 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Tue, 23 Jun 2026 15:15:56 +0000 Subject: [PATCH 43/54] fix: best-effort opposite-space sync on writer-scope exception When a FabricFrameView writer scope unwinds via an exception (a Python error in user code, a notebook cell interrupt, etc.), the opposite-space derive + wp.synchronize() now runs anyway so worldMatrix and localMatrix remain mutually consistent prim-by-prim on whatever partial-write state Fabric currently holds. The partial write itself is not rolled back -- callers needing transactional semantics should snapshot the matrices themselves before entering the scope. The recovery launch is itself wrapped in try/except: if it fails (e.g. the original exception came from a poisoned CUDA stream), the recovery error is logged and the original exception propagates -- masking it would hide the actual root cause. Hierarchy-tracking restore and the _is_rw flip happen in finally as before. Add a regression test that raises inside a writer scope and verifies: - tracking state restored - _is_rw back to False - world matrices reflect the partial write - local matrices derived from that partial-write world - the view is still usable for further writer scopes --- .../changelog.d/xform-space-writer.rst | 7 +++ .../sim/views/fabric_frame_view.py | 49 +++++++++++++-- .../test/sim/test_views_xform_prim_fabric.py | 60 ++++++++++++++++++- 3 files changed, 109 insertions(+), 7 deletions(-) diff --git a/source/isaaclab_physx/changelog.d/xform-space-writer.rst b/source/isaaclab_physx/changelog.d/xform-space-writer.rst index cc8bb4800bcd..25878c737b72 100644 --- a/source/isaaclab_physx/changelog.d/xform-space-writer.rst +++ b/source/isaaclab_physx/changelog.d/xform-space-writer.rst @@ -16,6 +16,13 @@ Changed ``updateWorldXforms()`` does not redundantly recompute matrices the user just wrote. The renderer's independent ``omni:fabric:worldMatrix`` listener is unaffected 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. The lazy-dirty-flag mechanism (the ``_DirtyFlag`` enum, ``_dirty`` field, ``_sync_*_if_dirty`` helpers, and the one-time 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 5e41ccaf66b9..5258c05e006c 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 @@ -77,6 +77,14 @@ class FabricFrameView(BaseFrameView): attribute and a single ``wp.synchronize()`` runs. After the scope exits, both Fabric matrices are self-consistent; getters read directly from Fabric storage without any further synchronization. + + 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. * **Hierarchy listeners are paused while a writer scope is active.** The writer's ``__enter__`` calls :meth:`IFabricHierarchy.track_local_xform_changes(False)` / @@ -713,9 +721,23 @@ class _FabricWriterMixin: ``_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: derives the opposite-space matrix, synchronizes, flips ``_is_rw`` - back to ``False`` (RO bundle for steady-state reads), restores - hierarchy-tracking state. + 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: @@ -735,9 +757,24 @@ def _enter_impl(self) -> None: def _exit_impl(self, exc_type, exc_val, exc_tb) -> None: view: FabricFrameView = self._view # type: ignore[assignment] try: - if self._wrote_anything and exc_type is None: - self._derive_opposite() - wp.synchronize() + 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. 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 3786d0ae271d..10eb883940f3 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 # 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 @@ -221,6 +221,64 @@ def test_fabric_rebuild_after_topology_change(device, view_factory): 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 From 54e26c1a6fc09896456a89d8292ba464c34e31d8 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Tue, 23 Jun 2026 15:59:05 +0000 Subject: [PATCH 44/54] docs: clarify writer scope is synchronous; no step/render inside The previous wording suggested IFabricHierarchy.update_world_xforms() could fire interleaved with our writes (a "between our write and the renderer's read" race). That cannot happen: the scope is synchronous Python code and no simulation step or render tick runs while it is open. The risk the tracking pause and the RO steady-state selection actually defend against is the next render/sim tick after the scope exits, not the scope itself. Update FabricFrameView's class docstring to reflect this, and add an explicit contract to FrameViewSpaceWriterBase: callers must not advance the simulation or render from inside a scope, because the matrices may be mid-write until exit and rendering would read torn data. --- .../isaaclab/sim/views/xform_space_writer.py | 13 +++++++++++ .../sim/views/fabric_frame_view.py | 23 ++++++++++++------- 2 files changed, 28 insertions(+), 8 deletions(-) diff --git a/source/isaaclab/isaaclab/sim/views/xform_space_writer.py b/source/isaaclab/isaaclab/sim/views/xform_space_writer.py index 17c437294a0b..2d657b031995 100644 --- a/source/isaaclab/isaaclab/sim/views/xform_space_writer.py +++ b/source/isaaclab/isaaclab/sim/views/xform_space_writer.py @@ -22,6 +22,14 @@ ``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 @@ -46,6 +54,11 @@ class FrameViewSpaceWriterBase(abc.ABC): 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): 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 5258c05e006c..88197249f581 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 @@ -86,11 +86,17 @@ class FabricFrameView(BaseFrameView): transactional all-or-nothing semantics, snapshot the matrices yourself before entering the scope. * **Hierarchy listeners are paused while a writer scope is active.** - The writer's ``__enter__`` calls + The writer 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 risk that this pause defends against is the *next* tick that + runs after the scope exits. On enter, the writer calls :meth:`IFabricHierarchy.track_local_xform_changes(False)` / - :meth:`track_world_xform_changes(False)` (saving the prior state) so - that Kit's per-tick ``updateWorldXforms()`` does not redundantly - recompute matrices we just wrote. ``__exit__`` restores the prior + :meth:`track_world_xform_changes(False)` (saving the prior state) + so that Kit's hierarchy tracker does not observe our writes as + user edits and queue propagation work for the next + ``update_world_xforms()`` tick. ``__exit__`` restores the prior tracking state (so we do not re-enable listeners the caller had previously paused). The renderer's own independent worldMatrix listener is unaffected and still observes our writes. @@ -111,10 +117,11 @@ class FabricFrameView(BaseFrameView): flip -- both bundles are always kept consistent via independent ``PrepareForReuse()`` polls in the accessors. - The renderer-clobber protection comes from the hierarchy - tracking-pause above; the RO/RW split tells Fabric which steady-state - attribute is user-authored (none, in the RO case) once the scope - closes and the matrices are mutually consistent. + 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. From d9617443e35d6f725d0ea900264ca823280fbd24 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Tue, 23 Jun 2026 17:07:24 +0000 Subject: [PATCH 45/54] docs: clarify what the writer-scope tracking pause actually prevents Previous wording suggested the pause was about keeping the renderer from seeing half-written data. That is wrong: the scope is synchronous Python code, so the renderer cannot run mid-scope and cannot see torn state. The "torn data" concern only motivates the separate rule that callers must not advance the simulation from inside the scope. What the pause actually prevents is Fabric updating the dependent xforms itself, in two ways the scope wants to keep exclusive: 1. Per-write, inside the scope -- the tracker would fire synchronously on each set_* and propagate to the opposite-space matrix (and to descendants, where applicable). 2. On the next render tick -- update_world_xforms() would replay queued tracker work and potentially derive the wrong direction (e.g. recompute world from local even though we just wrote world). We do the dependent-xform update ourselves in one batched kernel at scope exit, so the pause stops Fabric from doing the same work redundantly during the scope and incorrectly after it. --- .../sim/views/fabric_frame_view.py | 36 +++++++++++++------ 1 file changed, 25 insertions(+), 11 deletions(-) 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 88197249f581..f9b6af037a0b 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 @@ -86,20 +86,34 @@ class FabricFrameView(BaseFrameView): transactional all-or-nothing semantics, snapshot the matrices yourself before entering the scope. * **Hierarchy listeners are paused while a writer scope is active.** - The writer scope is synchronous Python code, so no simulation step + On enter, the writer calls + :meth:`IFabricHierarchy.track_local_xform_changes(False)` / + :meth:`track_world_xform_changes(False)` (saving the prior state). + With tracking on, Fabric would update the dependent xforms itself + in two ways the scope wants to keep exclusive: + + 1. *Per-write, inside the scope* -- each ``set_*`` call would fire + the tracker synchronously and propagate to the opposite-space + matrix (and to descendants, where applicable). + 2. *On the next render tick* -- ``update_world_xforms()`` would + replay queued tracker work and potentially derive the wrong + direction (e.g. recompute world from local even though we + just wrote world). + + We do the dependent-xform update ourselves in one batched kernel + at scope exit (see the "eager dual-write" bullet above), so the + pause prevents Fabric from doing the same work redundantly during + the scope and incorrectly after it. ``__exit__`` restores the + prior tracking state (so we do not re-enable listeners the caller + had previously paused). The renderer's own independent + worldMatrix listener is unaffected and still observes our writes. + + 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 risk that this pause defends against is the *next* tick that - runs after the scope exits. On enter, the writer calls - :meth:`IFabricHierarchy.track_local_xform_changes(False)` / - :meth:`track_world_xform_changes(False)` (saving the prior state) - so that Kit's hierarchy tracker does not observe our writes as - user edits and queue propagation work for the next - ``update_world_xforms()`` tick. ``__exit__`` restores the prior - tracking state (so we do not re-enable listeners the caller had - previously paused). The renderer's own independent worldMatrix - listener is unaffected and still observes our writes. + 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: From c3fc5068ca2eda0c7f62c1a44834102bf4fb6db6 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Tue, 23 Jun 2026 17:20:34 +0000 Subject: [PATCH 46/54] docs: correct mental model of the tracking pause (pull-based, not push) Fabric change tracking is pull-based: a per-attribute listener records writes into a private changelog, and Kit 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. Previous docstring claimed the pause prevented "per-write synchronous propagation through the tracker". There is no such synchronous path: nothing fires until the next update_world_xforms() tick. Rewrite the bullet to describe the actual mechanism, and explain why an empty changelog at scope exit is the property we need (the next tick has no entries to process for our prims and therefore can't pick a canonical direction and derive the other from it). Source: kit/runtime/source/plugins/usdrt.hierarchy/FabricHierarchy.cpp - trackLocalXformChanges_abi -> pauseChangeTracking/resumeChangeTracking - updateWorldXforms_abi reads getChanges()/popChanges() on each listener --- .../sim/views/fabric_frame_view.py | 41 +++++++++++-------- 1 file changed, 23 insertions(+), 18 deletions(-) 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 f9b6af037a0b..f1d483b2c342 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 @@ -89,24 +89,29 @@ class FabricFrameView(BaseFrameView): On enter, the writer calls :meth:`IFabricHierarchy.track_local_xform_changes(False)` / :meth:`track_world_xform_changes(False)` (saving the prior state). - With tracking on, Fabric would update the dependent xforms itself - in two ways the scope wants to keep exclusive: - - 1. *Per-write, inside the scope* -- each ``set_*`` call would fire - the tracker synchronously and propagate to the opposite-space - matrix (and to descendants, where applicable). - 2. *On the next render tick* -- ``update_world_xforms()`` would - replay queued tracker work and potentially derive the wrong - direction (e.g. recompute world from local even though we - just wrote world). - - We do the dependent-xform update ourselves in one batched kernel - at scope exit (see the "eager dual-write" bullet above), so the - pause prevents Fabric from doing the same work redundantly during - the scope and incorrectly after it. ``__exit__`` restores the - prior tracking state (so we do not re-enable listeners the caller - had previously paused). The renderer's own independent - worldMatrix listener is unaffected and still observes our writes. + Fabric's change tracking is pull-based: a per-attribute listener + records writes into a private changelog, and Kit 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 + renderer's own independent worldMatrix listener is unaffected and + still observes our writes. Note: the scope is synchronous Python code, so no simulation step and no render tick can run while it is open -- callers must not From 45fd8ab5e96854153a2a8991e4ab52338daf1f89 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Tue, 23 Jun 2026 17:23:44 +0000 Subject: [PATCH 47/54] docs: use precise Omniverse terminology (Fabric Hierarchy vs Fabric) Fabric is just the flat attribute data store. The plugin that owns the changelog listeners, the track_*_xform_changes toggles, and the update_world_xforms() step is Fabric Hierarchy (usdrt::hierarchy:: IFabricHierarchy), not Fabric itself. The renderer reads Fabric attributes through the Fabric Scene Delegate (FSD), not via some generic "renderer worldMatrix listener". Rework the writer-scope's listener bullet to use these names precisely: - the tracking pause and the changelog belong to Fabric Hierarchy - update_world_xforms() is a method on Fabric Hierarchy, not on "Kit" - FSD reads omni:fabric:worldMatrix from Fabric storage on the render path The IsaacLab project-level shorthand "Fabric backend" / FabricFrameView / _use_fabric is unchanged -- those are API labels, not technical claims. --- .../sim/views/fabric_frame_view.py | 24 ++++++++++++------- 1 file changed, 15 insertions(+), 9 deletions(-) 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 f1d483b2c342..df1f13eaecac 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 @@ -85,17 +85,22 @@ class FabricFrameView(BaseFrameView): The partial write itself is not rolled back -- if you need transactional all-or-nothing semantics, snapshot the matrices yourself before entering the scope. - * **Hierarchy listeners are paused while a writer scope is active.** - On enter, the writer calls + * **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's change tracking is pull-based: a per-attribute listener - records writes into a private changelog, and Kit drains and - processes that changelog on the next call to + 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. + 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 @@ -110,8 +115,9 @@ class FabricFrameView(BaseFrameView): ``__exit__`` restores the prior tracking state (so we do not re-enable listeners the caller had previously paused). The - renderer's own independent worldMatrix listener is unaffected and - still observes our writes. + 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 From 714bc7b576c0d377f1272f7a30082f92ea29ce58 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Tue, 23 Jun 2026 17:37:36 +0000 Subject: [PATCH 48/54] docs: drop intra-branch history from changelog fragments Changelog fragments are read against develop, not against an in-branch interim state. Remove every entry that describes a feature added then removed within this PR -- net effect vs develop is zero, and mentioning the round-trip only confuses reviewers. Specifically: - isaaclab/xform-space-writer: drop the "Removed: set_world_scales / set_local_scales" section. Those methods were added in fabric-local-poses and removed here; they never existed on develop. - isaaclab/fabric-local-poses: drop set_world_scales / set_local_scales from the Added list for the same reason; keep only the surviving getters and route scale writes via the writer scope. - isaaclab_physx/fabric-local-poses: drop the "lazy dirty tracking" and "interleave detection" entries; both mechanisms were added and then removed by the writer-scope migration. - isaaclab_physx/xform-space-writer: drop the paragraph that lists the removed _DirtyFlag enum / _sync_*_if_dirty helpers. Also switch the remaining prose to the correct Omniverse terminology -- Fabric Hierarchy owns update_world_xforms() and the change tracking; FSD feeds the renderer. - isaaclab_newton + isaaclab_ovphysx/fabric-local-poses: drop set_world_scales / set_local_scales from the Added/Deprecated lists. Also include the prior wording fix "Will be used by" -> "Used by" in the utils docs (the kernels are now in use). While I'm in this file, drop the deprecation-replacement clauses that told users to use set_world_scales / set_local_scales -- those methods don't exist; the writer scope is the right pointer. Plus: add a comment in scripts/benchmarks/benchmark_xform_prim_view.py explaining why each wp.clone() now goes through .warp (ProxyArray introduced by PR #5304 changed the FrameView getter return type). --- docs/source/api/lab/isaaclab.utils.rst | 2 +- .../benchmarks/benchmark_xform_prim_view.py | 5 +++ .../changelog.d/fabric-local-poses.rst | 22 +++++++------ .../changelog.d/xform-space-writer.rst | 11 ------- .../changelog.d/fabric-local-poses.rst | 20 ++++++------ .../changelog.d/fabric-local-poses.rst | 21 +++++++------ .../changelog.d/fabric-local-poses.rst | 28 ++++++----------- .../changelog.d/xform-space-writer.rst | 31 ++++++++----------- 8 files changed, 61 insertions(+), 79 deletions(-) diff --git a/docs/source/api/lab/isaaclab.utils.rst b/docs/source/api/lab/isaaclab.utils.rst index f236ebcb6a15..59c78b3aef48 100644 --- a/docs/source/api/lab/isaaclab.utils.rst +++ b/docs/source/api/lab/isaaclab.utils.rst @@ -194,7 +194,7 @@ 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`. Will be used by +: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. diff --git a/scripts/benchmarks/benchmark_xform_prim_view.py b/scripts/benchmarks/benchmark_xform_prim_view.py index ae6656e3d8d0..947a9328d56e 100644 --- a/scripts/benchmarks/benchmark_xform_prim_view.py +++ b/scripts/benchmarks/benchmark_xform_prim_view.py @@ -166,6 +166,11 @@ 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.warp) wp.to_torch(new_positions)[:, 2] += 0.1 diff --git a/source/isaaclab/changelog.d/fabric-local-poses.rst b/source/isaaclab/changelog.d/fabric-local-poses.rst index 1ca61cfda97b..b33982e3961f 100644 --- a/source/isaaclab/changelog.d/fabric-local-poses.rst +++ b/source/isaaclab/changelog.d/fabric-local-poses.rst @@ -1,12 +1,12 @@ Added ^^^^^ -* Added explicit local/world scale methods - :meth:`~isaaclab.sim.views.BaseFrameView.get_local_scales`, - :meth:`~isaaclab.sim.views.BaseFrameView.set_local_scales`, - :meth:`~isaaclab.sim.views.BaseFrameView.get_world_scales`, and - :meth:`~isaaclab.sim.views.BaseFrameView.set_world_scales` to the FrameView - API, implemented for :class:`~isaaclab.sim.views.UsdFrameView`. +* 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`, @@ -20,9 +20,11 @@ Deprecated ^^^^^^^^^^ * Deprecated :meth:`~isaaclab.sim.views.BaseFrameView.get_scales` and - :meth:`~isaaclab.sim.views.BaseFrameView.set_scales` in favor of the explicit - ``get_local_scales`` / ``set_local_scales`` (operates on ``xformOp:scale``) or - ``get_world_scales`` / ``set_world_scales`` (operates on composed world-space - scale). The deprecated methods still work but emit a ``DeprecationWarning``; + :meth:`~isaaclab.sim.views.BaseFrameView.set_scales`. For reads, use + the explicit ``get_local_scales`` (operates on ``xformOp:scale``) or + ``get_world_scales`` (composed world-space scale). For writes, use + ``with view.xform_world_space_writer() as w: w.set_scales(...)`` (or + ``xform_local_space_writer``). The deprecated methods still work but + emit a ``DeprecationWarning``; :class:`~isaaclab.sim.views.UsdFrameView` preserves prior behavior by defaulting to local scales. diff --git a/source/isaaclab/changelog.d/xform-space-writer.rst b/source/isaaclab/changelog.d/xform-space-writer.rst index 41c9c97a55cc..51810d1f8585 100644 --- a/source/isaaclab/changelog.d/xform-space-writer.rst +++ b/source/isaaclab/changelog.d/xform-space-writer.rst @@ -28,14 +28,3 @@ Deprecated instead. The deprecated methods still work but emit a one-time ``DeprecationWarning`` per class and open a single-statement writer scope internally. - -Removed -^^^^^^^ - -* **Breaking:** Removed ``set_world_scales`` and ``set_local_scales`` - from :class:`~isaaclab.sim.views.BaseFrameView` (and all subclasses). - These were introduced in this release cycle without a stable downstream - user, so they are removed outright (no deprecation cycle). Use - ``with view.xform_world_space_writer() as w: w.set_scales(...)`` (or - :meth:`~isaaclab.sim.views.BaseFrameView.xform_local_space_writer`) - instead. diff --git a/source/isaaclab_newton/changelog.d/fabric-local-poses.rst b/source/isaaclab_newton/changelog.d/fabric-local-poses.rst index 4e18af38eb2b..c5c0b0f3c7d2 100644 --- a/source/isaaclab_newton/changelog.d/fabric-local-poses.rst +++ b/source/isaaclab_newton/changelog.d/fabric-local-poses.rst @@ -1,19 +1,19 @@ Added ^^^^^ -* Added :meth:`~isaaclab_newton.sim.views.NewtonSiteFrameView.get_local_scales`, - :meth:`~isaaclab_newton.sim.views.NewtonSiteFrameView.set_local_scales`, - :meth:`~isaaclab_newton.sim.views.NewtonSiteFrameView.get_world_scales`, and - :meth:`~isaaclab_newton.sim.views.NewtonSiteFrameView.set_world_scales` for - transform (xform) scales. These explicit APIs are intentionally separate from - Newton collision shape geometry sizes. +* 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 xform-scale ``get_world_scales`` / ``set_world_scales`` (or - their local equivalents). The deprecated methods still work but emit a - ``DeprecationWarning`` and preserve Newton's legacy collision shape - geometry-scale behavior. + 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_ovphysx/changelog.d/fabric-local-poses.rst b/source/isaaclab_ovphysx/changelog.d/fabric-local-poses.rst index eb8caa5d254f..fdf33f0b7d94 100644 --- a/source/isaaclab_ovphysx/changelog.d/fabric-local-poses.rst +++ b/source/isaaclab_ovphysx/changelog.d/fabric-local-poses.rst @@ -1,18 +1,19 @@ Added ^^^^^ -* Added :meth:`~isaaclab_ovphysx.sim.views.OvPhysxFrameView.get_local_scales`, - :meth:`~isaaclab_ovphysx.sim.views.OvPhysxFrameView.set_local_scales`, - :meth:`~isaaclab_ovphysx.sim.views.OvPhysxFrameView.get_world_scales`, and - :meth:`~isaaclab_ovphysx.sim.views.OvPhysxFrameView.set_world_scales`, which - delegate to the internal :class:`~isaaclab.sim.views.UsdFrameView`. +* 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` in favor of the - explicit ``get_local_scales`` / ``set_local_scales`` (operates on - ``xformOp:scale``) or ``get_world_scales`` / ``set_world_scales``. The - deprecated methods still work but emit a ``DeprecationWarning`` and default to - local scales, preserving prior behavior. + :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_physx/changelog.d/fabric-local-poses.rst b/source/isaaclab_physx/changelog.d/fabric-local-poses.rst index 546d395cb78e..94c48cf83690 100644 --- a/source/isaaclab_physx/changelog.d/fabric-local-poses.rst +++ b/source/isaaclab_physx/changelog.d/fabric-local-poses.rst @@ -15,32 +15,22 @@ Added Warp kernels that propagate ``local = world * inv(parent)`` and ``world = local * parent`` directly on Fabric storage matrices. -* Added Fabric-accelerated ``get_local_poses`` / ``set_local_poses`` to - :class:`~isaaclab_physx.sim.views.FabricFrameView`. - - Local-pose operations now use ``wp.indexedfabricarray`` to read/write +* 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 lazy per-view dirty tracking: ``set_local_poses`` marks the world - matrix dirty and vice-versa, triggering automatic re-propagation only on - the next read (no eager kernel launches on the write path). - -* Added interleave detection: interleaving ``set_world_poses`` and - ``set_local_poses`` on the same view within a frame flushes the stale - direction automatically and emits a one-time performance warning. - * Added topology-change recovery via automatic ``PrepareForReuse`` detection and per-selection index rebuild. Deprecated ^^^^^^^^^^ -* Deprecated ``get_scales`` / ``set_scales`` on all ``BaseFrameView`` subclasses. - Use the new explicit ``get_local_scales`` / ``set_local_scales`` (operates on - ``xformOp:scale`` / ``localMatrix``) or ``get_world_scales`` / - ``set_world_scales`` (operates on composed world-space scale) instead. - The deprecated methods still work but emit a ``DeprecationWarning``; - ``UsdFrameView`` defaults to local, ``FabricFrameView`` defaults to world - (preserving prior behavior). +* 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 index 25878c737b72..c87325bf69a5 100644 --- a/source/isaaclab_physx/changelog.d/xform-space-writer.rst +++ b/source/isaaclab_physx/changelog.d/xform-space-writer.rst @@ -12,10 +12,11 @@ Changed - 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 Kit's per-tick - ``updateWorldXforms()`` does not redundantly recompute matrices the - user just wrote. The renderer's independent ``omni:fabric:worldMatrix`` - listener is unaffected and observes the writes. + 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`` @@ -24,18 +25,12 @@ Changed semantics should snapshot the matrices themselves before entering the scope. - The lazy-dirty-flag mechanism (the ``_DirtyFlag`` enum, ``_dirty`` field, - ``_sync_*_if_dirty`` helpers, and the one-time - ``interleaved set_world_poses / set_local_poses`` warning) has been - removed -- the eager dual-write inside the scope makes all of that - unnecessary. - - Two persistent Fabric selections are built once during - ``_initialize_fabric`` and kept for the view's lifetime: ``_sel_ro`` + 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). The - writer flips a single ``_is_rw`` flag on enter/exit; neither selection - is rebuilt on the flip. Renderer correctness comes from the hierarchy - tracking pause above; the RO/RW split tells Kit's per-tick - ``updateWorldXforms()`` which matrices the user is currently - authoring. + (``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. From a12cce96531ee9d4bb4ccd28d3bef3f4024f9e8c Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Wed, 24 Jun 2026 09:11:51 +0000 Subject: [PATCH 49/54] fix(bench): exclude one-time init from Total / Overall speedup The benchmark mixes a one-time view-construction cost (`init`) with per-iteration steady-state ops in the same dict. The Total row and Overall speedup were both summing every value in that dict, so a backend whose first call materializes the view (NewtonSiteFrameView took ~1 s of stage population in the reported run) crushed the per-iter ops in the totals and printed a misleading 0.13x overall. Separate the two: keep Initialization in the per-op table for visibility, but compute Total and Overall over the per-iter operations only. The new label makes the scope explicit -- "Per-iter total" and "SPEEDUP vs USD (per-iter ops; one-time init excluded)". With the user's Blackwell RTX PRO 5000 numbers, this turns the nonsensical 0.13x Newton overall into 207x and the dampened 25x Fabric overall into 164x. --- .../benchmarks/benchmark_xform_prim_view.py | 26 ++++++++++++------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/scripts/benchmarks/benchmark_xform_prim_view.py b/scripts/benchmarks/benchmark_xform_prim_view.py index 947a9328d56e..821844d11c27 100644 --- a/scripts/benchmarks/benchmark_xform_prim_view.py +++ b/scripts/benchmarks/benchmark_xform_prim_view.py @@ -346,8 +346,14 @@ 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"), @@ -359,6 +365,7 @@ def print_results(results_dict: dict[str, dict[str, float]], num_prims: int, num ("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}" @@ -369,15 +376,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: @@ -387,7 +395,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: @@ -400,11 +408,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: From 4e9297623236c961b0dcb36141c430a165bedf64 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Wed, 24 Jun 2026 09:24:17 +0000 Subject: [PATCH 50/54] test: drop tests that assert against past code, rename stale survivors Tests should describe and exercise the CURRENT API, not history. Removed - test_set_world_scales_method_no_longer_exists (whitebox assertion that two never-shipped attributes are absent from BaseFrameView). Renamed + docstrings rewritten (names and docstrings referenced ``set_world_scales`` / ``set_local_scales`` which do not exist; the test bodies already use the writer scope): - frame_view_contract_utils.py: test_set_local_scales_roundtrip -> test_local_scales_roundtrip test_set_world_scales_roundtrip -> test_world_scales_roundtrip - test_views_xform_prim.py: test_set_local_scales_then_get_world_scales -> test_world_scale_composes_with_parent_scale test_set_world_scales_then_get_local_scales -> test_local_scale_inverts_parent_when_writing_world_scale Stale docstring references to the removed dirty-tracking and the deprecated set_*_scales API also corrected in three fabric tests (test_local_scales_roundtrip, test_world_scales_roundtrip, test_fabric_cuda1_scales_roundtrip, test_initial_seed_with_scaled_parent). --- .../test/sim/frame_view_contract_utils.py | 8 +++--- .../test/sim/test_views_xform_prim.py | 17 +++++++++--- .../test/sim/test_views_xform_prim_fabric.py | 26 ++++++------------- 3 files changed, 25 insertions(+), 26 deletions(-) diff --git a/source/isaaclab/test/sim/frame_view_contract_utils.py b/source/isaaclab/test/sim/frame_view_contract_utils.py index e67330f36791..7ff53025415e 100644 --- a/source/isaaclab/test/sim/frame_view_contract_utils.py +++ b/source/isaaclab/test/sim/frame_view_contract_utils.py @@ -469,8 +469,8 @@ def test_world_scales_default_identity(device, view_factory): @pytest.mark.parametrize("device", ["cpu", "cuda:0"]) -def test_set_local_scales_roundtrip(device, view_factory): - """set_local_scales -> get_local_scales returns the same values.""" +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) @@ -484,8 +484,8 @@ def test_set_local_scales_roundtrip(device, view_factory): @pytest.mark.parametrize("device", ["cpu", "cuda:0"]) -def test_set_world_scales_roundtrip(device, view_factory): - """set_world_scales -> get_world_scales returns the same values.""" +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) diff --git a/source/isaaclab/test/sim/test_views_xform_prim.py b/source/isaaclab/test/sim/test_views_xform_prim.py index 2fa404f26a7f..61b281ca5e9e 100644 --- a/source/isaaclab/test/sim/test_views_xform_prim.py +++ b/source/isaaclab/test/sim/test_views_xform_prim.py @@ -260,8 +260,12 @@ def _make_scaled_parent_child_view(device, parent_scale, child_scale=None): @pytest.mark.parametrize("device", ["cpu", "cuda"]) -def test_set_local_scales_then_get_world_scales(device): - """Under a scaled parent, world scale == parent_scale * local_scale.""" +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") @@ -276,8 +280,13 @@ def test_set_local_scales_then_get_world_scales(device): @pytest.mark.parametrize("device", ["cpu", "cuda"]) -def test_set_world_scales_then_get_local_scales(device): - """Under a scaled parent, set_world_scales writes local = world / parent_scale.""" +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") 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 10eb883940f3..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 @@ -389,7 +389,7 @@ def test_get_scales_fabric_path(device, view_factory): @pytest.mark.parametrize("device", ["cuda:0"]) def test_local_scales_roundtrip(device, view_factory): - """set_local_scales -> get_local_scales roundtrip via localMatrix.""" + """Writing scales through the local-space writer roundtrips via ``localMatrix``.""" bundle = view_factory(num_envs=2, device=device) view = bundle.view @@ -409,7 +409,7 @@ def test_local_scales_roundtrip(device, view_factory): @pytest.mark.parametrize("device", ["cuda:0"]) def test_world_scales_roundtrip(device, view_factory): - """set_world_scales -> get_world_scales roundtrip via worldMatrix.""" + """Writing scales through the world-space writer roundtrips via ``worldMatrix``.""" bundle = view_factory(num_envs=2, device=device) view = bundle.view @@ -522,8 +522,8 @@ def test_initial_seed_with_scaled_parent(device): 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, after ``_sync_world_from_local_if_dirty`` the world - scale collapses to (2, 1, 1). This test catches both regressions. + 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() @@ -702,11 +702,11 @@ 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_world_scales -> get_world_scales roundtrip works on cuda:1. + """World-space scale writes roundtrip via ``worldMatrix`` on ``cuda:1``. - Both write paths (``set_world_poses`` and ``set_world_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 @@ -1022,13 +1022,3 @@ def test_view_getter_inside_scope_raises(device, view_factory): view.get_local_scales() # After the scope exits, view getters work again. view.get_world_poses() - - -def test_set_world_scales_method_no_longer_exists(): - """``set_world_scales`` / ``set_local_scales`` were deleted from this PR's surface.""" - # The deprecated set_world_poses / set_local_poses shims remain (with warnings), - # but the never-shipped set_world_scales / set_local_scales were removed. - from isaaclab.sim.views import BaseFrameView # noqa: PLC0415 - - assert not hasattr(BaseFrameView, "set_world_scales") - assert not hasattr(BaseFrameView, "set_local_scales") From d6104a18b5e4bc4d033bccbb5204618d0e5cce88 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Wed, 24 Jun 2026 09:30:27 +0000 Subject: [PATCH 51/54] fix: sync after every getter, not only on the cached path FabricFrameView's getters (_get_world_poses_impl, _get_local_poses_impl, and the shared _decompose_scales for both scale getters) only called wp.synchronize() on the cached path (indices is None / slice(None)). The per-indices path returned a fresh ProxyArray without syncing, so a caller that did .numpy() or wp.to_torch().cpu() on it could observe zeros from the wp.zeros initializer (the kernel hadn't completed). Move the sync below the kernel launch and call it unconditionally. This: - removes the asymmetry between the cached and the indexed paths, - matches what the class docstring claims about getters being immediately readable, - is essentially free for the cached path (was already syncing) and cheap for the indexed path (the kernel itself is the dominant cost). Update the class docstring to describe the new behaviour: every getter launches its decompose kernel and synchronizes before returning, so the returned ProxyArray is immediately readable from GPU or host without any caller-side sync. Reported by Greptile on PR #5677. --- .../sim/views/fabric_frame_view.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) 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 df1f13eaecac..657db26a13ba 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 @@ -75,8 +75,10 @@ class FabricFrameView(BaseFrameView): (``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 read directly - from Fabric storage without any further synchronization. + 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), @@ -303,8 +305,13 @@ def _get_world_poses_impl(self, indices: wp.array | None = None) -> tuple[ProxyA 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) @@ -339,8 +346,9 @@ def _get_local_poses_impl(self, indices: wp.array | None = None) -> tuple[ProxyA device=self._device, ) + # See note in _get_world_poses_impl: sync regardless of caching path. + wp.synchronize() if use_cached: - wp.synchronize() return self._fabric_local_translations_ta, self._fabric_local_orientations_ta return ProxyArray(translations_wp), ProxyArray(orientations_wp) @@ -386,8 +394,9 @@ def _decompose_scales(self, ro_array, indices) -> ProxyArray: device=self._device, ) + # See note in _get_world_poses_impl: sync regardless of caching path. + wp.synchronize() if use_cached: - wp.synchronize() return self._fabric_scales_ta return ProxyArray(scales_wp) From 15646c66d78b18c9399b7e5bbea67d9d21fc302b Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Wed, 24 Jun 2026 09:55:37 +0000 Subject: [PATCH 52/54] refactor: dedupe _compute_fabric_indices via the existing primitive _compute_fabric_indices and _compute_parent_fabric_indices both walked selection.GetPaths(), built the same path->index dict, then iterated self.prim_paths to look up either the prim itself or its parent. The selection-walk + dict + lookup loop is exactly what _compute_fabric_indices_for(selection, paths) already does for one-off index arrays. Have the two specialised helpers delegate to it. The parent variant keeps its stage-root precondition inline (where it reads naturally, next to the rsplit) and builds the parent path list before calling the shared primitive. The child variant is a one-line passthrough. No behavioural change: same selection walks, same indices, same ordering. Saves ~18 lines and removes a duplicated dict construction. The shared primitive's docstring is updated to reflect its new role as the canonical lookup, no longer just a one-shot. --- .../sim/views/fabric_frame_view.py | 39 +++++++------------ 1 file changed, 14 insertions(+), 25 deletions(-) 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 657db26a13ba..6e2a91763cd9 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 @@ -512,39 +512,22 @@ def _rebuild_rw_arrays(self) -> None: # ------------------------------------------------------------------ def _compute_fabric_indices(self, selection) -> wp.array: - fabric_paths = selection.GetPaths() - path_to_fabric_idx: dict[str, int] = {str(p): i for i, p in enumerate(fabric_paths)} - indices: list[int] = [] - for prim_path in self.prim_paths: - fabric_idx = path_to_fabric_idx.get(prim_path) - if fabric_idx is None: - raise RuntimeError( - f"Prim '{prim_path}' not found in Fabric selection. Ensure the hierarchy has been populated." - ) - indices.append(fabric_idx) - return wp.array(indices, dtype=wp.int32, device=self._device) + """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: - """For each child in this view, look up the parent prim's fabric index.""" - fabric_paths = selection.GetPaths() - path_to_fabric_idx: dict[str, int] = {str(p): i for i, p in enumerate(fabric_paths)} - indices: list[int] = [] + """View-side indices that map each managed prim's parent into ``selection``.""" + parent_paths: list[str] = [] for prim_path in self.prim_paths: parent_path = prim_path.rsplit("/", 1)[0] - if parent_path == "": + if not parent_path: 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." ) - fabric_idx = path_to_fabric_idx.get(parent_path) - if fabric_idx is None: - raise RuntimeError( - f"Parent prim '{parent_path}' (for child '{prim_path}') not found in Fabric selection. " - "Ensure parents have Fabric world+local matrices populated." - ) - indices.append(fabric_idx) - return wp.array(indices, dtype=wp.int32, device=self._device) + parent_paths.append(parent_path) + return self._compute_fabric_indices_for(selection, parent_paths) def _build_indexed_array(self, selection, attribute_name: str, fabric_indices: wp.array) -> wp.indexedfabricarray: fa = wp.fabricarray(selection, attribute_name) @@ -737,7 +720,13 @@ def _sync_fabric_from_usd_initial(self) -> None: wp.synchronize() def _compute_fabric_indices_for(self, selection, paths: list[str]) -> wp.array: - """Path-dict lookup helper used to build one-shot indexed arrays for a custom path set.""" + """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`. + """ fabric_paths = selection.GetPaths() path_to_idx = {str(p): i for i, p in enumerate(fabric_paths)} indices: list[int] = [] From 70f14004e7cbaebf27216c787ab9397b7e519286 Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Wed, 24 Jun 2026 10:00:49 +0000 Subject: [PATCH 53/54] refactor: comprehensions over imperative loops in fabric index helpers Two helpers were building lists with .append() inside an imperative for-loop: - _compute_fabric_indices_for: walked paths -> looked up indices, raised on miss, appended to a list. - _compute_parent_fabric_indices: walked self.prim_paths -> derived parent path, validated stage-root, appended to a list. Factor the per-element step into a named local function whose body documents the intent (lookup + validation, parent derivation + validation), then drive it from a list comprehension. Same control flow, denser at the call site, the loop body's purpose is now a one-token name. --- .../sim/views/fabric_frame_view.py | 25 ++++++++++--------- 1 file changed, 13 insertions(+), 12 deletions(-) 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 6e2a91763cd9..898973c9c59d 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 @@ -517,17 +517,18 @@ def _compute_fabric_indices(self, selection) -> wp.array: def _compute_parent_fabric_indices(self, selection) -> wp.array: """View-side indices that map each managed prim's parent into ``selection``.""" - parent_paths: list[str] = [] - for prim_path in self.prim_paths: - parent_path = prim_path.rsplit("/", 1)[0] - if not parent_path: + + 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." ) - parent_paths.append(parent_path) - return self._compute_fabric_indices_for(selection, parent_paths) + 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) @@ -727,15 +728,15 @@ def _compute_fabric_indices_for(self, selection, paths: list[str]) -> wp.array: index arrays such as the parent-world seed in :meth:`_sync_fabric_from_usd_initial`. """ - fabric_paths = selection.GetPaths() - path_to_idx = {str(p): i for i, p in enumerate(fabric_paths)} - indices: list[int] = [] - for path in paths: + 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.") - indices.append(idx) - return wp.array(indices, dtype=wp.int32, device=self._device) + return idx + + return wp.array([lookup(p) for p in paths], dtype=wp.int32, device=self._device) # ---------------------------------------------------------------------- From 78242170095a27a05c512af0fb84403a6c7e777d Mon Sep 17 00:00:00 2001 From: Peter Verswyvelen Date: Thu, 25 Jun 2026 12:05:53 +0000 Subject: [PATCH 54/54] fix: restore clone-template suffix stripping for Newton site ref prim The world-body fallback in NewtonSiteFrameView._resolve_source_prim resolved the reference prim directly from source_root, instead of stripping the clone-template suffix as the prior code did. For heterogeneous (multi-asset spawner) scenes this resolved the wrong reference frame, offsetting site world poses and breaking rendering-correctness (dexsuite kuka hetero kitless). Restore the split_clone_template/get_suffix-based ref_path computation (and the dropped imports). Homogeneous scenes were unaffected; hetero rgb/depth now match golden images again. --- .../sim/views/newton_site_frame_view.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) 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 d28d27f34730..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,7 +14,7 @@ from pxr import UsdPhysics import isaaclab.sim as sim_utils -from isaaclab.cloner.cloner_utils import 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 @@ -359,7 +359,13 @@ def _resolve_source_prim( return body_patterns, wp.transform(pos, quat), scale, False, env_ids body_prim = body_prim.GetParent() - ref_prim = stage.GetPrimAtPath(source_root) if source_root is not None else None + ref_path = source_root + if source_root is not None and destination_template is not None: + 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), scale, source_root is not None, env_ids