refactor: replace rotation vector with body up vector for tilt input

- Replace body_attitude() with body_up_vector(): rotate world-up [0,0,1]
  by corrected world→body quaternion to get body up vector (pitch/roll only,
  no yaw). Matches DiffPhysDrone's env.R[:, 2] approach.
- Update ComputeTilt transform to use body_up_vector_np
- Update visualize_dataset.py to display Euler angles and body up vector
- Update model.py comments and disable CNN (zero output)
- Sync AGENTS.md with new architecture description

Generated by Mistral Vibe (ds-v4-flash).
Co-Authored-By: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
2026-06-04 21:02:08 +08:00
parent 8e1a98207e
commit 0a504d648e
6 changed files with 218 additions and 36 deletions
+8 -5
View File
@@ -1,9 +1,9 @@
"""
VelocityPredictionModel: CNN + PoseMLP → concat → GRU → Head → [vx_body, vy_body].
VelocityPredictionModel: CNN + PoseMLP → concat → GRU → Head → [v_right, v_forward].
Architecture:
Event frame (1, H, W) ──► CNN ──┐
Tilt angles (3,) ──► MLP ──┤──► concat ──► GRU ──► Head ──► [vx, vy]
Tilt angles (3,) ──► MLP ──┤──► concat ──► GRU ──► Head ──► [v_right, v_forward]
"""
import torch
@@ -92,7 +92,7 @@ class VelocityPredictionModel(nn.Module):
events: (B, S, 1, H, W)
tilt: (B, S, 3)
Output:
v_body: (B, 2) — body-frame [vx, vy] for the last frame in the sequence
v_body: (B, 2) — body-frame [v_forward, v_lateral] for the last frame in the sequence
"""
def __init__(self, cnn_cfg=model_cfg.cnn, pose_cfg=model_cfg.pose_mlp,
@@ -129,10 +129,13 @@ class VelocityPredictionModel(nn.Module):
tilt: (B, S, 3)
Returns:
v_body: (B, 2) predicted body-frame [vx, vy] at the last timestep
v_body: (B, 2) predicted body-frame [v_forward, v_lateral] at the last timestep
"""
# Per-frame encoding
cnn_feat = self.cnn(events) # (B, S, 256)
# cnn_feat = self.cnn(events) # (B, S, 256)
B, S = events.shape[:2]
cnn_feat = events.new_zeros(B, S, self.cnn.out_dim) # 全零替代
pose_feat = self.pose_mlp(tilt) # (B, S, 64)
# Fuse per frame