Files
mindbot/scripts/environments/teleoperation/xr_utils/geometry.py
2026-03-05 22:41:56 +08:00

204 lines
6.4 KiB
Python

# Copyright (c) 2022-2026, The Mindbot Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
#
# Adapted from XRoboToolkit (https://github.com/NVIDIA-AI-IOT/XRoboToolkit-Teleop-Sample-Python)
# Original geometry utilities by Zhigen Zhao.
"""Geometry utilities for XR-to-World coordinate transformations.
This module provides the core coordinate transformation matrices and functions
needed to convert PICO 4 Ultra (OpenXR) data into Isaac Sim world coordinates.
Coordinate Systems:
- OpenXR (PICO 4): +X Right, +Y Up, -Z Forward
- Isaac Sim (World): +X Forward, +Y Left, +Z Up
"""
import numpy as np
from scipy.spatial.transform import Rotation as R
# =====================================================================
# Core Coordinate Transform: OpenXR -> Isaac Sim World
# =====================================================================
R_HEADSET_TO_WORLD = np.array(
[
[0, 0, -1], # XR -Z (forward) -> World +X (forward)
[-1, 0, 0], # XR -X (left) -> World +Y (left)
[0, 1, 0], # XR +Y (up) -> World +Z (up)
],
dtype=np.float64,
)
"""3x3 rotation matrix that maps OpenXR coordinates to Isaac Sim world coordinates."""
# =====================================================================
# Quaternion Utilities
# =====================================================================
def is_valid_quaternion(quat: np.ndarray, tol: float = 1e-6) -> bool:
"""Check if a quaternion is valid (unit norm).
Args:
quat: Quaternion array of length 4 (any convention).
tol: Tolerance for norm deviation from 1.0.
Returns:
True if the quaternion is valid.
"""
if not isinstance(quat, (list, tuple, np.ndarray)) or len(quat) != 4:
return False
if not np.all(np.isfinite(quat)):
return False
magnitude = np.sqrt(np.sum(np.square(quat)))
return abs(magnitude - 1.0) <= tol
def quaternion_to_angle_axis(quat_wxyz: np.ndarray, eps: float = 1e-6) -> np.ndarray:
"""Convert a quaternion (w, x, y, z) to angle-axis representation.
Args:
quat_wxyz: Quaternion in (w, x, y, z) format.
eps: Tolerance for small angles.
Returns:
Angle-axis vector [ax*angle, ay*angle, az*angle], shape (3,).
"""
q = np.array(quat_wxyz, dtype=np.float64, copy=True)
if q[0] < 0.0:
q = -q
w = q[0]
vec_part = q[1:]
angle = 2.0 * np.arccos(np.clip(w, -1.0, 1.0))
if angle < eps:
return np.zeros(3, dtype=np.float64)
sin_half_angle = np.sin(angle / 2.0)
if sin_half_angle < eps:
return np.zeros(3, dtype=np.float64)
axis = vec_part / sin_half_angle
return axis * angle
def quat_diff_as_angle_axis(
q1_wxyz: np.ndarray, q2_wxyz: np.ndarray, eps: float = 1e-6
) -> np.ndarray:
"""Calculate the rotation from q1 to q2 as an angle-axis vector.
Computes DeltaQ such that q2 = DeltaQ * q1.
Args:
q1_wxyz: Source quaternion (w, x, y, z).
q2_wxyz: Target quaternion (w, x, y, z).
eps: Tolerance for small rotations.
Returns:
Angle-axis vector representing DeltaQ, shape (3,).
"""
# Use scipy for robust quaternion math (scipy uses xyzw convention)
r1 = R.from_quat([q1_wxyz[1], q1_wxyz[2], q1_wxyz[3], q1_wxyz[0]])
r2 = R.from_quat([q2_wxyz[1], q2_wxyz[2], q2_wxyz[3], q2_wxyz[0]])
delta = r2 * r1.inv()
rotvec = delta.as_rotvec()
if np.linalg.norm(rotvec) < eps:
return np.zeros(3, dtype=np.float64)
return rotvec
def quat_diff_as_rotvec_xyzw(
q1_xyzw: np.ndarray, q2_xyzw: np.ndarray, eps: float = 1e-6
) -> np.ndarray:
"""Calculate the rotation from q1 to q2 as a rotation vector.
Same as quat_diff_as_angle_axis but takes (x, y, z, w) quaternions
directly (scipy convention), avoiding extra conversions.
Args:
q1_xyzw: Source quaternion (x, y, z, w).
q2_xyzw: Target quaternion (x, y, z, w).
eps: Tolerance for small rotations.
Returns:
Rotation vector (angle * axis), shape (3,).
"""
r1 = R.from_quat(q1_xyzw)
r2 = R.from_quat(q2_xyzw)
delta = r2 * r1.inv()
rotvec = delta.as_rotvec()
if np.linalg.norm(rotvec) < eps:
return np.zeros(3, dtype=np.float64)
return rotvec
def apply_delta_pose(
source_pos: np.ndarray,
source_quat_xyzw: np.ndarray,
delta_pos: np.ndarray,
delta_rot: np.ndarray,
eps: float = 1.0e-6,
) -> tuple:
"""Apply a delta pose to a source pose.
Args:
source_pos: Position of source frame, shape (3,).
source_quat_xyzw: Quaternion of source frame (x, y, z, w), shape (4,).
delta_pos: Position displacement, shape (3,).
delta_rot: Rotation displacement in angle-axis format, shape (3,).
eps: Tolerance to consider rotation as zero.
Returns:
Tuple of (target_pos, target_quat_xyzw), both np.ndarray.
"""
target_pos = source_pos + delta_pos
angle = np.linalg.norm(delta_rot)
if angle > eps:
delta_r = R.from_rotvec(delta_rot)
else:
delta_r = R.identity()
source_r = R.from_quat(source_quat_xyzw)
target_r = delta_r * source_r
target_quat_xyzw = target_r.as_quat()
return target_pos, target_quat_xyzw
# =====================================================================
# XR Pose Transform
# =====================================================================
def transform_xr_pose(
pos_xr: np.ndarray, quat_xyzw_xr: np.ndarray
) -> tuple:
"""Transform an XR device pose from OpenXR frame to Isaac Sim world frame.
Uses the R_HEADSET_TO_WORLD rotation matrix to transform both position
and orientation, following the same approach as XRoboToolkit.
Args:
pos_xr: Position in OpenXR frame, shape (3,).
quat_xyzw_xr: Quaternion in (x, y, z, w) format from XR SDK, shape (4,).
Returns:
Tuple of (pos_world, quat_xyzw_world):
- pos_world: Position in world frame, shape (3,).
- quat_xyzw_world: Quaternion in (x, y, z, w) for scipy, shape (4,).
"""
# Transform position
pos_world = R_HEADSET_TO_WORLD @ pos_xr
# Transform orientation: R_world = R_transform * R_xr * R_transform^T
r_xr = R.from_quat(quat_xyzw_xr) # scipy uses (x, y, z, w)
r_transform = R.from_matrix(R_HEADSET_TO_WORLD)
r_world = r_transform * r_xr * r_transform.inv()
quat_xyzw_world = r_world.as_quat()
return pos_world, quat_xyzw_world