使用 Pi0 模型微调数据

Pi0 是一个强大的视觉-语言-动作（VLA）模型，通过少量数据微调，就能快速适配您的机器人任务。

本文档基于 LeRobot 框架，使用预训练的 Pi0 模型（lerobot/pi0）在您自己的机器人数据集上进行微调（Fine-tuning）。

准备条件

• 您已经拥有一个符合 LeRobot 数据格式的数据集（参见上一节的导出LeRobot数据集）
• 希望让 Pi0 模型适配您的机器人硬件、任务或控制风格
• 熟悉 PyTorch 和 Hugging Face 生态进行训练
• 具备基本的深度学习训练经验

什么是微调（Fine-tuning）？

微调（Fine-tuning）是指在预训练模型的基础上，用您自己的数据继续训练模型的一小段时间。

Pi0 是一个在多个通用机器人任务上预训练的模型，它已经学到了很多通用的视觉理解、指令理解、动作预测能力。

通过微调，您可以让 Pi0：

• 适应您的机器人环境：适配您的摄像头视角、光照环境、机械结构
• 适配特定任务：针对您的具体任务（比如拿笔、分类、放置）进行优化
• 提高控制精度：显著提升模型在您任务中的控制精度与成功率

简单来说：预训练模型懂很多，但不一定懂您；微调之后，它就懂您了。

环境准备

系统要求

请确保您的环境满足以下要求：

• Python ≥ 3.8（推荐 3.10 或 3.11）
• GPU：建议至少 32GB 显存（Pi0是一个参数量较大的模型，推荐 V100 以上级别的显卡）
• 内存：建议至少 64GB RAM
• 存储：确保有足够的磁盘空间存储数据集和模型

安装依赖

# 克隆 LeRobot 仓库
git clone https://github.com/huggingface/lerobot.git
cd lerobot/

# 安装 LeRobot 框架（包含 Pi0 支持）
pip install -e ".[pi0]"

# 验证安装
python -c "from lerobot.policies import Pi0Policy; print('Pi0 安装成功！')"

准备您的数据集

导出 LeRobot 格式数据

可以通过艾欧数据平台的导出页面，将标注后的数据导出为 LeRobot 格式数据集，参见 导出LeRobot数据集。

重要提示：不需要将您的数据上传到 HuggingFace 也能训练，通过参数 --dataset.root= 来指定使用本地目录的数据进行微调训练。

数据集结构示例

假设您的导出数据存储在 ~/DualPiper_Pickup_Pen：

$ ls ~/DualPiper_Pickup_Pen
data  meta  videos

$ ls ~/DualPiper_Pickup_Pen/data
episode_000  episode_001  episode_002  ...

$ ls ~/DualPiper_Pickup_Pen/data/episode_000
observations  actions.npy  language_instruction.txt  metadata.json

开始微调（Fine-tune）

基本训练命令

使用 train.py 脚本进行训练：

# 设置 CUDA 内存分配策略（推荐）
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

# 启动训练
python3 -m lerobot.scripts.train \
  --dataset.root=~/DualPiper_Pickup_Pen \
  --dataset.repo_id=io-ai-data/DualPiper_Pickup_Pen \
  --policy.path=lerobot/pi0 \
  --policy.repo_id=lerobot/pi0 \
  --policy.device=cuda \
  --output_dir=./checkpoints/pi0_finetuned \
  --batch_size=1 \
  --policy.attention_implementation=flash_attention_2 \
  --training.learning_rate=1e-5 \
  --training.num_epochs=10

提示

Pi0模型本身比较大，所以单卡训练至少要数十个小时。比如我们对采集的2小时数据做微调，在单卡V100上训练时长50个小时，8卡V100速度会提升，但也需要将近10个小时。

参数详细说明

参数	含义	推荐值	说明
--dataset.root	本地数据集路径	~/your_dataset	指向您的 LeRobot 数据集目录
--dataset.repo_id	Hugging Face 数据集名称	your-username/dataset	用于标识数据集
--policy.path	预训练模型路径	lerobot/pi0	使用官方预训练模型
--policy.repo_id	模型仓库 ID	lerobot/pi0	Hugging Face 上的模型仓库
--policy.device	训练设备	cuda	使用 GPU 加速训练
--output_dir	模型保存目录	./checkpoints/pi0_finetuned	微调后模型的保存路径
--batch_size	批次大小	1	根据显存大小调整
--policy.attention_implementation	注意力机制实现	flash_attention_2	使用 Flash Attention 2 加速
--training.learning_rate	学习率	1e-5	微调时使用较小的学习率
--training.num_epochs	训练轮数	10	根据数据量调整

多卡加速训练

如果您有多张显卡，可以并行训练来大幅加快速度。

# 使用 8张 GPU 训练
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
torchrun --nproc_per_node=8 --master_port=29500 \
    -m lerobot.scripts.train \
    --dataset.root=/data/nfs2/export/lerobot/DualPiper_Pickup_Pen/ \
    --dataset.repo_id=io-ai-data/DualPiper_Pickup_Pen \
    --policy.path=lerobot/pi0 \
    --policy.repo_id=lerobot/pi0 \
    --policy.device=cuda \
    --output_dir=/home/puxk/DualPiper_Pickup_Pen_pi0_model \
    --batch_size=1 \
    --num_workers=4 \
    --policy.attention_implementation=flash_attention_2

