Pi0 模型训练

Pi0 是一个通用机器人策略基础模型，专为解决机器人学习中的数据稀缺、泛化能力差和鲁棒性不足等核心挑战而设计。借鉴大语言模型的训练方法，Pi0 通过大规模预训练掌握广泛的机器人操作技能，并能通过微调快速适应具体任务需求。该模型支持高效的数据利用与快速部署，在有限数据下也能实现良好性能，为机器人智能化提供了一种可扩展的解决方案。

1. 环境安装

在LeRobot安装及使用教程的基础上，训练Pi0需要显存较大的显卡，我租赁的实例是 H20，使用VS Code SSH远程连接服务器：

# ssh -p 17825 root@region-42.seetacloud.com
Host region-42
    HostName region-42.seetacloud.com
    User root
    Port 17825

在安装Pi0的过程中需要访问Github和HuggingFace等国外网站，建议开启AutoDL内置的学术加速服务，由于该加速服务可能对正常网络造成一定影响，安装过程中如果失败可以多次尝试，当不再需要时建议取消学术加速。

# 启用学术加速
source /etc/network_turbo

cd /root/autodl-tmp
git clone https://github.com/enpeizhao/openpi_episode1_student.git
cd openpi_episode1_student/
pip install uv
GIT_LFS_SKIP_SMUDGE=1 uv sync
GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .

# 取消学术加速
unset http_proxy && unset https_proxy

2. 转换数据

Pi0训练必须是弧度制数据，安装ffmpeg转换为openpi格式：

sudo apt update
sudo apt install ffmpeg

uv run ./examples/libero/lerobot2oppi.py \
--source-repo-id=hou/demo_move_toy \
--target-repo-id=hou/demo_move_toy_openpi \
--output-path=./demo_move_toy_openpi \
--source-dataset-root=/root/autodl-tmp/demo_move_toy \
--max-episodes=100

3. 修改配置

以下是我的配置文件，为了节省训练时间num_train_steps设置为10_000：

# 单臂配置：openpi_episode1_student/src/openpi/training/config.py
TrainConfig(
    name="enpei_robot_demo_move_toy_low_mem_finetune",
    # Here is an example of loading a pi0 model for LoRA fine-tuning.
    model=pi0.Pi0Config(paligemma_variant="gemma_2b_lora", action_expert_variant="gemma_300m_lora"),
    data=LeRobotLiberoDataConfig(
        repo_id="hou/demo_move_toy_openpi", # 数据集repo_od
        root="./demo_move_toy_openpi", # 数据集路径
        base_config=DataConfig(prompt_from_task=True),
    ),
    weight_loader=weight_loaders.CheckpointWeightLoader("/root/autodl-tmp/pi0_base/params"),
    num_train_steps=10_000,
    # The freeze filter defines which parameters should be frozen during training.
    # We have a convenience function in the model config that returns the default freeze filter
    # for the given model config for LoRA finetuning. Just make sure it matches the model config
    # you chose above.
    freeze_filter=pi0.Pi0Config(
        paligemma_variant="gemma_2b_lora", action_expert_variant="gemma_300m_lora"
    ).get_freeze_filter(),
    # Turn off EMA for LoRA finetuning.
    ema_decay=None,
),

4. 模型训练

# 计算 normalization
uv run scripts/compute_norm_stats.py --config-name enpei_robot_demo_move_toy_low_mem_finetune
# 开启训练
XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 uv run scripts/train.py enpei_robot_demo_move_toy_low_mem_finetune --exp-name=my_experiment --overwrite

在训练过程中可以看到模型的LOSS在逐渐降低，也可以在终端中查看GPU的利用率，采集100个episodes数据使用H20迭代10_000步大概需要花费9个小时：

训练过程中Pi0会保存多个时间点的权重checkpoint，每次约消耗磁盘12G，因此建议将代码目录放到/root/autodl-tmp数据盘目录下，这样训练过程中产生的权重文件会自动保存在/root/autodl-tmp/checkpoints下，避免系统盘因空间不足导致训练失败。

5. 本文小结

本文介绍了在H20实例上安装和训练Pi0模型的完整流程，包括环境配置、数据转换、修改训练配置以及模型训练等关键步骤，并提供了相关命令和注意事项，帮助读者快速上手Pi0模型训练。