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..baaf33ceb8a3 --- /dev/null +++ b/source/isaaclab/changelog.d/indexed-fabric-kernels.rst @@ -0,0 +1,18 @@ +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/isaaclab/utils/warp/fabric.py b/source/isaaclab/isaaclab/utils/warp/fabric.py index a48f773f4991..f665323cbe34 100644 --- a/source/isaaclab/isaaclab/utils/warp/fabric.py +++ b/source/isaaclab/isaaclab/utils/warp/fabric.py @@ -18,12 +18,14 @@ if TYPE_CHECKING: FabricArrayUInt32 = Any FabricArrayMat44d = Any + IndexedFabricArrayMat44d = Any ArrayUInt32 = Any ArrayUInt32_1d = Any ArrayFloat32_2d = Any else: FabricArrayUInt32 = wp.fabricarray(dtype=wp.uint32) FabricArrayMat44d = wp.fabricarray(dtype=wp.mat44d) + IndexedFabricArrayMat44d = wp.indexedfabricarray(dtype=wp.mat44d) ArrayUInt32 = wp.array(ndim=1, dtype=wp.uint32) ArrayUInt32_1d = wp.array(dtype=wp.uint32) ArrayFloat32_2d = wp.array(ndim=2, dtype=wp.float32) @@ -163,6 +165,192 @@ def compose_fabric_transformation_matrix_from_warp_arrays( ) +@wp.kernel(enable_backward=False) +def decompose_indexed_fabric_transforms( + fabric_matrices: IndexedFabricArrayMat44d, + array_positions: ArrayFloat32_2d, + array_orientations: ArrayFloat32_2d, + array_scales: ArrayFloat32_2d, + indices: ArrayUInt32, +): + """Decompose indexed Fabric transformation matrices into position, orientation, and scale. + + Like :func:`decompose_fabric_transformation_matrix_to_warp_arrays` but operates on a + :class:`wp.indexedfabricarray` that already encodes the view-to-fabric mapping, removing + the need for a separate ``mapping`` array. + + Args: + fabric_matrices: Indexed fabric array containing 4x4 transformation matrices. + array_positions: Output array for positions [m], shape (N, 3). + array_orientations: Output array for quaternions in xyzw format, shape (N, 4). + array_scales: Output array for scales, shape (N, 3). + indices: View indices to process (subset selection). + """ + output_index = wp.tid() + view_index = indices[output_index] + + position, rotation, scale = _decompose_transformation_matrix(wp.mat44f(fabric_matrices[view_index])) + + if array_positions.shape[0] > 0: + array_positions[output_index, 0] = position[0] + array_positions[output_index, 1] = position[1] + array_positions[output_index, 2] = position[2] + if array_orientations.shape[0] > 0: + array_orientations[output_index, 0] = rotation[0] + array_orientations[output_index, 1] = rotation[1] + array_orientations[output_index, 2] = rotation[2] + array_orientations[output_index, 3] = rotation[3] + if array_scales.shape[0] > 0: + array_scales[output_index, 0] = scale[0] + array_scales[output_index, 1] = scale[1] + array_scales[output_index, 2] = scale[2] + + +@wp.kernel(enable_backward=False) +def compose_indexed_fabric_transforms( + fabric_matrices: IndexedFabricArrayMat44d, + array_positions: ArrayFloat32_2d, + array_orientations: ArrayFloat32_2d, + array_scales: ArrayFloat32_2d, + broadcast_positions: bool, + broadcast_orientations: bool, + broadcast_scales: bool, + indices: ArrayUInt32, +): + """Compose indexed Fabric transformation matrices from position, orientation, and scale. + + Like :func:`compose_fabric_transformation_matrix_from_warp_arrays` but operates on a + :class:`wp.indexedfabricarray` that already encodes the view-to-fabric mapping, removing + the need for a separate ``mapping`` array. + + Args: + fabric_matrices: Indexed fabric array containing 4x4 transformation matrices to update. + array_positions: Input array for positions [m], shape (N, 3). + array_orientations: Input array for quaternions in xyzw format, shape (N, 4). + array_scales: Input array for scales, shape (N, 3). + broadcast_positions: If True, use first position for all prims. + broadcast_orientations: If True, use first orientation for all prims. + broadcast_scales: If True, use first scale for all prims. + indices: View indices to process (subset selection). + """ + i = wp.tid() + view_index = indices[i] + position, rotation, scale = _decompose_transformation_matrix(wp.mat44f(fabric_matrices[view_index])) + + if array_positions.shape[0] > 0: + if broadcast_positions: + index = 0 + else: + index = i + position[0] = array_positions[index, 0] + position[1] = array_positions[index, 1] + position[2] = array_positions[index, 2] + if array_orientations.shape[0] > 0: + if broadcast_orientations: + index = 0 + else: + index = i + rotation[0] = array_orientations[index, 0] + rotation[1] = array_orientations[index, 1] + rotation[2] = array_orientations[index, 2] + rotation[3] = array_orientations[index, 3] + if array_scales.shape[0] > 0: + if broadcast_scales: + index = 0 + else: + index = i + scale[0] = array_scales[index, 0] + scale[1] = array_scales[index, 1] + scale[2] = array_scales[index, 2] + + fabric_matrices[view_index] = wp.mat44d( # type: ignore[arg-type] + wp.transpose(wp.transform_compose(position, rotation, scale)) # type: ignore[arg-type] + ) + + +@wp.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, + parent_world_matrices: IndexedFabricArrayMat44d, + child_local_matrices: IndexedFabricArrayMat44d, + indices: ArrayUInt32, +): + """Recompute child localMatrix from (parent worldMatrix, child worldMatrix). + + Computes ``child_local = inv(parent_world) * child_world`` per prim and writes the + result back to the child's :data:`omni:fabric:localMatrix` so that subsequent + ``get_local_poses`` calls see consistent values after a world-pose write. + + All three indexed arrays are expected to be indexed by the same per-view indices + (i.e. ``view_to_child_fabric``, ``view_to_parent_fabric``, ``view_to_child_fabric``) + so the kernel only needs the view-side indices. + + Storage convention: Fabric matrices are stored as the transpose of the standard + column-major math convention. Math is ``local = inv(parent) * world``; under + the transpose identity ``(A * B)^T = B^T * A^T`` (and ``inv(A^T) = inv(A)^T``) + that is equivalent to storage-side ``local^T = world^T * inv(parent^T)``, so we + can compute it directly on the stored matrices without explicit transposes. + + Args: + child_world_matrices: Indexed fabric array of child world matrices (read). + parent_world_matrices: Indexed fabric array of parent world matrices (read). + child_local_matrices: Indexed fabric array of child local matrices (written). + indices: View indices to process. + """ + i = wp.tid() + view_index = indices[i] + child_world = wp.mat44f(child_world_matrices[view_index]) + parent_world = wp.mat44f(parent_world_matrices[view_index]) + child_local_matrices[view_index] = wp.mat44d( # type: ignore[arg-type] + _local_from_world_transposed(child_world, parent_world) + ) + + +@wp.kernel(enable_backward=False) +def update_indexed_world_matrix_from_local( + child_local_matrices: IndexedFabricArrayMat44d, + parent_world_matrices: IndexedFabricArrayMat44d, + child_world_matrices: IndexedFabricArrayMat44d, + indices: ArrayUInt32, +): + """Recompute child worldMatrix from (parent worldMatrix, child localMatrix). + + Computes ``child_world = parent_world * child_local`` per prim and writes the + result back to the child's :data:`omni:fabric:worldMatrix`. Used after a + ``set_local_poses`` write so that subsequent ``get_world_poses`` calls see + consistent values. Mirror of :func:`update_indexed_local_matrix_from_world`. + + Args: + child_local_matrices: Indexed fabric array of child local matrices (read). + parent_world_matrices: Indexed fabric array of parent world matrices (read). + child_world_matrices: Indexed fabric array of child world matrices (written). + indices: View indices to process. + + Storage convention: same as :func:`update_indexed_local_matrix_from_world`. + Math is ``world = parent * local``; under the transpose identity that becomes + storage-side ``world^T = local^T * parent^T``, no explicit transposes needed. + """ + i = wp.tid() + view_index = indices[i] + child_local = wp.mat44f(child_local_matrices[view_index]) + parent_world = wp.mat44f(parent_world_matrices[view_index]) + child_world_matrices[view_index] = wp.mat44d( # type: ignore[arg-type] + _world_from_local_transposed(child_local, parent_world) + ) + + @wp.func def _decompose_transformation_matrix(m: Any): # -> tuple[wp.vec3f, wp.quatf, wp.vec3f] """Decompose a 4x4 transformation matrix into position, orientation, and scale. diff --git a/source/isaaclab/test/utils/warp/test_fabric_kernels.py b/source/isaaclab/test/utils/warp/test_fabric_kernels.py new file mode 100644 index 000000000000..de48afadab24 --- /dev/null +++ b/source/isaaclab/test/utils/warp/test_fabric_kernels.py @@ -0,0 +1,164 @@ +# 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)