π0 (Pi-Zero)

核心思路

设计哲学：机器人领域的「基础模型」

π0 把大语言模型的「预训练 + 微调」范式搬到机器人领域：先在海量多任务、多机器人数据上预训练，再在小量下游数据上微调到具体任务。目标是训练一个通用机器人基础策略，能处理折衣服、收拾桌子、装盒子等复杂长程任务。

为什么用 Flow Matching 而不是传统方法

π0 选择 Flow Matching 的理由：

• 直接在连续动作空间回归「速度场」，无需量化，精度高
• Flow Matching 是扩散模型的「近亲」，但训练/推理更简单：训练只需一次前向，推理用 10 步欧拉积分
• 与 VLM 骨干天然兼容：action expert 作为独立「专家」分支挂在主 Transformer 上

四大核心创新

1. VLM 预训练继承语义知识：以 PaliGemma (3B VLM) 为骨干，天生拥有图像-语言理解能力
2. Flow Matching + Action Expert 解耦架构：主 VLM 处理感知/语言，小专家网络 (~300M) 专门处理动作
3. Action Chunking（动作分块）：一次预测未来 50 步动作，开环执行若干步后再推理
4. KV-cache 推理优化：前缀（图像+语言）只前向一次并缓存，10 步去噪只重算后缀

技术方案特点

模型架构：双专家单 Transformer

输入输出格式

Action Horizon（动作时域）：π0 默认 50 步，DROID 配置 10 步。Action Dim 因平台不同：DROID=8，ALOHA 双臂=14。

动作表达：Flow Matching 如何工作

训练阶段

# pi0.py → compute_loss

# 1. 采样噪声和流匹配时间步 τ
noise = randn(actions.shape)
time = beta(1.5, 1) * 0.999 + 0.001     # β(1.5,1)采样，偏向高 τ

# 2. 构造中间状态动作 x_τ（线性插值路径）
x_t = time * noise + (1 - time) * actions   # τ=0 干净动作，τ=1 纯噪声

# 3. 回归目标 = 速度场 u_t = noise - actions
u_t = noise - actions

# 4. 一次前向：前缀(图+语言)+后缀(x_τ, τ) 送入 Transformer
v_t = model(observation, x_t, time)         # 模型预测速度场

# 5. 损失 = MSE(v_t, u_t)
loss = mean(square(v_t - u_t), axis=-1)

推理阶段（10 步欧拉积分 + KV cache）

# pi0.py → sample_actions

dt = -1.0 / num_steps      # num_steps=10
x = randn(batch, action_horizon, action_dim)   # 从纯噪声出发

# ★ 第 1 步：前缀前向，缓存 KV cache（图像+语言只算一次！）
kv_cache = model.forward_prefix(observation)

# ★ 第 2~10 步：仅重算后缀（动作 token），复用 KV cache
for step in range(10):
    v_t = model.forward_suffix(x, t, kv_cache)
    x = x + dt * v_t          # 欧拉步
    t = t + dt
return x   # x_0 = 生成的动作块

训练数据

代码框架

仓库结构（openpi）

openpi/
├── src/openpi/                    # 核心库
│   ├── models/                    # ★ 模型实现
│   │   ├── pi0.py                 #   π0 主模型（flow matching 核心）
│   │   ├── pi0_config.py          #   Pi0Config（action_dim/horizon）
│   │   ├── pi0_fast.py            #   π0-FAST（自回归 + FAST tokenizer）
│   │   ├── gemma.py               #   ★ 双专家 Gemma Transformer + SigLIP
│   │   ├── model.py               #   Observation / Actions 数据结构
│   │   └── lora.py                #   LoRA 适配器
│   ├── policies/                  # ★ 策略封装（输入输出映射）
│   │   ├── policy.py              #   Policy.infer() 推理管线
│   │   └── policy_config.py       #   create_trained_policy()
│   ├── training/
│   │   ├── config.py              #   ★★ 所有 TrainConfig 注册表
│   │   ├── data_loader.py         #   LeRobot / RLDS 数据加载
│   │   └── checkpoints.py         #   checkpoint + norm_stats
│   ├── serving/                   # 推理服务（websocket）
│   └── transforms.py              # 数据变换管线
├── scripts/
│   ├── train.py                   # ★ JAX 训练入口
│   ├── train_pytorch.py           # ★ PyTorch 训练入口
│   ├── serve_policy.py            # ★ 启动推理服务器
│   └── compute_norm_stats.py      # ★ 计算归一化统计量
├── packages/openpi-client/        # ★ 轻量客户端（机器人端安装）
└── examples/                      # 各平台示例（DROID/ALOHA/LIBERO/UR5）

