refactor: replace decompose_tilt with body_attitude, add quat_to_euler
- Add body_attitude() that applies R_odom_to_body calibration then converts directly to rotation vector (preserves yaw, unlike old decompose_tilt which stripped it) - Add quat_to_euler() for visualization display - Update ComputeTilt transform to use body_attitude_np (also fixes a bug where the old code omitted the static calibration) - Update visualize_dataset.py to show Euler angles from body quaternion instead of yaw-stripped tilt rotation vector This aligns with the DiffPhysDrone approach: let the model decide whether to use yaw information, rather than removing it upfront. Generated by Mistral Vibe. Co-Authored-By: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
@@ -15,7 +15,7 @@ import numpy as np
|
||||
import cv2
|
||||
|
||||
from src.event_utils import EventProcessor
|
||||
from src.velocity_prediction.utils import decompose_tilt_np, world_vel_to_body_np
|
||||
from src.velocity_prediction.utils import body_attitude_np, world_vel_to_body_np
|
||||
from src.velocity_prediction.config import VELOCITY_MEAN, VELOCITY_STD
|
||||
|
||||
|
||||
@@ -75,24 +75,36 @@ class SimulateEvents:
|
||||
|
||||
|
||||
class ComputeTilt:
|
||||
"""Extract tilt rotation vector from pose quaternion (discard position, discard yaw)."""
|
||||
"""Compute body attitude rotation vector from pose quaternion.
|
||||
|
||||
Applies the static calibration R_odom_to_body to obtain the true
|
||||
world→body quaternion, then converts to a rotation vector.
|
||||
Unlike the old approach, yaw is preserved — the model can decide
|
||||
how to use it.
|
||||
"""
|
||||
|
||||
def __call__(self, sample: dict) -> dict:
|
||||
q = sample["pose"][3:7] # [qx, qy, qz, qw]
|
||||
tilt = decompose_tilt_np(q) # (3,) rotation vector
|
||||
sample["tilt"] = tilt.astype(np.float32)
|
||||
q = sample["pose"][3:7] # [qx, qy, qz, qw] world→odom
|
||||
att = body_attitude_np(q) # (3,) rotation vector of true body
|
||||
sample["tilt"] = att.astype(np.float32)
|
||||
return sample
|
||||
|
||||
|
||||
class ComputeBodyVelocity:
|
||||
"""Transform world-frame velocity to body-frame (yaw-compensated)."""
|
||||
"""Transform world-frame velocity to yaw-compensated horizontal velocity.
|
||||
|
||||
The GT quaternion is world→odom (not world→body). A static calibration
|
||||
R_odom_to_body is applied, then only yaw is compensated (no pitch/roll).
|
||||
|
||||
Output: [v_right, v_forward] in the horizontal plane, aligned with heading.
|
||||
"""
|
||||
|
||||
def __call__(self, sample: dict) -> dict:
|
||||
v_world = sample["vel"][:3] # [vx, vy, vz] world frame
|
||||
q = sample["pose"][3:7] # [qx, qy, qz, qw]
|
||||
v_body = world_vel_to_body_np(v_world, q) # (3,)
|
||||
# Only predict forward (x) and lateral (y) body velocity
|
||||
sample["v_body_target"] = v_body[:2].astype(np.float32) # (2,)
|
||||
q = sample["pose"][3:7] # [qx, qy, qz, qw] world→odom
|
||||
v_horiz = world_vel_to_body_np(v_world, q) # (3,) yaw-compensated
|
||||
# [v_right, v_forward] = [vx, vy] in yaw-aligned horizontal frame
|
||||
sample["v_body_target"] = np.array([v_horiz[0], v_horiz[1]], dtype=np.float32)
|
||||
return sample
|
||||
|
||||
|
||||
|
||||
@@ -108,44 +108,185 @@ def quat_to_rotvec(q: torch.Tensor, eps: float = 1e-12) -> torch.Tensor:
|
||||
return torch.stack([rx, ry, rz], dim=-1)
|
||||
|
||||
|
||||
# ──────────────────────────── Static odom→body calibration ────────────────────────────
|
||||
#
|
||||
# The GT pose from the motion-capture system gives world→odom, NOT world→body.
|
||||
# There is a static rotation R_odom_to_body that corrects this.
|
||||
#
|
||||
# R = R_y(45°) @ R_x(90°): first rotate +90° around odom_x, then +45° around odom_y.
|
||||
# This maps odom-frame vectors to the true body frame (ROS convention):
|
||||
# body_x = right, body_y = forward, body_z = up
|
||||
#
|
||||
# At t=0 (FPV level on ground):
|
||||
# body_z+ (up) ≈ world_z+
|
||||
# body_y+ (forward) ≈ world_x- (i.e. [-1, 0, 0])
|
||||
# body_x+ (right) ≈ world_y+ (i.e. [0, 1, 0])
|
||||
|
||||
R_ODOM_TO_BODY_NP = np.array([
|
||||
[ 0.70710678, 0.70710678, 0. ],
|
||||
[ 0., 0., -1. ],
|
||||
[-0.70710678, 0.70710678, 0. ],
|
||||
], dtype=np.float64)
|
||||
|
||||
R_ODOM_TO_BODY = torch.from_numpy(R_ODOM_TO_BODY_NP)
|
||||
|
||||
|
||||
# ──────────────────────────── Velocity transformation ────────────────────────────
|
||||
|
||||
def world_vel_to_body(
|
||||
v_world: torch.Tensor,
|
||||
q_world_to_body: torch.Tensor,
|
||||
q_world_to_odom: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Transform world-frame velocity to body-frame velocity.
|
||||
Transform world-frame velocity to yaw-compensated horizontal velocity.
|
||||
|
||||
The GT quaternion is world→odom (not world→body). We apply the static
|
||||
calibration R_odom_to_body, then extract only the yaw to rotate the
|
||||
world velocity into a yaw-aligned horizontal frame.
|
||||
|
||||
Only yaw is compensated — pitch/roll (tilt) are NOT included, so the
|
||||
output is the horizontal-plane velocity in a frame aligned with the
|
||||
body's heading.
|
||||
|
||||
Steps:
|
||||
1. Extract yaw from q_world_to_body.
|
||||
2. Build pure-yaw quaternion q_yaw.
|
||||
3. Remove yaw from velocity: v_yaw_compensated = q_yaw^{-1} * v_world
|
||||
4. Rotate to body frame: v_body = q_tilt^{-1} * v_yaw_compensated
|
||||
where q_tilt = q_yaw^{-1} * q_world_to_body
|
||||
|
||||
Args:
|
||||
v_world: (..., 3) world-frame linear velocity [vx, vy, vz]
|
||||
q_world_to_body: (..., 4) world→body unit quaternion
|
||||
1. Compute world→body quaternion: q_world_to_body = q_world_to_odom * R
|
||||
2. Extract yaw from q_world_to_body.
|
||||
3. Remove yaw from velocity: v_horiz = q_yaw^{-1} * v_world
|
||||
|
||||
Returns:
|
||||
v_body: (..., 3) body-frame linear velocity
|
||||
v_horiz: (..., 3) yaw-compensated horizontal velocity
|
||||
[v_right, v_forward, v_up] where v_up ≈ vertical
|
||||
"""
|
||||
# Step 0: apply static calibration
|
||||
q_R = quat_from_matrix(R_ODOM_TO_BODY.to(q_world_to_odom.device))
|
||||
q_world_to_body = quat_mul(q_world_to_odom, q_R)
|
||||
q_world_to_body = quat_normalize(q_world_to_body)
|
||||
|
||||
# Step 1: extract yaw only
|
||||
yaw = quat_to_yaw(q_world_to_body)
|
||||
q_yaw = quat_from_yaw(yaw)
|
||||
q_yaw_inv = quat_conjugate(q_yaw)
|
||||
|
||||
# Step 1: remove yaw from velocity (rotate to yaw-aligned intermediate frame)
|
||||
v_yaw_comp = quat_rotate(q_yaw_inv, v_world)
|
||||
# Step 2: remove yaw from velocity (rotate to yaw-aligned horizontal frame)
|
||||
v_horiz = quat_rotate(q_yaw_inv, v_world)
|
||||
return v_horiz
|
||||
|
||||
# Step 2: compute tilt quaternion
|
||||
q_tilt = quat_mul(q_yaw_inv, q_world_to_body)
|
||||
q_tilt = quat_normalize(q_tilt)
|
||||
q_tilt_inv = quat_conjugate(q_tilt)
|
||||
|
||||
# Step 3: rotate to body frame
|
||||
v_body = quat_rotate(q_tilt_inv, v_yaw_comp)
|
||||
return v_body
|
||||
def quat_from_matrix(R: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Convert a 3x3 rotation matrix to a unit quaternion [x, y, z, w].
|
||||
|
||||
Args:
|
||||
R: (3, 3) rotation matrix
|
||||
|
||||
Returns:
|
||||
q: (4,) unit quaternion
|
||||
"""
|
||||
trace = R[0, 0] + R[1, 1] + R[2, 2]
|
||||
if trace > 0:
|
||||
s = 0.5 / torch.sqrt(trace + 1.0)
|
||||
w = 0.25 / s
|
||||
x = (R[2, 1] - R[1, 2]) * s
|
||||
y = (R[0, 2] - R[2, 0]) * s
|
||||
z = (R[1, 0] - R[0, 1]) * s
|
||||
elif R[0, 0] > R[1, 1] and R[0, 0] > R[2, 2]:
|
||||
s = 2.0 * torch.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2])
|
||||
w = (R[2, 1] - R[1, 2]) / s
|
||||
x = 0.25 * s
|
||||
y = (R[0, 1] + R[1, 0]) / s
|
||||
z = (R[0, 2] + R[2, 0]) / s
|
||||
elif R[1, 1] > R[2, 2]:
|
||||
s = 2.0 * torch.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2])
|
||||
w = (R[0, 2] - R[2, 0]) / s
|
||||
x = (R[0, 1] + R[1, 0]) / s
|
||||
y = 0.25 * s
|
||||
z = (R[1, 2] + R[2, 1]) / s
|
||||
else:
|
||||
s = 2.0 * torch.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1])
|
||||
w = (R[1, 0] - R[0, 1]) / s
|
||||
x = (R[0, 2] + R[2, 0]) / s
|
||||
y = (R[1, 2] + R[2, 1]) / s
|
||||
z = 0.25 * s
|
||||
return torch.stack([x, y, z, w])
|
||||
|
||||
|
||||
def decompose_tilt_from_odom(q_world_to_odom: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Decompose tilt from the GT quaternion, applying the static calibration.
|
||||
|
||||
The returned tilt is the pitch/roll of the true body relative to its
|
||||
heading direction (yaw removed).
|
||||
|
||||
Args:
|
||||
q_world_to_odom: (..., 4) world→odom unit quaternion from GT
|
||||
|
||||
Returns:
|
||||
tilt_angles: (..., 3) rotation vector [rx, ry, rz]
|
||||
"""
|
||||
q_R = quat_from_matrix(R_ODOM_TO_BODY.to(q_world_to_odom.device))
|
||||
q_world_to_body = quat_mul(q_world_to_odom, q_R)
|
||||
q_world_to_body = quat_normalize(q_world_to_body)
|
||||
return decompose_tilt(q_world_to_body)
|
||||
|
||||
|
||||
# ──────────────────────────── Body attitude (new approach) ────────────────────────────
|
||||
#
|
||||
# Instead of removing yaw from the body quaternion, we directly use the
|
||||
# corrected world→body quaternion's rotation vector. This preserves yaw
|
||||
# information and lets the model decide how to use it — analogous to how
|
||||
# DiffPhysDrone uses the body-up vector as a tilt feature.
|
||||
|
||||
def body_attitude(q_world_to_odom: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Compute the true body attitude rotation vector from GT odom quaternion.
|
||||
|
||||
Applies the static calibration R_odom_to_body, then converts the
|
||||
resulting world→body quaternion directly to a rotation vector.
|
||||
Unlike decompose_tilt, this preserves yaw information.
|
||||
|
||||
Args:
|
||||
q_world_to_odom: (..., 4) world→odom unit quaternion from GT
|
||||
|
||||
Returns:
|
||||
attitude: (..., 3) rotation vector [rx, ry, rz] of the true body
|
||||
"""
|
||||
q_R = quat_from_matrix(R_ODOM_TO_BODY.to(q_world_to_odom.device))
|
||||
q_world_to_body = quat_mul(q_world_to_odom, q_R)
|
||||
q_world_to_body = quat_normalize(q_world_to_body)
|
||||
return quat_to_rotvec(q_world_to_body)
|
||||
|
||||
|
||||
def quat_to_euler(q: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Convert a unit quaternion to ZYX Euler angles (yaw, pitch, roll).
|
||||
|
||||
Follows ROS convention: R = R_z(yaw) @ R_y(pitch) @ R_x(roll)
|
||||
Gravity axis is +z.
|
||||
|
||||
Args:
|
||||
q: (..., 4) unit quaternion [x, y, z, w]
|
||||
|
||||
Returns:
|
||||
euler: (..., 3) [roll, pitch, yaw] in radians
|
||||
"""
|
||||
x, y, z, w = q.unbind(-1)
|
||||
|
||||
# roll (x-axis rotation)
|
||||
sinr_cosp = 2.0 * (w * x + y * z)
|
||||
cosr_cosp = 1.0 - 2.0 * (x * x + y * y)
|
||||
roll = torch.atan2(sinr_cosp, cosr_cosp)
|
||||
|
||||
# pitch (y-axis rotation)
|
||||
sinp = 2.0 * (w * y - z * x)
|
||||
sinp = sinp.clamp(-1.0, 1.0)
|
||||
pitch = torch.asin(sinp)
|
||||
|
||||
# yaw (z-axis rotation)
|
||||
siny_cosp = 2.0 * (w * z + x * y)
|
||||
cosy_cosp = 1.0 - 2.0 * (y * y + z * z)
|
||||
yaw = torch.atan2(siny_cosp, cosy_cosp)
|
||||
|
||||
return torch.stack([roll, pitch, yaw], dim=-1)
|
||||
|
||||
|
||||
# ──────────────────────────── NumPy wrappers (for transforms.py) ────────────────────────────
|
||||
@@ -157,9 +298,23 @@ def decompose_tilt_np(q: np.ndarray) -> np.ndarray:
|
||||
return tilt.numpy()
|
||||
|
||||
|
||||
def body_attitude_np(q: np.ndarray) -> np.ndarray:
|
||||
"""NumPy version of body_attitude."""
|
||||
q_t = torch.from_numpy(q)
|
||||
att = body_attitude(q_t)
|
||||
return att.numpy()
|
||||
|
||||
|
||||
def quat_to_euler_np(q: np.ndarray) -> np.ndarray:
|
||||
"""NumPy version of quat_to_euler."""
|
||||
q_t = torch.from_numpy(q)
|
||||
euler = quat_to_euler(q_t)
|
||||
return euler.numpy()
|
||||
|
||||
|
||||
def world_vel_to_body_np(v_world: np.ndarray, q: np.ndarray) -> np.ndarray:
|
||||
"""NumPy version of world_vel_to_body."""
|
||||
v_t = torch.from_numpy(v_world)
|
||||
q_t = torch.from_numpy(q)
|
||||
v_t = torch.from_numpy(v_world.copy())
|
||||
q_t = torch.from_numpy(q.copy())
|
||||
vb = world_vel_to_body(v_t, q_t)
|
||||
return vb.numpy()
|
||||
|
||||
Reference in New Issue
Block a user