diff --git a/actpred/README.md b/actpred/README.md index 0ec3e3e..5d3bd6f 100644 --- a/actpred/README.md +++ b/actpred/README.md @@ -10,15 +10,16 @@ Transformer（pre-LN）自注意力融合，取 [PRED] token 回归下一步动 ``` F 帧图像 (F×3×H×W) ─► 共享 CNN 编码器 ─► F 个 token ─► proj → d_model ┐ - ├─► [img×F, act×H, PRED] + 位置编码 -动作历史 (H×D) ─► Linear proj → d_model ─► H 个 token ────────┤ │ -可学习 [PRED] token ────────────────────────────────────────┘ ▼ - TransformerEncoder (pre-LN, 2 层, 4 头, d_model=128) + ├─► [img×F, act×H, (IMU), PRED] +动作历史 (H×D) ─► Linear proj → d_model ─► H 个 token ────────┤ │ + 位置编码 +IMU 姿态 (imu_dim) ─► Linear proj → 1 个 token（R2，可选）────────┤ ▼ +可学习 [PRED] token ────────────────────────────────────────┘ TransformerEncoder (pre-LN, 2 层, 4 头, d_model=128) │ 取 PRED 位置输出 ─► MLP 头 ─► 预测动作 (action_dim) ``` -token 序列长度固定为 F+H+1（默认 3+5+1=9），除 batch 外无动态 shape，对 NPU（如树莓派 +token 序列长度固定为 F+H+1（默认 3+5+1=9；`--imu-dim`>0 时 IMU token 插到 +PRED 之前，长度 F+H+2），除 batch 外无动态 shape，对 NPU（如树莓派 AI 扩展板）部署友好。 ## 快速开始 @@ -40,6 +41,9 @@ bash _run.sh train --data-dir ./data --epochs 100 --resume checkpoints/last.pt bash _run.sh train --data-dir ./data --epochs 100 --action-dim 18 \ --predict-delta --rollout-len 4 --ss-prob 0.7 +# R2 姿态反馈：--imu-dim 8 启用独立 IMU token（数据需含 imu.npy，见"数据格式"） +bash _run.sh train --data-dir ./data --epochs 100 --action-dim 18 --imu-dim 8 + # 推理（--images 按时间顺序旧→新传入 frame_len 个图像路径，冷启动时可重复传同一帧； # --action-history 格式：冒号分隔各步，每步 action_dim 个值，共 history_len 步） bash _run.sh --model checkpoints/best_inference.pt \ @@ -109,6 +113,7 @@ data/ 0001.jpg ... actions.npy # shape: (T, action_dim) float32 + imu.npy # 可选，shape: (T, imu_dim) float32（MPU6050 姿态，R2） episode_002/ ... ``` @@ -116,6 +121,13 @@ data/ 每个 timestep 的样本：帧 `[t-frame_len+1, t]`（旧→新）+ 前 `history_len` 步动作 → 当前步动作。 采样起点为 `max(history_len, frame_len-1)`。 +`imu.npy` 的键序为 `[roll, pitch, ax, ay, az, gx, gy, gz]`（roll/pitch 角度制， +其余为 MPU6050 原始 int16），由 pybot `--record` 在脚本含 `@MPU6050_BIND` 块时 +逐帧同步落盘。**旧 episode 没有姿态数据，训练 IMU 模型需用新版 pybot 重新录制**； +`--imu-dim`>0 时缺失 `imu.npy` 的 episode 按零向量兜底（向后兼容，但无姿态先验收益）。 +姿态归一化按物理量程固定缩放（roll/pitch ÷180，原始 int16 ÷32768），训练与推理 +共用 `imu_norm.py`，无需统计文件。 + ### 数据采集 ```python @@ -138,6 +150,7 @@ recorder.flush() | `--action-dim` | 18 | 动作向量维度（对齐 bot4cj 18 列：10 舵机 + 2 云台 + 扬声器频率 + 延时 + 4 预留）| | `--history-len` | 5 | 历史动作窗口长度 | | `--frame-len` | 3 | 输入图像帧数（含当前帧，旧帧在前）| +| `--imu-dim` | 0 | R2 姿态输入维度（>0 启用独立 IMU token；8=roll,pitch,ax,ay,az,gx,gy,gz）| | `--backbone` | simple | CNN 类型 (simple/mobilenet) | | `--attn-dim` | 128 | 注意力 token 维度 d_model（须能被 attn-heads 整除）| | `--attn-heads` | 4 | 注意力头数 | @@ -177,6 +190,8 @@ recorder.flush() 物理可行性用数据分布的经验界（min/max + 10% margin）而非硬编码范围——18 列混含 舵机角度/频率/延时/预留，无统一硬限位。`--compare` 用于 A/B 对比（如 R1 开/关）。 +滚动复现闭环口径：模型带 imu 输入时按 episode 的 `imu.npy` 真值姿态逐步喂入 +（缺失则零向量，与 pybot 闭环无 `@MPU6050_BIND` 块时的行为一致）。 ## 部署注意 @@ -195,13 +210,16 @@ print(predictor.action_dim, predictor.history_len, predictor.frame_len) # 公 # images: 长度为 frame_len 的 RGB 图像列表（旧→新，末帧为当前帧；冷启动可重复同一帧） action = predictor.predict(images, action_history) # np.ndarray (action_dim,) # action_history 形状须为 (history_len, action_dim) +# 模型含 IMU 输入时（predictor.imu_dim 非 None）可传 imu= 原始姿态向量 +# （roll,pitch,ax,ay,az,gx,gy,gz；内部自动归一化，None 等价零向量） +action = predictor.predict(images, action_history, imu=imu_vec) ``` ## ONNX Runtime 后端（CPU 加速） `onnx_infer.py` 的 `OnnxPredictor` 走 ONNX Runtime CPU 推理，接口与 -`ActionPredictorInference` 对齐，维度从 ONNX 输入元数据自动推导（无需 -checkpoint）。Pi 5 实测（224×224，2026-08-13 复测）11.4ms/次含预处理 +`ActionPredictorInference` 对齐（含可选 `imu=` 入参），维度从 ONNX 输入元数据 +自动推导（无需 checkpoint）。Pi 5 实测（224×224，2026-08-13 复测）11.4ms/次含预处理 （纯推理 10.7ms），PyTorch 21ms，约 2 倍加速，与 torch 前向输出差异 <1e-5。数据见 [测速报告](../docs/reports/actpred_inference_speed_benchmark_20260813.md) §6。 @@ -232,7 +250,8 @@ reshape）**无法被 DFC 解析**，必须用 `export_deploy.py` 的部署版 当前 NPU 目标分辨率为 **224×224**（pretrain checkpoint 配置），HEF 输入 为 frame0~2 各 (1,224,224,3) NHWC + action_history (1,1,5,18) + -pred_token (1,1,1,128)，形状与量化参数由 `hailo_infer.py` 从 vstream +pred_token (1,1,1,128)（`--imu-dim`>0 时追加 imu (1,1,1,imu_dim)，共 6 输入）， +形状与量化参数由 `hailo_infer.py` 从 vstream info 自动推导，不硬编码。 ```bash @@ -269,9 +288,10 @@ python3 dfc_compile.py # 产物 checkpoints/actpred_pretrain.hef > 遗留：全零历史角（z≈-3.5）超出校准输入范围（-2.7）会被截断， > 该角仅作连通性验证，非真实部署场景。 -`hailo_infer.py` 的 `HailoPredictor` 适配 5 输入 HEF，输入布局与量化 +`hailo_infer.py` 的 `HailoPredictor` 适配 5/6 输入 HEF（含可选 imu 输入， +按形状自动归类绑定），输入布局与量化 参数（帧归一化由量化网格 dequant(0) 自动判定）从 vstream info 自动推导， 单元测试含 fake hailo_platform 的 5 输入与乱序报告顺序用例。CLI 支持 -`--pred-token` 喂入学习值、`--stats` 喂入归一化统计。 +`--pred-token` 喂入学习值、`--stats` 喂入归一化统计、`--imu` 喂入姿态。 环境调研与 API 踩坑记录见 [docs/reports/actpred_hailo_hat2_inference_research_20260813.md](../docs/reports/actpred_hailo_hat2_inference_research_20260813.md)。 diff --git a/actpred/_npu_test.py b/actpred/_npu_test.py index 8be7add..63da6c7 100644 --- a/actpred/_npu_test.py +++ b/actpred/_npu_test.py @@ -78,30 +78,39 @@ tf = transforms.Compose([ ]) -def torch_pred(imgs, hist, use_learned_pred=True): +def torch_pred(imgs, hist, use_learned_pred=True, imu=None): x = torch.stack([tf(Image.fromarray(im)) for im in imgs]).unsqueeze(0) h = torch.tensor(hist).unsqueeze(0) if stats_path is not None: stats = np.load(stats_path) h = (h - torch.tensor(stats[0])) / torch.tensor(stats[1]) + ii = None + if imu_dim > 0: + # 与 HailoPredictor 同口径：原始姿态 → 内部归一化前需手动归一 + from imu_norm import normalize_imu + ii = torch.from_numpy(normalize_imu(imu, imu_dim)).unsqueeze(0) if use_learned_pred: with torch.no_grad(): - return model(x, h).squeeze(0).numpy() + return model(x, h, ii).squeeze(0).numpy() saved = model.pred_token.data.clone() model.pred_token.data.zero_() with torch.no_grad(): - out = model(x, h).squeeze(0).numpy() + out = model(x, h, ii).squeeze(0).numpy() model.pred_token.data.copy_(saved) return out +# R2 方案一：imu_dim>0 时每步传原始姿态（HailoPredictor 内部归一化） +imu_dim = cfg.data.imu_dim +imu_real = np.array([5.0, -3.0] + [0.0] * 6, dtype=np.float32) + with HailoPredictor(HEF_PATH, stats_path=stats_path) as p: print(f"\ndims: frame_len={p.frame_len} img={p.image_width}x{p.image_height} " f"history={p.history_len} action_dim={p.action_dim}") t0 = time.perf_counter() - a = p.predict(imgs_real, hist_real, pred_token=pred_learned) - b = p.predict(imgs_zero, hist_real, pred_token=pred_learned) - c = p.predict(imgs_real, hist_zero, pred_token=pred_learned) + a = p.predict(imgs_real, hist_real, pred_token=pred_learned, imu=imu_real) + b = p.predict(imgs_zero, hist_real, pred_token=pred_learned, imu=imu_real) + c = p.predict(imgs_real, hist_zero, pred_token=pred_learned, imu=imu_real) d = p.predict(imgs_zero, hist_zero, pred_token=pred_learned) t1 = time.perf_counter() print(f"4x predict: {(t1 - t0) / 4 * 1000:.1f} ms/次") @@ -113,9 +122,14 @@ with HailoPredictor(HEF_PATH, stats_path=stats_path) as p: hist_sens = np.abs(a - c).max() print(f"\nimg sensitivity = {img_sens:.4f}（>0.1 视为输入连通）") print(f"hist sensitivity = {hist_sens:.4f}（>0.1 视为输入连通）") + ok = img_sens > 0.1 and hist_sens > 0.1 + if imu_dim > 0: + e = p.predict(imgs_real, hist_real, pred_token=pred_learned, + imu=np.zeros_like(imu_real)) + print(f"imu sensitivity = {np.abs(a - e).max():.4f}（>0 视为姿态连通）") - ref = torch_pred(imgs_real, hist_real, True) - ref_zero_pred = torch_pred(imgs_real, hist_real, False) + ref = torch_pred(imgs_real, hist_real, True, imu_real) + ref_zero_pred = torch_pred(imgs_real, hist_real, False, imu_real) # torch 参考是模型原始 z 空间输出；HailoPredictor 带 stats 时已反归一化 # 到原始动作空间，比较前须同空间对齐，否则归一化模型下 max_diff 无意义 if stats_path is not None: @@ -127,6 +141,5 @@ with HailoPredictor(HEF_PATH, stats_path=stats_path) as p: print(f"HEF vs torch(learned) max_diff={np.abs(a - ref).max():.4f}" f"（原始动作空间）") - ok = img_sens > 0.1 and hist_sens > 0.1 print("\nPASS: HEF 输入连通，可继续与 torch 对比量化误差" if ok else "\nFAIL: HEF 输出对输入不敏感，编译仍有问题") diff --git a/actpred/_onnx_validate.py b/actpred/_onnx_validate.py index e24896f..a40bd37 100644 --- a/actpred/_onnx_validate.py +++ b/actpred/_onnx_validate.py @@ -69,6 +69,17 @@ if action_stats is not None: pred_learned = model.pred_token.data.unsqueeze(1).numpy() # (1,1,1,128) pred_none = np.zeros((1, 1, 1, 128), dtype=np.float32) +# R2 方案一：imu_dim>0 时模型/ONNX 多了归一化姿态输入 +imu_dim = cfg.data.imu_dim +imu_real = None +imu_zero = None +if imu_dim > 0: + from imu_norm import normalize_imu + # 原始姿态（roll=5°, pitch=-3°, 其余 0）→ 部署版输入布局 (1,1,1,imu_dim) + imu_real = normalize_imu([5.0, -3.0] + [0.0] * 6, imu_dim).reshape( + 1, 1, 1, imu_dim) + imu_zero = np.zeros((1, 1, 1, imu_dim), dtype=np.float32) + def frames_nhwc(imgs): """列表帧 → (1,H,W,3) 归一化 float32 NHWC × frame_len""" @@ -76,35 +87,45 @@ def frames_nhwc(imgs): return [chw[i].transpose(1, 2, 0)[None] for i in range(len(imgs))] -def run_onnx(imgs, hist, pred): +def run_onnx(imgs, hist, pred, imu=None): feeds = {f"frame{i}": f for i, f in enumerate(frames_nhwc(imgs))} feeds["action_history"] = hist[None, None] feeds["pred_token"] = pred + if imu_dim > 0: + feeds["imu"] = imu if imu is not None else imu_zero return sess.run(None, feeds)[0].reshape(-1) -def run_torch(imgs, hist): +def run_torch(imgs, hist, imu=None): x = torch.stack([tf(Image.fromarray(im)) for im in imgs]).unsqueeze(0) hh = torch.tensor(hist).unsqueeze(0) + ii = None + if imu_dim > 0: + ii = torch.tensor( + imu if imu is not None else imu_zero).reshape(1, imu_dim).float() with torch.no_grad(): - return model(x, hh).squeeze(0).numpy() + return model(x, hh, ii).squeeze(0).numpy() -a = run_onnx(imgs_real, hist_real, pred_learned) -b = run_onnx(imgs_zero, hist_real, pred_learned) -c = run_onnx(imgs_real, hist_zero, pred_learned) +a = run_onnx(imgs_real, hist_real, pred_learned, imu_real) +b = run_onnx(imgs_zero, hist_real, pred_learned, imu_real) +c = run_onnx(imgs_real, hist_zero, pred_learned, imu_real) d = run_onnx(imgs_zero, hist_zero, pred_learned) -p0 = run_onnx(imgs_real, hist_real, pred_none) +p0 = run_onnx(imgs_real, hist_real, pred_none, imu_real) print(f"onnx real_img+real_hist: {np.round(a[:6], 3)}...") print(f"onnx pred_token=0 diff : max={np.abs(a - p0).max():.4f}") print(f"img sensitivity (a-b): max={np.abs(a - b).max():.4f}") print(f"hist sensitivity (a-c): max={np.abs(a - c).max():.4f}") -ref = run_torch(imgs_real, hist_real) +ok = np.abs(a - b).max() > 1e-4 and np.abs(a - c).max() > 1e-4 +if imu_dim > 0: + e = run_onnx(imgs_real, hist_real, pred_learned, imu_zero) + print(f"imu sensitivity (a-e): max={np.abs(a - e).max():.4f}") + +ref = run_torch(imgs_real, hist_real, imu_real) print(f"torch reference : {np.round(ref[:6], 3)}...") print(f"onnx vs torch max_diff : {np.abs(a - ref).max():.6e}") -ok = (np.abs(a - b).max() > 1e-4 and np.abs(a - c).max() > 1e-4 - and np.abs(a - ref).max() < 1e-4) +ok = ok and np.abs(a - ref).max() < 1e-4 print("PASS" if ok else "FAIL") diff --git a/actpred/config.py b/actpred/config.py index 4cb7171..13d2de0 100644 --- a/actpred/config.py +++ b/actpred/config.py @@ -13,6 +13,10 @@ class DataConfig: data_dir: str = "./data" normalize_actions: bool = False action_stats_path: Optional[str] = None + # R2 方案一：MPU6050 姿态作为独立 IMU token 注入（序列 F+H+2）。 + # 0 = 关闭（默认，序列保持 F+H+1，旧 checkpoint 完全兼容）； + # 8 = roll,pitch,ax,ay,az,gx,gy,gz（pybot recorder 的 imu.npy 格式）。 + imu_dim: int = 0 @dataclass @@ -79,6 +83,9 @@ def build_argparser(description: str = "") -> argparse.ArgumentParser: parser.add_argument("--data-dir", type=str, default=None) parser.add_argument("--normalize-actions", action=argparse.BooleanOptionalAction, default=None) parser.add_argument("--action-stats-path", type=str, default=None) + parser.add_argument("--imu-dim", type=int, default=None, + help="R2 方案一：MPU6050 姿态维度（0=关闭，默认；" + "8=roll,pitch,ax,ay,az,gx,gy,gz，序列变 F+H+2）") parser.add_argument("--backbone", type=str, default=None, choices=["simple", "mobilenet"]) parser.add_argument("--cnn-feature-dim", type=int, default=None) parser.add_argument("--fused-hidden-dim", type=int, default=None) @@ -126,6 +133,7 @@ def validate_config_values(data_cfg: DataConfig, model_cfg: ModelConfig, ("history_len", data_cfg.history_len, 1, 1024, int), ("frame_len", data_cfg.frame_len, 1, 1024, int), ("action_dim", data_cfg.action_dim, 1, 1024, int), + ("imu_dim", data_cfg.imu_dim, 0, 64, int), ("cnn_feature_dim", model_cfg.cnn_feature_dim, 1, 4096, int), ("fused_hidden_dim", model_cfg.fused_hidden_dim, 1, 4096, int), ("dropout", model_cfg.dropout, 0.0, 1.0, float), @@ -190,6 +198,11 @@ def warn_init_from_mismatch(init_cfg: dict, current: Config): print(f"WARNING: init-from checkpoint 的 action_dim={init_dim} " f"与当前运行 {current.data.action_dim} 不一致，" f"act_proj 权重将随机初始化") + init_imu = init_data.get("imu_dim", 0) # 旧 checkpoint 无此键 → 0 + if init_imu != current.data.imu_dim: + print(f"WARNING: init-from checkpoint 的 imu_dim={init_imu} " + f"与当前运行 {current.data.imu_dim} 不一致，" + f"imu_proj/pos_embed 权重将随机初始化") def parse_checkpoint_config(cfg_dict: dict) -> Config: @@ -235,6 +248,7 @@ def apply_overrides(cfg: Config, args: argparse.Namespace) -> Config: "data_dir": ("data", "data_dir"), "normalize_actions": ("data", "normalize_actions"), "action_stats_path": ("data", "action_stats_path"), + "imu_dim": ("data", "imu_dim"), "backbone": ("model", "cnn_backbone"), "cnn_feature_dim": ("model", "cnn_feature_dim"), "fused_hidden_dim": ("model", "fused_hidden_dim"), diff --git a/actpred/dataset.py b/actpred/dataset.py index d6c93d5..82068ed 100644 --- a/actpred/dataset.py +++ b/actpred/dataset.py @@ -10,6 +10,8 @@ from torch.utils.data import Dataset from torchvision import transforms from torchvision.transforms import functional as TF +from imu_norm import IMU_SCALES + def _natural_key(p): """自然排序键：'1.jpg' < '2.jpg' < ... < '10.jpg'""" @@ -22,9 +24,12 @@ class ActionDataset(Dataset): normalize_actions: bool = False, action_stats_path: str = None, episode_indices: list = None, frame_len: int = 3, augment: bool = False, - predict_delta: bool = False, rollout_len: int = 1): + predict_delta: bool = False, rollout_len: int = 1, + imu_dim: int = 0): self.history_len = history_len self.action_dim = action_dim + # R2 方案一：>0 时样本带 "imu" 键（归一化姿态向量，缺失 episode 补零） + self.imu_dim = imu_dim self.image_width = image_width self.image_height = image_height self.frame_len = frame_len @@ -65,12 +70,13 @@ class ActionDataset(Dataset): print(f"WARNING: {ep_dir} 动作帧数 {n_frames} 与图像数 {n_images} 不一致，" f"按较少者（{min(n_frames, n_images)}）使用，请检查录制完整性") n = min(n_frames, n_images) + imus = self._load_imu(ep_dir, n) # 第 k 步（k=0..K-1）目标时刻 t=t_start+k：帧 [t-frame_len+1, t]、 # 历史 [t-history_len, t)、目标 actions[t]。t_start 下界保证第 0 步窗口完整， # 上界保证第 K-1 步目标不越界。 start = max(history_len, frame_len - 1) for t_start in range(start, n - (self.rollout_len - 1)): - self.samples.append((image_files, actions, t_start)) + self.samples.append((image_files, actions, imus, t_start)) self._action_mean = None self._action_std = None @@ -94,7 +100,7 @@ class ActionDataset(Dataset): # 归一化统计基于绝对动作（与 Δ 无关）：覆盖所有样本 rollout 范围内 # 出现的 history 与 target 动作，保持与旧版一致的统计口径。 all_actions = np.stack([ - s[1][s[2] - self.history_len: s[2] + self.rollout_len] + s[1][s[3] - self.history_len: s[3] + self.rollout_len] for s in self.samples ]).reshape(-1, self.action_dim) self._action_mean = all_actions.mean(axis=0) @@ -134,22 +140,57 @@ class ActionDataset(Dataset): return self._norm(actions[t]) - self._norm(actions[t - 1]) return self._norm(actions[t]) + def _load_imu(self, ep_dir, n): + """加载 episode 的 imu.npy 并按物理量程归一化（R2） + + imu_dim=0（未启用）或文件缺失 → None（样本侧补零向量，等价无姿态 + 先验，混合数据集可训）。形状 pad/truncate 到 (n, imu_dim)。 + """ + if self.imu_dim <= 0: + return None + path = ep_dir / "imu.npy" + if not path.exists(): + print(f"WARNING: {ep_dir} 无 imu.npy，该 episode 姿态按零向量处理") + return None + imus = np.load(path).astype(np.float32) + if imus.ndim != 2: + raise ValueError( + f"{path} 形状应为 (T, imu_dim)，实际为 {imus.shape}") + if imus.shape[1] != self.imu_dim: + fixed = np.zeros((imus.shape[0], self.imu_dim), dtype=np.float32) + m = min(imus.shape[1], self.imu_dim) + fixed[:, :m] = imus[:, :m] + imus = fixed + if len(imus) < n: # 与帧数对齐：不足补零 + imus = np.concatenate([ + imus, np.zeros((n - len(imus), self.imu_dim), dtype=np.float32)]) + return imus[:n] / IMU_SCALES[:self.imu_dim] + + def _imu_at(self, imus, t): + """时刻 t 的归一化姿态向量（与最新帧对齐）；无数据 → 零向量""" + if imus is None: + return torch.zeros(self.imu_dim, dtype=torch.float32) + return torch.tensor(imus[t], dtype=torch.float32) + def __getitem__(self, idx): - image_files, actions, t_start = self.samples[idx] + image_files, actions, imus, t_start = self.samples[idx] if self.rollout_len == 1: - # 单步：与旧版输出结构完全一致 + # 单步：与旧版输出结构完全一致（imu_dim>0 时附加 "imu" 键） images = self._load_frames(image_files, t_start) history = torch.tensor(self._norm( actions[t_start - self.history_len: t_start]), dtype=torch.float32) target = torch.tensor(self._step_target(actions, t_start), dtype=torch.float32) - return {"images": images, "action_history": history, "target": target} + out = {"images": images, "action_history": history, "target": target} + if self.imu_dim > 0: + out["imu"] = self._imu_at(imus, t_start) + return out # 多步：返回 K 步序列，scheduled sampling 在 train 循环内按步取用 steps = [] for k in range(self.rollout_len): t = t_start + k - steps.append({ + step = { "images": self._load_frames(image_files, t), # 真实历史（scheduled sampling 的真值来源；首步必用） "action_history": torch.tensor( @@ -157,7 +198,10 @@ class ActionDataset(Dataset): "target": torch.tensor(self._step_target(actions, t), dtype=torch.float32), # 上一步真实动作（Δ 残差的基准 a_{t-1}，归一化空间） "prev_action": torch.tensor(self._norm(actions[t - 1]), dtype=torch.float32), - }) + } + if self.imu_dim > 0: + step["imu"] = self._imu_at(imus, t) + steps.append(step) return {"steps": steps} def get_action_stats(self): diff --git a/actpred/dfc_compile.py b/actpred/dfc_compile.py index 4f54702..267b159 100644 --- a/actpred/dfc_compile.py +++ b/actpred/dfc_compile.py @@ -36,10 +36,14 @@ import tempfile NET_NAME = "actpred_deploy" -def start_nodes(frame_count: int) -> list: - """export_deploy.py 的输入命名约定：frame{i} × frame_count + action_history + pred_token""" - return ([f"frame{i}" for i in range(frame_count)] - + ["action_history", "pred_token"]) +def start_nodes(frame_count: int, with_imu: bool = False) -> list: + """export_deploy.py 的输入命名约定：frame{i} × frame_count + action_history + + pred_token（R2 imu_dim>0 时末尾追加 imu）""" + nodes = ([f"frame{i}" for i in range(frame_count)] + + ["action_history", "pred_token"]) + if with_imu: + nodes.append("imu") + return nodes def hailo(*args, venv_bin: str): @@ -118,6 +122,9 @@ def main(): parser.add_argument("--frame-count", type=int, default=3, help="帧输入个数（与导出 ONNX 的 frame_len 一致，" "默认 3）") + parser.add_argument("--with-imu", action="store_true", + help="R2：导出 ONNX 含 imu 输入时开启（start nodes " + "末尾追加 imu，与 export_deploy.py 命名一致）") parser.add_argument("--venv-bin", default="venv_dfc/bin", help="安装了 hailo-sdk 的 venv bin 目录") args = parser.parse_args() @@ -126,7 +133,7 @@ def main(): if not os.path.exists(os.path.join(venv_bin, "hailo")): sys.exit(f"错误: {venv_bin}/hailo 不存在，请先安装 hailo-sdk") - nodes = start_nodes(args.frame_count) + nodes = start_nodes(args.frame_count, args.with_imu) input_formats = ["NHWC"] * len(nodes) workdir = tempfile.mkdtemp(prefix="dfc_") diff --git a/actpred/evaluate.py b/actpred/evaluate.py index 6b12052..997f0e2 100644 --- a/actpred/evaluate.py +++ b/actpred/evaluate.py @@ -28,7 +28,11 @@ def _episode_ids(data_dir): def _load_episode(data_dir, ep_name): - """加载单个 episode 的图像路径（自然序）与动作真值 (T, D)""" + """加载单个 episode 的图像路径（自然序）、动作真值 (T, D)、姿态真值 (T, imu_dim) + + imu.npy 缺失时第三项返回 None（旧数据兼容，滚动时按零向量喂入）。 + 姿态保持原始量纲——predict 端内部归一化，与 pybot 闭环口径一致。 + """ ep_dir = [d for d in ActionDataset._valid_episode_dirs(data_dir) if d.name == ep_name][0] from dataset import _natural_key @@ -37,7 +41,14 @@ def _load_episode(data_dir, ep_name): list(images_dir.glob("*.png")), key=_natural_key) actions = np.load(ep_dir / "actions.npy").astype(np.float32) n = min(len(image_files), len(actions)) - return image_files[:n], actions[:n] + imus = None + imu_path = ep_dir / "imu.npy" + if imu_path.exists(): + imus = np.load(imu_path).astype(np.float32) + if imus.ndim != 2: + raise ValueError(f"{imu_path} 形状应为 (T, imu_dim)，实际为 {imus.shape}") + imus = imus[:n] + return image_files[:n], actions[:n], imus def _action_bounds(data_dir, action_dim, margin_ratio=0.1): @@ -58,16 +69,19 @@ def _action_bounds(data_dir, action_dim, margin_ratio=0.1): @np.errstate(all="ignore") def rollout_episode(predictor, image_files, actions, horizon=None, - lo=None, hi=None): + lo=None, hi=None, imus=None): """单 episode 开环滚动，复现闭环口径（图像真值滚动、动作自预测回写）。 从 t0 = max(history_len, frame_len-1) 起步（窗口完整处），每步： 帧窗 = 真值图像 [t-frame_len+1, t]（旧→新） 历史 = 前 H 步动作（首步用真值填充，之后用自预测回写，复刻 run_loop） + 姿态 = imus[t]（predictor.imu_dim 非 None 时传入；imus 为 None 或 + 长度不足按零向量 = 无姿态先验） 预测 â_t → 与真值 a_t 比较 → 回写历史 返回逐步记录 dict。horizon=None 表示滚到 episode 末尾。 """ F, H = predictor.frame_len, predictor.history_len + use_imu = getattr(predictor, "imu_dim", None) is not None t0 = max(H, F - 1) n = len(actions) if horizon is None: @@ -88,8 +102,12 @@ def rollout_episode(predictor, image_files, actions, horizon=None, for k in range(horizon): t = t0 + k frames = [load(t - F + 1 + i) for i in range(F)] + kwargs = {} + if use_imu: + kwargs["imu"] = (imus[t] if imus is not None and t < len(imus) + else None) pred = np.asarray( - predictor.predict(frames, np.array(history)), + predictor.predict(frames, np.array(history), **kwargs), dtype=np.float32).reshape(-1) true = actions[t] err = pred - true @@ -158,14 +176,16 @@ def evaluate(checkpoint, data_dir, episodes=None, horizon=None, device=None): lo, hi = _action_bounds(data_dir, predictor.action_dim) all_steps = [] for name in ep_names: - image_files, actions = _load_episode(data_dir, name) - steps = rollout_episode(predictor, image_files, actions, horizon, lo, hi) + image_files, actions, imus = _load_episode(data_dir, name) + steps = rollout_episode(predictor, image_files, actions, horizon, + lo, hi, imus=imus) if steps: all_steps.append(steps) rep = summarize(all_steps, lo, hi) rep["checkpoint"] = checkpoint rep["frame_len"] = predictor.frame_len rep["history_len"] = predictor.history_len + rep["imu_dim"] = getattr(predictor, "imu_dim", None) rep["predict_delta"] = getattr(predictor, "_predict_delta", False) return rep @@ -174,6 +194,7 @@ def _fmt(rep): lines = [ f"checkpoint : {rep['checkpoint']}", f"predict_delta : {rep['predict_delta']}", + f"imu_dim : {rep.get('imu_dim')}", f"episodes/steps : {rep['n_episodes']} / {rep['n_steps']}", f"overall_mse : {rep['overall_mse']:.6f}", f"error_growth : {rep['error_growth']:.3f} (>1 = 误差随步数发散)", diff --git a/actpred/export_deploy.py b/actpred/export_deploy.py index 7a7f294..7d35d97 100644 --- a/actpred/export_deploy.py +++ b/actpred/export_deploy.py @@ -114,6 +114,11 @@ class DeployWrapper(nn.Module): # 避免 DFC 对 3D Linear 的错误卷积融合 self.act_w = nn.Parameter(m.act_proj.weight.data.t().clone()) # (D, d) self.act_b = nn.Parameter(m.act_proj.bias.data.clone()) + # R2 方案一：IMU 独立 token（imu_dim>0 时追加图输入，同 MatMul+Add 手法） + self._imu_dim = m._imu_dim + if self._imu_dim > 0: + self.imu_w = nn.Parameter(m.imu_proj.weight.data.t().clone()) + self.imu_b = nn.Parameter(m.imu_proj.bias.data.clone()) # pred_token 改为图输入（常量 initializer 会被 DFC 误转置）， # 位置编码经 concat 后的 ew_add 广播可被正确推断 self.pos_embed4 = nn.Parameter(m.pos_embed.data.unsqueeze(1).clone()) # (1, 1, S, d) @@ -128,14 +133,20 @@ class DeployWrapper(nn.Module): return self.pool(x) # (1, d, 1, 1) NCHW def forward(self, frame0: torch.Tensor, frame1: torch.Tensor, frame2: torch.Tensor, - action_history: torch.Tensor, pred_token: torch.Tensor) -> torch.Tensor: + action_history: torch.Tensor, pred_token: torch.Tensor, + imu: torch.Tensor = None) -> torch.Tensor: # 全程 (1, 1, S, d) 布局：seq 占 W 维，避免任何 transpose frames = [frame0, frame1, frame2] feats = [self._frame_feat(frames[i], i).permute(0, 2, 3, 1) for i in range(self._frame_len)] # (1, 1, 1, d) NHWC img_tokens = torch.cat(feats, dim=2) + self.proj_bias # (1, 1, F, d) act_tokens = torch.matmul(action_history, self.act_w) + self.act_b # (1, 1, H, d) - tokens = torch.cat([img_tokens, act_tokens, pred_token], dim=2) # (1, 1, S, d) + parts = [img_tokens, act_tokens] + if self._imu_dim > 0: + # imu: (1, 1, 1, imu_dim) 归一化姿态 → (1, 1, 1, d)，位于 PRED 之前 + parts.append(torch.matmul(imu, self.imu_w) + self.imu_b) + parts.append(pred_token) + tokens = torch.cat(parts, dim=2) # (1, 1, S, d) tokens = tokens + self.pos_embed4 for blk in self.blocks: tokens = blk(tokens) @@ -232,18 +243,25 @@ def main(): dummy_frames = [torch.randn(1, h, w, 3) for _ in range(f)] dummy_history = torch.randn(1, 1, cfg.data.history_len, cfg.data.action_dim) dummy_pred = model.pred_token.data.unsqueeze(1).clone() # (1, 1, 1, d) + imu_dim = cfg.data.imu_dim + # imu_dim>0 时追加归一化姿态输入（均匀采样的归一化值域） + dummy_imu = torch.randn(1, 1, 1, imu_dim) if imu_dim > 0 else None + + extra_inputs = (dummy_imu,) if imu_dim > 0 else () + extra_names = ["imu"] if imu_dim > 0 else [] with torch.no_grad(): - ref = wrapper(*dummy_frames, dummy_history, dummy_pred) + ref = wrapper(*dummy_frames, dummy_history, dummy_pred, *extra_inputs) torch.onnx.export( - wrapper, (*dummy_frames, dummy_history, dummy_pred), + wrapper, (*dummy_frames, dummy_history, dummy_pred, *extra_inputs), args.output, - input_names=[f"frame{i}" for i in range(f)] + ["action_history", "pred_token"], + input_names=[f"frame{i}" for i in range(f)] + ["action_history", "pred_token"] + extra_names, output_names=["predicted_action"], opset_version=18, ) - print(f"部署版 ONNX 已导出: {args.output}（{f} 个帧输入 (1,{h},{w},3) NHWC + action_history + pred_token）") + print(f"部署版 ONNX 已导出: {args.output}（{f} 个帧输入 (1,{h},{w},3) NHWC + action_history + pred_token" + f"{' + imu' if imu_dim > 0 else ''}）") if args.pred_token_out: import numpy as np @@ -252,7 +270,8 @@ def main(): stacked = torch.stack([fr.permute(0, 3, 1, 2) for fr in dummy_frames], dim=1) with torch.no_grad(): - orig_out = model(stacked, dummy_history.squeeze(1)) + orig_out = model(stacked, dummy_history.squeeze(1), + dummy_imu.squeeze(1).squeeze(1) if imu_dim > 0 else None) print(f"wrapper 与原模型最大差异: {(ref.squeeze(1) - orig_out).abs().max():.6e}") try: @@ -262,6 +281,8 @@ def main(): feeds = {f"frame{i}": dummy_frames[i].numpy() for i in range(f)} feeds["action_history"] = dummy_history.numpy() feeds["pred_token"] = dummy_pred.numpy() + if imu_dim > 0: + feeds["imu"] = dummy_imu.numpy() ort_out = sess.run(None, feeds)[0] print(f"ONNX 一致性验证通过（最大差异: {np.abs(ort_out - ref.numpy()).max():.6e}）") except ImportError: diff --git a/actpred/export_onnx.py b/actpred/export_onnx.py index 9c36c00..f117547 100644 --- a/actpred/export_onnx.py +++ b/actpred/export_onnx.py @@ -23,31 +23,38 @@ def main(): dummy_images = torch.randn( 1, data_cfg.frame_len, 3, data_cfg.image_height, data_cfg.image_width) dummy_history = torch.randn(1, data_cfg.history_len, data_cfg.action_dim) + # R2 方案一：imu_dim>0 时追加归一化姿态输入 (1, imu_dim) + imu_dim = data_cfg.imu_dim + extra_inputs = (torch.randn(1, imu_dim),) if imu_dim > 0 else () + extra_names = ["imu"] if imu_dim > 0 else [] batch = Dim("batch") + dynamic = {"images": {0: batch}, "action_history": {0: batch}} + if imu_dim > 0: + dynamic["imu"] = {0: batch} torch.onnx.export( - model, (dummy_images, dummy_history), + model, (dummy_images, dummy_history, *extra_inputs), args.output, - input_names=["images", "action_history"], + input_names=["images", "action_history"] + extra_names, output_names=["predicted_action"], - dynamic_shapes={ - "images": {0: batch}, - "action_history": {0: batch}, - }, + dynamic_shapes=dynamic, opset_version=18, ) - print(f"ONNX model exported to: {args.output}") + print(f"ONNX model exported to: {args.output}" + f"{'（含 imu 输入）' if imu_dim > 0 else ''}") # 可选：验证导出的 ONNX 模型一致性 try: import onnxruntime as ort import numpy as np sess = ort.InferenceSession(args.output) - img_np = dummy_images.numpy() - hist_np = dummy_history.numpy() - ort_out = sess.run(None, {"images": img_np, "action_history": hist_np})[0] + feeds = {"images": dummy_images.numpy(), + "action_history": dummy_history.numpy()} + for name, t in zip(extra_names, extra_inputs): + feeds[name] = t.numpy() + ort_out = sess.run(None, feeds)[0] with torch.no_grad(): - torch_out = model(dummy_images, dummy_history).numpy() + torch_out = model(dummy_images, dummy_history, *extra_inputs).numpy() max_diff = np.abs(ort_out - torch_out).max() print(f"ONNX 一致性验证通过（最大差异: {max_diff:.6e}）") except ImportError: diff --git a/actpred/gen_calibration.py b/actpred/gen_calibration.py index 5472cb3..d7b1893 100644 --- a/actpred/gen_calibration.py +++ b/actpred/gen_calibration.py @@ -137,14 +137,24 @@ def build_deploy_set(args): size=jitter.shape).astype(np.float32) pred = np.broadcast_to(v, (args.samples, 1, 1, d)).astype(np.float32) + jitter + inputs = [*frames, history, pred] + if cfg.imu_dim > 0: + # R2 方案一：imu 输入为归一化姿态（roll,pitch/180，ax..gz/32768）。 + # 校准覆盖典型工作范围：归一化值域 ±0.5（≈±90°/±16384）均匀采样 + imu = rng.uniform(-0.5, 0.5, + size=(args.samples, 1, 1, cfg.imu_dim)).astype(np.float32) + inputs.append(imu) + np.savez_compressed( args.output, **{f"actpred_deploy/input_layer{i + 1}": arr - for i, arr in enumerate([*frames, history, pred])}) + for i, arr in enumerate(inputs)}) print(f"部署校准集已保存: {args.output}") print(f" frame0~2: {frames[0].shape}（分辨率 {w}x{h}，与 checkpoint 一致）") print(f" action_history: {history.shape} pred_token: {pred.shape}" f"（学习值 ±{args.pred_jitter:.0%} 抖动）") + if cfg.imu_dim > 0: + print(f" imu: {inputs[-1].shape}（归一化值域 ±0.5 均匀采样）") def main(): diff --git a/actpred/hailo_infer.py b/actpred/hailo_infer.py index 445a50b..0a0dc67 100644 --- a/actpred/hailo_infer.py +++ b/actpred/hailo_infer.py @@ -6,9 +6,10 @@ HEF 由 actpred ONNX 经 Hailo DFC 编译生成（量化校准后）。本封装 vstream 名称/形状/量化参数，全部从 HEF 的 input/output vstream info 动态推导， 支持两种布局： - 新版 5 输入（export_deploy.py + dfc_compile.py 产物）： + 新版 5/6 输入（export_deploy.py + dfc_compile.py 产物）： frame0~2 各 (H,W,3) NHWC uint8（ImageNet 归一化校准，负零偏自动识别） + action_history (1,H_len,D) + pred_token (1,1,d) + + imu (1,1,1,imu_dim)（R2 方案一可选输入，原始姿态内部归一化） 旧版单图像输入： (1, F, 3, H, W) → NHWC uint8 (1, F, H, W, 3) 或 (1, H, W, F*3) → 帧沿通道拼接；action_history (1, H, D) @@ -34,6 +35,8 @@ import re import numpy as np from PIL import Image +from imu_norm import normalize_imu + IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) @@ -87,8 +90,8 @@ class HailoPredictor: raise ValueError( f"HEF 应只有一个输出 vstream，实际 {len(self._out_infos)}") self._out_info = self._out_infos[0] - self._frame_infos, self._history_info, self._pred_info = \ - self._classify_inputs() + self._frame_infos, self._history_info, self._pred_info, \ + self._imu_info = self._classify_inputs() self._derive_dims() self._action_mean = None @@ -110,11 +113,12 @@ class HailoPredictor: # ---------------- 输入分类与形状推导 ---------------- def _classify_inputs(self): - """将 HEF 输入分为 帧列表 / 动作历史 / pred_token（可能为 None） + """将 HEF 输入分为 帧列表 / 动作历史 / pred_token / imu（后两者可为 None） 多帧输入的绑定顺序按 input_layerN 数字后缀排序，而非 vstream 报告顺序：hailort 5.3.0 实测报告顺序被打乱（1,5,4,3,2）， 按报告顺序绑定会导致 frame1/frame2 错位。 + imu 为 R2 方案一新增输入（1,1,1,imu_dim），旧 HEF 无此输入。 """ ins = list(self._in_infos) out_dim = tuple(self._out_info.shape)[-1] @@ -127,24 +131,26 @@ class HailoPredictor: # 旧版：单一 4/5 维图像输入 frames = [i for i in ins if len(i.shape) in (4, 5)] rest = [i for i in ins if i not in frames] - history, pred = None, None + history, pred, imu = None, None, None for i in rest: s = tuple(i.shape) if len(s) == 3 and s[-1] == out_dim: history = i elif len(s) == 3 and s[1] == 1: pred = i + elif len(s) == 4 and s[0] == 1 and s[1] == 1 and s[2] == 1: + imu = i # R2：imu (1,1,1,imu_dim) else: raise ValueError( f"无法识别 HEF 输入 {i.name} 形状 {s}（期望动作历史 " - f"(1,H_len,D) 或 pred_token (1,1,d)）") + f"(1,H_len,D)、pred_token (1,1,d) 或 imu (1,1,1,imu_dim)）") if not frames: raise ValueError("HEF 输入中未找到图像输入") if history is None: raise ValueError( f"HEF 输入中未找到动作历史输入（输出维度 {out_dim}），" f"实际: {[(i.name, tuple(i.shape)) for i in ins]}") - return frames, history, pred + return frames, history, pred, imu def _derive_dims(self): # 帧输入布局 @@ -185,6 +191,11 @@ class HailoPredictor: h_shape = tuple(self._history_info.shape) self.history_len, self.action_dim = h_shape[1], h_shape[2] + # R2：imu 输入存在时暴露 imu_dim（pybot 闭环鸭子类型据此传 imu=）， + # 无则 None（旧 HEF 保持旧两参调用路径） + self.imu_dim = tuple(self._imu_info.shape)[-1] \ + if self._imu_info is not None else None + out_shape = tuple(self._out_info.shape) out_dim = out_shape[-1] if out_dim != self.action_dim: @@ -206,8 +217,11 @@ class HailoPredictor: # ---------------- 推理 ---------------- def predict(self, images, action_history: np.ndarray, - pred_token: np.ndarray = None) -> np.ndarray: + pred_token: np.ndarray = None, + imu: np.ndarray = None) -> np.ndarray: """images: 长度为 frame_len 的 RGB 图像列表（旧→新，末帧为当前帧） + imu: 可选原始姿态向量（roll,pitch,ax,ay,az,gx,gy,gz，HEF 含 imu + 输入时生效，内部按物理量程归一化；None 等价零向量） 返回 (action_dim,) float32 动作向量。pred_token 仅新版 5 输入 HEF 需要；默认喂零，抖动校准的新 HEF 建议喂 checkpoint 学习值。 @@ -254,6 +268,16 @@ class HailoPredictor: 0, 255).astype(np.uint8) bindings.input(self._pred_info.name).set_buffer(values) + if self._imu_info is not None: + i_shape = tuple(self._imu_info.shape) + # 原始姿态 → 物理量程归一化（与训练 dataset 同口径） + values = normalize_imu(imu, i_shape[-1]).reshape(i_shape) + if self._imu_info.format.type == FormatType.UINT8: + q = self._imu_info.quant_info + values = np.clip(np.round(values / q.qp_scale + q.qp_zp), + 0, 255).astype(np.uint8) + bindings.input(self._imu_info.name).set_buffer(values) + out_shape = tuple(self._out_info.shape) if self._out_info.format.type == FormatType.UINT8: out_buf = np.zeros(out_shape, dtype=np.uint8) @@ -325,11 +349,18 @@ def main(): parser.add_argument("--pred-token", default=None, help="pred_token .npy 路径（export_deploy.py " "--pred-token-out 产物；缺省喂零）") + parser.add_argument("--imu", type=str, default=None, + help="可选姿态向量（逗号分隔，原始量纲：roll,pitch 度 / " + "ax..gz int16 原始值）；HEF 含 imu 输入且缺省时按零向量") parser.add_argument("--output-format", choices=["json", "raw"], default="json") args = parser.parse_args() pred_token = np.load(args.pred_token) if args.pred_token else None + imu = None + if args.imu is not None: + imu = np.array([float(v) for v in args.imu.split(",")], + dtype=np.float32) with HailoPredictor(args.hef, stats_path=args.stats) as predictor: if len(args.images) != predictor.frame_len: raise SystemExit( @@ -337,7 +368,8 @@ def main(): f"实际为 {len(args.images)} 个") images = [np.array(Image.open(p).convert("RGB")) for p in args.images] history = parse_action_history(args.action_history, predictor.action_dim) - predicted = predictor.predict(images, history, pred_token=pred_token) + predicted = predictor.predict(images, history, pred_token=pred_token, + imu=imu) if args.output_format == "json": print(json.dumps({"predicted_action": predicted.tolist()}, indent=2)) diff --git a/actpred/imu_norm.py b/actpred/imu_norm.py new file mode 100644 index 0000000..992e683 --- /dev/null +++ b/actpred/imu_norm.py @@ -0,0 +1,37 @@ +"""IMU 姿态归一化 —— 训练（dataset）与推理（predict / hailo_infer）共用 + +量纲差异大（pybot MPU6050 原始输出，见 pybot/hd/mpu6050.py）： + roll/pitch 角度制，±180° + ax/ay/az 加速度计原始 int16（±8g 量程，1g≈4096 LSB） + gx/gy/gz 陀螺仪原始 int16（±500°/s 量程） + +按物理量程做固定缩放归一化（不依赖数据统计文件，部署链零状态）： + 索引 0,1 /180 → ≈[-1, 1] + 索引 2..7 /32768 → ≈[-1, 1] + +键序与 pybot robot.IMU_KEYS / recorder 一致： + [roll, pitch, ax, ay, az, gx, gy, gz] + +仅依赖 numpy，保证 hailo_infer.py（HAT 部署环境无 torch）可直接复用。 +""" + +import numpy as np + +# 各索引归一化除数（imu_dim=8 全量；更短的向量取前缀） +IMU_SCALES = np.array([180.0, 180.0, + 32768.0, 32768.0, 32768.0, + 32768.0, 32768.0, 32768.0], dtype=np.float32) + + +def normalize_imu(imu, imu_dim: int = 8) -> np.ndarray: + """原始姿态向量 → 归一化向量（pad/truncate 到 imu_dim） + + imu: list/np.ndarray，长度任意（None 等价全零）。 + 返回 (imu_dim,) float32。 + """ + vec = np.zeros(imu_dim, dtype=np.float32) + if imu is not None: + src = np.asarray(imu, dtype=np.float32).ravel() + n = min(len(src), imu_dim) + vec[:n] = src[:n] + return vec / IMU_SCALES[:imu_dim] diff --git a/actpred/model.py b/actpred/model.py index 50b229c..704cff3 100644 --- a/actpred/model.py +++ b/actpred/model.py @@ -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]) diff --git a/actpred/onnx_infer.py b/actpred/onnx_infer.py index 4f56f4b..2f282e0 100644 --- a/actpred/onnx_infer.py +++ b/actpred/onnx_infer.py @@ -1,10 +1,11 @@ """actpred ONNX Runtime 推理封装 —— CPU 后端（Pi 5 实测 5.9ms/次，PyTorch 的 ~5 倍） -接口对齐 predict.ActionPredictorInference：predict(images, action_history)。 +接口对齐 predict.ActionPredictorInference：predict(images, action_history, imu=None)。 模型维度从 ONNX session 输入元数据动态推导（images: [1,F,3,H,W]， -action_history: [1,H,D]），无需 checkpoint。 -图像预处理复现训练时的 dataset.py 流水线：Resize → /255 → ImageNet 归一化。 +action_history: [1,H,D]；R2 imu_dim>0 的模型另有 imu: [1,imu_dim]）， +无需 checkpoint。图像预处理复现训练时的 dataset.py 流水线： +Resize → /255 → ImageNet 归一化。 背景见 docs/reports/actpred_inference_speed_benchmark_20260813.md。 """ @@ -16,6 +17,8 @@ import os import numpy as np from PIL import Image +from imu_norm import normalize_imu + IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) @@ -54,11 +57,16 @@ class OnnxPredictor: inputs = {i.name: [d if isinstance(d, int) else -1 for d in i.shape] for i in self._session.get_inputs()} - img_shape, hist_shape = self._locate_inputs(inputs) + img_shape, hist_shape, imu_shape = self._locate_inputs(inputs) _, self.frame_len, _, self.image_height, self.image_width = img_shape _, self.history_len, self.action_dim = hist_shape self._image_name = next(n for n, s in inputs.items() if len(s) == 5) self._history_name = next(n for n, s in inputs.items() if len(s) == 3) + # R2 方案一：2 维输入 [1, imu_dim] 为姿态（无则 None）。 + # pybot 闭环鸭子类型约定：None 时按旧两参调用。 + self.imu_dim = imu_shape[-1] if imu_shape is not None else None + self._imu_name = (next(n for n, s in inputs.items() if len(s) == 2) + if imu_shape is not None else None) self._action_mean = None self._action_std = None @@ -75,17 +83,19 @@ class OnnxPredictor: @staticmethod def _locate_inputs(inputs): - """按维度数识别图像(5维)/动作历史(3维)输入，返回各自形状""" + """按维度数识别图像(5维)/动作历史(3维)/姿态(2维,可选)输入，返回各自形状""" img = [s for s in inputs.values() if len(s) == 5] hist = [s for s in inputs.values() if len(s) == 3] - if len(img) != 1 or len(hist) != 1: + imu = [s for s in inputs.values() if len(s) == 2] + if len(img) != 1 or len(hist) != 1 or len(imu) > 1: raise ValueError( - "ONNX 输入应为 1 个 5 维图像 + 1 个 3 维动作历史，" + "ONNX 输入应为 1 个 5 维图像 + 1 个 3 维动作历史" + "（R2 可再加 1 个 2 维 imu），" f"实际: {inputs}") if img[0][2] != 3: raise ValueError( f"图像输入形状应为 [1,F,3,H,W]，实际为 {img[0]}") - return img[0], hist[0] + return img[0], hist[0], (imu[0] if imu else None) # ---------------- 生命周期 ---------------- @@ -101,8 +111,11 @@ class OnnxPredictor: # ---------------- 推理 ---------------- - def predict(self, images, action_history: np.ndarray) -> np.ndarray: + def predict(self, images, action_history: np.ndarray, + imu: np.ndarray = None) -> np.ndarray: """images: 长度为 frame_len 的 RGB 图像列表（旧→新，末帧为当前帧） + imu: 可选原始姿态向量（roll,pitch,ax,ay,az,gx,gy,gz，模型含 imu + 输入时生效，内部按物理量程归一化；None 等价零向量） 返回 (action_dim,) float32 动作向量。 """ @@ -128,6 +141,9 @@ class OnnxPredictor: self._image_name: frames[np.newaxis], # (1,F,3,S,S) self._history_name: action_history[np.newaxis], # (1,H,D) } + if self._imu_name is not None: + feeds[self._imu_name] = normalize_imu( + imu, self.imu_dim)[np.newaxis] # (1,imu_dim) result = self._session.run(None, feeds)[0].reshape(-1).astype(np.float32) if self._action_mean is not None: result = result * self._action_std + self._action_mean @@ -152,6 +168,9 @@ def main(): help="冒号分隔各步，每步 action_dim 个值") parser.add_argument("--stats", default=None, help="action_stats.npy 路径（训练时启用归一化才需要）") + parser.add_argument("--imu", type=str, default=None, + help="可选姿态向量（逗号分隔，原始量纲：roll,pitch 度 / " + "ax..gz int16 原始值）；模型含 imu 输入且缺省时按零向量") parser.add_argument("--output-format", choices=["json", "raw"], default="json") args = parser.parse_args() @@ -163,7 +182,11 @@ def main(): f"实际为 {len(args.images)} 个") images = [np.array(Image.open(p).convert("RGB")) for p in args.images] history = parse_action_history(args.action_history, predictor.action_dim) - predicted = predictor.predict(images, history) + imu = None + if args.imu is not None: + imu = np.array([float(v) for v in args.imu.split(",")], + dtype=np.float32) + predicted = predictor.predict(images, history, imu=imu) if args.output_format == "json": print(json.dumps({"predicted_action": predicted.tolist()}, indent=2)) diff --git a/actpred/predict.py b/actpred/predict.py index e37775b..e671e43 100644 --- a/actpred/predict.py +++ b/actpred/predict.py @@ -7,6 +7,7 @@ from PIL import Image from torchvision import transforms from config import parse_checkpoint_config +from imu_norm import normalize_imu from model import ActionPredictor @@ -29,6 +30,10 @@ class ActionPredictorInference: self.action_dim = self._data_cfg.action_dim self.history_len = self._data_cfg.history_len self.frame_len = self._data_cfg.frame_len + # R2 方案一：imu_dim>0 时 predict 接受 imu= 原始姿态向量（内部归一化）。 + # pybot 闭环鸭子类型约定：此属性为 None 时按旧两参调用， + # 故 imu_dim=0（旧 checkpoint）必须映射为 None。 + self.imu_dim = self._data_cfg.imu_dim or None # R1-C 残差预测：checkpoint 记录 predict_delta=True 时，模型输出 Δ̂， # 需在归一化空间加回 a_{t-1}（history 末步）得到绝对动作。 # 旧 checkpoint 无该字段，getattr 默认 False，行为与旧版一致。 @@ -67,8 +72,12 @@ class ActionPredictorInference: self._action_std = torch.from_numpy(stats[1].astype(np.float32)).to(self._device) @torch.no_grad() - def predict(self, images, action_history: np.ndarray) -> np.ndarray: - """images: 长度为 frame_len 的 RGB 图像序列（旧→新，末帧为当前帧）""" + def predict(self, images, action_history: np.ndarray, + imu: np.ndarray = None) -> np.ndarray: + """images: 长度为 frame_len 的 RGB 图像序列（旧→新，末帧为当前帧） + imu: 可选原始姿态向量（roll,pitch,ax,ay,az,gx,gy,gz，imu_dim>0 时 + 生效，内部按物理量程归一化；None 等价零向量 = 无姿态先验） + """ if not isinstance(images, (list, tuple)): images = [images] if len(images) != self._data_cfg.frame_len: @@ -88,7 +97,12 @@ class ActionPredictorInference: if self._action_mean is not None: history = (history - self._action_mean) / self._action_std history = history.unsqueeze(0).to(self._device) - pred = self._model(img_tensor, history).squeeze(0) # 归一化空间 + imu_tensor = None + if self._data_cfg.imu_dim > 0: + imu_tensor = torch.from_numpy( + normalize_imu(imu, self._data_cfg.imu_dim) + ).unsqueeze(0).to(self._device) + pred = self._model(img_tensor, history, imu_tensor).squeeze(0) # 归一化空间 if self._predict_delta: # 残差：Δ̂ 加回 a_{t-1}（history 末步，同为归一化空间）→ 绝对动作 pred = history[0, -1] + pred @@ -118,6 +132,9 @@ def main(): parser.add_argument("--action-history", required=True, help="冒号分隔各步，每步 action_dim 个值，共 history_len 步。" "e.g. '0,0,0,0:1,0,0,0' (action_dim=4, history_len=2)") + parser.add_argument("--imu", type=str, default=None, + help="可选姿态向量（逗号分隔，原始量纲：roll,pitch 度 / " + "ax..gz int16 原始值）；模型 imu_dim>0 且缺省时按零向量") parser.add_argument("--output-format", choices=["json", "raw"], default="json") parser.add_argument("--device", type=str, default=None) args = parser.parse_args() @@ -131,7 +148,10 @@ def main(): f"实际为 {len(args.images)} 个。冷启动时可重复传同一帧路径。") images = [np.array(Image.open(p).convert("RGB")) for p in args.images] history = parse_action_history(args.action_history, action_dim) - predicted = predictor.predict(images, history) + imu = None + if args.imu is not None: + imu = np.array([float(v) for v in args.imu.split(",")], dtype=np.float32) + predicted = predictor.predict(images, history, imu=imu) if args.output_format == "json": print(json.dumps({"predicted_action": predicted.tolist()}, indent=2)) diff --git a/actpred/test_actpred.py b/actpred/test_actpred.py index 3be9a7d..333a62e 100644 --- a/actpred/test_actpred.py +++ b/actpred/test_actpred.py @@ -1714,7 +1714,8 @@ class TestR4Evaluate(unittest.TestCase): with tempfile.TemporaryDirectory() as tmpdir: actions = np.arange(20, dtype=np.float32).reshape(10, 2) self._create_episode(Path(tmpdir), "episode_0001", actions) - image_files, acts = evaluate._load_episode(tmpdir, "episode_0001") + image_files, acts, _ = evaluate._load_episode( + tmpdir, "episode_0001") pred = self._ConstPredictor([100.0, 100.0]) steps = evaluate.rollout_episode(pred, image_files, acts) # t0 = max(H=2, F-1=1) = 2，horizon = 10-2 = 8 @@ -1729,7 +1730,8 @@ class TestR4Evaluate(unittest.TestCase): with tempfile.TemporaryDirectory() as tmpdir: actions = np.zeros((10, 2), dtype=np.float32) self._create_episode(Path(tmpdir), "episode_0001", actions) - image_files, acts = evaluate._load_episode(tmpdir, "episode_0001") + image_files, acts, _ = evaluate._load_episode( + tmpdir, "episode_0001") pred = self._ConstPredictor([0.0, 0.0]) steps = evaluate.rollout_episode(pred, image_files, acts, horizon=3) self.assertEqual(len(steps), 3) @@ -1739,7 +1741,8 @@ class TestR4Evaluate(unittest.TestCase): with tempfile.TemporaryDirectory() as tmpdir: actions = np.zeros((2, 2), dtype=np.float32) # n=2 <= t0=2 self._create_episode(Path(tmpdir), "episode_0001", actions) - image_files, acts = evaluate._load_episode(tmpdir, "episode_0001") + image_files, acts, _ = evaluate._load_episode( + tmpdir, "episode_0001") pred = self._ConstPredictor([0.0, 0.0]) self.assertIsNone(evaluate.rollout_episode(pred, image_files, acts)) @@ -1758,7 +1761,8 @@ class TestR4Evaluate(unittest.TestCase): with tempfile.TemporaryDirectory() as tmpdir: actions = np.zeros((10, 2), dtype=np.float32) # 真值全 0 self._create_episode(Path(tmpdir), "episode_0001", actions) - image_files, acts = evaluate._load_episode(tmpdir, "episode_0001") + image_files, acts, _ = evaluate._load_episode( + tmpdir, "episode_0001") pred = self._ConstPredictor([5.0, 5.0]) # 远超 [0,0] 范围 lo = np.zeros(2, dtype=np.float32) hi = np.zeros(2, dtype=np.float32) @@ -1797,5 +1801,489 @@ class TestR4Evaluate(unittest.TestCase): np.testing.assert_allclose(hi, [11.0, 11.0], atol=1e-6) +class TestR2ImuNorm(unittest.TestCase): + """R2 姿态归一化：物理量程除数、pad/truncate、None 兜底""" + + def test_scales_applied(self): + from imu_norm import normalize_imu, IMU_SCALES + raw = [180.0, -180.0, 32768.0, -32768.0, 16384.0, 655.36, -655.36, 0.0] + out = normalize_imu(raw, 8) + np.testing.assert_allclose(out, np.asarray(raw) / IMU_SCALES, atol=1e-6) + + def test_none_gives_zeros(self): + from imu_norm import normalize_imu + out = normalize_imu(None, 8) + self.assertEqual(out.shape, (8,)) + np.testing.assert_array_equal(out, np.zeros(8, np.float32)) + + def test_short_input_pads_zero(self): + from imu_norm import normalize_imu + out = normalize_imu([90.0, -90.0], 8) + self.assertEqual(out.shape, (8,)) + np.testing.assert_allclose(out[:2], [0.5, -0.5], atol=1e-6) + np.testing.assert_array_equal(out[2:], np.zeros(6, np.float32)) + + def test_long_input_truncates(self): + from imu_norm import normalize_imu + out = normalize_imu([180.0] * 10, 8) + self.assertEqual(out.shape, (8,)) + np.testing.assert_allclose(out[:2], [1.0, 1.0], atol=1e-6) + + def test_imu_dim_prefix_of_scales(self): + """imu_dim<8 时取 IMU_SCALES 前缀（roll/pitch 子集可用）""" + from imu_norm import normalize_imu + out = normalize_imu([180.0, -180.0], 2) + self.assertEqual(out.shape, (2,)) + np.testing.assert_allclose(out, [1.0, -1.0], atol=1e-6) + + +class TestR2ImuModel(unittest.TestCase): + """R2 方案一模型：独立 IMU token 注入（序列 F+H+2）""" + + @staticmethod + def _make_cfg(imu_dim): + cfg = Config() + cfg.data.image_width = 32 + cfg.data.image_height = 32 + cfg.data.action_dim = 4 + cfg.data.history_len = 3 + cfg.data.frame_len = 3 + cfg.data.imu_dim = imu_dim + return cfg + + def test_imu_token_changes_output(self): + """imu_dim>0：不同姿态输入应产生不同输出（token 真的进了序列）""" + cfg = self._make_cfg(imu_dim=8) + model = ActionPredictor(cfg.model, cfg.data).eval() + imgs = torch.randn(1, 3, 3, 32, 32) + hist = torch.randn(1, 3, 4) + with torch.no_grad(): + out_zero = model(imgs, hist, torch.zeros(1, 8)).numpy() + out_imu = model(imgs, hist, torch.randn(1, 8)).numpy() + self.assertGreater(np.abs(out_zero - out_imu).max(), 1e-6) + + def test_none_imu_equals_zero_vector(self): + """imu=None 应与显式零向量一致（无姿态先验语义）""" + cfg = self._make_cfg(imu_dim=8) + model = ActionPredictor(cfg.model, cfg.data).eval() + imgs = torch.randn(1, 3, 3, 32, 32) + hist = torch.randn(1, 3, 4) + with torch.no_grad(): + a = model(imgs, hist).numpy() + b = model(imgs, hist, torch.zeros(1, 8)).numpy() + np.testing.assert_allclose(a, b, atol=1e-6) + + def test_imu_dim_zero_has_no_imu_proj(self): + """imu_dim=0（默认）不建 imu_proj，旧 checkpoint 完全兼容""" + cfg = self._make_cfg(imu_dim=0) + model = ActionPredictor(cfg.model, cfg.data) + self.assertFalse(hasattr(model, "imu_proj")) + # 序列长度保持 F+H+1 + self.assertEqual(model.pos_embed.shape[1], 3 + 3 + 1) + + def test_imu_dim_positive_seq_len(self): + """imu_dim>0 时序列长度 F+H+2""" + cfg = self._make_cfg(imu_dim=8) + model = ActionPredictor(cfg.model, cfg.data) + self.assertEqual(model.pos_embed.shape[1], 3 + 3 + 2) + + def test_dynamic_batch_with_imu(self): + """batch 维动态：imu 输入随 batch 变化""" + cfg = self._make_cfg(imu_dim=8) + model = ActionPredictor(cfg.model, cfg.data).eval() + with torch.no_grad(): + out = model(torch.randn(2, 3, 3, 32, 32), + torch.randn(2, 3, 4), torch.randn(2, 8)) + self.assertEqual(out.shape, (2, 4)) + + +class TestR2ImuDataset(unittest.TestCase): + """R2 数据集：imu.npy 加载、归一化、缺失补零、形状修正""" + + @staticmethod + def _create_episode(base_dir, name, n, action_dim, image_size=16, + imus=None): + ep_dir = base_dir / name + images_dir = ep_dir / "images" + images_dir.mkdir(parents=True, exist_ok=True) + for i in range(n): + Image.fromarray(np.random.randint(0, 255, (image_size, image_size, 3), + dtype=np.uint8)).save( + images_dir / f"{i:04d}.jpg") + actions = np.arange(n * action_dim, dtype=np.float32).reshape( + n, action_dim) + np.save(ep_dir / "actions.npy", actions) + if imus is not None: + np.save(ep_dir / "imu.npy", np.asarray(imus, dtype=np.float32)) + + def test_imu_loaded_and_normalized(self): + """imu_dim>0 且 imu.npy 存在 → 样本带归一化 imu 键""" + with tempfile.TemporaryDirectory() as tmpdir: + # roll=180 → 归一化 1.0；ax=32768 → 1.0 + imus = [[180.0, -180.0, 32768.0, 0, 0, 0, 0, 0]] * 10 + self._create_episode(Path(tmpdir), "episode_0001", 10, 2, + imus=imus) + ds = ActionDataset(tmpdir, history_len=2, action_dim=2, + image_width=16, image_height=16, frame_len=2, + imu_dim=8) + sample = ds[0] + self.assertIn("imu", sample) + np.testing.assert_allclose( + sample["imu"].numpy(), + [1.0, -1.0, 1.0, 0, 0, 0, 0, 0], atol=1e-6) + + def test_missing_imu_fills_zero(self): + """无 imu.npy 的 episode → 零向量（混合数据集可训）""" + with tempfile.TemporaryDirectory() as tmpdir: + self._create_episode(Path(tmpdir), "episode_0001", 10, 2) + ds = ActionDataset(tmpdir, history_len=2, action_dim=2, + image_width=16, image_height=16, frame_len=2, + imu_dim=8) + np.testing.assert_array_equal( + ds[0]["imu"].numpy(), np.zeros(8, np.float32)) + + def test_imu_dim_zero_no_imu_key(self): + """imu_dim=0 时样本无 imu 键（旧行为不变）""" + with tempfile.TemporaryDirectory() as tmpdir: + self._create_episode(Path(tmpdir), "episode_0001", 10, 2, + imus=[[1.0] * 8] * 10) + ds = ActionDataset(tmpdir, history_len=2, action_dim=2, + image_width=16, image_height=16, frame_len=2) + self.assertNotIn("imu", ds[0]) + + def test_imu_dim_mismatch_padded(self): + """imu.npy 列数与 imu_dim 不一致 → pad/truncate 不炸""" + with tempfile.TemporaryDirectory() as tmpdir: + imus = [[180.0, -180.0]] * 10 # 只有 2 列，imu_dim=8 + self._create_episode(Path(tmpdir), "episode_0001", 10, 2, + imus=imus) + ds = ActionDataset(tmpdir, history_len=2, action_dim=2, + image_width=16, image_height=16, frame_len=2, + imu_dim=8) + imu = ds[0]["imu"].numpy() + self.assertEqual(imu.shape, (8,)) + np.testing.assert_allclose(imu[:2], [1.0, -1.0], atol=1e-6) + np.testing.assert_array_equal(imu[2:], np.zeros(6, np.float32)) + + def test_short_imu_aligned_with_zero_pad(self): + """imu.npy 行数少于帧数 → 尾部补零对齐""" + with tempfile.TemporaryDirectory() as tmpdir: + imus = [[180.0] + [0.0] * 7] * 5 # 5 行 < 10 帧 + self._create_episode(Path(tmpdir), "episode_0001", 10, 2, + imus=imus) + ds = ActionDataset(tmpdir, history_len=2, action_dim=2, + image_width=16, image_height=16, frame_len=2, + imu_dim=8) + # 第 0 样本 t_start=2（有 imu），末样本 t_start=8（超出 5 行 → 0） + np.testing.assert_allclose(ds[0]["imu"].numpy()[0], 1.0, atol=1e-6) + np.testing.assert_array_equal(ds[-1]["imu"].numpy(), + np.zeros(8, np.float32)) + + def test_rollout_steps_carry_imu(self): + """多步展开时每步带对应时刻的 imu""" + with tempfile.TemporaryDirectory() as tmpdir: + imus = [[float(t), 0, 0, 0, 0, 0, 0, 0] for t in range(10)] + self._create_episode(Path(tmpdir), "episode_0001", 10, 2, + imus=imus) + ds = ActionDataset(tmpdir, history_len=2, action_dim=2, + image_width=16, image_height=16, frame_len=2, + rollout_len=3, imu_dim=8) + steps = ds[0]["steps"] + for k, step in enumerate(steps): + t = 2 + k + self.assertIn("imu", step) + np.testing.assert_allclose(step["imu"].numpy()[0], + t / 180.0, atol=1e-6) + + +@unittest.skipUnless( + importlib.util.find_spec("onnxscript") is not None + and importlib.util.find_spec("onnxruntime") is not None, + "OnnxPredictor 测试需要 onnxscript + onnxruntime") +class TestR2OnnxPredictorImu(unittest.TestCase): + """R2 ONNX 后端：含 imu 输入的导出模型推理对齐""" + + @staticmethod + def _export_imu_onnx(tmpdir): + cfg = Config() + cfg.data.image_width = 32 + cfg.data.image_height = 32 + cfg.data.action_dim = 4 + cfg.data.history_len = 3 + cfg.data.frame_len = 3 + cfg.data.imu_dim = 8 + model = ActionPredictor(cfg.model, cfg.data) + model.eval() + out_path = os.path.join(tmpdir, "model_imu.onnx") + torch.onnx.export( + model, + (torch.randn(1, 3, 3, 32, 32), torch.randn(1, 3, 4), + torch.randn(1, 8)), + out_path, + input_names=["images", "action_history", "imu"], + output_names=["predicted_action"], + opset_version=18) + return out_path, model + + def _frames(self, n): + return [np.random.randint(0, 255, (48, 64, 3), dtype=np.uint8) + for _ in range(n)] + + def test_imu_dim_detected(self): + from onnx_infer import OnnxPredictor + with tempfile.TemporaryDirectory() as tmpdir: + path, _ = self._export_imu_onnx(tmpdir) + with OnnxPredictor(path) as p: + self.assertEqual(p.imu_dim, 8) + + def test_imu_predict_matches_torch(self): + """predict(imu=原始姿态) 应与 torch（手动归一化）前向一致""" + from imu_norm import normalize_imu + from onnx_infer import OnnxPredictor + with tempfile.TemporaryDirectory() as tmpdir: + path, model = self._export_imu_onnx(tmpdir) + frames = self._frames(3) + history = np.random.randn(3, 4).astype(np.float32) + imu_raw = np.array([5.0, -3.0, 100, 200, 16384, 10, -20, 30], + dtype=np.float32) + with OnnxPredictor(path) as p: + onnx_out = p.predict(frames, history, imu=imu_raw) + pre = np.stack([p._preprocess(f) for f in frames]) + with torch.no_grad(): + torch_out = model( + torch.from_numpy(pre[np.newaxis]), + torch.from_numpy(history[np.newaxis]), + torch.from_numpy( + normalize_imu(imu_raw, 8)[np.newaxis])).numpy() + np.testing.assert_allclose(onnx_out, torch_out.reshape(-1), + atol=1e-4) + + def test_imu_changes_output(self): + """姿态输入应影响输出（确认 imu 真的喂进了图）""" + from onnx_infer import OnnxPredictor + with tempfile.TemporaryDirectory() as tmpdir: + path, _ = self._export_imu_onnx(tmpdir) + frames = self._frames(3) + history = np.random.randn(3, 4).astype(np.float32) + with OnnxPredictor(path) as p: + a = p.predict(frames, history, imu=None) + b = p.predict(frames, history, + imu=np.array([180.0] + [0.0] * 7, np.float32)) + self.assertGreater(np.abs(a - b).max(), 1e-6) + + def test_no_imu_model_imu_dim_none(self): + """旧两输入 ONNX：imu_dim 为 None（pybot 鸭子类型走旧两参路径）""" + from onnx_infer import OnnxPredictor + with tempfile.TemporaryDirectory() as tmpdir: + path, _ = TestOnnxPredictor._export_small_onnx(tmpdir) + with OnnxPredictor(path) as p: + self.assertIsNone(p.imu_dim) + + +class TestR2EvaluateImu(unittest.TestCase): + """R2 评估：滚动时姿态真值按时刻传给 predict""" + + class _ImuRecordingPredictor: + """记录每步收到的 imu 的假推理器""" + frame_len = 2 + history_len = 2 + action_dim = 2 + imu_dim = 8 + + def __init__(self): + self.imus_received = [] + + def predict(self, images, history, imu=None): + self.imus_received.append( + None if imu is None else np.asarray(imu).copy()) + return np.zeros(2, dtype=np.float32) + + def test_rollout_passes_imu_per_step(self): + """imus 存在时逐步传 imus[t]（t0 起步与帧窗对齐）""" + import evaluate + with tempfile.TemporaryDirectory() as tmpdir: + actions = np.zeros((10, 2), dtype=np.float32) + TestR4Evaluate._create_episode(Path(tmpdir), "episode_0001", + actions) + imus = np.arange(80, dtype=np.float32).reshape(10, 8) + np.save(Path(tmpdir) / "episode_0001" / "imu.npy", imus) + image_files, acts, loaded = evaluate._load_episode( + tmpdir, "episode_0001") + np.testing.assert_array_equal(loaded, imus) + pred = self._ImuRecordingPredictor() + steps = evaluate.rollout_episode(pred, image_files, acts, + imus=loaded) + self.assertEqual(len(steps), 8) # t0=2 → t=2..9 + for k, rec in enumerate(steps): + t = 2 + k + np.testing.assert_array_equal(pred.imus_received[k], imus[t]) + + def test_rollout_without_imu_passes_none(self): + """imu.npy 缺失 → 每步 imu=None（等价零向量，不炸）""" + import evaluate + with tempfile.TemporaryDirectory() as tmpdir: + actions = np.zeros((10, 2), dtype=np.float32) + TestR4Evaluate._create_episode(Path(tmpdir), "episode_0001", + actions) + image_files, acts, loaded = evaluate._load_episode( + tmpdir, "episode_0001") + self.assertIsNone(loaded) + pred = self._ImuRecordingPredictor() + evaluate.rollout_episode(pred, image_files, acts, imus=loaded) + self.assertTrue(all(v is None for v in pred.imus_received)) + + def test_legacy_predictor_not_given_imu_kwarg(self): + """imu_dim=None 的旧推理器不应收到 imu 关键字（鸭子类型兼容）""" + import evaluate + with tempfile.TemporaryDirectory() as tmpdir: + actions = np.zeros((10, 2), dtype=np.float32) + TestR4Evaluate._create_episode(Path(tmpdir), "episode_0001", + actions) + image_files, acts, _ = evaluate._load_episode( + tmpdir, "episode_0001") + pred = TestR4Evaluate._ConstPredictor([0.0, 0.0]) # 无 imu_dim + steps = evaluate.rollout_episode(pred, image_files, acts) + self.assertEqual(len(steps), 8) + self.assertEqual(pred.calls, 8) + + +class TestR2HailoImu(unittest.TestCase): + """R2 HEF 封装：6 输入布局（frame×3 + history + pred_token + imu）""" + + def _frames(self, n, size=8): + return [np.random.randint(0, 255, (size, size, 3), dtype=np.uint8) + for _ in range(n)] + + def _install_six_input_fake(self, seen): + infos_in = [ + _FakeVStreamInfo("net/input_layer1", (8, 8, 3), + _FakeFormatType.FLOAT32), + _FakeVStreamInfo("net/input_layer2", (8, 8, 3), + _FakeFormatType.FLOAT32), + _FakeVStreamInfo("net/input_layer3", (8, 8, 3), + _FakeFormatType.FLOAT32), + _FakeVStreamInfo("net/input_layer4", (1, 2, 3), + _FakeFormatType.FLOAT32), + _FakeVStreamInfo("net/input_layer5", (1, 1, 4), + _FakeFormatType.FLOAT32), + _FakeVStreamInfo("net/input_layer6", (1, 1, 1, 8), + _FakeFormatType.FLOAT32), + ] + infos_out = [_FakeVStreamInfo("net/conv49", (1, 1, 3), + _FakeFormatType.FLOAT32)] + + def transform(b): + seen["imu"] = np.asarray( + b.input("net/input_layer6").get_buffer()).copy() + _first_row_transform("net/input_layer4", "net/conv49")(b) + + _install_fake_hailo(self, infos_in, infos_out, transform) + + def test_imu_input_classified_and_bound(self): + """imu (1,1,1,8) 输入应被识别为姿态而非 pred_token，绑定归一化值""" + from hailo_infer import HailoPredictor + seen = {} + self._install_six_input_fake(seen) + with HailoPredictor("fake.hef") as p: + self.assertEqual(p.imu_dim, 8) + # roll=180° → 归一化 1.0 + p.predict(self._frames(3), np.zeros((2, 3), np.float32), + imu=np.array([180.0] + [0.0] * 7, np.float32)) + self.assertEqual(seen["imu"].shape, (1, 1, 1, 8)) + np.testing.assert_allclose(seen["imu"][0, 0, 0, 0], 1.0, atol=1e-6) + np.testing.assert_array_equal(seen["imu"][0, 0, 0, 1:], + np.zeros(7, np.float32)) + + def test_imu_none_binds_zeros(self): + """imu=None → 零向量绑定（无姿态先验）""" + from hailo_infer import HailoPredictor + seen = {} + self._install_six_input_fake(seen) + with HailoPredictor("fake.hef") as p: + p.predict(self._frames(3), np.zeros((2, 3), np.float32)) + np.testing.assert_array_equal(seen["imu"], np.zeros((1, 1, 1, 8))) + + def test_imu_uint8_quantized(self): + """uint8 量化边缘：归一化值按 quant_info 量化""" + from hailo_infer import HailoPredictor + q_imu = _FakeQuantInfo(qp_zp=128.0, qp_scale=0.01) + infos_in = [ + _FakeVStreamInfo("net/input_layer1", (8, 8, 3), + _FakeFormatType.FLOAT32), + _FakeVStreamInfo("net/input_layer2", (8, 8, 3), + _FakeFormatType.FLOAT32), + _FakeVStreamInfo("net/input_layer3", (8, 8, 3), + _FakeFormatType.FLOAT32), + _FakeVStreamInfo("net/input_layer4", (1, 2, 3), + _FakeFormatType.FLOAT32), + _FakeVStreamInfo("net/input_layer5", (1, 1, 4), + _FakeFormatType.FLOAT32), + _FakeVStreamInfo("net/input_layer6", (1, 1, 1, 8), + _FakeFormatType.UINT8, q_imu), + ] + infos_out = [_FakeVStreamInfo("net/conv49", (1, 1, 3), + _FakeFormatType.FLOAT32)] + seen = {} + + def transform(b): + seen["imu"] = np.asarray( + b.input("net/input_layer6").get_buffer()).copy() + _first_row_transform("net/input_layer4", "net/conv49")(b) + + _install_fake_hailo(self, infos_in, infos_out, transform) + with HailoPredictor("fake.hef") as p: + # roll=18° → 归一化 0.1 → 量化 round(0.1/0.01+128)=138 + p.predict(self._frames(3), np.zeros((2, 3), np.float32), + imu=np.array([18.0] + [0.0] * 7, np.float32)) + self.assertEqual(seen["imu"].dtype, np.uint8) + self.assertEqual(seen["imu"][0, 0, 0, 0], 138) + # 零值 → round(0/0.01+128)=128 + np.testing.assert_array_equal(seen["imu"][0, 0, 0, 1:], + np.full(7, 128, np.uint8)) + + +class TestR2DfcCompileImu(unittest.TestCase): + """R2 dfc_compile：--with-imu 时 start_nodes 末尾追加 imu""" + + def test_start_nodes_with_imu(self): + from dfc_compile import start_nodes + self.assertEqual(start_nodes(3, with_imu=True), + ["frame0", "frame1", "frame2", + "action_history", "pred_token", "imu"]) + + def test_start_nodes_without_imu_unchanged(self): + from dfc_compile import start_nodes + self.assertEqual(start_nodes(3), + ["frame0", "frame1", "frame2", + "action_history", "pred_token"]) + self.assertEqual(start_nodes(3, with_imu=False), start_nodes(3)) + + +class TestR2ConfigImu(unittest.TestCase): + """R2 配置：imu_dim 校验与覆盖""" + + def test_default_imu_dim_zero(self): + self.assertEqual(DataConfig().imu_dim, 0) + + def test_imu_dim_out_of_range_raises(self): + data = DataConfig(imu_dim=65) + with self.assertRaises(ValueError): + validate_config_values(data, ModelConfig()) + data = DataConfig(imu_dim=-1) + with self.assertRaises(ValueError): + validate_config_values(data, ModelConfig()) + + def test_imu_dim_cli_override(self): + args = build_argparser().parse_args(["--imu-dim", "8"]) + cfg = apply_overrides(Config(), args) + self.assertEqual(cfg.data.imu_dim, 8) + + def test_checkpoint_config_roundtrip_with_imu(self): + cfg = Config() + cfg.data.imu_dim = 8 + restored = parse_checkpoint_config(cfg.to_dict()) + self.assertEqual(restored.data.imu_dim, 8) + + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/actpred/train.py b/actpred/train.py index 89d756e..d2ca160 100644 --- a/actpred/train.py +++ b/actpred/train.py @@ -57,7 +57,8 @@ def train_one_epoch(model, loader, optimizer, device, max_grad_norm: float = Non images = batch["images"].to(device) history = batch["action_history"].to(device) target = batch["target"].to(device) - pred = model(images, history) + imu = batch["imu"].to(device) if "imu" in batch else None + pred = model(images, history, imu) loss = torch.nn.functional.mse_loss(pred, target) bs = images.size(0) else: @@ -71,10 +72,11 @@ def train_one_epoch(model, loader, optimizer, device, max_grad_norm: float = Non true_history = step["action_history"].to(device) target = step["target"].to(device) prev_action = step["prev_action"].to(device) + imu = step["imu"].to(device) if "imu" in step else None if k == 0 or rng.random_sample() < ss_prob: history = true_history # 否则沿用上一轮回传的 history（自预测构造） - pred = model(images, history) + pred = model(images, history, imu) loss = loss + torch.nn.functional.mse_loss(pred, target) bs = images.size(0) # 回传：残差模式先把 Δ̂ 加回 a_{t-1} 得到绝对动作，再移入历史 @@ -101,7 +103,8 @@ def validate(model, loader, device, rollout_len: int = 1, predict_delta: bool = images = batch["images"].to(device) history = batch["action_history"].to(device) target = batch["target"].to(device) - pred = model(images, history) + imu = batch["imu"].to(device) if "imu" in batch else None + pred = model(images, history, imu) loss = torch.nn.functional.mse_loss(pred, target) bs = images.size(0) else: @@ -113,7 +116,8 @@ def validate(model, loader, device, rollout_len: int = 1, predict_delta: bool = images = step["images"].to(device) history = step["action_history"].to(device) target = step["target"].to(device) - pred = model(images, history) + imu = step["imu"].to(device) if "imu" in step else None + pred = model(images, history, imu) loss = loss + torch.nn.functional.mse_loss(pred, target) bs = images.size(0) loss = loss / len(batch["steps"]) @@ -151,7 +155,8 @@ def main(): cfg.data.image_width, cfg.data.image_height, cfg.data.normalize_actions, cfg.data.action_stats_path, train_indices, cfg.data.frame_len, augment=cfg.train.augment, - predict_delta=cfg.train.predict_delta, rollout_len=cfg.train.rollout_len) + predict_delta=cfg.train.predict_delta, rollout_len=cfg.train.rollout_len, + imu_dim=cfg.data.imu_dim) if cfg.data.normalize_actions and not cfg.data.action_stats_path: # 归一化统计量必须只用训练集计算，再供验证集复用，否则 val loss 不可信 stats_path = os.path.join(cfg.train.checkpoint_dir, "action_stats.npy") @@ -163,7 +168,8 @@ def main(): cfg.data.data_dir, cfg.data.history_len, cfg.data.action_dim, cfg.data.image_width, cfg.data.image_height, cfg.data.normalize_actions, cfg.data.action_stats_path, val_indices, cfg.data.frame_len, - predict_delta=cfg.train.predict_delta, rollout_len=cfg.train.rollout_len) + predict_delta=cfg.train.predict_delta, rollout_len=cfg.train.rollout_len, + imu_dim=cfg.data.imu_dim) print(f"Samples: {len(train_ds)} train, {len(val_ds)} val") if len(train_ds) == 0: diff --git a/pybot/README.md b/pybot/README.md index 64de207..bcad940 100644 --- a/pybot/README.md +++ b/pybot/README.md @@ -31,6 +31,10 @@ cat scripts/std.mt | ./_run.sh # 与 bot4cj 相同的管道用法 ./_run.sh --record data/ep scripts/std.mt # 录制 actpred 格式 episode ``` +`--record` 时若脚本含 `@MPU6050_BIND` 块，逐帧同步落盘 `imu.npy` +（T×8：roll,pitch,ax,ay,az,gx,gy,gz，与图像/动作按帧时刻天然对齐）； +无 IMU 硬件则保持旧格式（仅 images + actions.npy）。 + `-i` / `-o` 参数位保留（reserved），未来对接跨进程 FIFO。 ### 实时预测闭环（`--loop`） @@ -49,6 +53,10 @@ cat scripts/std.mt | ./_run.sh # 与 bot4cj 相同的管道用法 - **帧源**：闭环从 camera 守护线程队列取最近 F 帧（无阻塞），`@Camera` 的 `framerate` 建议 ≥ 2×`hz` 以保证帧新鲜度；`capacity` 应 ≥ F（帧窗长度）。 部署时若日志出现 `loop overrun` 告警，说明 `hz` 超过推理+执行能力，应降低 `hz` +- **姿态回流（R2）**：脚本含 `@MPU6050_BIND` 且模型带 imu 输入 + （`predictor.imu_dim` 非 None）时，每步读取 MPU6050 最新姿态随帧/历史 + 一同传给推理器（原始量纲，归一化在推理端完成）；模型无 imu 输入时 + 不接入推理（若带 `--record` 仍会读取落盘 imu.npy） ## 目录 diff --git a/pybot/pybot/recorder.py b/pybot/pybot/recorder.py index d701ed9..fe21d0a 100644 --- a/pybot/pybot/recorder.py +++ b/pybot/pybot/recorder.py @@ -6,6 +6,7 @@ images/ 0000.jpg, 0001.jpg, ... actions.npy # shape (T, action_dim) float32 + imu.npy # 可选，shape (T, imu_dim) float32（MPU6050 姿态回流，R2） numpy / Pillow 为可选依赖（lazy import），仅 --record 时需要， 不影响 pybot 的零依赖运行时承诺。 @@ -18,19 +19,26 @@ log = logging.getLogger("pybot.recorder") class EpisodeRecorder: - """录制 actpred 格式的 episode（images/*.jpg + actions.npy） + """录制 actpred 格式的 episode（images/*.jpg + actions.npy [+ imu.npy]） 用法： rec = EpisodeRecorder("./actpred/data", action_dim=18) for frame, action in stream: - rec.record(frame, action) + rec.record(frame, action, imu_vec=imu_or_None) rec.flush() + + imu_vec：可选 MPU6050 姿态向量（roll,pitch,ax,ay,az,gx,gy,gz）。 + 只要同一 record 调用内同时传入帧/动作/姿态，三者即按帧时刻天然对齐。 + 全程 None（无 IMU 硬件）则不产出 imu.npy，保持旧格式兼容。 """ - def __init__(self, output_dir, action_dim=18): + def __init__(self, output_dir, action_dim=18, imu_dim=8): self.action_dim = action_dim + self.imu_dim = imu_dim self._frame_idx = 0 self._actions = [] + self._imus = [] + self._has_imu = False output = Path(output_dir) output.mkdir(parents=True, exist_ok=True) @@ -51,14 +59,17 @@ class EpisodeRecorder: except FileExistsError: ep_idx += 1 - log.info("recording to %s (action_dim=%d)", - self.episode_dir, action_dim) + log.info("recording to %s (action_dim=%d imu_dim=%d)", + self.episode_dir, action_dim, imu_dim) - def record(self, frame, action_vec): - """记录一帧 + 动作向量 + def record(self, frame, action_vec, imu_vec=None): + """记录一帧 + 动作向量 [+ 姿态向量] frame: np.ndarray (H,W,3) uint8 或 None（跳过图像保存） action_vec: list/np.array，长度任意，内部 pad/truncate 到 action_dim + imu_vec: 可选 list/np.array/dict，长度任意，pad/truncate 到 imu_dim； + dict 时按 (roll,pitch,ax,ay,az,gx,gy,gz) 顺序取值。 + 为 None 时该帧记零向量（混合数据集中标记"无姿态"）。 """ import numpy as np @@ -68,6 +79,19 @@ class EpisodeRecorder: vec[:n] = src[:n] self._actions.append(vec) + if imu_vec is None: + ivec = np.zeros(self.imu_dim, dtype=np.float32) + else: + self._has_imu = True + if isinstance(imu_vec, dict): + imu_vec = [imu_vec.get(k, 0.0) for k in + ("roll", "pitch", "ax", "ay", "az", "gx", "gy", "gz")] + ivec = np.zeros(self.imu_dim, dtype=np.float32) + isrc = np.asarray(imu_vec, dtype=np.float32).ravel() + m = min(len(isrc), self.imu_dim) + ivec[:m] = isrc[:m] + self._imus.append(ivec) + if frame is not None: from PIL import Image img = Image.fromarray(frame) @@ -77,11 +101,16 @@ class EpisodeRecorder: self._frame_idx += 1 def flush(self): - """保存 actions.npy，返回 episode 路径""" + """保存 actions.npy（及有姿态时的 imu.npy），返回 episode 路径""" import numpy as np if self._actions: actions = np.array(self._actions, dtype=np.float32) np.save(self.episode_dir / "actions.npy", actions) + # 仅当确实采到非零姿态才落盘 imu.npy，避免全程无 IMU 时产生误导性文件 + if self._has_imu: + imus = np.array(self._imus, dtype=np.float32) + np.save(self.episode_dir / "imu.npy", imus) + log.info("recorded imu.npy shape=%s", imus.shape) log.info("recorded %d frames to %s", self._frame_idx, self.episode_dir) return str(self.episode_dir) diff --git a/pybot/pybot/robot.py b/pybot/pybot/robot.py index 83a1f04..c49cfb6 100644 --- a/pybot/pybot/robot.py +++ b/pybot/pybot/robot.py @@ -21,6 +21,10 @@ from .parser import MatrixParser, MatrixSaver, get_mt, parse_mt log = logging.getLogger("pybot.robot") +# MPU6050 姿态向量键序（recorder.record 的 dict 取值顺序与此一致， +# mpu6050.MPU6050._latest 的键集合相同） +IMU_KEYS = ("roll", "pitch", "ax", "ay", "az", "gx", "gy", "gz") + # 已知块名 → 硬件变量工厂（Action 特殊处理；未知块 → 插件） _FACTORIES = { "I2C": I2CVar, @@ -156,11 +160,15 @@ class RobotStruct: delay_idx = self._find_delay_idx(columns) camera_var = None + imu_var = None if recorder is not None: camera_var = self._find_camera() if camera_var is None and not self.dry_run: raise ValueError( "--record 需要 @Camera 块（或使用 --dry-run 生成合成帧）") + imu_var = self._find_imu() + if imu_var is None: + log.info("--record: 无 @MPU6050_BIND 块，不产出 imu.npy") prev_vals = None @@ -180,7 +188,8 @@ class RobotStruct: recorder.record( self._grab_frame(camera_var), self._build_action_vec( - step_vals, recorder.action_dim)) + step_vals, recorder.action_dim), + imu_vec=self._read_imu(imu_var)) self._execute_step(step_vals, columns) if self.dry_run: import time @@ -189,7 +198,8 @@ class RobotStruct: if recorder is not None: recorder.record( self._grab_frame(camera_var), - self._build_action_vec(vals, recorder.action_dim)) + self._build_action_vec(vals, recorder.action_dim), + imu_vec=self._read_imu(imu_var)) self._execute_step(vals, columns) prev_vals = vals @@ -207,6 +217,9 @@ class RobotStruct: frame_len / history_len / action_dim 属性与 predict(images, action_history) -> (action_dim,) 方法。 pybot 不 import 具体推理类，保持零依赖承诺。 + 若额外声明 imu_dim 属性（R2 方案一），则以 + predict(images, action_history, imu=(imu_dim,)) 调用； + 脚本无 @MPU6050_BIND 块时 imu 为零向量。 hz: 控制频率（周期 = 1/hz 秒，须 ≥ 推理 + 执行耗时） max_steps: 步数上限，None 表示无限循环（Ctrl-C 退出） recorder: 可选 EpisodeRecorder，闭环中同步录制 (frame, action) @@ -226,6 +239,19 @@ class RobotStruct: if camera_var is None and not self.dry_run: raise ValueError("闭环需要 @Camera 块（或使用 --dry-run 生成合成帧）") + imu_var = self._find_imu() + # 鸭子类型约定：predictor 声明 imu_dim 属性即接受 imu= 关键字 + # （方案一：独立 IMU token）；未声明则保持旧两参调用，向后兼容 + use_imu = getattr(predictor, "imu_dim", None) is not None + if imu_var is None: + if use_imu: + log.warning("predictor 声明 imu_dim 但脚本无 @MPU6050_BIND " + "块，imu 将以零向量输入（等价于无姿态先验）") + else: + log.info("closed loop imu: %s → predictor %s", + imu_var.name, + "imu= 接入" if use_imu else "不接入（仅录制）") + action = self.find_mt("Action") if action is None: raise ValueError("no @Action matrix in script") @@ -248,25 +274,40 @@ class RobotStruct: overrun = 0 log.info("closed loop start: hz=%.1f F=%d H=%d action_dim=%d", hz, F, H, action_dim) - while max_steps is None or step < max_steps: - t0 = time.monotonic() - frames.extend(self._loop_recent_frames(camera_var, F)) # 1 采样 - act = predictor.predict(list(frames), np.array(history)) # 2 推理 - act = np.asarray(act, dtype=np.float32).reshape(-1) - self._execute_action_vec(act, exec_cols) # 3 执行 - history.append(act.copy()) # 4 回写 + try: + while max_steps is None or step < max_steps: + t0 = time.monotonic() + frames.extend(self._loop_recent_frames(camera_var, F)) # 1 采样 + imu = self._read_imu(imu_var) + if use_imu: # 2 推理 + act = predictor.predict(list(frames), np.array(history), + imu=self._imu_to_vec( + imu, predictor.imu_dim)) + else: + act = predictor.predict(list(frames), np.array(history)) + act = np.asarray(act, dtype=np.float32).reshape(-1) + self._execute_action_vec(act, exec_cols) # 3 执行 + history.append(act.copy()) # 4 回写 + if recorder is not None: + recorder.record( + frames[-1], + self._build_action_vec(act, recorder.action_dim), + imu_vec=imu) + step += 1 + remain = dt - (time.monotonic() - t0) # 节拍 + if remain > 0: + time.sleep(remain) + else: + overrun += 1 + log.warning("loop overrun: step %d took %.1fms > %.1fms " + "（推理+执行超过周期，实际频率已低于 %.1fHz）", + step, (dt - remain) * 1000, dt * 1000, hz) + finally: + # max_steps 正常退出与 Ctrl-C（KeyboardInterrupt）都要落盘， + # 否则闭环录制的 actions.npy/imu.npy 全部丢失 if recorder is not None: - recorder.record(frames[-1], - self._build_action_vec(act, recorder.action_dim)) - step += 1 - remain = dt - (time.monotonic() - t0) # 节拍 - if remain > 0: - time.sleep(remain) - else: - overrun += 1 - log.warning("loop overrun: step %d took %.1fms > %.1fms " - "（推理+执行超过周期，实际频率已低于 %.1fHz）", - step, (dt - remain) * 1000, dt * 1000, hz) + path = recorder.flush() + log.info("episode saved: %s", path) if overrun: log.warning("closed loop overrun %d/%d steps (%.0f%%)，" "建议降低 --hz 或检查推理后端延迟", @@ -378,6 +419,39 @@ class RobotStruct: return var return None + def _find_imu(self): + """在已创建的硬件变量中查找第一个 MPU6050Var（无则 None）""" + for var in self.M.values(): + if isinstance(var, MPU6050Var): + return var + return None + + def _read_imu(self, imu_var): + """读 MPU6050 最新姿态 dict（键序见 IMU_KEYS） + + 无硬件 / 读失败返回 None（recorder 记零向量、predictor 侧零输入）。 + latest() 冷启动返回全零 dict，语义安全无需特判。 + """ + if imu_var is None: + return None + try: + return imu_var.play({}) + except Exception: + log.exception("imu read failed") + return None + + @staticmethod + def _imu_to_vec(imu, imu_dim): + """姿态 dict → np 向量（按 IMU_KEYS 顺序 pad/truncate 到 imu_dim）""" + import numpy as np + vec = np.zeros(imu_dim, dtype=np.float32) + if imu is not None: + src = np.asarray([imu.get(k, 0.0) for k in IMU_KEYS], + dtype=np.float32) + n = min(len(src), imu_dim) + vec[:n] = src[:n] + return vec + def _build_action_vec(self, vals, action_dim): """从值列表构建动作向量，pad/truncate 到 action_dim""" import numpy as np diff --git a/pybot/tests/test_loop.py b/pybot/tests/test_loop.py index c687a0f..ba867fc 100644 --- a/pybot/tests/test_loop.py +++ b/pybot/tests/test_loop.py @@ -58,9 +58,19 @@ name alias; A angle; @ +@I2C + name path; +i2c-1 /dev/i2c-1; +@ + +@PCA9685 + name I2C address; +pca9685-0 i2c-1 0x40; +@ + @Servo name max min offset controller channel; - s1 90 - - - 0; + s1 90 - - pca9685-0 0; @ @Action @@ -281,3 +291,98 @@ def test_loop_empty_queue_falls_back_synthetic(mock_i2c): # dry_run 合成帧是 numpy 数组（HWC uint8），非 int assert all(isinstance(f, np.ndarray) for f in frames) assert frames[0].shape == (240, 320, 3) # LOOP_MT @Camera 320x240 + + +# ---------------- IMU 姿态回流（R2 方案一：独立 IMU token） ---------------- + +# 在 LOOP_MT 基础上加 @MPU6050_BIND（i2c-1 已定义，相机块之前插入） +LOOP_IMU_MT = LOOP_MT.replace( + "@Camera", + """@MPU6050_BIND + name I2C addr; +mpu6050-0 i2c-1 0x69; +@ + +@Camera""") + + +class ImuPredictor(MockPredictor): + """声明 imu_dim 的推理器桩：闭环应以 imu= 关键字调用""" + + def __init__(self, *args, imu_dim=8, **kwargs): + super().__init__(*args, **kwargs) + self.imu_dim = imu_dim + self.imus = [] # 每步收到的 imu 向量 + + def predict(self, images, action_history, imu=None): + self.imus.append(None if imu is None else np.asarray(imu).copy()) + return super().predict(images, action_history) + + +def _freeze_imu(robot, name="mpu6050-0", **vals): + """停掉采样线程并注入固定姿态（避免 10ms 采样覆盖测试值）""" + sensor = robot.M[name].sensor + sensor.stop_daemon() + sensor.wait_daemon() + sensor._latest.update(**vals) + + +def test_loop_imu_passed_to_predictor(mock_i2c): + """predictor 声明 imu_dim + 脚本含 @MPU6050_BIND → 每步 imu= 向量""" + robot = _make_robot(LOOP_IMU_MT, mock_i2c) + robot.active() + _freeze_imu(robot, roll=1.5, pitch=-2.0, gz=30.0) + pred = ImuPredictor(frame_len=3, history_len=2, action_dim=4) + robot.run_loop(pred, hz=1000.0, max_steps=3) + robot.deactive() + + assert len(pred.imus) == 3 + for imu in pred.imus: + assert imu is not None and imu.shape == (8,) + # IMU_KEYS 顺序：roll pitch ax ay az gx gy gz + assert imu[0] == 1.5 and imu[1] == -2.0 and imu[7] == 30.0 + + +def test_loop_imu_zero_without_bind(mock_i2c): + """predictor 声明 imu_dim 但脚本无 @MPU6050_BIND → imu 零向量输入""" + robot = _make_robot(LOOP_MT, mock_i2c) + robot.active() + pred = ImuPredictor(frame_len=3, history_len=2, action_dim=4) + robot.run_loop(pred, hz=1000.0, max_steps=2) + robot.deactive() + + assert len(pred.imus) == 2 + for imu in pred.imus: + assert imu is not None and imu.shape == (8,) + assert np.allclose(imu, 0.0) + + +def test_loop_imu_not_passed_to_legacy_predictor(mock_i2c): + """未声明 imu_dim 的旧推理器：保持两参调用（向后兼容）""" + robot = _make_robot(LOOP_IMU_MT, mock_i2c) + robot.active() + _freeze_imu(robot, roll=9.0) + pred = MockPredictor(frame_len=3, history_len=2, action_dim=4) + robot.run_loop(pred, hz=1000.0, max_steps=2) # 不炸即兼容 + robot.deactive() + assert len(pred.calls) == 2 + + +def test_loop_imu_recorded(mock_i2c, tmp_path): + """闭环录制：含 IMU 时 imu.npy 与 actions.npy 帧对齐（T 一致）""" + from pybot.recorder import EpisodeRecorder + robot = _make_robot(LOOP_IMU_MT, mock_i2c) + robot.active() + _freeze_imu(robot, roll=3.0, ay=0.5) + rec = EpisodeRecorder(str(tmp_path), action_dim=18) + pred = MockPredictor(frame_len=3, history_len=2, action_dim=4) + robot.run_loop(pred, hz=1000.0, max_steps=4, recorder=rec) + robot.deactive() + + ep = tmp_path / "episode_0000" + actions = np.load(ep / "actions.npy") + imus = np.load(ep / "imu.npy") + assert actions.shape == (4, 18) + assert imus.shape == (4, 8) # 与 actions 帧数对齐 + assert np.allclose(imus[:, 0], 3.0) # roll + assert np.allclose(imus[:, 3], 0.5) # ay diff --git a/pybot/tests/test_recorder.py b/pybot/tests/test_recorder.py index 81e30e5..a225da8 100644 --- a/pybot/tests/test_recorder.py +++ b/pybot/tests/test_recorder.py @@ -285,6 +285,95 @@ def test_delay_not_multiple_of_10(tmp_path, mock_i2c): robot.deactive() +# ---------------- IMU 姿态录制（R2） ---------------- + +# @Camera + @MPU6050_BIND + 2 行 @Action —— 脚本驱动录制姿态 +IMU_MT = """ +@I2C + name path; +i2c-1 /dev/i2c-1; +@ + +@MPU6050_BIND + name I2C addr; +mpu6050-0 i2c-1 0x69; +@ + +@Camera + name width height framerate capacity; + cam0 320 240 15 2; +@ + +@Action +cam0.frame; +0; +0; +@ +""" + + +def test_imu_npy_written_with_dict(tmp_path): + """record 传姿态 dict → flush 产出 imu.npy（T×8），键序 roll..gz""" + rec = EpisodeRecorder(str(tmp_path), action_dim=18) + imu = {"roll": 1.0, "pitch": 2.0, "ax": 3.0, "ay": 4.0, + "az": 5.0, "gx": 6.0, "gy": 7.0, "gz": 8.0} + rec.record(None, [0.0], imu_vec=imu) + rec.record(None, [0.0], imu_vec=imu) + rec.flush() + + imus = np.load(tmp_path / "episode_0000" / "imu.npy") + assert imus.shape == (2, 8) + assert list(imus[0]) == [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0] + + +def test_imu_mixed_none_marks_zero(tmp_path): + """混合数据集：部分帧无姿态 → 该帧记零向量，仍产出 imu.npy""" + rec = EpisodeRecorder(str(tmp_path), action_dim=18) + rec.record(None, [0.0], imu_vec=None) + rec.record(None, [0.0], imu_vec={"roll": 9.0}) + rec.flush() + + imus = np.load(tmp_path / "episode_0000" / "imu.npy") + assert imus.shape == (2, 8) + assert np.allclose(imus[0], 0.0) + assert imus[1, 0] == 9.0 + + +def test_no_imu_no_file(tmp_path): + """全程无姿态（旧数据路径）→ 不产出 imu.npy，保持旧格式兼容""" + robot = RobotStruct(CAMERA_MT, dry_run=True) + robot.active() + rec = EpisodeRecorder(str(tmp_path), action_dim=18) + robot.run(recorder=rec) + robot.deactive() + + ep_dir = tmp_path / "episode_0000" + assert (ep_dir / "actions.npy").exists() + assert not (ep_dir / "imu.npy").exists() + + +def test_run_with_imu_block(tmp_path, mock_i2c): + """脚本含 @MPU6050_BIND 时 run() 录制产出 imu.npy 且与 actions 帧对齐""" + robot = RobotStruct(IMU_MT, dry_run=True) + robot.active() + # 停采样线程注入固定姿态，避免 mock 采样值干扰断言 + sensor = robot.M["mpu6050-0"].sensor + sensor.stop_daemon() + sensor.wait_daemon() + sensor._latest.update(roll=-4.0, gz=12.0) + rec = EpisodeRecorder(str(tmp_path), action_dim=18) + robot.run(recorder=rec) + robot.deactive() + + ep_dir = tmp_path / "episode_0000" + actions = np.load(ep_dir / "actions.npy") + imus = np.load(ep_dir / "imu.npy") + assert actions.shape == (2, 18) # 2 行 @Action，无 Speaker 不分解 + assert imus.shape == (2, 8) + assert np.allclose(imus[:, 0], -4.0) + assert np.allclose(imus[:, 7], 12.0) + + def test_actpred_compatibility(tmp_path, mock_i2c): """用 actpred 的 ActionDataset 加载录制的 episode""" pytest.importorskip("torch")