PI0.5 微调训练与数据流水线踩坑简记

发表于 21 小时前  21 次阅读

文章目录
------
数据task_5867（排除第 355 集）+ task_5867-1，合并后 **681 集 / 305,425 帧**（30fps，3 路相机）
配置global batch 32（4/卡 × 8 卡）、30000 步、action horizon 32、bf16
状态/动作各 **16 维**（14 关节 + 2 夹爪）
loss1.47 → **0.03 量级**（step 7000 时 0.029–0.043）
吞吐0.838 s/step（扫描）、38.1 samples/s
显存每卡 22–48 GB（64 GB 卡），batch 上限 = **每卡 4**

二、环境与权重

  • 容器：`drivingsdk:26.1.0-9.1.0-torch_npu2.7.1.post6-910b-openeuler24.03-py3.10`
  • 代码：`/DrivingSDK/model_examples/Pi-0.5`（lerobot 0.3.4 + pi05 适配）
  • 环境：`conda activate torch2.7.1`（README 里写的 `pi05` 环境并不存在）
  • 权重：`/data/huawei/models/pi05_base`
  • 离线机器（访问不了 huggingface.co），必须先把所有依赖本地化

**一个容易忽略的点**：pi05 的分词器是硬编码的 `google/paligemma-3b-pt-224`，离线环境会直接报

`OSError: We couldn't connect to 'https://huggingface.co'`。而且 PaliGemma 是 HF 的 gated 仓库，镜像站也拿不到。

最后从 **ModelScope** 下载了全套（`tokenizer.model` 4.26MB + `tokenizer.json` 17.5MB 等 4 个文件），

把 `processor_pi05.py` 里的 `tokenizer_name` 指向本地路径解决。

三、数据准备：11 个坑

3.1 为什么不能直接用

两件事必须先做：

  1. **两份数据集要一起训** → 但这份 lerobot 的 `make_dataset` 对多数据集直接

`raise NotImplementedError("The MultiLeRobotDataset isn't supported for now.")`，只能合并成一份。

  1. **数据是 v2.1，训练只吃 v3.0** → 官方提示 `Please update your dataset to the new format`。

3.2 官方合并路径踩死（重要）

按官方 `merge_datasets` 走，连续崩了多次：

ValueError: cannot convert float NaN to integer          # 删集时
TypeError: loop of ufunc does not support argument 0 of type numpy.ndarray   # np.sqrt（numpy2 兼容性）
TypeError: Object of type ndarray is not JSON serializable
ValueError: Same features is expected, but got ...
ValueError: all input arrays must have the same shape
ValueError: Shape of 'count' must be (1), but is (1, 1) instead.

逐个修完之后，**合并终于成功（681 集 / 305425 帧），但视频元数据是坏的**——校验发现有

**414/681 集的视频时间戳超出了实际文件时长**，训练直接报：

AssertionError: One or several query timestamps unexpectedly violate the tolerance
queried timestamps: tensor([1422.6666])  →  实际文件只到 766.5s

也就是说官方 `aggregate_videos` 的拼接逻辑有问题（写进去的时间戳与产出的视频对不上）。

**最终方案：绕开官方聚合，改用手工文件级合并**——在 v2.1 层面把两份数据按 episode 逐文件复制并重编号

（`data/chunk-000/episode_XXXXXX.parquet` + 三路 `videos/.../episode_XXXXXX.mp4`），改写

`episode_index` / `index` 列与 `meta/{episodes,episodes_stats,tasks,info}.json`，再用官方转换器转 v3.0。

这样**视频拼接交给转换器做**（已验证正确），结果：**681 集 / 305,425 帧，视频元数据 0 异常**，而且合并本身只花了约 1 分钟。

3.3 官方分位数脚本不可用

pi05 的 STATE/ACTION 强制使用 **QUANTILES 归一化**（缺 q01/q99 直接报错），官方给了

`augment_dataset_quantile_stats.py`。但它是**逐集解码 3 路视频**：

Processing episodes: 6%|▌ | 44/681 [4:20:13<62:47:22, 354.85s/it]   ← 330+ 秒/集，外推 62 小时

跑到第 44 集还崩了。**改用直接读 parquet 计算**（215K 帧 × 199 维，秒级完成）：

for q in (0.01, 0.10, 0.50, 0.90, 0.99):
    st[f'q{int(q*100):02d}'] = np.percentile(X, q * 100, axis=0)

顺便补一句：图像特征不需要数据集统计——`DatasetConfig.use_imagenet_stats=True` 会在建数据集时自动套 ImageNet 均值方差。

**所以视频特征那份 `count=0` 的假统计根本没必要算**，而它正是把官方合并路径搞崩的元凶之一。

3.4 坑清单（11 项）

