前言
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的微调
而截止到25年3月底,在我司「七月在线」的具身落地中,过去半年用π0居多,其次idp3和其他模型——也是目前国内具身落地经验最丰富的团队之一了(每个月都在并行开发多个具身订单,全职开发者和合作开发者总计百余人,还有我联合带的来自多个985的十几个具身研究生)
PS,有兴趣或也在对π0做微调的,欢迎私我一两句自我简介,邀请进:『七月具身:π0复现微调交流群』,更多见《具身智能机器人复现实战营 [含UMI/DexCap/iDP3/π0]》
第一部分 如何基于各种开源数据集微调π0
1.1 π0及π0-FAST对外开源的具体内容:2套代码、4+3套权重且支持远程推理
1.1.1 开源基础模型π0及π0-FAST,可自行微调
- π0基础模型的代码及权重
可用于微调的扩散π0「标准的预训练π 0模型,该模型在 OXE 和他们的 7 个机器人平台上进行训练」
其对应的checkpoint路径为:s3://openpi-assets/checkpoints/pi0_base
注意,这个π0基础版 本身就已经具备了一定的指令跟随能力 - π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,当然,它们可能无法推广到你的特定设置
- π0 DROID:在DROID 数据集上微调过的扩散π0
DROID数据集由 Franka 机械臂在不同环境中执行的不同任务组成的开源数据集,且他们通过视频展示了 openpi 在训练数据中从未见过的环境中运行,这些环境遍布世界各地,包括蒙特利尔大学、华盛顿大学、韩国科学技术研究院等
其对应的checkpoint路径为:s3://openpi-assets/checkpoints/pi0_droid
推理速度比π0-FAST-DROID快,但可能不遵循语言命令 - π0-FAST DROID:在DROID数据集微调过的π0-FAST
可以在DROID机器人平台上的新场景中执行各种简单的零样本桌面操控任务,例如“从烤面包机中取出面包”任务
其对应的checkpoint路径为:s3://openpi-assets/checkpoints/pi0_fast_droid - π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 - π0 Libero
此检查点针对 Libero 基准进行了微调,并且可以在 Libero 任务上进行开箱即用的评估
1.1.3 模型推理步骤:可连远程GPU进行推理
他们还开源了在多个显示世界和仿真机器人平台上推理运行的示例代码(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机器人上运行预先训练的检查点的推理的详细分步示例
此外
- 远程推理:他们提供了示例和代码,用于远程运行模型推理,详见远程运行π0模型
模型可以在不同的服务器上运行,并通过 websocket 连接将操作传输到机器人。这样可以轻松在机器人外使用更强大的 GPU,并将机器人和策略环境分开
要启动远程策略服务器,只需运行以下命令:
其中,参数env指定哪个π0 checkpoint 应该被加载,比如是ALOHA还是LIBEROuv run scripts/serve_policy.py --env=[DROID | ALOHA | LIBERO]
事实上,在后台,这个脚本将执行如下命令,可以使用它来启动策略服务器,例如用于你自己训练的检查点——这里是 DROID 环境的一个示例
其中,config这将启动一个策略服务器,该服务器将执行和参数指定的策略dir。该策略将在指定端口(默认值:8000)上执行uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi0_fast_droid --policy.dir=s3://openpi-assets/checkpoints/pi0_fast_droid
- 无需机器人即可测试推理:提供无需机器人即可测试推理的脚本。此脚本将生成随机观察并使用模型运行推理。有关更多详细信息,请参阅此处
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基础模型,只需以下三个步骤即可:
- 第一步,将Libero数据集转换为LeRobot数据集v2.0格式
作者提供了一个脚本convert_libero_data_to_lerobot.py「下面,我会详细分析一下这个脚本」,用于将 Libero 数据转换为 LeRobot 数据集 v2.0 格式,且可以轻松修改它以转换你自己的数据
比如从此处下载原始 Libero 数据集,然后使用以下命令运行脚本
且作为示例,已将 BiPlay 代码库中的 aloha_pen_uncap_diverse_raw 数据集转换,并将其上传至 HuggingFace Hub,地址为 physical-intelligence/aloha_pen_uncap_diverseuv run examples/libero/convert_libero_data_to_lerobot.py --data_dir /path/to/your/libero/data
- 第二步,定义使用自定义数据集的训练配置,并运行训练
提供pi0_aloha_pen_uncap 配置作为示例,参考根README以了解如何使用新配置运行训练
比如下面是 Libero 示例配置,可以根据自己的数据集进行修改:
LiberoInputs和LiberoOutputs:定义从 Libero 环境到模型的数据映射,反之亦然——将用于训练和推理
LeRobotLiberoDataConfig:定义了如何处理已经转换为LeRobot格式的Libero数据以用于训练,涉及数据配置、模型转换、数据转换
可能马上就有同学疑问了,不是已经在上面第一步做了数据转换了么——把Libero数据转换成LeRobot格式的 咋还要转呢?
原因很简单,这两个转换 不一样
一者,对LeRobotLiberoDataConfig源码的分析 请见此文《π0源码剖析——从π0模型架构的实现(如何基于PaLI-Gemma和扩散策略去噪生成动作),到基于C/S架构下的模型训练与部署》的「2.1.2 数据集配置:包含ALOHA、Libero两套数据集——LeRobotLiberoDataConfig」
二者 下文会对比下 这两个转换的不同
TrainConfig:定义微调超参数、数据配置和权重加载器
在运行训练之前,可以计算训练数据的标准化统计数据,使用训练配置的名称运行以下脚本
接下来,可以使用以下命令开始训练(uv run scripts/compute_norm_stats.py --config-name pi0_fast_libero
--overwrite
如果使用相同的配置重新运行微调,则该标志用于覆盖现有检查点)
该命令会将训练进度记录到控制台并将检查点保存到checkpoints目录中。还可以在 Weights & Biases 仪表板上监控训练进度。为了最大限度地利用 GPU 内存,请XLA_PYTHON_CLIENT_MEM_FRACTION=0.9在运行训练之前设置 - 这使 JAX 能够使用高达 90% 的 GPU 内存(而默认值为 75%)XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 uv run scripts/train.py pi0_fast_libero --exp-name=my_experiment --overwrite
- 第三步,启动策略服务器并运行推理
训练完成后,可以通过启动策略服务器,然后从 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 (文本)接下来,按如下步骤逐一执行
- 关键常量定义
# 输出数据集名称,也用于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数据集合并为一个训练数据集
- 定义主函数结构
主函数接受两个参数:def main(data_dir: str, *, push_to_hub: bool = False):
data_dir: 原始数据目录路径
push_to_hub: 是否推送到HuggingFace Hub的标志- 清理现有数据
确保输出目录干净,删除任何已存在的数据output_path = LEROBOT_HOME / REPO_NAME if output_path.exists(): shutil.rmtree(output_path)
- 创建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维浮点数向量
- 多线程/进程配置用于图像处理- 数据转换循环
这部分实现了数据转换的核心逻辑:# 遍历原始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数据和语言指令- 数据集的整合与发布
最后的处理步骤:
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这样的模型提供数据支持
二方面,这个上面第一步提到的convert_libero_data_to_lerobot与上面第二步提到的LeRobotLiberoDataConfig 中的各自转换的不同
- 不同的处理阶段:
- `convert_libero_data_to_lerobot` 是离线数据准备阶段的转换
- `LeRobotLiberoDataConfig` 中的转换是在训练和推理过程中实时应用的,至于详情,参见上面提到过的此文的「2.1.2 数据集配置:包含ALOHA、Libero两套数据集——LeRobotLiberoDataConfig」 - 键名映射的特殊性:
- LeRobotLiberoDataConfig中的`repack_transform` 专门负责将数据集键名映射到推理环境期望的键名
- repack_transform的注释明确说明:"重新打包转换仅应用于来自数据集的数据,不用于推理阶段",这确保了训练数据和推理时输入格式的一致性 - 模型特定转换
LeRobotLiberoDataConfig中的处理超出了简单的数据格式转换范围,比如其中的`data_transforms` 和 `model_transforms` 包含了许多模型特定的处理,比如data_transforms中包含的LiberoInputs/LiberoOutputs 处理 - 动作表示转换
- LeRobotLiberoDataConfig代码中有将绝对动作转换为相对动作的特殊逻辑
- 注释解释:"pi0模型是在相对动作上训练的",这是模型训练需求,而非数据格式问题 - 训练推理一致性
- LeRobotLiberoDataConfig中的`data_transforms` 在注释中明确表示:"应用于来自数据集的数据*和*推理过程中的数据"
- 这保证了训练和推理阶段处理的一致性
简言之,convert_libero_data_to_lerobot只涉及到数据格式的简单转换,但LeRobotLiberoDataConfig转换、处理更多
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 上测试过,其他操作系统可能不支持
以下是安装过程
- 克隆此 repo 时,请确保更新子模块:
git clone --recurse-submodules git@github.com:Physical-Intelligence/openpi.git # Or if you already cloned the repo: git submodule update --init --recursive
- 使用uv来管理 Python 依赖项。可参阅uv 安装说明进行设置
安装 uv 后,运行以下命令设置环境:
注意:GIT_LFS_SKIP_SMUDGE=1需要将 LeRobot 作为依赖项GIT_LFS_SKIP_SMUDGE=1 uv sync
Docker:作为 uv 安装的替代方案,他们还提供了使用 Docker 安装 openpi 的说明
如果在系统设置中遇到问题,还可以考虑使用 Docker 来简化安装,详情请参阅Docker 设置
1.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:建立策略客户端
注意:如果看到 EGL 错误,则可能需要安装以下依赖项# 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
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——含我司的微调实践
2.1 如何在真实环境中采集数据,一步一步微调π0然后部署到机械臂上
要在真实环境中采集数据并用于微调π0模型,需要完成以下几个步骤
2.1.1 准备环境、硬件、安装依赖、采集环境
- 设置开发环境和安装所需依赖:
# 克隆仓库(假设已经完成) # cd openpi-main # 创建虚拟环境 uv venv --python 3.10 .venv source .venv/bin/activate # 安装项目依赖 uv pip install -e . uv pip install -e packages/openpi-client
- 准备数据采集环境
根据`examples/aloha_real/README.md`,需要先设置Aloha机器人环境:
1. 按照Aloha仓库中的硬件安装说明设置机器人
2. 根据你的相机序列号修改`third_party/aloha/aloha_scripts/realsense_publisher.py`文件
2.1.2 数据的采集与转换
从`convert_aloha_data_to_lerobot.py`脚本中可以看出,Aloha机器人的数据是以HDF5格式存储的,并且需要转换成LeRobot数据集格式才能用于训练。数据包含了以下内容:
- 机器人关节状态(qpos)
- 动作(action)
- 可选的速度(qvel)和力矩(effort)数据
- 多个相机视角的图像(包括机器人基座相机和手腕相机)
采集步骤如下
- 设置数据采集环境
# 在一个终端窗口启动ROS节点 cd /Users/mac/Desktop/openpi-main source examples/aloha_real/.venv/bin/activate roslaunch aloha ros_nodes.launch
- 执行示范任务并记录数据
使用Aloha机器人的现有数据收集工具(在aloha目录中)来记录示范数据。这通常涉及:
- 设置你想要机器人执行的任务场景
- 使用示教模式记录人类演示
- 将数据保存为HDF5格式的episode文件 - 数据转换
使用提供的转换脚本将Aloha格式的数据转换为LeRobot格式:uv run examples/aloha_real/convert_aloha_data_to_lerobot.py --raw-dir /path/to/your/episodes --repo-id your-org/your-dataset-name
2.1.3 定义训练配置
下面是一个很好的示例配置 `pi0_aloha_pen_uncap`,它展示了如何为自定义的Aloha数据集创建训练配置。根据这个示例,你可以创建自己的配置:
- 创建自定义训练配置
在 config.py 中添加一个新的训练配置,参考 `pi0_aloha_pen_uncap` 示例:TrainConfig( name="pi0_your_custom_task", # 为你的任务设置一个唯一名称 model=pi0.Pi0Config(), # 使用标准π0模型或根据需要修改 data=LeRobotAlohaDataConfig( repo_id="your-org/your-dataset-name", # 你转换后的数据集名称 assets=AssetsConfig( assets_dir="s3://openpi-assets/checkpoints/pi0_base/assets", asset_id="trossen", # 根据你的机器人平台选择 ), # 例如:"拿起杯子" default_prompt="你的任务指令",
- 根据你的数据集结构定义重包装转换
# 根据你的数据集结构定义重包装转换 repack_transforms=_transforms.Group( inputs=[ _transforms.RepackTransform( { "images": { "cam_high": "observation.images.cam_high", "cam_left_wrist": "observation.images.cam_left_wrist", "cam_right_wrist": "observation.images.cam_right_wrist", }, "state": "observation.state", "actions": "action", } ) ] ),
- 如果是本地数据集
base_config=DataConfig( # 如果是本地数据集,设为True local_files_only=False, ), ), weight_loader=weight_loaders.CheckpointWeightLoader("s3://openpi-assets/checkpoints/pi0_base/params"), # 根据数据集大小调整 num_train_steps=20_000, )
- 选择合适的模型变体
- 对于标准微调,使用 `pi0.Pi0Config()`
- 对于低内存微调,使用 `pi0.Pi0Config(paligemma_variant="gemma_2b_lora", action_expert_variant="gemma_300m_lora")`
- 如果要冻结VLM参数只调整动作专家的参数,记得添加 `freeze_filter` 参数
但有一些注意事项:
- 对于 `assets=AssetsConfig` 部分,需要确认 `asset_id="trossen"` 是否匹配你的机器人平台
- `repo_id` 需要确保已正确上传到HuggingFace或配置为本地数据集
2.1.4 开始训练
训练可以通过以下命令执行:
# 使用你的自定义配置名称
uv run scripts/train.py --config pi0_your_custom_task --exp_name your_experiment_name
训练参数说明:
- `--config`:你的配置名称
- `--exp_name`:实验名称,用于区分不同的训练运行
- 还可以添加 `--batch_size`、`--num_train_steps` 等参数覆盖配置中的默认值
对于训练进度的监控,如果你在配置中启用了wandb(`wandb_enabled=True`),可以通过Weights & Biases平台监控训练进度
2.1.5 测试和部署训练好的模型
- 先启动ROS节点
- 启动策略服务器
# 加载训练好的模型 uv run scripts/serve_policy.py --env ALOHA --checkpoint_path ./checkpoints/pi0_your_custom_task/your_experiment_name/latest/params --default_prompt "你的任务指令"
- 同时,在另一个终端运行机器人控制程序:
# 在另一个终端运行机器人控制程序 python -m examples.aloha_real.main
2.2 我司七月微调π0的实践
七月内部从24年q3开始,一直给各种工厂做场景落地和定制开发,25年q3开始,一个个过了保密期之后,可以逐一拿出部分 分享下
// 待更