code/so101-pi05-erythromycin-v1.patch

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

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 | <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.
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 | <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 的成功/失败和安全事件完整记录。
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 <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.
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<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()
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