PI05 策略说明（OpenPI 风格）
PI05 与 PI0 同属「大视觉语言模型 + 专家头」范式，接口与配置采用 OpenPI 风格。基于 PaliGemma 编码图像与可选语言，通过 flow matching 专家头生成动作。支持纯 RGB 或 RGB+Depth（Cross-Attention 融合），可选语言条件；默认使用 QUANTILES 归一化，需事先补充 quantile 统计。本项目通过 CustomPI05ConfigWrapper 与 CustomPI05Pytorch 扩展，支持 depth 的 Cross-Attention 融合。
1. 核心思想
- VLM + Expert：与 PI0 类似，PaliGemma + 专家头；
- Flow Matching：动作通过 flow matching 生成，推理时多步去噪；
- QUANTILES 归一化：state 与 action 使用分位数归一化，需
augment_dataset_quantile_stats； - 可选冻结：
freeze_vision_encoder、train_expert_only可减少显存、加速训练（PI05 特有）。
2. 模型架构概览
观测（RGB + 可选 Depth + state + 可选 language）
↓
┌─────────────────────────────────────────────────────────────┐
│ RGB: PaliGemma Vision Tower → RGB features │
│ Depth（若 use_depth）: 复制3通道 → PaliGemma Vision Tower │
│ → CrossModalAttentionFusion 双向 Cross-Attention │
│ → [fused_rgb, fused_depth] 加入 prefix embeddings │
└─────────────────────────────────────────────────────────────┘
↓
[ image_embeddings, (可选) language_embeddings, state_embeddings ]
↓
PaliGemma LLM + Expert Head（OpenPI 风格）
↓
Flow Matching 解码 → (B, chunk_size, action_dim)
3. 输入与输出
3.1 输入
| 类型 | 键名 | 说明 |
|---|---|---|
| RGB | image_features |
多视角 RGB，如 observation.images.head_cam_h、wrist_cam_l、wrist_cam_r |
| Depth | depth_features |
可选，Cross-Attention 时与 RGB 共享 PaliGemma |
| State | robot_state_feature |
关节角度、夹爪等 |
| Language | task / 语言键 |
可选，部署时可注入空字符串 "" |
3.2 输出
- 形状：
(B, chunk_size, action_dim) - 推理：
select_action返回 action chunk，按n_action_steps逐步执行。
4. 关键配置参数
4.1 结构与动作
| 参数 | 说明 | 典型值 |
|---|---|---|
n_obs_steps |
观测历史步数 | 1 |
chunk_size |
动作块长度 | 50 |
n_action_steps |
每步执行动作数 | 50 |
max_state_dim |
状态维度上限 | 32 |
max_action_dim |
动作维度上限 | 32 |
image_resolution |
图像分辨率 | [224, 224] |
empty_cameras |
空相机数量 | 0 |
4.2 模型变体
| 参数 | 说明 | 典型值 |
|---|---|---|
paligemma_variant |
PaliGemma 变体 | gemma_2b |
action_expert_variant |
专家头变体 | gemma_300m |
dtype |
精度 | float32 / bfloat16 |
4.3 Flow Matching
| 参数 | 说明 | 典型值 |
|---|---|---|
num_inference_steps |
推理去噪步数 | 10 |
time_sampling_beta_alpha |
时间采样参数 | 1.5 |
time_sampling_beta_beta |
时间采样参数 | 1.0 |
min_period / max_period |
周期范围 | 4e-3 / 4.0 |
4.4 Depth（可选）
| 参数 | 说明 | 典型值 |
|---|---|---|
use_depth |
是否启用 depth | false |
depth_features |
depth 观测键 | [] 或与 RGB 对应 |
depth_fusion_method |
融合方法 | cross_attention |
depth_preprocessing |
预处理 | minmax |
depth_scale |
缩放因子 | 0.0001 |
4.5 归一化（PI05 默认 QUANTILES）
- VISUAL：
IDENTITY - STATE：
QUANTILES（需 quantile 统计） - ACTION：
QUANTILES（需 quantile 统计）
4.6 训练优化（含 PI05 特有）
| 参数 | 说明 | 典型值 |
|---|---|---|
gradient_checkpointing |
梯度检查点 | true |
compile_model |
torch.compile | false |
optimizer_lr |
学习率 | 2.5e-5 |
optimizer_weight_decay |
权重衰减 | 0.01 |
optimizer_grad_clip_norm |
梯度裁剪 | 1.0 |
tokenizer_max_length |
Tokenizer 最大长度 | 200 |
freeze_vision_encoder |
冻结视觉编码器（PI05 特有） | false |
train_expert_only |
仅训练专家头（PI05 特有） | false |
5. 预训练模型下载（HuggingFace）
训练与部署 PI05 需要两类模型：PaliGemma（视觉-语言底座，含 tokenizer）与 LeRobot π₀.₅ 预训练策略（pi05_base）。以下为从 HuggingFace 获取的常用方式。
5.1 PaliGemma 开源模型（含 Tokenizer）
PI05 使用的 tokenizer 与视觉编码器来自 Google 的 PaliGemma。部署时若使用本地 tokenizer，需单独下载。
- HuggingFace 模型 ID：google/paligemma-3b-pt-224
- 说明：使用前需在 HuggingFace 模型页 同意 Google 的使用条款。
方式一：命令行下载到本地目录
pip install huggingface_hub
huggingface-cli login
huggingface-cli download google/paligemma-3b-pt-224 \
--local-dir /path/to/models/paligemma-3b-pt-224 \
--local-dir-use-symlinks False
方式二：Python 预下载
from huggingface_hub import snapshot_download
snapshot_download(
repo_id="google/paligemma-3b-pt-224",
local_dir="/path/to/models/paligemma-3b-pt-224",
local_dir_use_symlinks=False,
)
部署时指定 tokenizer 路径：
export PALIGEMMA_TOKENIZER_PATH=/path/to/models/paligemma-3b-pt-224
5.2 LeRobot π₀.₅ 预训练底座（pi05_base）
PI05 的完整预训练策略由 LeRobot 提供，训练时 pretrained_path 可直接使用 Hub ID，首次运行会自动下载。
- HuggingFace 模型 ID：lerobot/pi05_base
- 可选：lerobot/pi05_libero（在 Libero 上微调）
方式一：训练时直接使用 Hub ID（自动下载）
--policy.pretrained_path=lerobot/pi05_base
方式二：预先下载到本地
huggingface-cli download lerobot/pi05_base \
--local-dir /path/to/pi05_base \
--local-dir-use-symlinks False
from huggingface_hub import snapshot_download
snapshot_download(repo_id="lerobot/pi05_base", local_dir="/path/to/pi05_base", local_dir_use_symlinks=False)
然后设置 pretrained_path=/path/to/pi05_base。离线环境可设 HF_HUB_OFFLINE=1 使用本地缓存。
6. 预训练与 Quantile
- pretrained_path：必填，指向 PI05 预训练底座（如
lerobot/pi05_base或本地路径）； - Quantile：PI05 默认使用 QUANTILES，必须先运行
augment_dataset_quantile_stats.py为stats.json补充分位数； - freeze_vision_encoder：冻结 PaliGemma 视觉编码器，可省显存、加速训练；
- train_expert_only：冻结整个 VLM，仅训练专家头与投影层。
7. Depth 支持（Cross-Attention）
与 PI0 相同：启用 use_depth 且 depth_fusion_method: "cross_attention" 时，Depth 经 PaliGemma 编码后与 RGB 做 CrossModalAttentionFusion。详见 RGB-深度融合。
8. 训练
8.1 配置与入口
- Config：
configs/policy/pi05_config.yaml（RGB-only）或pi05_crossattn_config.yaml（RGB+Depth） - 入口：
third_party/lerobot/src/lerobot/scripts/lerobot_train.py
8.2 训练命令示例（RGB-only）
# 先补充 quantile 统计
python third_party/lerobot/src/lerobot/datasets/v30/augment_dataset_quantile_stats.py \
--repo-id sim_task2_lerobot \
--root /path/to/sim_task2_lerobot/lerobot
CUDA_VISIBLE_DEVICES=0 python third_party/lerobot/src/lerobot/scripts/lerobot_train.py \
--dataset.repo_id=sim_task2_lerobot \
--dataset.root=/path/to/sim_task2_lerobot/lerobot \
--policy.type=pi05 \
--output_dir=./outputs/train/sim_task2/pi05_rgb \
--job_name=pi05_training \
--policy.pretrained_path=/path/to/pi05_base \
--policy.push_to_hub=false \
--steps=30000 \
--batch_size=16 \
--policy.device=cuda
8.3 可选冻结策略
# 仅训练专家头（省显存、 faster）
--policy.train_expert_only=true
# 冻结视觉编码器
--policy.freeze_vision_encoder=true
9. 部署
- policy_type：
pi05 - 仿真 / 真机：
sim_auto_test、real_single_test通过setup_policy(..., "pi05", device)加载； - RGB-only 部署：自动裁剪 depth、注入
observation["task"] = [""]； - Tokenizer：可设置
PALIGEMMA_TOKENIZER_PATH覆盖 checkpoint 中的路径； - dtype / compile：部署时强制
float32、compile_model=false。
10. PI0 与 PI05 对比
| 项目 | PI0 | PI05 |
|---|---|---|
| 归一化 | MEAN_STD 或 QUANTILES | 默认 QUANTILES |
| Quantile 统计 | 可选 | 必填 |
| tokenizer_max_length | 48 | 200 |
| freeze_vision_encoder | 无 | 有 |
| train_expert_only | 无 | 有 |
| Config 类 | CustomPI0ConfigWrapper | CustomPI05ConfigWrapper |