feat(pybot×actpred): R2 多 MPU6050 支持——8×N 姿态向量拼接与归一化平铺
- robot：_find_imus/_read_imus 按脚本块行序收集拼接 ≤4 个 MPU6050，_pad_imu 对齐 predictor.imu_dim - recorder/imu_norm/dataset/config：imu.npy 宽度 8×N，归一化除数按传感器数平铺 - 测试：双传感器闭环拼接/pad/truncate、宽布局录制、16 维归一化加载
This commit is contained in:
qing committed
1 parent
b289e7d17c
commit
cd4d512ecd
10 files changed
+234
-56
No files matched your search
+10
-8
@@ -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 | 注意力头数 |
|
||||
|
||||
+3
-2
@@ -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)
|
||||
|
||||
+2
-2
@@ -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 的归一化姿态向量（与最新帧对齐）；无数据 → 零向量"""
|
||||
|
||||
+14
-3
@@ -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)
|
||||
@@ -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:
|
||||
|
||||
+6
-4
@@ -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/`。
|
||||
|
||||
## 开发
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
+45
-35
@@ -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"""
|
||||
|
||||
@@ -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
|
||||
@@ -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")
|
||||
|
||||
Reference in new issue
Block a user