源码解读：拿下顶会最佳论文的重建式VLA，是如何实现的！

“如果模型能重建它就说明它真正注意到了它”——源码级解析目录01 问题的起点为什么 VLA 需要“重建”02 系统架构总览03 核心技术一DiT 扩散去噪与 adaLN-Zero 条件注入条件的构建与融合adaLN-Zero 注入逻辑04 核心技术二动作离散化与自回归生成归一化与分箱映射到词表末尾05 工程落地两阶段训练流程阶段一多数据源预训练pre_train_vla_action.py阶段二任务特定微调train_vla.py06 总结ReconVLA 作为 AAAI 2026 最佳论文提出了基于重建机制的VLA模型为具身智能中复杂场景下的细粒度感知与稳健动作生成提供了新的技术范式。此前针对该工作的解读多集中于论文整体框架、实验结论与核心思想概述较少触及模型内部的模块实现、代码逻辑与训练细节。因此本文打算换个视角从任务特性出发结合架构实现与工程实践对 ReconVLA 展开一次面向代码与底层机制的深度解析。01 问题的起点为什么 VLA 需要“重建”现有的 VLA 模型如 RT-2、OpenVLA通常采用“编码器-投影器-LLM”的标准架构。这种架构虽然能很好地利用 LLM 的推理能力但也带来了副作用视觉空间信息的丢失。为了对齐语言特征高维的视觉特征被严重压缩。当面对小目标或复杂背景时LLM 的注意力往往会偏离真正的交互区域Gaze Region。▲图1 | 注意力涣散。上行热力图揭示了传统 VLA 的致命缺陷——在堆叠积木这类长时序任务中模型的视觉注意力始终“漫散”在整个场景而非精准锁定当前操作目标中行红框展示了 ReconVLA 动态追踪的 Gaze Region随任务进展自适应切换目标物体下行则是机器人的真实观测视角。这组对比正是 ReconVLA 提出重建辅助监督的出发点。ReconVLA 的解法是在 LLM 的输出端增加一个重建损失Reconstruction Loss。模型不仅要输出动作还要用它当前的隐状态去“复原”目标区域的图像——这就逼着模型在隐状态中保留足够的视觉细节。▲图2 | 无需任何显式标注仅凭重建 Token 驱动扩散去噪让 VLA 在看不见的监督下学会聚焦目标。02 系统架构总览我们先来看 ReconVLA 的整体架构。它在标准 VLA 的基础上增加了一个并行的重建分支。▲图3 | ReconVLA 的双轨并行。左侧重建分支Recon. PartGaze Region 图像经冻结的 Visual Tokenizer 编码为场景 Token z₀加噪后得到 zₜ再由可训练的 DiT Denoiser 以 hᵣ 为条件预测噪声右侧动作分支Action Part多视角图像与文本指令经 Vision Encoder 和 Textual Tokenizer 送入 LLM输出四类隐状态——hᵢ图像、h_S场景、hᵣ重建同时作为左侧扩散的条件、h_A动作。两条分支共享同一个 LLM重建任务的梯度信号反向传播迫使 LLM 学会看准目标。从代码实现recon_arch.py来看整个前向传播过程非常清晰# recon_arch.py - compute_vm_loss完整核心逻辑 def compute_vm_loss(self, images, hidden_states, boi_ids, eoi_ids, eps1e-6, origin_textNone): batch_size hidden_states.shape[0] vm_loss_mask torch.zeros((batch_size,), devicehidden_states.device).bool() # ① 从 LLM 输出的 hidden_states 中按 boi/eoi 位置切片提取图像 Token 的隐状态 # 这就是论文中的 h_R —— 并非独立输入而是图像 Token 在 LLM 输出端的隐状态 image_hidden_states torch.zeros( (batch_size, self.model.image_embed_len, hidden_states.shape[-1]), dtypehidden_states.dtype, devicehidden_states.device ) for batch_index, (cur_boi_id, cur_eoi_id, cur_hidden_state) in enumerate( zip(boi_ids, eoi_ids, hidden_states) ): if (cur_boi_id is not None) and (cur_eoi_id is not None): assert cur_eoi_id - cur_boi_id 1 self.model.image_embed_len # 关键切片h_R hidden_states[boi_id : eoi_id 1] image_hidden_states[batch_index] cur_hidden_state[cur_boi_id: cur_eoi_id 1] vm_loss_mask[batch_index] True # ② 对目标图像Gaze Region进行预处理转换到 VAE 输入范围 [-1, 1] images_std torch.tensor(self.config.image_std, ...).view(1, -1, 1, 1) images_mean torch.tensor(self.config.image_mean, ...).view(1, -1, 1, 1) images_vae ((images * images_std images_mean - 0.5) / 0.5).clamp(-1., 1.) images_vae F.interpolate(images_vae, size(self.config.decode_image_size, ...), modebilinear) with torch.no_grad(): # ③ 冻结的 Flux VAE Encoder 将目标图像编码为 Latent z_0 posterior self.model.pixel_decoder.encode(images_vae).latent_dist z_q (posterior.sample() - self.model.pixel_decoder.shift_factor) \ * self.model.pixel_decoder.scaling_factor # ④ 2×2 窗口分组将 z_0 从 [B, 4, H/8, W/8] 重排为 [B, 16, H/16, W/16] z_q z_q.unfold(2, 2, 2).unfold(3, 2, 2) z_q rearrange(z_q, b c h w p1 p2 - b (c p1 p2) h w).contiguous() with torch.amp.autocast(cuda, dtypetorch.float32): # ⑤ mm_inv_projectorDiT Denoiser以 h_R 为条件计算扩散损失 # 注意image_hidden_states 先经 ln_pre 归一化再 reshape 为空间特征图 image_hidden_states self.model.mm_inv_projector.ln_pre(image_hidden_states) h w int(image_hidden_states.shape[1] ** 0.5) image_hidden_states rearrange(image_hidden_states, b (h w) c - b c h w, hh, ww) # repeat(4) 对应 DiT 训练时的 classifier-free guidance 数据增强 vm_loss self.model.mm_inv_projector( zimage_hidden_states.repeat(4, 1, 1, 1).contiguous().float(), targetz_q.repeat(4, 1, 1, 1).contiguous().float(), ) # ⑥ 用 vm_loss_mask 过滤无效样本target_image 缺失的样本不参与损失计算 vm_loss vm_loss.float() vm_loss_mask vm_loss_mask.repeat(4) vm_loss (vm_loss.view(batch_size, -1).mean() * vm_loss_mask).sum() \ / (vm_loss_mask.sum() eps) return vm_lossStep1多模态输入构建将多视角 RGB 图像通过冻结的 SigLIP 视觉编码器再经过 MLP Projector 投影到 LLM 的维度。Step2序列拼接构建 [System] [Image Tokens] [Instruction] [Recon Tokens] [Action Tokens] 的输入序列。Step3LLM 推理将序列送入 Qwen2-7b 模型得到完整的隐状态 hidden_states。Step4动作输出LLM 的 lm_head 直接输出动作 Token 的 Logits。Step5重建条件提取从 hidden_states 中切片提取出重建 TokenRecon Tokens位置对应的隐状态在代码中通过 boi_ids 和 eoi_ids 定位。这段隐状态就是论文中提到的 hᵣ。Step6扩散去噪将 hᵣ 作为条件指导 DiT 模型从纯噪声中还原出目标图像的 Latent 特征。说明hᵣ 对应于 LLM 序列中专用 Recon Tokens 位置的输出隐状态与图4紫色区域一致。如需确认可对照 recon_arch.py 中 boi_ids / eoi_ids 变量的实际定位范围。03 核心技术一DiT 扩散去噪与 adaLN-Zero 条件注入ReconVLA 的重建分支并没有直接生成像素而是生成了 VAE 的 Latent 特征。具体来说它使用了冻结的 Flux VAE 将目标图像编码为连续的 Latent z₀然后训练一个 DiTDiffusion Transformer来进行去噪。▲图4 | 双分支如何共享一个大脑。本图以数据流视角重绘了 ReconVLA 的完整前向过程右侧动作分支中LLM Token 序列依次为System蓝→ Image红×N→ Instruction黄→ Recon紫→ Action绿×nLLM 输出的紫色 hᵣ 跨越两个分支以Condition身份注入左侧 Diffusion Denoiser完成从语义理解到像素重建的信息传递。火焰图标标注可训练模块雪花图标标注冻结模块。在 denoiser_dit.py 中我们可以看到 DiT 是如何将 LLM 的隐状态 hᵣ 作为条件注入的。与传统的 Cross-Attention 不同ReconVLA 使用了更高效的 adaLN-Zero自适应层归一化机制。条件的构建与融合# denoiser_dit.py - DiT.forward x self.x_embedder(x) self.pos_embed # (N, T, D) t self.t_embedder(t) # (N, D) # h_R (context) 经过 reshape 和线性层映射 z rearrange(context, b c h w - b (h w) c).contiguous() z self.z_embedder(z) # (N, T, D) # 核心时间步嵌入与 h_R 嵌入直接相加 c t.unsqueeze(1) z # (N, T, D)这里有一个非常巧妙的设计DiT 的条件 c 是由时间步嵌入Timestep Embedding和 hᵣ 的线性映射直接元素级相加Element-wise Add得到的而不是拼接或交叉注意力。adaLN-Zero 注入逻辑在每一个 DiT Block 中条件 c 被用来生成 6 个调制参数# denoiser_dit.py - DiTBlock.forward shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp self.adaLN_modulation(c).chunk(6, dim-1) # 调制 Attention x x gate_msa * self.attn(modulate(self.norm1(x), shift_msa, scale_msa)) # 调制 MLP x x gate_mlp * self.mlp(modulate(self.norm2(x), shift_ml代码中的 adaLN_modulation 的最后一层被初始化为全 0。这意味着在训练初期所有的 shift、scale 和 gate 都是 0DiT Block 退化为一个恒等映射Identity Transform。这种 Zero-init 策略极大地稳定了训练初期的梯度。04 核心技术二动作离散化与自回归生成作为一个 VLA 模型最终的输出必须是机器人的控制动作。ReconVLA 采用了动作离散化Action Tokenization策略将连续的物理动作映射为 LLM 词表中的离散 Token。▲图5 | 六把调音旋钮让 LLM 语义精准控制扩散去噪。左侧展示完整的训练与推理流程Gaze Region 经冻结的 Flux VAE Encoder 编码为 形状 z₀经 2×2 窗口分组后加噪得到 zₜ送入 L 层 DiT Block条件向量 c 由时间步嵌入 τₜ 与 hᵣ 元素级相加而非 cross-attention构成。右侧放大单个 adaLN-Zero Blockc 经 adaLN_modulation 线性层 chunk(6) 分裂出 6 个参数分别以 shift/scale 调制 LayerNorm、以 gate 缩放残差Zero-init 保证训练初期每个 Block 均为恒等变换。在 action_tokenizer.py 中动作的编解码过程如下归一化与分箱首先利用 statistics.yaml 中的统计数据将原始的 7 自由度动作3 维平移 3 维旋转 1 维夹爪归一化到 [-1, 1] 的区间。接着使用 np.digitize 将连续值落入预设的 Bins 中。代码支持两种分箱策略均匀分箱默认np.linspace(-1, 1, 256)简单直接。非均匀分箱可选在 0 附近微小动作设置密集的 Bins在两端设置稀疏的 Bins以提高精细操作的精度。映射到词表末尾为了不干扰 LLM 原有的语言能力ReconVLA 将动作 Token 映射到 Qwen2 词表的最后 256 个位置# action_tokenizer.py Token ID vocab_size − num_bins bin_idx在推理时LLM 以自回归的方式依次预测出 7 个动作 Token然后再通过查找 Bin Center箱体中心值并反归一化还原为物理动作。05 工程落地两阶段训练流程为了让模型既能学习到通用的视觉-动作表征又能适应具体的下游任务ReconVLA 设计了严谨的两阶段训练pipeline。▲图6 | 把连续动作塞进语言模型词表的三步魔法。Stage 1编码7 自由度连续动作向量先经 statistics.yaml 归一化至 [-1, 1]再经 np.digitize 离散化为 256 个分箱的索引——默认均匀分箱可选非均匀分箱在 0 附近密集采样以提升小幅运动精度最终映射到 Qwen2-7B 词表末尾 256 个 Token。Stage 2生成LLM 以自回归方式逐 Token 预测 7 个动作 Token。、Stage 3解码Token ID 反查分箱中心值再经反归一化还原为可执行的连续动作。阶段一多数据源预训练pre_train_vla_action.py在这个阶段模型会混合 BridgeData V2、LIBERO 和 CALVIN 三个开源数据集。缺失图像回退如果某些数据集没有预先裁剪好的目标区域图像Target Image代码会自动回退使用原始的输入图像作为重建目标。差异化学习率在 recon_trainer.py 中为 LLM Backbone 和 DiT Denoiser 设置了不同的学习率通常 DiT 的学习率会更高以加速重建分支的收敛。# pre_train_vla_action.py - LazySupervisedDataset.__init__多数据源混合 # DataArguments 支持传入多个数据路径和对应的采样比例 # data_path: List[str] -- 多个 JSON 数据文件路径 # data_proportions: List[float] -- 每个数据集的采样比例0.0 ~ 1.0 if data_args.data_proportions is not None: for path, proportion in zip(data_path, data_args.data_proportions): with megfile.smart_open(path, r, encodingutf-8) as file: source_data json.load(file) # 按比例截取count ceil(total * proportion) count_to_sample int(math.ceil(len(source_data) * proportion)) list_data_dict.extend(source_data[:count_to_sample]) rank0_print(f Available: {len(source_data)}, Sampling: {count_to_sample} ({proportion:.0%})) else: # 未指定比例时合并所有数据集的全量数据 for path in data_path: with megfile.smart_open(path, r, encodingutf-8) as file: list_data_dict.extend(json.load(file)) # 混合后全局随机打乱 random.shuffle(list_data_dict) rank0_print(fLoaded a total of {len(list_data_dict)} example阶段二任务特定微调train_vla.py在微调阶段模型专注于单一的下游任务如 CALVIN 操控任务。此时所有的数据都必须包含精确裁剪的 Gaze Region 图像作为 Target Image。损失函数为# recon_trainer.py - ReconTrainer.compute_loss def compute_loss(self, model, inputs, return_outputsFalse, *args, **kwargs): # 调用父类 Trainer.compute_loss获取 total_loss 和完整 outputs loss, outputs super().compute_loss(model, inputs, return_outputsTrue) # 当重建分支激活时分别记录 lm_loss 和 vm_loss 到 TensorBoard if outputs.get(vm_loss, None) is not None: assert outputs.get(lm_loss, None) is not None vm_loss outputs[vm_loss] lm_loss outputs[lm_loss] # 每 logging_steps 步记录一次避免频繁 I/O if self.state.global_step % (self.args.logging_steps * self.args.gradient_accumulation_steps) 0: self.log({vm_loss: round(vm_loss.item(), 4), lm_loss: round(lm_loss.item(), 4)}) return (loss, outputs) if return_outputs else lossReconVLA官方的 README.md 提供了非常清晰的步骤。conda create -n reconvla python3.10 pip install -r recon_requirements.txt # 下载预训练的 Flux VAE 和 Qwen2-7B 权重 启动微调训练 Bash torchrun --nnodes1 --nproc_per_node8 --master_port29505 \ reconvla/train_vla.py \ --model_name_or_path stage1_checkpoint \ --data_path calvin_data_path \ --recon_enable True \ --reconstruct_image_num 1 \ --action_stat reconvla/statistics.yaml \ --output_dir ./checkpoints/reconvla-calvin▲环境与数据准备06 总结ReconVLA 用一行极简的理念——“如果模型能重建它就说明它真正注意到了它”解决了 VLA 模型的注意力失焦问题通过本次深度阅读代码我们发现它的实现的确非常扎实没有凭空创造复杂的架构而是通过在 LLM 序列中引入专用 Recon Tokens将其输出隐状态作为扩散条件优雅地连通了语言理解与像素重建。引入 Flux VAE 和 DiT通过 adaLN-Zero 实现了高效的条件注入。动作离散化与词表映射逻辑清晰最大程度保护了 LLM 的原有权重。这种“生成式辅助感知”的思路不仅在 CALVIN 基准测试上取得了 SOTA也为未来具身智能大模型的设计提供了一个非常值得借鉴的范本。RefReconVLA: Reconstructive Vision-Language-Action Modelhttps://github.com/OpenHelix-Team/ReconVLA