对开源VLA sota π0的微调——如何基于各种开源数据集、以及你自己的私有数据集微调π0(含我司的微调实践)

前言

25年2.4日，几个月前推出π0的公司Physical Intelligence (π)宣布正式开源π0及π0-FAST，如之前所介绍的，他们对用超过 10,000 小时的机器人数据进行了预训练

该GitHub代码仓库「 π0及π0-FAST的GitHub地址：github.com/Physical-Intelligence/openpi」包括4个方面：简言之，就是π0本身的代码和权重、特定平台上特定任务的微调checkpoint、推理代码、微调代码

本文接上一篇文章《π0源码剖析——从π0模型架构的实现(如何基于PaLI-Gemma和扩散策略去噪生成动作)，到基于C/S架构下的模型训练与部署》而来，但本文侧重对π0的微调

第一部分 如何基于各种开源数据集微调π0

1.1 π0及π0-FAST对外开源的具体内容

1.1.1 开源基础模型π0及π0-FAST，可自行微调

1. π0基础模型的代码及权重
   可用于微调的扩散π0「标准的预训练π 0模型，该模型在 OXE 和他们的 7 个机器人平台上进行训练」
   其对应的checkpoint路径为：s3://openpi-assets/checkpoints/pi0_base
2. π0-FAST基础模型的代码及权重
   可用于微调的π0_FAST——基于FAST分词器的自回归模型
   该模型使用FAST tokenizer通过自回归离散化实现控制
   它提供了更好的语言跟踪性能，但推理成本更高(根据他们的经验，大约高出 4-5 倍)，如果您更喜欢使用discretization而不是流匹配，这是一个不错的选择
   
   其对应的checkpoint路径为：s3://openpi-assets/checkpoints/pi0_fast_base

1.1.2 已经微调好的模型 可直接推理：π0 DROID/π0 ALOHA/π0 Libero

且他们还提供了专门为ALOHA和DROID平台上一些简单任务做了微调的checkpoint，相当于在ALOHA 和 DROID 收集的相对较小的数据集上进行了微调

即several checkpoints that are fine-tuned for a few simple tasks on a few widely available platforms such as ALOHA and DROID，当然，它们可能无法推广到您的特定设置

1. π0 DROID：在DROID 数据集上微调过的扩散π0
   DROID数据集由 Franka 机械臂在不同环境中执行的不同任务组成的开源数据集，且他们通过视频展示了 openpi 在训练数据中从未见过的环境中运行，这些环境遍布世界各地，包括蒙特利尔大学、华盛顿大学、韩国科学技术研究院等
   
   其对应的checkpoint路径为：s3://openpi-assets/checkpoints/pi0_droid
   推理速度比π0-FAST-DROID快，但可能不遵循语言命令
2. π0-FAST DROID：在DROID数据集微调过的π0-FAST
   可以在DROID机器人平台上的新场景中执行各种简单的零样本桌面操控任务，例如“从烤面包机中取出面包”任务
   
   其对应的checkpoint路径为：s3://openpi-assets/checkpoints/pi0_fast_droid
3. π0 ALOHA
   根据 ALOHA(适合灵巧操作的低成本双臂系统) 数据进行了微调，可以在ALOHA机器人平台上进行毛巾折叠、食物舀取和其他任务，相当于提供了一套针对 ALOHA 平台上的任务进行微调的检查点「这些检查点可能对整体机器人设置非常敏感，但能够在完全未出现在训练数据中的全新 ALOHA 站点上运行它们」
   
   不同任务对应的不同checkpoint路径分别为
   折叠毛巾：s3://openpi-assets/checkpoints/pi0_aloha_towel
   从容器中取出食物：s3://openpi-assets/checkpoints/pi0_aloha_tupperware
   打开笔帽：s3://openpi-assets/checkpoints/pi0_aloha_pen_uncap
4. π0 Libero
   此检查点针对 Libero 基准进行了微调，并且可以在 Libero 任务上进行开箱即用的评估

1.1.3 模型推理与私有数据自行微调