核心模块

微调流程（以 LIBERO 为标准示例）

Step 1 — 数据转 LeRobot 格式

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

你的数据需包含：图像帧、本体状态、动作、语言指令。

Step 2 — 定义配置 + 训练

在 training/config.py 的 _CONFIGS 中添加配置：

TrainConfig(
    name="pi05_my_robot",
    model=pi0_config.Pi0Config(pi05=True, action_horizon=10),
    data=LeRobotLiberoDataConfig(repo_id="your_org/your_dataset"),
    batch_size=256,
    lr_schedule=CosineDecaySchedule(warmup_steps=10_000, peak_lr=5e-5),
    ema_decay=0.999,
    weight_loader=CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base"),
    num_train_steps=30_000,
)

# (a) 计算归一化统计量（必做）
uv run scripts/compute_norm_stats.py --config-name pi05_my_robot

# (b) 训练
XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 uv run scripts/train.py pi05_my_robot \
    --exp-name=my_experiment --overwrite

Step 3 — 部署推理

# 启动策略服务器
uv run scripts/serve_policy.py policy:checkpoint \
    --policy.config=pi05_my_robot \
    --policy.dir=checkpoints/pi05_my_robot/my_experiment/20000

LoRA 微调（低显存方案）

model=pi0_config.Pi0Config(
    paligemma_variant="gemma_2b_lora",       # rank=16
    action_expert_variant="gemma_300m_lora", # rank=32
)

LoRA 冻结主干、只训练低秩适配器，显存降至 22.5GB。

推理部署

方式 A — 进程内推理（最简单）

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

config = _config.get_config("pi05_droid")
ckpt = download.maybe_download("gs://openpi-assets/checkpoints/pi05_droid")
policy = policy_config.create_trained_policy(config, ckpt)

action_chunk = policy.infer(example)["actions"]   # → (action_horizon, action_dim)

方式 B — 远程推理（推荐，机器人端零依赖）

# GPU 服务器端
uv run scripts/serve_policy.py --env DROID

# 机器人端（pip install openpi-client）
from openpi_client import websocket_client_policy, image_tools

client = websocket_client_policy.WebsocketClientPolicy(host="GPU_IP", port=8000)
observation = {
    "observation/image": image_tools.convert_to_uint8(
        image_tools.resize_with_pad(img, 224, 224)),
    "observation/state": state,
    "prompt": "pick up the fork",
}
action_chunk = client.infer(observation)["actions"]   # (action_horizon, action_dim)

快速上手

安装

git clone --recurse-submodules https://github.com/Physical-Intelligence/openpi.git
cd openpi
# 用 uv 管理依赖（需先装 uv: curl -LsSf https://astral.sh/uv/install.sh | sh）
GIT_LFS_SKIP_SMUDGE=1 uv sync
GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .

加载模型 + 推理（最短路径）

from openpi.training import config as _config
from openpi.policies import policy_config
from openpi.shared import download
import numpy as np

# 1. 选配置 + 自动下载 checkpoint
config = _config.get_config("pi05_droid")
ckpt   = download.maybe_download("gs://openpi-assets/checkpoints/pi05_droid")

# 2. 创建策略（自动组装 transforms + norm_stats）
policy = policy_config.create_trained_policy(config, ckpt)

# 3. 构造观测（DROID 格式）
example = {
    "observation/exterior_image_1_left": np.random.randint(256, (224,224,3), dtype=np.uint8),
    "observation/wrist_image_left":      np.random.randint(256, (224,224,3), dtype=np.uint8),
    "observation/joint_position":        np.random.rand(7),
    "observation/gripper_position":      np.random.rand(1),
    "prompt": "pick up the fork",
}

# 4. 推理 → 动作块
actions = policy.infer(example)["actions"]
print(actions.shape)   # (action_horizon, action_dim) e.g. (15, 8)

一条命令启动推理服务器

uv run scripts/serve_policy.py --env DROID   # 自动用 pi05_droid 默认 checkpoint

硬件需求

多卡可用 FSDP（fsdp_devices=N）分摊显存。远程推理方案下机器人端无需 GPU。大规模 DROID 训练（100k 步）需 8×H100 约 2 天。

关键数字