技术博客
VLA教程原理训练

AAAI 2026 最佳论文-ReconVLA详解(三)--ReconVLA代码详解

coffee co.2026-08-04 10:41
AAAI 2026 最佳论文-ReconVLA详解(三)--ReconVLA代码详解

ReconVLA 项目代码详解:重建式视觉-语言-动作模型如何提升机器人感知能力

相信读者们在阅读前面两篇文章后,已经对文章有一个大概的了解,废话少说,我们现在来讲解项目代码。

ReconVLA 的核心贡献正在于它提出了一种重建式 VLA 范式,在动作预测之外引入目标区域重建任务,利用重建监督强化模型对操作相关视觉区域的关注。模型不仅学习生成动作 token,还要从内部视觉表征中恢复 gaze region 或 manipulated target 对应的图像区域。这样一来,视觉 grounding 不再只是隐含在动作损失中的副产物,而成为训练过程中被显式约束的能力。

从项目代码看,ReconVLA 的实现可以拆解为四个关键部分:

  • reconvla/train_vla.py:训练入口、数据读取、动作 token 化和 Trainer 初始化。

  • reconvla/action_tokenizer.py:连续机器人动作与语言模型 token 之间的映射。

  • reconvla/recon/model/recon_arch.py:多模态输入拼接、视觉 token 插入、重建损失计算。

  • reconvla/recon/model/language_model/recon_qwen.py:基于 Qwen2 扩展出的 ReconVLA CausalLM。


1. 从普通 VLA 到重建式 VLA

典型 VLA 模型通常采用如下范式:

instruction + image -> vision-language backbone -> action tokens

ReconVLA 在这条主链路之外增加了视觉重建分支:

instruction + image
        -> vision encoder
        -> multimodal projector
        -> Qwen2 hidden states
        -> action token prediction
        -> visual hidden states
        -> denoising reconstruction head
        -> target / gaze region reconstruction

其中,动作预测仍然是策略学习的核心目标;重建分支则承担感知约束的作用。训练时,模型需要利用视觉 token 对应的 hidden states 重建目标区域图像。如果内部表征没有充分保留与任务相关的空间和语义信息,重建损失就会变大。由此,模型被迫学习更加细粒度、更加任务相关的视觉表示。

在代码中,这一机制最终体现为 vm_loss。模型的总训练目标可以概括为:

loss = language_modeling_loss + visual_modeling_loss

language_modeling_loss 负责动作 token 预测,visual_modeling_loss 负责目标区域重建。两者共同塑造策略模型的行为能力和视觉感知能力。


2. 动作离散化:将连续控制问题纳入语言建模框架

机器人动作通常是连续向量,例如末端执行器位姿变化、夹爪开合状态等。ReconVLA 没有为动作单独设计连续回归头,而是将动作离散化为 token,使动作生成可以复用语言模型的 next-token prediction 机制。

相关实现位于 reconvla/action_tokenizer.py

class ActionTokenizer:
    def __init__(
        self,
        tokenizer,
        bins=256,
        min_action=-1,
        max_action=1,
        use_norm_bins=False,
    ):
        self.tokenizer = tokenizer
        self.n_bins = bins
        self.min_action = min_action
        self.max_action = max_action

        self.bins = self.get_bins(min_action, max_action, self.n_bins, use_norm_bins)
        self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0

        self.action_token_begin_idx = int(self.tokenizer.vocab_size - (self.n_bins + 1))

它的基本流程是:

  1. 将连续动作裁剪到指定区间,默认是 [-1, 1]

  2. 使用 np.digitize 将连续值映射到离散 bin。

  3. 将 bin index 映射到 tokenizer 词表末尾的 token id。

核心调用如下:

def __call__(self, action):
    action = np.clip(action, a_min=float(self.min_action), a_max=float(self.max_action))
    discretized_action = np.digitize(action, self.bins)

    real_action_token = list(self.tokenizer.vocab_size - discretized_action)
    decode_action_token = self.tokenizer.decode(real_action_token)
    return real_action_token, decode_action_token

因此,动作序列在训练时可以像普通文本一样进入 causal language modeling 流程。模型需要预测的不是连续数值,而是动作 bin 对应的 token。

推理阶段再将 token id 还原为连续动作:

