pi0_soarm101_lora

pi0_soarm101_lora

SO-ARM101 Pro 抓取放置策略：pi0 基座 + LoRA 微调（10,000 步）。

一、模型信息

• 基座：lerobot/pi0_base（OpenPI pi0，约 3.5B 参数）
• 微调方式：LoRA（r=16），冻结全部参数仅训练低秩适配器
• 训练步数：10,000 步（batch=16，约 4 epochs，40,923 帧）
• 训练数据：franecki/pick_place_100_train（95 集）+ _eval（5 集）
• 输出：完整权重（bf16，7GB），adapter 已合并进基座
• 任务指令（prompt，必须逐字符一致）：Place the white cylinder into the black circle

二、相对基座我们做了什么修改（重要）

2.1 对 pi0_base 基座配置的修改

policy_preprocessor.json / policy_postprocessor.json 已做两处修改：

1. tokenizer 本地化：tokenizer_name 从 google/paligemma-3b-pt-224（gated 仓库）改为本地路径 /root/data/paligemma_tokenizer。本仓库不包含 tokenizer 文件，部署端须按 3.3 节自行准备并放到该路径，否则加载 preprocessor 会报错。
2. 删除 relative_actions_processor / absolute_actions_processor：这两个是官方新版 LeRobot 的处理器，本项目用的 LeRobot 0.4.4（Seeed fork）注册表里没有。本项目使用绝对动作，删除后数据流与 0.4.4 原生 pi0 一致，不影响训练与推理。

2.2 对代码的修改（LeRobot 0.4.4 Seeed fork）

• 文件：src/lerobot/policies/pi0/modeling_pi0.py 的 _fix_pytorch_state_dict_keys
• 修改：增加 key 映射，把 checkpoint 的 paligemma.lm_head.weight 复制为 paligemma.model.language_model.embed_tokens.weight
• 原因：pi0 的 lm_head 与 embed_tokens 是 tied 权重，OpenPI checkpoint 只存 lm_head；0.4.4 代码缺该映射会导致 embed_tokens 加载缺失、模型以随机初始化运行。
• 部署端仅推理无需此补丁（权重已合并）；若重新训练或加载 adapter 必须保留。

2.3 模型合并方式

LoRA 训练产物为 adapter-only（5.4MB）。本仓库已用 PeftModel.merge_and_unload() 合并进 pi0_base 并导出完整 bf16 权重（7GB），部署端直接加载，无需 peft。

三、部署端操作步骤

3.1 环境要求

• Ubuntu 20.04/22.04（x86_64 或 aarch64），NVIDIA GPU 8GB 以上显存
• Python 3.10，PyTorch 大于等于 2.2 且小于 2.8（JetPack 6.x 自带 torch 亦可）
• LeRobot 0.4.4（Seeed fork）：git clone https://github.com/Seeed-Projects/lerobot
• 安装：pip install -e ".[pi]" 与 pip install peft

3.2 下载模型

export HF_ENDPOINT=https://hf-mirror.com export HF_HUB_DISABLE_XET=1 huggingface-cli download franecki/pi0_soarm101_lora --local-dir /root/models/pi0_soarm101_lora （大文件 7GB 若断流：重跑同命令续传；仍失败可用 aria2 多线程下载）

3.3 关键：准备 tokenizer（否则加载必失败）

模型配置的 tokenizer_name 指向 /root/data/paligemma_tokenizer，两种获取方式：

方式 A（推荐）：从训练服务器拷贝整个目录（含 tokenizer.json、tokenizer.model、tokenizer_config.json、special_tokens_map.json、config.json、preprocessor_config.json 共 6 个文件）： mkdir -p /root/data rsync -avzP 训练服务器IP:/root/data/paligemma_tokenizer /root/data/

方式 B（从 HF 下载）：google/paligemma-3b-pt-224 是 gated 仓库，需先在网页点击 Acknowledge license 同意许可，再用 read token 下载： export HF_ENDPOINT=https://hf-mirror.com huggingface-cli login --token hf_你的token huggingface-cli download google/paligemma-3b-pt-224 --local-dir /root/data/paligemma_tokenizer --include "tokenizer*" "special_tokens_map.json" "tokenizer_config.json" "config.json" "preprocessor_config.json"

若想放到其他路径，须同步修改 policy_preprocessor.json 的 tokenizer_name。

3.4 加载与推理

from lerobot.policies.pi0.modeling_pi0 import PI0Policy policy = PI0Policy.from_pretrained("/root/models/pi0_soarm101_lora") # 用绝对路径，勿用 ~ policy.eval()

观察输入（与训练一致）：

• observation.images.top：顶部相机 640x480（模型内部 resize 224x224）
• observation.images.gripper：腕部相机 640x480
• observation.state：6 维关节角度，与训练数据同单位
• prompt 必须严格一致：Place the white cylinder into the black circle

动作输出：action chunk，50 步（chunk_size=50），6 维。

3.5 Jetson（AGX Orin 64GB）部署要点

• 推理约 0.5 到 2 秒每次；配合 action chunk 执行，控制频率 0.5 到 1Hz 可行
• 推理慢的降级顺序：num_inference_steps 10 改 5；再慢启用 RTC；最后用远程推理
• 视频解码：aarch64 自动回退 pyav，无需配置
• 相机画面、光照、物体摆放尽量与训练数据一致，否则成功率下降

3.6 常见坑

• 加载报 tokenizer 路径错误：tokenizer 未放到 /root/data/paligemma_tokenizer，见 3.3
• 推理乱动或不动：prompt 与训练不一致，逐字符核对
• 模型加载报 repo id 错误：用了 ~ 相对路径，改用绝对路径
• 夹爪抖动：gripper 归一化 std 极小（训练数据特征），已知现象，观察是否影响成功率

四、训练详情

• 训练 loss：10k 步收敛至约 0.15，grdn 约 0.5，无发散
• 学习率：cosine，峰值 2.5e-5
• 数据：100 集示教（单任务），6 维动作，2 路相机，30fps，AV1 编码（训练用 pyav 解码）
• 对比参考：franecki/pi0_soarm101_expert_only（train_expert_only 8,000 步），真机对比后择优部署

五、相关仓库

• 数据集：franecki/pick_place_100（原始 100 集）、franecki/pick_place_100_train 和 _eval（切分后）
• 模型：franecki/pi0_soarm101_expert_only、franecki/pi0_soarm101_lora（本仓库）