| """Small numpy quaternion/rotation helpers (quaternions are wxyz, as in Isaac Lab).""" | |
| from __future__ import annotations | |
| import math | |
| import numpy as np | |
| def quat_normalize(q: np.ndarray) -> np.ndarray: | |
| n = np.linalg.norm(q) | |
| return q / n if n > 1e-8 else np.array([1.0, 0.0, 0.0, 0.0], dtype=q.dtype) | |
| def quat_mul(a: np.ndarray, b: np.ndarray) -> np.ndarray: | |
| w1, x1, y1, z1 = a | |
| w2, x2, y2, z2 = b | |
| return np.array( | |
| [ | |
| w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, | |
| w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, | |
| w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2, | |
| w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2, | |
| ], | |
| dtype=np.float64, | |
| ) | |
| def quat_conj(q: np.ndarray) -> np.ndarray: | |
| return np.array([q[0], -q[1], -q[2], -q[3]], dtype=np.float64) | |
| def quat_to_matrix(q: np.ndarray) -> np.ndarray: | |
| w, x, y, z = quat_normalize(np.asarray(q, dtype=np.float64)) | |
| return np.array( | |
| [ | |
| [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)], | |
| [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)], | |
| [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)], | |
| ], | |
| dtype=np.float64, | |
| ) | |
| def quat_apply(q: np.ndarray, v: np.ndarray) -> np.ndarray: | |
| return quat_to_matrix(q) @ np.asarray(v, dtype=np.float64) | |
| def quat_apply_inverse(q: np.ndarray, v: np.ndarray) -> np.ndarray: | |
| return quat_to_matrix(q).T @ np.asarray(v, dtype=np.float64) | |
| def quat_from_euler_xyz(roll: float, pitch: float, yaw: float) -> np.ndarray: | |
| """Isaac Lab ``quat_from_euler_xyz`` convention (intrinsic x-y-z), wxyz.""" | |
| cr, sr = math.cos(0.5 * roll), math.sin(0.5 * roll) | |
| cp, sp = math.cos(0.5 * pitch), math.sin(0.5 * pitch) | |
| cy, sy = math.cos(0.5 * yaw), math.sin(0.5 * yaw) | |
| return np.array( | |
| [ | |
| cr * cp * cy + sr * sp * sy, | |
| sr * cp * cy - cr * sp * sy, | |
| cr * sp * cy + sr * cp * sy, | |
| cr * cp * sy - sr * sp * cy, | |
| ], | |
| dtype=np.float64, | |
| ) | |
| def yaw_of(q: np.ndarray) -> float: | |
| w, x, y, z = q | |
| return math.atan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z)) | |
| def rot6d_interleaved(q: np.ndarray) -> np.ndarray: | |
| """``matrix_from_quat(q)[..., :2].reshape(6)`` as used by the ``motion_root_ori`` observation. | |
| Row-major reshape of the first two columns: ``[m00, m01, m10, m11, m20, m21]``. | |
| """ | |
| return quat_to_matrix(q)[:, :2].reshape(6) | |
| def quat_slerp_sign_aligned(a: np.ndarray, b: np.ndarray, alpha: float) -> np.ndarray: | |
| """Linear blend with sign alignment then renormalisation (training blend convention).""" | |
| b = b if float(np.dot(a, b)) >= 0.0 else -b | |
| return quat_normalize(a + alpha * (b - a)) | |