总 batch size 是 nproc_per_node × batch_size ，设置了 --batch_size=1，8 卡等于总 batch size 为 8，如果内存足够可以调大。

调试时建议先用 --nproc_per_node=2 测试是否能跑通，再放开全部 8 卡。

特别说明

1. 网络加速：可能需要网络加速来下载 HuggingFace 上的 Pi0 模型，参考 hg-mirror 等工具
2. 显存监控：训练过程中请监控 GPU 显存使用情况
3. 断点续训：如果训练中断，可以从 --output_dir 中的 checkpoint 继续训练

Pi0 模型内部使用了 Google PaLI-Gemma 模型结构进行语言处理。由于它是最近新增的 Hugging Face 模型：

版本要求

• Transformers：确保使用最新版 transformers（推荐 ≥ 4.37.0）
• 兼容性检查：如果遇到 embed_tokens 相关报错，请升级 transformers

解决方案

# 升级到最新版 transformers
pip install -U transformers

# 或者指定版本
pip install transformers>=4.37.0

查看训练结果

模型保存位置

训练完成后，模型将保存在：

./checkpoints/pi0_finetuned/
├── config.json
├── pytorch_model.bin
├── tokenizer.json
└── training_args.bin

加载微调后的模型

from lerobot.policies import Pi0Policy

# 加载微调后的模型
policy = Pi0Policy.from_pretrained("./checkpoints/pi0_finetuned")

# 进行推理
action = policy.select_action(batch)

训练日志分析

训练过程中会生成详细的日志，包括：

• 损失曲线：训练损失和验证损失的变化
• 学习率调度：学习率的变化情况
• 梯度信息：梯度范数和梯度裁剪情况
• 内存使用：GPU 显存和系统内存使用情况

如何评估效果？

艾欧数据平台评测

艾欧数据平台支持模型推理质量评测，可以用 任意真实数据 与 模型推理结果 做比较，可以对每个关节的指令做可视化的对比。

其他评估方法

您也可以通过以下方式评估 Pi0 在您任务上的效果：

1. 定量评估

• 动作误差：对比预测动作与真实动作（如角度误差、位置误差）
• 轨迹对比：可视化预测 vs ground truth 行为轨迹
• 损失指标：记录模型在验证集上的 loss 变化趋势

2. 定性评估

• 实机部署：在真实机器人上验证任务成功率
• 行为观察：观察模型生成的动作是否符合预期
• 用户反馈：收集实际使用中的用户反馈

3. 评估脚本示例

import numpy as np
from lerobot.policies import Pi0Policy

# 加载模型
policy = Pi0Policy.from_pretrained("./checkpoints/pi0_finetuned")

# 评估函数
def evaluate_model(policy, test_dataset):
    total_error = 0
    num_samples = 0
    
    for batch in test_dataset:
        predicted_action = policy.select_action(batch)
        true_action = batch['actions']
        
        # 计算动作误差
        error = np.mean(np.abs(predicted_action - true_action))
        total_error += error
        num_samples += 1
    
    return total_error / num_samples

# 执行评估
avg_error = evaluate_model(policy, test_dataset)
print(f"平均动作误差: {avg_error:.4f}")

相关资源

• LeRobot 项目主页
• Hugging Face Pi0 模型
• 原始 Pi0 实现 (JAX)
• Pi0 论文
• LeRobot 文档

常见问题

Q: 微调过程中模型报错 embed_tokens 不存在？

A: 这是因为 transformers 版本太老，升级即可：

pip install -U transformers

Q: 数据集没有上传 Hugging Face，可以用本地的吗？

A: 可以，使用 --dataset.root=你的本地路径 参数即可。

Q: 训练时显存不足怎么办？

A: 可以尝试以下方法：

# 1. 减小批次大小
--batch_size=1

还有其它参数，需要根据实际情况决定。

Q: 训练时间大概需要多久？

A: 训练时间取决于：

• 数据量：通常 10 小时以上
• 硬件配置：GPU 性能影响很大
• 训练轮数：建议 5-20 轮

Q: 如何判断训练是否收敛？

A: 观察以下指标：

• 损失曲线：训练损失和验证损失趋于稳定
• 验证指标：在验证集上的性能不再提升
• 过拟合：验证损失开始上升时停止训练

Q: 微调后的模型如何部署？

A: 艾欧数据平台已经支持自动部署 Pi0 和 SmolVLA 等机器人领域常见模型，具体请咨询我们的技术专家。

您也可以参考以下步骤：

# 1. 加载模型
policy = Pi0Policy.from_pretrained("./checkpoints/pi0_finetuned")

# 2. 准备输入数据
observation = {
    'images': camera_images,
    'language_instruction': "拿起红色的笔"
}

# 3. 生成动作指令
action = policy.select_action(observation)

祝您训练顺利！