跳到主要内容

Pi0 与 Pi0.5 模型微调

Pi0 与 Pi0.5 来自 Physical Intelligence 的 OpenPI 项目。官方仓库 Physical-Intelligence/openpi 明确提供 pi0_base、pi05_base 等基座权重，并以 LeRobot 数据集作为自定义数据微调的主要入口。Pi0.5 在该仓库中当前训练与推理支持 flow matching head；如果你的任务依赖更完整的分层规划能力，需要先确认官方实现是否已经覆盖。

艾欧智能已经基于 OpenPI 训练链路构建并发布 ioaitech/train_openpi:pi0 与 ioaitech/train_openpi:pi05 Docker 镜像。镜像包含 OpenPI 依赖环境、基座权重下载逻辑、归一化统计计算，以及 LeRobot 数据适配层。建议先用镜像跑通微调闭环，再按需回到 OpenPI 官方源码修改数据映射或训练配置。

国内镜像

文档默认使用 Docker Hub 地址。若访问较慢，在镜像名前加 swr.cn-east-3.myhuaweicloud.com/：

  • swr.cn-east-3.myhuaweicloud.com/ioaitech/train_openpi:pi0
  • swr.cn-east-3.myhuaweicloud.com/ioaitech/train_openpi:pi05

何时使用​

Pi0/Pi0.5 适合已经有较高质量多模态数据、希望在大模型基座上微调的场景。相较 ACT 或 Diffusion Policy，它的环境复杂度更高：JAX、FSDP、LoRA、基座权重、归一化统计和语言任务字段都需要正确配置。

模型镜像基座权重建议场景
Pi0ioaitech/train_openpi:pi0gs://openpi-assets/checkpoints/pi0_base先跑通 OpenPI 微调链路，或任务对语言泛化要求中等。
Pi0.5ioaitech/train_openpi:pi05gs://openpi-assets/checkpoints/pi05_base需要更强开放环境泛化，且具备较充足算力与数据。

数据要求​

镜像入口会检查 /data/input/meta/info.json。数据目录至少应包含：

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

封装脚本会自动完成以下处理：

  • 读取 LeRobot 数据集 schema。
  • 将数据集链接到容器内 LeRobot cache。
  • 对 v2.1 数据补齐 episodes_stats.jsonl。
  • 处理部分 Hugging Face parquet metadata 兼容问题。
  • 计算 OpenPI 训练所需的归一化统计。
  • 根据 GPU 数量自动选择单卡 LoRA 或多卡 FSDP。

语言任务字段优先来自数据集本身。只有数据集缺少任务文本时，--prompt 才作为默认提示生效。

使用镜像训练​

先确认 GPU 容器可用：

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

最小验证​

这条命令只用于检查挂载、数据读取和 OpenPI 训练入口是否工作：

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

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 4 \
  --steps 20000 \
  --learning_rate 2.5e-5 \
  --save_interval 1000 \
  --action_horizon 50

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 30000 \
  --learning_rate 2.5e-5 \
  --save_interval 1000 \
  --action_horizon 50

多卡时可以显式指定：

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

输出默认写入：

/path/to/output/docker_train/train/

这是封装脚本中 TrainConfig(name="docker_train", exp_name="train") 的结果。

常用参数​

参数默认值说明
--batch_size1全局 batch；脚本会调整到能被 JAX 设备数整除。
--steps1000训练步数。
--gpusall使用全部 GPU，或填 0,1。
--prompt空数据集缺少任务文本时的默认指令。
--save_interval500checkpoint 保存间隔。
--learning_rate2.5e-5不传时使用封装默认峰值学习率。
--fsdp_devicesauto多卡时自动使用可见 GPU 数；单卡为 1。
--loraauto单卡默认启用 LoRA，多卡默认关闭。
--ema_decay空LoRA 模式下关闭 EMA 以节省显存。
--action_horizon50动作序列长度。
--num_workers8训练 DataLoader worker 数。

从官方源码复现​

如果需要修改数据映射、模型配置或训练策略，应回到 OpenPI 官方源码。官方 README 的核心流程是：安装 uv，同步依赖，定义或修改训练配置，先计算归一化统计，再启动训练。

git clone --recurse-submodules https://github.com/Physical-Intelligence/openpi.git
cd openpi

GIT_LFS_SKIP_SMUDGE=1 uv sync
GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .

官方以 LIBERO 为示例：

uv run scripts/compute_norm_stats.py --config-name pi05_libero

XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
  uv run scripts/train.py pi05_libero \
  --exp-name=my_experiment \
  --overwrite

要换成自己的机器人数据，需要修改官方 src/openpi/training/config.py 中的数据配置与 TrainConfig，并补齐输入/输出 transform。上述镜像内置一层通用 LeRobot 数据结构适配，适合先跑通训练；若要深度修改数据映射，仍应回到 OpenPI 官方源码实现自己的配置。

复现边界​

  • 镜像内基座权重在构建阶段下载到 /models/openpi-assets/checkpoints/*_base/params。
  • 运行时默认关闭 W&B 与 Hugging Face Hub 网络访问，适合离线训练环境。
  • 单卡默认启用 LoRA；多卡默认走 FSDP。若要做严格对比，请在实验记录中写明 --lora 与 --fsdp_devices。
  • Pi0.5 当前官方仓库说明训练/推理支持 flow matching head。不要把它理解为完整复现所有内部训练阶段。

常见问题​

报错找不到 LeRobot 数据集​

确认宿主机路径挂载到 /data/input 后，容器内存在 /data/input/meta/info.json。多数问题来自多挂了一层目录。

单卡和多卡结果为什么不同​

单卡默认 LoRA，多卡默认 FSDP，训练参数量和优化路径不同。若要比较模型本身，请固定 --lora、--fsdp_devices、--batch_size 和随机种子。

是否需要手动计算归一化统计​

使用 ioaitech/train_openpi:* 不需要，封装脚本会在训练前计算。使用官方源码时，需要按官方流程先运行 scripts/compute_norm_stats.py。

输出目录为什么固定​

当前封装固定 name=docker_train、exp_name=train。建议在宿主机输出路径中区分实验，例如 /outputs/pi05_pick_block_v1。

参考链接​