def decode_token_ids_to_actions(self, action_token_ids):
    discretized_actions = self.tokenizer.vocab_size - action_token_ids
    discretized_actions = np.clip(
        discretized_actions - 1,
        a_min=0,
        a_max=self.bin_centers.shape[0] - 1,
    )
    return self.bin_centers[discretized_actions]

这种设计的价值在于统一建模接口。语言指令、图像占位符、机器人观测和动作输出都可以组织为同一个 autoregressive 序列,VLA 模型也因此可以最大限度复用大语言模型已有的序列建模能力。


3. 数据流:从图像、目标区域到动作 token

训练数据由 LazySupervisedDataset 懒加载。每个样本在 __getitem__ 中完成图像读取、目标图像读取和对话格式化。

关键代码如下:

if 'image' in sources[0]:
    image_file = self.list_data_dict[i]['image']
    image_folder = self.data_args.image_folder
    processor = self.data_args.image_processor

    image = Image.open(...).convert('RGB')
    image = processor.preprocess(image, return_tensors='pt')['pixel_values'][0]

    target_image_file = self.list_data_dict[i]['image_target']
    target_image_folder = self.data_args.target_image_folder
    target_image = Image.open(...).convert('RGB')
    target_image = processor.preprocess(target_image, return_tensors='pt')['pixel_values'][0]

这里需要区分两类图像:

  • image 是模型输入的当前观察图像。

  • target_image 是重建分支的监督信号,对应 gaze region 或操作目标区域。

随后,样本进入动作预处理流程:

data_dict = preprocess_action(
    sources,
    self.tokenizer,
    self.action_tokenizer,
    has_image=('image' in self.list_data_dict[i]),
    has_embody=('embody' in self.list_data_dict[i]),
)

当样本属于 embodied robot 数据时,代码会将 GPT 回复中的连续动作替换为离散动作 token,并将 human 输入中的机器人观测替换为观测 token:

if sentence["from"] == "gpt":
    if has_embody:
        real_action_token, sent_value = action_to_lang(sentence["value"], action_tokenizer)
else:
    if has_embody:
        sent_value_parts = sent_value.split("\n")
        real_obs_token, robot_obs = robot_obs_lang(sent_value_parts[-1], action_tokenizer)
        sent_value_parts[-1] = robot_obs

处理后的训练样本可以抽象为:

Human: <image>
       instruction
       robot observation tokens

Assistant: action tokens

这种格式使机器人控制任务能够自然转化为多模态对话建模问题。输入侧包含视觉观察、任务指令和机器人状态,输出侧则是离散化后的动作序列。


4. 图像 token 如何进入 Qwen2

ReconVLA 基于 Qwen2 构建语言模型主干,但 Qwen2 原生只能处理 token embedding。图像需要先经过视觉编码器,再通过 projector 映射到语言模型的 hidden size。

视觉模块初始化位于 ReconMetaModel.initialize_vision_modules

self.vision_tower = build_vision_tower(model_args)
self.image_embed_len = (
    self.vision_tower.config.image_size // self.vision_tower.config.patch_size
) ** 2

self.config.mm_hidden_size = vision_tower.hidden_size
self.mm_projector = build_vision_projector(self.config)

mm_projector 的构建逻辑在 reconvla/recon/model/multimodal_projector/builder.py

def build_vision_projector(config, delay_load=False, **kwargs):
    projector_type = getattr(config, 'mm_projector_type', 'linear')

    if projector_type == 'linear':
        return nn.Linear(config.mm_hidden_size, config.hidden_size)

    mlp_gelu_match = re.match(r'^mlp(\d+)x_gelu$', projector_type)
    if mlp_gelu_match:
        modules = [nn.Linear(config.mm_hidden_size, config.hidden_size)]
        for _ in range(1, mlp_depth):
            modules.append(nn.GELU())
            modules.append(nn.Linear(config.hidden_size, config.hidden_size))
        return nn.Sequential(*modules)

默认训练脚本中采用如下配置:

--vision_tower ./siglip-so400m-patch14-384
--mm_projector_type mlp2x_gelu
--mm_vision_select_layer -2

这意味着图像首先由 SigLIP 编码为 patch features,再通过两层 MLP 对齐到 Qwen2 的 hidden space。

真正完成图像 token 插入的是 prepare_inputs_labels_for_multimodal

