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
1 parent
3e773debc9
commit
a2ba797410
22 files changed
+1236
-133
No files matched your search
@@ -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）
|
||||
|
||||
## 目录
|
||||
|
||||
|
||||
+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)
|
||||
+94
-20
@@ -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
|
||||
|
||||
+106
-1
@@ -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
|
||||
@@ -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")
|
||||
|
||||
Reference in new issue
Block a user