Pi0 与 Pi0.5 模型微调指南

本文介绍如何使用艾欧智能发布的 Docker 镜像，在 pi0_base / pi05_base 等预训练权重基础上微调 Pi0 与 Pi0.5 模型。文中的命令、目录约定和参数名与镜像内 OpenPI 封装入口 train_lerobot.py 一致。

如果你的目标是复现官方 OpenPI 微调链路，同时又希望直接使用已经准备好的镜像环境，那么这条路径最省事，也最接近实际可复用的部署方式。

为什么推荐这条路径

Pi0 / Pi0.5 基于 OpenPI 官方栈与公开基座检查点进行微调（底层为 JAX 实现）。对于已经准备好 LeRobot 数据集的使用者来说，最直接的做法不是先手工搭建完整环境，而是直接使用：

• ioaitech/train_openpi:pi0
• ioaitech/train_openpi:pi05

这两个镜像已经包含了微调所需依赖，并约定：

• 输入数据挂载到 /data/input
• 输出（检查点等）写入 /data/output

国内镜像加速（华为云 SWR）

文档默认以 Docker Hub 上的 ioaitech/train_openpi:pi0 / ioaitech/train_openpi:pi05 为例。若你在中国大陆访问 Docker Hub 较慢，可将镜像前缀替换为 swr.cn-east-3.myhuaweicloud.com/ioaitech/，例如：

• swr.cn-east-3.myhuaweicloud.com/ioaitech/train_openpi:pi0
• swr.cn-east-3.myhuaweicloud.com/ioaitech/train_openpi:pi05 docker run 的其余参数不变。

一键开始微调

前置条件

• Linux 主机
• NVIDIA 驱动安装正常
• Docker 可用
• docker run --gpus all 可正常访问 GPU

建议先做一次 GPU 自检：

docker run --rm --gpus all nvidia/cuda:12.1.0-base-ubuntu22.04 nvidia-smi

Pi0 一键微调

docker run --rm --gpus all \
  -v /path/to/lerobot_dataset:/data/input \
  -v /path/to/output:/data/output \
  ioaitech/train_openpi:pi0 \
  --batch_size 8 \
  --steps 20000 \
  --save_interval 1000 \
  --learning_rate 2.5e-5 \
  --action_horizon 50 \
  --prompt "pick up the object"

Pi0.5 一键微调

docker run --rm --gpus all \
  -v /path/to/lerobot_dataset:/data/input \
  -v /path/to/output:/data/output \
  ioaitech/train_openpi:pi05 \
  --batch_size 8 \
  --steps 20000 \
  --save_interval 1000 \
  --learning_rate 2.5e-5 \
  --action_horizon 50 \
  --prompt "pick up the object"

上面两条命令的主要区别在于镜像标签，对应不同的基座权重与模型配置。

最小可运行命令

如果你只是想先验证数据挂载和微调链路是否打通，可以把命令缩减为：

docker run --rm --gpus all \
  -v /path/to/lerobot_dataset:/data/input \
  -v /path/to/output:/data/output \
  ioaitech/train_openpi:pi0 \
  --batch_size 1 \
  --steps 1000

首次跑通之后，再逐步增加 batch_size、steps 和其它微调超参数。

数据要求

容器入口会检查 /data/input/meta/info.json 是否存在；缺少该文件时会直接退出。因此请确保你的 LeRobot 数据集根目录至少包含：

your_dataset/
├── meta/
│   └── info.json
├── data/
└── videos/

当前微调封装会自动处理以下事项：

• 识别 LeRobot 数据集版本
• 必要时将 v3 数据转换为兼容微调流程的 v2.1 布局
• 为 v2.1 数据补齐 episodes_stats.jsonl
• 生成归一化统计信息
• 将挂载的数据集链接到 LeRobot cache，供 OpenPI 读取

Pi0 与 Pi0.5 的区别

这两种模型共用一套微调封装，但底层配置并不相同。

在当前代码中，差异主要体现在以下几点：

• Pi0 使用 pi0_base 预训练权重
• Pi0.5 使用 pi05_base 预训练权重
• Pi0.5 会启用 Pi0Config(pi05=True)，其状态输入和 token 配置与 Pi0 不同
• Pi0.5 默认 max_token_len 更长，更适合需要更强表达能力的场景

若你希望先快速复现微调流程，建议从 Pi0 开始；若你已经有明确的任务需求，并且计划直接使用 Pi0.5 的能力边界，再切到 ioaitech/train_openpi:pi05。

常用参数

以下参数与镜像内 train_lerobot.py 的定义一致。