image_features = self.encode_images(images)

image_position = torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0].tolist()
boi_ids[batch_idx] = image_position[0]
eoi_ids[batch_idx] = image_position[0] + image_features.shape[1] - 1

其中,IMAGE_TOKEN_INDEX 对应文本序列中的 <image> 占位符。模型会找到该占位符的位置,并将单个占位 token 替换为一整段视觉 patch embeddings:

cur_new_input_embeds.append(cur_input_embeds_no_im[i])

if i < num_images:
    cur_image_features = image_features[cur_image_idx]
    cur_new_input_embeds.append(cur_image_features)
    cur_new_labels.append(
        torch.full(
            (cur_image_features.shape[0],),
            IGNORE_INDEX,
            device=cur_labels.device,
            dtype=cur_labels.dtype,
        )
    )

这里有一个重要细节:视觉 patch 对应的 labels 被全部设置为 IGNORE_INDEX。因此,语言模型损失不会要求模型预测图像 token。图像 token 作为上下文参与注意力计算,但不作为文本生成目标。

与此同时,代码记录了视觉 token 在序列中的起止位置:

boi_ids = begin of image ids
eoi_ids = end of image ids

这两个位置会在重建损失中使用。它们标记了哪些 hidden states 对应视觉区域,使重建分支能够从语言模型内部表示中提取视觉信息。


5. Qwen2 前向过程:语言损失与视觉重建损失

ReconVLA 的语言模型定义在 reconvla/recon/model/language_model/recon_qwen.py

class ReconQwen2ForCausalLM(Qwen2ForCausalLM, ReconMetaForCausalLM):
    config_class = ReconConfig

    def __init__(self, config):
        super(Qwen2ForCausalLM, self).__init__(config)
        self.model = ReconQwen2Model(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

该类保留了 Qwen2 的 causal language modeling 能力,同时通过 ReconMetaForCausalLM 注入多模态输入处理与视觉重建逻辑。

在 forward 中,首先构造多模态 embedding:

(
    input_ids,
    position_ids,
    attention_mask,
    past_key_values,
    inputs_embeds,
    labels,
    boi_ids,
    eoi_ids,
    cache_position,
) = self.prepare_inputs_labels_for_multimodal(
    input_ids,
    position_ids,
    attention_mask,
    past_key_values,
    labels,
    images,
    image_sizes,
    cache_position,
    recon_return_flag=0,
)

随后进入 Qwen2 主干:

outputs = self.model(
    input_ids=input_ids,
    attention_mask=attention_mask,
    position_ids=position_ids,
    past_key_values=past_key_values,
    inputs_embeds=inputs_embeds,
    use_cache=use_cache,
    output_attentions=output_attentions,
    output_hidden_states=output_hidden_states,
    return_dict=return_dict,
    cache_position=cache_position,
)

hidden_states = outputs[0]
logits = self.lm_head(hidden_states)

语言建模损失采用标准 next-token cross entropy:

shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
loss_fct = nn.CrossEntropyLoss()
loss = loss_fct(
    shift_logits.view(-1, self.config.vocab_size),
    shift_labels.view(-1),
)

当重建分支启用时,forward 会额外计算 vm_loss

vm_loss = None
if self.training and getattr(self.config, 'recon_enable', False):
    if self.config.reconstruct_image_num == 1:
        vm_loss = self.compute_vm_loss(
            target_images,
            hidden_states,
            boi_ids,
            eoi_ids,
            eps,
            origin_text,
        )
    loss = loss + vm_loss

因此,ReconVLA 的训练目标不是单纯拟合动作 token,而是同时优化动作生成与视觉区域重建。前者保证策略输出,后者约束模型内部视觉表示,使其更贴近操作目标。


6. 重建损失:从语言模型 hidden states 中恢复目标区域

compute_vm_loss 位于 reconvla/recon/model/recon_arch.py。它首先根据 boi_ids/eoi_ids 提取视觉 token 对应的 hidden states:

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):
        image_hidden_states[batch_index] = cur_hidden_state[cur_boi_id: cur_eoi_id + 1]
        vm_loss_mask[batch_index] = True

