From cd4d512ecd5070565b893ab82239db717fca9886 Mon Sep 17 00:00:00 2001 From: qing <1776177662@qq.com> Date: Fri, 14 Aug 2026 09:34:39 +0800 Subject: [PATCH] =?UTF-8?q?feat(pybot=C3=97actpred):=20R2=20=E5=A4=9A=20MP?= =?UTF-8?q?U6050=20=E6=94=AF=E6=8C=81=E2=80=94=E2=80=948=C3=97N=20?= =?UTF-8?q?=E5=A7=BF=E6=80=81=E5=90=91=E9=87=8F=E6=8B=BC=E6=8E=A5=E4=B8=8E?= =?UTF-8?q?=E5=BD=92=E4=B8=80=E5=8C=96=E5=B9=B3=E9=93=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - robot：_find_imus/_read_imus 按脚本块行序收集拼接 ≤4 个 MPU6050，_pad_imu 对齐 predictor.imu_dim - recorder/imu_norm/dataset/config：imu.npy 宽度 8×N，归一化除数按传感器数平铺 - 测试：双传感器闭环拼接/pad/truncate、宽布局录制、16 维归一化加载 --- actpred/README.md | 18 ++++---- actpred/config.py | 5 ++- actpred/dataset.py | 4 +- actpred/imu_norm.py | 17 ++++++-- actpred/test_actpred.py | 30 ++++++++++++++ pybot/README.md | 10 +++-- pybot/pybot/recorder.py | 7 +++- pybot/pybot/robot.py | 80 ++++++++++++++++++++---------------- pybot/tests/test_loop.py | 72 ++++++++++++++++++++++++++++++++ pybot/tests/test_recorder.py | 47 +++++++++++++++++++++ 10 files changed, 234 insertions(+), 56 deletions(-) diff --git a/actpred/README.md b/actpred/README.md index 5d3bd6f..8efaf97 100644 --- a/actpred/README.md +++ b/actpred/README.md @@ -113,7 +113,7 @@ data/ 0001.jpg ... actions.npy # shape: (T, action_dim) float32 - imu.npy # 可选，shape: (T, imu_dim) float32（MPU6050 姿态，R2） + imu.npy # 可选，shape: (T, 8×N) float32（N 个 MPU6050 姿态拼接，R2） episode_002/ ... ``` @@ -121,12 +121,14 @@ 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`，无需统计文件。 +`imu.npy` 每传感器占 8 槽 `[roll, pitch, ax, ay, az, gx, gy, gz]`（roll/pitch +角度制，其余为 MPU6050 原始 int16），多传感器（最多 4 个）按脚本 +`@MPU6050_BIND` 行顺序拼接，由 pybot `--record` 逐帧同步落盘。**旧 episode +没有姿态数据，训练 IMU 模型需用新版 pybot 重新录制**；`--imu-dim`>0 时缺失 +`imu.npy` 的 episode 按零向量兜底（向后兼容，但无姿态先验收益）。 +姿态归一化按物理量程固定缩放（roll/pitch ÷180，原始 int16 ÷32768，8 槽模式 +按传感器数平铺），训练与推理共用 `imu_norm.py`，无需统计文件。 +录制/训练/闭环脚本的传感器行顺序必须一致，否则槽位错位。 ### 数据采集 @@ -150,7 +152,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）| +| `--imu-dim` | 0 | R2 姿态输入维度（>0 启用独立 IMU token；8×N，N=传感器数≤4）| | `--backbone` | simple | CNN 类型 (simple/mobilenet) | | `--attn-dim` | 128 | 注意力 token 维度 d_model（须能被 attn-heads 整除）| | `--attn-heads` | 4 | 注意力头数 | diff --git a/actpred/config.py b/actpred/config.py index 13d2de0..b58fcdf 100644 --- a/actpred/config.py +++ b/actpred/config.py @@ -15,7 +15,8 @@ class DataConfig: 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 格式）。 + # 8×N = N 个传感器的 roll,pitch,ax,ay,az,gx,gy,gz 按 pybot 脚本 + # @MPU6050_BIND 行顺序拼接（最多 4 个 → 32，pybot recorder 的 imu.npy 格式）。 imu_dim: int = 0 @@ -85,7 +86,7 @@ def build_argparser(description: str = "") -> argparse.ArgumentParser: 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）") + "8×N=N 个传感器拼接（≤4 个），序列变 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) diff --git a/actpred/dataset.py b/actpred/dataset.py index 82068ed..ac06bc1 100644 --- a/actpred/dataset.py +++ b/actpred/dataset.py @@ -10,7 +10,7 @@ from torch.utils.data import Dataset from torchvision import transforms from torchvision.transforms import functional as TF -from imu_norm import IMU_SCALES +from imu_norm import imu_scales def _natural_key(p): @@ -164,7 +164,7 @@ class ActionDataset(Dataset): 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] + return imus[:n] / imu_scales(self.imu_dim) def _imu_at(self, imus, t): """时刻 t 的归一化姿态向量（与最新帧对齐）；无数据 → 零向量""" diff --git a/actpred/imu_norm.py b/actpred/imu_norm.py index 992e683..59b87cb 100644 --- a/actpred/imu_norm.py +++ b/actpred/imu_norm.py @@ -9,20 +9,31 @@ 索引 0,1 /180 → ≈[-1, 1] 索引 2..7 /32768 → ≈[-1, 1] -键序与 pybot robot.IMU_KEYS / recorder 一致： +单传感器键序与 pybot robot.IMU_KEYS / recorder 一致： [roll, pitch, ax, ay, az, gx, gy, gz] +多传感器（最多 4 个）按 pybot 脚本 @MPU6050_BIND 行顺序拼接， +每传感器占 8 槽，imu_dim = 8×N，归一化除数按 8 槽模式平铺。 仅依赖 numpy，保证 hailo_infer.py（HAT 部署环境无 torch）可直接复用。 """ import numpy as np -# 各索引归一化除数（imu_dim=8 全量；更短的向量取前缀） +# 各索引归一化除数（单传感器 8 槽；多传感器用 imu_scales() 平铺） IMU_SCALES = np.array([180.0, 180.0, 32768.0, 32768.0, 32768.0, 32768.0, 32768.0, 32768.0], dtype=np.float32) +def imu_scales(imu_dim: int) -> np.ndarray: + """归一化除数向量：8 槽模式按传感器数平铺（imu_dim = 8×N，N≤4） + + imu_dim ≤ 8 时等价 IMU_SCALES 前缀（roll/pitch 子集仍可用）。 + """ + reps = (imu_dim + len(IMU_SCALES) - 1) // len(IMU_SCALES) + return np.tile(IMU_SCALES, reps)[:imu_dim] + + def normalize_imu(imu, imu_dim: int = 8) -> np.ndarray: """原始姿态向量 → 归一化向量（pad/truncate 到 imu_dim） @@ -34,4 +45,4 @@ def normalize_imu(imu, imu_dim: int = 8) -> np.ndarray: src = np.asarray(imu, dtype=np.float32).ravel() n = min(len(src), imu_dim) vec[:n] = src[:n] - return vec / IMU_SCALES[:imu_dim] + return vec / imu_scales(imu_dim) diff --git a/actpred/test_actpred.py b/actpred/test_actpred.py index 333a62e..352913c 100644 --- a/actpred/test_actpred.py +++ b/actpred/test_actpred.py @@ -1836,6 +1836,20 @@ class TestR2ImuNorm(unittest.TestCase): self.assertEqual(out.shape, (2,)) np.testing.assert_allclose(out, [1.0, -1.0], atol=1e-6) + def test_multi_sensor_tiled_scales(self): + """多传感器（8×N，N≤4）：归一化除数按 8 槽模式平铺""" + from imu_norm import normalize_imu, imu_scales, IMU_SCALES + scales = imu_scales(32) + self.assertEqual(scales.shape, (32,)) + for k in range(4): + np.testing.assert_array_equal(scales[k * 8:(k + 1) * 8], + IMU_SCALES) + # 单传感器 8 槽模式 [±180, ∓180, 32768×6] 拼接两组 → 归一化重复 + raw = ([180.0, -180.0] + [32768.0] * 6) * 2 + out = normalize_imu(raw, 16) + np.testing.assert_allclose( + out, ([1.0, -1.0] + [1.0] * 6) * 2, atol=1e-6) + class TestR2ImuModel(unittest.TestCase): """R2 方案一模型：独立 IMU token 注入（序列 F+H+2）""" @@ -1965,6 +1979,22 @@ class TestR2ImuDataset(unittest.TestCase): 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_multi_sensor_16wide_normalized(self): + """双传感器 imu.npy（16 列）：两组 8 槽分别按物理量程归一化""" + with tempfile.TemporaryDirectory() as tmpdir: + imus = [[180.0, 0, 0, 0, 0, 0, 0, 0, # sensor0 roll + 0, 0, 0, 0, 0, 0, 0, 32768.0]] * 10 # sensor1 gz + 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=16) + imu = ds[0]["imu"].numpy() + self.assertEqual(imu.shape, (16,)) + np.testing.assert_allclose(imu[0], 1.0, atol=1e-6) + np.testing.assert_allclose(imu[15], 1.0, atol=1e-6) + np.testing.assert_array_equal(imu[1:15], np.zeros(14, np.float32)) + def test_short_imu_aligned_with_zero_pad(self): """imu.npy 行数少于帧数 → 尾部补零对齐""" with tempfile.TemporaryDirectory() as tmpdir: diff --git a/pybot/README.md b/pybot/README.md index bcad940..916c50c 100644 --- a/pybot/README.md +++ b/pybot/README.md @@ -32,7 +32,8 @@ cat scripts/std.mt | ./_run.sh # 与 bot4cj 相同的管道用法 ``` `--record` 时若脚本含 `@MPU6050_BIND` 块，逐帧同步落盘 `imu.npy` -（T×8：roll,pitch,ax,ay,az,gx,gy,gz，与图像/动作按帧时刻天然对齐）； +（每传感器 8 槽 roll,pitch,ax,ay,az,gx,gy,gz；多个传感器（≤4）按块行顺序 +拼接，shape (T, 8×N)，与图像/动作按帧时刻天然对齐）； 无 IMU 硬件则保持旧格式（仅 images + actions.npy）。 `-i` / `-o` 参数位保留（reserved），未来对接跨进程 FIFO。 @@ -55,8 +56,9 @@ cat scripts/std.mt | ./_run.sh # 与 bot4cj 相同的管道用法 部署时若日志出现 `loop overrun` 告警，说明 `hz` 超过推理+执行能力，应降低 `hz` - **姿态回流（R2）**：脚本含 `@MPU6050_BIND` 且模型带 imu 输入 （`predictor.imu_dim` 非 None）时，每步读取 MPU6050 最新姿态随帧/历史 - 一同传给推理器（原始量纲，归一化在推理端完成）；模型无 imu 输入时 - 不接入推理（若带 `--record` 仍会读取落盘 imu.npy） + 一同传给推理器（原始量纲，归一化在推理端完成）；多个传感器（≤4）按 + 块行顺序拼接为 8×N 向量，再 pad/truncate 到 `predictor.imu_dim`； + 模型无 imu 输入时不接入推理（若带 `--record` 仍会读取落盘 imu.npy） ## 目录 @@ -68,7 +70,7 @@ scripts -> ../bot4cj/scripts # .mt 脚本与仓颉版共享 tests/ # pytest 单元测试（Mock 总线，无需硬件） ``` -插件查找顺序：`$PYBOT_MODULES_PATH`（冒号分隔）→ 包内置 `modules/` → `./modules/`。 +插件查找顺序：`$PYBOT_MODULES_PATH`（`os.pathsep` 分隔：Linux `:` / Windows `;`）→ 包内置 `modules/` → `./modules/`。 ## 开发 diff --git a/pybot/pybot/recorder.py b/pybot/pybot/recorder.py index fe21d0a..698b920 100644 --- a/pybot/pybot/recorder.py +++ b/pybot/pybot/recorder.py @@ -27,7 +27,9 @@ class EpisodeRecorder: rec.record(frame, action, imu_vec=imu_or_None) rec.flush() - imu_vec：可选 MPU6050 姿态向量（roll,pitch,ax,ay,az,gx,gy,gz）。 + imu_vec：可选 MPU6050 姿态向量（单传感器 roll,pitch,ax,ay,az,gx,gy,gz； + 多传感器按脚本 @MPU6050_BIND 行顺序拼接为 8×N，最多 4 个）。 + RobotStruct 会在录制前把 imu_dim 同步为 8×N（N=传感器数）。 只要同一 record 调用内同时传入帧/动作/姿态，三者即按帧时刻天然对齐。 全程 None（无 IMU 硬件）则不产出 imu.npy，保持旧格式兼容。 """ @@ -68,7 +70,8 @@ class EpisodeRecorder: 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) 顺序取值。 + dict 时按 (roll,pitch,ax,ay,az,gx,gy,gz) 顺序取值（单传感器； + 多传感器请传 RobotStruct._read_imus 拼好的 8×N 向量）。 为 None 时该帧记零向量（混合数据集中标记"无姿态"）。 """ import numpy as np diff --git a/pybot/pybot/robot.py b/pybot/pybot/robot.py index c49cfb6..8dbcd21 100644 --- a/pybot/pybot/robot.py +++ b/pybot/pybot/robot.py @@ -160,15 +160,18 @@ class RobotStruct: delay_idx = self._find_delay_idx(columns) camera_var = None - imu_var = None + imu_vars = [] 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: + imu_vars = self._find_imus() + if not imu_vars: log.info("--record: 无 @MPU6050_BIND 块，不产出 imu.npy") + else: + # 槽位宽度 8×N（按块顺序拼接），录制前同步给 recorder + recorder.imu_dim = len(IMU_KEYS) * len(imu_vars) prev_vals = None @@ -189,7 +192,7 @@ class RobotStruct: self._grab_frame(camera_var), self._build_action_vec( step_vals, recorder.action_dim), - imu_vec=self._read_imu(imu_var)) + imu_vec=self._read_imus(imu_vars)) self._execute_step(step_vals, columns) if self.dry_run: import time @@ -199,7 +202,7 @@ class RobotStruct: recorder.record( self._grab_frame(camera_var), self._build_action_vec(vals, recorder.action_dim), - imu_vec=self._read_imu(imu_var)) + imu_vec=self._read_imus(imu_vars)) self._execute_step(vals, columns) prev_vals = vals @@ -219,7 +222,8 @@ class RobotStruct: pybot 不 import 具体推理类，保持零依赖承诺。 若额外声明 imu_dim 属性（R2 方案一），则以 predict(images, action_history, imu=(imu_dim,)) 调用； - 脚本无 @MPU6050_BIND 块时 imu 为零向量。 + 脚本无 @MPU6050_BIND 块时 imu 为零向量。多个 MPU6050 + （最多 4 个）按脚本块行顺序拼接为 8×N 原始向量。 hz: 控制频率（周期 = 1/hz 秒，须 ≥ 推理 + 执行耗时） max_steps: 步数上限，None 表示无限循环（Ctrl-C 退出） recorder: 可选 EpisodeRecorder，闭环中同步录制 (frame, action) @@ -239,17 +243,19 @@ class RobotStruct: if camera_var is None and not self.dry_run: raise ValueError("闭环需要 @Camera 块（或使用 --dry-run 生成合成帧）") - imu_var = self._find_imu() + imu_vars = self._find_imus() # 鸭子类型约定：predictor 声明 imu_dim 属性即接受 imu= 关键字 # （方案一：独立 IMU token）；未声明则保持旧两参调用，向后兼容 use_imu = getattr(predictor, "imu_dim", None) is not None - if imu_var is None: + if not imu_vars: if use_imu: log.warning("predictor 声明 imu_dim 但脚本无 @MPU6050_BIND " "块，imu 将以零向量输入（等价于无姿态先验）") else: + if recorder is not None: + recorder.imu_dim = len(IMU_KEYS) * len(imu_vars) log.info("closed loop imu: %s → predictor %s", - imu_var.name, + ",".join(v.name for v in imu_vars), "imu= 接入" if use_imu else "不接入（仅录制）") action = self.find_mt("Action") @@ -278,10 +284,10 @@ class RobotStruct: 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) + imu = self._read_imus(imu_vars) if use_imu: # 2 推理 act = predictor.predict(list(frames), np.array(history), - imu=self._imu_to_vec( + imu=self._pad_imu( imu, predictor.imu_dim)) else: act = predictor.predict(list(frames), np.array(history)) @@ -419,38 +425,42 @@ 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 _find_imus(self): + """按脚本块顺序收集全部 MPU6050Var（无则空列表，最多 4 个）""" + return [v for v in self.M.values() if isinstance(v, MPU6050Var)] - def _read_imu(self, imu_var): - """读 MPU6050 最新姿态 dict（键序见 IMU_KEYS） + def _read_imus(self, imu_vars): + """读全部 MPU6050 最新姿态，按块顺序拼接为 (8×N,) 原始向量 - 无硬件 / 读失败返回 None（recorder 记零向量、predictor 侧零输入）。 - latest() 冷启动返回全零 dict，语义安全无需特判。 + imu_vars 为空返回 None（无传感器：recorder 不产出 imu.npy）； + 单个传感器读失败其 8 槽记零并 log。latest() 冷启动返回全零 + dict，语义安全无需特判。 """ - if imu_var is None: - return None - try: - return imu_var.play({}) - except Exception: - log.exception("imu read failed") + if not imu_vars: return None + import numpy as np + nk = len(IMU_KEYS) + out = np.zeros(nk * len(imu_vars), dtype=np.float32) + for k, var in enumerate(imu_vars): + try: + latest = var.play({}) + except Exception: + log.exception("imu read failed: %s", var.name) + continue + out[k * nk:(k + 1) * nk] = [latest.get(key, 0.0) + for key in IMU_KEYS] + return out @staticmethod - def _imu_to_vec(imu, imu_dim): - """姿态 dict → np 向量（按 IMU_KEYS 顺序 pad/truncate 到 imu_dim）""" + def _pad_imu(vec, imu_dim): + """拼接姿态向量 → pad/truncate 到 predictor 的 imu_dim（None → 全零）""" 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) + out = np.zeros(imu_dim, dtype=np.float32) + if vec is not None: + src = np.asarray(vec, dtype=np.float32).ravel() n = min(len(src), imu_dim) - vec[:n] = src[:n] - return vec + out[:n] = src[:n] + return out def _build_action_vec(self, vals, action_dim): """从值列表构建动作向量，pad/truncate 到 action_dim""" diff --git a/pybot/tests/test_loop.py b/pybot/tests/test_loop.py index ba867fc..51de120 100644 --- a/pybot/tests/test_loop.py +++ b/pybot/tests/test_loop.py @@ -386,3 +386,75 @@ def test_loop_imu_recorded(mock_i2c, tmp_path): assert imus.shape == (4, 8) # 与 actions 帧数对齐 assert np.allclose(imus[:, 0], 3.0) # roll assert np.allclose(imus[:, 3], 0.5) # ay + + +# ---------------- 多 MPU6050（最多 4 个，8×N 槽位拼接） ---------------- + +# 双传感器脚本（槽位顺序对齐 std.mt 前两行：i2c-1 上 0x69 / 0x68） +LOOP_IMU2_MT = LOOP_MT.replace( + "@Camera", + """@MPU6050_BIND + name I2C addr; +mpu6050-0 i2c-1 0x69; +mpu6050-1 i2c-1 0x68; +@ + +@Camera""") + + +def test_loop_imu_multi_sensor_concat(mock_i2c): + """多 MPU6050：按块行顺序拼接 8×N 向量喂 predictor（槽位不串）""" + robot = _make_robot(LOOP_IMU2_MT, mock_i2c) + robot.active() + _freeze_imu(robot, "mpu6050-0", roll=1.0, gz=7.0) + _freeze_imu(robot, "mpu6050-1", roll=2.0, pitch=-3.0) + pred = ImuPredictor(frame_len=3, history_len=2, action_dim=4, imu_dim=16) + 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 == (16,) + assert imu[0] == 1.0 and imu[7] == 7.0 # sensor0: roll, gz + assert imu[8] == 2.0 and imu[9] == -3.0 # sensor1: roll, pitch + + +def test_loop_imu_multi_sensor_pad_truncate(mock_i2c): + """传感器数与 predictor.imu_dim 不匹配：pad 零 / 截断，槽位不漂移""" + robot = _make_robot(LOOP_IMU2_MT, mock_i2c) + robot.active() + _freeze_imu(robot, "mpu6050-0", roll=1.0) + _freeze_imu(robot, "mpu6050-1", roll=2.0) + + pred24 = ImuPredictor(frame_len=3, history_len=2, action_dim=4, + imu_dim=24) + robot.run_loop(pred24, hz=1000.0, max_steps=1) # 16 → 24 补零 + imu = pred24.imus[0] + assert imu.shape == (24,) + assert imu[0] == 1.0 and imu[8] == 2.0 + assert np.allclose(imu[16:], 0.0) + + pred8 = ImuPredictor(frame_len=3, history_len=2, action_dim=4, + imu_dim=8) + robot.run_loop(pred8, hz=1000.0, max_steps=1) # 16 → 8 截断（仅 sensor0） + imu = pred8.imus[0] + assert imu.shape == (8,) and imu[0] == 1.0 + robot.deactive() + + +def test_loop_imu_multi_recorded(mock_i2c, tmp_path): + """多传感器闭环录制：imu.npy 宽度 8×N 且槽位按块行顺序""" + from pybot.recorder import EpisodeRecorder + robot = _make_robot(LOOP_IMU2_MT, mock_i2c) + robot.active() + _freeze_imu(robot, "mpu6050-0", roll=1.0) + _freeze_imu(robot, "mpu6050-1", pitch=4.0) + 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=3, recorder=rec) + robot.deactive() + + imus = np.load(tmp_path / "episode_0000" / "imu.npy") + assert imus.shape == (3, 16) + assert np.allclose(imus[:, 0], 1.0) # sensor0 roll + assert np.allclose(imus[:, 9], 4.0) # sensor1 pitch diff --git a/pybot/tests/test_recorder.py b/pybot/tests/test_recorder.py index a225da8..dd6be13 100644 --- a/pybot/tests/test_recorder.py +++ b/pybot/tests/test_recorder.py @@ -374,6 +374,53 @@ def test_run_with_imu_block(tmp_path, mock_i2c): assert np.allclose(imus[:, 7], 12.0) +# 双传感器脚本（槽位顺序对齐 std.mt 前两行：i2c-1 上 0x69 / 0x68） +IMU2_MT = """ +@I2C + name path; +i2c-1 /dev/i2c-1; +@ + +@MPU6050_BIND + name I2C addr; +mpu6050-0 i2c-1 0x69; +mpu6050-1 i2c-1 0x68; +@ + +@Camera + name width height framerate capacity; + cam0 320 240 15 2; +@ + +@Action +cam0.frame; +0; +0; +@ +""" + + +def test_run_multi_imu_wide_layout(tmp_path, mock_i2c): + """多 MPU6050：run() 录制 imu.npy 宽度 8×N，槽位按块行顺序""" + robot = RobotStruct(IMU2_MT, dry_run=True) + robot.active() + for name, vals in (("mpu6050-0", {"roll": -4.0}), + ("mpu6050-1", {"roll": 6.0, "gz": 12.0})): + sensor = robot.M[name].sensor + sensor.stop_daemon() + sensor.wait_daemon() + sensor._latest.update(**vals) + rec = EpisodeRecorder(str(tmp_path), action_dim=18) + robot.run(recorder=rec) + robot.deactive() + + imus = np.load(tmp_path / "episode_0000" / "imu.npy") + assert imus.shape == (2, 16) + assert np.allclose(imus[:, 0], -4.0) # sensor0 roll + assert np.allclose(imus[:, 8], 6.0) # sensor1 roll（槽位 8 起） + assert np.allclose(imus[:, 15], 12.0) # sensor1 gz + + def test_actpred_compatibility(tmp_path, mock_i2c): """用 actpred 的 ActionDataset 加载录制的 episode""" pytest.importorskip("torch")