Speed up UR5e analytical Jacobian path for OmniReset RelCartesianOSC

CONTRIBUTORS.md

@@ -19,6 +19,7 @@ Guidelines for modifications:
 
 ---
 
+* AbdelStark
 * Feng Yu
 * Mateo Guaman Castro
 * Patrick Yin

source/uwlab_assets/config/extension.toml

@@ -1,6 +1,6 @@
 [package]
 # Semantic Versioning is used: https://semver.org/
-version = "0.5.2"
+version = "0.5.3"
 
 # Description
 title =  "UW Lab Assets"

source/uwlab_assets/docs/CHANGELOG.rst

@@ -1,6 +1,15 @@
 Changelog
 ---------
 
+0.5.3 (2026-03-26)
+~~~~~~~~~~~~~~~~~~
+
+Changed
+^^^^^^^
+
+* cached calibrated UR5e analytical Jacobian constants to reduce per-step RelCartesianOSC overhead
+
+
 0.5.2 (2025-03-23)
 ~~~~~~~~~~~~~~~~~~
 

source/uwlab_assets/test/test_ur5e_kinematics.py

@@ -0,0 +1,174 @@
+# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md).
+# All Rights Reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+# ignore private usage of cache helpers in unit tests
+# pyright: reportPrivateUsage=none
+
+import importlib.util
+import sys
+import types
+from pathlib import Path
+
+import pytest
+import torch
+
+MODULE_PATH = (
+    Path(__file__).resolve().parents[1] / "uwlab_assets" / "robots" / "ur5e_robotiq_gripper" / "kinematics.py"
+)
+
+
+def _load_kinematics_module(monkeypatch):
+    """Load the kinematics module with a minimal isaaclab stub for unit testing."""
+    isaaclab_module = types.ModuleType("isaaclab")
+    utils_module = types.ModuleType("isaaclab.utils")
+    assets_module = types.ModuleType("isaaclab.utils.assets")
+    assets_module.retrieve_file_path = lambda path, download_dir=None: path
+    utils_module.assets = assets_module
+    isaaclab_module.utils = utils_module
+    monkeypatch.setitem(sys.modules, "isaaclab", isaaclab_module)
+    monkeypatch.setitem(sys.modules, "isaaclab.utils", utils_module)
+    monkeypatch.setitem(sys.modules, "isaaclab.utils.assets", assets_module)
+
+    module_name = "uwlab_test_ur5e_kinematics"
+    spec = importlib.util.spec_from_file_location(module_name, MODULE_PATH)
+    module = importlib.util.module_from_spec(spec)
+    monkeypatch.setitem(sys.modules, module_name, module)
+    assert spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+def _make_calibration(dtype: torch.dtype = torch.float32) -> dict[str, torch.Tensor]:
+    return {
+        "joints_xyz": torch.tensor(
+            [
+                [0.0, 0.0, 0.1625],
+                [-0.4250, 0.0, 0.0],
+                [-0.3922, 0.0, 0.0],
+                [0.0, -0.1333, 0.1333],
+                [0.0, 0.0997, 0.0996],
+                [0.0, -0.0996, 0.0997],
+            ],
+            dtype=dtype,
+        ),
+        "joints_rpy": torch.tensor(
+            [
+                [0.0, 0.0, 0.0],
+                [0.0, 1.5708, 0.0],
+                [0.0, 0.0, 0.0],
+                [1.5708, 0.0, 0.0],
+                [-1.5708, 0.0, 0.0],
+                [1.5708, 0.0, 0.0],
+            ],
+            dtype=dtype,
+        ),
+        "link_masses": torch.ones(6, dtype=dtype),
+        "link_coms": torch.zeros(6, 3, dtype=dtype),
+        "link_inertias": torch.zeros(6, 6, dtype=dtype),
+    }
+
+
+def _reference_compute_jacobian(
+    module, joint_angles: torch.Tensor, calibration: dict[str, torch.Tensor]
+) -> torch.Tensor:
+    """Reference implementation copied from the pre-optimization Jacobian path."""
+    device = joint_angles.device
+    dtype = joint_angles.dtype
+    batch_size = joint_angles.shape[0]
+    xyz_all = calibration["joints_xyz"].to(device=device, dtype=dtype)
+    rpy_all = calibration["joints_rpy"].to(device=device, dtype=dtype)
+    base_rotation = module.R_180Z.to(device=device, dtype=dtype)
+
+    transform = torch.eye(4, device=device, dtype=dtype).unsqueeze(0).repeat(batch_size, 1, 1)
+    for joint_idx in range(module.NUM_ARM_JOINTS):
+        fixed_rot = module.rpy_to_matrix_torch(rpy_all[joint_idx])
+        fixed_tf = torch.eye(4, device=device, dtype=dtype)
+        fixed_tf[:3, :3] = fixed_rot
+        fixed_tf[:3, 3] = xyz_all[joint_idx]
+        fixed_tf = fixed_tf.unsqueeze(0).repeat(batch_size, 1, 1)
+        theta = joint_angles[:, joint_idx]
+        ct, st = torch.cos(theta), torch.sin(theta)
+        joint_tf = torch.eye(4, device=device, dtype=dtype).unsqueeze(0).repeat(batch_size, 1, 1)
+        joint_tf[:, 0, 0] = ct
+        joint_tf[:, 0, 1] = -st
+        joint_tf[:, 1, 0] = st
+        joint_tf[:, 1, 1] = ct
+        transform = torch.bmm(torch.bmm(transform, fixed_tf), joint_tf)
+    ee_pos = transform[:, :3, 3]
+
+    jacobian = torch.zeros(batch_size, 6, module.NUM_ARM_JOINTS, device=device, dtype=dtype)
+    transform = torch.eye(4, device=device, dtype=dtype).unsqueeze(0).repeat(batch_size, 1, 1)
+    for joint_idx in range(module.NUM_ARM_JOINTS):
+        fixed_rot = module.rpy_to_matrix_torch(rpy_all[joint_idx])
+        fixed_tf = torch.eye(4, device=device, dtype=dtype)
+        fixed_tf[:3, :3] = fixed_rot
+        fixed_tf[:3, 3] = xyz_all[joint_idx]
+        fixed_tf = fixed_tf.unsqueeze(0).repeat(batch_size, 1, 1)
+        joint_frame_tf = torch.bmm(transform, fixed_tf)
+        joint_axis = joint_frame_tf[:, :3, 2]
+        joint_pos = joint_frame_tf[:, :3, 3]
+        jacobian[:, :3, joint_idx] = torch.cross(joint_axis, ee_pos - joint_pos, dim=1)
+        jacobian[:, 3:, joint_idx] = joint_axis
+        theta = joint_angles[:, joint_idx]
+        ct, st = torch.cos(theta), torch.sin(theta)
+        joint_rot_tf = torch.eye(4, device=device, dtype=dtype).unsqueeze(0).repeat(batch_size, 1, 1)
+        joint_rot_tf[:, 0, 0] = ct
+        joint_rot_tf[:, 0, 1] = -st
+        joint_rot_tf[:, 1, 0] = st
+        joint_rot_tf[:, 1, 1] = ct
+        transform = torch.bmm(joint_frame_tf, joint_rot_tf)
+
+    base_rotation = base_rotation.unsqueeze(0).repeat(batch_size, 1, 1)
+    jacobian[:, :3, :] = torch.bmm(base_rotation, jacobian[:, :3, :])
+    jacobian[:, 3:, :] = torch.bmm(base_rotation, jacobian[:, 3:, :])
+    return jacobian
+
+
+@pytest.fixture
+def kinematics_module(monkeypatch):
+    return _load_kinematics_module(monkeypatch)
+
+
+@pytest.mark.parametrize("batch_size", [1, 32])
+@pytest.mark.parametrize(
+    "device",
+    [
+        "cpu",
+        pytest.param("cuda", marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable")),
+    ],
+)
+def test_compute_jacobian_matches_reference(kinematics_module, monkeypatch, batch_size, device):
+    calibration = _make_calibration()
+    monkeypatch.setattr(kinematics_module, "_load_calibration", lambda: calibration)
+    kinematics_module._get_calibrated_kinematics_cache.cache_clear()
+
+    torch.manual_seed(7)
+    joint_angles = torch.randn(batch_size, kinematics_module.NUM_ARM_JOINTS, device=device, dtype=torch.float32)
+
+    fast_jacobian = kinematics_module.compute_jacobian_analytical(joint_angles)
+    reference_jacobian = _reference_compute_jacobian(kinematics_module, joint_angles, calibration)
+
+    assert torch.allclose(fast_jacobian, reference_jacobian, atol=1e-6, rtol=1e-5)
+
+
+def test_compute_jacobian_reuses_cached_fixed_rotations(kinematics_module, monkeypatch):
+    calibration = _make_calibration()
+    monkeypatch.setattr(kinematics_module, "_load_calibration", lambda: calibration)
+    kinematics_module._get_calibrated_kinematics_cache.cache_clear()
+
+    original_rpy_to_matrix = kinematics_module.rpy_to_matrix_torch
+    call_count = {"value": 0}
+
+    def counting_rpy_to_matrix(rpy: torch.Tensor) -> torch.Tensor:
+        call_count["value"] += 1
+        return original_rpy_to_matrix(rpy)
+
+    monkeypatch.setattr(kinematics_module, "rpy_to_matrix_torch", counting_rpy_to_matrix)
+
+    joint_angles = torch.randn(8, kinematics_module.NUM_ARM_JOINTS, dtype=torch.float32)
+    kinematics_module.compute_jacobian_analytical(joint_angles)
+    kinematics_module.compute_jacobian_analytical(joint_angles.clone())
+
+    assert call_count["value"] == 1

source/uwlab_assets/uwlab_assets/robots/ur5e_robotiq_gripper/kinematics.py

@@ -23,6 +23,8 @@
 import functools
 import os
 import tempfile
+from dataclasses import dataclass
+
 import torch
 import yaml
 
@@ -47,6 +49,15 @@
 R_180Z = torch.tensor([[-1, 0, 0], [0, -1, 0], [0, 0, 1]], dtype=torch.float32)
 
 
+@dataclass(frozen=True)
+class _CalibratedKinematicsCache:
+    """Device-local calibrated constants reused across Jacobian evaluations."""
+
+    joints_xyz: torch.Tensor
+    fixed_rotations: torch.Tensor
+    base_rotation: torch.Tensor
+
+
 # ============================================================================
 # Lazy-loaded calibration data (from metadata.yaml next to the robot USD)
 # ============================================================================
@@ -111,73 +122,99 @@ def rpy_to_matrix_torch(rpy: torch.Tensor) -> torch.Tensor:
         return R
 
 
+def _resolve_kinematics_device(device: torch.device | str | None, fallback: torch.device) -> torch.device:
+    """Resolve an optional explicit device override for analytical kinematics."""
+    if device is None:
+        return fallback
+    return torch.device(device)
+
+
+@functools.lru_cache(maxsize=None)
+def _get_calibrated_kinematics_cache(
+    device_type: str, device_index: int, dtype: torch.dtype
+) -> _CalibratedKinematicsCache:
+    """Materialize calibrated UR5e constants once per device/dtype."""
+    resolved_device = torch.device(device_type if device_index < 0 else f"{device_type}:{device_index}")
+    calibration = _load_calibration()
+    joints_xyz = calibration["joints_xyz"].to(device=resolved_device, dtype=dtype)
+    joints_rpy = calibration["joints_rpy"].to(device=resolved_device, dtype=dtype)
+    return _CalibratedKinematicsCache(
+        joints_xyz=joints_xyz,
+        fixed_rotations=rpy_to_matrix_torch(joints_rpy),
+        base_rotation=R_180Z.to(device=resolved_device, dtype=dtype),
+    )
+
+
 # ============================================================================
 # Analytical Jacobian
 # ============================================================================
 
 
-def compute_jacobian_analytical(joint_angles: torch.Tensor, device: str = "cuda") -> torch.Tensor:
+def compute_jacobian_analytical(
+    joint_angles: torch.Tensor, device: torch.device | str | None = None
+) -> torch.Tensor:
     """Compute geometric Jacobian using calibrated kinematics (batched).
 
     Computes to wrist_3_link frame origin (NOT COM), matching real robot code.
 
     Args:
         joint_angles: (N, 6) joint angles in radians.
+        device: Optional explicit device override. Defaults to ``joint_angles.device``.
     Returns:
         J: (N, 6, 6) Jacobian [linear; angular].
     """
-    N = joint_angles.shape[0]
-    cal = _load_calibration()
-    xyz_all = cal["joints_xyz"].to(device)
-    rpy_all = cal["joints_rpy"].to(device)
-    R_180Z_dev = R_180Z.to(device)
-
-    # FK to get EE position
-    T = torch.eye(4, device=device, dtype=torch.float32).unsqueeze(0).repeat(N, 1, 1)
-    for i in range(6):
-        R_fixed = rpy_to_matrix_torch(rpy_all[i])
-        T_fixed = torch.eye(4, device=device, dtype=torch.float32)
-        T_fixed[:3, :3] = R_fixed
-        T_fixed[:3, 3] = xyz_all[i]
-        T_fixed = T_fixed.unsqueeze(0).repeat(N, 1, 1)
-        theta = joint_angles[:, i]
-        ct, st = torch.cos(theta), torch.sin(theta)
-        T_joint = torch.eye(4, device=device, dtype=torch.float32).unsqueeze(0).repeat(N, 1, 1)
-        T_joint[:, 0, 0] = ct
-        T_joint[:, 0, 1] = -st
-        T_joint[:, 1, 0] = st
-        T_joint[:, 1, 1] = ct
-        T = torch.bmm(torch.bmm(T, T_fixed), T_joint)
-    p_ee = T[:, :3, 3]
-
-    # Jacobian columns
-    J = torch.zeros(N, 6, 6, device=device, dtype=torch.float32)
-    T = torch.eye(4, device=device, dtype=torch.float32).unsqueeze(0).repeat(N, 1, 1)
-    for i in range(6):
-        R_fixed = rpy_to_matrix_torch(rpy_all[i])
-        T_fixed = torch.eye(4, device=device, dtype=torch.float32)
-        T_fixed[:3, :3] = R_fixed
-        T_fixed[:3, 3] = xyz_all[i]
-        T_fixed = T_fixed.unsqueeze(0).repeat(N, 1, 1)
-        T_joint_frame = torch.bmm(T, T_fixed)
-        z_i = T_joint_frame[:, :3, 2]
-        p_i = T_joint_frame[:, :3, 3]
-        J[:, :3, i] = torch.cross(z_i, p_ee - p_i, dim=1)
-        J[:, 3:, i] = z_i
-        theta = joint_angles[:, i]
-        ct, st = torch.cos(theta), torch.sin(theta)
-        T_joint_rot = torch.eye(4, device=device, dtype=torch.float32).unsqueeze(0).repeat(N, 1, 1)
-        T_joint_rot[:, 0, 0] = ct
-        T_joint_rot[:, 0, 1] = -st
-        T_joint_rot[:, 1, 0] = st
-        T_joint_rot[:, 1, 1] = ct
-        T = torch.bmm(T_joint_frame, T_joint_rot)
-
-    # Rotate from base_link_inertia to base_link (REP-103)
-    R_180Z_batch = R_180Z_dev.unsqueeze(0).repeat(N, 1, 1)
-    J[:, :3, :] = torch.bmm(R_180Z_batch, J[:, :3, :])
-    J[:, 3:, :] = torch.bmm(R_180Z_batch, J[:, 3:, :])
-    return J
+    if joint_angles.ndim != 2 or joint_angles.shape[1] != NUM_ARM_JOINTS:
+        raise ValueError(f"Expected joint_angles to have shape (N, {NUM_ARM_JOINTS}), got {tuple(joint_angles.shape)}")
+    if not joint_angles.is_floating_point():
+        raise TypeError("joint_angles must be a floating point tensor")
+
+    resolved_device = _resolve_kinematics_device(device, joint_angles.device)
+    joint_angles = joint_angles.to(device=resolved_device)
+    batch_size = joint_angles.shape[0]
+    device_index = resolved_device.index if resolved_device.index is not None else -1
+    constants = _get_calibrated_kinematics_cache(resolved_device.type, device_index, joint_angles.dtype)
+
+    # Use a compact recurrence over rotation/translation instead of rebuilding 4x4 transforms
+    # at every step. This is the hot path used by RelCartesianOSCAction.
+    rot_world = (
+        torch.eye(3, device=resolved_device, dtype=joint_angles.dtype)
+        .unsqueeze(0)
+        .expand(batch_size, -1, -1)
+        .clone()
+    )
+    pos_world = joint_angles.new_zeros(batch_size, 3)
+    joint_axes = []
+    joint_positions = []
+
+    for joint_idx in range(NUM_ARM_JOINTS):
+        joint_frame_rot = torch.matmul(rot_world, constants.fixed_rotations[joint_idx])
+        joint_frame_pos = pos_world + torch.matmul(rot_world, constants.joints_xyz[joint_idx])
+        joint_axes.append(joint_frame_rot[:, :, 2])
+        joint_positions.append(joint_frame_pos)
+
+        theta = joint_angles[:, joint_idx]
+        ct = torch.cos(theta).unsqueeze(-1)
+        st = torch.sin(theta).unsqueeze(-1)
+        rot_world = torch.stack(
+            (
+                ct * joint_frame_rot[:, :, 0] + st * joint_frame_rot[:, :, 1],
+                -st * joint_frame_rot[:, :, 0] + ct * joint_frame_rot[:, :, 1],
+                joint_frame_rot[:, :, 2],
+            ),
+            dim=-1,
+        )
+        pos_world = joint_frame_pos
+
+    joint_axes_tensor = torch.stack(joint_axes, dim=-1)
+    joint_positions_tensor = torch.stack(joint_positions, dim=-1)
+    ee_offsets = pos_world.unsqueeze(-1) - joint_positions_tensor
+    linear_jacobian = torch.cross(joint_axes_tensor, ee_offsets, dim=1)
+
+    base_rotation = constants.base_rotation.unsqueeze(0).expand(batch_size, -1, -1)
+    jacobian = joint_angles.new_empty(batch_size, 6, NUM_ARM_JOINTS)
+    jacobian[:, :3, :] = torch.bmm(base_rotation, linear_jacobian)
+    jacobian[:, 3:, :] = torch.bmm(base_rotation, joint_axes_tensor)
+    return jacobian
 
 
 # ============================================================================

source/uwlab_tasks/config/extension.toml

@@ -1,7 +1,7 @@
 [package]
 
 # Semantic Versioning is used: https://semver.org/
-version = "0.13.8"
+version = "0.13.9"
 
 # Description
 title = "UW Lab Tasks"