这一设计值得注意:重建分支使用的不是原始视觉 encoder 输出,而是视觉 token 经过 Qwen2 多层注意力交互后的 hidden states。也就是说,重建任务约束的是“被语言和上下文调制过的视觉表示”。这正是 ReconVLA 实现 implicit grounding 的关键。

接下来,目标图像会经过冻结 VAE 编码到 latent 空间:

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, self.config.decode_image_size),
    mode='bilinear',
)

with torch.no_grad():
    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

pixel_decoder 来自 diffusers.AutoencoderKL,定义在 flux_decoder.py

class FluxDecoder(nn.Module):
    def __init__(self, config, **kwargs):
        super().__init__()
        self.pixel_decoder = AutoencoderKL.from_pretrained(config.mm_pixel_decoder)
        self.pixel_decoder.requires_grad_(False)
        self.pixel_decoder.float()
        self.pixel_decoder.eval()

VAE 在这里保持冻结。其作用是提供稳定的 latent 表示空间,使重建分支专注于学习从 LLM hidden states 到目标区域 latent 的条件映射,而不是重新学习像素级图像生成。

随后,视觉 hidden states 会被送入 inverse projector,即 mm_inv_projector

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,
).contiguous()

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(),
)

在默认配置中,mm_inv_projector 被设置为:

--mm_inv_projector_type denoiser_vit3x

这表明 ReconVLA 使用一个轻量级 DiT denoiser 作为重建头,而不是简单的线性投影。


7. DiT 重建头:条件去噪视角下的视觉监督

mm_inv_projector 的构建逻辑位于 reconvla/recon/model/multimodal_projector/builder.py

def build_inv_projector(config, delay_load=False, **kwargs):
    projector_type = getattr(config, 'mm_inv_projector_type', 'linear')

    if projector_type.startswith("denoiser"):
        vit_match = re.match(r'^denoiser_vit(\d+)x$', projector_type)
        depth = int(vit_match.group(1))

        return ReconDenoiser(
            x_channel=config.mm_inv_hidden_size,
            z_channel=config.hidden_size,
            embed_dim=width,
            depth=depth,
            timesteps='1000',
            learn_sigma=False,
            n_patches=config.image_embed_len,
        )

ReconDenoiser 内部采用 DiT 结构。它将带噪声的 VAE latent 作为输入,将 LLM 视觉 hidden states 作为条件,用扩散训练方式学习目标区域重建。

DiT 的 forward 如下:

def forward(self, x, t, context):
    x = self.x_embedder(x) + self.pos_embed
    t = self.t_embedder(t)

    z = rearrange(context, 'b c h w -> b (h w) c').contiguous()
    z = self.z_embedder(z)
    c = t.unsqueeze(1) + z

    for block in self.blocks:
        x = block(x, c)

    x = self.final_layer(x, c)
    x = self.unpatchify(x)
    return x

其中,扩散时间步 embedding t 与视觉上下文 z 共同构成条件 c。每个 DiT block 通过 adaLN-Zero 接收条件调制:

shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = \
    self.adaLN_modulation(c).chunk(6, dim=-1)

x = x + gate_msa * self.attn(modulate(self.norm1(x), shift_msa, scale_msa))
x = x + gate_mlp * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))

从训练目标看,该过程可以理解为:

noisy target latent z_t + visual hidden context -> denoising target

如果 Qwen2 内部的视觉 hidden states 没有捕获与任务相关的目标区域信息,denoiser 将难以完成重建,进而产生更高的 vm_loss。因此,重建任务通过梯度反向传播直接约束多模态主干,使其形成更可操作的视觉表示。


8. 训练配置:重建分支如何启用

默认预训练脚本 reconvla/scripts/train_vla/pretrain.sh 给出了关键配置:

torchrun --nproc-per-node=8 \
    pre_train_vla_action.py \
    --model_name_or_path ./qwen2-7b \
    --vision_tower ./siglip-so400m-patch14-384 \
    --version qwen_2 \
    --mm_pixel_decoder ./pretrained_vae \
    --reconstruct_image_num 1 \
    --mm_projector_type mlp2x_gelu \
    --mm_inv_projector_type denoiser_vit3x \
    --image_aspect_ratio pad \
    --model_max_length 32768 \
    --gradient_checkpointing True \
    --bf16 True

