源码解读:拿下顶会最佳论文的重建式VLA,是如何实现的!-夜雨聆风

源码解读:拿下顶会最佳论文的重建式VLA,是如何实现的!

“如果模型能重建它，就说明它真正注意到了它”

——源码级解析

ReconVLA 作为 AAAI 2026 最佳论文，提出了基于重建机制的VLA模型，为具身智能中复杂场景下的细粒度感知与稳健动作生成提供了新的技术范式。

此前针对该工作的解读多集中于论文整体框架、实验结论与核心思想概述，较少触及模型内部的模块实现、代码逻辑与训练细节。

因此，本文打算换个视角，从任务特性出发，结合架构实现与工程实践，对 ReconVLA 展开一次面向代码与底层机制的深度解析。

我们开设此账号，除了想要向各位对【具身智能】感兴趣的人传递前沿权威的知识讯息外，也想和大家一起见证它到底是泡沫还是又一场热浪？

欢迎关注【深蓝具身智能】👇

问题的起点：为什么 VLA 需要“重建”？

现有的 VLA 模型（如 RT-2、OpenVLA）通常采用“编码器-投影器-LLM”的标准架构。

这种架构虽然能很好地利用 LLM 的推理能力，但也带来了副作用：视觉空间信息的丢失。

为了对齐语言特征，高维的视觉特征被严重压缩。当面对小目标或复杂背景时，LLM 的注意力往往会偏离真正的交互区域（Gaze Region）。

▲图1 | 注意力涣散。©【深蓝具身智能】编译

上行热力图揭示了传统 VLA 的致命缺陷——在”堆叠积木”这类长时序任务中，模型的视觉注意力始终“漫散”在整个场景，而非精准锁定当前操作目标；

中行红框展示了 ReconVLA 动态追踪的 Gaze Region，随任务进展自适应切换目标物体；

下行则是机器人的真实观测视角。

这组对比，正是 ReconVLA 提出重建辅助监督的出发点。

ReconVLA 的解法是：在 LLM 的输出端增加一个重建损失（Reconstruction Loss）。

模型不仅要输出动作，还要用它当前的隐状态去“复原”目标区域的图像——这就逼着模型在隐状态中保留足够的视觉细节。

▲图2 | 无需任何显式标注，仅凭重建 Token 驱动扩散去噪，让 VLA 在”看不见的监督”下学会聚焦目标。©【深蓝具身智能】编译

系统架构总览

我们先来看 ReconVLA 的整体架构。它在标准 VLA 的基础上，增加了一个并行的重建分支。

▲图3 | ReconVLA 的”双轨并行”。©【深蓝具身智能】编译

左侧重建分支（Recon. Part）：Gaze 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, eps=1e-6, origin_text=None):    batch_size = hidden_states.shape[0]    vm_loss_mask = torch.zeros((batch_size,), device=hidden_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]),        dtype=hidden_states.dtype, device=hidden_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, ...), mode='bilinear')    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', dtype=torch.float32):        # ⑤ mm_inv_projector（DiT 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', h=h, w=w)        # repeat(4) 对应 DiT 训练时的 classifier-free guidance 数据增强        vm_loss = self.model.mm_inv_projector(            z=image_hidden_states.repeat(4, 1, 1, 1).contiguous().float(),            target=z_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_loss

Step1：多模态输入构建

将多视角 RGB 图像通过冻结的 SigLIP 视觉编码器，再经过 MLP Projector 投影到 LLM 的维度。

Step2：序列拼接

构建 [System] [Image Tokens] [Instruction] [Recon Tokens] [Action Tokens] 的输入序列。

Step3：LLM 推理

将序列送入 Qwen2-7b 模型，得到完整的隐状态 hidden_states。

Step4：动作输出

LLM 的 lm_head 直接输出动作 Token 的 Logits。

Step5：重建条件提取

从 hidden_states 中，切片提取出重建 Token（Recon Tokens）位置对应的隐状态（在代码中通过 boi_ids 和 eoi_ids 定位）。这段隐状态就是论文中提到的 hᵣ。

