From 7af1458fb7a8ad68227f7db618917e77d89f17ca Mon Sep 17 00:00:00 2001 From: AbdelStark Date: Thu, 26 Mar 2026 18:46:53 +0100 Subject: [PATCH 1/2] Optimize UR5e analytical Jacobian path for OmniReset OSC Signed-off-by: AbdelStark --- source/uwlab_assets/config/extension.toml | 2 +- source/uwlab_assets/docs/CHANGELOG.rst | 9 + .../uwlab_assets/test/test_ur5e_kinematics.py | 174 ++++++++++++++++++ .../robots/ur5e_robotiq_gripper/kinematics.py | 143 ++++++++------ source/uwlab_tasks/config/extension.toml | 2 +- source/uwlab_tasks/docs/CHANGELOG.rst | 9 + .../mdp/actions/task_space_actions.py | 2 +- 7 files changed, 285 insertions(+), 56 deletions(-) create mode 100644 source/uwlab_assets/test/test_ur5e_kinematics.py diff --git a/source/uwlab_assets/config/extension.toml b/source/uwlab_assets/config/extension.toml index dd819024..f6dc8019 100644 --- a/source/uwlab_assets/config/extension.toml +++ b/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" diff --git a/source/uwlab_assets/docs/CHANGELOG.rst b/source/uwlab_assets/docs/CHANGELOG.rst index cb39cdbe..9f32983f 100644 --- a/source/uwlab_assets/docs/CHANGELOG.rst +++ b/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) ~~~~~~~~~~~~~~~~~~ diff --git a/source/uwlab_assets/test/test_ur5e_kinematics.py b/source/uwlab_assets/test/test_ur5e_kinematics.py new file mode 100644 index 00000000..65935af4 --- /dev/null +++ b/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 diff --git a/source/uwlab_assets/uwlab_assets/robots/ur5e_robotiq_gripper/kinematics.py b/source/uwlab_assets/uwlab_assets/robots/ur5e_robotiq_gripper/kinematics.py index cd4d3991..6ca551ec 100644 --- a/source/uwlab_assets/uwlab_assets/robots/ur5e_robotiq_gripper/kinematics.py +++ b/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 # ============================================================================ diff --git a/source/uwlab_tasks/config/extension.toml b/source/uwlab_tasks/config/extension.toml index 37dbd4f7..b31ffb48 100644 --- a/source/uwlab_tasks/config/extension.toml +++ b/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" diff --git a/source/uwlab_tasks/docs/CHANGELOG.rst b/source/uwlab_tasks/docs/CHANGELOG.rst index 44050103..579b06e0 100644 --- a/source/uwlab_tasks/docs/CHANGELOG.rst +++ b/source/uwlab_tasks/docs/CHANGELOG.rst @@ -1,6 +1,15 @@ Changelog --------- +0.13.9 (2026-03-26) +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* updated OmniReset RelCartesianOSCAction to use the cached UR5e analytical Jacobian path + + 0.13.8 (2025-10-24) ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/omnireset/mdp/actions/task_space_actions.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/omnireset/mdp/actions/task_space_actions.py index 857610bb..b501c13d 100644 --- a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/omnireset/mdp/actions/task_space_actions.py +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/omnireset/mdp/actions/task_space_actions.py @@ -146,7 +146,7 @@ def apply_actions(self): joint_vel = self._asset.data.joint_vel[:, self._joint_ids] # Analytical Jacobian (base_link frame, matching EE pose frame) - jacobian = compute_jacobian_analytical(joint_pos, device=str(self.device)) + jacobian = compute_jacobian_analytical(joint_pos) # EE velocity from J @ dq (consistent with analytical Jacobian) ee_vel = torch.bmm(jacobian, joint_vel.unsqueeze(-1)).squeeze(-1) # (N, 6) From 0567b465bb3aaaeb5bd2c8fb371f7a8a5bd9ee47 Mon Sep 17 00:00:00 2001 From: AbdelStark Date: Thu, 26 Mar 2026 18:47:54 +0100 Subject: [PATCH 2/2] Add contributor entry for AbdelStark Signed-off-by: AbdelStark --- CONTRIBUTORS.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 73998c5d..5045d97f 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -19,6 +19,7 @@ Guidelines for modifications: --- +* AbdelStark * Feng Yu * Mateo Guaman Castro * Patrick Yin