他们还开源了在多个显示世界和仿真机器人平台上推理运行的示例代码(example code to run inference on several real-world and simulated robot platforms)

 以下是跑π0-FAST-DROID模型的预训练检查点

from openpi.training import config
from openpi.policies import policy_config
from openpi.shared import download

config = config.get_config("pi0_fast_droid")
checkpoint_dir = download.maybe_download("s3://openpi-assets/checkpoints/pi0_fast_droid")

# Create a trained policy.
policy = policy_config.create_trained_policy(config, checkpoint_dir)

# Run inference on a dummy example.
example = {
    "observation/exterior_image_1_left": ...,
    "observation/wrist_image_left": ...,
    ...
    "prompt": "pick up the fork"
}
action_chunk = policy.infer(example)["actions"]

且也可以在示例笔记本中测试这一点，他们还提供了在DROID和ALOHA机器人上运行预先训练的检查点的推理的详细分步示例

此外

1. 远程推理：他们提供了示例和代码，用于远程运行模型推理，详见远程运行π0模型
   模型可以在不同的服务器上运行，并通过 websocket 连接将操作传输到机器人。这样可以轻松在机器人外使用更强大的 GPU，并将机器人和策略环境分开
   
   要启动远程策略服务器，您只需运行以下命令：

   uv run scripts/serve_policy.py --env=[DROID | ALOHA | LIBERO]

   其中，参数env指定哪个π0 checkpoint 应该被加载，比如是ALOHA还是LIBERO。在后台，这个脚本将执行如下命令，你可以使用它来启动策略服务器，例如用于你自己训练的检查点——这里是 DROID 环境的一个示例

   uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi0_fast_droid --policy.dir=s3://openpi-assets/checkpoints/pi0_fast_droid

   其中，config这将启动一个策略服务器，该服务器将执行和参数指定的策略dir。该策略将在指定端口(默认值：8000)上执行
2. 无需机器人即可测试推理：提供无需机器人即可测试推理的脚本。此脚本将生成随机观察并使用模型运行推理。有关更多详细信息，请参阅此处

1.2 基于自己的数据集微调π0基础模型

1.2.1 利用Libero数据集微调π0基础模型的三个步骤

此外，他们还提供了用于根据用户自身任务和平台微调π0的代码(code for fine-tuning the base π0 model for your own tasks and platforms)，这个微调π0的代码 个人觉得很有价值，且huggingface上有相应的pytorch接口

Physical Intelligence (π)认为，1 到 20 小时的数据足以微调各种任务，具体而言，如果想利用自己的数据上微调π0基础模型，只需以下三个步骤即可：

1. 将Libero数据集转换为LeRobot数据集v2.0格式
   作者提供了一个脚本convert_libero_data_to_lerobot.py「下一小节，我会详细分析一下这个脚本」，用于将 Libero 数据转换为 LeRobot 数据集 v2.0 格式，且可以轻松修改它以转换您自己的数据
   
   比如从此处下载原始 Libero 数据集，然后使用以下命令运行脚本

   uv run examples/libero/convert_libero_data_to_lerobot.py --data_dir /path/to/your/libero/data

   且作为示例，已将 BiPlay 代码库中的 aloha_pen_uncap_diverse_raw 数据集转换，并将其上传至 HuggingFace Hub，地址为 physical-intelligence/aloha_pen_uncap_diverse
2. 定义使用自定义数据集的训练配置，并运行训练
   提供pi0_aloha_pen_uncap 配置作为示例，您应该参考根README以了解如何使用新配置运行训练
   
   比如下面是 Libero 示例配置，可以根据自己的数据集进行修改：
     LiberoInputs和LiberoOutputs：定义从 Libero 环境到模型的数据映射，反之亦然——将用于训练和推理
     LeRobotLiberoDataConfig：定义了如何处理来自 LeRobot 数据集的原始 Libero 数据以用于训练
     TrainConfig：定义微调超参数、数据配置和权重加载器
