| From 78183844fcc4b1a090d244585659e6a7d2e5b9b8 Mon Sep 17 00:00:00 2001 |
| From: CodeChildCZJ <CodeChildCZJ@users.noreply.github.com> |
| Date: Tue, 4 Aug 2026 23:30:18 +1000 |
| Subject: [PATCH] Add SO-101 pi0.5 fine-tuning and deployment release |
|
|
| |
| HF_MODEL_CARD.md | 160 +++++++++ |
| NOTICE | 1 + |
| SO101_PI05_HANDOFF.md | 410 +++++++++++++++++++++++ |
| examples/so101/README.md | 158 +++++++++ |
| scripts/compute_so101_norm_stats.py | 74 ++++ |
| scripts/eval_so101_checkpoint.py | 74 ++++ |
| scripts/plot_so101_training.py | 131 ++++++++ |
| scripts/run_so101_training.sh | 39 +++ |
| scripts/so101_preflight.py | 99 ++++++ |
| scripts/train.py | 16 +- |
| src/openpi/policies/so101_policy.py | 78 +++++ |
| src/openpi/policies/so101_policy_test.py | 35 ++ |
| src/openpi/training/config.py | 113 +++++++ |
| src/openpi/training/data_loader.py | 23 +- |
| src/openpi/training/data_loader_test.py | 27 ++ |
| 15 files changed, 1430 insertions(+), 8 deletions(-) |
| create mode 100644 HF_MODEL_CARD.md |
| create mode 100644 NOTICE |
| create mode 100644 SO101_PI05_HANDOFF.md |
| create mode 100644 examples/so101/README.md |
| create mode 100644 scripts/compute_so101_norm_stats.py |
| create mode 100644 scripts/eval_so101_checkpoint.py |
| create mode 100644 scripts/plot_so101_training.py |
| create mode 100755 scripts/run_so101_training.sh |
| create mode 100644 scripts/so101_preflight.py |
| create mode 100644 src/openpi/policies/so101_policy.py |
| create mode 100644 src/openpi/policies/so101_policy_test.py |
|
|
| |
| new file mode 100644 |
| |
| |
| |
| @@ -0,0 +1,160 @@ |
| +--- |
| +license: other |
| +language: |
| + - en |
| +library_name: openpi |
| +tags: |
| + - robotics |
| + - vision-language-action |
| + - openpi |
| + - pi0.5 |
| + - so101 |
| + - lerobot |
| + - orbax |
| +--- |
| + |
| +# pi0.5 SO-101 Erythromycin-on-Tea |
| + |
| +OpenPI pi0.5 fully fine-tuned for a dual-camera SO-101 follower on one tabletop manipulation task: |
| + |
| +> Pick up the red erythromycin ointment box and place it on top of the green Rizhao tea tin. |
| + |
| +Author: **CodeChild** |
| + |
| +This repository contains inference-only OpenPI/Orbax artifacts. It is not a standard Transformers checkpoint and cannot be loaded with `transformers.AutoModel.from_pretrained()`. |
| + |
| +## Model details |
| + |
| +| Field | Value | |
| +| --- | --- | |
| +| Base checkpoint | `gs://openpi-assets/checkpoints/pi05_base/params` | |
| +| Fine-tuning | Full parameters, 8000 optimizer updates | |
| +| Inference weights | EMA parameters, decay 0.99 | |
| +| Cameras | Fixed RGB + wrist RGB, 640x480 at collection time | |
| +| State/action | 6-D calibrated SO-101 position space | |
| +| Action horizon | 50 steps at 30 Hz | |
| +| Output | `(50, 6)` absolute SO-101 position targets | |
| +| Code release tag | `so101-pi05-erythromycin-v1` | |
| +| W&B | <https://wandb.ai/99087192-zhejiang-university/openpi/runs/xgk0h74f> | |
| + |
| +The full deployment and safety handoff is in [`SO101_PI05_HANDOFF.md`](./SO101_PI05_HANDOFF.md). Read it before connecting this policy to motors. |
| + |
| +## Repository contents |
| + |
| +```text |
| +params/ # EMA inference parameters, about 12 GiB |
| +assets/ # clean-train normalization statistics |
| +_CHECKPOINT_METADATA # original Orbax checkpoint metadata |
| +code/so101-pi05-erythromycin-v1.patch |
| +SO101_PI05_HANDOFF.md |
| +LICENSE_OPENPI.txt |
| +LICENSE_GEMMA.txt |
| +NOTICE |
| +``` |
| + |
| +The optimizer state is intentionally not published. This repository is suitable for inference, not direct training resume. |
| + |
| +## Code setup |
| + |
| +The SO-101 adapter was developed from OpenPI commit: |
| + |
| +```text |
| +15a9616a00943ada6c20a0f158e3adb39df2ccac |
| +``` |
| + |
| +The release commit is tagged locally as `so101-pi05-erythromycin-v1`. Because the source checkout only has the upstream Physical Intelligence remote, the exact code delta is also included in this model repository: |
| + |
| +```bash |
| +git clone https://github.com/Physical-Intelligence/openpi.git |
| +cd openpi |
| +git checkout 15a9616a00943ada6c20a0f158e3adb39df2ccac |
| +git apply /path/to/so101-pi05-erythromycin-v1.patch |
| +GIT_LFS_SKIP_SMUDGE=1 UV_LINK_MODE=copy uv sync |
| +``` |
| + |
| +The data config expects the portable dataset package next to the OpenPI checkout when running data-dependent scripts. Policy inference only needs this repository's `params/` and `assets/`. |
| + |
| +## Download and serve |
| + |
| +```python |
| +from huggingface_hub import snapshot_download |
| + |
| +checkpoint_dir = snapshot_download("CodeChild/pi05-so101-erythromycin-tea") |
| +print(checkpoint_dir) |
| +``` |
| + |
| +From the patched OpenPI checkout: |
| + |
| +```bash |
| +CUDA_VISIBLE_DEVICES=<GPU_ID> .venv/bin/python scripts/serve_policy.py \ |
| + policy:checkpoint \ |
| + --policy.config pi05_so101_erythromycin \ |
| + --policy.dir <HF_SNAPSHOT_DIR> |
| +``` |
| + |
| +Policy input: |
| + |
| +```python |
| +observation = { |
| + "observation/state": state_float32_6, |
| + "observation/fixed_image": fixed_rgb_uint8_hwc, |
| + "observation/wrist_image": wrist_rgb_uint8_hwc, |
| + "prompt": "Pick up the red erythromycin ointment box and place it on top of the green Rizhao tea tin.", |
| +} |
| +``` |
| + |
| +The server returns `result["actions"]` with shape `(50, 6)`. |
| + |
| +## Critical action semantics |
| + |
| +During training, dimensions 0-4 are represented relative to the same current state and dimension 5 remains absolute: |
| + |
| +```text |
| +delta[t, 0:5] = absolute_target[t, 0:5] - current_state[0:5] |
| +delta[t, 5] = absolute_target[t, 5] |
| +``` |
| + |
| +This is not a step-to-step increment. OpenPI applies the inverse transform before returning actions, so the policy server output is already absolute. A robot client must **not** take a cumulative sum and must **not** add the current state again. |
| + |
| +The trajectory was collected at 30 Hz. Fifty predicted steps correspond to approximately 1.67 seconds of control ticks. For an initial supervised robot test, execute a short prefix and replan; the accompanying handoff recommends starting with 5 steps at 30 Hz. This is a deployment recommendation, not a robot-validated hyperparameter. |
| + |
| +## Training data |
| + |
| +The source dataset has 90 episodes and two synchronized camera streams. It is not redistributed in this model repository. The authoritative split was `splits/split_manifest.json`: |
| + |
| +| Split | Episodes | Frames | Use | |
| +| --- | ---: | ---: | --- | |
| +| `clean_train` | 67 | 15070 | Training and normalization statistics | |
| +| `clean_val` | 18 | 3972 | Offline validation only | |
| +| `recovery` | 5 | 1515 | Excluded | |
| + |
| +The default `train: 0:90` field in the source LeRobot metadata was not used because it would leak validation and recovery episodes into training. |
| + |
| +OpenPI's standard training augmentation was active: crop/rotation/color augmentation for the fixed camera and color augmentation for the wrist camera. Evaluation and inference use no random augmentation. |
| + |
| +## Offline evaluation |
| + |
| +| Model | Split | Samples | Flow-matching loss | |
| +| --- | --- | ---: | ---: | |
| +| Original pi0.5 base | `clean_val` | 3972 | 0.04753249 | |
| +| Fine-tuned checkpoint | `clean_train` | 15068 | 0.00407172 | |
| +| Fine-tuned checkpoint | `clean_val` | 3972 | 0.01467515 | |
| + |
| +The held-out validation loss is 69.126% lower than the base checkpoint under this evaluation. The validation/train ratio is 3.604, indicating a generalization gap. There is no independent test split. |
| + |
| +Flow-matching loss is not a robot task-success metric. No closed-loop real-robot success rate has been measured for this checkpoint yet. |
| + |
| +## Intended use and limitations |
| + |
| +- Intended for research and supervised evaluation on the stated SO-101 task. |
| +- Requires the same joint order, direction, zero points, gripper calibration, camera assignment, RGB convention and 30 Hz timing used during collection. |
| +- Before motor execution, validate finite values and shape, enforce hardware joint/gripper limits, maximum target deltas, velocity/workspace limits, timeouts and an emergency stop. |
| +- Start with motors-off shadow inference, then low-speed supervised closed-loop tests. |
| +- The model was trained on one task with a small dataset and may fail under new layouts, lighting, camera movement, object appearance or calibration drift. |
| +- Do not infer safety or reliability from the offline flow-matching loss. |
| + |
| +## Licenses |
| + |
| +OpenPI code is provided under Apache-2.0; see `LICENSE_OPENPI.txt`. |
| + |
| +The model is derived from pi0.5, which includes Gemma components. Gemma use and redistribution are subject to the Gemma Terms of Use in `LICENSE_GEMMA.txt`, and the required notice is provided in `NOTICE`. For this reason the Hugging Face metadata uses `license: other` rather than describing the complete artifact as Apache-2.0 only. |
| |
| new file mode 100644 |
| |
| |
| |
| @@ -0,0 +1 @@ |
| +Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms |
| |
| new file mode 100644 |
| |
| |
| |
| @@ -0,0 +1,410 @@ |
| +# SO-101 pi0.5 训练与真机部署交接 |
| + |
| +更新时间：2026-08-04 |
| +状态：8000-step 微调和离线评估已完成；**尚未进行真机闭环验证**。 |
| + |
| +本文是当前 checkpoint 的权威交接说明。文中将训练时已经确定的事实和首次上真机的建议明确分开。任何真机客户端都必须先读完“动作语义”和“首次上机流程”，尤其不能把服务端返回值再次当作 delta 累加。 |
| + |
| +## 1. 一页摘要 |
| + |
| +| 项目 | 当前结果 | |
| +| --- | --- | |
| +| 任务 | `Pick up the red erythromycin ointment box and place it on top of the green Rizhao tea tin.` | |
| +| 模型 | OpenPI pi0.5 base，全参数微调 | |
| +| 最终 checkpoint | `checkpoints/pi05_so101_erythromycin/pi05_so101_stage1_wandb/7999` | |
| +| 训练量 | 8000 次 optimizer update，batch size 8 | |
| +| 推理参数 | EMA 参数，`ema_decay=0.99` | |
| +| 训练数据 | `clean_train`：67 episodes / 15070 frames | |
| +| 验证数据 | `clean_val`：18 episodes / 3972 frames | |
| +| 未使用数据 | `recovery`：5 episodes / 1515 frames | |
| +| 相机 | 固定相机 `fixed` + 腕部相机 `wrist`，RGB 640x480 @ 30 FPS | |
| +| 状态/动作 | 6 维 SO-101 校准后位置空间 | |
| +| 动作表示 | 前 5 维在模型内部为相对当前 state 的 delta；第 6 维夹爪为 absolute | |
| +| 服务端输出 | `(50, 6)`，已经还原为 SO-101 **绝对目标** | |
| +| 数据/控制频率 | 30 Hz，约 33.33 ms/step | |
| +| 预测 horizon | 50 steps；50 个控制 tick 约 1.667 s | |
| +| 首次真机建议 | 30 Hz 执行，每次只执行前 5 steps 后重规划；这是建议，不是已验证参数 | |
| +| W&B | <https://wandb.ai/99087192-zhejiang-university/openpi/runs/xgk0h74f> | |
| + |
| +最重要的三点： |
| + |
| +1. policy server 返回的动作已经是 absolute，机器人端不要 `cumsum`，也不要再次加当前 state。 |
| +2. 训练数据是 30 Hz。不能把 50 个目标按 10 Hz 执行，否则同一段轨迹会被拉长约 3 倍。 |
| +3. 50 是预测长度，不代表首次测试应开环执行完整 50 steps。先执行 5 steps 后重新观测和规划。 |
| + |
| +## 2. Artifact 和代码来源 |
| + |
| +仓库： |
| + |
| +```text |
| +/home2/czj/AutoResearch/real_machine/so_arm101/openpi_so101_pi05 |
| +``` |
| + |
| +最终 checkpoint： |
| + |
| +```text |
| +/home2/czj/AutoResearch/real_machine/so_arm101/openpi_so101_pi05/checkpoints/ |
| + pi05_so101_erythromycin/pi05_so101_stage1_wandb/7999 |
| +``` |
| + |
| +目录内容和磁盘占用： |
| + |
| +| 目录 | 用途 | 大小 | |
| +| --- | --- | ---: | |
| +| `params/` | 推理必需；保存的是 EMA 参数 | 约 12 GiB | |
| +| `assets/` | 推理必需；包含仅由 `clean_train` 计算的 norm stats | 约 16 KiB | |
| +| `train_state/` | 仅继续训练需要；包含 optimizer state 等 | 约 31 GiB | |
| +| 全部 | 可推理并可续训 | 约 42 GiB | |
| + |
| +`7999` 是从 0 开始计数的最终保存 step，对应已经完成 8000 次更新。checkpoint 保存逻辑在存在 EMA 时会将 EMA 参数放进 `params/`，所以 policy server 加载的就是 EMA 推理权重。 |
| + |
| +当前 OpenPI 上游基准 commit： |
| + |
| +```text |
| +15a9616a00943ada6c20a0f158e3adb39df2ccac |
| +``` |
| + |
| +**可移植性提醒：** SO-101 policy、数据 split 读取和训练脚本目前包含本地尚未提交的适配代码。仅上传权重并指向干净的上游 commit，不能保证能识别 `pi05_so101_erythromycin` 配置。上传 Hugging Face 前，应把这些修改形成一个可检出的 Git commit/tag，或随模型仓库提供完整 patch，至少覆盖： |
| + |
| +```text |
| +src/openpi/policies/so101_policy.py |
| +src/openpi/training/config.py |
| +src/openpi/training/data_loader.py |
| +scripts/compute_so101_norm_stats.py |
| +scripts/eval_so101_checkpoint.py |
| +scripts/run_so101_training.sh |
| +scripts/so101_preflight.py |
| +``` |
| + |
| +## 3. 数据、划分和任务 |
| + |
| +数据包： |
| + |
| +```text |
| +/home2/czj/AutoResearch/real_machine/so_arm101/ |
| + so101_erythromycin_on_tea_grid90_v2_portable |
| +``` |
| + |
| +总数据为 90 episodes / 20557 frames / 685.233 s，SO-101 follower，单一语言任务，两路原始 AV1 视频。训练必须以此文件为准： |
| + |
| +```text |
| +splits/split_manifest.json |
| +``` |
| + |
| +实际划分： |
| + |
| +| split | episodes | frames | 是否用于本次训练 | |
| +| --- | ---: | ---: | --- | |
| +| `clean_train` | 67 | 15070 | 是；也只用它计算 norm stats | |
| +| `clean_val` | 18 | 3972 | 只用于离线验证 | |
| +| `recovery` | 5 | 1515 | 否；保留给后续 recovery 实验 | |
| + |
| +**绝不能直接使用 `dataset/meta/info.json` 中的 `train: 0:90`。** 那是原始录制范围，不是实验划分；使用它会把 validation 和 recovery 都混入训练。 |
| + |
| +`clean_val` 是六个完整 held-out layout settings：`setting_01`、`setting_09`、`setting_12`、`setting_17`、`setting_21`、`setting_29`。当前没有独立 test split，`clean_val` 仍可能参与 checkpoint 选择，因此不能把它称为最终无偏测试集。 |
| + |
| +任务 prompt 必须保持完全一致： |
| + |
| +```text |
| +Pick up the red erythromycin ointment box and place it on top of the green Rizhao tea tin. |
| +``` |
| + |
| +## 4. 输入 schema 和相机处理 |
| + |
| +policy server 的单次输入： |
| + |
| +```python |
| +observation = { |
| + "observation/state": state_float32_6, |
| + "observation/fixed_image": fixed_rgb_uint8_hwc, |
| + "observation/wrist_image": wrist_rgb_uint8_hwc, |
| + "prompt": "Pick up the red erythromycin ointment box and place it on top of the green Rizhao tea tin.", |
| +} |
| +``` |
| + |
| +约束： |
| + |
| +- `state_float32_6.shape == (6,)`。 |
| +- 图像使用 RGB，而不是 OpenCV 默认的 BGR。 |
| +- 最稳妥的图像格式是 `uint8`、HWC、范围 `[0, 255]`；分辨率按采集配置为 640x480。 |
| +- `fixed` 必须对应训练时的固定相机视角，`wrist` 必须对应腕部视角，不得互换、镜像或旋转。 |
| +- OpenPI 内部使用等比例 `resize_with_pad` 变成 224x224，不应在客户端做会改变宽高比的强制拉伸。 |
| +- pi0.5 需要三个图像槽；SO-101 适配会把第三个 `right_wrist_0_rgb` 塞零并设置 `mask=false`。客户端不需要发送第三路图像。 |
| + |
| +### 训练中实际使用的图像增强 |
| + |
| +本次没有添加 SO-101 专用的自定义增强，但 OpenPI 的标准 JAX 训练预处理在 `train=True` 时确实启用： |
| + |
| +- 固定相机：95% 随机裁剪后 resize、随机旋转 `[-5 deg, +5 deg]`、颜色扰动； |
| +- 腕部相机：颜色扰动，不做上述随机裁剪和旋转； |
| +- 颜色扰动参数：brightness `0.3`、contrast `0.4`、saturation `0.5`； |
| +- 离线评估和 policy inference 使用 `train=False`，不做随机增强。 |
| + |
| +因此复现实验时不能将本次训练描述成“完全无图像增强”。 |
| + |
| +## 5. 状态、动作顺序和校准 |
| + |
| +state 和 action 的六维顺序完全相同： |
| + |
| +| index | LeRobot 名称 | 含义 | |
| +| ---: | --- | --- | |
| +| 0 | `shoulder_pan.pos` | shoulder pan | |
| +| 1 | `shoulder_lift.pos` | shoulder lift | |
| +| 2 | `elbow_flex.pos` | elbow flex | |
| +| 3 | `wrist_flex.pos` | wrist flex | |
| +| 4 | `wrist_roll.pos` | wrist roll | |
| +| 5 | `gripper.pos` | gripper | |
| + |
| +这些值不是电机原始 encoder tick，而是 LeRobot SO-101 校准后的 position space：前五维按关节角度使用，夹爪使用线性校准空间，通常映射到约 `[0, 100]`。部署机器必须使用与采集机器一致的关节顺序、方向、零点、角度定义和夹爪标定。 |
| + |
| +不要只因为数值“看起来在范围内”就假设两台机器人标定一致。首次连接时应逐关节读取 state，与已知安全姿态和采集数据样本对照；发现符号、offset 或夹爪开合方向不一致时禁止下发模型动作。 |
| + |
| +## 6. Delta 的精确定义 |
| + |
| +训练配置中的 mask 是： |
| + |
| +```python |
| +delta_action_mask = (True, True, True, True, True, False) |
| +``` |
| + |
| +设发起推理时当前状态为 `s`，数据中的第 `t` 个绝对目标为 `a[t]`。训练输入模型前执行： |
| + |
| +```text |
| +d[t, 0:5] = a[t, 0:5] - s[0:5] |
| +d[t, 5] = a[t, 5] |
| +``` |
| + |
| +这里 50 个未来目标全部减去同一个当前状态 `s`。它不是： |
| + |
| +```text |
| +a[t] - a[t-1] |
| +``` |
| + |
| +也不是每步速度。因此绝对不能沿时间轴对模型结果做 cumulative sum。 |
| + |
| +推理时 OpenPI 的输出 transform 会执行相反操作： |
| + |
| +```text |
| +a_hat[t, 0:5] = d_hat[t, 0:5] + s[0:5] |
| +a_hat[t, 5] = d_hat[t, 5] |
| +``` |
| + |
| +随后移除模型内部 padding，只返回前 6 维。所以外部收到： |
| + |
| +```python |
| +result["actions"].shape == (50, 6) |
| +``` |
| + |
| +`result["actions"]` 已经是机器人校准空间中的 absolute position targets。机器人客户端只需验证、限幅并按顺序下发；不要再次加 state，不要 `cumsum`。 |
| + |
| +## 7. 30 Hz 和 50-step action chunk |
| + |
| +训练 action chunk 按数据集的 30 Hz 采样： |
| + |
| +```text |
| +control period = 1 / 30 s = 33.33 ms |
| +action horizon = 50 steps |
| +50 control ticks = 1.667 s |
| +last sampled target offset = 49 / 30 s = 1.633 s |
| +``` |
| + |
| +### 已确定的事实 |
| + |
| +- 模型每次预测 50 个顺序目标。 |
| +- 每个相邻目标的训练时间间隔是 33.33 ms。 |
| +- policy 不会自动决定机器人端执行其中多少个动作。 |
| +- 50-step prediction horizon 不等于必须 50-step open-loop execution。 |
| + |
| +### 首次真机建议，尚未验证 |
| + |
| +- 机器人目标下发循环保持 30 Hz。 |
| +- 初始设置执行前缀 `K=5`，即每次预测后执行约 167 ms，再用新图像和新 state 重规划。 |
| +- shadow 和低速测试稳定后，可根据实测推理延迟尝试 `K=5..10`；不要一开始开环执行完整 50 steps。 |
| +- 若采用异步推理，记录 observation timestamp、response timestamp、p50/p95 round-trip latency 和实际控制 jitter。执行前缀至少应覆盖正常的 p95 推理时间；可用 `ceil(p95_latency_seconds * 30)` 估算最低 K，再留少量调度余量。 |
| +- 如果覆盖 p95 延迟所需的 K 已经大于 10，首次上机不应简单增大开环窗口来掩盖问题；应先降低推理/网络延迟，或采用可安全 hold 的同步流程。 |
| +- 丢弃明显过期、乱序或基于旧 observation 的 response。切换到新 chunk 时记录其 observation 序号，避免旧结果覆盖新结果。 |
| +- 不要按 10 Hz 直接执行这 50 个动作；那会把约 1.67 s 的训练轨迹拉成约 5 s。 |
| + |
| +`K=5` 是保守起点，不是已经通过真机成功率验证的超参数。最终 K 应由推理延迟、安全性和真实 rollout 数据共同决定。 |
| + |
| +## 8. Policy server 启动 |
| + |
| +在 145 上先检查 GPU，再选择空闲设备： |
| + |
| +```bash |
| +nvidia-smi |
| +``` |
| + |
| +从仓库根目录启动： |
| + |
| +```bash |
| +CUDA_VISIBLE_DEVICES=<GPU_ID> .venv/bin/python scripts/serve_policy.py \ |
| + policy:checkpoint \ |
| + --policy.config pi05_so101_erythromycin \ |
| + --policy.dir checkpoints/pi05_so101_erythromycin/pi05_so101_stage1_wandb/7999 |
| +``` |
| + |
| +部署前至少做一次 motors-off 请求，并断言： |
| + |
| +```python |
| +actions = np.asarray(result["actions"]) |
| +assert actions.shape == (50, 6) |
| +assert np.isfinite(actions).all() |
| +``` |
| + |
| +若从 Hugging Face 下载的是 inference-only snapshot，snapshot 根目录应直接包含 `params/` 和 `assets/`，此时 `--policy.dir` 指向 snapshot 根目录即可。 |
| + |
| +## 9. 真机安全门和首次 rollout 流程 |
| + |
| +以下保护应在机器人客户端实现，不能依赖模型自己学会： |
| + |
| +1. **形状和数值检查：** 必须是 `(50, 6)` 且全部 finite；出现 NaN、Inf、缺帧或超时立即 hold/stop。 |
| +2. **硬件绝对限位：** 按该台 SO-101 的校准和物理限制 clamp 五个关节；夹爪限制在有效线性标定区间。训练集 q01/q99 不是硬件安全限位。 |
| +3. **单步变化限制：** 对每个相邻 absolute target 应用 `max_relative_target` 或等价检查，拒绝突跳；同时限制速度和必要的加速度。 |
| +4. **工作空间约束：** 禁止桌面穿透、自碰撞、相机线缆拉扯和进入人员区域。 |
| +5. **时序保护：** 30 Hz monotonic scheduler；监测 missed deadline、queue underrun、过期 chunk 和相机/state 时间差。 |
| +6. **失联行为：** policy server、相机或网络超时后进入定义好的 safe hold/stop，而不是继续无限执行旧 chunk。 |
| +7. **现场保护：** 低速/低力矩起步、急停可触达、单人专职观察、首次 rollout 不无人值守。 |
| + |
| +推荐按以下顺序放行： |
| + |
| +### 阶段 A：离线接口检查 |
| + |
| +- 用一条已录制 observation 请求模型；保存输入图像、state 和 `(50, 6)` 输出。 |
| +- 检查 RGB/BGR、相机顺序、图像方向和 prompt。 |
| +- 画出六维 action chunk，并比较 `actions[0, :5] - state[:5]`；确认没有明显跳变。 |
| + |
| +### 阶段 B：shadow inference，电机不执行 |
| + |
| +- 真机按 30 Hz 采集相机和 state，但仅打印/记录模型动作。 |
| +- 统计各维 min/max、最大单步变化、相对当前 state 的最大偏差和推理 p95 latency。 |
| +- 人工确认夹爪开合方向、所有关节符号和目标姿态合理。 |
| + |
| +### 阶段 C：低速闭环 |
| + |
| +- 从安全 home pose 开始，桌面清空危险障碍物。 |
| +- `K=5`，30 Hz，开启全部限位和急停，人工全程监护。 |
| +- 先做短时运动并主动停止，再做完整任务。 |
| + |
| +### 阶段 D：正式评估 |
| + |
| +- 固定初始姿态、物体布局、光照和相机位置，并记录每次实验配置。 |
| +- 每个 layout 做多次独立 rollout；建议至少 10 次，报告成功数和总次数，而不只展示最好视频。 |
| +- 同时记录抓取成功、最终放置稳定、掉落、碰撞、人工干预、超时和完成时间。 |
| + |
| +## 10. 离线评估结果 |
| + |
| +评估脚本使用 batch size 4、固定随机 seed、无随机图像增强。结果如下： |
| + |
| +| 模型 | split | 实际样本数 | flow-matching loss | |
| +| --- | --- | ---: | ---: | |
| +| 原始 pi0.5 base | `clean_val` | 3972 | 0.04753249 | |
| +| 最终微调模型 | `clean_train` | 15068 | 0.00407172 | |
| +| 最终微调模型 | `clean_val` | 3972 | 0.01467515 | |
| + |
| +说明： |
| + |
| +- `clean_train` 原有 15070 frames；评估按完整 batch 统计，因此使用 15068 个样本。 |
| +- 相对原始 pi0.5 base，微调模型的 validation loss 降低约 69.126%。 |
| +- 微调模型的 val/train loss 比约为 3.604，存在明显泛化差距。 |
| +- validation 是完整 held-out settings，不与 train 重叠；但没有独立 test split。 |
| +- flow-matching loss 衡量训练目标，不等于抓取成功率、放置成功率或安全性。 |
| +- 在完成受控真机 rollout 前，不能声称该模型已经可以可靠完成任务。 |
| + |
| +复现最终模型离线验证： |
| + |
| +```bash |
| +CUDA_VISIBLE_DEVICES=<GPU_ID> .venv/bin/python scripts/eval_so101_checkpoint.py \ |
| + --checkpoint-dir checkpoints/pi05_so101_erythromycin/pi05_so101_stage1_wandb/7999 \ |
| + --split-name clean_val |
| +``` |
| + |
| +## 11. 训练配置记录 |
| + |
| +W&B run 中记录的实际运行参数优先于源码中的默认值： |
| + |
| +| 参数 | 值 | |
| +| --- | --- | |
| +| base weights | `gs://openpi-assets/checkpoints/pi05_base/params` | |
| +| fine-tuning | full parameters，`freeze_filter=Nothing()` | |
| +| updates | 8000 | |
| +| batch size | 8 | |
| +| seed | 42 | |
| +| precision | bfloat16 | |
| +| optimizer | Adam，`b1=0.9`，`b2=0.95`，gradient clip 1.0 | |
| +| LR | warmup 1000，peak `2.5e-5`，配置的 decay horizon 30000 | |
| +| EMA | 0.99 | |
| +| checkpoint interval | 1000；manager 只保留最新 regular checkpoint | |
| +| W&B run ID | `xgk0h74f` | |
| + |
| +训练只运行到 8000 updates，因此 30000-step LR decay schedule 没有完整走完。源码默认训练步数后来仍可显示 30000，不要据此误称本 checkpoint 已训练 30000 steps。 |
| + |
| +## 12. Hugging Face 发布建议 |
| + |
| +### 推荐默认：inference-only，约 12 GiB |
| + |
| +模型 repo 根目录至少包含： |
| + |
| +```text |
| +params/ |
| +assets/ |
| +_CHECKPOINT_METADATA # 建议保留原始 checkpoint 元数据 |
| +README.md # Hugging Face model card |
| +SO101_PI05_HANDOFF.md |
| +LICENSE_GEMMA.txt |
| +NOTICE |
| +LICENSE_OPENPI.txt # 或等价保留 OpenPI Apache-2.0 文本 |
| +``` |
| + |
| +还应提供一个可复现的 OpenPI code commit/tag 或 patch。该 checkpoint 是 Orbax/OpenPI 格式，不是可直接用 `transformers.AutoModel.from_pretrained()` 加载的标准 Transformers 权重；model card 必须明确要求通过 OpenPI policy loader 使用。 |
| + |
| +### 可选：resumable，约 42 GiB |
| + |
| +只有在确实需要继续训练时才额外上传： |
| + |
| +```text |
| +train_state/ |
| +``` |
| + |
| +上传 `train_state` 会增加约 31 GiB，并且仍需完全匹配的代码、配置和 optimizer 定义。默认不建议为了推理上传它。 |
| + |
| +### 许可证和数据权限 |
| + |
| +- OpenPI 代码为 Apache-2.0。 |
| +- 权重由包含 Gemma 的 pi0.5 派生，发布时不能把整个模型简单标成纯 Apache-2.0。 |
| +- HF model card 建议使用 `license: other`，正文同时说明 OpenPI Apache-2.0 和 Gemma Terms。 |
| +- 分发时保留完整 `LICENSE_GEMMA.txt`，并在 `NOTICE` 中包含： |
| + |
| +```text |
| +Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms |
| +``` |
| + |
| +- 数据包当前明确写着“未分配数据许可证”。除非数据拥有方确认可再分发，否则不要顺手把原始双相机视频或整个数据集上传到公开 HF repo。 |
| +- 发布前由发布者确认 Gemma 条款和数据授权；本文不构成法律意见。 |
| + |
| +### 上传前需要确认的信息 |
| + |
| +1. HF namespace：使用个人账号 `CodeChild`，还是某个 organization。 |
| +2. model repo 名，例如 `pi05-so101-erythromycin-tea`。 |
| +3. repo 是 `public` 还是 `private`。 |
| +4. 上传 inference-only（约 12 GiB，推荐）还是 resumable（约 42 GiB）。 |
| +5. model card 上的作者、机构、联系方式和希望展示的模型名称。 |
| +6. 是否有允许公开链接的数据集 repo；若没有，model card 只描述数据，不上传原始数据。 |
| +7. 是否确认按 Gemma Terms 分发派生权重并保留要求的 notice。 |
| +8. SO-101 适配代码是发布为 Git commit/tag，还是随模型提供 patch。 |
| + |
| +本机已检测到可用的 Hugging Face 登录，当前身份为 `CodeChild`。不要在聊天中粘贴 access token；如果要换账号，应在本机运行 `hf auth login`，并使用具有目标 namespace write 权限的 token。 |
| + |
| +## 13. 发布和上机前最终 checklist |
| + |
| +- [ ] HF repo 中同时有 `params/` 和 `assets/`，norm stats 没有遗漏。 |
| +- [ ] code commit/tag 或 patch 可以构造 `pi05_so101_erythromycin` policy。 |
| +- [ ] model card 明确 Orbax/OpenPI 加载方法、delta 语义、30 Hz 和 50-step horizon。 |
| +- [ ] model card 不把 flow-matching loss 写成真机成功率。 |
| +- [ ] Gemma license 和 `NOTICE` 完整，数据发布权限已确认。 |
| +- [ ] 真机校准、关节顺序、方向和 gripper range 与采集端一致。 |
| +- [ ] fixed/wrist 图像为 RGB、方向正确、时间同步。 |
| +- [ ] 客户端确认输出是 absolute，没有二次加 state 或 `cumsum`。 |
| +- [ ] 30 Hz 调度、`K=5` 初始前缀、超时 hold、限位和急停均已实现。 |
| +- [ ] 完成 motors-off shadow inference 后才进入低速真机测试。 |
| +- [ ] 真机结果按多次 rollout 的成功/失败和安全事件完整记录。 |
| |
| new file mode 100644 |
| |
| |
| |
| @@ -0,0 +1,158 @@ |
| +# SO-101 pi0.5 fine-tuning |
| + |
| +This setup fine-tunes `pi05_base` on the local dual-camera SO-101 dataset for the task: |
| + |
| +> Pick up the red erythromycin ointment box and place it on top of the green Rizhao tea tin. |
| + |
| +The implementation intentionally reads `splits/split_manifest.json`. It never uses the `train: 0:90` entry in |
| +`dataset/meta/info.json` as an experimental split. |
| + |
| +## Configuration |
| + |
| +The `pi05_so101_erythromycin` config uses: |
| + |
| +- full pi0.5 fine-tuning from `gs://openpi-assets/checkpoints/pi05_base/params`; |
| +- the 67 `clean_train` episodes and no validation or recovery episodes; |
| +- fixed and wrist RGB cameras, with the unused right-wrist image slot masked; |
| +- six native SO-101 state/action dimensions padded to the model's 32 dimensions; |
| +- delta actions for the first five arm joints and an absolute sixth gripper action; |
| +- a 50-step action horizon at the dataset's 30 Hz rate; |
| +- batch size 8, a 30,000-step configured ceiling, and checkpoints every 1,000 steps; |
| +- fresh quantile normalization statistics computed only from `clean_train`; |
| +- no SO-101-specific custom augmentation; OpenPI's standard training-time crop/rotation/color augmentation remains enabled. |
| + |
| +The local dataset path and exact episode list are validated whenever the config is created. The data loader uses PyAV |
| +because the original videos are AV1. |
| + |
| +## Environment |
| + |
| +From the repository root: |
| + |
| +```bash |
| +GIT_LFS_SKIP_SMUDGE=1 UV_LINK_MODE=copy uv sync |
| +``` |
| + |
| +The checked setup uses the repository-local `.venv`; it does not use the separate `/home2/czj/openpi` environment. |
| + |
| +Local paths on `ZjuServer145` are: |
| + |
| +```text |
| +repository: /home2/czj/AutoResearch/real_machine/so_arm101/openpi_so101_pi05 |
| +environment: /home2/czj/AutoResearch/real_machine/so_arm101/openpi_so101_pi05/.venv |
| +dataset: /home2/czj/AutoResearch/real_machine/so_arm101/so101_erythromycin_on_tea_grid90_v2_portable |
| +``` |
| + |
| +The released stage-1 artifact completed 8,000 updates and is stored at: |
| + |
| +```text |
| +checkpoints/pi05_so101_erythromycin/pi05_so101_stage1_wandb/7999 |
| +``` |
| + |
| +Its W&B run is <https://wandb.ai/99087192-zhejiang-university/openpi/runs/xgk0h74f>. |
| + |
| +## Preflight and normalization |
| + |
| +```bash |
| +.venv/bin/python scripts/so101_preflight.py |
| +.venv/bin/python scripts/compute_so101_norm_stats.py |
| +``` |
| + |
| +The normalization script reads Parquet directly so it does not waste time decoding images. It still constructs the |
| +same clamped 50-step action chunks used by LeRobot and applies the same five-joint delta-action transform. |
| + |
| +Expected source counts are: |
| + |
| +```text |
| +clean_train: 67 episodes / 15070 frames |
| +clean_val: 18 episodes / 3972 frames |
| +recovery: 5 episodes / 1515 frames |
| +``` |
| + |
| +## One-step GPU smoke test |
| + |
| +Use a nearly empty 96 GB GPU for full fine-tuning. By default, the runner refuses to start with less than 90,000 MiB |
| +free; override `SO101_MIN_FREE_MEMORY_MIB` only when intentionally using a smaller batch. This command initializes the base checkpoint, |
| +compiles the training step, and performs one update without writing a large checkpoint: |
| + |
| +```bash |
| +SO101_GPU_ID=0 scripts/run_so101_training.sh pi05_so101_smoke \ |
| + --num-train-steps 1 \ |
| + --no-save-final-checkpoint |
| +``` |
| + |
| +This exact smoke test passed on an RTX PRO 6000 Blackwell with batch size 8. Its first compiled update reported loss |
| +`0.0470` and gradient norm `0.4075`; JAX reserved about 88.2 GiB under the configured 90% allocator limit. No checkpoint |
| +was written. |
| + |
| +## Training |
| + |
| +Start with an 8,000-step stage instead of committing immediately to all 30,000 steps: |
| + |
| +```bash |
| +SO101_GPU_ID=0 scripts/run_so101_training.sh pi05_so101_stage1_wandb \ |
| + --num-train-steps 8000 \ |
| + --wandb-enabled |
| +``` |
| + |
| +To enable Weights & Biases explicitly, add `--wandb-enabled`. To continue the same experiment to 30,000 total steps: |
| + |
| +```bash |
| +SO101_GPU_ID=0 scripts/run_so101_training.sh pi05_so101_stage1 \ |
| + --resume \ |
| + --num-train-steps 30000 |
| +``` |
| + |
| +The local loss curve and its CSV source can be generated from the training tmux session with: |
| + |
| +```bash |
| +.venv/bin/python scripts/plot_so101_training.py --refresh-seconds 30 |
| +``` |
| + |
| +The checkpoint manager keeps only the latest regular checkpoint by default to avoid filling the local filesystem. |
| +Preserve an inference checkpoint separately before resuming if it is a candidate selected by validation or robot |
| +rollouts. |
| + |
| +## Offline validation |
| + |
| +Evaluate a saved checkpoint on the held-out `clean_val` episodes using the training normalization statistics: |
| + |
| +```bash |
| +CUDA_VISIBLE_DEVICES=0 .venv/bin/python scripts/eval_so101_checkpoint.py \ |
| + --checkpoint-dir checkpoints/pi05_so101_erythromycin/pi05_so101_stage1_wandb/7999 |
| +``` |
| + |
| +Use the actual final step directory printed by training. Add `--max-batches 100` for a faster diagnostic. The reported |
| +flow-matching loss is useful for comparing checkpoints, but it is not a substitute for closed-loop robot success rate. |
| + |
| +## Policy server input |
| + |
| +Start the server with a selected checkpoint: |
| + |
| +```bash |
| +CUDA_VISIBLE_DEVICES=0 .venv/bin/python scripts/serve_policy.py \ |
| + policy:checkpoint \ |
| + --policy.config pi05_so101_erythromycin \ |
| + --policy.dir checkpoints/pi05_so101_erythromycin/pi05_so101_stage1_wandb/7999 |
| +``` |
| + |
| +The robot client must send: |
| + |
| +```python |
| +observation = { |
| + "observation/state": state_float32_6, |
| + "observation/fixed_image": fixed_rgb_uint8_hwc, |
| + "observation/wrist_image": wrist_rgb_uint8_hwc, |
| + "prompt": "Pick up the red erythromycin ointment box and place it on top of the green Rizhao tea tin.", |
| +} |
| +``` |
| + |
| +The server returns a `(50, 6)` `actions` array in the original SO-101 absolute control space. A real robot loop should |
| +execute only a short prefix, obtain a new observation, and replan. Before any unattended rollout, enforce joint and |
| +gripper limits, workspace limits, an emergency stop, a low initial speed, and human supervision. |
| + |
| +## Evaluation protocol |
| + |
| +`clean_val` contains six held-out layout combinations and is used for checkpoint selection. It is not an independent |
| +test set. The five recovery episodes remain excluded from this clean baseline. Final model quality should be measured |
| +with repeated real-robot rollouts per held-out layout, including task success, stable placement, collision, drop, and |
| +human-intervention rates. |
| |
| new file mode 100644 |
| |
| |
| |
| @@ -0,0 +1,74 @@ |
| +"""Compute SO-101 normalization stats from Parquet without decoding videos.""" |
| + |
| +import numpy as np |
| +import tyro |
| + |
| +import openpi.shared.normalize as _normalize |
| +import openpi.training.config as _config |
| +import openpi.training.data_loader as _data_loader |
| + |
| + |
| +def _unwrap_dataset(dataset): |
| + while isinstance(dataset, _data_loader.TransformedDataset): |
| + dataset = dataset._dataset # noqa: SLF001 |
| + return dataset |
| + |
| + |
| +def _stats(values: np.ndarray) -> _normalize.NormStats: |
| + values = values.reshape(-1, values.shape[-1]).astype(np.float64) |
| + return _normalize.NormStats( |
| + mean=np.mean(values, axis=0), |
| + std=np.std(values, axis=0), |
| + q01=np.quantile(values, 0.01, axis=0), |
| + q99=np.quantile(values, 0.99, axis=0), |
| + ) |
| + |
| + |
| +def main(config_name: str = "pi05_so101_erythromycin") -> None: |
| + config = _config.get_config(config_name) |
| + if not isinstance(config.data, _config.LeRobotSO101DataConfig): |
| + raise TypeError(f"Config {config_name!r} does not use LeRobotSO101DataConfig") |
| + if config.data.split_name != "clean_train": |
| + raise ValueError("Normalization statistics must be computed from clean_train") |
| + |
| + data_config = config.data.create(config.assets_dirs, config.model) |
| + dataset = _data_loader.create_torch_dataset(data_config, config.model.action_horizon, config.model) |
| + raw_dataset = _unwrap_dataset(dataset) |
| + if raw_dataset.num_episodes != 67 or raw_dataset.num_frames != 15_070: |
| + raise ValueError( |
| + f"Expected clean_train with 67 episodes/15070 frames, got " |
| + f"{raw_dataset.num_episodes}/{raw_dataset.num_frames}" |
| + ) |
| + |
| + states = np.stack(raw_dataset.hf_dataset["observation.state"]).astype(np.float32) |
| + actions = np.stack(raw_dataset.hf_dataset["action"]).astype(np.float32) |
| + episode_indices = np.asarray(raw_dataset.hf_dataset["episode_index"], dtype=np.int64) |
| + |
| + episode_ends = np.empty(len(raw_dataset), dtype=np.int64) |
| + for episode_id in data_config.episodes or (): |
| + locations = np.flatnonzero(episode_indices == episode_id) |
| + if locations.size == 0: |
| + raise ValueError(f"Episode {episode_id} is missing from the selected LeRobot dataset") |
| + episode_ends[locations] = locations[-1] + 1 |
| + |
| + offsets = np.arange(config.model.action_horizon, dtype=np.int64) |
| + query_indices = np.arange(len(raw_dataset), dtype=np.int64)[:, None] + offsets[None, :] |
| + query_indices = np.minimum(query_indices, episode_ends[:, None] - 1) |
| + action_chunks = actions[query_indices] |
| + if config.data.use_delta_joint_actions: |
| + action_chunks[..., :5] -= states[:, None, :5] |
| + |
| + norm_stats = { |
| + "state": _stats(states), |
| + "actions": _stats(action_chunks), |
| + } |
| + output_path = config.assets_dirs / data_config.repo_id |
| + _normalize.save(output_path, norm_stats) |
| + |
| + print(f"source_split=clean_train episodes={raw_dataset.num_episodes} frames={raw_dataset.num_frames}") |
| + print(f"state_samples={states.shape[0]} action_samples={action_chunks.shape[0] * action_chunks.shape[1]}") |
| + print(f"wrote={output_path / 'norm_stats.json'}") |
| + |
| + |
| +if __name__ == "__main__": |
| + tyro.cli(main) |
| |
| new file mode 100644 |
| |
| |
| |
| @@ -0,0 +1,74 @@ |
| +"""Measure deterministic pi0.5 flow-matching loss on the held-out SO-101 split.""" |
| + |
| +import dataclasses |
| +import pathlib |
| + |
| +import jax |
| +import jax.numpy as jnp |
| +import numpy as np |
| +import tyro |
| + |
| +from openpi.models import model as _model |
| +from openpi.shared import nnx_utils |
| +from openpi.training import config as _config |
| +from openpi.training import data_loader as _data_loader |
| + |
| + |
| +def _join_checkpoint_path(checkpoint_dir: str, child: str) -> str: |
| + if checkpoint_dir.startswith("gs://"): |
| + return f"{checkpoint_dir.rstrip('/')}/{child}" |
| + return str(pathlib.Path(checkpoint_dir).expanduser().resolve() / child) |
| + |
| + |
| +def main( |
| + checkpoint_dir: str, |
| + config_name: str = "pi05_so101_erythromycin", |
| + split_name: str = "clean_val", |
| + batch_size: int = 4, |
| + max_batches: int | None = None, |
| + seed: int = 0, |
| +) -> None: |
| + config = _config.get_config(config_name) |
| + if not isinstance(config.data, _config.LeRobotSO101DataConfig): |
| + raise TypeError(f"Config {config_name!r} does not use LeRobotSO101DataConfig") |
| + if split_name not in {"clean_val", "clean_train"}: |
| + raise ValueError("Offline checkpoint evaluation supports clean_train or clean_val only") |
| + |
| + data_factory = dataclasses.replace(config.data, split_name=split_name) |
| + eval_config = dataclasses.replace( |
| + config, |
| + data=data_factory, |
| + batch_size=batch_size, |
| + num_workers=min(config.num_workers, 4), |
| + ) |
| + data_config = data_factory.create(eval_config.assets_dirs, eval_config.model) |
| + raw_dataset = _data_loader.create_torch_dataset(data_config, eval_config.model.action_horizon, eval_config.model) |
| + num_batches = len(raw_dataset) // batch_size |
| + if max_batches is not None: |
| + num_batches = min(num_batches, max_batches) |
| + if num_batches < 1: |
| + raise ValueError("No complete validation batches are available") |
| + |
| + loader = _data_loader.create_data_loader( |
| + eval_config, |
| + shuffle=False, |
| + num_batches=num_batches, |
| + ) |
| + params_path = _join_checkpoint_path(checkpoint_dir, "params") |
| + model = eval_config.model.load(_model.restore_params(params_path, dtype=jnp.bfloat16)) |
| + model.eval() |
| + compute_loss = nnx_utils.module_jit(model.compute_loss) |
| + |
| + rng = jax.random.key(seed) |
| + losses = [] |
| + for batch_index, (observation, actions) in enumerate(loader): |
| + batch_rng = jax.random.fold_in(rng, batch_index) |
| + loss = compute_loss(batch_rng, observation, actions) |
| + losses.append(float(np.asarray(jnp.mean(loss)))) |
| + |
| + print(f"split={split_name} batches={len(losses)} samples={len(losses) * batch_size}") |
| + print(f"flow_matching_loss={np.mean(losses):.8f}") |
| + |
| + |
| +if __name__ == "__main__": |
| + tyro.cli(main) |
| |
| new file mode 100644 |
| |
| |
| |
| @@ -0,0 +1,131 @@ |
| +"""Plot local SO-101 training metrics captured from a tmux pane.""" |
| + |
| +import argparse |
| +import csv |
| +import pathlib |
| +import re |
| +import subprocess |
| +import time |
| + |
| +import matplotlib.pyplot as plt |
| + |
| +STEP_PATTERN = re.compile(r"^Step (?P<step>\d+): (?P<metrics>.+)$", re.MULTILINE) |
| +METRIC_PATTERN = re.compile(r"(?P<name>[a-z_]+)=(?P<value>[-+0-9.eE]+)") |
| + |
| + |
| +def _capture_metrics(session: str) -> list[dict[str, float]]: |
| + result = subprocess.run( |
| + ["tmux", "capture-pane", "-p", "-t", session, "-S", "-100000"], |
| + check=True, |
| + capture_output=True, |
| + text=True, |
| + ) |
| + metrics_by_step: dict[int, dict[str, float]] = {} |
| + for match in STEP_PATTERN.finditer(result.stdout): |
| + step = int(match.group("step")) |
| + metrics = { |
| + item.group("name"): float(item.group("value")) for item in METRIC_PATTERN.finditer(match.group("metrics")) |
| + } |
| + if "loss" in metrics: |
| + metrics_by_step[step] = {"step": float(step), **metrics} |
| + return [metrics_by_step[step] for step in sorted(metrics_by_step)] |
| + |
| + |
| +def _write_csv(metrics: list[dict[str, float]], output_path: pathlib.Path) -> None: |
| + fieldnames = ["step", "loss", "grad_norm", "param_norm"] |
| + with output_path.open("w", encoding="utf-8", newline="") as file: |
| + writer = csv.DictWriter(file, fieldnames=fieldnames) |
| + writer.writeheader() |
| + for row in metrics: |
| + writer.writerow({name: int(row[name]) if name == "step" else row.get(name, "") for name in fieldnames}) |
| + |
| + |
| +def _read_csv(input_path: pathlib.Path) -> list[dict[str, float]]: |
| + if not input_path.is_file(): |
| + return [] |
| + with input_path.open(encoding="utf-8", newline="") as file: |
| + return [{name: float(value) for name, value in row.items() if value} for row in csv.DictReader(file)] |
| + |
| + |
| +def _merge_metrics(*metric_groups: list[dict[str, float]]) -> list[dict[str, float]]: |
| + metrics_by_step = {int(item["step"]): item for group in metric_groups for item in group} |
| + return [metrics_by_step[step] for step in sorted(metrics_by_step)] |
| + |
| + |
| +def _plot(metrics: list[dict[str, float]], output_path: pathlib.Path, session: str) -> None: |
| + steps = [int(item["step"]) for item in metrics] |
| + losses = [item["loss"] for item in metrics] |
| + grad_norms = [item.get("grad_norm", float("nan")) for item in metrics] |
| + |
| + plt.rcParams.update( |
| + { |
| + "axes.facecolor": "#f6f2e8", |
| + "axes.edgecolor": "#27251f", |
| + "axes.labelcolor": "#27251f", |
| + "figure.facecolor": "#eee7d8", |
| + "font.family": "DejaVu Sans", |
| + "grid.color": "#c9bfaa", |
| + "text.color": "#27251f", |
| + "xtick.color": "#27251f", |
| + "ytick.color": "#27251f", |
| + } |
| + ) |
| + figure, axes = plt.subplots(2, 1, figsize=(10, 7), sharex=True, constrained_layout=True) |
| + figure.suptitle(f"SO-101 pi0.5 training | {session}", fontsize=16, fontweight="bold") |
| + |
| + axes[0].plot(steps, losses, color="#c7432b", marker="o", markersize=4, linewidth=2) |
| + axes[0].fill_between(steps, losses, color="#c7432b", alpha=0.12) |
| + axes[0].set_ylabel("Training loss") |
| + axes[0].set_title(f"Latest: {losses[-1]:.4f} at step {steps[-1]} | Best logged: {min(losses):.4f}", loc="left") |
| + axes[0].grid(alpha=0.7, linestyle="--") |
| + |
| + axes[1].plot(steps, grad_norms, color="#146b66", marker="o", markersize=4, linewidth=2) |
| + axes[1].set_xlabel("Optimizer step") |
| + axes[1].set_ylabel("Gradient norm") |
| + axes[1].grid(alpha=0.7, linestyle="--") |
| + |
| + figure.savefig(output_path, dpi=160) |
| + plt.close(figure) |
| + |
| + |
| +def _pane_is_dead(session: str) -> bool: |
| + result = subprocess.run( |
| + ["tmux", "display-message", "-p", "-t", session, "#{pane_dead}"], |
| + check=True, |
| + capture_output=True, |
| + text=True, |
| + ) |
| + return result.stdout.strip() == "1" |
| + |
| + |
| +def main() -> None: |
| + parser = argparse.ArgumentParser() |
| + parser.add_argument("--session", default="pi05_so101_stage1") |
| + parser.add_argument( |
| + "--output-dir", |
| + type=pathlib.Path, |
| + default=pathlib.Path("checkpoints/pi05_so101_erythromycin/pi05_so101_stage1"), |
| + ) |
| + parser.add_argument("--refresh-seconds", type=float, default=0.0) |
| + args = parser.parse_args() |
| + |
| + args.output_dir.mkdir(parents=True, exist_ok=True) |
| + csv_path = args.output_dir / "training_metrics.csv" |
| + image_path = args.output_dir / "loss_curve.png" |
| + |
| + while True: |
| + metrics = _merge_metrics(_read_csv(csv_path), _capture_metrics(args.session)) |
| + if not metrics: |
| + raise RuntimeError(f"No training metrics found in tmux session {args.session!r}") |
| + _write_csv(metrics, csv_path) |
| + _plot(metrics, image_path, args.session) |
| + latest = metrics[-1] |
| + print(f"step={int(latest['step'])} loss={latest['loss']:.4f} wrote={image_path}", flush=True) |
| + |
| + if args.refresh_seconds <= 0 or _pane_is_dead(args.session): |
| + return |
| + time.sleep(args.refresh_seconds) |
| + |
| + |
| +if __name__ == "__main__": |
| + main() |
| |
| new file mode 100755 |
| |
| |
| |
| @@ -0,0 +1,39 @@ |
| +#!/usr/bin/env bash |
| +set -euo pipefail |
| + |
| +if [[ $# -lt 1 ]]; then |
| + echo "Usage: $0 EXPERIMENT_NAME [additional train.py arguments...]" >&2 |
| + exit 2 |
| +fi |
| + |
| +so101_script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" |
| +so101_repo_root="$(dirname -- "$so101_script_dir")" |
| +so101_python="$so101_repo_root/.venv/bin/python" |
| +so101_gpu_id="${SO101_GPU_ID:-0}" |
| +so101_memory_fraction="${SO101_XLA_MEMORY_FRACTION:-0.90}" |
| +so101_min_free_memory_mib="${SO101_MIN_FREE_MEMORY_MIB:-90000}" |
| +so101_experiment_name="$1" |
| +shift |
| + |
| +cd "$so101_repo_root" |
| +so101_free_memory_mib="$(nvidia-smi --id="$so101_gpu_id" --query-gpu=memory.free --format=csv,noheader,nounits)" |
| +if (( so101_free_memory_mib < so101_min_free_memory_mib )); then |
| + echo "GPU $so101_gpu_id has only ${so101_free_memory_mib} MiB free; ${so101_min_free_memory_mib} MiB is required." >&2 |
| + echo "Choose another GPU with SO101_GPU_ID, or explicitly lower SO101_MIN_FREE_MEMORY_MIB for a smaller batch." >&2 |
| + exit 1 |
| +fi |
| + |
| +export CUDA_VISIBLE_DEVICES="$so101_gpu_id" |
| +export XLA_PYTHON_CLIENT_MEM_FRACTION="$so101_memory_fraction" |
| + |
| +nvidia-smi --id="$so101_gpu_id" --query-gpu=index,name,memory.used,memory.free,utilization.gpu --format=csv,noheader |
| +"$so101_python" scripts/so101_preflight.py |
| + |
| +so101_stats_path="assets/pi05_so101_erythromycin/local/so101_erythromycin_on_tea_grid90_v2/norm_stats.json" |
| +if [[ ! -f "$so101_stats_path" ]]; then |
| + "$so101_python" scripts/compute_so101_norm_stats.py |
| +fi |
| + |
| +exec "$so101_python" scripts/train.py pi05_so101_erythromycin \ |
| + --exp-name "$so101_experiment_name" \ |
| + "$@" |
| |
| new file mode 100644 |
| |
| |
| |
| @@ -0,0 +1,99 @@ |
| +"""Validate the authoritative SO-101 splits and the pi0.5 input pipeline.""" |
| + |
| +import dataclasses |
| + |
| +import numpy as np |
| +import tyro |
| + |
| +import openpi.training.config as _config |
| +import openpi.training.data_loader as _data_loader |
| + |
| +EXPECTED_SPLITS = { |
| + "clean_train": (67, 15_070), |
| + "clean_val": (18, 3_972), |
| + "recovery": (5, 1_515), |
| +} |
| + |
| + |
| +def _unwrap_dataset(dataset): |
| + while isinstance(dataset, _data_loader.TransformedDataset): |
| + dataset = dataset._dataset # noqa: SLF001 |
| + return dataset |
| + |
| + |
| +def _validate_action_chunk_boundaries(dataset, split_name: str) -> None: |
| + episode_indices = np.asarray(dataset.hf_dataset["episode_index"], dtype=np.int64) |
| + if len(episode_indices) != dataset.num_frames: |
| + raise ValueError(f"{split_name} episode index length does not match its frame count") |
| + |
| + for episode_id in dataset.episodes: |
| + episode_start = int(dataset.episode_data_index["from"][episode_id]) |
| + episode_end = int(dataset.episode_data_index["to"][episode_id]) |
| + if episode_start < 0 or episode_end <= episode_start: |
| + raise ValueError(f"{split_name} episode {episode_id} has invalid compact frame boundaries") |
| + if not np.all(episode_indices[episode_start:episode_end] == episode_id): |
| + raise ValueError(f"{split_name} episode {episode_id} frame boundaries point to another episode") |
| + |
| + for frame_index in range(episode_start, episode_end): |
| + query_indices, _ = dataset._get_query_indices(frame_index, episode_id) # noqa: SLF001 |
| + action_indices = np.asarray(query_indices["action"]) |
| + if action_indices.min() < episode_start or action_indices.max() >= episode_end: |
| + raise ValueError(f"{split_name} episode {episode_id} action chunk crosses an episode boundary") |
| + |
| + |
| +def main(config_name: str = "pi05_so101_erythromycin") -> None: |
| + config = _config.get_config(config_name) |
| + if not isinstance(config.data, _config.LeRobotSO101DataConfig): |
| + raise TypeError(f"Config {config_name!r} does not use LeRobotSO101DataConfig") |
| + |
| + train_dataset = None |
| + train_data_config = None |
| + for split_name, (expected_episodes, expected_frames) in EXPECTED_SPLITS.items(): |
| + factory = dataclasses.replace(config.data, split_name=split_name) |
| + data_config = factory.create(config.assets_dirs, config.model) |
| + dataset = _data_loader.create_torch_dataset(data_config, config.model.action_horizon, config.model) |
| + raw_dataset = _unwrap_dataset(dataset) |
| + if raw_dataset.num_episodes != expected_episodes or raw_dataset.num_frames != expected_frames: |
| + raise ValueError( |
| + f"{split_name} resolved to {raw_dataset.num_episodes} episodes/{raw_dataset.num_frames} frames; " |
| + f"expected {expected_episodes}/{expected_frames}" |
| + ) |
| + _validate_action_chunk_boundaries(raw_dataset, split_name) |
| + print(f"{split_name}: episodes={raw_dataset.num_episodes} frames={raw_dataset.num_frames}") |
| + if split_name == "clean_train": |
| + train_dataset = dataset |
| + train_data_config = data_config |
| + |
| + assert train_dataset is not None |
| + assert train_data_config is not None |
| + sample = train_dataset[0] |
| + for transform in ( |
| + *train_data_config.repack_transforms.inputs, |
| + *train_data_config.data_transforms.inputs, |
| + *train_data_config.model_transforms.inputs, |
| + ): |
| + sample = transform(sample) |
| + |
| + expected_shapes = { |
| + "state": (32,), |
| + "actions": (50, 32), |
| + "tokenized_prompt": (200,), |
| + } |
| + for key, expected_shape in expected_shapes.items(): |
| + if np.asarray(sample[key]).shape != expected_shape: |
| + raise ValueError(f"{key} has shape {np.asarray(sample[key]).shape}; expected {expected_shape}") |
| + for image_name, image in sample["image"].items(): |
| + if np.asarray(image).shape != (224, 224, 3): |
| + raise ValueError(f"{image_name} has shape {np.asarray(image).shape}; expected (224, 224, 3)") |
| + if bool(sample["image_mask"]["right_wrist_0_rgb"]): |
| + raise ValueError("The nonexistent right-wrist camera must be masked") |
| + if np.any(np.asarray(sample["actions"])[..., 6:]): |
| + raise ValueError("Padded action dimensions must be zero") |
| + |
| + print("model_input: state=(32,) actions=(50, 32) images=3x(224, 224, 3)") |
| + print("camera_masks: base=True left_wrist=True right_wrist=False") |
| + print("preflight: PASS") |
| + |
| + |
| +if __name__ == "__main__": |
| + tyro.cli(main) |
| |
| |
| |
| |
| @@ -226,12 +226,12 @@ def main(config: _config.TrainConfig): |
| batch = next(data_iter) |
| logging.info(f"Initialized data loader:\n{training_utils.array_tree_to_info(batch)}") |
| |
| - # Log images from first batch to sanity check. |
| - images_to_log = [ |
| - wandb.Image(np.concatenate([np.array(img[i]) for img in batch[0].images.values()], axis=1)) |
| - for i in range(min(5, len(next(iter(batch[0].images.values()))))) |
| - ] |
| - wandb.log({"camera_views": images_to_log}, step=0) |
| + if config.wandb_enabled and config.wandb_log_images: |
| + images_to_log = [ |
| + wandb.Image(np.concatenate([np.array(img[i]) for img in batch[0].images.values()], axis=1)) |
| + for i in range(min(5, len(next(iter(batch[0].images.values()))))) |
| + ] |
| + wandb.log({"camera_views": images_to_log}, step=0) |
| |
| train_state, train_state_sharding = init_train_state(config, init_rng, mesh, resume=resuming) |
| jax.block_until_ready(train_state) |
| @@ -269,7 +269,9 @@ def main(config: _config.TrainConfig): |
| infos = [] |
| batch = next(data_iter) |
| |
| - if (step % config.save_interval == 0 and step > start_step) or step == config.num_train_steps - 1: |
| + if (step % config.save_interval == 0 and step > start_step) or ( |
| + config.save_final_checkpoint and step == config.num_train_steps - 1 |
| + ): |
| _checkpoints.save_state(checkpoint_manager, train_state, data_loader, step) |
| |
| logging.info("Waiting for checkpoint manager to finish") |
| |
| new file mode 100644 |
| |
| |
| |
| @@ -0,0 +1,78 @@ |
| +import dataclasses |
| + |
| +import einops |
| +import numpy as np |
| + |
| +from openpi import transforms |
| +from openpi.models import model as _model |
| + |
| +SO101_ACTION_DIM = 6 |
| + |
| + |
| +def make_so101_example() -> dict: |
| + """Create an example observation using the policy-server input schema.""" |
| + return { |
| + "observation/state": np.zeros((SO101_ACTION_DIM,), dtype=np.float32), |
| + "observation/fixed_image": np.zeros((480, 640, 3), dtype=np.uint8), |
| + "observation/wrist_image": np.zeros((480, 640, 3), dtype=np.uint8), |
| + "prompt": "Pick up the red erythromycin ointment box and place it on top of the green Rizhao tea tin.", |
| + } |
| + |
| + |
| +def _parse_image(image: np.ndarray) -> np.ndarray: |
| + image = np.asarray(image) |
| + if np.issubdtype(image.dtype, np.floating): |
| + image = np.clip(image * 255.0, 0.0, 255.0).astype(np.uint8) |
| + if image.ndim != 3: |
| + raise ValueError(f"Expected a 3-D image, got shape {image.shape}") |
| + if image.shape[0] == 3 and image.shape[-1] != 3: |
| + image = einops.rearrange(image, "c h w -> h w c") |
| + if image.shape[-1] != 3: |
| + raise ValueError(f"Expected an RGB image, got shape {image.shape}") |
| + return image |
| + |
| + |
| +@dataclasses.dataclass(frozen=True) |
| +class SO101Inputs(transforms.DataTransformFn): |
| + """Map SO-101 observations to the three image slots expected by pi0.5.""" |
| + |
| + model_type: _model.ModelType |
| + |
| + def __call__(self, data: dict) -> dict: |
| + state = np.asarray(data["observation/state"], dtype=np.float32) |
| + if state.shape[-1] != SO101_ACTION_DIM: |
| + raise ValueError(f"Expected {SO101_ACTION_DIM}-D state, got shape {state.shape}") |
| + |
| + fixed_image = _parse_image(data["observation/fixed_image"]) |
| + wrist_image = _parse_image(data["observation/wrist_image"]) |
| + inputs = { |
| + "state": state, |
| + "image": { |
| + "base_0_rgb": fixed_image, |
| + "left_wrist_0_rgb": wrist_image, |
| + "right_wrist_0_rgb": np.zeros_like(fixed_image), |
| + }, |
| + "image_mask": { |
| + "base_0_rgb": np.True_, |
| + "left_wrist_0_rgb": np.True_, |
| + "right_wrist_0_rgb": np.True_ if self.model_type == _model.ModelType.PI0_FAST else np.False_, |
| + }, |
| + } |
| + |
| + if "actions" in data: |
| + actions = np.asarray(data["actions"], dtype=np.float32) |
| + if actions.shape[-1] != SO101_ACTION_DIM: |
| + raise ValueError(f"Expected {SO101_ACTION_DIM}-D actions, got shape {actions.shape}") |
| + inputs["actions"] = actions |
| + |
| + if "prompt" in data: |
| + inputs["prompt"] = data["prompt"] |
| + return inputs |
| + |
| + |
| +@dataclasses.dataclass(frozen=True) |
| +class SO101Outputs(transforms.DataTransformFn): |
| + """Remove the model's padded action dimensions before robot execution.""" |
| + |
| + def __call__(self, data: dict) -> dict: |
| + return {"actions": np.asarray(data["actions"])[..., :SO101_ACTION_DIM]} |
| |
| new file mode 100644 |
| |
| |
| |
| @@ -0,0 +1,35 @@ |
| +import numpy as np |
| + |
| +from openpi.models import model as _model |
| +from openpi.policies import so101_policy |
| + |
| + |
| +def test_so101_inputs_map_two_cameras_and_actions() -> None: |
| + transform = so101_policy.SO101Inputs(model_type=_model.ModelType.PI05) |
| + result = transform( |
| + { |
| + "observation/state": np.arange(6, dtype=np.float32), |
| + "observation/fixed_image": np.zeros((3, 12, 16), dtype=np.float32), |
| + "observation/wrist_image": np.zeros((12, 16, 3), dtype=np.uint8), |
| + "actions": np.zeros((50, 6), dtype=np.float32), |
| + "prompt": "test prompt", |
| + } |
| + ) |
| + |
| + assert result["state"].shape == (6,) |
| + assert result["actions"].shape == (50, 6) |
| + assert result["image"]["base_0_rgb"].shape == (12, 16, 3) |
| + assert result["image"]["left_wrist_0_rgb"].shape == (12, 16, 3) |
| + assert result["image"]["right_wrist_0_rgb"].shape == (12, 16, 3) |
| + assert result["image_mask"] == { |
| + "base_0_rgb": np.True_, |
| + "left_wrist_0_rgb": np.True_, |
| + "right_wrist_0_rgb": np.False_, |
| + } |
| + assert result["prompt"] == "test prompt" |
| + |
| + |
| +def test_so101_outputs_remove_padding() -> None: |
| + actions = np.arange(50 * 32, dtype=np.float32).reshape(50, 32) |
| + result = so101_policy.SO101Outputs()({"actions": actions}) |
| + np.testing.assert_array_equal(result["actions"], actions[:, :6]) |
| |
| |
| |
| |
| @@ -4,6 +4,7 @@ import abc |
| from collections.abc import Sequence |
| import dataclasses |
| import difflib |
| +import json |
| import logging |
| import pathlib |
| from typing import Any, Literal, Protocol, TypeAlias |
| @@ -20,6 +21,7 @@ import openpi.models.tokenizer as _tokenizer |
| import openpi.policies.aloha_policy as aloha_policy |
| import openpi.policies.droid_policy as droid_policy |
| import openpi.policies.libero_policy as libero_policy |
| +import openpi.policies.so101_policy as so101_policy |
| import openpi.shared.download as _download |
| import openpi.shared.normalize as _normalize |
| import openpi.training.droid_rlds_dataset as droid_rlds_dataset |
| @@ -90,6 +92,13 @@ class DataConfig: |
| # If true, will use the LeRobot dataset task to define the prompt. |
| prompt_from_task: bool = False |
| |
| + # Optional local LeRobot dataset root. If unset, LeRobot uses its standard cache. |
| + dataset_root: str | None = None |
| + # Optional episode subset. This must be explicit for datasets with held-out validation episodes. |
| + episodes: Sequence[int] | None = None |
| + # Video backend passed to LeRobot. PyAV is required for the SO-101 AV1 videos. |
| + video_backend: str | None = None |
| + |
| # Only used for RLDS data loader (ie currently only used for DROID). |
| rlds_data_dir: str | None = None |
| # Action space for DROID dataset. |
| @@ -462,6 +471,83 @@ class LeRobotDROIDDataConfig(DataConfigFactory): |
| ) |
| |
| |
| +_SO101_PACKAGE_ROOT = str( |
| + pathlib.Path(__file__).resolve().parents[3].parent / "so101_erythromycin_on_tea_grid90_v2_portable" |
| +) |
| + |
| + |
| +@dataclasses.dataclass(frozen=True) |
| +class LeRobotSO101DataConfig(DataConfigFactory): |
| + """Data config for the local SO-101 erythromycin-on-tea dataset.""" |
| + |
| + package_root: str = _SO101_PACKAGE_ROOT |
| + split_name: Literal["clean_train", "clean_val", "recovery", "clean_all"] = "clean_train" |
| + use_delta_joint_actions: bool = True |
| + |
| + def _load_and_validate_split(self, package_root: pathlib.Path) -> tuple[int, ...]: |
| + manifest_path = package_root / "splits" / "split_manifest.json" |
| + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) |
| + |
| + expected_counts = {"clean_train": 67, "clean_val": 18, "recovery": 5, "clean_all": 85} |
| + splits = {name: tuple(int(index) for index in manifest[name]) for name in expected_counts} |
| + for name, expected_count in expected_counts.items(): |
| + if len(splits[name]) != expected_count or len(set(splits[name])) != expected_count: |
| + raise ValueError(f"Invalid {name} split in {manifest_path}: expected {expected_count} unique episodes") |
| + |
| + train, val, recovery = (set(splits[name]) for name in ("clean_train", "clean_val", "recovery")) |
| + if train & val or train & recovery or val & recovery: |
| + raise ValueError(f"Train, validation, and recovery splits overlap in {manifest_path}") |
| + if train | val | recovery != set(range(90)): |
| + raise ValueError(f"Authoritative splits do not cover exactly episodes 0 through 89 in {manifest_path}") |
| + if set(splits["clean_all"]) != train | val: |
| + raise ValueError(f"clean_all is not clean_train union clean_val in {manifest_path}") |
| + return splits[self.split_name] |
| + |
| + @override |
| + def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig: |
| + package_root = pathlib.Path(self.package_root).expanduser().resolve() |
| + dataset_root = package_root / "dataset" |
| + if not (dataset_root / "meta" / "info.json").is_file(): |
| + raise FileNotFoundError(f"SO-101 dataset not found at {dataset_root}") |
| + episode_ids = self._load_and_validate_split(package_root) |
| + |
| + repack_transform = _transforms.Group( |
| + inputs=[ |
| + _transforms.RepackTransform( |
| + { |
| + "observation/fixed_image": "observation.images.fixed", |
| + "observation/wrist_image": "observation.images.wrist", |
| + "observation/state": "observation.state", |
| + "actions": "action", |
| + "prompt": "prompt", |
| + } |
| + ) |
| + ] |
| + ) |
| + data_transforms = _transforms.Group( |
| + inputs=[so101_policy.SO101Inputs(model_type=model_config.model_type)], |
| + outputs=[so101_policy.SO101Outputs()], |
| + ) |
| + if self.use_delta_joint_actions: |
| + # The first five dimensions are arm joints; the sixth is the absolute gripper command. |
| + delta_action_mask = _transforms.make_bool_mask(5, -1) |
| + data_transforms = data_transforms.push( |
| + inputs=[_transforms.DeltaActions(delta_action_mask)], |
| + outputs=[_transforms.AbsoluteActions(delta_action_mask)], |
| + ) |
| + |
| + return dataclasses.replace( |
| + self.create_base_config(assets_dirs, model_config), |
| + dataset_root=str(dataset_root), |
| + episodes=episode_ids, |
| + video_backend="pyav", |
| + repack_transforms=repack_transform, |
| + data_transforms=data_transforms, |
| + model_transforms=ModelTransformFactory()(model_config), |
| + action_sequence_keys=("action",), |
| + ) |
| + |
| + |
| @dataclasses.dataclass(frozen=True) |
| class TrainConfig: |
| # Name of the config. Must be unique. Will be used to reference this config. |
| @@ -514,6 +600,8 @@ class TrainConfig: |
| log_interval: int = 100 |
| # How often (in steps) to save checkpoints. |
| save_interval: int = 1000 |
| + # Save a checkpoint on the final step even if it is not on the regular interval. |
| + save_final_checkpoint: bool = True |
| # If set, any existing checkpoints matching step % keep_period == 0 will not be deleted. |
| keep_period: int | None = 5000 |
| |
| @@ -524,6 +612,8 @@ class TrainConfig: |
| |
| # If true, will enable wandb logging. |
| wandb_enabled: bool = True |
| + # If true, upload a small first-batch camera preview to wandb. |
| + wandb_log_images: bool = True |
| |
| # Used to pass metadata to the policy server. |
| policy_metadata: dict[str, Any] | None = None |
| @@ -916,6 +1006,29 @@ _CONFIGS = [ |
| num_train_steps=20_000, |
| batch_size=32, |
| ), |
| + # Full pi0.5 fine-tuning on the local dual-camera SO-101 dataset. |
| + TrainConfig( |
| + name="pi05_so101_erythromycin", |
| + model=pi0_config.Pi0Config(pi05=True, action_horizon=50), |
| + data=LeRobotSO101DataConfig( |
| + repo_id="local/so101_erythromycin_on_tea_grid90_v2", |
| + base_config=DataConfig(prompt_from_task=True), |
| + ), |
| + weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"), |
| + batch_size=8, |
| + num_workers=4, |
| + num_train_steps=30_000, |
| + save_interval=1_000, |
| + keep_period=None, |
| + wandb_enabled=False, |
| + wandb_log_images=False, |
| + policy_metadata={ |
| + "robot_type": "so101_follower", |
| + "action_dim": so101_policy.SO101_ACTION_DIM, |
| + "action_horizon": 50, |
| + "cameras": ["fixed", "wrist"], |
| + }, |
| + ), |
| # |
| # ALOHA Sim configs. This config is used to demonstrate how to train on a simple simulated environment. |
| # |
| |
| |
| |
| |
| @@ -137,13 +137,17 @@ def create_torch_dataset( |
| if repo_id == "fake": |
| return FakeDataset(model_config, num_samples=1024) |
| |
| - dataset_meta = lerobot_dataset.LeRobotDatasetMetadata(repo_id) |
| + dataset_meta = lerobot_dataset.LeRobotDatasetMetadata(repo_id, root=data_config.dataset_root) |
| dataset = lerobot_dataset.LeRobotDataset( |
| data_config.repo_id, |
| + root=data_config.dataset_root, |
| + episodes=None if data_config.episodes is None else list(data_config.episodes), |
| delta_timestamps={ |
| key: [t / dataset_meta.fps for t in range(action_horizon)] for key in data_config.action_sequence_keys |
| }, |
| + video_backend=data_config.video_backend, |
| ) |
| + _align_selected_episode_data_index(dataset) |
| |
| if data_config.prompt_from_task: |
| dataset = TransformedDataset(dataset, [_transforms.PromptFromLeRobotTask(dataset_meta.tasks)]) |
| @@ -151,6 +155,23 @@ def create_torch_dataset( |
| return dataset |
| |
| |
| +def _align_selected_episode_data_index(dataset: lerobot_dataset.LeRobotDataset) -> None: |
| + """Make compact LeRobot subset boundaries indexable by preserved episode IDs.""" |
| + if dataset.episodes is None: |
| + return |
| + |
| + compact_index = dataset.episode_data_index |
| + if len(compact_index["from"]) != len(dataset.episodes): |
| + raise ValueError("LeRobot returned an unexpected episode boundary table") |
| + |
| + max_episode_id = max(dataset.episodes) |
| + aligned_index = {key: values.new_full((max_episode_id + 1,), -1) for key, values in compact_index.items()} |
| + for compact_position, episode_id in enumerate(dataset.episodes): |
| + for key, values in compact_index.items(): |
| + aligned_index[key][episode_id] = values[compact_position] |
| + dataset.episode_data_index = aligned_index |
| + |
| + |
| def create_rlds_dataset( |
| data_config: _config.DataConfig, |
| action_horizon: int, |
| |
| |
| |
| |
| @@ -1,12 +1,39 @@ |
| import dataclasses |
| +from types import SimpleNamespace |
| |
| import jax |
| +from lerobot.common.datasets.lerobot_dataset import LeRobotDataset |
| +import torch |
| |
| from openpi.models import pi0_config |
| from openpi.training import config as _config |
| from openpi.training import data_loader as _data_loader |
| |
| |
| +def test_align_selected_episode_data_index_preserves_original_ids(): |
| + dataset = SimpleNamespace( |
| + episodes=[0, 2, 68, 89], |
| + episode_data_index={ |
| + "from": torch.tensor([0, 5, 12, 20]), |
| + "to": torch.tensor([5, 12, 20, 30]), |
| + }, |
| + delta_indices={"action": list(range(50))}, |
| + ) |
| + |
| + _data_loader._align_selected_episode_data_index(dataset) # noqa: SLF001 |
| + |
| + assert dataset.episode_data_index["from"][68].item() == 12 |
| + assert dataset.episode_data_index["to"][68].item() == 20 |
| + assert dataset.episode_data_index["from"][89].item() == 20 |
| + assert dataset.episode_data_index["to"][89].item() == 30 |
| + assert dataset.episode_data_index["from"][67].item() == -1 |
| + |
| + query_indices, padding = LeRobotDataset._get_query_indices(dataset, idx=19, ep_idx=68) # noqa: SLF001 |
| + assert min(query_indices["action"]) == 19 |
| + assert max(query_indices["action"]) == 19 |
| + assert padding["action_is_pad"].tolist() == [False] + [True] * 49 |
| + |
| + |
| def test_torch_data_loader(): |
| config = pi0_config.Pi0Config(action_dim=24, action_horizon=50, max_token_len=48) |
| dataset = _data_loader.FakeDataset(config, 16) |
| -- |
| 2.43.0 |
|
|
|
|
