qing/blackbird
3
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(pybot): 运行 .mt 脚本时录制 actpred 格式自监督数据集

Browse Source 
新增 EpisodeRecorder，在 run() 每行 @Action 执行前抓取摄像头帧
+ 构建动作向量，产出 actpred ActionDataset 可直接加载的
episode_XXXX/images/*.jpg + actions.npy 格式。dry_run 下生成
随机噪声帧替代真实图像。

CLI 新增 --record DIR 和 --action-dim 选项。
This commit is contained in:
qing committed 2026-08-13 16:08:47 +08:00
1 parent eadba81c26
commit b9f8abdfeb
4 files changed
+365 -3
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

pybot/pybot/main.py
+11 -1
View File
@@ -38,6 +38,11 @@ def _build_argparser():
help="(reserved) 传感器数据输出管道，对接 actpred")
p.add_argument("-o", metavar="FIFO", default=None,
help="(reserved) 动作指令输入管道，对接 actpred")
p.add_argument("--record", metavar="DIR", default=None,
help="录制 actpred 格式 episode 到 DIR（需要 @Camera 块；"
"--dry-run 下用合成帧）")
p.add_argument("--action-dim", type=int, default=18,
help="动作向量维度（默认 18，与 actpred 一致）")
return p
@@ -89,8 +94,13 @@ def main(argv=None):
source = open(args.script, encoding="utf-8") if args.script else sys.stdin
try:
robot = RobotStruct(source, dry_run=args.dry_run)
recorder = None
if args.record:
from .recorder import EpisodeRecorder
recorder = EpisodeRecorder(args.record,
action_dim=args.action_dim)
with robot:
robot.run()
robot.run(recorder=recorder)
finally:
if args.script:
source.close()
pybot/pybot/recorder.py
+87
View File
@@ -0,0 +1,87 @@
"""Episode 录制器 —— 运行 .mt 脚本时同步产出 actpred 格式的自监督数据集
产出格式（与 actpred/dataset.py 的 DataRecorder 兼容）：
output_dir/
episode_XXXX/
images/
0000.jpg, 0001.jpg, ...
actions.npy # shape (T, action_dim) float32
numpy / Pillow 为可选依赖（lazy import），仅 --record 时需要，
不影响 pybot 的零依赖运行时承诺。
"""
import logging
from pathlib import Path
log = logging.getLogger("pybot.recorder")
class EpisodeRecorder:
"""录制 actpred 格式的 episode（images/*.jpg + actions.npy）
用法：
rec = EpisodeRecorder("./actpred/data", action_dim=18)
for frame, action in stream:
rec.record(frame, action)
rec.flush()
"""
def __init__(self, output_dir, action_dim=18):
self.action_dim = action_dim
self._frame_idx = 0
self._actions = []
output = Path(output_dir)
output.mkdir(parents=True, exist_ok=True)
ep_idx = 0
for d in output.glob("episode_*"):
try:
ep_idx = max(ep_idx, int(d.name.split("_")[1]) + 1)
except (ValueError, IndexError):
pass
while True:
self.episode_dir = output / f"episode_{ep_idx:04d}"
self.images_dir = self.episode_dir / "images"
try:
self.images_dir.mkdir(parents=True, exist_ok=False)
break
except FileExistsError:
ep_idx += 1
log.info("recording to %s (action_dim=%d)",
self.episode_dir, action_dim)
def record(self, frame, action_vec):
"""记录一帧 + 动作向量
frame: np.ndarray (H,W,3) uint8 或 None（跳过图像保存）
action_vec: list/np.array，长度任意，内部 pad/truncate 到 action_dim
"""
import numpy as np
vec = np.zeros(self.action_dim, dtype=np.float32)
src = np.asarray(action_vec, dtype=np.float32).ravel()
n = min(len(src), self.action_dim)
vec[:n] = src[:n]
self._actions.append(vec)
if frame is not None:
from PIL import Image
img = Image.fromarray(frame)
img.save(self.images_dir / f"{self._frame_idx:04d}.jpg",
quality=90)
self._frame_idx += 1
def flush(self):
"""保存 actions.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)
log.info("recorded %d frames to %s",
self._frame_idx, self.episode_dir)
return str(self.episode_dir)
pybot/pybot/robot.py
+57 -2
View File
@@ -135,13 +135,31 @@ class RobotStruct:
return item
return None
def run(self):
"""逐行执行 @Action 矩阵"""
def run(self, recorder=None):
"""逐行执行 @Action 矩阵
recorder: 可选 EpisodeRecorder，启用后每行执行前录制 (frame, action)
"""
action = self.find_mt("Action")
if action is None:
raise ValueError("no @Action matrix in script")
columns = self._validate_action(action) # Q05: 执行前整表校验
camera_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 生成合成帧）")
for i in range(1, action.size):
# 录制：在执行动作前抓帧（observation → action 语义）
if recorder is not None:
action_vec = self._build_action_vec(
action, i, columns, recorder.action_dim)
frame = self._grab_frame(camera_var)
recorder.record(frame, action_vec)
row_sets = {} # {设备名: {参数名: float}}
for j, (dev, param) in enumerate(columns):
# 可能的参数别名转换（A → angle）
@@ -154,6 +172,43 @@ class RobotStruct:
for dev, params in row_sets.items():
self.M[dev].play(params)
if recorder is not None:
path = recorder.flush()
log.info("episode saved: %s", path)
def _find_camera(self):
"""在已创建的硬件变量中查找第一个 CameraVar"""
for var in self.M.values():
if isinstance(var, CameraVar):
return var
return None
def _build_action_vec(self, action, row, columns, action_dim):
"""按 @Action 列顺序构建动作向量，pad/truncate 到 action_dim
与 actpred/pretrain_mt.py 的 MTActionDataset 语义一致：
列值按顺序填入向量，不足补零，超出截断。
"""
import numpy as np
vec = np.zeros(action_dim, dtype=np.float32)
for j in range(min(len(columns), action_dim)):
vec[j] = parse_mt(action, row, j)
return vec
def _grab_frame(self, camera_var):
"""从摄像头取最新帧；dry_run 或无帧时返回随机噪声帧"""
import numpy as np
if camera_var is not None:
result = camera_var.play({})
if result is not None:
return result[1] # (timestamp, frame)
# dry_run 或队列空：随机噪声帧（用 camera 配置尺寸）
h, w = 480, 640
if camera_var is not None:
cam = camera_var.camera
h, w = cam.height, cam.width
return np.random.randint(0, 256, (h, w, 3), dtype=np.uint8)
def _validate_action(self, action):
"""校验 @Action 全部列名为 '设备.参数' 且设备已定义，返回 [(dev, param)]"""
columns = []
pybot/tests/test_recorder.py
+210
View File
@@ -0,0 +1,210 @@
"""EpisodeRecorder 测试：actpred 格式数据集录制
全部在 dry_run / mock I2C 下运行，不碰真实硬件。
"""
import sys
from pathlib import Path
import numpy as np
import pytest
from PIL import Image
from pybot.recorder import EpisodeRecorder
from pybot.robot import RobotStruct
# ---- 测试用 .mt 脚本 ----
# 仅 @Camera + @Action(cam0.frame) —— 不需要 I2C
CAMERA_MT = """
@Camera
name width height framerate capacity;
cam0 320 240 15 2;
@
@Action
cam0.frame;
0;
@
"""
# @Camera + @Servo + 多行 @Action（含占位符）—— 需要 mock I2C
SERVO_CAMERA_MT = """
@Alias
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 - - pca9685-0 0;
s2 180 - - - 1;
@
@Camera
name width height framerate capacity;
cam0 320 240 15 2;
@
@Action
s1.A s2.A;
45 0;
- 10;
- 20;
- 30;
- 0;
- 10;
- 20;
- 30;
- 0;
- 10;
@
"""
# @Servo + @Action 但无 @Camera —— 用于测试无摄像头报错
NO_CAMERA_MT = """
@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 - - pca9685-0 0;
@
@Action
s1.A;
45;
@
"""
# ---- 测试 ----
def test_episode_format(tmp_path):
"""录制后目录结构正确（images/*.jpg + actions.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.is_dir()
assert (ep_dir / "images").is_dir()
assert (ep_dir / "actions.npy").exists()
jpgs = list((ep_dir / "images").glob("*.jpg"))
assert len(jpgs) == 1 # CAMERA_MT 有 1 行 @Action
def test_action_vector_padding(tmp_path, mock_i2c):
"""2 列 .mt → 18 维向量，末尾 16 维为零"""
robot = RobotStruct(SERVO_CAMERA_MT, dry_run=True)
robot.active()
rec = EpisodeRecorder(str(tmp_path), action_dim=18)
robot.run(recorder=rec)
robot.deactive()
actions = np.load(tmp_path / "episode_0000" / "actions.npy")
assert actions.shape == (10, 18) # 10 行 @Action × 18 维
assert np.all(actions[:, 2:] == 0) # 列 2-17 全为零
def test_placeholder_resolution(tmp_path, mock_i2c):
"""`-` 占位符正确递归解析"""
robot = RobotStruct(SERVO_CAMERA_MT, dry_run=True)
robot.active()
rec = EpisodeRecorder(str(tmp_path), action_dim=18)
robot.run(recorder=rec)
robot.deactive()
actions = np.load(tmp_path / "episode_0000" / "actions.npy")
# s1.A: row0=45, row1..9 全是 '-' → 递归取上一行 → 全 45
assert np.all(actions[:, 0] == 45)
# s2.A: 0,10,20,30,0,10,20,30,0,10
expected_s2 = [0, 10, 20, 30, 0, 10, 20, 30, 0, 10]
assert list(actions[:, 1]) == expected_s2
def test_dry_run_synthetic_frame(tmp_path):
"""dry_run 下生成随机噪声帧，尺寸与 camera 配置一致"""
robot = RobotStruct(CAMERA_MT, dry_run=True)
robot.active()
rec = EpisodeRecorder(str(tmp_path), action_dim=18)
robot.run(recorder=rec)
robot.deactive()
img_path = tmp_path / "episode_0000" / "images" / "0000.jpg"
assert img_path.exists()
img = Image.open(img_path)
assert img.size == (320, 240) # (width, height)
arr = np.array(img)
assert arr.shape == (240, 320, 3)
assert not np.all(arr == 0) # 随机噪声，非全黑
def test_no_camera_error(tmp_path, mock_i2c):
"""非 dry_run 无 @Camera 时 --record 报错"""
robot = RobotStruct(NO_CAMERA_MT, dry_run=False)
robot.active()
rec = EpisodeRecorder(str(tmp_path), action_dim=18)
with pytest.raises(ValueError, match="--record 需要 @Camera"):
robot.run(recorder=rec)
robot.deactive()
def test_episode_auto_increment(tmp_path):
"""多次录制自动递增 episode 编号"""
for _ in range(2):
robot = RobotStruct(CAMERA_MT, dry_run=True)
robot.active()
rec = EpisodeRecorder(str(tmp_path), action_dim=18)
robot.run(recorder=rec)
robot.deactive()
assert (tmp_path / "episode_0000").is_dir()
assert (tmp_path / "episode_0001").is_dir()
def test_actpred_compatibility(tmp_path, mock_i2c):
"""用 actpred 的 ActionDataset 加载录制的 episode"""
pytest.importorskip("torch")
pytest.importorskip("torchvision")
robot = RobotStruct(SERVO_CAMERA_MT, dry_run=True)
robot.active()
rec = EpisodeRecorder(str(tmp_path), action_dim=18)
robot.run(recorder=rec)
robot.deactive()
actpred_dir = str(Path(__file__).resolve().parent.parent.parent / "actpred")
sys.path.insert(0, actpred_dir)
try:
from dataset import ActionDataset
ds = ActionDataset(str(tmp_path), history_len=5,
action_dim=18, frame_len=3)
# 10 行，start=max(5,2)=5，samples: t=5..9 → 5 个样本
assert len(ds) == 5
sample = ds[0]
# ActionDataset 默认 resize 到 640×480，与相机原始尺寸无关
assert sample["images"].shape == (3, 3, 480, 640) # (F, C, H, W)
assert sample["action_history"].shape == (5, 18)
assert sample["target"].shape == (18,)
finally:
sys.path.remove(actpred_dir)