参数	默认值	说明
--batch_size	1	每个优化 step 的 batch；最终会自动调整为能整除 JAX 设备数
--steps	1000	优化步数
--gpus	all	使用全部 GPU，或传 0,1 这类列表
--prompt	空	当数据集没有任务文本时使用的默认语言提示
--save_interval	500	checkpoint 保存间隔
--learning_rate	空	不传时使用默认峰值学习率 2.5e-5
--fsdp_devices	auto	自动决定是否启用 FSDP，以及分片设备数
--lora	auto	单卡默认启用 LoRA，多卡默认关闭
--ema_decay	空	EMA 衰减；LoRA 模式下默认关闭，以节省显存
--action_horizon	50	动作序列长度

微调行为说明

当前封装对单卡和多卡做了自动策略选择：

• 单卡时，--lora auto 默认会启用 LoRA，以降低显存压力
• 多卡时，--fsdp_devices auto 会自动切换到 FSDP 路径
• batch_size 若小于设备数，脚本会自动上调到可整除的值

这意味着文档中的一键命令首先强调“能跑起来”，而不是要求用户一开始就处理全部 JAX 细节。

关于 prompt

如果数据集中本身包含任务文本，微调时会优先使用数据中的任务信息。只有当数据集缺少任务字段时，--prompt 才会作为默认指令生效。

因此，建议把 --prompt 理解为兜底参数，而不是所有场景都必须手工填写的参数。

输出结果

微调输出会写入挂载的 /data/output。当前 TrainConfig 固定使用：

• name = docker_train
• exp_name = train

因此 checkpoint 默认会落到：

/path/to/output/docker_train/train/

微调过程中还会生成归一化统计信息，保存在资产目录下，供后续微调与推理使用。

推荐用法

1. 先做小规模验证

第一次跑新数据集时，建议先用下面这类配置做链路验证：

docker run --rm --gpus all \
  -v /path/to/lerobot_dataset:/data/input \
  -v /path/to/output:/data/output \
  ioaitech/train_openpi:pi0 \
  --batch_size 1 \
  --steps 1000 \
  --save_interval 200

确认日志正常、输出目录出现 checkpoint 后，再把微调规模调大。

2. 单卡优先稳妥配置

如果你只有单张 GPU，建议先保留默认自动策略，不要一开始就强行关闭 LoRA：

docker run --rm --gpus all \
  -v /path/to/lerobot_dataset:/data/input \
  -v /path/to/output:/data/output \
  ioaitech/train_openpi:pi0 \
  --batch_size 4 \
  --steps 20000 \
  --save_interval 1000 \
  --action_horizon 50

3. 多卡微调

若你已经具备多卡环境，可以显式指定 GPU 和 FSDP 设备数：

docker run --rm --gpus all \
  -v /path/to/lerobot_dataset:/data/input \
  -v /path/to/output:/data/output \
  ioaitech/train_openpi:pi05 \
  --gpus 0,1,2,3 \
  --batch_size 16 \
  --steps 30000 \
  --fsdp_devices 4 \
  --save_interval 1000

常见问题

1. 容器启动后提示找不到 LeRobot 数据集

先检查：

• 宿主机路径是否正确挂载到 /data/input
• /data/input/meta/info.json 是否存在

容器入口会在最开始做这个检查，路径不对时不会进入微调阶段。

2. 为什么同样的命令在单卡和多卡上行为不同

这是封装脚本的设计使然。train_lerobot.py 会根据设备数量自动决定是否启用 LoRA、是否使用 FSDP，并对 batch_size 做兼容调整，以减少初次使用的试错成本。

3. 为什么输出目录不是自己命名的

当前封装中输出目录命名固定为 /data/output/docker_train/train。如果后续需要按实验名区分目录，可以在平台侧或后续版本里再增加更细的命名入口；目前文档以镜像内行为为准。

4. 需要手工先跑归一化统计吗

不需要。当前封装会在微调开始前自动计算并保存归一化统计信息，这也是推荐直接使用镜像的原因之一。

5. Pi0 和 Pi0.5 该选哪个

如果你当前的目标是先稳定复现完整微调流程，建议先用 Pi0；如果你已经确认要使用 Pi0.5，并且希望直接沿着 Pi0.5 的基座继续微调，就直接使用 ioaitech/train_openpi:pi05。

实践建议

• 先验证数据集结构，再启动正式微调
• 先跑短流程，再增加优化步数
• 单卡场景先接受默认自动策略，避免过早手工干预 LoRA/FSDP
• Pi0 与 Pi0.5 分别保留独立输出目录，便于后续比较效果

参考资料

• OpenPI 官方 GitHub
• Pi0 论文: π₀: A Flow-based Vision-Language-Action Model
• LeRobot 数据集规范