其中几个参数尤为重要:

  • --model_name_or_path ./qwen2-7b:语言模型主干。

  • --vision_tower ./siglip-so400m-patch14-384:视觉编码器。

  • --mm_projector_type mlp2x_gelu:将视觉特征映射到 Qwen2 hidden size。

  • --mm_pixel_decoder ./pretrained_vae:冻结 VAE,用于目标图像 latent 编码。

  • --mm_inv_projector_type denoiser_vit3x:基于 DiT 的重建头。

  • --reconstruct_image_num 1:启用单目标区域重建。

train() 中,视觉模块初始化后会写入重建相关配置:

model.config.reconstruct_image_num = model_args.reconstruct_image_num
model.config.reconstruct_image = data_args.reconstruct_image

而在 initialize_vision_modules 中,只要提供了 mm_pixel_decoder,重建分支就会被启用:

self.config.recon_enable = False
if getattr(model_args, 'mm_pixel_decoder', False):
    self.config.recon_enable = True
    self.pixel_decoder = build_pixel_decoder(self.config)
    self.config.mm_inv_hidden_size = self.pixel_decoder.latent_dim
    self.mm_inv_projector = build_inv_projector(self.config)

这说明 ReconVLA 的重建监督并不是独立后处理流程,而是直接嵌入主模型训练的 forward-loss 路径中,与动作 token 预测共同优化。


9. 推理阶段:重建分支的训练收益如何保留

推理阶段不再需要 target_image,也不会计算 vm_loss。模型只需完成标准 VLA 推理流程:

  1. 输入图像、语言指令和必要的机器人状态。

  2. 将图像编码为视觉 patch embeddings,并插入文本序列。

  3. 调用 generate() 生成动作 token。

  4. 将动作 token 反离散化为连续动作。

ReconQwen2ForCausalLM.generate 中仍会处理图像输入:

if images is not None:
    (
        inputs,
        position_ids,
        attention_mask,
        _,
        inputs_embeds,
        _,
        boi_ids,
        eoi_ids,
        cache_positions,
    ) = self.prepare_inputs_labels_for_multimodal(
        inputs,
        position_ids,
        attention_mask,
        None,
        None,
        images,
        image_sizes=image_sizes,
        recon_return_flag=1,
    )

尽管重建头在推理时不必参与动作生成,训练阶段的重建监督已经改变了主干模型的视觉表征学习方式。换言之,ReconVLA 将目标区域重建作为训练时的结构化约束,用以提升部署时的感知和操作泛化能力。


10. 总结:ReconVLA 的工程思想

从代码实现看,ReconVLA 可以概括为如下系统:

Qwen2 sequence backbone
    + SigLIP vision encoder
    + MLP multimodal projector
    + action tokenizer
    + frozen VAE
    + DiT reconstruction head

其中,Qwen2 负责统一的序列建模,SigLIP 提供视觉 patch 表征,动作 tokenizer 将连续控制变量转化为 token 序列,VAE 和 DiT 则构成目标区域重建分支。最终,模型通过 lm_loss + vm_loss 同时学习动作生成和视觉 grounding。

ReconVLA 的关键价值不在于简单增加一个辅助任务,而在于它将“机器人是否真正关注了任务相关区域”转化为可优化的训练目标。对于机器人操作而言,目标物体、接触区域、抓取边界和局部几何结构往往决定动作成败。通过重建 gaze region,ReconVLA 让这些关键视觉信息在模型内部表示中得到更强约束,从而提升精细操作和跨场景泛化能力。


附:建议的代码阅读顺序

如果希望继续深入该项目,建议按照以下路径阅读:

README.md
  -> reconvla/scripts/train_vla/pretrain.sh
  -> reconvla/train_vla.py
  -> reconvla/action_tokenizer.py
  -> reconvla/recon/model/language_model/recon_qwen.py
  -> reconvla/recon/model/recon_arch.py
  -> reconvla/recon/model/multimodal_projector/builder.py
  -> reconvla/recon/model/multimodal_denoiser/denoiser_dit.py
  -> reconvla/recon/model/pixel_decoder/flux_decoder.py

这一路径从训练脚本进入数据与动作表示,再进入模型 forward、视觉重建损失和 DiT 重建头,基本覆盖了 ReconVLA 从输入到损失函数的完整实现链路。

点赞收藏
// 评论0
0 / 500
还没有评论,快来抢沙发