3. 启动策略服务器并运行推理
   训练完成后，可以通过启动策略服务器，然后从 Libero 评估脚本中查询它来运行推理。启动模型服务器很容易（他们在此示例中使用迭代 20,000 的检查点，根据需要进行修改）

   uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi0_fast_libero --policy.dir=checkpoints/pi0_fast_libero/my_experiment/20000

Libero数据集转换脚本convert_libero_data_to_lerobot.py的解析

---

​接下来，我们使用Libero数据集（存储在RLDS格式中）作为示例——将RLDS格式转换为LeRobot格式，如何修改以适应其他自定义格式的数据

首先，咱们明确下Libero数据集的层级结构

Libero数据集
├── libero_10_no_noops/
├── libero_goal_no_noops/
├── libero_object_no_noops/
└── libero_spatial_no_noops/
    └── train/
                └── episode
                            └── steps
                                      ├── observation
                                      │      ├── image (256x256x3)
                                      │      ├── wrist_image (256x256x3)
                                      │      └── state (8维向量)
                                      ├── action (7维向量)
                                      └── language_instruction (文本)

接下来，按如下步骤逐一执行

1. 关键常量定义

   # 输出数据集名称，也用于Hugging Face Hub
   REPO_NAME = "your_hf_username/libero"  
   RAW_DATASET_NAMES = [
       "libero_10_no_noops",
       "libero_goal_no_noops",
       "libero_object_no_noops",
       "libero_spatial_no_noops",
   ]  # 将多个Libero数据集合并为一个训练数据集

2. 定义主函数结构

   def main(data_dir: str, *, push_to_hub: bool = False):

   主函数接受两个参数：
   data_dir: 原始数据目录路径
   push_to_hub: 是否推送到HuggingFace Hub的标志
3. 清理现有数据

       output_path = LEROBOT_HOME / REPO_NAME
       if output_path.exists():
           shutil.rmtree(output_path)

   确保输出目录干净，删除任何已存在的数据
4. 创建LeRobot数据集：包含处理图像、状态和动作数据

       # 创建LeRobot数据集，定义要存储的特征
       # OpenPi假设本体感知数据存储在`state`中，动作存储在`action`中
       # LeRobot假设图像数据的dtype为`image`
       dataset = LeRobotDataset.create(
           repo_id=REPO_NAME,           # 数据集的仓库ID
           robot_type="panda",          # 机器人类型
           fps=10,                      # 帧率
           features={                   # 特征定义
               "image": {                  # image数据
                   "dtype": "image",                             # 数据类型
                   "shape": (256, 256, 3),                       # 数据形状
                   "names": ["height", "width", "channel"],      # 维度名称
               },
               "wrist_image": {            # 手腕image数据
                   "dtype": "image", 
                   "shape": (256, 256, 3),  
                   "names": ["height", "width", "channel"], 
               },
               "state": {                  # 状态数据
                   "dtype": "float32",  
                   "shape": (8,), 
                   "names": ["state"],  
               },
               "actions": {                # 动作数据
                   "dtype": "float32", 
                   "shape": (7,),  
                   "names": ["actions"],  
               },
           },
           image_writer_threads=10,           # 图像写入线程数
           image_writer_processes=5,          # 图像写入进程数
       )

   定义数据集结构，包括：
   - 机器人类型：panda
   - 帧率：10fps
   - 特征定义：
     - 图像：256x256x3的RGB图像
     - 手腕图像：同样的格式
     - 状态：8维浮点数向量
     - 动作：7维浮点数向量
   - 多线程/进程配置用于图像处理
5. 数据转换循环

       # 遍历原始Libero数据集并将episode写入LeRobot数据集
       # 你可以根据自己的数据格式修改此部分
       for raw_dataset_name in RAW_DATASET_NAMES:
           raw_dataset = tfds.load(raw_dataset_name, data_dir=data_dir, split="train")  # 加载原始数据集
           for episode in raw_dataset:  # 遍历每个episode
               for step in episode["steps"].as_numpy_iterator():      # 遍历每个步骤
                   dataset.add_frame(
                       {
                           "image": step["observation"]["image"],      # 添加图像数据
                           "wrist_image": step["observation"]["wrist_image"],  # 添加手腕图像数据
                           "state": step["observation"]["state"],      # 添加状态数据
                           "actions": step["action"],                  # 添加动作数据
                       }
                   )
               dataset.save_episode(task=step["language_instruction"].decode())  # 保存episode并解码语言指令

   这部分实现了数据转换的核心逻辑：
   1. 遍历所有原始数据集
   2. 对每个数据集中的每个episode
   3. 处理episode中的每一步
   4. 保存frame数据和语言指令
