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

+37 -8
View File
@@ -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)