feat(pybot×actpred): R2 姿态反馈方案一——MPU6050 独立 IMU token 全链路

- pybot:recorder 逐帧落盘 imu.npy(T×8,roll,pitch,ax,ay,az,gx,gy,gz,与帧/动作同调用对齐);run_loop 对声明 imu_dim 的 predictor 以 predict(..., imu=) 调用(鸭子类型,旧两参调用向后兼容)

- actpred:新增 --imu-dim;model 序列 [img×F, act×H, (IMU), PRED](imu_dim>0 时 F+H+2,IMU token 位于 PRED 前);dataset/train 读 imu.npy(缺失按零向量兜底);predict/ONNX/Hailo 三后端与 export/gen_calibration/dfc_compile/_npu_test 部署链适配 imu 输入

- 姿态按物理量程固定缩放归一化(imu_norm.py,训练/推理共用,零状态)

- 测试:actpred 128 passed / pybot loop+recorder 28 passed
This commit is contained in:
qing committed 2026-08-14 09:07:24 +08:00
1 parent 3e773debc9
commit a2ba797410
22 files changed
+1236 -133

No files matched your search

+18 -4
View File
@@ -89,7 +89,8 @@ class PredictionHead(nn.Module):
class ActionPredictor(nn.Module):
"""多帧图像 + 历史动作的 Transformer 注意力融合预测器
token 序列:[img×F, act×H, PRED],长度固定为 F+H+1,
token 序列:[img×F, act×H, (IMU), PRED],长度固定为 F+H+1(imu_dim=0)
或 F+H+2(imu_dim>0,R2 方案一:IMU 独立 token 位于 PRED 之前),
取 PRED 位置输出进预测头回归下一时间步动作。
除 batch 维外全部维度静态,保证 ONNX/NPU 部署友好。
"""
@@ -107,7 +108,12 @@ class ActionPredictor(nn.Module):
self.img_proj = nn.Linear(model_cfg.cnn_feature_dim, d)
self.act_proj = nn.Linear(data_cfg.action_dim, d)
seq_len = data_cfg.frame_len + data_cfg.history_len + 1
self._imu_dim = data_cfg.imu_dim
if self._imu_dim > 0:
self.imu_proj = nn.Linear(self._imu_dim, d)
seq_len = data_cfg.frame_len + data_cfg.history_len + 1 \
+ (1 if self._imu_dim > 0 else 0)
self.pred_token = nn.Parameter(torch.zeros(1, 1, d))
self.pos_embed = nn.Parameter(torch.zeros(1, seq_len, d))
nn.init.trunc_normal_(self.pred_token, std=0.02)
@@ -131,8 +137,10 @@ class ActionPredictor(nn.Module):
self._image_height = data_cfg.image_height
self._cnn_feature_dim = model_cfg.cnn_feature_dim
def forward(self, images: torch.Tensor, action_history: torch.Tensor) -> torch.Tensor:
def forward(self, images: torch.Tensor, action_history: torch.Tensor,
imu: torch.Tensor = None) -> torch.Tensor:
# images: (B, F, 3, H, W) action_history: (B, H_len, D)
# imu: (B, imu_dim) 已归一化;None 时填零(无姿态先验)
b = images.size(0)
frames = images.reshape(
b * self._frame_len, 3, self._image_height, self._image_width)
@@ -141,7 +149,13 @@ class ActionPredictor(nn.Module):
img_tokens = self.img_proj(img_feat)
act_tokens = self.act_proj(action_history)
pred_token = self.pred_token.expand(b, -1, -1)
tokens = torch.cat([img_tokens, act_tokens, pred_token], dim=1)
parts = [img_tokens, act_tokens]
if self._imu_dim > 0:
if imu is None:
imu = images.new_zeros(b, self._imu_dim)
parts.append(self.imu_proj(imu).unsqueeze(1))
parts.append(pred_token)
tokens = torch.cat(parts, dim=1)
tokens = tokens + self.pos_embed
out = self.encoder(tokens)
return self.head(out[:, -1])