6. 数据集的整合与发布
   最后的处理步骤：
   1. 整合数据集（不计算统计信息）
   2. 可选：推送到HuggingFace Hub，包括：
      - 添加标签
      - 设置为公开
      - 包含视频
      - 指定许可证

脚本可以通过以下命令运行：

uv run examples/libero/convert_libero_data_to_lerobot.py --data_dir /path/to/your/data

如果要推送到HuggingFace Hub：

uv run examples/libero/convert_libero_data_to_lerobot.py --data_dir /path/to/your/data --push_to_hub

总之，这个脚本展示了如何构建标准化的机器人学习数据集，为训练像Pi0这样的模型提供数据支持

1.2.2 在 UR5 数据集上微调 pi0

// 待更

1.3 安装与运行

1.3.1 如何安装本openpi开源库

要运行此存储库中的模型，需要至少具有以下规格的 NVIDIA GPU。这些估算假设单个 GPU，但您也可以通过fsdp_devices在训练配置中进行配置来使用具有模型并行性的多个 GPU，以减少每个 GPU 的内存要求。另请注意，当前的训练脚本尚不支持多节点训练

模式	所需内存	示例 GPU
推理	> 8 GB	RTX 4090
微调（LoRA）	> 22.5 GB	RTX 4090
微调（完整版）	> 70 GB	A100（80GB）/H100

PS，他们说该 repo 已在 Ubuntu 22.04 上测试过，其他操作系统可能不支持

以下是安装过程

1. 克隆此 repo 时，请确保更新子模块：

   git clone --recurse-submodules [email protected]:Physical-Intelligence/openpi.git
   
   # Or if you already cloned the repo:
   git submodule update --init --recursive

2. 使用uv来管理 Python 依赖项。可参阅uv 安装说明进行设置
   安装 uv 后，运行以下命令设置环境：

   GIT_LFS_SKIP_SMUDGE=1 uv sync

   注意：GIT_LFS_SKIP_SMUDGE=1需要将 LeRobot 作为依赖项
   Docker：作为 uv 安装的替代方案，他们还提供了使用 Docker 安装 openpi 的说明
   如果在系统设置中遇到问题，还可以考虑使用 Docker 来简化安装，详情请参阅Docker 设置

2.3.2 如何把该库的ALOHA sim环境跑起来

此外，该开源库还提供ALOHA SIM

• 如使用docker

  export SERVER_ARGS="--env ALOHA_SIM"
  docker compose -f examples/aloha_sim/compose.yml up --build

• 如没有docker，则采用C/S架构「我在上一篇文章中也解读过了这个C/S架构了，详见此文的《π0源码剖析——从π0模型架构的实现(如何基于PaLI-Gemma和扩散策略去噪生成动作)，到基于C/S架构下的模型训练与部署》的第三部分 模型的训练与部署：基于客户端-服务器C/S架构」
    终端窗口1：建立策略客户端

  # Create virtual environment
  uv venv --python 3.10 examples/aloha_sim/.venv
  source examples/aloha_sim/.venv/bin/activate
  uv pip sync examples/aloha_sim/requirements.txt
  uv pip install -e packages/openpi-client
  
  # Run the simulation
  MUJOCO_GL=egl python examples/aloha_sim/main.py

  注意：如果您看到 EGL 错误，则可能需要安装以下依赖项

  sudo apt-get install -y libegl1-mesa-dev libgles2-mesa-dev

    终端窗口2：建立策略服务器

  # Run the server
  uv run scripts/serve_policy.py --env ALOHA_SIM

// 待更

第二部分 如何基于你自己的私有数据集微调π0：含我司的微调实践

// 待更