File size: 2,770 Bytes
d8290eb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | """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))
|