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:
1 parent
3e773debc9
commit
a2ba797410
22 files changed
+1236
-133
No files matched your search
+37
-8
@@ -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)
|
||||
Reference in new issue
Block a user