【VLA】Pi0＋ReconVLA 联合微调项目：技术实现与思考笔记

type

status

date

slug

summary

Pi0相关源码深度解析与侵入式修改策略

1.1 问题背景

原版pi0的代码设计初衷是用于推理（Inference）和标准微调。其核心类 PI0Pytorch 的 forward 函数是一个黑盒：

输入：图像 + 文本 + 动作 GT。

输出：仅返回一个标量 loss。

痛点：ReconVLA 任务的核心逻辑是“利用 VLM 提取的特征进行图像重建”。我们需要拿到 VLM 主干网络中间层的 Visual Embeddings（视觉特征） 和 Text Embeddings（文本特征）。原版代码在计算完 Action Loss 后，就把这些中间特征丢弃了。

1.2 “白盒化”修改策略

为了截获这些特征，我采用了侵入式修改（Intrusive Modification）策略，而非外部 Hook（因为 OpenPI 模块嵌套过深，Hook 难以定位且容易断开梯度）。

修改路径： src/pi0_core/pi0_pytorch.py

定位特征源头 (embed_prefix)：

分析发现 embed_prefix 函数负责调用 SigLIP 和 Gemma 的 Embedding 层。

修改前：它将视觉和文本特征拼接后返回一个混合的 embs。

修改后：强制它额外返回原始的 img_emb (Visual) 和 lang_emb (Text)。

思考：必须在这里截获，因为一旦拼接并通过 Transformer 层，视觉和文本信息就会深度混合，难以剥离用于独立的图像重建任务。

打通传输管道 (forward)：

forward 函数调用 embed_prefix。我修改了解包逻辑，接收新增的特征变量。

关键点：在返回值中，将单一的 loss 扩展为字典 {'loss': ..., 'visual_features': ..., 'text_embeds': ...}。

梯度流向：注意这里直接返回 Tensor，没有使用 .detach()。这是为了保证 Recon Head 计算出的 Loss 梯度能够反向传播回 VLM Backbone，从而优化 Backbone 的表征能力。

动态维度适配：

发现源码中硬编码了 nn.Linear(32, ...)，导致无法适配不同自由度（如 14DoF）的机器人数据。

修改：将其改为 nn.Linear(config.action_dim, ...)，增强了模型的通用性。

胶水层 Wrapper 的设计与实现

src/wrapper.py 是本项目的核心，它充当了 Pi0 和 ReconVLA 之间的“翻译官”和“调度员”。

2.1 主要功能

双流控制：同时运行 Pi0（动作生成）和 ReconHead（图像重建）。

维度对齐（最难点）：解决 VLM 序列化输出与 Diffusion 图像化输入之间的矛盾。

Loss 融合：实现。

2.2 关键数据结构变换 (Dimension Shape Logic)

这是调试过程中报错最多的地方，也是理解多模态融合的关键。

数据阶段	形状 (Shape)	含义	处理逻辑
Pi0 输出	`[B, N, D]` 例如 `[2, 768, 64]`	序列化特征。包含多个摄像头(3个)拼接后的长序列。	原始输出，无法直接进 DiT。
维度重排	`[B*K, D, H, W]` 例如 `[6, 64, 16, 16]`	空间化特征。使用 `einops.rearrange` 将序列还原为图片网格。	`K=3` (视角数), `N=KHW`。将 Batch 和 View 维度合并，视为独立样本处理。
文本条件	`[B, L, D]` -> `[B*K, D, 1, 1]`	全局条件向量。需广播到每个视角。	1. `pool`: 序列取平均变 `[B, D]`。 2. `repeat`: 复制 K 份。 3. `reshape`: 伪装成 1x1 图片以适配 DiT 接口。

PEFT / LoRA 微调实战解析

本项目并未重新训练所有参数，而是使用了 PEFT (Parameter-Efficient Fine-Tuning)。

3.1 为什么要用 LoRA？

显存限制：全量微调 2B+ 模型需要 80GB+ 显存，而 LoRA 仅需 24GB 左右。

灾难性遗忘：保留预训练模型的通用知识，仅学习适应机器人控制和物理重建的“增量知识”。

3.2 代码中的使用范式 (`scripts/train.py`)

配置图纸 (LoraConfig)：

理解：我们在 Transformer 最核心的注意力投影层旁挂载了低秩矩阵。

实施手术 (get_peft_model)：

效果：这一行代码在内存中修改了模型结构。原本的 Linear 层变成了 Linear + LoRAAdapter。同时，自动将原参数设为 requires_grad=False。

差异化供电 (Optimizer)：

理解：这是一个混合优化策略。Recon 分支是新加的，所以全参都需调整；Pi0 主干是成熟的，所以 LoRA 微调。

Debug 历程与问题排查 (Troubleshooting Log)

4.1 环境依赖问题

现象：ModuleNotFoundError: openpi, ImportError: flax。

原因：直接移植的代码保留了原项目的绝对路径引用和JAX依赖。

解决：

将绝对引用 from openpi... 改为相对引用 from . ...。
安装 jax CPU 版和 beartype 等库“骗”过解释器。
利用pi0项目提供的方案覆盖 transformers 源码，解决底层补丁问题。

4.2 内存溢出 (OOM/Killed)

现象：WSL 运行脚本直接被 Kill。

原因：默认配置加载了 2B 模型，且开启了 torch.compile，超出了开发机内存。

解决：

引入 Dummy Mode：修改 Config，将模型宽度从 2048 强降为 64。
注释掉 torch.compile。

4.3 维度与广播冲突 (The Shape Mismatches)

这是最耗时的部分，主要有三次报错：

768 != 729：OpenPI 将多视角图片Flatten了，而Wrapper以为是单张正方形图片。

修正：引入 num_cams 变量，先切分视角再 Reshape。

256 != 16：DiT 的 AdaLN 模块试图将文本序列 (16) 乘到图像 Patch (256) 上。

修正：对文本进行 Pooling (平均化)，使其变为全局向量。

14 != 6：计算总 Loss 时，Action Loss 是 [B, T, D]，Recon Loss 是 [B*K]，无法相加。

修正：分别对其 .mean() 取标量后再相加。

小结

已验证：数据通路畅通，梯度回传逻辑正确，环境配置脚本化。

在架构上，我通过改写 Pi0 源码实现了中间层特征的截获，让 Recon 任务能直接优化 Backbone。
在工程上，我解决了多视角数据的维度对齐问题，并封装了 setup_env.sh 脚本，解决了环境依赖和 transformers 补丁问题。
在验证上，我通过 Dummy 模式跑通了全流程（Loss 计算正常），目前代码已经由 Mock 数据和轻量模型验证通过。

待完成：

数据侧：目前使用 Mock 数据，需接入真实的 LeRobot/Aloha 数据集（需包含 3 个视角）。
算力侧：需在服务器上关闭 Dummy 模式，加载真实权重进行训练。
监控侧：接入可视化观察 Action Loss 和 Recon Loss 的下降曲线，验证辅助任务的有效性。