Step6：扩散去噪

将 hᵣ 作为条件，指导 DiT 模型从纯噪声中还原出目标图像的 Latent 特征。

说明：

hᵣ 对应于 LLM 序列中专用 Recon Tokens 位置的输出隐状态（与图4紫色区域一致）。如需确认，可对照 recon_arch.py 中 boi_ids / eoi_ids 变量的实际定位范围。

核心技术一：DiT 扩散去噪与 adaLN-Zero 条件注入

ReconVLA 的重建分支并没有直接生成像素，而是生成了 VAE 的 Latent 特征。

具体来说，它使用了冻结的 Flux VAE 将目标图像编码为连续的 Latent z₀，然后训练一个 DiT（Diffusion Transformer）来进行去噪。

▲图4 | 双分支如何”共享一个大脑”。©【深蓝具身智能】编译

本图以数据流视角重绘了 ReconVLA 的完整前向过程：

右侧动作分支中，LLM Token 序列依次为：

System（蓝）→ Image（红，×N）→ Instruction（黄）→ Recon（紫）→ Action（绿，×n）

LLM 输出的紫色 hᵣ 跨越两个分支，以”Condition”身份注入左侧 Diffusion Denoiser，完成从语义理解到像素重建的信息传递。（火焰图标标注可训练模块，雪花图标标注冻结模块）。

在 denoiser_dit.py 中，我们可以看到 DiT 是如何将 LLM 的隐状态 hᵣ 作为条件注入的。

与传统的 Cross-Attention 不同，ReconVLA 使用了更高效的 adaLN-Zero（自适应层归一化）机制。

条件的构建与融合

# denoiser_dit.py -> DiT.forwardx = 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.forwardshift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=-1)# 调制 Attentionx = x + gate_msa * self.attn(modulate(self.norm1(x), shift_msa, scale_msa))# 调制 MLPx = x + gate_mlp * self.mlp(modulate(self.norm2(x), shift_ml

代码中的 adaLN_modulation 的最后一层被初始化为全 0。

这意味着在训练初期，所有的 shift、scale 和 gate 都是 0，DiT Block 退化为一个恒等映射（Identity Transform）。

这种 Zero-init 策略极大地稳定了训练初期的梯度。

核心技术二：动作离散化与自回归生成

作为一个 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 Block：c 经 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.pyToken ID = vocab_size − num_bins + bin_idx

在推理时，LLM 以自回归的方式依次预测出 7 个动作 Token，然后再通过查找 Bin Center（箱体中心值）并反归一化，还原为物理动作。

工程落地：两阶段训练流程

为了让模型既能学习到通用的视觉-动作表征，又能适应具体的下游任务，ReconVLA 设计了严谨的两阶段训练pipeline。

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", encoding="utf-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", encoding="utf-8") as file:            list_data_dict.extend(json.load(file))# 混合后全局随机打乱random.shuffle(list_data_dict)rank0_print(f"Loaded a total of {len(list_data_dict)} example

阶段二：任务特定微调（train_vla.py）

在微调阶段，模型专注于单一的下游任务（如 CALVIN 操控任务）。

此时，所有的数据都必须包含精确裁剪的 Gaze Region 图像作为 Target Image。

损失函数为：

# recon_trainer.py - ReconTrainer.compute_lossdef compute_loss(self, model, inputs, return_outputs=False, *args, **kwargs):    # 调用父类 Trainer.compute_loss，获取 total_loss 和完整 outputs    loss, outputs = super().compute_loss(model, inputs, return_outputs=True)    # 当重建分支激活时，分别记录 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 loss

ReconVLA官方的 README.md 提供了非常清晰的步骤。

conda create -n reconvla python=3.10pip install -r recon_requirements.txt# 下载预训练的 Flux VAE 和 Qwen2-7B 权重启动微调训练：Bashtorchrun --nnodes=1 --nproc_per_node=8 --master_port=29505 \    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