diff --git a/pybot/pybot/main.py b/pybot/pybot/main.py index 5cd324a..b9a53ed 100644 --- a/pybot/pybot/main.py +++ b/pybot/pybot/main.py @@ -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() diff --git a/pybot/pybot/recorder.py b/pybot/pybot/recorder.py new file mode 100644 index 0000000..d701ed9 --- /dev/null +++ b/pybot/pybot/recorder.py @@ -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) diff --git a/pybot/pybot/robot.py b/pybot/pybot/robot.py index d94852a..2ba2b82 100644 --- a/pybot/pybot/robot.py +++ b/pybot/pybot/robot.py @@ -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 = [] diff --git a/pybot/tests/test_recorder.py b/pybot/tests/test_recorder.py new file mode 100644 index 0000000..fdd8296 --- /dev/null +++ b/pybot/tests/test_recorder.py @@ -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)