跳到主要内容

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 格式:

导出流程:

  1. 选择导出格式: 在艾欧平台中选择 LeRobot 或 HDF5 格式导出
  2. 筛选数据: 根据项目、时间、质量等级等维度筛选需要的数据
  3. 批量导出: 支持大规模数据集的批量导出处理
  4. 格式转换: 将导出的数据转换为 OpenPI 所需 LeRobot 格式

导出 LeRobot 格式数据

数据格式要求

OpenPI 支持多种数据格式,数据集需要包含:

  • 观察数据: 图像、状态信息等
  • 动作数据: 机器人执行的动作序列
  • 语言指令: 任务描述(可选,用于语言条件化任务)

数据集配置

OpenPI 使用配置文件来定义数据集和训练参数。您需要:

  1. 创建数据集配置: 在配置文件中定义数据路径、格式等
  2. 数据预处理: 计算归一化统计信息
  3. 验证数据: 确保数据格式正确
# 计算数据集的归一化统计信息(必需步骤)
uv run scripts/compute_norm_stats.py <config_name>

# 示例:为自定义配置计算统计信息
uv run scripts/compute_norm_stats.py my_custom_config

支持的数据格式

OpenPI 支持多种机器人数据格式,包括:

  • HuggingFace 数据集
  • LeRobot 格式数据
  • HDF5 格式数据
  • 从艾欧平台导出的数据(通过格式转换支持)

模型训练实践

使用艾欧平台进行训练(推荐)

艾欧数据平台提供完整的模型训练功能,支持 OpenPI 等多种机器人学习模型的训练:

艾欧平台训练功能概览

平台训练优势:

  • 无代码训练: 通过可视化界面完成整个训练流程,无需编程经验
  • 灵活算力选择: 支持私有云和公有云算力资源,按需分配
  • 实时监控: 提供训练指标、模型输出、系统日志的实时监控
  • 自动化管理: 支持检查点管理、断点续训、参数调整等功能

训练流程

  1. 数据准备: 选择已导出的数据集或上传外部数据 选择训练数据源

  2. 模型配置: 选择 Pi0/Pi0.5 模型并配置训练参数 训练参数配置

  3. 训练监控: 实时查看训练进度和模型表现 训练详情监控

  4. 模型导出: 训练完成后获取模型检查点 模型检查点管理

本地 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文件测试、离线边缘部署
  • 实时监控: 提供服务状态、资源使用、性能指标的实时监控
  • 智能适配: 自动识别并适应不同模型的输入输出要求

推理服务功能

  1. 服务状态监控: 查看推理服务的运行状态和配置信息 推理服务详情

  2. 模拟推理测试: 使用随机数据快速验证模型功能 模拟推理测试

  3. MCAP文件测试: 使用真实机器人数据验证推理效果 MCAP文件测试

  4. 离线边缘部署: 将推理服务部署到机器人本地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_simPi0 模型,ALOHA 仿真环境ALOHA 机器人仿真训练
pi05_droidPi0.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 版本:

  1. 推理: 大部分权重和计算使用 bfloat16,少数计算使用 float32 保证稳定性
  2. 训练: 默认混合精度,权重和梯度使用 float32,激活和计算使用 bfloat16

PyTorch 版本:

  1. 推理: 与 JAX 版本匹配,大部分使用 bfloat16
  2. 训练: 支持全 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

常见问题解决

内存和性能优化

  1. GPU 内存不足

    # 设置 JAX 内存分配比例
    export XLA_PYTHON_CLIENT_MEM_FRACTION=0.9

    # 使用 FSDP 减少内存使用
    uv run scripts/train.py <config> --exp_name <name> --fsdp-devices <n>
  2. 训练速度优化

    • 使用 SSD 存储提升数据加载速度
    • 适当增加 batch size 和 gradient accumulation
    • 考虑禁用 EMA(指数移动平均)以节省内存