#现象根因处置上游 bug
1要求 v3.0 格式数据是 v2.1官方转换器转换—
2`Couldn't cast list to uint64``ts`/`action_src_status` 存成 `list`[1]压平为标量—
3删集崩 `int(NaN)`**转换器把 ep312 的 `data/file_index` 写错**按实际文件修正✅
4聚合除零 → NaN → `np.sqrt` 崩视频特征 stats 是 `count=0` 占位符清除占位—
5`ndarray is not JSON serializable`stats 含嵌套 ndarray`write_json` 递归 sanitize✅
6`np.sqrt(object array)` 崩numpy 2 兼容性显式 `float64`✅
7`Same features is expected`两份数据 features 的 `fps` 字段不一致统一字段—
8`all input arrays must have the same shape`同一个特征的 std/count 形状 (1,) vs (1,1)按 `min` 形状对齐—
9**414/681 集视频时间戳错误**官方 `aggregate_videos` 拼接 bug**改手工文件级合并**✅
10分位数脚本 330 秒/集逐集解码视频直接读 parquet 计算工具性能
11**训练 loss 爆炸到 8×10⁵**见下节裁剪维度—

四、最隐蔽的一个坑：loss 爆炸

数据都对齐之后开训，结果：

step:20 loss:1.733
step:40 loss:813270.345        ← 直接炸了

根因不在超参，而在**数据量纲**。AgiBot 导出的 `observation.state` 是 **163 维**，里面除了关节角还混着：

  • 力/力矩（wrench，±50）
  • 关节电流（±3000）
  • 相机外参（±1000 量级）
  • 笛卡尔遥操作量

pi05 用 QUANTILES 归一化：把 [q01, q99] 映到 [-1, 1]，**超出范围的值做线性外推**。

于是这些"极端量纲"通道被放大成几百上千的输入，训练立刻发散。

**处置**：把 state/action 裁到真正有意义的控制量——

state  = joint/position[28:42) + effector[0:2)   → 16 维（14 关节 + 左右夹爪）
action = joint/position[16:30) + effector[0:2)   → 16 维

顺带一个好处：16 维落回 pi05 预训练头原生的 `max_state_dim = max_action_dim = 32`，

不再需要 `--policy.max_state_dim=256 --policy.max_action_dim=64` 这种补丁

（那个补丁虽然能跑通——它解决的是"36 维 action > 32 上限"导致的

`aclnnAddmm: k-axis different`——但超出预训练维度的部分是随机初始化的，质量不保证）。

裁完重训，loss 立刻正常：

step:160  loss:1.473
step:200  loss:1.403
step:220  loss:1.346   updt_s:0.846  data_s:0.018
...
step:7K   loss:0.029~0.043

五、batch size 与显存实测

换硬件或换配置时，建议先跑一遍扫描（脚本：单卡 `pi05_sweep.sh`、8 卡 `pi05_sweep8.sh`）。

5.1 单卡（910B3 64GB）

batchs/step吞吐峰值 HBM结果
20.5293.7 samples/s38.4 GB✅
**4****0.705****5.6 samples/s****47.5 GB**✅ 上限
8——62.8 GB❌ 崩（近 OOM）
16——65.3 GB❌ OOM

5.2 8 卡

每卡 batchglobal batchs/step吞吐峰值 HBM结果
2160.74821.3 samples/s27.9 GB✅
**4****32****0.838****38.1 samples/s**22.2 GB✅ 上限
864———❌ OOM

**结论**：本平台（64 GB/卡）的上限是**每卡 4 → global batch 32**。这恰好等于客户给定配置。

5.3 一个有意思的对照

客户用 **RTX PRO 6000（96 GB）** 跑同样 batch 32，但**显存占用明显偏低**（有大半富余）。

我们在 64 GB 卡上跑同样的 global batch 32 已经贴着上限。

差异来自**单样本开销**：我们这份数据的 3 路相机是 1280×800 + 848×480（头部鱼眼 + 双手 D405），

而对方那套分辨率/相机数更省。**换句话说，他们那张卡的显存还有很大空间，加大 batch 是提升吞吐最直接的手段。**

六、小结

  1. **具身模型的适配成本，大头在数据不是模型**。这次 11 个坑里有 9 个在数据流水线。
  2. **对客交付的数据要按"能直接喂训练"的标准验收**：格式版本、类型、统计量、视频时间戳，缺一项都会在训练时才炸。
  3. **官方工具链不能盲信**：`aggregate_videos`（视频时间戳错）、`augment_dataset_quantile_stats`（330 秒/集）、

转换器的 `file_index`、`np.sqrt` 的 numpy2 兼容性、`write_json` 的 ndarray——建议反馈上游。

  1. **loss 爆炸先看数据量纲**，别急着调学习率：极端量纲通道在量化归一化下会被线性外推放大。
  2. 16 维（关节 + 夹爪）的控制量既够用、又恰好落在预训练头维度内，是这次最"顺"的选择。

**附**：完整复现步骤、脚本清单与校验标准见《PI0.5 微调训练复现手册》。

scanz个人博客