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]。使用
np.digitize将连续值映射到离散 bin。将 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 推理流程:
输入图像、语言指令和必要的机器人状态。
将图像编码为视觉 patch embeddings,并插入文本序列。
调用
generate()生成动作 token。将动作 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 从输入到损失函数的完整实现链路。
