| """Rotation and gravity transform utilities for data collection.""" |
|
|
| import numpy as np |
| from scipy.spatial.transform import Rotation as R |
|
|
|
|
| def quat_to_rot6d(q): |
| """Convert scalar-first quaternion(s) (wxyz) to 6D rotation representation. |
| |
| The 6D representation consists of the first two columns of the rotation |
| matrix, flattened (Zhou et al., CVPR 2019). |
| |
| Accepted input shapes: |
| * ``(4,)`` -- single quaternion -> returns ``(6,)`` |
| * ``(N, 4)`` -- batch of quats -> returns ``(N, 6)`` |
| * ``(N*4,)`` -- flat concatenated -> returns ``(N*6,)`` |
| """ |
| q = np.asarray(q) |
| if q.ndim == 1 and q.shape[0] > 4: |
| assert q.shape[0] % 4 == 0, f"Flat quat length {q.shape[0]} is not divisible by 4" |
| q = q.reshape(-1, 4) |
| rot_6d = quat_to_rot6d(q) |
| return rot_6d.ravel() |
|
|
| single = q.ndim == 1 |
| q = np.atleast_2d(q) |
| q_xyzw = q[:, [1, 2, 3, 0]] |
| rot_mat = R.from_quat(q_xyzw).as_matrix() |
| rot_6d = rot_mat[:, :, :2].transpose(0, 2, 1).reshape(-1, 6) |
| if single: |
| return rot_6d[0].astype(q.dtype) |
| return rot_6d.astype(q.dtype) |
|
|
|
|
| def rot6d_to_quat(r): |
| """Convert 6D rotation representation to scalar-first quaternion(s) (wxyz). |
| |
| Accepted input shapes: |
| * ``(6,)`` -- single rot6d -> returns ``(4,)`` |
| * ``(N, 6)`` -- batch -> returns ``(N, 4)`` |
| * ``(N*6,)`` -- flat concat -> returns ``(N*4,)`` |
| (length must be divisible by 6) |
| |
| Args: |
| r: 6D rotation array (first two columns of rotation matrix, row-major). |
| |
| Returns: |
| Quaternion array in wxyz order. |
| """ |
| r = np.asarray(r, dtype=np.float64) |
| if r.ndim == 1 and r.shape[0] > 6: |
| assert r.shape[0] % 6 == 0, f"Flat rot6d length {r.shape[0]} is not divisible by 6" |
| r = r.reshape(-1, 6) |
| quats = rot6d_to_quat(r) |
| return quats.ravel() |
|
|
| single = r.ndim == 1 |
| r = np.atleast_2d(r) |
| col0 = r[:, :3] |
| col1 = r[:, 3:] |
| col0 = col0 / (np.linalg.norm(col0, axis=1, keepdims=True) + 1e-8) |
| dot = np.sum(col0 * col1, axis=1, keepdims=True) |
| col1 = col1 - dot * col0 |
| col1 = col1 / (np.linalg.norm(col1, axis=1, keepdims=True) + 1e-8) |
| col2 = np.cross(col0, col1) |
| rot_mat = np.stack([col0, col1, col2], axis=-1) |
| q_xyzw = R.from_matrix(rot_mat).as_quat() |
| q_wxyz = q_xyzw[:, [3, 0, 1, 2]] |
| if single: |
| return q_wxyz[0].astype(np.float32) |
| return q_wxyz.astype(np.float32) |
|
|
|
|
| def compute_projected_gravity(base_quat: np.ndarray) -> np.ndarray: |
| """Compute projected gravity vector in robot's body frame from base quaternion. |
| |
| Projects the world gravity vector [0, 0, -1] into the robot's body frame by |
| rotating it by the inverse of the base quaternion. |
| |
| Args: |
| base_quat: Base quaternion [qw, qx, qy, qz] of shape (4,) |
| |
| Returns: |
| Projected gravity vector [gx, gy, gz] of shape (3,) in robot's body frame |
| """ |
| base_quat = np.asarray(base_quat, dtype=np.float64) |
| if base_quat.shape != (4,): |
| raise ValueError(f"base_quat must have shape (4,), got {base_quat.shape}") |
|
|
| gravity_vec_world = np.array([0.0, 0.0, -1.0]) |
| base_rotation = R.from_quat(base_quat, scalar_first=True) |
| projected_gravity = base_rotation.inv().apply(gravity_vec_world) |
|
|
| return projected_gravity.astype(np.float32) |
|
|