数据相关问题

  1. 归一化统计错误

    # 重新计算归一化统计
    uv run scripts/compute_norm_stats.py <config_name>
  2. 数据集下载失败

    • 检查网络连接
    • 对于 HuggingFace 数据集,确保已登录:huggingface-cli login

环境配置问题

  1. 依赖冲突

    # 清理虚拟环境重新安装
    rm -rf .venv
    uv sync
  2. 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 模型的集成支持,可以:

  • 可视化对比:真实数据与模型输出的可视化比较
  • 关节级分析:每个关节的预测精度分析
  • 任务成功率统计:实际部署后的任务完成情况

评估指标

定量指标

  • 动作精度:预测动作与真实动作的误差统计
  • 轨迹一致性:生成轨迹的平滑性和合理性
  • 收敛分析:训练过程中损失函数的收敛情况

定性评估

  • 实机验证:在真实机器人上测试模型表现
  • 泛化能力:在未见过的场景中的适应能力
  • 鲁棒性测试:在噪声和干扰下的稳定性

艾欧平台完整工作流

端到端机器人学习流程

艾欧数据平台为 OpenPI 模型提供完整的端到端解决方案,从数据采集到模型部署的全链路支持:

平台集成优势

相比纯 OpenPI 开发的优势:

  1. 降低技术门槛:无需深度学习和编程经验,通过可视化界面操作
  2. 提升开发效率:从数据到模型的全流程自动化,大幅缩短开发周期
  3. 保证数据质量:专业的数据标注和质量控制流程
  4. 简化部署流程:一键部署到生产环境,支持多种部署方式
  5. 持续优化支持:提供模型性能监控和持续改进机制

使用建议

推荐使用场景:

  • 初学者或非技术用户:使用艾欧平台的可视化界面进行完整流程操作
  • 快速原型验证:需要快速验证想法和概念的研发团队
  • 生产环境部署:需要稳定可靠的生产级部署解决方案
  • 大规模数据处理:需要处理大量机器人数据的项目

本地开发场景:

  • 深度定制需求:需要修改模型架构或训练算法的高级用户
  • 研究实验:需要进行算法研究和实验对比的学术用户
  • 离线环境:网络受限或安全要求较高的环境

相关资源

常见问题解答

Q: 如何选择 JAX 还是 PyTorch 版本?

A:

  • JAX 版本:推荐用于生产环境,更稳定成熟
  • PyTorch 版本:适合需要 PyTorch 生态系统集成的场景

Q: Pi0.5 模型在 OpenPI 中如何使用?

A: Pi0.5 模型在 OpenPI 中完全支持,使用配置 pi05_droid 或类似的 Pi0.5 配置即可。

Q: 如何处理内存不足问题?

A:

# 设置 JAX 内存分配
export XLA_PYTHON_CLIENT_MEM_FRACTION=0.9

# 使用 FSDP 分布式训练
uv run scripts/train.py <config> --exp_name <name> --fsdp-devices <n>

Q: 训练数据需要什么格式?

A: OpenPI 支持多种数据格式,包括 HuggingFace 数据集格式。艾欧平台导出的数据需要进行格式转换。

Q: 如何监控训练进度?

A: OpenPI 会自动保存训练日志和检查点到 experiments/<exp_name>/ 目录,可以通过日志文件监控训练进度。

Q: 模型微调需要多长时间?

A: 训练时间取决于:

  • 数据集大小
  • 硬件配置(GPU 数量和型号)
  • 模型复杂度
  • 训练配置参数

通常在单张 A100 上需要数小时到数天不等。

Q: 如何将训练好的模型集成到艾欧平台?

A: 艾欧平台提供 OpenPI 模型的集成支持,包括模型注册、部署和监控功能。具体集成方式请咨询技术支持团队。


通过本指南,您应该能够成功使用 OpenPI 框架训练和部署 Pi0/Pi0.5 模型。如有任何问题,请参考官方文档或寻求技术支持。