Pi0 与 Pi0.5 模型微调:Physical Intelligence 开源机器人学习平台
Pi0 和 Pi0.5 是 Physical Intelligence 团队开发的业界领先的视觉-语言-动作(Vision-Language-Action, VLA)模型系列。这些模型通过大规模机器人数据预训练,能够通过微调快速适应特定的机器人任务和应用场景。
本指南基于 Physical Intelligence 官方的 OpenPI 框架, 详细介绍如何使用预训练的 Pi0/Pi0.5 模型在自定义数据集上进行高效微调。
模型系列概述
OpenPI 仓库目前提供三种类型的模型:
模型类型 | 描述 | 特点 | 适用场景 |
---|---|---|---|
π₀ (Pi0) | 基于流匹配的视觉-语言-动作模型 | 流式生成,高质量动作预测 | 复杂操作任务,高精度要求 |
π₀-FAST | 基于 FAST 动作标记器的自回归 VLA | 自回归生成,快速推理 | 实时控制,低延迟需求 |
π₀.₅ (Pi0.5) | π₀ 的升级版本,具备更好的开放世界泛化能力 | 知识隔离训练,增强泛化 | 多样化环境,跨域应用 |
注意:所有模型都基于 10,000+ 小时的机器人数据进行预训练,提供基础模型检查点供微调使用。
先决条件与环境准备
系统要求
最低配置:
- Python 3.11+(推荐使用 uv 包管理器)
- GPU: 训练需要 NVIDIA GPU(推荐 A100/H100)
- 内存: 32GB+ 系统 RAM
- 存储: 100GB+ 可用磁盘空间
推荐配置:
- 硬件: NVIDIA A100/H100 或多卡设置
- 存储: NVMe SSD 固态硬盘
- 网络: 稳定的网络连接用于模型和数据下载
环境安装与配置
使用 uv 包管理器(推荐)
# 克隆 OpenPI 官方仓库
git clone https://github.com/Physical-Intelligence/openpi.git
cd openpi
# 使用 uv 安装依赖(自动创建虚拟环境)
uv sync
# 验证安装
uv run python -c "from openpi.policies import policy_config; print('OpenPI 环境配置成功!')"
传统 pip 安装
# 克隆仓库
git clone https://github.com/Physical-Intelligence/openpi.git
cd openpi
# 创建虚拟环境
python -m venv venv
source venv/bin/activate # Linux/macOS
# 或 venv\Scripts\activate # Windows
# 安装依赖
pip install -e .
# 验证安装
python -c "from openpi.policies import policy_config; print('OpenPI 环境配置成功!')"
数据集准备
使用艾欧平台导出数据
艾欧数据平台提供完整的数据导出功能,一键将标注数据导出为 OpenPI 支持的 LeRobot 格式:
导出流程:
- 选择导出格式: 在艾欧平台中选择 LeRobot 或 HDF5 格式导出
- 筛选数据: 根据项目、时间、质量等级等维度筛选需要的数据
- 批量导出: 支持大规模数据集的批量导出处理
- 格式转换: 将导出的数据转换为 OpenPI 所需 LeRobot 格式
数据格式要求
OpenPI 支持多种数据格式,数据集需要包含:
- 观察数据: 图像、状态信息等
- 动作数据: 机器人执行的动作序列
- 语言指令: 任务描述(可选,用于语言条件化任务)
数据集配置
OpenPI 使用配置文件来定义数据集和训练参数。您需要:
- 创建数据集配置: 在配置文件中定义数据路径、格式等
- 数据预处理: 计算归一化统计信息
- 验证数据: 确保数据格式正确
# 计算数据集的归一化统计信息(必需步骤)
uv run scripts/compute_norm_stats.py <config_name>
# 示例:为自定义配置计算统计信息
uv run scripts/compute_norm_stats.py my_custom_config
支持的数据格式
OpenPI 支持多种机器人数据格式,包括:
- HuggingFace 数据集
- LeRobot 格式数据
- HDF5 格式数据
- 从艾欧平台导出的数据(通过格式转换支持)
模型训练实践
使用艾欧平台进行训练(推荐)
艾欧数据平台提供完整的模型训练功能,支持 OpenPI 等多种机器人学习模型的训练:
平台训练优势:
- 无代码训练: 通过可视化界面完成整个训练流程,无需编程经验
- 灵活算力选择: 支持私有云和公有云算力资源,按需分配
- 实时监控: 提供训练指标、模型输出、系统日志的实时监控
- 自动化管理: 支持检查点管理、断点续训、参数调整等功能
训练流程
-
数据准备: 选择已导出的数据集或上传外部数据
-
模型配置: 选择 Pi0/Pi0.5 模型并配置训练参数
-
训练监控: 实时查看训练进度和模型表现
-
模型导出: 训练完成后获取模型检查点
本地 JAX 版本训练
如需在本地环境进行训练,OpenPI 主要使用 JAX 进行训练,提供最佳性能和稳定性:
# 设置 JAX 内存分配(推荐)
export XLA_PYTHON_CLIENT_MEM_FRACTION=0.9
# 基础训练命令
uv run scripts/train.py <config_name> --exp_name <实验名称>
# 示例:训练 Pi0 模型
uv run scripts/train.py pi0_aloha_sim --exp_name my_pi0_experiment
# 多GPU 训练(使用 FSDP)
uv run scripts/train.py <config_name> --exp_name <实验名称> --fsdp-devices <GPU数量>
# 示例:使用 4 张 GPU 训练
uv run scripts/train.py pi0_aloha_sim --exp_name my_experiment --fsdp-devices 4
PyTorch 版本训练
OpenPI 也支持 PyTorch 训练,但需要先进行环境配置:
PyTorch 环境配置
# 安装 PyTorch 支持(需要替换 transformers 文件)
cp -r src/openpi/models_pytorch/transformers_replace/* .venv/lib/python3.11/site-packages/transformers/
# 警告:这会永久修改 transformers 库,可能影响其他项目
PyTorch 训练命令
# 单 GPU 训练
uv run scripts/train_pytorch.py <config_name> --exp_name <实验名称>
# 多 GPU 训练(单节点)
uv run torchrun --standalone --nnodes=1 --nproc_per_node=<GPU数量> \
scripts/train_pytorch.py <config_name> --exp_name <实验名称>
# 示例:使用 2 张 GPU 训练
uv run torchrun --standalone --nnodes=1 --nproc_per_node=2 \
scripts/train_pytorch.py pi0_aloha_sim --exp_name pytorch_experiment
模型推理与部署
使用艾欧平台进行推理部署(推荐)
艾欧数据平台提供完整的模型推理服务,支持从云端验证到边缘部署的全场景AI推理 能力:
平台推理优势:
- 一键部署: 将训练完成的模型一键部署为生产级推理服务
- 多种测试方式: 支持模拟推理、MCAP文件测试、离线边缘部署
- 实时监控: 提供服务状态、资源使用、性能指标的实时监控
- 智能适配: 自动识别并适应不同模型的输入输出要求
推理服务功能
-
服务状态监控: 查看推理服务的运行状态和配置信息
-
模拟推理测试: 使用随机数据快速验证模型功能
-
MCAP文件测试: 使用真实机器人数据验证推理效果
-
离线边缘部署: 将推理服务部署到机器人本地GPU
部署方式对比
推理方式 | 适用场景 | 说明 |
---|---|---|
模拟推理测试 | 快速验证 | 使用随机数据或自定义输入,快速验证模型推理功能和性能 |
MCAP文件测试 | 真实数据验证 | 使用录制好的机器人演示数据,验证模型在真实场景下的推理效果 |
离线边缘部署 | 生产环境应用 | 将推理服务部署到机器人本地GPU,实现低延迟的实时控制 |
本地策略服务器部署
如需在本地环境部署,OpenPI 提供策略服务器用于模型推理部署:
# 启动策略服务器(JAX 版本)
uv run scripts/serve_policy.py policy:checkpoint \
--policy.config=<配置名称> \
--policy.dir=<检查点路径>
# 示例:部署 Pi0.5 模型
uv run scripts/serve_policy.py policy:checkpoint \
--policy.config=pi05_droid \
--policy.dir=/path/to/trained/checkpoint
# PyTorch 版本部署
uv run scripts/serve_policy.py policy:checkpoint \
--policy.config=pi05_droid \
--policy.dir=/path/to/converted/pytorch/checkpoint
Python API 使用
from openpi.training import config as _config
from openpi.policies import policy_config
# 加载配置
config = _config.get_config("pi05_droid")
checkpoint_dir = "/path/to/trained/checkpoint"
# 创建训练好的策略
policy = policy_config.create_trained_policy(config, checkpoint_dir)
# 执行推理
action_chunk = policy.infer(example)["actions"]
可用配置
OpenPI 提供多个预定义配置,适用于不同的机器人平台和任务:
配置名称 | 描述 | 适用场景 |
---|---|---|
pi0_aloha_sim | Pi0 模型,ALOHA 仿真环境 | ALOHA 机器人仿真训练 |
pi05_droid | Pi0.5 模型,DROID 数据集 | 真实机器人数据训练 |
debug | 调试配置,小数据集 | 快速测试和调试 |
自定义配置
您可以创建自定义配置文件来适应特定的数据集和训练需求:
# 在 src/openpi/training/config.py 中添加新配置
@config_flags.DEFINE_config_file("config")
def get_config(config_name: str):
# 基于现有配置修改
config = get_base_config()
# 自定义数据集路径
config.dataset.path = "/path/to/your/dataset"
# 调整训练参数
config.training.learning_rate = 1e-4
config.training.batch_size = 8
return config
JAX 与 PyTorch 版本对比
精度设置
JAX 版本:
- 推理: 大部分权重和计算使用 bfloat16,少数计算使用 float32 保证稳定性
- 训练: 默认混合精度,权重和梯度使用 float32,激活和计算使用 bfloat16
PyTorch 版本:
- 推理: 与 JAX 版本匹配,大部分使用 bfloat16
- 训练: 支持全 bfloat16(默认)或全 float32,通过
pytorch_training_precision
配置
性能比较
- 推理速度: 使用 torch.compile 后,PyTorch 与 JAX 性能相当
- 内存使用: JAX 版本通常更高效
- 稳定性: JAX 版本更成熟稳定
模型转换
如需在 JAX 和 PyTorch 之间转换模型:
# JAX 模型转换为 PyTorch
uv run examples/convert_jax_model_to_pytorch.py \
--checkpoint_dir /path/to/jax/checkpoint \
--config_name <配置名称> \
--output_path /path/to/pytorch/checkpoint
常见问题解决
内存和性能优化
-
GPU 内存不足:
# 设置 JAX 内存分配比例
export XLA_PYTHON_CLIENT_MEM_FRACTION=0.9
# 使用 FSDP 减少内存使用
uv run scripts/train.py <config> --exp_name <name> --fsdp-devices <n> -
训练速度优化:
- 使用 SSD 存储提升数据加载速度
- 适当增加 batch size 和 gradient accumulation
- 考虑禁用 EMA(指数移动平均)以节省内存
数据相关问题
-
归一化统计错误:
# 重新计算归一化统计
uv run scripts/compute_norm_stats.py <config_name> -
数据集下载失败:
- 检查网络连接
- 对于 HuggingFace 数据集,确保已登录:
huggingface-cli login
环境配置问题
-
依赖冲突:
# 清理虚拟环 境重新安装
rm -rf .venv
uv sync -
CUDA 错误:
- 验证 NVIDIA 驱动安装正确
- 对于 Docker,确保安装了 nvidia-container-toolkit
- 可能需要卸载系统级 CUDA 库以避免冲突
训练监控与评估
训练日志
OpenPI 训练过程会自动记录详细日志,包括:
- 损失曲线:训练损失变化趋势
- 学习率调度:学习率变化情况
- 资源使用:GPU 内存和计算利用率
- 训练指标:各种评估指标
检查点管理
# 训练检查点保存在实验目录中
ls experiments/<exp_name>/checkpoints/
# 恢复训练(从最新检查点继续)
uv run scripts/train.py <config> --exp_name <name> --resume
# 从特定检查点恢复
uv run scripts/train.py <config> --exp_name <name> --resume --checkpoint_path <path>
模型评估
使用训练好的模型进行推理评估:
from openpi.policies import policy_config
from openpi.training import config as _config
# 加载配置和模型
config = _config.get_config("pi05_droid")
policy = policy_config.create_trained_policy(config, checkpoint_dir)
# 执行推理
result = policy.infer(observation_data)
actions = result["actions"]
艾欧平台集成评估
艾欧平台提供与 OpenPI 模型的集成支持,可以:
- 可视化对比:真实数据与模型输出的可视化比较
- 关节级分析:每个关节的预测精度分析
- 任务成功率统计:实际部署后的任务完成情况