From 78183844fcc4b1a090d244585659e6a7d2e5b9b8 Mon Sep 17 00:00:00 2001 From: CodeChildCZJ 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 diff --git a/HF_MODEL_CARD.md b/HF_MODEL_CARD.md new file mode 100644 index 0000000..0680386 --- /dev/null +++ b/HF_MODEL_CARD.md @@ -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 | | + +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= .venv/bin/python scripts/serve_policy.py \ + policy:checkpoint \ + --policy.config pi05_so101_erythromycin \ + --policy.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. diff --git a/NOTICE b/NOTICE new file mode 100644 index 0000000..27fe1d3 --- /dev/null +++ b/NOTICE @@ -0,0 +1 @@ +Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms diff --git a/SO101_PI05_HANDOFF.md b/SO101_PI05_HANDOFF.md new file mode 100644 index 0000000..f2e3db6 --- /dev/null +++ b/SO101_PI05_HANDOFF.md @@ -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 | | + +最重要的三点： + +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= .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= .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 的成功/失败和安全事件完整记录。 diff --git a/examples/so101/README.md b/examples/so101/README.md new file mode 100644 index 0000000..1415ca4 --- /dev/null +++ b/examples/so101/README.md @@ -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 . + +## 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. diff --git a/scripts/compute_so101_norm_stats.py b/scripts/compute_so101_norm_stats.py new file mode 100644 index 0000000..174716c --- /dev/null +++ b/scripts/compute_so101_norm_stats.py @@ -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) diff --git a/scripts/eval_so101_checkpoint.py b/scripts/eval_so101_checkpoint.py new file mode 100644 index 0000000..3fbe957 --- /dev/null +++ b/scripts/eval_so101_checkpoint.py @@ -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) diff --git a/scripts/plot_so101_training.py b/scripts/plot_so101_training.py new file mode 100644 index 0000000..918c174 --- /dev/null +++ b/scripts/plot_so101_training.py @@ -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\d+): (?P.+)$", re.MULTILINE) +METRIC_PATTERN = re.compile(r"(?P[a-z_]+)=(?P[-+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() diff --git a/scripts/run_so101_training.sh b/scripts/run_so101_training.sh new file mode 100755 index 0000000..3a3670f --- /dev/null +++ b/scripts/run_so101_training.sh @@ -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" \ + "$@" diff --git a/scripts/so101_preflight.py b/scripts/so101_preflight.py new file mode 100644 index 0000000..4614a7f --- /dev/null +++ b/scripts/so101_preflight.py @@ -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) diff --git a/scripts/train.py b/scripts/train.py index 5d28941..aa70a08 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -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") diff --git a/src/openpi/policies/so101_policy.py b/src/openpi/policies/so101_policy.py new file mode 100644 index 0000000..7cd2566 --- /dev/null +++ b/src/openpi/policies/so101_policy.py @@ -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]} diff --git a/src/openpi/policies/so101_policy_test.py b/src/openpi/policies/so101_policy_test.py new file mode 100644 index 0000000..85397a6 --- /dev/null +++ b/src/openpi/policies/so101_policy_test.py @@ -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]) diff --git a/src/openpi/training/config.py b/src/openpi/training/config.py index 4ca47e1..98273b5 100644 --- a/src/openpi/training/config.py +++ b/src/openpi/training/config.py @@ -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. # diff --git a/src/openpi/training/data_loader.py b/src/openpi/training/data_loader.py index e2ee7dd..9cbc551 100644 --- a/src/openpi/training/data_loader.py +++ b/src/openpi/training/data_loader.py @@ -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, diff --git a/src/openpi/training/data_loader_test.py b/src/openpi/training/data_loader_test.py index d15a735..5288dfa 100644 --- a/src/openpi/training/data_loader_test.py +++ b/src/openpi/training/data_loader_test.py @@ -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